VQ-Logits: Compressing the Output Bottleneck of Large Language Models via Vector Quantized Logits

TL;DR

VQ-Logits addresses the output bottleneck in large language models by replacing the expansive $d_{model} \times V$ output projection with a compact codebook of size $K$ and a fixed vocabulary-to-codebook mapping. The model computes logits over the codebook and then scatters them to the full vocabulary, achieving up to $99\%$ reduction in output-layer parameters and up to $6\times$ speedups in logit computation, with only modest perplexity degradation on standard benchmarks. The paper provides two initialization/learning strategies, extensive experiments on WikiText-103, PTB, and C4, and detailed ablations showing robustness to codebook size and initialization. This approach offers a simple, effective path to efficient inference for large models and can complement existing compression techniques.

Abstract

Large Language Models (LLMs) have achieved remarkable success but face significant computational and memory challenges, particularly due to their extensive output vocabularies. The final linear projection layer, mapping hidden states to vocabulary-sized logits, often constitutes a substantial portion of the model's parameters and computational cost during inference. Existing methods like adaptive softmax or hierarchical softmax introduce structural complexities. In this paper, we propose VQ-Logits, a novel approach that leverages Vector Quantization (VQ) to drastically reduce the parameter count and computational load of the LLM output layer. VQ-Logits replaces the large V * dmodel output embedding matrix with a small, shared codebook of K embedding vectors (K << V ). Each token in the vocabulary is mapped to one of these K codebook vectors. The LLM predicts logits over this compact codebook, which are then efficiently "scattered" to the full vocabulary space using the learned or preassigned mapping. We demonstrate through extensive experiments on standard language modeling benchmarks (e.g., WikiText-103, C4) that VQ-Logits can achieve up to 99% parameter reduction in the output layer and 6x speedup in logit computation, with only a marginal 4% increase in perplexity compared to full softmax baselines. We further provide detailed ablation studies on codebook size, initialization, and learning strategies, showcasing the robustness and effectiveness of our approach.

Figures (2)

• Figure 1: Overview of the VQ-Logits architecture. The LLM's final hidden state $\mathbf{h}$ is projected onto a small codebook $\mathbf{C}$ (transposed) to get codebook logits $\mathbf{L}_{\text{c}}$. These are then scattered to the full vocabulary size $V$ using a vocabulary-to-codebook mapping $\mathcal{M}$ to produce $\mathbf{L}_{\text{v}}$, before the final softmax.
• Figure 2: Impact of codebook size $K$ on perplexity and logit computation speedup for VQ-Logits on WikiText-103. A clear trade-off is visible (Data corresponds to WikiText-103 with our Transformer model; plots are illustrative based on trends from Table \ref{['tab:main_results_wikitext103']}).