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