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

May 7, 2024

FlashBack: Efficient Retrieval-Augmented Language Modeling for Fast Inference

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.

Runheng Liu, Xingchen Xiao, Heyan Huang +2

School of Computer Science and Technology, Beijing Institute of Technology

Knowledge Augmentation for Language Models LLM Inference Acceleration

Ask PDF

May 16, 2026

Lever: Speculative LLM Inference on Smartphones

Large language models (LLMs) are increasingly needed for interactive mobile applications, but high-quality models exceed the limited DRAM available on smartphones. Flash storage can hold larger models, yet flash-backed inference is slow because autoregressive decoding repeatedly invokes the target model and incurs costly I/O. We observe that speculative decoding is a natural fit for this setting: a small draft model can remain in DRAM, while a larger flash-resident target model verifies multiple candidate tokens per invocation. However, existing methods assume server-class accelerators and fail to account for prolonged I/O latency, limited computation parallelism, and irregular speculation execution. We present Lever, an end-to-end system for efficient flash-backed LLM inference on smartphones. Lever jointly optimizes the three stages of speculative decoding under mobile constraints. For drafting, it builds token trees using an I/O- and compute-aware gain-cost objective. For verification, it prunes low-value branches through early-exit prediction to reduce target-model computation. For execution, it maps speculation efficiently across mobile CPU-NPU hardware to improve utilization. Comprehensive evaluations show that Lever reduces inference latency by an average of 2.93x over baseline flash-offloaded inference and 1.50x over conventional speculative decoding, narrowing the latency gap between flash-backed and memory-resident LLM inference.

Tuowei Wang, Fengzu Li, Yanfan Sun +2

Tsinghua University · Beijing, China · Beihang University +1

On-Device Language Model Inference LLM Inference Speculative Decoding LLM Inference Acceleration

Ask PDF

Sep 11, 2026

RoofLang: Enabling AI-Driven Architecting of LLM Inference Systems

AI is beginning to make substantive contributions to LLM inference optimization. Existing AI optimizations are predominantly profiling-based. Profiling-bound feedback confines the search to the capabilities and performance of an existing software stack, preventing a fundamentally better architecture of LLM inference systems from being identified. To enable the AI-driven LLM inference system architecting loop, we argue that a general workload representation, a verifiable mutation space, and an implementation-independent evaluator are required. We present the RoofLang domain-specific language (DSL) that provides these features. In our evaluation, RoofLang reveals that DeepSeek V4-series models could achieve 3.5-39.5

\times

higher peak decode throughput than other representative models. This gap is disproportionate to their total parameter counts and arises largely from compact KV-cache designs that support larger batches and reduce memory traffic. A persistent optimizer agent further discovered several new architectures that improved both throughput and interactivity of DeepSeek V4 Pro on NVIDIA B300 by 6.23-50.1%.

Ziyue Yang, Yuting Jiang, Lei Qu +1

Shanghai Xingyunzhili Artificial Intelligence Institute, Shanghai, China · Microsoft Research, Redmond, WA, USA

LLM Inference Large Language Model-Guided Optimization LLM Inference Acceleration

Ask PDF

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