Learning Compact Embedding Layers via Differentiable Product Quantization

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 FW : V → Rd , where V denotes the vocabulary of discrete symbols , andW ∈ Rn×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 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 tableW . 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 : Rd → Rd . 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 φ ( · ) : Rd → { 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 → Rd 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 ∈ Rn×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 ∈ Rn×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 ∈ RK×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 ) ∈ RK×d/D and Q ( j ) ∈ Rn×d/D . We can compute each of D dimensional KD codes separately . The j-th dimension of a KD code Ci for the i-th symbol is computed as follows . C ( j ) i = arg min k dist ( Q ( j ) i , K ( j ) k ) ( 1 ) The dist ( · , · ) computes distance measure between two vectors , and use it to decide which discrete code to take . 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 ∈ RK×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 ) ∈ RK×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 . Hi = [ V ( 1 ) c ( 1 ) i , · · · , V ( j ) c ( j ) i , · · · , V ( D ) c ( D ) i ] ( 2 ) 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 log2K 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 log2K < 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 ) ∈ RK×d/D , ∀j , are all full-rank . 132Kd/D bits if we share the weights among D groups/subspaces . 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 log2K < 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 .

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