LLM Vocabulary Compression for Low-Compute Environments

TL;DR

A method to compress the final linear layer of language models, reducing memory usage by up to 3.4x without significant performance loss and significantly improving throughput by up to 3x is presented, making it suitable for low-compute environments.

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 and GPT2 while significantly improving throughput by up to 3x, making it suitable for low-compute environments.

Figures (7)

• Figure 1: Example images from the synthetic dataset for image classification.
• Figure 2: Validation accuracy vs Train steps for the synthetic image classification dataset.
• Figure 3: Ablation studies
• Figure 4: Slow implementation which requires looping over the groups
• Figure 5: Fast implementation applying the scale and shift transformation.