Direction is what you need: Improving Word Embedding Compression in Large Language Models
This work addresses deployment constraints for large language models on edge devices, offering a task-agnostic compression method that is incremental but shows strong specific gains.
The paper tackles the problem of compressing token embeddings in Transformer-based models to reduce memory and improve inference time, presenting a novel loss objective that leverages an AutoEncoder to emphasize embedding direction, which significantly outperforms SVD-based methods in language model perplexity and most downstream tasks.
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.