Tensorized Embedding Layers (2020)

Abstract: The embedding layers transforming input words into real vectors are the key components of deep neural networks used in natural language processing. However, when the vocabulary is large, the corresponding weight matrices can be enormous, which precludes their deployment in a limited resource setting. We introduce a novel way of parameterizing embedding layers based on the Tensor Train decomposition, which allows compressing the model significantly at the cost of a negligible drop or even a slight gain in performance. We evaluate our method on a wide range of benchmarks in natural language processing and analyze the trade-off between performance and compression ratios for a wide range of architectures, from MLPs to LSTMs and Transformers.

## introduction deep neural networks (dnns) typically used in natural language processing (nlp) employ large embeddings layers, which map the input words into continuous representations and usually have the form of lookup tables. despite such simplicity, these layers often occupy a large portion of model weights which may cause problems in training and deploying them in a limited resource setting. thus, the compression of large neural networks and the development of novel lightweight architectures have become essential problems in nlp research. one way to reduce the number of parameters in the trained model is to imply a specific structure on its weight matrices, e.g., assume that they are low-rank or can be well approximated by lowrank tensor networks @xcite . such approaches are successful at compressing the pretrained models, but they do not facilitate the training itself. furthermore, they usually require an * the first three authors contributed equally to this work additional fine-tuning stage to recover the performance of the original model @xcite . in this paper, we introduce a new, parameter efficient embedding layer, termed tt-embedding, which can be plugged in into any model and trained end-to-end. the benefits of our compressed ttlayer are twofold. firstly, instead of storing huge embedding matrix, we store a sequence of much smaller 2-dimensional and 3-dimensional tensors, necessary for reconstructing the required embeddings, which allows compressing the model significantly at the cost of a negligible performance drop. secondly, the overall number of parameters can be relatively small (and constant) during training. the main contributions of our paper are: ## related work in recent years, a large body of research was devoted to compressing and speeding up various components of neural networks used in nlp tasks. @xcite adapted the framework of product quantization to reduce the number of parameters in linear models used for text classification. @xcite proposed to compress lstm-based neural machine translation models with pruning algorithms. @xcite showed that the recurrent models could be significantly sparsified with the help of variational dropout @xcite . @xcite successfully compressed the transformer architecture with the combination of pruning and quantization. there is a plethora of prior work on compressing the embedding layers used in nlp models. @xcite proposed more compact kway d-dimensional discrete encoding scheme to replace the "one-hot" encoding of categorical features, such as words in nlp taks. @xcite introduced west, a compression method based on structured sparse and structured dense decomposition of the embedding matrix. @xcite proposed to compress the pre-trained embedding matrix by capitalizing on the powerlaw distribution of words and using smaller dimensionality (lower rank) for the embeddings of less frequent words. @xcite used a similar idea in end-to-end fashion by training such structured low-rank embeddings from scratch. however, both of these methods rely on the assumption of power-law distribution of tokens and are not efficient when dealing with other popular tokenizations, such as wordpieces @xcite or bpes @xcite . the effectiveness of simple low-rank factorized embeddings has been recently re-discovered by @xcite , and we refer to this method as to important baseline. also, lam (2018) proposed a quantization algorithm for compressing word vectors, but its benefits are orthogonal to those of lowrank matrix and tensor factorizations and they can be used together, complementing each other. tensor methods have also been already successfully applied to neural networks compression. @xcite coined the idea of reshaping weights of fully-connected layers into high-dimensional tensors and representing them in tensor train (tt) @xcite format. this approach was later extended to convolutional @xcite and recurrent @xcite @xcite neural networks. furthermore, @xcite showed that convolutional layers could be also compressed with canonical (cp) tensor decomposition @xcite . finally, @xcite compressed both fully-connected and convolutional layers with tensor ring decomposition @xcite . recently, @xcite succesfully applied block-term tensor decomposition to the compression of self-attention modules in the transformer @xcite architecture. in this work, we show the benefits of applying tensor machinery to the compression of embedding layers, which are an essential component of all models used in nlp. ## motivation since most of the parameters in the nlp models occupy the embedding layers, we can greatly reduce size of the entire model by compressing these layers. our goal is to replace the standard embedding matrix with a more compact, yet powerful and trainable, representation which would allow us to efficiently map words into vectors. in this section, we briefly discuss our motivation of using tensorized embedding layers instead of both standard embedding layers and their low-rank factorized counterpart. ## tensor train embedding in this section, we briefly introduce the necessary notation and present the algorithm for training the tt-embedding layer. hereinafter, by n -way tensor x we mean a multidimensional array: ## experiments code we have implemented tt-embeddings described in section 4 in python using pytorch @xcite . the code is available at the anonymous repository https://github.com/ttembedding/tt-embeddings . experimental setup we tested our approach on several popular nlp tasks: to prove the generality and wide applicability of the proposed approach, we tested it on various architectures, such as mlps (ctr), lstms (sentiment analysis), and transformers (nmt, lm). the baselines we compare with are 1. standard embedding layer parametrized by a matrix @xmath0. note that transformers in lm and nmt use the same weight matrix for their embedding and softmax layers @xcite which already significantly reduces model size. untying weights and tensorizing the embedding layer only will lead to the increase in the number of parameters instead of compression. in our experiments, we use two separate tt-decompositions of the same shape for embedding and softmax layers and report the compression ratios as |v |×d model 2×tt-params . ## click through rate prediction among other applications of the tt-embedding layer, we chose to focus on ctr prediction, a pop-4 https://github.com/kimiyoung/transformer-xl ular task in digital advertising @xcite . we consider open dataset provided by criteo for kaggle display advertising challenge (criteo labs, 2014) which consists of 39 categorical features, 45.8m samples and is binary labeled according to whether the user clicked on the given advertisement. unique values of categorical features are bijectively mapped into integers. to reduce the memory footprint, if the size of a corresponding vocabulary is immense (e.g., a cardinality of some features in this dataset is of order 10 6 ), these integers are further hashed by taking modulus with respect to some fixed number such as 10 5 . however, due to strong compression properties of tt-embeddings, this is not necessary for our approach, and we consider both full and hashed datasets in our experiments. ctr with the baseline algorithm the task at hand can be treated as a binary classification problem. as a baseline algorithm, we consider the neural network with the following architecture. first, each of the categorical features is passed through a separate embedding layer with embedding size j. after that, the embedded features are concatenated and passed through 4 fully-connected layers of 1024 neurons and relu activation functions. in all experiments, we used adam optimizer with the learning rate equal to 0.0005. since many input features have a large number of unique values (e.g., 10131227) and storing the corresponding embedding matrices would be costly, we employ the hashing procedure mentioned earlier. ctr with tt-embeddings we substitute the embedding layers with the tt-embedding layers. besides that, we leave the overall structure of the neural network unchanged with the same parameters as in the baseline approach. to the best of our knowledge, our loss value is very close to the state-of-the-art result @xcite . these experiments indicate that the substitution of large embedding layers with tt-embeddings leads to significant compression ratios (up to 2011 times) with a slight improvement in the test loss, and up to 4200 with a small drop in the test loss. the total size of the compressed model does not exceed 20 mb, while the baseline model weighs about 160 mb. the obtained compression ratio suggests that the usage of tt-embedding layers may be beneficial in ctr prediction. ## discussion and future work we propose a novel embedding layer, the ttembedding, for compressing huge lookup tables used for encoding categorical features of significant cardinality, such as the index of a token in natural language processing tasks. the proposed approach, based on the tt-decomposition, experimentally proved to be effective, as it heavily decreases the number of training parameters at the cost of a small deterioration in performance. in addition, our method can be easily integrated into any deep learning framework and trained via backpropagation, while capitalizing on reduced memory requirements and increased training batch size. our experimental results suggest several appealing directions for future work. first of all, ttembeddings impose a concrete tensorial low-rank structure on the embedding matrix, which was shown to improve the generalization ability of the networks acting as a regularizer. the properties and conditions of applicability of this regularizer are subject to more rigorous analysis. secondly, unlike standard embedding, we can introduce nonlinearity into tt-cores to improve their expressive power @xcite . additionally, it is important to understand how the order of tokens in the vocabulary affects the properties of the networks with tt-embedding. we hypothesize that there exists the optimal order of tokens which better exploits the particular structure of tt-embedding and leads to a boost in performance and/or compression ratio. finally, the idea of applying higherorder tensor decompositions to reduce the number of parameters in neural nets is complementary to more traditional methods such as pruning @xcite and quantization @xcite . thus, it would be interesting to make a thorough comparison of all these methods and investigate whether their combination may lead to even stronger compression.