# FlashHead: Efficient Drop-In Replacement for the Classification Head in Language Model Inference

arXiv: https://arxiv.org/abs/2603.14591 · HTML: https://mlsanity.com/paper/2603.14591/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 44 references that are in the map.
- First posted on arXiv 15 Mar 2026.

## Builds on (most cited first, 29 of 29)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them](https://mlsanity.com/paper/2210.09261/) (2022)
- [Instruction-Following Evaluation for Large Language Models](https://mlsanity.com/paper/2311.07911/) (2023)
- [XNLI: Evaluating Cross-lingual Sentence Representations](https://mlsanity.com/paper/1809.05053/) (2018)
- [AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback](https://mlsanity.com/paper/2305.14387/) (2023)
- [Small Language Models are the Future of Agentic AI](https://mlsanity.com/paper/2506.02153/) (2025)
- [Lessons from the Trenches on Reproducible Evaluation of Language Models](https://mlsanity.com/paper/2405.14782/) (2024)
- [On-Device Training Under 256KB Memory](https://mlsanity.com/paper/2206.15472/) (2022)
- [Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies](https://mlsanity.com/paper/2407.13623/) (2024)
- [OWQ: Outlier-Aware Weight Quantization for Efficient Fine-Tuning and Inference of Large Language Models](https://mlsanity.com/paper/2306.02272/) (2023)
- [FastDepth: Fast Monocular Depth Estimation on Embedded Systems](https://mlsanity.com/paper/1903.03273/) (2019)
- ["Give Me BF16 or Give Me Death"? Accuracy-Performance Trade-Offs in LLM Quantization](https://mlsanity.com/paper/2411.02355/) (2024)
- [Text Clustering with Large Language Model Embeddings](https://mlsanity.com/paper/2403.15112/) (2024)
- [An Efficient Multilingual Language Model Compression through Vocabulary Trimming](https://mlsanity.com/paper/2305.15020/) (2023)
- [FR-Spec: Accelerating Large-Vocabulary Language Models via Frequency-Ranked Speculative Sampling](https://mlsanity.com/paper/2502.14856/) (2025)
- [Strategies for Training Large Vocabulary Neural Language Models](https://mlsanity.com/paper/1512.04906/) (2015)
- [Learning to Screen for Fast Softmax Inference on Large Vocabulary Neural Networks](https://mlsanity.com/paper/1810.12406/) (2018)
- [SHAKTI: A 2.5 Billion Parameter Small Language Model Optimized for Edge AI and Low-Resource Environments](https://mlsanity.com/paper/2410.11331/) (2024)
- [Adaptive Sampled Softmax with Inverted Multi-Index: Methods, Theory and Applications](https://mlsanity.com/paper/2501.08563/) (2025)

## Built on by (most cited first, 2 of 2)

- [SoftWater: Class-Aware Rate Allocation for Softmax Quantization](https://mlsanity.com/paper/2608.12026/) (2026)
- [Two-Level Softmax Sampling Done Right: Correcting Bias from Size Imbalance and Dispersion](https://mlsanity.com/paper/2610.10483/) (2026)

## Abstract (arXiv)

Language models are increasingly adopting smaller architectures optimized for consumer devices. In this setting, inference efficiency is the primary constraint. Meanwhile, vocabulary sizes continue to grow rapidly, making the classification head a critical bottleneck that accounts for up to 60% of model parameters, and 50% of inference compute. We introduce FlashHead, the first efficient drop-in replacement for the dense classification head that is training-free and hardware-friendly. FlashHead builds on principles from information retrieval, reframing that computation at the output head as a retrieval problem rather than a dense classification over the full vocabulary. FlashHead introduces four key innovations: (1) a balanced clustering scheme that structures vocabulary partitions into compact hardware-efficient tensors, (2) extending multiprobe retrieval to language model heads, enabling thousands of clusters to be scored in parallel, (3) a novel inference-time sampling mechanism that extends retrieval beyond top tokens, enabling probabilistic sampling across the full vocabulary, and (4) selective quantization, enabling effective low-bit computation in the head. Experiments on Llama-3.2, Gemma-3, and Qwen-3 show that FlashHead delivers model-level inference speedups of up to 1.75x which maintaining output accuracy compared to the original head. By overcoming the classification head bottleneck, FlashHead establishes a new benchmark for efficient inference and removes a key barrier to developing smaller, capable models for consumer hardware.