Direction is what you need: Improving Word Embedding Compression in Large Language Models (2021)

Abstract: The adoption of Transformer-based models in natural language processing (NLP) has led to great success using a massive number of parameters. However, due to deployment constraints in edge devices, there has been a rising interest in the compression of these models to improve their inference time and memory footprint. This paper presents a novel loss objective to compress token embeddings in the Transformer-based models by leveraging an AutoEncoder architecture. More specifically, we emphasize the importance of the direction of compressed embeddings with respect to original uncompressed embeddings. The proposed method is task-agnostic and does not require further language modeling pre-training. Our method significantly outperforms the commonly used SVD-based matrix-factorization approach in terms of initial language model Perplexity. Moreover, we evaluate our proposed approach over SQuAD v1.1 dataset and several downstream tasks from the GLUE benchmark, where we also outperform the baseline in most scenarios. Our code is public.

## introduction pretraining deep transformer models @xcite with language modeling and fine-tuning these models over downstream tasks have led to great success in recent years @xcite @xcite , and even enabled researchers to design models that outperform human baselines in the glue benchmark @xcite . although these models are empirically powerful in many figure 1 : this figure presents a two-dimensional visualization of a token embedding vector v with its two approximations: v and v . vector v has a larger euclidean distance error than v , but its direction is more similar to the reference vector. our experiments show that v generally provides a better approximation of the original token compared to v . natural language understanding (nlu) tasks, they often require a massive number of parameters, making them hard to use for memory-constrained applications (e.g., edge devices). therefore, there have been efforts to compress bert-like models while preserving comparable performance with the original model. many of these compression methods are based on knowledge distillation @xcite to help the compressed model (student) to perform close to the original model in different nlu tasks. however, these approaches often need high computation resources due to e.g., the necessity of retraining the expensive language modeling on a huge corpus @xcite or the use of expensive augmentation techniques to make the distillation effectively work @xcite . moreover, compression techniques that rely on training/fine-tuning language models are becoming less feasible due to its ever-increasing cost for current state-of-the-art architectures with hundreds of millions of parameters @xcite @xcite . more recently, there have been efforts to compress transformer-based models for more resource-constrained scenarios @xcite by using offline methods, such as matrix factorization @xcite @xcite , weight pruning @xcite , and also weight quantization @xcite . this paper focuses on token embedding matrix compression due to being one of the largest matrices in bert-based architectures. we specifically question the effectiveness of current low-rank matrix factorization methods in recent literature @xcite by comparing them with the performance of a linear autoencoder over different compression ratios 2 . we define a new loss objective which is not only dependent on the commonly used mean absolute error (mae) or mean squared error (mse) loss between input embeddings and autoencoder reconstruction, but is also sensitive to the noise in reconstructed embeddings "direction" (measured by cosine distance). we present the intuition behind the importance of embedding vector direction in the figure 1 . in the following sections we show that cosine distance indeed plays a more critical role than mae/mse (figure 3 ) as measured by the perplexity of the entire model in language modeling. in section 4, we demonstrate that our compression algorithm is superior or competitive to the singular value decomposition (svd) baseline over several natural language understanding tasks from glue @xcite benchmark, as well as the squad dataset @xcite for question answering. we also compare our performance with the svd-based compression over different compression ratios, and specifically show that our model performs consistently better in higher compression ratios. our contribution can be summarized as follows: the current mostly used compression methods can be roughly categorized into four classes, namely knowledge distillation @xcite , weight pruning @xcite , matrix factorization @xcite @xcite and weight quantization @xcite . this section focuses on matrix factorization-based methods that are currently used for token embedding compression in the literature. ## model description although svd matrix-factorization is one of the most popular methods for matrix compression, we believe it is not an optimal method for compressing token embeddings in bert-like architectures. the objective of svd is to minimize the l 2 norm between the original matrix and the reconstructed one; however, focusing on l 2 norm optimization prioritizes the reduction of larger errors, and it may end up ignoring more minor vector differences. it is also sensitive to the influence of outliers. the most crucial reason for the l 2 norm not being the best choice is that it only considers the distance between the original and reconstructed token vector, and it does not necessarily pay attention to the orientation difference between them. in section 4, we demonstrate that vectors representing language tokens are more sensitive to noise in their direction rather than to changes in euclidean distance from the reference vector. we also discuss the motivation behind it further in this section. in order to mitigate the problem of focusing only on the largest errors between two vectors, we propose replacing the l 2 norm objective with the l 1 norm raised to the power of α: taking into consideration all points above, we propose to replace the l 2 norm objective with a new multi-objective function consisting of l 1 norm (raised to the power of α, where α is a hyper-parameter that can be changed during training) and cosine distance: φ α,β (a, a) = l α 1 (a, a) + β * c d(a, a), (5) in addition to the new loss function, we propose leveraging auto-encoder architecture for φ α,β and ψ β loss optimization (equation 5 and 6). we use a simple autoencoder consisting of a one-layer encoder/decoder without any activation function in order to have a fair comparison with the svd baseline. using auto-encoder enables efficient multi-objective optimization, but it also allows to select the appropriate level of model complexity when needed. at the end of the auto-encoder training, we extract an approximation of the original matrix, as shown in figure 2 . we substitute the original embedding matrix with a new module consisting of latent representation of vocabulary tokens along with the decoder module. ## results in this section, we evaluate our approach, which is based on using autoencoder model with a multi-objective loss function that incorporates cosine distance with l 1 or l 2 norm (equation 5 and equation 6 ) on the task of bert-like token embedding matrix compression. we compare our results versus the commonly used randomized svd method @xcite to perform low-rank matrix factorization. we have implemented our token embeddings compression with the pytorch backend @xcite and as an extension of huggingface's transformers library @xcite , enabling researchers to apply our compression method in most of the existing transformer architectures. it is worth noting that 5 ) and ψ β (equation 6 ). in all configurations we select a final model based on the best perplexity achieved during training. the term [t 1 , t 2 ] indicates linearly decreasing α parameter (equation 4 ). setting β = 0 represents not including cosine distance component in the loss function. we may observe that not including cosine distance in the loss function as well as making it a too dominant component (very big β) is not optimal for achieving good perplexity. we also present the best perplexity achieved by the baseline svd method for three compression ratios: 2.5, 5.0, 10.0. our approach significantly outperforms the baseline in the studied scenarios. the offline training of our compression method on bert-base @xcite token embedding matrix takes only few minutes on a single gpu device. ## conclusion in this work, we propose a simple linear autoencoder model with a multi-objective loss function for bert-like token embeddings compression. we emphasize the importance of the direction component (measured by the cosine distance between the original and the reconstructed token embeddings) in the compression objective function. we challenge the commonly used svd-based matrix-factorization method and show that our approach achieves significantly better zero-shot language model perplexity. moreover, we show that bert-like models with our compressed token embeddings submodule converge much faster and outperform the svd baseline on squad v1.1 and on glue benchmark tasks in most scenarios.