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

• URL: http://arxiv.org/abs/2603.14591v1
• Date: Sun, 15 Mar 2026 20:26:42 GMT
• Title: FlashHead: Efficient Drop-In Replacement for the Classification Head in Language Model Inference
• Abstract summary: We introduce FlashHead, a drop-in replacement for the dense classification head that is training-free and hardware-friendly.<n> FlashHead builds on principles from information retrieval, reframing that computation at the output head as a retrieval problem.<n>We show that FlashHead delivers model-level inference speedups of up to textbf1.75x which maintaining output accuracy compared to the original head.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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 \textbf{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.

Related papers

• Explicit Multi-head Attention for Inter-head Interaction in Large Language Models [70.96854312026319]
  Multi-head Explicit Attention (MEA) is a simple yet effective attention variant that explicitly models cross-head interaction.<n>MEA shows strong robustness in pretraining, which allows the use of larger learning rates that lead to faster convergence.<n>This enables a practical key-value cache compression strategy that reduces KV-cache memory usage by 50% with negligible performance loss.
  arXiv  Detail & Related papers  (2026-01-27T13:45:03Z)
• An Enhanced Model-based Approach for Short Text Clustering [58.60681789677676]
  Short text clustering has become increasingly important with the popularity of social media like Twitter, Google+, and Facebook.<n>Existing methods can be broadly categorized into two paradigms: topic model-based approaches and deep representation learning-based approaches.<n>We propose a collapsed Gibbs Sampling algorithm for the Dirichlet Multinomial Mixture model (GSDMM), which effectively handles the sparsity and high dimensionality of short texts.<n>Based on several aspects of GSDMM that warrant further refinement, we propose an improved approach, GSDMM+, designed to further optimize its performance.
  arXiv  Detail & Related papers  (2025-07-18T10:07:42Z)
• Comprehensive Multi-Modal Prototypes are Simple and Effective Classifiers for Vast-Vocabulary Object Detection [68.26282316080558]
  Current open-world detectors can recognize a broader range of vocabularies, despite being trained on limited categories.<n>We introduce Prova, a prototype classifier for vast-vocabulary object detection.
  arXiv  Detail & Related papers  (2024-12-23T18:57:43Z)
• Dynamic layer selection in decoder-only transformers [21.18795712840146]
  We empirically examine two common dynamic inference methods for natural language generation. We find that a pre-trained decoder-only model is significantly more robust to layer removal via layer skipping. We also show that dynamic computation allocation on a per-sequence basis holds promise for significant efficiency gains.
  arXiv  Detail & Related papers  (2024-10-26T00:44:11Z)
• Class-Imbalanced Semi-Supervised Learning for Large-Scale Point Cloud Semantic Segmentation via Decoupling Optimization [64.36097398869774]
  Semi-supervised learning (SSL) has been an active research topic for large-scale 3D scene understanding. The existing SSL-based methods suffer from severe training bias due to class imbalance and long-tail distributions of the point cloud data. We introduce a new decoupling optimization framework, which disentangles feature representation learning and classifier in an alternative optimization manner to shift the bias decision boundary effectively.
  arXiv  Detail & Related papers  (2024-01-13T04:16:40Z)
• Gramian Attention Heads are Strong yet Efficient Vision Learners [26.79263390835444]
  We introduce a novel architecture design that enhances expressiveness by incorporating multiple head classifiers (ie, classification heads) Our approach employs attention-based aggregation, utilizing pairwise feature similarity to enhance multiple lightweight heads with minimal resource overhead. Our models eventually surpass state-of-the-art CNNs and ViTs regarding the accuracy-grained trade-off on ImageNet-1K.
  arXiv  Detail & Related papers  (2023-10-25T09:08:58Z)
• Sentiment analysis with adaptive multi-head attention in Transformer [0.0]
  We propose a novel framework based on the attention mechanism to identify the sentiment of a movie review document. In this paper, we propose an adaptive multi-head attention architecture (AdaptAttn) which varies the number of attention heads based on length of sentences.
  arXiv  Detail & Related papers  (2023-10-23T02:32:30Z)
• Learning to Detect and Segment for Open Vocabulary Object Detection [6.678101044494558]
  We propose a principled dynamic network design to better generalize the box regression and mask segmentation for open vocabulary setting. CondHead is composed of two streams of network heads, the dynamically aggregated head and the dynamically generated head. Our method brings significant improvement to the state-of-the-art open vocabulary object detection methods with very minor overhead.
  arXiv  Detail & Related papers  (2022-12-23T03:54:59Z)
• Confident Adaptive Language Modeling [95.45272377648773]
  CALM is a framework for dynamically allocating different amounts of compute per input and generation timestep. We demonstrate the efficacy of our framework in reducing compute -- potential speedup of up to $times 3$ -- while provably maintaining high performance.
  arXiv  Detail & Related papers  (2022-07-14T17:00:19Z)
• Clustering by Maximizing Mutual Information Across Views [62.21716612888669]
  We propose a novel framework for image clustering that incorporates joint representation learning and clustering. Our method significantly outperforms state-of-the-art single-stage clustering methods across a variety of image datasets.
  arXiv  Detail & Related papers  (2021-07-24T15:36:49Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.