Embedding layers are commonly used to map discrete symbols into continuous embedding vectors that reflect their semantic meanings. Despite their effectiveness, the number of parameters in an embedding layer increases linearly with the number of symbols and poses a critical challenge on memory and storage constraints. In this work, we propose a generic and end-to-end learnable compression framework termed differentiable product quantization (DPQ). We present two instantiations of DPQ that leverage different approximation techniques to enable differentiability in end-to-end learning. Our method can readily serve as a drop-in alternative for any existing embedding layer. Empirically, DPQ offers significant compression ratios (14-238x) at negligible or no performance cost on 10 datasets across three different language tasks.
This paper works on methods for compressed embedding layers for low memory inference, where the compressed embedding are learned together with the task-specific models in a differentiable end-to-end fashion. The methods in the paper build on a slight variant of the K-way D-dimensional discrete code representation proposed by Chen et al.. Specifically, the two methods in the paper are motivated by the idea that the K-way D-dimensional code can be viewed as a form of product quantization. The first proposed method (DPQ-SX) uses softmax-based approximation to allow for differentiable learning, while the second proposed methods uses clustering centroid-based approximation. Empirically, the authors demonstrate that the proposed methods for generating compressed embedding can achieve matching inference accuracy with the corresponding uncompressed full embedding across 3 NLP tasks; these proposed approach can outperform the pre-trained word embedding and the K-way D-dimensional code baselines in language modeling tasks.
This paper builds on an existing embedding representation approach---K-way D-dimensional code. But I think the perspective on viewing k-way D-dim approach as product quantization (which motivate the differentiable learning approach in the paper) is very interesting. Also I think the empirical performance of the proposed method is promising. I gave weak rejection because 1) the proof of theorem 1 is flawed; 2) The experiment might need additional validation to fully support the merit of the proposed methods.
I list below the major concern / questions I have. I am happy to raise the score if the following questions are properly resolved in the rebuttal:
1. Correct me if I am wrong, I think the proof of theorem 1 is wrong---if the integer based code book C is full rank , it does not necessarily imply the one-hot vector based code book B is full rank. E.g. assume K = 2 D = 2,
C = [1 1;1 2;1 1; 1 2] is full-rank (rank = 2), but the corresponding B = [1 0 1 0 ; 1 0  0 1; 1 0 1 0 ; 1 0  0 1] is not full rank (rank < 4).
2. As the proposed methods advocate training the inference-oriented compressed embedding together with the task models (such as translation models), I think the following naive baseline is necessary to fully evaluate the merit of the proposed approach: one can train the full embedding with the task model as usual, compress the task-specific full embedding using the K-way D-dim approach by Chen et al. or using the deep compositional code learning approach by Shu et al., and then use it for the inference. This provides an alternative way to use product quantization based approach for embedding with low inference memory, without training together with task models. Without this, I can not evaluate the necessity to use the train-together approach the author proposed.
3. The proposed DPQ-SX approach performs better in the two proposed approaches. However this approach uses different K and V matrix where in the original K-way D-dim approach we have K = V. This makes it hard to say if the better performance is due to the decoupling of K and V, or because of the training method inspired from the product quantization perspective. It needs ablation study here.
4. In Table 4, the authors only compare to baselines on the LM task, I am wondering how it compares to the the baselines on the other two translation and text classification models.
For improving the paper, the relatively minor comments are as the following:
1. In equation 4, the partition function Z is not explicitly defined.
2. In the second paragraph of section 3.1, it is not clear what exactly is the pre-trained embedding used as baseline.
3. For better readability, it is better to inflate the caption of figure and tables to provide useful take-away message there.