Learning Compact Embedding Layers via Differentiable Product Quantization
Authors: Ting Chen
Lala Li
Yizhou Sun
Venue:  (2020)
Keywords: efficient modeling
compact embedding
embedding table compression
differentiable product quantization

Abstract: 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.

Under review a conference paper at ICLR 2020
LEARNING COMPACT EMBEDDING LAYERS VIA DIFFERENTIABLE PRODUCT QUANTIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
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 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 end-to-end learning. Our method readily serve as drop-in alternative for any existing embedding layer. Empirically, DPQ offers significant compression ratios (14-238x) negligible or no performance cost on 10 datasets across three different language tasks.
INTRODUCTION
The embedding layer is a basic neural network module which maps a discrete symbol/word into a continuous hidden vector. It is widely used in NLP related applications, including language modeling, machine translation and text classification. With large vocabulary sizes, embedding layers consume large amounts of storage and memory. For example, in the medium-sized LSTM-based model on the PTB dataset (Zaremba et al., 2014), the embedding table accounts for more than 95% of the total number of parameters. Even with sub-words encoding (e.g. Byte-pair encoding), the size of the embedding layer is still very significant. In addition to words/sub wordss models in the text domain (Mikolov al., 2013 Devlin al.) (2088), embedding layers are also used in wide range of applications such as knowledge graphs (Bordes et al., 2013 Socher et al. (2033) and recommender systems Koren et al. 2009), where the vocabulary sizes are even larger.
Recent efforts to reduce the size of embedding layers have been made (Chen et al. 2018b Shu and Nakayama 2017). where the authors proposed first learn to encode symbols/words with K K-way D-dimensional discrete codes (KD codes, such 5-1-2-4 for "cat" and 5-1-2-3 for "dog"), and then compose the codes to form the output symbol embedding. However, in Shu and Nakayama (2017). the discrete codes fixed before training and therefore non-adaptive and limited to downstream tasks. [Chen et af.] (2018b) proposes to learn codes in an end-to-end fashion which leads to better task performance. However, their method employs an expensive embedding composition function to turn KD codes into embedding vectors, and requires a distillation procedure which incorporates pre-trained embedding table guidance, in order match the performance of the full embedding baseline.
In this work, we propose novel differentiable product quantization (DPQ) framework. The proposal is based on the observation that the discrete codes (KD codes) are naturally derived through the process of quantization (product quantization by Jegou al. (2010) in particular). We also provide two concrete approximation techniques that allow differentiable learning. By making the quantization process differentiable, we are able to learn the KD codes in an end-to-end fashion. Compared to the existing methods (Chen al. 2018b 8b Shu and Nakayama 2017), our framework 1) brings new and general perspective on how the discrete codes can be obtained in differentiable manner; 2) allows more flexible model designs (e.g. distance functions and approximation algorithms), and 3) achieves better task performance as well as compression efficiency (by leveraging the sizes product keys and values) while avoiding the cumbersome distillation procedure.

---

Under review as a conference paper at ICLR 2020
We conduct experiments on ten different datasets across three tasks, by simply replacing the original embedding layer with DPQ. The results show that DPQ can learn compact discrete embeddings with higher compression ratios than the existing methods, at the same time achieving the same performance results are obtained from end-to-end training where
as the original full embeddings. Furthermore, no extra procedures such as distillation are required. To the best of our knowledge, this is the first work to train compact discrete embeddings in end-to-end fashion without distillation.
2 METHOD
Problem setup. An embedding function can be defined as Fw : : V RR where V denotes the vocabulary of discrete symbols, and W E R is the embedding table with n = [V.. In standard end-to-end training, the embedding function is jointly trained with other neural net parameters to optimize given objective. The goal of this work is to learn a compact embedding function Fw in the same end-to-end fashion, but the number of bits used for the new parameterization W is substantially smaller than the original full embedding table W.
Motivation. To represent the embedding table in a more compact way, we can first associate each symbol with K-way D-dimensional discrete code (KD code), and then use an embedding composition function that turns the KD code into a continuous embedding vector (Chen et al. 2018b). However, it is not clear where the discrete KD codes come from. One could directly optimize them as free parameters, but is both ad-hoc and restrictive. Our key insight in this work is that discrete codes are naturally derived from the process of quantization (product quantization (Jegou et al. 2010) in particular) of a continuous space. It is flexible to specify the quantization process in various ways, and by making this quantization process differentiable, enable end-to-er learning of discrete codes via optimizing some task-specific objective.
2.1 DIFFERENTIABLE PRODUCTION QUANTIZATION FRAMEWORK
The proposed differentiable production quantization (DPQ) function is mapping between continuous spaces, i.e. T Rdd + Rdd In between the two continuous spaces, there is a discrete space (1..... KJD which can be seen as discrete bottleneck. To transform from continuous space to discrete space and back, two major functions are used: 1) a discretization function ((-)= R + {1.... K] that maps a continuous vector into a K-way D-dimensional discrete code (KD code). and 2) divrrs-discrttiiation function p(-): {1.... K] -, IRd that maps the KD code into continuous embedding vector. other words, the general DPQ mapping is T(-) = p ((-).
Compact embedding layer via DPQ. In order to obtain a compact embedding layer, we first take a raw embedding and put through DPQ function. More specifically, the raw embedding matrix can be presented as a Query matrix Q E Rndd where the number of rows equals to the vocabulary size. The discretization function of DPQ computes discrete codes C Φ Q) where C E (1.... K) n is the KD codebook. To construct the final embedding table for all symbols, the reverse diverse-diccretizatin function of DPQ is applied, i.e. H p(C) where H E Rn is the final symbol embedding matrix. In order to make it compact for the inference, we will discard the original embedding matrix Q and only store the codebook C and small parameters needed the exverseddisreetizatio function. They sufficient to (re (re)construct partial whole embedding table. In below, we specify the discretization function (-) and evvess-diicrretzattii function p(-) via product keys and values.
Product keys for discretization function ((.). Given the query matrix Q the discretization function computes the KD codebook C. While is possible to use complicated transformation, in order to make it efficient, we simply leverage Key matrix K E R&K xd with K rows where K is the number of choices for each code bit. In the spirit of product keys in product quantization, we further split columns of K and Q into D groups/subspace, such that KU) Kxd/D and QU) Daxd/D
We can compute each of D dimensional KD codes separately. The j-th dimension of a KD code C₁ for the i-th symbol is computed as follows.
cf arg min dist Q.") K!
(1)
k
2

---

Under review as a conference paper at ICLR 2020
Training only
Groups
D Groups
Key Matrix
Value Matrix
Groups
Query Matrix
Result Matrix
Codebook
(Raw Embedding
[Constructed Embedding Table)
Table)
Discretization (Eq. (1))
Index Concateration (Eq. (2))
Figure 1: The DPQ embedding framework. During training, differentiable product quantization is used to approximate the raw embedding table (i.e. the Query Matrix). At inference, only the codebook C E {I, KjnXD and the Value matrix V I xd are needed to construct the embedding table.
The -ist(..-) computes distance measure between two vectors, and use it to decide which discrete code to take.
Product values for reverse-diseretizatin function p(.). Given the codebook C, the reverse- discretization function computes the final continuous embedding vectors, While this can be another sophisticated transformation, we again opt for the most efficient design and employee a single Value matrix V IRK kd the parameter. Similarly, we leverage product keys, and split the columns of V into D groups/subspaces the same way as K and Q i.e. vû) E R k We use the code in each of D dimension to index the subspace in V, and concatenate the results to form the final embedding vector as follows.
H, (1)
V
( c (D)
We note that this is simplification. both conceptually and computationally. of the ones used in (Chen jet al. 2018b Shu and Nakayama, 2017), which reduces the computation overhead and eases the optimization.
Figure [ iilluttrtte the proposed framework. The proposed method can also seen as a learned hash function of finite input into a set of KD codes, and use lookup during the inference instead of re-compute the codes.
Storage complexity. Assuming the default 32-bit floating point is used, the original full embedding table requires 32nd bits. As for DPQ embedding, we only need to store the codebook and the Value matrix: 1) codebook C requires nD log, K bits, which is the only thing that depends on vocabulary size n, and 2) Value matrix V requires 32K bits which does not explicitly depend on n and is ignoble when is large. Since typically nDD log, K 32nd, the DPQ embedding is more compact.
Inference complexity. Since only indexing and concatenation (Eq. 22 are used during inference, both the extra computation complexity and memory footprint usually negligible compared to the regular full embedding (which directly indexes an embedding table).
Expressiveness. Although the DPQ embedding is compact than full embedding, it is not achieved by reducing the rank of the matrix (as in traditional low-rank factorization). Instead, introduces sparsity into the embedding matrix in two axis: (1) the product keys/values, and (2) top-1 selection in each group/subspace.
Theorem 1. The DPQ embedding matrix H is full rank given the following constraints are satisfied.
1) One-hot encoded C (I,..., K) denoted as B (0, 1J full-rankk
2) Sub-matrices of splitted V, i.e. v e Rxd//D Nj, are all full-rank,
32Kd/D D bits if we share the weights among 1) groups/subspaces.
3

---

Under review conference paper at ICLR 2020
3) KD2 2d d
The proof is given in the appendix [B] Note that it is easy to keep H full-rank while achieving good compression ratio, since is easy to achieve nD logz K < 32nd with KD d.
So far we have not specified some designs of the discretization function such as the distance function in Eq [ More importantly, how can we compute gradients through the arg min function in Eq. 7 While there could be many instantiations with different design choices, below we introduce two DPQ instantiations that use two different approximation schemes.
2.2 SOFTMAX-BASHE APPROXIMATION
The first instantiation of DPQ (named DPQ-SX) approximates the non-differentiable arg max opera- tion with differentiable softmax function. To do so, we first specify the distance function in Eq. with softmax function as follows.
exp((Q))) ( K()
c = arg max
Ex (xp((Q) K()
&
where (:, denotes dot product of two vectors (alternatively, other metrics such as Euclidean distance, cosine distance can also be used). To approximate the arg max, similar to (Chen et 2018b 8b Jang et al. 2016), we relax the softmax function with temperature T:
c exp((Q,)) K k////>//2
where Z = En e e
K 1/T). Note that now 0 A is a probabilistic vector (i.e. soft
one-hot vector) instead of an integer c. And meehoo(( (c(²)) ~ c(a)) or c. = arg max C(U)
With a one-hot code relaxed into soft one-hot vector, can replace index operation vūj with dot (
product to compute the output embedding vector, i.e. H v())
The softmax approximated computation defined above is fully differentiable when r * O. However, to compute discrete codes during the forward pass, we have to set r - 0, which turns the softmax function into spike concentrated on the ( -th dimension. This is equivalent to the arg max operation which does not have gradient.
To enable pseudo gradient while still be able to output discrete codes, we use different temperatures during forward and backward pass, i.e. set T -> in forward pass, and r - the backward pass. So the final DPQ function can be expressed as follows.
H. = T(Q = - sg T(Q/T= =)) - -(Qal (orm0)) 0))
(5)
Where sg is the stop gradient operator, which is identity function in forward pass, but drops gradient for variables inside it during the backward pass.
2.3 CENTROID-BASED APPROXIMATION
The second instantiation of DPQ (named DPQ-VQ) uses centroid-bard based approximation, which directly pass the gradient straight-through (Bengio et al. 2013) small set of centroids. In order to do so, we need to put Q K, V into the same space.
First, we treat rows in Key matrix K as centroids, and use them to approximate Query matrix Q. The approximation is based on the Euclidean distance as follows.
arg min |QQ K (
(6)
k
Secondly, we tie the Key and Value matrices, ie. = K, so that we can pass the gradient through.
We still have the non-differentiable arg min operation, and the input query Q (J) are different from
selected output centroid (5) c( However, since they are in the same space, it allows us to directly pass the gradient ttroigh as follows.
H; Qi -sg(Q; T(Q.))
4

---

Under review conference paper at ICLR 2020
Forward poss:
the diesest key to approximate the query
Forward
Backward solmax approximate
2
Backward pass:
original query
K 10
2,
1,2,3,
K
(a) Softmax-based (DPQ-SX)
(b) Centroid-based (DPQ-VQ)
Figure 2: Illustration of two types of approximation to enable differentiaaility in DPQ.
Table Summary differences between VQ and SX. DPQ-SX allows more flexibility in distance
metrics and whether to tie the Key and Value metrices. DPQ-VQ is more efficient during training and therefore is more scalable to larger K, D values.
Method
Dist. Metric
Key/Value matrices
Train
Inference
DPQ-SX
Dot product and more
Not tied, allows different sizes
Efficient
Efficient
DPQ-VQ
Euclidean only
Tied
More efficient
Efficient
Where sg is again the stop gradient operation. During the forward pass, the selected centroid is emitted, but during the backward pass, the gradient pass the query directly. This provides way to compute discrete codes in the forward pass (which are the indexes of the centroids), and update the Query matrix during the backward pass.
However, it is worth noting that the Eq. 7 only approximates gradient for Query matrix, but does the tied Key/Value matrix. Similar van den Oord al. (2017),
updates the centroids, add a regularization term: Lreg Σ, (IT)),) sg Q,) which makes entries of the Key/ Value matrix arithmetic mean of their members. Altemntively, one can also use Exponential Moving Average (Kaiser et al. 2018) to update the centroids.
comparison between DPQ-SX and DPQ-VQ. DPQ-VQ and DPQ-SX only differ during train- ing. They very different in how they approximate the gradient for the non-differentiable arg min function: DPQ-SX approximates the one-hot vector with softmax, while DPQ-VQ approximates the continuous vector using a set of centroids. Figure Dillustrates this difference. This suggests that when there is large gap between one-hot and probabilistic vectors (large K), DPQ-SX approximation could be poor; and when there is large gap between the continuous vector and the selected centroid (large subspace dimension, i.e. small D), DPQ-VQ could have big approximation error.
Table [ ssummarizes the comparisons between DPQ-SX and DPQ-VQ. DPQ-SX is more flexible does not constrain the distance metric, nor does it tie the Key/Value matrices in DPQ-VQ. Thus one could use different sizes of Key and Value matrices. Regarding to the computational cost during training, DPQ-SX propagarsss through the whole distribution of K choices, while DPQ-VQ only back propagatss through the nearest centroid, making it more scalable (to large K, D, and batch sizes).
EXPERIMENTS
We conduct experiments on ten datasets across three tasks: language modeling (LM), neural machine translation (NMT) and text classification (TextC) We adopt existing architectures for these tasks as base models and only replace the input embedding layer with DPQ embeddings. The details of datasets and base models are summarized in T
For the five text classification datasets Zhang et al. 2015). Yahoo! answers and AG news represent topic prediction, Yelp Polarity and Yelp Full represent sentiment analysis, and DBpedia represents ontology classification.
5

---

Under review conference paper at ICLR 2020
Table Datasets and models used in our experiments. More details in Appendix C
Task
Dataset
Vocab Size
Tokenization
Base Model
PTB
10,000
LSTM-based models from Zaremba ct al. 2014), three model sizes
LM
Words
Wikitext-2
33,278
IWSLT15 (En-Vi)
17.191
Seg2seq-based model from Luong al. 2017)
Words
NMT
IWSLTI5 (Vi-En)
7,709
WMT19 (En-De)
32,000
Sub-words
Transformer Base in Vaswani et al 2017
AG News
69,322
Yahoo! Ans.
477,522
One hidden layer after mean pooling word vectors, similar to fast from Joulin et al. (0017)
TextC
DBpedia Yelp P
612,530
Words
246,739
Yelp F
268,414
Table 3: Comparisons of DPQ variants vs. the full embedding baselines.
Task
Metric
Dataset
Baseline
DPQ-SX
(CR)
DPQ-VQ (CR)
PTB
83.38
83.17
(163.2)
83.27
(58.67)
LM
PPL
Wikitext-2
95.61
94.94
(59.25)
95.92
(95.25)
IWSLTIS (En-Vi)
25,4
25.3
(86.17)
25.3
(16.13)
NMT
BLEU
IWSLTIS (Vi-En)
23.0
(72.00)
22.5
(14.05)
WMT19 (En-De)
38.8
38.8
(18.00)
38.7
(18.23)
AG News
92.59
92.49
(19.26)
92.55
(23.95)
Yahoo! Ans.
69.41
69.62
(48,16)
69.15
(19.24)
TexIC
Acc(%)
DBpedia
98.12
98.13
(24.08)
98.14
(38.45)
Yelp
93,92
94.17
(38.52)
93.91
(24.04)
Yelp
60.33
60.10
(48.16)
60.22
(24.05)
We evaluate the models using two metrics: task performance and compression ratio. Task performance metrics perplexity for LM tasks, BLEU scores for NMT tasks, and accuracy in TextC tasks.
Compression ratios for the embedding layer is computed as follows:
# of bits used in the full embedding table
CR of bits used in the compressed model during inference
For DPQ in particular, this can be computed as CR 177 K K2 Further compression
32nd
can be achieved with subbspacessharigg as described in Appendix
With subspace-sharing,
CR DD loga K K+32Kd/D
32nd
3.1 COMPRESSION RATIOS AND TASK PERFORMANCE AGAINST BASELINES
Table 3 jummarizes the task performance and compression ratios of DPQ-SX and DPQ-VQ against baseline models that use the regular full embeddings In each task/dataset, we report results from configuration that gives as good task performance as the baseline (or as good as possible, if it does not match with the baseline) while providing largest compression ratio. In all tasks, both DPQ-SX and DPQ-VQ achieve comparable better task performance while providing compression ratio from 14x to 163 x. In 6 out of 10 datasets, DPQ-SX performs strictly better than DPQ- VQ both metrics. Remarkably, DPQ is able to further compress the already-compact sub- word representations. This shows great potential of DPQ to learn very compact embedding layers.
We also compare DPQ against the following recently proposed embedding compression meth- ods (Chen et al. 2018h Shu and Nakayama 2017). Pre-train: a three-step procedure where one firstly trains full model, secondly learns discrete codes reconstruct the pre-trained embedding
For LMM results from the medium-sized LST model.

---

Under review conference paper at ICLR 2020
Table 4: Comparison of DPQ against recently proposed embedding compression techniques on the PTB LM task (LSTMs with three model sizes are studied). Metrics are perplexity (PPL) and compression ratio (CR).
Small
Medium
Large
Method
PPL
CR
PPL
CR
PPL
CR
Full
114.5
1
83.4
1
78.7
1
Pre-train (Chen
2018h
108.0
4.8
84.9
11.7
80.7
18.5
E2E (Chen ct al. 2018b)
108.5
4.8
89.0
11.7
86.4
18.5
E2E-dist. (Chen ct al 2018b
107.8
4.8
83.1
11.7
77.7
18.5
DPQ-SX
105.8
85,5
82.0
82.9
78.5
238.3
DPQ-VQ
106.5
51.1
83.3
58.7
79.5
238.3
114.9
91.2
89.0
123.00
92.9
86
2000.0
327.8
415.9
7.5
20.0
25.4
9.0
13.6
23.4
1657.2
484
1264
-05.
.1.7
-04.5
-55.6
-3.7
138.6
26-7
36.0
15.7
519.1
130
82.5
24.9
25
50
18
25
70
10
50
170
135
128
nAg-99 (DPG-VG)
BLTU (XPgg-XX)
DN (3PQ-V)
Ex
(a) PTB
(b) IWSLT15 (En-Vi)
Figure 3: Heat-maps of task performance and compression ratio for various K and D values. Darker better. Key observations are: 1) increasing K or D typically improves the task performance at the expense of lower CRs; 2) the combination of small K and large D is better than the other way round.
layer and thirdly fixes the discrete codes and trains the model again: E2E: end-to-end training without distillation guidance from pre-trained embedding table; E2E-dist. end-to-end training with distillation procedure that uses pre-trained embedding guidance during training. Fable/4showe comparison between DPQ and the above methods on the PTB language modeling task using LSTMs with three different model sizes. We find that 1) both Pre-train and E2E achieve good com- pression ratios but with worse perplexity scores on the Medium and Large models, 2) the E2E-dist. method has the same compression ratio them and is able to achieve similar perplexity scores the full embedding baseline, with the downside that it requires the extra distillation procedure, 3) DPQ variants (particularly DPQ-SX) able obtain extremely competitive perplexity scores in all cases, while offering compression ratios that are an order of magnitude larger than the alternatives.
3.2 EFFECTS OF K AND
Among key hyper-parameters of DPQ are the code size: K the number of centroids per dimension and the code length. Figure 3 shows the task performance and compression ratios for different K and D values on PTB and IWSLT15 (En-Vo Vi). Firstly,
observe that the combination of small K and large is better configuration than the other way round. For example, in IWSLT15 (En-Vi), (K = 2, D 128) is better than (K 128, D 8) in both BLEU and CR, with both DPQ-SX and DPQ Secondly, increasing K or D would typically improve the task performance at the expense of lower CRs, which means one can adjust and D achieve the best task performance and compression ratio trade-off. Thirdly, we note that decreasing D has a much more traumatic effect on DPQ-VQ than on DPQ-SX in terms of task performance. This is because the dimension of each sub-space (d/ D) increases, the nearest neighbour approximation (that DPQ-VQ relies on) becomes less exact.
3.3 COMPUTATIONAL COST
DPQ incurs a slightly higher computational cost during training and no extra cost at inference. Figure A shows the training speed well the (GPU) memory required when using DPQ the medium
7

---

Under review as a conference paper at ICLR 2020
LSTM model, trained on Tesla- V100 GPUs. For most K and D values, the extra training time is within 10%, and the extra training memory is zero. For very large K and D values, DPQ-VQ has better computational efficiency than DPQ-SX (as expected). At inference, we do not observe any impact speed memory from DPQ.
1.%%
7.1%
7.2%
6.7
0.0%
0.0%
0.0%
3.0%
5.0%
0.2%
8.1%
11.0%
13.99
9%
80%
98%
8.2%
0.0%
6.0%
0.0%
0.0%
0.0%
0.0%
0.0%
4.5%
10
20
158
150
93
58
130
10
50
135
DPG-EX
DO
DFG-XX
DPG-V
(a) Extra training time used.
(b) Extra training memory used.
Figure 4: Extra training cost incurred by DPQ, measured on medium sized LSTM for LM trained on Tesla- GPUs. For most K and π) values, the extra training time is within 10%, and the extra memory usage is zero. For very large K and D values, DPQ-VQ has better computational efficiency than DPQ-SX in both memory and speed (as expected).
3.4
CODE STUDY
To better understand the KD codes learned end-to-end via DPQ, we investigated the codes and observed the following. Firstly, the centroids in all D groups are usually well utilized (Appendix D.I.. Secondly, the KD codebook changes as training progresses, but the rate of change decreases throughout training and converges to < 20% (Appendix(D2) D.2) Thirdly, the nearest neighbours in the continuous embedding space between DPQ and the baseline align very well (Appendix (D3)). Finally, we also list the learned codes for selected words in (pp
4 RELATED WORK
Modern neural networks have many parameters and redundancies. The compression of such models has attracted many research efforts (Han et al.) 2015 Howard et al., [0e7::Chen Chen al. [018ap. Most of these compression techniques focus on the weights that are shared among many examples, such as convolutional and dense layers (Howard et al. 2017 Chen et al. (018a)). The embedding layers different in the sense that they are tabular and very sparsely accessed, T.e. the pruning cannot remove rows/symbols in the embedding table, and only few symbols are accessed in each data sample. This makes the compression challenges different for the embedding layers.
Existing work compressing embedding layers includes (Shu and Nakayama, 2017 2018b), which also leverages discrete codes. However, we propose a new formulation from product quantization perspective, in which discrete codes are compute from product quantization on some continuous space. This formulation makes it more general and allows two types of instantiations with different gradient approximation. The product keys and values in our model also make it more efficient in both training and inference. Empirically, DPQ achieve better compression ratios without resorting to the extra distillation process.
Chen et al.
Our work differs from traditional quantization techniques (Jegou et al. 2010) in that they can be trained in an end-to-end fashion. The idea of utilizing multiple orthogonal subspaces/groups for quantization is used in product quantization (Jegou et al.) 2010 Norouzi and Fleet 2013) and multi-head attention (Vaswani et al.) 2017).
The two approximation techniques presented for DPQ in this work also share similarities with Gumbel-softmay (Jang et al. 2016) and VAE (van den Oord et al. 2017). However, we do not useful since aim to get deterministic codes.
find using stochastic noises in Gumbel-softmat is also worth pointing out that these techniques (Jang et al., 2016 yan den Oord et al. 2017) by themselves cannot be directly applied to compression.
8

---

Under review conference paper at ICLR 2020
5 CONCLUSION
In this work, we propose novel and general differentiable product quantization framework for learning compact embedding layers. We provide two instantiations of our framework, which can readily serve as drop-in replacement for existing embedding layers. Empirically, we evaluate the proposed method on ten datasets across three different language tasks, and show that our method surpasses existing compression methods and can compress the embedding table up 238x without suffering performance loss. In the future, we plan to apply the DPQ framework to a wider range of applications and architectures.
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory- Adaptive Optimization for Large-Scale Learning. In arXiv, 2019.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1.08.338... 34332, 2013.
Antoine Bordes, Nicolas Usunier, Alberto Garcia-Duran, Jason Weston, and Oksana Yakhnenko. Translating embeddings for modeling multi-relational data. In Advances in neural information processing systems, pages 2787-2795, 2013.
Ting Chen. Ji Lin, Tian Lin, Song Han, Chong Wang, and Denny Zhou. Adaptive mixture of low-rank factorizations for compact neural modeling. Neural Information Processing Systems (CDNNRIA workshop), 2018a.
Ting Chen, Martin Renqiang Min, and Yizhou Sun. Learning k-way d-dimensional discrete codes for compact embedding representations. In International Conference Machine Leurning, 201 8b.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 18Xiv:1810.04450 2018.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv: 2015.
Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: 1704. 07704.0466, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameteriization with gumbel-softmax. arXiv preprint arXiv:1611.. 1.01144, 2016.
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search.
IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128, 2010.
Armand Joulin, Edouard Grave, Piotr Bojanowski, and Tomas Mikolov. Bag of tricks for efficient text classification. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics, 2017.
Lukasz Kaiser, Aurko Roy, Ashish Vaswani, Niki Parmar, Samy Bengio, Jakob Uszkoreit, and Noam Shazeer. Fast decoding in sequence models using discrete latent variables. arXiv preprint arXiv: 03382, 2018.
Yehuda Koren, Robert Bell, and Chris Volinsky. Matrix factorization techniques for recommender systems. Computer, pages 30-37, 2009.
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods Natural Language Processing: System Demonstrations, pages 66-71, 2018.
Minh-Thang Thang Luong, Eugene Brevdo, and Rui Zhao. Neural machine translation (seq2seq) tutorial. http:://gihubb..omeenn///// 2017.

---

Under review conference paper at ICLR 2020
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representa- tions in vector space. arXiv preprint arXiv: 2013.
Mohammad Norouzi and David Fleet. Cartesian k-means. In Proceedings of the IEEE Conference computer Vision and Pattern Recognition, pages 3017-3024, 2013.
Raphael Shu and Hideki Nakayama. Compressing word embeddings via deep compositional code learning. arXiv preprint arXiv:1711.01068, 2017.
Richard Socher, Danqi Chen, Christopher D Manning, and Andrew Ng. Reasoning with neural tensor networks for knowledge base completion. In Advances in neural information processing systems, pages 926-934, 2013.
Aaron van den Oord, Oriol Vinyals, et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems, pages 6306-6315, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances in neural information processing systems, pages
2017.
Wojciech Zaremba, Ilya Sutskever, and Oriol Vinyals. Recurrent neural network regularization, arXiv preprint arXiv:1409.2329, 2014.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In Advances in neural information processing systems, pages 649-657, 2015.
10

---

Under review conference paper at ICLR 2020
A ALGORITHM PSEUDO-CODE
This section lays out the algorithm pseudo-code for the DPQ embedding layer during the forward training/inference pass.
Algorithm DPQ for the i-th token the vocab (training, forward pass)
Algorithm 2 DPQ for the i-th token in the vocab (inference)
h-params :K, D
h-params :K, D parameters: V E BXXDx(1/D) c e
parameters: Q E BaXDx(1ID), K.V E I DKDDx(1/DD, Ce (I,,, KjnxD
[I,.,,Knx)
for j in 1, ...,
for in 1,... do
(1) arg max; dist Q (1) kK)
h.00 v c
h v c
end for return concatenate h! h/ (
end for return concatenate (h) h, h (n)
B PROOF OF THEOREM
Proof. We first re-parameterize both the codebook C and the Value matrix V follows.
The original codebook E {1,..., """".",)"""" K]nxD, and we turn each code bit, which is an integer in [1..... K}, into small one-hot vector of length-K. K. This results in the new binary codebook B E 1 Per our constraint in theorem B is full rank matrix.
The original Value matrix is V E Kxd, and we turn it into a block-diagonal matrix U E RDxd where the j-th block diiegnaa is set to vOi E (d/ Given that each block diagonal, i.e. V(i, full rank, the resulting block diagonal matrix U also full rank.
With the above re-paameterizatton, we can write the output embedding matrix H = BU. Given both B and U are full rank and KD 2 d, the resulting embedding matrix H is also full rank.
C DETAILS OF MODEL TRAINING
We follow the training settings of the base models used, and most of the time, just tune the DPQ hyper-parmeters such K, D and/or cusspace-sharingg We also apply batch normalization for the distance measure in DPQ along the K-dimension, i.e. each centroid will have normalized distance distribution with batch samples.
For training the Transformer Model on WMT' 19 En-De dataset, the training set contains approxi- mately 27M parallel sentences. We generated vocabulary 32k sub-words from the training data using the SentencePiece tokenizer (Kudo and Richardson 2018). The architecture is the Transformer Base configuration described in Vaswani et al. (2017) with a context window size of 256 tokens. All models optimizer (Anil et al., 2019) with momentum 0.9 and a quadratic learning rate warm-up schedule with 10k warm-up steps. We searched the learning rate in (0.1,0.3).
trained with batch size of 2048 sentences for 250k steps, and with the SM3
D CODE STUDY
D.1 CODE DISTRIBUTION
DPQ discretizes the embedding space into the KD codebook {1, .... K) We examine the code distribution by computing the number of times each discrete code in each of the 1) groups is used in
11

---

Under review conference paper at ICLR 2020
the entire codebook:
Count ( &2 = -C = L.
K}
Figure 5 shows the code distribution heat-maps for the Transformer model on WMT 19 En-De, with K 32 and D 32 and Sebspace-sharingg We find that 1) DPQ-VQ has
evenly distributed code utilization, 2) DPQ-SX has a more concentrated and sparse code distribution: in each group, only few discrete codes used, and some codes are not used in the codebook.
4000
3200
2400
1100
Figure 5: Code heat- mapt. Left: DPQ-SX. Right: DPQ-VQ. æ-axis: codes per group. y-axis: D
groups. K D 32.
D.2 RATE OF CODE CHANGES
We investigate how the codebook changes during training by computing the percentage of code bits in the KD codebook C changed since the last saved checkpoint. An example is plotted in Figure for the Transformer on WMT' 19 En-De task, with D 128 and various K values. Checkpoints were saved every 600 iterations. Interestingly, for DPQ-SX, code convergence remains about the same for different K values; while for DPQ-VQ, the codes takes longer stabilize for larger K values.
DPp-VQ, K-d
DPO-VO, K-d2
bRp-V K-128
DRQ St, K-8
DPO SX, K-32
DPQ-XX K-128
0x
0.2
0.5
Figure 6: Percentage code bits in codebook which changed from the previous checkpoint. Trans- former WMT' 19 En-De. D 128 for all runs. Checkpoints saved every 600 iterations.
D.3 NEAREST NEIGHBOURS OF RECONSTRUCTED EMBEDDINGS
Table 5 and 7 show examples of nearest neighbours in the reconstructed continuous embedding space, trained the Transformer model on the WMT' 19 En-De task. Distance between two sub- words is measured by the cosine similarity of their embedding vectors. Baseline is the original full embeddings model. DPQ variants were trained with
D 128 with no subspace-sharing.
Taking the sub-word "_evolve' bours as the original full embedding: both have out of 10 overlapping top neighbours as the baseline model. However, in DPQ-SX the neighbours have closer distances than the baseline, hence tighter cluster; while DPQ-VC VQ the neighbours are further from the original word. We observe similar patterns in the other two examples.
example, DPQ variants give very similar top 10 nearest neigh-
12

---

Under review conference paper at ICLR 2020
Table 5: Nearest neighbours of _evolvve in the embedding space.
Baseline (Full)
Dist
DPQ-SX
Dist
DPQ-VQ
Dist
cvollve
1.000
cvolvv
1.000
_evolve
1.000
eevolved
0.533
cvoolved
0.571
evolved
0.506
evolving
0.493
evolution
0.499
develop
0.417
develop
0.434
develop
0.435
evolving
0.359
evolution
0.397
evolving
0.418
developed
0.320
developed
0.379
_arisc
0,405
development
0.307
.dveeloping
0.316
_eveloopd
0.405
developing
0.299
_arise
0,298
rresulted
0.394
evoluuion
0,282
unfolld
0.294
originate
0.361
changed
0.278
emerge
0.290
rresult
0.359
grrw
0.273
Table 6: Nearest neighbours of "_monppoly"" in the emhedding space.
Baseline
Dist
DPQ-SX
Dist
DPQ-VQ
Dist
monopoly
1.000
monopoly
1.000
monopoly
1.000
monopolies
0.613
monopolies
0.762
monopolies
0.509
monopol
0.552
monopol
0.714
monopol
0.483
Monopol
0.380
Monopol
0.531
Monopol
0.341
moratooium
0,271
zugestimmt
0,486
dominant
0.258
privileged
0.269
legitim
0.420
moortooiium
0.239
.niilatrrll
0.262
GroBunternehmen
0.401
Lautonomy
0.230
.iiracle
0.260
Figenkapital
0.400
_uggllasenn
0.227
privilege
0.254
wirkungsvoll
0.399
Limperial
0.226
dominant
0,250
UCLAF
0.388
capptaalsst
0.223
Table 7: Nearest neighbours of
Toronto' in the embedding space.
Baseline
Dist
DPQ-SX
Dist
DPQ-VQ
Dist
Toronto
1.000
Toronto
1.000
Toronto
1.000
Vancouver
0.390
Chicago
0.475
Orlando
0.307
Tokyo
0.378
Orleans
0.467
Detroit
0.306
Ottawa
0.372
Melbourne
0.435
Canada
0,280
Philadelphia
0.353
Miami
0.434
London
0,280
Orlando
0.345
Vancouver
0.415
Glasgow
0.276
Chicago
0.340
Tokyo
0.407
Montreal
0.272
Canada
0.330
Ottawa
0.405
Vancouver
0.271
Seoul
0.329
Azeroth
0.403
Philadelphia
0.267
Boston
0.325
Antonio
0.400
Hamilton
0.264
D.4 CODE VISUALIZZATOON
Table 8 shows some examples of compressed codes for both DPQ-SX and DPQ-VQ. Semantically related words share common codes more dimensions than unrelated words.
E ADDITIONAL HYPER-PARAMETERS STUDY
E. EFFECTS OF K AND D
Figure shows extra heatmaps with varied K and D in addition to those in Section 3.2
13

---

Under review conference paper at ICLR 2020
Table Examples of KD codes.
DPQ-SX
DPQ-VE VQ
Monday
2
5
7
0
6
6
5
2
4
7
Tuesday
6
0
7
0
7
1
7
0
2
0
3
7
Wednesday
6
5
3
0
6
6
2
2
0
7
Thursday
7
0
2
Friday
4
6
7
0
1
7
6
0
0
2
1
7
Saturday
4
0
7
0
0
6
2
0
2
3
3
7
Sunday
2
0
3
0
6
7
2
2
6
7
Obama
2
6
2
5
7
2
3
6
6
4
Clinton
3
7
3
Merkel
6
6
3
4
Sarkozy
7
6
1
4
5
0
0
3
7
5
7
7
4
Berlusconi
4
6
1
4
6
7
6
3
0
6
6
7
7
4
Putin
2
6
1
6
6
7
5
3
6
6
6
Trump
7
5
7
Toronto
4
6
3
7
Vancouver
2
6
6
3
6
1
Ottawa
2
5
1
6
2
7
6
3
6
6
2
0
4
Montreal
4
0
2
6
1
7
4
3
6
2
0
1
London
2
4
7
2
3
6
7
Paris
4
0
3
6
Munich
4
2
4
0
5
0
1
3
5
6
3
7
-226
300.9
302.8
342.6
311.7
102.2
2027.5
032.0
418.4
100.
97.5
39.1
107.7
J02.7
.8.8
$63.0
277.0
158.8
36.5
37.3
38
20-7
37
28
000
170.2
42.
A
100.0
06.5
15.7
200.0
37.3
37
26.7
27
101.
41.4
46.8
.44
-100
-5.0
267.1
116.9
562
25
23
20
nery (0R0)500000 5X0
1eg (DRO)VG)
R.EU
OLEU (BP0-V00)
WiKiTex12
(b) WMT19 (En-DE)
Figure 7: Heat-maps of task performance and compression ratio. Darker is better.
E.2 SUBSPACE-SHANE
Subspace-sharing refers to the option of whether to share parameters among the D groups in the Key/Value Matrices, i.e. constraining (() K(') and Viv) VO') v'jj'. For simplicity refer to this as "subspace-sharingg. Subspace-sharing improves the compression ratio to: CR = 32nd/ (n log, K + 32Kd/D).
Figure shows the trade-off curves task performance and compression ratio with different DPQ variants, K, D and subspace-sharing. We find that could vary the hyper-parameters to search for optimal performance and compression trade-off. We also observe the effect of subspace-sharing appears very much task-dependent: it improves perplexity scores in LM tasks but hurts BLEU scores in NMT tasks. For TexIC tasks, subspace-sharing seems beneficial for DPQ-SX but harmful for DPQ-VQ.
F RELATIONS TO CHEN ET AL. (2018B) AND OTHER CONVENTIONAL METHODS
Both this work and (Chen et al.) 2018b) are based on the idea representing symbols with discrete codes, but there some major differences which listed below:
14

---

Under review conference paper at ICLR 2020
SPOVO The
ERIGNG ThE
Aatedd Soartt
50 clamm
255
246
800
766
200
300
(a) PTB
IWSLTIS (En-Vi)
(c) DBpedia
TRO DE
100
Netormanc
Mirrer Jhand
Oergressiennu
Compr Couio 1.
Comprsssinnna rel
(d) Wikitext-2
(c) IWSLT15 (Vi-En)
(f) AG News
Aiciod 200/0
092 Fide
(R040 11.
299/00 7U6
510
Comsress rafo
Compresdin redb
Compresdon ratl
(g) Yahoo! Answers
(h) Yelp Polarity
(i) Yelp Full
Figure 8: Task performance compression ratio trade-off curves. Each subplot comes from one task/dataset and contains four configurations: (DPX-SX, DPX-VQ] x (subspace-sharing, NO- subspace-sharing).
In (Chen et al. 2018b), discrete codes are directly associated with each the symbols, in this work, discrete codes computed outcome of product quantization. This shift of perspective allows the proposed framework to generalize beyond a fixed set of vocabulary, and be applied in potentially in any other neural network layers stand-alone module.
Our formulation of discrete codes with product quantization allows us to derive two variants with different approximation techniques ssoftmax-based and vector iiaanizatioonbaselee while (Chen et al. 2018b) is only based on softmax approximation.
The product quantization has minimal overhead and very efficient compared to encoder functions used in (Chen al., (018b), i.e. MLP-based and RNN-based functions that compose codes into continuous embedding. Our DPQ has very small memory footprint and computation time overhead (Figure 4). Furthermore, the approximation error also reduced, and DPQ can be truly trained end-to-end without two pass training with distillation loss in (Chen al. 2018b).
Here are comparisons to more traditional approaches:
Scalar quantization: it quantize each floating number independently. and has very limited compression ratios. E.g. quantizing float32 into int8 would offer CR of 32/8=4, while likely dropping in task performance metrics (e.g. PPI.).
Product quantization: it generalizes scalar quantization and quantize sebtorc.ors. However, this approach is cindifferrentiale (cannot train end-to-end) and requires a post-training procedure. Small quantization errors accumulate and thus performances suffer.
15

---

Under review conference paper at ICLR 2020
Pruning: pruning effect reduces the embedding size for each symbol. Therefore its performance is usually less than ideal (Shu and Nakayama 2017).
Low-rank factorization: larger compression ratio requires smaller rank, which in effect reduces embedding table size and leads to worse results.
Different from these techniques, DPQ makes use of discrete codes, and uses product quantization to generate discrete codes. Unlike traditional product quantization, we propose techniques make it end-to-end differentiable so that the neural nets can adapt to quantization error. DPQ also relates to factorization-baat method (Theorem D]. but DPQ can produce high-rank embedding tables with sparse factorization.
G COMPARISONS TO MORE BASELINES
G.1 COMPARISONS TO TRADITIONAL. COMPRESSION TECHNIQUES
Table 9 shows comparisons on PTB language modeling task (medium-sized LSTM) with broader set of baselines (including methods that are not based on discrete codes). We find that 1) traditional compression techniques, such scalar and product quantization, as well as low-rank factorization, typically degenerates the performance significantly in order to achieve good compression ratios compared to discrete code learning-basse methods (Chen al. [00 8b Shu and Nakayama, (017); 2) the proposed method (DPQ) can largely improve the compression ratio while achieving similar or better task performance (perplexity this case).
Table 9: Performance comparison on PTB language modeling task. The proposed method provides significantly better compression ratio over baselines while achieving similar better/smaller PPI..
Method
PPL
Compression ratio
Full
83.38
1.0
Scalar quantization (8 bits)
84.06
4.0
Scalar quantization (6 bits)
87.73
5.3
Scalar quantization (4 bits)
92.86
8.3
Product quantization (64x325)
84.03
8.3
Product quantization( 128x
83.71
6.7
Product quantization(256x325
83.66
5.3
Low-rank (5X)
84.84
5.0
Low-rank (10X)
85.53
10.2
Shu and Nakayama (2017)
84.92
12.5
Chen et al. (2018h)
83.11
12.5
Ours (DPQ-VQ)
83.3
58.7
Ours (DPQ-SX)
82.0
82.9
G.2 COMPARISONS TO BASELINES ON TEXT CLASSIFICATION
Table TO provides performance comparisons on text classification task. We found that the proposed method (DPQ) usually achieve better accuracies than baselines, at the same time providing better compression ratios.
G.3
COMPARISONS TO POST-TRAINING RECONSTRUCTION BASED BASELINES ON NMT
The proposed method (DPQ) supports end-to-end compact embedding learning. An alternative is learning to reconstruct the learned full embedding table with discrete codes after the model is train. The reconstructed compact embedding table is then used to replace the original embedding table for inference. We name this Reconstruction baseline. In our experiment, we use auto-encoder and DPQ (with different K and D) to learn to reconstruct the trained full embedding table.
Table II shows performance comparisons between the proposed method and reconstruction baseline WMT19 (En-De) translation task based Transformer (Vaswani et al. 2017). We that
16

---

Under review conference paper at ICLR 2020
Table 10: Performance comparison on text classification task. The accuracy and compression ratios (in parenthesis) are shown below. The proposed method (DPQ) usually achieve better accuracies than baselines, at the same time providing better compression ratios.
Dataset
AG News
Yahoo!
DBPedia
Yelp
Yelp F
Full
92.6 (1.0)
69.4 (1.0)
98.1 (1.0)
93.9 (1.0)
60.3 (1.0)
(ow-rank(10x)
91.4 (10.4)
69.5 (10.2)
97.7 (10.3)
92.4 (10.4)
57.8 (10.3)
Low-rank x
91.5 (21.4)
69.1 (21.5)
97.9 (21.3)
92.4 (21.5)
57.3 (21.4)
Chen et al. (2018b)
91.6 (53.3)
(31.7)
98.0 (48.4)
93.1 (48.6)
59.0 (54.4)
DPQ-VQ VQ DPQ-SX
92.6 (24.0)
69.2 (19.2)
98.1 (38.5)
93.9 (24.0)
60.2 (24.1)
92.5 (19.3)
69.6 (48.2)
98.1 (24.1)
94.2 (38.5)
60.1 (48.2)
the reconstruction baseline degenerates the performance significantly. This is expected as small approximation errors in the embedding layer accumulate and can be amplified as the errors propagate through the deep neural nets, finally lead to large error in output space. Our method does not have this problem the whole system is jointly trained the later networks account for small approximation errors the early layer.
Table 11: Performance comparisons against the reconstruction baselines.
Method
BLEU
CR
Full
38.8
Reconstruction (K=128, D=64)
28.9
31.9
Reconstruction (K=32, D=128)
35.4
25.0
Reconstruction (K=128, D=128)
35.7
17.0
Reconstruction (K=32, D=256)
36.9
12.6
Reconstruction (K=128, D=256)
37.8
8.8
DPQ-VQ (K=32, D=128)
38.7
17.0
DPQ-SX (K=32, D=128)
38.8
17.0
G.4 MORE ABLATIONS ON DPQ-SX
Table 12 shows ablation study PTB language modeling task (medium-sized) LSTM), in which we choose to tie the K and V matrices in DPQ-SX (achieved by sharing single variable during the optimization process). We fix K=128, D=50. We find DPQ-SX (with untied K,V) perform the best, followed by DPQ-SX (tied K, K) and DPQ-VQ.
Table 12: Ablation study of DPQ-SX on whether or not to tie K and V matrices. By default, DPQ-SX does not tie these two matrices.
Method
PTB
Wikitext-2
DPQ-SX (untied K, V)
82.4
95.2
DPQ-SX (lied K, V) DPQ-VQ
83.5
95.8
83.5
97.0
H APPLYING DPQ TO BERT
BERT (Devlin et al. 2018) has shown excellent results on a wide range natural language tasks, therefore is important to demonstrate that DPQ can also achieve competitive performance on BERT. As our baseline, we pre-trained BERT-base on 512-token sequences for 1M iterations with batch size 1024. We the same optimizer (Adam) and learning rate schedule described in Devlin et al. (2018). For the DPQ experiments, we use DPQ-SX with no sabspace-sharing D = 128 and K 32; these choices not from hyper-parameter search, but inspired from results Transformer
17

---

Under review conference paper at ICLR 2020
WMT19 EnDe. We pre-trained BERT with embedding layer replaced by our DPQ, and then finetuned all model parameters for downstream tasks. Note that we do not perform additional tuning for the DPQ experiments either pre-training or finetuning: we used exactly the same configurations and hyperparameters in our baseline. Table [ that DPQ performs par with full embedding in most of the downstream tasks, while giving a compression ratio of 37 x on the embedding table. This equivalent to saving 24M parameters in the BERT-base model, decreasing the total model size by 22.2%.
Table 13: Effect using DPQ on BERT. DPQ gives a compression ratio 37x on the embedding table while the model's performance downstream tasks remains competitive.
Embeddings
CR
Squad 1.1
Squad 2.0
CoLA
MNLI
MRPC
XNLI
Full
1.0
90. 00.1/83.1
79.3/76.1
81.1
84.2
86.0
53.3
DPQ-SX
37.0
90.0/83.1
78.1/74.9
80.8
83.9
85.8
53.5
18


Reviews:
[{"review_id": "r1lYKM7SKB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this manuscript, authors improve the work in [1] by simplifying the reverse-discretization function. The empirical study demonstrates the effectiveness of the proposed algorithm.\nI’m not familiar with the area. The differences between this work and [1] should be elaborated more in the related work, since they are closely related.\nBesides, for the technical part, DPQ-SX outperforms DPQ-VQ while the softmax approximation seems identical to that developed in [1].\n\n[1] ICML’18: Learning K-way D-dimensional Discrete Codes for Compact Embedding Representations", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "Hyl2WZpf5H", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper considers the problem of having compact yet expressive KD code for NLP tasks. The authors claim that the proposed differentiable product quantization framework has better compression but similar performance compared to existing KD codes.The authors present two instances of the DPQ framework: DPQ-SX using softmax to make it differentiable, and DPQ-VQ using centroid based approximation. While DPQ-SX performs better in terms of performance and compression, DPQ-VQ has the advantage in scalability.\n\n- Significance\nIt's understandable that the size of the embedding is important, but there's been a lack of explanation as to why this should be done only through KD codes. Hence, it is doubtful how big the impact of the proposed framework is.\n\n- Novelty\nJust extending and making Chen et al., 2018b's distilling method to be differentiable has limited novelty.\n\n- Clarity\nThe paper is clearly written in most places, but there were some questions about the importance and logic of statements.\n\n- Pros and cons\nCompared to Chen et al., 2018b, there is no need to use expensive functions, and performance is better. But, the baseline consists only of algorithms using KD codes; there might be many disadvantages compared to other types of algorithms.\n\n- Detailed comments and questions\n1. It is true that the parameters for embedding make up a large part of the overall parameters, but I would like some additional explanation of how important they are to learning. It is usually not necessary to train the entire embedding vector on GPU, so it would not be a big issue in the actual learning process.\n2. In a similar vein, it would be nice to show which of the embedding vector size or the LSTM model size contributes significantly to performance improvements. If LSTM model size contributes more, the motivation would be weakened.\n3. It would be nice to add more baselines such as Nakayama 2017 as well as the standard compression/quantization methods used in other deep networks. And please explain why we should use KD codes to reduce embedding size. Also, why the distilling in Chen et al., 2018b is a problem?\n4. Did you run all experiments just one time? There is no confidence interval.\n5. DPQ models have different compression ratios depending on the size of K and D. It would be great to show the change in PPL according to the compression ratio of DPQ models.\n6. Can we apply it to pre-trained models like BERT?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "Hyet4X5p_S", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThis 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.\n\nI 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:\n\n1. 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,\n\nC = [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).\n\n2. 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.\n\n3. 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.\n\n4. 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.\n\nFor improving the paper, the relatively minor comments are as the following:\n\n1. In equation 4, the partition function Z is not explicitly defined.\n\n2. In the second paragraph of section 3.1, it is not clear what exactly is the pre-trained embedding used as baseline.\n\n3. For better readability, it is better to inflate the caption of figure and tables to provide useful take-away message there.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The presented paper gives a differentiable product quantization framework to compress embedding and support the claim by experiments (the supporting materials are as large as the paper itself). Reviewers agreed that the idea is simple is interesting, and also nice and positive discussion appeared. However, the main limiting factor is the small novelty over Chen 2018b, and I agree with that. Also, the comparison with low rank is rather formal: of course it would be of full rank , as the authors claim in the answer, but looking at singular values is needed to make this claim. Also, one can use low-rank tensor factorization to compress embeddings, and this can be compared. 
To summarize, I think the contribution is not enough to be accepted.
Decision: Reject