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

TL;DR

This paper introduces a novel autoencoder-based method focusing on embedding direction to effectively compress token embeddings in Transformer models, improving performance over traditional methods without additional pre-training.

Contribution

It proposes a task-agnostic embedding compression technique emphasizing directionality, outperforming SVD-based methods in language modeling and downstream tasks.

Findings

Outperforms SVD-based matrix factorization in perplexity

Achieves better results on SQuAD v1.1 and GLUE tasks

Does not require additional language model pre-training

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.