I mproving W ord E mbedding F actorization for C ompression U sing D istilled N onlinear N eural D ecomposition (2020)

Abstract: Word-embeddings are vital components of Natural Language Processing (NLP) models and have been extensively explored. However, they consume a lot of memory which poses a challenge for edge deployment. Embedding matrices, typically, contain most of the parameters for language models and about a third for machine translation systems. In this paper, we propose Distilled Embedding, an (input/output) embedding compression method based on low-rank matrix decomposition and knowledge distillation. First, we initialize the weights of our decomposed matrices by learning to reconstruct the full pre-trained word-embedding and then fine-tune end-to-end, employing knowledge distillation on the factorized embedding. We conduct extensive experiments with various compression rates on machine translation and language modeling, using different data-sets with a shared word-embedding matrix for both embedding and vocabulary projection matrices. We show that the proposed technique is simple to replicate, with one fixed parameter controlling compression size, has higher BLEU score on translation and lower perplexity on language modeling compared to complex, difficult to tune state-of-the-art methods.

## introduction deep learning models are the state-of-the-art in nlp, computer vision, speech recognition and many other fields in computer science and engineering. the remarkable deep learning revolution has been built on top of massive amounts of data (both labeled and unlabeled), and faster computation. in nlp, large pre-trained language models like bert @xcite are state-of-theart on a large number of downstream nlp prob-lems. the largest publicly available language models are trained with hundred of billions of parameters @xcite . in machine translation the state-of-the-art models have parameters in the order of billions. data privacy and server cost are some major issues, driving research towards deploying these models on edge-devices. however, running these models on edge-devices, faces memory and latency issues due to limitations of the hardware. thus, there has been considerable interest towards research in reducing the memory footprint and faster inference speed for these models @xcite @xcite @xcite . the architecture of deep-learning-based language generation models can be broken down into three components. the first component, represents the embedding, which maps words in the vocabulary to continuous dense vector representations of the words. in language modeling we typically have one dictionary but machine translation has at least two dictionaries corresponding to a translation pair. we model these as a single dictionary with a common embedding matrix. the second component, consists of a function f , typically a deep neural-network @xcite @xcite which maps the embedding representation for different nlp problems (machine-translation, summarization, question-answering and others), to the output-space of function f . the third component, is the output layer which maps the output of function f to the vocabulary-space, followed by a softmax function. since, the first and third components depend upon a large vocabulary-size, they require large number of parameters which results in higher latency and larger memory requirements. for instance, the transformer base model @xcite uses 37% of the parameters in the first and third components using a vocabulary size of 50k, and with parameter-tying between the components. the percentage of parameters increases to 54%, when parameters are not shared between the first and third components. thus, an obvious step towards model compression is to reduce the parameters used by the embedding matrices. recently, there has been considerable work on compressing word-embedding matrices @xcite @xcite @xcite . these techniques have proven to perform atpar with the uncompressed models, but still suffer from a number of issues. first, state-of-the-art embedding compression methods such as groupreduce, structured emebedding and tensor train decomposition @xcite @xcite , require multiple hyper-parameters to be fine-tuned to optimize performance on each dataset. these hyper-parameters influence the number of parameters in the model, and thus the compression rate. this leads to an additional layer of complexity for optimizing the model for different nlp problems. additionally, @xcite requires an additional optimization step for grouping words, and lacks end-to-end training through back-propagation. @xcite also requires an additional step for performing k-means clustering for generating the quantization matrix. thus, most of the current state-of-theart systems are much more complicated to fine-tune for different nlp problems and data-sets. second, all the state-of-the-art embedding compression models compress the input and output embedding separately. in practice, state-of-the-art nlp models @xcite have shown better performance with parameter sharing between the two (press and wolf, 2017). thus, there is a need for an exhaustive analysis of various embedding compression techniques, with parameter sharing. lastly, embedding compression models not based on linear svd @xcite require the reconstruction of the entire embedding matrix or additional computations, when used at the output-layer. thus during runtime, the model either uses the same amount of memory as the uncompressed model or pays a higher computation cost. this makes linear svd based techniques more desirable for running models on edge-devices. in this paper, we introduce distilled embedding, a matrix factorization method, based on singular value decomposition (svd) with two key changes a) a neural network decomposition instead of an eigenvalue decomposition and b) a distillation loss on the word embedding while fine-tuning. our method, first compresses the vocabulary-space to the desired size, then applies a non-linear activation function, before recovering the original embeddingdimension. additionally, we also introduce an embedding distillation method, which is similar to knowledge distillation @xcite but we apply it to distill knowledge from a pre-trained embedding matrix and use an l2 loss instead of cross-entropy loss. to summarize, our contributions are as follows: ## related work we can model the problem of compressing the embedding matrix as a matrix factorization problem. there is a considerable amount of work done in this field and some of the popular methods include singular value decomposition (svd) @xcite , product quantization @xcite and tensor decomposition @xcite . a number of prior works in embedding compression are influenced by these fields and have been applied to various nlp problems. in this section, we will discuss some of the significant works across different nlp problems. low-rank factorization low-rank approximation of weight matrices, using svd, is a natural way to compress deep learning based nlp models. @xcite apply this to a convolutional neural network for language modeling and acoustic modeling. winata et al. (2019) use svd on all the weight matrices of an lstm and demonstrate competitive results on question-answering, language modeling and text-entailment. acharya et al. (2019) use low-rank matrix factorization for word-embedding layer during training to compress a classification model. however, they do not study the effects of applying a non-linear function before reconstructing the original dimension. groupreduce chen et al. ( 2018 ) apply weighted low-rank approximation to the embedding matrix of an lstm. they first create a many-to-one mapping of all the words in the vocabulary into g groups based upon word frequency. for each group g they apply weighted svd to obtain a lower rank estimation, the rank is determined by setting a minimum rank and linearly increasing it based upon average frequency. finally, they update the groups by minimizing the reconstruction error from the weighted svd approximation. they demonstrate strong results on language modeling and machine translation compared to simple svd. in their models they use different embedding matrices for input and softmax layers and apply different compression ratios to each. product quantization @xcite introduced product quantization for compressing high dimensional vectors, by uniformly partitioning them into subvectors and quantizing each subvector using k-means clustering technique. basically, product quantization assumes that the subvectors share some underlying properties which can be used to group similar ones together and unify their representation. that being said, this approach breaks the original matrix into a set of codebooks coming from the center of the clusters in different partitions together with a separate index matrix which refers to the index of the clusters for each subvector. @xcite applied product quantization to a language model and were able to show better perplexity scores. @xcite extended this technique by first representing the product quantization as a matrix factorization problem, and then learning the quantization matrix in an end-to-end trainable neural network. @xcite implement product quantization through randomly sharing parameters in the embedding matrix, and show good results on perplexity for an lstm based language model. tensor @xcite introduced multilinear svd, which is a generalization of svd for higher order tensors. oseledets (2011) introduced an efficient algorithm tensor train (tt) for multilinear svd tensor. @xcite applied the tensor train decomposition on fully connected layers of deep neural networks. @xcite applied tensor train algorithm to the input embedding layer on different nlp problems like language modeling, machine translation and sentiment analysis. they demonstrate high compression rate with little loss of performance. however, they compress only the input embedding and not the softmax layer for language modeling and machine translation. knowledge distillation knowledge distillation @xcite . has been studied in model compression where knowledge of a large cumbersome model is transferred to a small model for easy deployment. in this paper, we propose an embedding factorization of wordembedding matrix using knowledge distillation to mimic the pre-trained word-embedding representation. ## conclusion and future work in this paper we proposed distilled embedding, a low-rank matrix decomposition with nonlinearity in the bottleneck layer for a shared wordembedding and vocabulary projection matrix. we also introduce knowledge distillation of the embedding during fine-tuning using the full embedding matrix as the teacher and the decomposed embedding as the student. we compared our proposed approach with state-of-the-art methods for compressing word-embedding matrix. we did extensive experiments using three different sizes of datasets and showed that our approach outperforms the state-of-the art methods on the challenging task of machine translation. our method also generalized well to the task of language modeling. for future work, we will apply our approach to compress feed-forward and multi-head attention layers of the transformer network.