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. 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. 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.
We propose a differentiable product quantization framework that can reduce the size of embedding layer in an end-to-end training at no performance cost. This paper works on methods for compressing embedding layers for low memory inference, where compressed embeddings are learned together with the task-specific models in a differentiable end-to-end fashion.