Replacing Dense Output Heads with Vector Search to Speed Up LLM Decoding

Introduction

At every decoding step, a language model normally projects its hidden state against the entire vocabulary and then selects or samples from the resulting logits. When the vocabulary is large, especially in multilingual models, this output head can become a memory-bandwidth problem. The issue is particularly visible for compact models running on CPUs with a batch size of one: the bottleneck may be moving the output embedding matrix rather than performing the arithmetic itself.

Turning the output head into a search task

The study proposes viewing the output projection followed by Top-k selection as a maximum inner product search problem. Instead of multiplying the hidden state by the complete vocabulary matrix, the system treats token embeddings as a vector collection and uses a Hierarchical Navigable Small World, or HNSW, index to retrieve tokens with high inner products.

The approach has several important components:

• Index the token embeddings: output embeddings are organized into an HNSW structure for approximate vector search.
• Retrieve a limited candidate set: only a small number of likely high-scoring tokens are evaluated, reducing full-vocabulary memory access.
• Preserve decoder compatibility: retrieved logits are scattered into a sparse tensor shaped like the full vocabulary, allowing existing decoding pipelines to continue using familiar sampling or Top-k interfaces.
• Trade exactness for speed: approximate search may not reproduce the exact global ranking, but it aims to retain the candidates that matter for generation.

Results and limitations

The authors evaluate the method on CPU inference with Gemma 3, Llama 3.2, and Qwen 3. The indexed output head substantially reduces the cost of the output projection. For end-to-end batch-size-one decoding, the reported throughput improvement reaches up to 82% on Gemma 3 270M. The study also uses AlpacaEval and reports that generation quality is preserved under that evaluation.

The result should not be read as evidence that dense output layers are obsolete in every serving environment. The proposed design is most relevant when the vocabulary is large, the batch is small, and the output projection consumes a meaningful share of latency. At larger batch sizes, dense matrix operations may benefit more from hardware parallelism. Approximate retrieval also introduces practical questions around index construction, memory layout, candidate recall, and the cost of scattering sparse logits. The supplied material does not provide a complete sweep of those trade-offs, so deployment decisions will require model- and hardware-specific testing.

Why it matters

The work broadens the way engineers can think about an LLM output head. Rather than treating it only as a fixed dense matrix multiplication, it can also be implemented as retrieval over a token-vector database. If future work reduces index overhead and evaluates the quality-latency trade-off across more models, sampling methods, and sequence lengths, this technique could become useful for CPU inference, edge deployment, and low-concurrency services.

Its broader lesson is straightforward: when decoding is limited by memory traffic, reducing the portion of the vocabulary that must be accessed may be more effective than optimizing the same full matrix operation. The approach is therefore best understood as a targeted systems optimization for latency-sensitive decoding, not a universal replacement for the conventional output projection.

Source: arXiv