用检索思想替换分类头,让小模型推理快1.75倍
FlashHead: Efficient Drop-In Replacement for the Classification Head in Language Model Inference
- 把输出头变成检索任务,分组压缩参数结构
- 支持千级聚类并行计算,实现1.75倍推理加速
- 无需训练、适配低比特硬件,适合移动端部署
语言模型正向小型化发展以适配消费设备,此时推理效率成为主要瓶颈。词汇表持续增长,分类头占模型参数高达60%,计算量达50%。本文提出FlashHead,首个无需训练、硬件友好的密集分类头替代方案。它基于信息检索思想,将输出头计算重构为检索问题:(1)采用平衡聚类方案,将词汇分区组织为紧凑的硬件友好张量;(2)拓展多探针检索至语言模型头,实现上千个聚类并行评分;(3)引入推理时采样机制,支持全词汇表概率采样;(4)结合选择性量化,实现头部低比特高效计算。在Llama-3.2、Gemma-3和Qwen-3上实验显示,FlashHead可带来最高1.75倍的模型级推理加速,且保持原始输出准确性。该方法突破分类头瓶颈,为小型化高性能模型在消费硬件上的落地树立新基准。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。