The final linear projection in a language model is easy to overlook. It is just a matrix multiplication, applied once per token. But for compact models with large multilingual vocabularies, this single operation can consume the majority of decoding time, because it requires streaming an enormous weight matrix from memory at every step. A 270M parameter model with a 256k-token vocabulary spends roughly 62% of its total parameter budget on that output head alone. This paper, available at arXiv:2608.27460, proposes a clean and practically motivated solution: replace the dense projection with an approximate nearest-neighbour search over the token embedding matrix.
The Core Idea
Selecting the top-k tokens from a vocabulary projection is, mathematically, a Maximum Inner Product Search (MIPS) problem. Given the final hidden state as a query vector, you want the k token embeddings with the highest dot product. Dense matrix multiplication computes this exhaustively across the entire vocabulary. For large vocabularies and small batch sizes, most of that computation is wasted, because only a handful of tokens will ever appear in the candidate set.
The authors replace this exhaustive step with a Hierarchical Navigable Small World (HNSW) graph index built over the output embedding matrix. HNSW is the standard approximate nearest-neighbour algorithm underlying most production vector databases, and it trades a small, controllable approximation error for a large reduction in the number of distance computations required. The retrieved candidate set is then scattered into a sparse full-vocabulary logit tensor, preserving compatibility with standard sampling pipelines. No retraining or architectural modification is needed; the index is constructed post-hoc from the existing weight matrix.
This is a sensible reframing. The connection between top-k sampling and MIPS has been noted before in the retrieval-augmented generation literature, but applying it directly to the output head of a standard transformer decoder is a practically useful contribution, particularly given the trend toward small, vocabulary-heavy multilingual models distilled from larger systems.
Methodology and Implementation
The experiments cover Gemma 3 (270M and 1B), Llama 3.2 (1B and 3B), and Qwen 3 (0.6B and 1.7B), all evaluated on CPU in float32. The authors use a custom fork of hnswlib tuned for high-dimensional vector access patterns and multithreaded search. Index construction uses M=32 and ef_construction=5000, which are relatively aggressive settings that favour recall at the cost of build time. During generation, ef=200 is the primary operating point, with ef=100 explored as a lower-latency alternative.
A few methodological choices are worth examining. First, the decision to benchmark exclusively on CPU is both a strength and a limitation. CPU inference is the realistic deployment target for edge scenarios, and it is where memory bandwidth constraints are most acute. On the other hand, it means the results say nothing about GPU throughput, which is where most production LLM serving happens. Second, the use of float32 throughout is somewhat conservative; the interaction between this approach and weight quantization (INT8, INT4) is acknowledged but not explored, and that interaction matters a great deal for practical deployment.
Quality evaluation is performed using AlpacaEval, which measures instruction-following quality through pairwise comparison. The authors report minimal degradation at ef=200, which is plausible given that HNSW recall at that setting is typically very high for well-separated embedding spaces. A more rigorous quality analysis across perplexity benchmarks and diverse sampling temperatures would strengthen this claim, but the AlpacaEval results are at least a reasonable sanity check.
Results
The headline result is an 82% end-to-end throughput improvement for Gemma 3 270M at batch size one. This is a large number, and it is credible given that the output projection accounts for 62% of that model's parameters. The gains are smaller but still positive for larger models where the output head is a smaller fraction of total compute.
Key findings from the experiments:
- All evaluated models show positive throughput gains at batch size one with ef=200.
- Gains persist at moderate batch sizes but erode as batch size increases beyond 32.
- At large batch sizes, the aggregate number of unique tokens retrieved across the batch forces loading a large fraction of the embedding matrix, negating the memory bandwidth advantage.
- Halving ef to 100 roughly doubles search throughput with some quality trade-off, offering a useful operating point for latency-critical applications.
- The speedup is most pronounced for models where the vocabulary is disproportionately large relative to the transformer backbone, which describes the current generation of small multilingual models accurately.
The crossover behaviour at large batch sizes is well-understood and honestly reported. Batched GEMM on modern hardware benefits from excellent cache locality and highly optimised kernel implementations; graph traversal in HNSW is fundamentally sequential and does not benefit from batching in the same way. This limits the method to the single-request, interactive inference regime, which is explicitly the target use case.
Limitations and Broader Implications
The CPU-only constraint is the most significant practical limitation. GPU-accelerated HNSW implementations exist in research form but are not mature, and the sequential graph traversal in HNSW is genuinely difficult to parallelise efficiently on GPU architectures. The authors acknowledge this and frame GPU acceleration as future work. Until that gap is closed, the method is restricted to edge and on-device scenarios, which is a real but narrow deployment context.
The interaction with quantization deserves more attention than it receives here. Heavily quantized dense baselines (e.g., INT4 with grouped quantization) have substantially lower memory bandwidth requirements than float32, which would reduce the relative advantage of the vector index approach. The paper's framing of quantization as orthogonal and complementary is technically correct but potentially optimistic about the magnitude of combined gains.
There is also a question about index construction cost. Building an HNSW index with ef_construction=5000 over a 256k-token embedding matrix is not instantaneous, and for deployment scenarios where model weights are updated or fine-tuned frequently, this rebuild cost needs to be factored in. For static deployed models this is a one-time cost, but it is worth quantifying.
The broader implication is worth taking seriously. As the field continues to produce small, distilled models with large multilingual vocabularies, the output projection bottleneck will become more common, not less. Approaches that decouple vocabulary size from decoding cost are genuinely useful for this class of model. The MIPS reframing is clean, the implementation is practical, and the reference code is publicly available at github.com/martinloretzzz/vector-index-embedding. For researchers working on CPU inference or edge deployment of multilingual models, this is a method worth evaluating.