LLM Vocabulary Compression for Low-Compute Environments

Sreeram Vennam, Anish R Joishy, Ponnurangam Kumaraguru

Keywords: vocabulary compression, low compute environments, language modelling

TL;DR: We present a simple method that compresses the vocabulary layer in language models resulting in significantly lower memory usage.

Abstract: We present a method to compress the final linear layer of language models, reducing memory usage by up to 3.4x without significant performance loss. By grouping tokens based on Byte Pair Encoding (BPE) merges, we prevent materialization of the memory-intensive logits tensor. Evaluations on the TinyStories dataset show that our method performs on par with GPT-Neo while significantly improving throughput by up to 3x, making it suitable for low-resource environments.

Submission Number: 105