混合线性注意力的正确做法：超长上下文的高效蒸馏与有效架构 Hybrid Linear Attention Done Right: Efficient Distillation and Effective Architectures for Extremely Long Contexts
仅用2.3B token将Transformer蒸馏为高效混合模型
前置知识
Softmax Attention
Transformer中的核心注意力机制，通过计算Query和Key的点积经过softmax归一化得到注意力权重，再与Value相加权求和。其计算复杂度为$O(T^2)$，其中$T$为序列长度，导致在长上下文场景下计算开销巨大。形式化表达为$Y = \text{softmax}(\frac{QK^\top}{\sqrt{d_h}} \odot M) V W_o^\top$，其中$Q, K, V$分别由输入$X$经过线性变换得到。
理解softmax attention的二次复杂度瓶颈是理解本文动机的关键，也是为何需要将其转换为线性复杂度的RNN层
RNN/线性注意力层
循环神经网络的一种形式，通过维护固定大小的隐藏状态$S_t$来实现线性复杂度的序列建模。其核心公式为状态更新规则$S_t = F_t S_{t-1} + k_t^\top v_t$和查询规则$y_t = q_t S_t W_o^\top$，其中$F_t$为转移矩阵。现代RNN变体包括Mamba2、Gated DeltaNet、Lightning Attention等，通常将$F_t$设计为对角矩阵或秩-1矩阵以支持并行计算。
RNN层的线性复杂度使其在长上下文推理时效率远高于softmax attention，但固定大小状态限制了其召回能力
RoPE (旋转位置编码)
Rotary Position Embedding，通过对Query和Key向量应用位置相关的旋转操作来注入位置信息。已成为Transformer模型的事实标准位置编码方案，能有效提升语言建模性能，但在超长上下文外推时存在长度泛化问题。
本文提出的HyPE方案正是基于对RoPE和NoPE各自优劣的深入分析，理解RoPE的特性是理解HyPE设计动机的基础
NoPE (无位置编码)
指不使用任何显式位置编码的注意力机制。研究表明NoPE在训练免长度泛化方面表现优异，因为不存在位置编码带来的分布外问题。但NoPE在短距离位置建模上可能不如RoPE。
HyPE的核心创新就是在注意力层使用NoPE、在RNN层使用RoPE，充分利用两者各自的优势
知识蒸馏
一种模型压缩技术，通过让小模型（学生模型）学习大模型（教师模型）的输出分布或中间表示来转移知识。本文中特指将预训练Transformer（教师）转换为混合架构（学生）的过程，采用KL散度作为蒸馏损失函数$\mathcal{L}_{stage2} = D_{KL}(f_{orig}(X) \| f_{hybrid}(X))$。
HALO的核心就是一套高效的跨架构蒸馏流程，理解知识蒸馏原理有助于把握本文的技术路线
研究动机
当前混合Transformer架构（交织softmax attention层和RNN层）在长上下文建模中展现出理想的性能-吞吐量权衡，但面临两大核心瓶颈。第一，大规模从头预训练成本极高，动辄需要数百B token的训练数据（如Jet-Nemotron使用400B token，KL-LS使用25B token），这使得绝大多数学术团队无法参与研究。第二，现有蒸馏方法虽然能够将预训练Transformer转换为混合模型，但转换后的模型在长上下文任务上表现严重退化——而这恰恰是混合模型相对于纯Transformer最具推理加速优势的场景。例如，Jet-Nemotron在128K上下文的NIAH任务上准确率降至0%，KL-LS在256K上下文上也仅有10-14%的准确率。
本文的目标是本文的直接目标是开发一套高效的数据蒸馏流程（HALO），仅需不到3B token即可将预训练Transformer模型转换为RNN-attention混合模型；同时设计一种新型混合架构（HypeNet），在保持短上下文性能的同时实现卓越的长度泛化能力，使转换后的模型在长上下文场景下不仅能享受推理加速，还能维持甚至超越原始Transformer的性能。
与已有工作不同的是，本文的独特切入点在于两个方面。一是蒸馏效率的极致追求：HALO通过精心设计的注意力层选择机制，仅需2.3B token（不到预训练数据的0.01%）就能完成转换，比之前最高效的方法（SMART的7B token）少3倍以上。二是对位置编码在混合模型中作用的深刻洞察：提出HyPE方案，在RNN层使用RoPE以获得丰富位置信息，在注意力层使用NoPE以实现优秀的长度泛化，这种"分而治之"的策略是对现有混合模型PE方案（通常注意力用RoPE、RNN用NoPE）的根本性颠覆。
核心方法
HALO的整体思路可以概括为'先对齐、再选择、后蒸馏'的三阶段流水线。直觉上，将Transformer转换为混合模型的关键挑战是：哪些注意力层应该保留（对召回任务至关重要），哪些应该替换为RNN层（提升效率）。HALO首先通过注意力权重转移初始化RNN层，然后独立训练每个RNN层使其输出与对应注意力层对齐（Stage 1），接着基于召回性能和常识推理性能的变化计算各层重要性分数来选择保留哪些注意力层，最后通过端到端知识蒸馏（Stage 2）和长上下文微调（Stage 3）优化整个混合模型。这种设计使得整个转换过程仅需2.3B token。
本文的核心创新在于HyPE（混合位置编码）方案及其背后的设计哲学。与现有混合模型中注意力层使用RoPE、RNN层使用NoPE的常规做法相反，HyPE在RNN层使用RoPE、在注意力层使用NoPE。这一设计的精妙之处在于：RNN层通过状态转移矩阵$F_t$天然具有位置感知能力，添加RoPE能进一步增强其对短距离依赖的建模能力；而注意力层使用NoPE则避免了RoPE带来的长度泛化限制，使得模型在推理时能处理远超训练长度的上下文。配合动态注意力logits缩放机制$st = \log_a(t + a)$，HyPE实现了远超现有PE方案的长度泛化能力。
方法步骤详情
HALO流程包含以下具体步骤：(1) 初始化阶段：对Transformer的每个注意力层$ATTN^{(l)}$，利用其预训练的投影权重$(W_q, W_k, W_v, W_o)$实例化对应的RNN层$RNN^{(l)}$，对于RNN特有的模块则随机初始化。(2) Stage 1隐状态对齐：独立训练每个RNN层，最小化其输出与教师注意力层输出的均方误差$\mathcal{L}_{stage1} = MSE(Y^{(l)}_{teacher}, RNN^{(l)}(X^{(l-1)}))$，此阶段仅训练RNN层，其他权重冻结。(3) 注意力层选择：计算每个注意力层的重要性分数$s_i = \frac{\max_i R(M^{(i)}) - R(M^{(i)})}{\max_i C(M^{(i)}) - C(M^{(i)}) + \epsilon}$，其中$R$为召回性能、$C$为常识推理性能，选择Top-k个最重要的注意力层保留（$k = \lfloor L/4 \rfloor$）。(4) Stage 2知识蒸馏：构建最终混合模型，使用KL散度进行端到端蒸馏$\mathcal{L}_{stage2} = D_{KL}(f_{orig}(X) \| f_{hybrid}(X))$，使用1B训练数据和余弦学习率调度。(5) Stage 3长上下文微调：使用更长上下文和更小学习率进行1B token的微调以优化长上下文性能。
技术新颖性
本文的技术新颖性体现在多个层面。首先是HyPE位置编码方案的颠覆性设计：与主流做法相反地在RNN层使用RoPE、注意力层使用NoPE，这一看似反直觉的设计却实现了最佳的长度泛化。其次是注意力层选择的重要性度量方法：同时考虑召回性能下降和常识推理性能变化的比值，而非单一指标，确保保留的注意力层对召回任务最关键。第三是GQA到MHA的转换：在RNN层初始化时解耦KV头（$W^{(i)}_{\square} \leftarrow W^{(\lfloor i/g \rfloor)}_{\square}$），增加RNN层的表达能力。第四是为RNN层和注意力层都添加输出门机制$y_t = (Norm(o_t) \odot z_t) W_o^\top$，这在转换场景中提供了稳定的性能提升。这些改进使得HypeNet在仅增加约10%参数量的情况下，实现了显著的效率提升和性能改善。
实验结果
本文的实验结果表明HALO和HypeNet在多个维度上取得了显著突破。在长上下文召回性能方面，HypeNet-2B在256K上下文的NIAH-Single-1任务上达到99.8%准确率，远超Jet-Nemotron的0%和KL-LS的14.8%；在128K上下文上，HypeNet保持99.8%准确率，而Jet-Nemotron降至0%、KL-LS降至68.4%。在效率方面，HypeNet在128K上下文下比Qwen3系列实现了更好的性能-效率权衡，解码速度提升最高达3.0倍，预填充速度提升最高达3.4倍（在512K上下文下）。在CSR（常识推理）任务上，HypeNet-2B达到55.9的平均分数，与Qwen3教师模型的58.5分相当。消融实验证明了HyPE的关键作用：去除RNN RoPE会导致256K NIAH从79.9降至47.9，添加注意力RoPE则会导致从79.9降至19.7。Lightning Attention在所有RNN混合器中展现出最佳的长度泛化能力。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| NIAH-Single-1 (256K上下文) | 准确率(%) | 99.8 | Jet-Nemotron: 0.0, KL-LS: 14.8 | 相对Jet-Nemotron提升99.8个百分点 |
| NIAH-Single-2 (256K上下文) | 准确率(%) | 86.2 | Jet-Nemotron: 0.0, KL-LS: 10.4 | 相对KL-LS提升75.8个百分点 |
| NIAH-Single-3 (256K上下文) | 准确率(%) | 48.8 | Jet-Nemotron: 0.0, KL-LS: 11.0 | 相对KL-LS提升37.8个百分点 |
| CSR (常识推理平均) | 分数 | 55.9 | Qwen3教师: 58.5, Jet-Nemotron: 55.0 | 接近教师模型，超越Jet-Nemotron 0.9分 |
| 解码速度 (512K上下文) | 倍数提升 | 3.0x | Qwen3-1.7B | 3.0倍解码加速 |
| 预填充速度 (512K上下文) | 倍数提升 | 3.4x | Qwen3-1.7B | 3.4倍预填充加速 |
| 训练数据量 | Token数 | 2.3B | Jet-Nemotron: 400B, KL-LS: 25B | 比KL-LS少10倍，比Jet-Nemotron少174倍 |
局限与改进
作者坦诚地指出了几个重要局限性。首先，HALO流程是在FineWeb-Edu语料上训练的，该数据集主要是预训练风格数据，因此预训练模型通过后训练引入的指令跟随和对齐行为可能会在转换过程中被削弱——这是所有现有蒸馏方法的共同问题。其次，HALO的转换协议专门针对Transformer架构设计，对于其他模型架构（如纯SSM模型）的适用性需要进一步研究，尽管绝大多数公开LLM都是Transformer架构。从个人观察来看，本文的评估主要集中在CSR和NIAH任务上，对于更复杂的推理任务（如数学推理、代码生成）的表现未被充分评估；此外，实验规模主要在1.7B-8B参数量级，对于更大规模模型（如70B+）的表现有待验证；HyPE方案在从头训练场景下表现出更好的长度泛化（64倍训练长度下93.5% NIAH），但本文主要聚焦于转换场景，从头训练的潜力未被充分挖掘。
独立分析的弱点
尽管本文取得了显著成果，仍存在几个可改进的方面。第一，注意力层选择过程需要在每个候选层上评估召回和CSR性能，这需要额外的评估开销；可以探索更高效的代理指标或一次性评估策略。第二，Stage 1的独立对齐可能导致层间协同效应的缺失；可以考虑联合优化多个RNN层或引入层间一致性损失。第三，输出门机制是随机初始化的，在转换场景中可能不是最优的；可以探索基于教师模型注意力权重的更智能初始化策略。第四，HyPE中的超参数$a$（用于注意力缩放）需要通过在预训练文档上最小化损失来确定，这个过程可能不够高效；可以考虑将其作为可学习参数或使用更自适应的缩放策略。第五，GQA到MHA的转换增加了约10%的参数量，虽然论文认为这是值得的，但在资源受限场景下可能需要更精细的权衡分析。
未来方向
本文为多个研究方向奠定了基础。首先，作者指出如何高效恢复基础模型的指令跟随能力是一个开放问题，可以探索在蒸馏后进行轻量级的对齐微调。其次，HyPE方案在从头训练场景下展现出比转换更好的长度泛化（93.5% NIAH在64倍训练长度），这表明混合架构的预训练研究有巨大潜力。第三，可以探索HALO对其他RNN变体（如RWKV-7、GLA）的适用性，以及不同RNN混合器在不同任务类型上的最优选择策略。第四，将HALO应用于更大规模的模型（如Qwen3-72B）以及更多样化的Transformer架构（如Mixture-of-Experts模型）是自然的扩展方向。第五，结合本文的效率优势，可以探索混合架构在长上下文推理、长文档理解、多轮对话等实际应用场景中的部署策略。
复现评估
本文在复现性方面做得较好。作者在GitHub上开源了代码和模型检查点（https://github.com/THUNLP/hybrid-linear-attention），提供了详细的训练配置（Appendix B）、注意力层选择结果（Appendix C）和评估细节（Appendix E）。训练数据使用的是FineWeb-Edu的10B子集，这是一个公开可用的高质量预训练语料。效率测量在单卡NVIDIA A800 GPU上进行，使用Flash-Attention-2 v2.8.3和Flash-Linear-Attention v0.4.1等标准库。然而，完整复现HALO流程仍需要一定的GPU资源进行Stage 1-3的训练，虽然论文声称只需2.3B token，但对于1B+参数模型的训练仍需要相当的算力。此外，消融实验在500M参数模型上进行，这使得学术团队能够以较低成本验证设计选择的有效性。
论文图表
