Title: LEARNING COMPACT EMBEDDING LAYERS VIA 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.

1. 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-words models in the text domain (Mikolov et al., 2013; Devlin et al., 2018) , embedding layers are also used in a wide range of applications such as knowledge graphs (Bordes et al., 2013; Socher et al., 2013) 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 to first learn to encode symbols/words with K-way D-dimensional discrete codes (KD codes, such as 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 are fixed before training and are therefore non-adaptive and limited to downstream tasks. Chen et al. (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 a pre-trained embedding table as guidance, in order to match the performance of the full embedding baseline.

In this work, we propose a 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 et 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 et al., 2018b; Shu and Nakayama, 2017) , our framework 1) brings a new and general perspective on how the discrete codes can be obtained in a 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 of product keys and values) while avoiding the cumbersome distillation procedure.

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 as the original full embeddings. Furthermore, our results are obtained from end-to-end training where 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 an end-to-end fashion without distillation.

2. METHOD:

Problem setup. An embedding function can be defined as F W : V → R d , where V denotes the vocabulary of discrete symbols, and W ∈ R n×d 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 a given objective. The goal of this work is to learn a compact embedding function F W 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 a 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 it 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, we enable end-to-end learning of discrete codes via optimizing some task-specific objective.

2.1. DIFFERENTIABLE PRODUCTION QUANTIZATION FRAMEWORK:

The proposed differentiable production quantization (DPQ) function is a mapping between continuous spaces, i.e. T : R d → R d . In between the two continuous spaces, there is a discrete space {1, • • • , K} D 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 d → {1, • • • , K} D that maps a continuous vector into a K-way D-dimensional discrete code (KD code), and 2) a reverse-discretization function ρ(•) : {1, • • • , K} D → R d that maps the KD code into a continuous embedding vector. In other words, the general DPQ mapping is

T (•) = ρ • φ(•).

Compact embedding layer via DPQ. In order to obtain a compact embedding layer, we first take a raw embedding and put it through DPQ function. More specifically, the raw embedding matrix can be presented as a Query matrix Q ∈ R n×d where the number of rows equals to the vocabulary size.

The discretization function of DPQ computes discrete codes:

C = φ(Q) where C ∈ {1, • • • , K} n×D is the KD codebook.

To construct the final embedding table for all symbols, the reverse-discretization function of DPQ is applied, i.e. H = ρ(C) where H ∈ R n×d 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 in the reverse-discretization function. They are sufficient to (re)construct partial or whole embedding table. In below, we specify the discretization function φ(•) and reverse-discretization function ρ(•) via product keys and values.

Product keys for discretization function φ(•). Given the query matrix Q, the discretization function computes the KD codebook C. While it is possible to use a complicated transformation, in order to make it efficient, we simply leverage a Key matrix K ∈ R K×d 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 K (j) ∈ R K×d/D and Q (j) ∈ R n×d/D . We can compute each of D dimensional KD codes separately. The j-th dimension of a KD code C i for the i-th symbol is computed as follows. The dist(•, •) computes distance measure between two vectors, and use it to decide which discrete code to take.

C (j) i = arg min k dist Q (j) i , K (j) k (1)

Product values for reverse-discretization function ρ(•). Given the codebook C, the reversediscretization 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 ∈ R K×d as 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 (j) ∈ R K×d/D . 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.

EQUATION

We note that this is a simplification, both conceptually and computationally, of the ones used in (Chen et al., 2018b; Shu and Nakayama, 2017) , which reduces the computation overhead and eases the optimization.

Figure 1 illustrates the proposed framework. The proposed method can also be 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 2 K bits, which is the only thing that depends on vocabulary size n, and 2) Value matrix V requires 32Kd bits1 , which does not explicitly depend on n and is ignoble when n is large. Since typically nD log 2 K < 32nd, the DPQ embedding is more compact.

Inference complexity. Since only indexing and concatenation (Eq. 2) are used during inference, both the extra computation complexity and memory footprint are usually negligible compared to the regular full embedding (which directly indexes an embedding table ).

Expressiveness. Although the DPQ embedding is more compact than full embedding, it is not achieved by reducing the rank of the matrix (as in traditional low-rank factorization). Instead, it 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 ∈ {1, ..., K} n×D , denoted as B ∈ {0, 1} n×KD , is full-rank.

2) Sub-matrices of splitted V, i.e. V (j) ∈ R K×d/D , ∀j, are all full-rank.

3) KD ≥ 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 it is easy to achieve nD log 2 K < 32nd with KD = d.

So far we have not specified some designs of the discretization function such as the distance function in Eq 1. More importantly, how can we compute gradients through the arg min function in Eq. 1? 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-BASED APPROXIMATION:

The first instantiation of DPQ (named DPQ-SX) approximates the non-differentiable arg max operation with a differentiable softmax function. To do so, we first specify the distance function in Eq. 1 with a softmax function as follows.

EQUATION

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 al., 2018b; Jang et al., 2016) , we relax the softmax function with temperature τ :

EQUATION

where

Z = k exp( Q (j) i , K (j) k /τ ). Note that now C(j) i ∈ ∆ K is a probabilistic vector (i.e. soft one-hot vector) instead of an integer C (j) i . And one_hot(C (j) i ) ≈ C(j) i , or C (j) i = arg max C(j)

i . With a one-hot code relaxed into soft one-hot vector, we can replace index operation V (j) C(j) i with dot product to compute the output embedding vector, i.e. j) . The softmax approximated computation defined above is fully differentiable when τ = 0. However, to compute discrete codes during the forward pass, we have to set τ → 0, which turns the softmax function into a spike concentrated on the C (j) i -th dimension. This is equivalent to the arg max operation which does not have gradient.

H (j) i = C(j) i V (

To enable a pseudo gradient while still be able to output discrete codes, we use a different temperatures during forward and backward pass, i.e. set τ → 0 in forward pass, and τ → 1 in the backward pass. So the final DPQ function can be expressed as follows.

H i = T (Q i |τ = 1) -sg T (Q i |τ = 1) -T (Q i |τ = 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 a centroid-based approximation, which directly pass the gradient straight-through (Bengio et al., 2013) a 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.

C (j) i = arg min k Q (j) i -K (j) k 2 (6)

Secondly, we tie the Key and Value matrices, i.e. V = K, so that we can pass the gradient through.

We still have the non-differentiable arg min operation, and the input query

Q (j) i are different from selected output centroid V (j) C (j) i

. However, since they are in the same space, it allows us to directly pass the gradient straight-through as follows.

H

i = Q i -sg(Q i -T (Q i )) (7) k = 1, 2, 3, …, K k = 1, 2, 3, …, K c i (j)

Forward pass: use argmax over k:

Backward pass: use softmax to approximate argmax

c i (j) (a)

Softmax-based (DPQ-SX) Where sg is again the stop gradient operation. During the forward pass, the selected centroid is emitted, but during the backward pass, the gradient is pass to the query directly. This provides a 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.

K 1 (j) K 2 (j) Q i (j) K 0 (j)

However, it is worth noting that the Eq. 7 only approximates gradient for Query matrix, but does not updates the centroids, i.e. the tied Key/Value matrix. Similar to van den Oord et al. ( 2017), we add a regularization term:

L reg = i T (Q i ) -sg(Q i ) 2

, which makes entries of the Key/Value matrix arithmetic mean of their members. Alternatively, one can also use Exponential Moving Average (Kaiser et al., 2018) to update the centroids.

A comparison between DPQ-SX and DPQ-VQ. DPQ-VQ and DPQ-SX only differ during training. They are 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 2 illustrates this difference. This suggests that when there is a large gap between one-hot and probabilistic vectors (large K), DPQ-SX approximation could be poor; and when there is a large gap between the continuous vector and the selected centroid (large subspace dimension, i.e. small D), DPQ-VQ could have a big approximation error. 

3. EXPERIMENTS:

We conduct experiments on ten datasets across three tasks: language modeling (LM), neural machine translation (NMT) and text classification (TextC) 2 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 Table 2 . We evaluate the models using two metrics: task performance and compression ratio. Task performance metrics are perplexity scores for LM tasks, BLEU scores for NMT tasks, and accuracy in TextC tasks.

Compression ratios for the embedding layer is computed as follows:

CR = # of bits used in the full embedding table # of bits used in the compressed model during inference For DPQ in particular, this can be computed as CR = 32nd nD log 2 K+32Kd . Further compression can be achieved with 'subspace-sharing' as described in Appendix E.2. With subspace-sharing, CR = 32nd nD log 2 K+32Kd/D .

3.1. COMPRESSION RATIOS AND TASK PERFORMANCE AGAINST BASELINES:

Table 3 summarizes the task performance and compression ratios of DPQ-SX and DPQ-VQ against baseline models that use the regular full embeddings3 . In each task/dataset, we report results from a 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 the largest compression ratio. In all tasks, both DPQ-SX and DPQ-VQ can achieve comparable or better task performance while providing a compression ratio from 14× to 163×. In 6 out of 10 datasets, DPQ-SX performs strictly better than DPQ-VQ in 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 methods (Chen et al., 2018b; Shu and Nakayama, 2017) . Pre-train: a three-step procedure where one firstly trains a full model, secondly learns discrete codes to reconstruct the pre-trained embedding Figure 3 : Heat-maps of task performance and compression ratio for various K and D values. Darker is better. Key observations are: 1) increasing K or D typically improves the task performance at the expense of lower CRs; 2) the combination of a small K and a 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 a pre-trained embedding table; E2E-dist.: end-to-end training with a distillation procedure that uses a pre-trained embedding as guidance during training. Table 4 shows the 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 compression ratios but with worse perplexity scores on the Medium and Large models, 2) the E2E-dist. method has the same compression ratio as them and is able to achieve similar perplexity scores as the full embedding baseline, with the downside that it requires the extra distillation procedure, 3) DPQ variants (particularly DPQ-SX) are able to 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 D:

Among key hyper-parameters of DPQ are the code size: K the number of centroids per dimension and D the code length. Figure 3 shows the task performance and compression ratios for different K and D values on PTB and IWSLT15 (En-Vi). Firstly, we observe that the combination of a small K and a large D is a 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-VQ. Secondly, increasing K or D would typically improve the task performance at the expense of lower CRs, which means one can adjust K and D to 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 as 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 4 shows the training speed as well as the (GPU) memory required when using DPQ on the medium LSTM model, trained on Tesla-V100 GPUs. 

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.1). Secondly, the KD codebook changes as training progresses, but the rate of change decreases throughout training and converges to < 20% (Appendix D.2). Thirdly, the nearest neighbours in the continuous embedding space between DPQ and the baseline align very well (Appendix D.3). Finally, we also list the learned codes for selected words in Appendix D.4.

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., 2017; Chen et al., 2018a) . 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., 2018a) . The embedding layers are different in the sense that they are tabular and very sparsely accessed, i.e. the pruning cannot remove rows/symbols in the embedding table, and only a few symbols are accessed in each data sample. This makes the compression challenges different for the embedding layers.

Existing work on compressing embedding layers includes (Shu and Nakayama, 2017; Chen et al., 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.

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-softmax (Jang et al., 2016) and VQ-VAE (van den Oord et al., 2017) . However, we do not find using stochastic noises (as in Gumbel-softmax) useful since we aim to get deterministic codes.

It is also worth pointing out that these techniques (Jang et al., 2016; van den Oord et al., 2017 ) by themselves cannot be directly applied to compression.

5. CONCLUSION:

In this work, we propose a novel and general differentiable product quantization framework for learning compact embedding layers. We provide two instantiations of our framework, which can readily serve as a 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 to 238× without suffering a performance loss. In the future, we plan to apply the DPQ framework to a wider range of applications and architectures.

A ALGORITHM PSEUDO-CODE:

This section lays out the algorithm pseudo-code for the DPQ embedding layer during the forward training/inference pass.

Algorithm 1 DPQ for the i-th token in the vocab (training, forward pass) h-params :K, D parameters :

EQUATION

i , ..., h

(D) i ) B PROOF OF THEOREM 1

Proof. We first re-parameterize both the codebook C and the Value matrix V as follows.

The original codebook is C ∈ {1, • • • , K} n×D , and we turn each code bit, which is an integer in {1, • • • , K}, into a small one-hot vector of length-K. This results in the new binary codebook B ∈ {0, 1} n×KD . Per our constraint in theorem 1, B is a full rank matrix.

The original Value matrix is V ∈ R K×d , and we turn it into a block-diagonal matrix U ∈ R KD×d where the j-th block-diagonal is set to V (j) ∈ R K×(d/D) . Given that each block diagonal, i.e. V (j) , is full rank, the resulting block diagonal matrix U is also full rank.

With the above re-parameterization, we can write the output embedding matrix H = BU. Given both B and U are full rank and KD ≥ 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 as K, D and/or subspace-sharing. We also apply batch normalization for the distance measure in DPQ along the K-dimension, i.e. each centroid will have a normalized distance distribution with batch samples.

For training the Transformer Model on WMT'19 En-De dataset, the training set contains approximately 27M parallel sentences. We generated a vocabulary of 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 were trained with a batch size of 2048 sentences for 250k steps, and with the SM3 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}. Taking the sub-word '_evolve' as an example, DPQ variants give very similar top 10 nearest neighbours as the original full embedding: both have 7 out of 10 overlapping top neighbours as the baseline model. However, in DPQ-SX the neighbours have closer distances than the baseline, hence a tighter cluster; while in DPQ-VQ the neighbours are further from the original word. We observe similar patterns in the other two examples. 6 0 0 7 0 6 1 7  1 7 0 2 0 3 1 7  _Wednesday  6 5 0 3 0 6 1 6  6 2 3 2 0 2 1 7  _Thursday  5 5 0 3 0 6 1 7  7 2 0 2 0 3 1 2  _Friday  4 6 0 7 0 6 1 7  6 0 0 2 1 6 1 7  _Saturday  4 0 6 7 0 6 1 0  6 2 0 2 3 3 1 7  _Sunday  2 0 0 3 0 6 1 6  7 2 0 2 6 3 1 7 _Obama 2 6 7 2 5 7 3 7 2 3 1 6 6 1 7 4 _Clinton 2 4 7 2 3 5 6 7 5 3 5 6 6 0 7 4 _Merkel 4 1 7 2 6 2 2 6 6 3 1 1 4 6 7 4 _Sarkozy 7 6 7 1 4 2 5 0 0 3 1 7 5 7 7 4 _Berlusconi 4 6 5 1 4 2 6 7 6 3 0 6 6 7 7 4 _Putin 2 6 7 1 6 7 6 7 5 3 1 6 6 7 7 6 _Trump 7 6 7 2 0 7 6 7 2 3 1 6 5 7 7 7 _Toronto 6 2 3 2 4 2 2 6 4 3 4 7 6 2 0 7 _Vancouver 2 1 3 2 6 2 5 6 7 3 6 6 6 2 3 1 _Ottawa 2 5 6 1 6 2 2 7 6 3 1 6 6 2 0 4 _Montreal 4 0 0 2 6 2 1 7 4 3 1 1 6 2 0 1 _London 1 2 0 2 4 7 1 7 2 3 0 2 6 3 3 7 _Paris 4 0 3 5 4 2 1 0 5 3 0 0 6 3 2 7 _Munich 4 2 0 4 0 7 5 0 1 3 3 5 6 3 1 7 Figure 7 : Heat-maps of task performance and compression ratio. Darker is better.

D CODE STUDY:



E.2 SUBSPACE-SHARING:

Subspace-sharing refers to the option of whether to share parameters among the D groups in the Key/Value Matrices, i.e. constraining K (j) = K (j ) and V (j) = V (j ) , ∀j, j . For simplicity we refer to this as "subspace-sharing". Subspace-sharing improves the compression ratio to: CR = 32nd/(nD log 2 K + 32Kd/D).

Figure 8 shows the trade-off curves of task performance and compression ratio with different DPQ variants, K, D and subspace-sharing. We find that one 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 TextC tasks, subspace-sharing seems beneficial for DPQ-SX but harmful for DPQ-VQ. • In (Chen et al., 2018b) , discrete codes are directly associated with each of the symbols, in this work, discrete codes are computed as 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 as a stand-alone module. • Our formulation of discrete codes with product quantization allows us to derive two variants with different approximation techniques (softmax-based and vector quantization-based), while (Chen et al., 2018b) is only based on softmax approximation. • The product quantization has minimal overhead and is very efficient compared to encoder functions used in (Chen et al., 2018b) , 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 are also reduced, and DPQ can be truly trained end-to-end without two pass training with distillation loss as in (Chen et 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 a CR of 32/8=4, while likely dropping in task performance metrics (e.g. PPL). • Product quantization: it generalizes scalar quantization and quantize sub-vectors. However, this approach is non-differentiable (cannot train end-to-end) and requires a post-training procedure. Small quantization errors accumulate and thus performances suffer. 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 as the whole system is jointly trained so the later networks can account for small approximation errors in the early layer. 

G.4 MORE ABLATIONS ON DPQ-SX:

Table 12 shows an ablation study on PTB language modeling task (medium-sized LSTM), in which we choose to tie the K and V matrices in DPQ-SX (achieved by sharing a single variable during the optimization process). We fix K=128, D=50. We find DPQ-SX (with untied K,V) to perform the best, followed by DPQ-SX (tied K,V) 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. BERT (Devlin et al., 2018) has shown excellent results on a wide range of natural language tasks, therefore it 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 use the same optimizer (Adam) and learning rate schedule as described in Devlin et al. (2018) . For the DPQ experiments, we use DPQ-SX with no subspace-sharing, D = 128 and K = 32; these choices are not from hyper-parameter search, but inspired from our results on Transformer on

:

32Kd/D bits if we share the weights among D groups/subspaces.

For the five text classification datasetsZhang et al. (2015), Yahoo! answers and AG news represent topic prediction, Yelp Polarity and Yelp Full represent sentiment analysis, and DBpedia represents ontology classification.

For LM, results are from the medium-sized LSTM model.

Figures and Tables:

Figure 2: Illustration of two types of approximation to enable differentiability in DPQ.

Figure 4: Extra training cost incurred by DPQ, measured on a medium sized LSTM for LM trained on Tesla-V100 GPUs. For most K and D 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).

Figure5shows the code distribution heat-maps for the Transformer model on WMT'19 En-De, with K = 32 and D = 32 and no subspace-sharing. We find that 1) DPQ-VQ has a more evenly distributed code utilization, 2) DPQ-SX has a more concentrated and sparse code distribution: in each group, only a few discrete codes are used, and some codes are not used in the codebook.

Figure 8: Task performance vs compression ratio trade-off curves. Each subplot comes from one task/dataset and contains four configurations: {DPX-SX, DPX-VQ} × {subspace-sharing, NOsubspace-sharing}.

..., K} n×D and the Value matrix V ∈ R K×d are needed to construct the embedding table.

Table content: <table><tr><td/><td>Training only</td><td>D Groups</td><td>D Groups</td><td/></tr><tr><td/><td>K</td><td>Key Matrix</td><td>Value Matrix</td><td>K</td></tr><tr><td/><td/><td>K</td><td>V</td><td/></tr><tr><td/><td>D Groups</td><td/><td/><td/></tr><tr><td/><td>Query Matrix</td><td/><td/><td/><td>Result Matrix</td></tr><tr><td>n</td><td>Q (Raw Embedding</td><td>Codebook C</td><td/><td/><td>H (Constructed</td><td>n</td></tr><tr><td/><td>Table)</td><td>Discretization</td><td colspan="2">Index &amp; Concatenation</td><td>Embedding Table)</td></tr><tr><td/><td/><td>(Eq. (1))</td><td>(Eq. (2))</td><td/></tr><tr><td/><td>d</td><td>D</td><td/><td/><td>d</td></tr><tr><td colspan="6">Figure 1: The DPQ embedding framework. During training, differentiable product quantization is</td></tr><tr><td colspan="6">used to approximate the raw embedding table (i.e. the Query Matrix). At inference, only the codebook</td></tr><tr><td>C ∈ {1,</td><td/><td/><td/><td/></tr></table>

Summary of 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.

Table content: <table><tr><td>Method</td><td>Dist. Metric</td><td>Key/Value matrices</td><td>Train</td><td>Inference</td></tr><tr><td colspan="4">DPQ-SX Dot product and more Not tied, allows different sizes Efficient</td><td>Efficient</td></tr><tr><td colspan="2">DPQ-VQ Euclidean only</td><td>Tied</td><td colspan="2">More efficient Efficient</td></tr></table>

summarizes the comparisons between DPQ-SX and DPQ-VQ. DPQ-SX is more flexible as it does not constrain the distance metric, nor does it tie the Key/Value matrices as in DPQ-VQ. Thus one could use different sizes of Key and Value matrices. Regarding to the computational cost during training, DPQ-SX back-propagates through the whole distribution of K choices, while DPQ-VQ only back-propagates through the nearest centroid, making it more scalable (to large K, D, and batch sizes).

Table content: <table/>

Datasets and models used in our experiments. More details in Appendix C.

Table content: <table><tr><td>Task</td><td>Dataset</td><td colspan="3">Vocab Size Tokenization Base Model</td></tr><tr><td>LM</td><td>PTB Wikitext-2</td><td>10,000 33,278</td><td>Words</td><td>LSTM-based models from Zaremba et al. (2014), three model sizes</td></tr><tr><td>NMT</td><td colspan="2">IWSLT15 (En-Vi) 17,191 IWSLT15 (Vi-En) 7,709</td><td>Words</td><td>Seq2seq-based model from Luong et al. (2017)</td></tr><tr><td/><td>WMT19 (En-De)</td><td>32,000</td><td>Sub-words</td><td>Transformer Base in Vaswani et al. (2017)</td></tr><tr><td/><td>AG News</td><td>69,322</td><td/><td/></tr><tr><td/><td>Yahoo! Ans.</td><td>477,522</td><td/><td>One hidden layer after mean pooling of</td></tr><tr><td>TextC</td><td>DBpedia</td><td>612,530</td><td>Words</td><td>word vectors, similar to fastText</td></tr><tr><td/><td>Yelp P</td><td>246,739</td><td/><td>from Joulin et al. (2017)</td></tr><tr><td/><td>Yelp F</td><td>268,414</td><td/><td/></tr></table>

Comparisons of DPQ variants vs. the full embedding baselines.

Table content: <table><tr><td>Task</td><td>Metric</td><td>Dataset</td><td colspan="2">Baseline DPQ-SX (CR)</td><td>DPQ-VQ (CR)</td></tr><tr><td>LM</td><td>PPL</td><td>PTB Wikitext-2</td><td>83.38 95.61</td><td>83.17 (163.2) 94.94 (59.25)</td><td>83.27 (58.67) 95.92 (95.25)</td></tr><tr><td/><td/><td>IWSLT15 (En-Vi)</td><td>25.4</td><td>25.3 (86.17)</td><td>25.3 (16.13)</td></tr><tr><td>NMT</td><td>BLEU</td><td>IWSLT15 (Vi-En)</td><td>23.0</td><td>23.1 (72.00)</td><td>22.5 (14.05)</td></tr><tr><td/><td/><td>WMT19 (En-De)</td><td>38.8</td><td>38.8 (18.00)</td><td>38.7 (18.23)</td></tr><tr><td/><td/><td>AG News</td><td>92.59</td><td>92.49 (19.26)</td><td>92.55 (23.95)</td></tr><tr><td/><td/><td>Yahoo! Ans.</td><td>69.41</td><td>69.62 (48.16)</td><td>69.15 (19.24)</td></tr><tr><td colspan="2">TextC Acc(%)</td><td>DBpedia</td><td>98.12</td><td>98.13 (24.08)</td><td>98.14 (38.45)</td></tr><tr><td/><td/><td>Yelp P</td><td>93.92</td><td>94.17 (38.52)</td><td>93.91 (24.04)</td></tr><tr><td/><td/><td>Yelp F</td><td>60.33</td><td>60.10 (48.16)</td><td>60.22 (24.05)</td></tr></table>

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).

Table content: <table><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="2">Small</td><td/><td/><td colspan="2">Medium</td><td/><td colspan="2">Large</td><td/></tr><tr><td/><td/><td colspan="2">Method</td><td/><td/><td/><td/><td colspan="2">PPL</td><td>CR</td><td/><td colspan="3">PPL CR</td><td/><td>PPL</td><td>CR</td><td/></tr><tr><td/><td/><td colspan="2">Full</td><td/><td/><td/><td/><td colspan="2">114.5</td><td>1</td><td/><td colspan="2">83.4</td><td>1</td><td/><td>78.7</td><td>1</td><td/></tr><tr><td/><td/><td colspan="6">Pre-train (Chen et al., 2018b)</td><td colspan="2">108.0</td><td>4.8</td><td/><td colspan="3">84.9 11.7</td><td/><td>80.7</td><td>18.5</td><td/></tr><tr><td/><td/><td colspan="5">E2E (Chen et al., 2018b)</td><td/><td colspan="2">108.5</td><td>4.8</td><td/><td colspan="3">89.0 11.7</td><td/><td>86.4</td><td>18.5</td><td/></tr><tr><td/><td/><td colspan="6">E2E-dist. (Chen et al., 2018b)</td><td colspan="2">107.8</td><td>4.8</td><td/><td colspan="3">83.1 11.7</td><td/><td>77.7</td><td>18.5</td><td/></tr><tr><td/><td/><td colspan="2">DPQ-SX</td><td/><td/><td/><td/><td colspan="3">105.8 85.5</td><td/><td colspan="3">82.0 82.9</td><td/><td colspan="2">78.5 238.3</td><td/></tr><tr><td/><td/><td colspan="2">DPQ-VQ</td><td/><td/><td/><td/><td colspan="3">106.5 51.1</td><td/><td colspan="3">83.3 58.7</td><td/><td colspan="2">79.5 238.3</td><td/></tr><tr><td>2</td><td colspan="2">-104.9 -91.2</td><td>-89.0</td><td colspan="2">-123.0 -92.9</td><td>-86.2</td><td colspan="3">2000.0 827.8 415.9</td><td>2</td><td>7.5</td><td>20.0</td><td>25.4</td><td>9.0</td><td>13.6</td><td>23.4</td><td colspan="3">1657.2 484.2 126.4</td></tr><tr><td>8</td><td>-88.3</td><td>-84.7</td><td>-84.5</td><td>-96.1</td><td>-85.6</td><td>-83.7</td><td colspan="3">657.9 275.3 138.6</td><td>8</td><td>19.0</td><td>24.7</td><td>26.0</td><td>13.6</td><td>16.7</td><td>20.5</td><td colspan="2">519.1 158.4</td><td>41.9</td></tr><tr><td>K</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>K</td><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td>32</td><td>-86.4</td><td>-83.2</td><td>-82.0</td><td>-91.4</td><td>-84.5</td><td>-83.4</td><td colspan="2">367.7 163.2</td><td>82.9</td><td>32</td><td>21.5</td><td>25.3</td><td>25.3</td><td>14.7</td><td>19.8</td><td>24.5</td><td>232.9</td><td>86.2</td><td>24.5</td></tr><tr><td>128</td><td>-85.3</td><td>-83.8</td><td>-82.0</td><td>-91.6</td><td>-85.1</td><td>-83.3</td><td colspan="2">215.6 112.2</td><td>58.6</td><td>128</td><td>21.9</td><td>24.2</td><td>24.9</td><td>18.1</td><td>21.6</td><td>25.3</td><td>92.2</td><td>47.4</td><td>16.1</td></tr><tr><td/><td>10</td><td>25</td><td>50</td><td>10</td><td>25</td><td>50</td><td>10</td><td>25</td><td>50</td><td/><td>8</td><td>32</td><td>128</td><td>8</td><td>32</td><td>128</td><td>8</td><td>32</td><td>128</td></tr><tr><td/><td colspan="3">D neg-PPL (DPQ-SX)</td><td colspan="3">D neg-PPL (DPQ-VQ)</td><td/><td>D CR</td><td/><td/><td colspan="3">D BLEU (DPQ-SX)</td><td colspan="3">D BLEU (DPQ-VQ)</td><td/><td>D CR</td></tr><tr><td/><td/><td/><td/><td colspan="2">(a) PTB</td><td/><td/><td/><td/><td/><td/><td/><td colspan="4">(b) IWSLT15 (En-Vi)</td><td/><td/></tr></table>

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 on speed or memory from DPQ.

Table content: <table><tr><td>32</td><td>5.2%</td><td>7.1%</td><td>7.7%</td><td>9.1%</td><td>6.7%</td><td>6.0%</td><td>32</td><td>0.0%</td><td>0.0%</td><td>0.0%</td><td>0.0%</td><td>0.0%</td><td>0.0%</td></tr><tr><td>128 K</td><td colspan="3">8.1% 11.0% 13.1%</td><td>9.0%</td><td>9.6%</td><td>8.2%</td><td>128 K</td><td>0.0%</td><td>0.0%</td><td>0.0%</td><td>0.0%</td><td>0.0%</td><td>0.0%</td></tr><tr><td>512</td><td colspan="3">9.9% 23.6% 46.9%</td><td colspan="3">10.2% 16.8% 37.5%</td><td>512</td><td>0.0%</td><td colspan="2">0.0% 83.2%</td><td>0.0%</td><td>0.0%</td><td>4.9%</td></tr><tr><td/><td>10</td><td>50 D DPQ-SX</td><td>130</td><td>10</td><td>50 D DPQ-VQ</td><td>130</td><td/><td>10</td><td>50 D DPQ-SX</td><td>130</td><td>10</td><td>50 D DPQ-VQ</td><td>130</td></tr><tr><td/><td colspan="5">(a) Extra training time used.</td><td/><td/><td colspan="6">(b) Extra training memory used.</td></tr></table>

Nearest neighbours of '_evolve' in the embedding space.

Table content: <table><tr><td>Baseline (Full)</td><td>Dist DPQ-SX</td><td>Dist DPQ-VQ</td><td>Dist</td></tr><tr><td>_evolve</td><td>1.000 _evolve</td><td>1.000 _evolve</td><td>1.000</td></tr><tr><td>_evolved</td><td>0.533 _evolved</td><td>0.571 _evolved</td><td>0.506</td></tr><tr><td>_evolving</td><td>0.493 _evolution</td><td>0.499 _develop</td><td>0.417</td></tr><tr><td>_develop</td><td>0.434 _develop</td><td>0.435 _evolving</td><td>0.359</td></tr><tr><td>_evolution</td><td>0.397 _evolving</td><td>0.418 _developed</td><td>0.320</td></tr><tr><td>_developed</td><td>0.379 _arise</td><td colspan="2">0.405 _development 0.307</td></tr><tr><td>_developing</td><td colspan="2">0.316 _developed 0.405 _developing</td><td>0.299</td></tr><tr><td>_arise</td><td>0.298 _resulted</td><td>0.394 _evolution</td><td>0.282</td></tr><tr><td>_unfold</td><td>0.294 _originate</td><td>0.361 _changed</td><td>0.278</td></tr><tr><td>_emerge</td><td>0.290 _result</td><td>0.359 _grew</td><td>0.273</td></tr></table>

Nearest neighbours of '_monopoly' in the embedding space.

Table content: <table><tr><td>Baseline</td><td>Dist DPQ-SX</td><td>Dist DPQ-VQ</td><td>Dist</td></tr><tr><td>_monopoly</td><td>1.000 _monopoly</td><td>1.000 _monopoly</td><td>1.000</td></tr><tr><td colspan="2">_monopolies 0.613 _monopolies</td><td colspan="2">0.762 _monopolies 0.509</td></tr><tr><td>monopol</td><td>0.552 monopol</td><td>0.714 monopol</td><td>0.483</td></tr><tr><td>_Monopol</td><td>0.380 _Monopol</td><td>0.531 _Monopol</td><td>0.341</td></tr><tr><td colspan="2">_moratorium 0.271 _zugestimmt</td><td>0.486 _dominant</td><td>0.258</td></tr><tr><td>_privileged</td><td>0.269 legitim</td><td colspan="2">0.420 _moratorium 0.239</td></tr><tr><td>_unilateral</td><td colspan="2">0.262 _Großunternehmen 0.401 _autonomy</td><td>0.230</td></tr><tr><td>_miracle</td><td>0.260 _Eigenkapital</td><td>0.400 _zugelassen</td><td>0.227</td></tr><tr><td>_privilege</td><td>0.254 _wirkungsvoll</td><td>0.399 _imperial</td><td>0.226</td></tr><tr><td>_dominant</td><td>0.250 _UCLAF</td><td>0.388 _capitalist</td><td>0.223</td></tr></table>

Nearest neighbours of '_Toronto' in the embedding space.Figure7shows extra heatmaps with varied K and D in addition to those in Section 3.2.

Table content: <table><tr><td>Baseline</td><td>Dist DPQ-SX</td><td>Dist DPQ-VQ</td><td>Dist</td></tr><tr><td>_Toronto</td><td>1.000 _Toronto</td><td>1.000 _Toronto</td><td>1.000</td></tr><tr><td>_Vancouver</td><td>0.390 _Chicago</td><td>0.475 _Orlando</td><td>0.307</td></tr><tr><td>_Tokyo</td><td>0.378 _Orleans</td><td>0.467 _Detroit</td><td>0.306</td></tr><tr><td>_Ottawa</td><td colspan="2">0.372 _Melbourne 0.435 _Canada</td><td>0.280</td></tr><tr><td colspan="2">_Philadelphia 0.353 _Miami</td><td>0.434 _London</td><td>0.280</td></tr><tr><td>_Orlando</td><td colspan="2">0.345 _Vancouver 0.415 _Glasgow</td><td>0.276</td></tr><tr><td>_Chicago</td><td>0.340 _Tokyo</td><td>0.407 _Montreal</td><td>0.272</td></tr><tr><td>_Canada</td><td>0.330 _Ottawa</td><td>0.405 _Vancouver</td><td>0.271</td></tr><tr><td>_Seoul</td><td>0.329 _Azeroth</td><td colspan="2">0.403 _Philadelphia 0.267</td></tr><tr><td>_Boston</td><td>0.325 _Antonio</td><td>0.400 _Hamilton</td><td>0.264</td></tr><tr><td>D.4 CODE VISUALIZATION</td><td/><td/><td/></tr><tr><td colspan="4">Table 8 shows some examples of compressed codes for both DPQ-SX and DPQ-VQ. Semantically</td></tr><tr><td colspan="4">related words share common codes in more dimensions than unrelated words.</td></tr><tr><td colspan="3">E ADDITIONAL HYPER-PARAMETERS STUDY</td><td/></tr><tr><td>E.1 EFFECTS OF K AND D</td><td/><td/><td/></tr></table>

Examples of KD codes.

Table content: <table><tr><td/><td>DPQ-SX</td><td>DPQ-VQ</td></tr><tr><td>_Monday</td><td>2 5 0 7 0 6 1 6</td><td>6 5 0 2 4 3 1 7</td></tr><tr><td>_Tuesday</td><td/><td/></tr></table>

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.

Table content: <table><tr><td>Dataset</td><td>AG News</td><td>Yahoo!</td><td>DBPedia</td><td>Yelp P</td><td>Yelp F</td></tr><tr><td>Full</td><td>92.6 (1.0)</td><td>69.4 (1.0)</td><td>98.1 (1.0)</td><td>93.9 (1.0)</td><td>60.3 (1.0)</td></tr><tr><td>Low-rank(10×)</td><td colspan="5">91.4 (10.4) 69.5 (10.2) 97.7 (10.3) 92.4 (10.4) 57.8 (10.3)</td></tr><tr><td>Low-rank(20×)</td><td colspan="5">91.5 (21.4) 69.1 (21.5) 97.9 (21.3) 92.4 (21.5) 57.3 (21.4)</td></tr><tr><td colspan="6">Chen et al. (2018b) 91.6 (53.3) 69.5 (31.7) 98.0 (48.4) 93.1 (48.6) 59.0 (54.4)</td></tr><tr><td>DPQ-VQ</td><td colspan="5">92.6 (24.0) 69.2 (19.2) 98.1 (38.5) 93.9 (24.0) 60.2 (24.1)</td></tr><tr><td>DPQ-SX</td><td colspan="5">92.5 (19.3) 69.6 (48.2) 98.1 (24.1) 94.2 (38.5) 60.1 (48.2)</td></tr></table>

Performance comparisons against the reconstruction baselines.

Table content: <table><tr><td>Method</td><td>BLEU</td><td>CR</td></tr><tr><td>Full</td><td>38.8</td><td>1</td></tr><tr><td>Reconstruction (K=128, D=64)</td><td>28.9</td><td>31.9</td></tr><tr><td>Reconstruction (K=32, D=128)</td><td>35.4</td><td>25.0</td></tr><tr><td>Reconstruction (K=128, D=128)</td><td>35.7</td><td>17.0</td></tr><tr><td>Reconstruction (K=32, D=256)</td><td>36.9</td><td>12.6</td></tr><tr><td>Reconstruction (K=128, D=256)</td><td>37.8</td><td>8.8</td></tr><tr><td>DPQ-VQ (K=32, D=128)</td><td>38.7</td><td>17.0</td></tr><tr><td>DPQ-SX (K=32, D=128)</td><td>38.8</td><td>17.0</td></tr></table>

Peer review (iclr2020):
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. I’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. Besides, for the technical part, DPQ-SX outperforms DPQ-VQ while the softmax approximation seems identical to that developed in [1]. [1] ICML’18: Learning K-way D-dimensional Discrete Codes for Compact Embedding Representations