用KL引导选择层,高效蒸馏大模型为混合注意力架构。
Distilling to Hybrid Attention Models via KL-Guided Layer Selection
- 基于少量通用文本训练计算层重要性得分,指导哪层转为线性注意力。
- 在LLaMA-3-8B上实现90%推理速度提升,损失仅0.5%性能。
- 适合追求推理效率的部署场景,无需重新预训练。
将预训练的softmax注意力Transformer蒸馏为更高效的混合架构(交错使用softmax与线性注意力层),是提升大语言模型推理效率而不需从头预训练的有前景方法。转换过程中的关键因素是层选择,即决定哪些层转换为线性注意力变体。本文提出一种简单高效的层选择方案:通过在少量通用文本数据上微调,计算各层的重要性分数。选定层后,采用近期提出的蒸馏流程(RADLADS),包括注意力权重转移、隐藏状态对齐、基于KL的分布匹配,最后进行少量微调。实验表明,该方法优于现有层选择策略,包括基于固定比例均匀交错线性注意力的启发式方法,以及依赖专用诊断数据的复杂方法。
原文摘要 · Abstract (English)
Distilling pretrained softmax attention Transformers into more efficient hybrid architectures that interleave softmax and linear attention layers is a promising approach for improving the inference efficiency of LLMs without requiring expensive pretraining from scratch. A critical factor in the conversion process is layer selection, i.e., deciding on which layers to convert to linear attention variants. This paper describes a simple and efficient recipe for layer selection that uses layer importance scores derived from a small amount of training on generic text data. Once the layers have been selected we use a recent pipeline for the distillation process itself \citep[RADLADS;][]{goldstein2025radlads}, which consists of attention weight transfer, hidden state alignment, KL-based distribution matching, followed by a small amount of finetuning. We find that this approach is more effective than existing approaches for layer selection, including heuristics that uniformly interleave linear attentions based on a fixed ratio, as well as more involved approaches that rely on specialized diagnostic datasets.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。