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

Wilhelm Tranheden; Shahnawaz Ahmed; Devdatt Dubhashi; Jonna Matthiesen; Hannes von Essen

cs.LGMar 15, 2026

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

Authors: Wilhelm Tranheden, Shahnawaz Ahmed, Devdatt Dubhashi, Jonna Matthiesen, Hannes von Essen

Organizations: Embedl Gamla Almedalsvägen 21, 412 63 Gothenburg, Sweden

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Several clusters of tokens are selected conditioned on the hidden vector produced by the model body. The reduced subset of token embeddings are multiplied with the hidden vector to make the final token selection.. Illustration of the proposed FlashHead algorithm, highlighting the difference between greedy approach of selecting the most likely next token, and the sampling approach allowing for probabilistic sampling of new tokens.

Method	MMLU Pro	Hella Swag	IFEval	BoolQ	BBH	Truthful QA	GSM8K
Baseline	0.18	0.59	0.45	0.69	0.38	0.36	0.46
Vocab. Trimming	0.18	0.53	0.35	0.65	0.37	0.36	0.46
SVDSoftmax	0.16	0.44	0.44	0.69	0.13	0.36	0.26
FGD	—	N/A	0.32	N/A	—	0.42	—
FlashHead	0.18	0.59	0.45	0.69	0.38	0.36	0.46

Table 1: Evaluation on LM benchmarks. MMLU-Pro (exact_match, custom-extract), HellaSwag (acc_norm), IFEval (avg. of four metrics), BoolQ (acc_norm), BBH (exact_match), TruthfulQA-gen (bleu_acc), GSM8K (exact_match) for Llama 3.2 1B Instruct as the baseline.

Method	Top‑1			Top‑3
Method	Alpaca	MATH‑Hard	XNLI	Alpaca	MATH‑Hard	XNLI
Baseline	1.00	1.00	1.00	1.00	1.00	1.00
Vocab. Trimming	0.99	0.99	0.51	1.00	1.00	0.69
SVDSoftmax	0.94	0.96	0.96	0.99	0.99	0.99
FlashHead	1.00	1.00	0.97	1.00	1.00	1.00

Table 2: Top‑

k

Containment (

c{=}8016

p{=}512

) vs. training‑free baselines on Llama‑3.2‑1B‑Instruct . For

k{=}1

, containment = accuracy.

Method	TPOT ${}^{\text{H}}$ $\downarrow$	TPOT (BF16) $\downarrow$	TPOT (INT4) $\downarrow$
Baseline	1.94	7.69	3.60
Vocab. Trimming	1.07 (1.81 $\times$ )	6.82 (1.13 $\times$ )	2.73 (1.32 $\times$ )
SVDSoftmax	0.61 (3.18 $\times$ )	6.36 (1.21 $\times$ )	2.27 (1.59 $\times$ )
FGD	N/A	N/A	N/A
FlashHead	0.40 (4.85 $\times$ )	6.15 (1.25 $\times$ )	2.06 (1.75 $\times$ )

Table 3: GPU latency (ms) for head (

\mathrm{TPOT}^{\mathrm{H}}

) and full model (

\mathrm{TPOT}

) under bfloat16 and int4 heads. Note that FGD lacks GPU support. Extended CPU and parameter results appear in the Appendix A

Model	BBH $\uparrow$	TPOT ${}^{\text{H}}$ $\downarrow$	TPOT (BF16) $\downarrow$	TPOT (INT4) $\downarrow$
Llama‑3.2‑1B	0.38 $\rightarrow$ 0.38	1.94 $\rightarrow$ 0.40 (4.9 $\times$ )	7.69 $\rightarrow$ 6.15 (1.2 $\times$ )	3.60 $\rightarrow$ 2.06 (1.8 $\times$ )
Llama‑3.2‑3B	0.57 $\rightarrow$ 0.57	2.13 $\rightarrow$ 0.68 (3.1 $\times$ )	18.60 $\rightarrow$ 17.15 (1.1 $\times$ )	7.11 $\rightarrow$ 5.66 (1.3 $\times$ )
Llama‑3.1‑8B	0.71 $\rightarrow$ 0.70	2.72 $\rightarrow$ 1.21 (2.2 $\times$ )	OOM	13.55 $\rightarrow$ 12.04 (1.1 $\times$ )
Qwen‑3‑1.7B	0.45 $\rightarrow$ 0.45	1.61 $\rightarrow$ 0.45 (3.6 $\times$ )	9.97 $\rightarrow$ 8.81 (1.1 $\times$ )	4.85 $\rightarrow$ 3.69 (1.3 $\times$ )
Gemma‑3‑270M	0.27 $\rightarrow$ 0.27	0.99 $\rightarrow$ 0.37 (2.7 $\times$ )	2.52 $\rightarrow$ 1.90 (1.3 $\times$ )	2.38 $\rightarrow$ 1.76 (1.4 $\times$ )
Gemma‑3‑1B	0.38 $\rightarrow$ 0.38	1.66 $\rightarrow$ 0.52 (3.2 $\times$ )	6.77 $\rightarrow$ 5.63 (1.2 $\times$ )	4.12 $\rightarrow$ 2.98 (1.4 $\times$ )

Table 4: BBH and GPU latency improvements across models. Values in parentheses are speedup factors vs. the dense head. A comprehensive accuracy table and CPU metrics are included in Appendix B . Note that the 8 billion parameter Llama 3.1 model is too big for inference in bfloat16 precision (Out-of-Memory, OOM).

Method	Precision	BBH $\uparrow$	GPU TPOT ${}^{\text{H}}$ $\downarrow$
Baseline	BF16	0.381	1.433
	INT4	0.369	0.486
FlashHead	BF16	0.381	0.320
	INT4	0.379	0.258

Table 5: Impact of head quantization on BBH and GPU

\mathrm{TPOT}^{\mathrm{H}}

. Latencies are measured and quantization is performed with the transformers Wolf et al. (2020) and HQQ Badri & Shaji (2023) libraries, due to their support of quantizing language model heads.

Method	BBH $\uparrow$	GPU TPOT ${}^{\text{H}}$ $\downarrow$
Unequal clusters	0.371	0.520
Equal clusters	0.381	0.320

Table 6: Equal vs. Unequal clustering on BBH and GPU

\mathrm{TPOT}^{\mathrm{H}}

#clusters	#probes	Top‑1			GPU Latency (ms)
#clusters	#probes	Alpaca	MATH‑Hard	XNLI	TPOT ${}^{\text{H}}$ $\downarrow$	TPOT (BF16) $\downarrow$	TPOT (INT4) $\downarrow$
Baseline		1.00	1.00	1.00	1.94	7.69	3.60
4008	128	0.99	0.99	0.93	0.18 (10.78 $\times$ )	5.93 (1.30 $\times$ )	1.84 (1.96 $\times$ )
	256	0.99	0.99	0.95	0.33 (5.88 $\times$ )	6.08 (1.26 $\times$ )	1.99 (1.80 $\times$ )
	512	1.00	1.00	0.97	0.58 (3.34 $\times$ )	6.33 (1.21 $\times$ )	2.24 (1.61 $\times$ )
8016	128	0.98	0.99	0.93	0.12 (16.17 $\times$ )	5.87 (1.31 $\times$ )	1.78 (2.02 $\times$ )

Table 7: Top‑1 containment and GPU latency as a function of clusters

c

and probes

p

. We report head latency (

\mathrm{TPOT}^{\mathrm{H}}

) and full‑model

\mathrm{TPOT}

in BF16/INT4.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Table 8: On-device latencies and parameter counts for all methods. Latencies are reported in time-per-output-token (ms), averaged over 128 generated tokens. Factors in parentheses indicate relative speedup or reduction compared to the baseline.

Table 9: Evaluation of models across (a) LM benchmarks and (b) efficiency metrics. We exclude BoolQ and HellaSwag due to the high cost of Monte Carlo simulations. Params

{}^{\text{H}}

and Params indicate active and total parameters.

MMLU-Pro	BBH	TruthfulQA (gen)	IFEval	GSM8K (chat)
0.181 $\pm$ 0.001	0.377 $\pm$ 0.004	0.363 $\pm$ 0.002	0.452 $\pm$ 0.003	0.465 $\pm$ 0.003

Appendix

Table 10: Mean and standard deviation of FlashHead performance on LM-Eval benchmarks (Llama-3.2-1B-Instruct); scores are averaged over five independent clusterings.

Figure 2: Convergence of the Monte-Carlo estimator when used to evaluate accuracy on the BoolQ and HellaSwag datasets. An average is taken over five independent clustering procedures with different random seeds. Shaded ribbons show

\pm 1

s.e.

Method	Hyper-parameters
Baseline (LLaMA-3.2 1B Instruct)	num_hidden_layers: 16
	hidden_size: 2048
	intermediate_size: 8192
	num_attention_heads: 32
	num_key_value_heads: 8
	head_dim: 64

Appendix

Table 11: Methods and hyper-parameters used in the various implementations.

Dataset name	Reference	Link
Alpaca (train and eval)	Taori et al. (2023) ; Dubois et al. (2023)	tatsu-lab/alpaca , tatsu-lab/alpaca_eval
MATH-Hard	Hendrycks et al. (2021)	lighteval/MATH-Hard
XNLI	Conneau et al. (2018)	facebook/xnli
MMLU-Pro	Wang et al. (2024)	TIGER-Lab/MMLU-Pro
Hella-Swag	Zellers et al. (2019)	hellaswag
IFEval	Zhou et al. (2023)	google/IFEval

Appendix

Table 12: Overview of the datasets used in this paper.

Benchmark	Reported metric	Settings
MMLU-Pro	exact_match,custom-extract	max_gen_toks: 2048
		num_fewshot: 5
HellaSwag	acc_norm	num_fewshot: 0
IFEval	average of (prompt_strict, instruction_strict, prompt_loose, instruction_loose acc)	num_fewshot: 0
		max_gen_toks: 1280
BoolQ	acc	num_fewshot: 0

Appendix

Table 13: Evaluation settings for LM-Eval tasks.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

FlashBack: Efficient Retrieval-Augmented Language Modeling for Fast Inference
May 7, 2024Runheng Liu, Xingchen Xiao, Heyan Huang +2Knowledge Augmentation for Language Models LLM Inference Acceleration
School of Computer Science and Technology, Beijing Institute of Technology
Lever: Speculative LLM Inference on Smartphones
May 16, 2026Tuowei Wang, Fengzu Li, Yanfan Sun +2On-Device Language Model Inference LLM Inference
Tsinghua University · Beijing, China · Beihang University +1
RoofLang: Enabling AI-Driven Architecting of LLM Inference Systems
Sep 11, 2026Ziyue Yang, Yuting Jiang, Lei Qu +1LLM Inference Large Language Model-Guided Optimization
Shanghai Xingyunzhili Artificial Intelligence Institute, Shanghai, China · Microsoft Research, Redmond, WA, USA

May 7, 2024J/K move · Enter open · S save

FlashBack: Efficient Retrieval-Augmented Language Modeling for Fast Inference

Runheng Liu, Xingchen Xiao, Heyan Huang, Zewen Chi+1

School of Computer Science and Technology, Beijing Institute of Technology

Knowledge Augmentation for Language Models LLM Inference Acceleration

Open paper AskPDF arXiv

Retrieval-Augmented Language Modeling (RALM) by integrating large language models (LLM) with relevant documents from an external corpus is a proven method for enabling the LLM to generate information beyond the scope of its pre-training corpus. Previous work utilizing retrieved content by simply prepending it to the input poses a high runtime issue, which degrades the inference efficiency of the LLMs because they fail to use the Key-Value (KV) cache efficiently. In this paper, we propose FlashBack, a modular RALM designed to improve the inference efficiency of RALM with appending context pattern while maintaining decent performance after fine-tuning by Low-Rank Adaption. FlashBack appends retrieved documents at the end of the context for efficiently utilizing the KV cache instead of prepending them. And we introduce Marking Token as two special prompt tokens for marking the boundary of the appending context during fine-tuning. Our experiments on testing generation quality show that FlashBack can remain decent generation quality in perplexity. And the inference speed of FlashBack is up to

4\times

faster than the prepending counterpart on a 7B LLM (Llama 2) in the runtime test. Via bypassing unnecessary re-computation, it demonstrates an advancement by achieving significantly faster inference speed, and this heightened efficiency will substantially reduce inferential cost.

Abstract

LLM Inference Acceleration

Explore similar work

FlashBack: Efficient Retrieval-Augmented Language Modeling for Fast Inference

Lever: Speculative LLM Inference on Smartphones

RoofLang: Enabling AI-Driven Architecting of LLM Inference Systems