新浪新闻客户端

NeurIPS 2026 | 层内混合注意力:告别位置「伪混合」,STILL首个内容感知的层内混合注意力架构

NeurIPS 2026 | 层内混合注意力:告别位置「伪混合」,STILL首个内容感知的层内混合注意力架构
2026年10月17日 07:44 新浪网 作者
avatar
新闻眼

  (来源:机器之心)

  本文一作孟维康是哈尔滨工业大学(深圳)博士研究生,主要研究方向是高效基础模型架构,线性 / 混合注意力机制。通讯作者张正教授是哈尔滨工业大学(深圳)长聘教授、博士生导师,长期从事高效能多模态机器学习研究。

  Softmax 注意力的平方复杂度,一直是大模型迈向长上下文的最大拦路虎。线性注意力虽快,从头训练一个大型线性模型却动辄需要成千上万 GPU 时。于是,「把现成的预训练 LLM 直接线性化」成为最具性价比的路线 —— 但现有方法要么按位置粗暴划窗、把远处关键证据挡在高质量注意力之外,要么在特征映射中破坏预训练模型的范数结构,长文能力断崖式下跌。

  STILL近日被 NeurIPS 2026 正式接收。STILL 的核心贡献可以概括为两点:

  •   首个真正的同层混合注意力框架。 此前的层内方案(LoLCATs、Liger 等)只是把固定滑窗内的 token 交给 Softmax、窗外的交给线性注意力 —— 在作者看来,这种按位置一刀切的做法只是两种算子的「拼接」,算不上真正的混合。STILL 首次让 token 凭内容重要性在同一层内分工:重要的进 Softmax,其余的进线性注意力;

  •   用局部窗口预测未来有用的历史信息。 STILL 提出「自显著性分数(Self-Saliency Score)」,只需在局部滑窗内计算,就能预判一个 token 成为远端历史之后「未来还会不会被关注」,理论上保证局部 - 全局排名一致,实证重合率高达 84.8%。

  配合保范数的特征映射 NP-Map 与统一训练 - 推理的分块并行实现,STILL 仅用 0.04B 训练 token 就完成线性化,在 16K 上下文下实现 80% 显存削减、83% 预填充加速,并在长上下文基准上相对已有线性化方法最高取得 86.2% 的提升。

  •   论文题目:STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs

  •   作者:Weikang Meng, Liangyu Huo, Yadan Luo, Jiawen Guan, Jingyi Zhang, Yingjian Li, Zheng Zhang

  •   单位:哈尔滨工业大学(深圳),鹏城实验室,昆士兰大学,瑞汀灵动(上海)

  •   论文链接:https://arxiv.org/abs/2602.02180

  •   代码开源:https://github.com/ZacharyMeng/STILL

  引入:层内混合,真的「混」了吗?

  Transformer 自注意力需要计算所有 token 两两之间的交互,时间与显存复杂度均为序列长度的平方 O (N²)。为绕开从头预训练的天价成本,LoLCATs、Liger、KL-LS 等工作提出了层内混合注意力(Intra-Layer Hybrid Attention):在同一个注意力层内,用 Softmax 注意力(SA)处理一小部分 token,用线性注意力(LA)总结其余上下文,只需少量蒸馏微调即可让预训练 LLM「变身」线性模型。

  然而,这一范式存在两个关键缺陷:

  1. 路由只看位置,不看内容 —— 滑窗拼接不是真正的混合。 现有方法普遍依赖滑动窗口做路由:窗口内的 token 走 SA,窗外的全部交给 LA 笼统压缩(LoLCATs/Liger),或者像压缩稀疏注意力(CSA)那样把未选中的 token 直接丢弃(NULL)。在作者看来,这种按固定窗口划分的方式只是 Softmax 与线性注意力在位置上的「拼接」,算不上真正的混合 —— 哪些 token 值得高保真处理,应当由内容重要性决定,而不是由它离当前位置有多远决定。这一位置偏见隐含假设了「关键依赖都是局部的」—— 但关键证据往往远在窗口之外。一旦重要 token 被挤出 SA 分支,模型对它们的记忆就只剩线性注意力的有损压缩。

  图 1:层内混合注意力对比。LoLCATs/Liger 对窗外上下文统一施加线性注意力,长文容量受限;CSA 使用压缩稀疏注意力,但未选中的 token 被完全丢弃(NULL);STILL 为重要 token 保留高保真 Softmax 注意力,同时用线性注意力无损地保留全部上下文。

  2. 特征映射破坏预训练范数结构。 此前研究(如 NaLaFormer)已指出,预训练 Softmax 注意力是「范数感知」的:向量模长调制表征强度,直接影响注意力分布的尖锐度。而主流线性化方案(如 Hedgehog)在特征映射前引入可学习 MLP,训练中会自由缩放特征范数 —— 即使后续的 softmax 映射名义上是范数感知的,它面对的也已经是被「拧歪」的范数分布,与预训练模型的表征强度失配。

  这两个问题共同指明:有效的线性化需要 (i) 反映 token 重要性而非位置的路由,以及 (ii) 保持预训练范数统计、而非重新学习的核映射。STILL 正是围绕这两点展开。

  亮点一:同层混合注意力框架

  什么才算真正的「混合」?STILL 的答案是:让每个 token 凭内容重要性决定自己进入哪条分支,而不是凭位置。论文的名字即点题 ——Selecting Tokens for Intra-Layer Hybrid Attention:先选 token,再谈混合。

  图 2:STILL 层内混合线性注意力架构。先在滑窗内为每个 token 计算自显著性分数,Top-k 高分 token 路由至 Softmax 注意力,其余 token 由线性注意力处理;逐 token 选择被分块选择取代,支持并行训练与推理。

  固定预算路由。 为保证显存有界、复杂度线性,STILL 为每个注意力头设定与序列长度无关的固定 Softmax KV 缓存预算 m:其中 W 个槽位留给最近的滑窗 token(保留细粒度局部交互),其余 m−W 个槽位全部分配给按分数选出的历史显著 token,每个已完成的 chunk 保留 λ=⌊(m−W)/T_max⌋ 个。显著 token 一旦选中即永久保留,未选中的则累积进线性注意力的循环状态,全程无需重新打分或驱逐。换言之,滑窗只负责「近处」,「远处」的名额全部由内容竞争产生 —— 这是与此前位置式拼接的本质区别。

  延迟选择与分块并行。 逐 token 路由会引入不可忽视的开销、破坏并行性。STILL 改为以 chunk 为粒度做「延迟选择」:当新 chunk 开始时,才对上一个已完成 chunk 内的 token 按分数做 Top-λ 选择。训练与预填充阶段则采用分块并行形式,与解码阶段共享同一套路规则与记忆更新行为,训练 - 推理架构完全统一。

  亮点二:自显著性分数

  用局部窗口预测「未来有用」的历史

  真正的混合需要一个回答得了关键问题的路由信号:一个 token 在变成遥远的上下文之后,未来的 query 还会关注它吗? 全局注意力可以回答这个问题,但那是平方复杂度、不可部署的。STILL 的洞察是:这件事可以在局部窗口内低成本地预测。

  STILL 提出自显著性分数(Self-Saliency Score):对每个 token,在其局部滑窗内分别计算「含自身项」与「去除自身项」的两个滑窗注意力分布 a^diag 与 a^nodiag,然后用二者的 KL 型敏感度作为该 token 的分数:

  直觉上,一个 token 对自己的注意力分布扰动越敏感(自交互越强),它在全上下文注意力下也越倾向于保持高排名。分数高的 token 被路由到 SA 分支,其余交给 LA,路由决策由内容而非位置驱动。

  这一「局部预测全局」的能力有严格的理论保证:

  命题 1(局部 - 全局一致性):设 p_local 与 p_global 分别为局部与全局自注意力概率,当 ε→0⁺ 时,自显著性分数满足:(1) 单调性 ——∂Score_t/∂p_local > 0;(2) 保序性 ——Score_i > Score_j 蕴含 E [p^(i)_global] > E [p^(j)_global]。即局部算出的分数保留了全局排名的期望序。

  实验可视化也证实了这一点:在滑窗内被标红的高分 token,放到全上下文注意力下依然高分 —— 大多数 token 分数近零、适合 LA,只有极少数脱颖而出、值得 SA 精算。

  图 3:自显著性分数分布。左:滑窗注意力下分数最高的 token(红点)被路由到 SA;右:把这些 token 放到全上下文注意力中依然高亮,展现出强局部 - 全局一致性。

  一个有趣的佐证是词云分析:被 SA 分支最频繁选中的,是否定词(not, no, never)、情态词(will, would, may)与量化词(all, any, nothing)这类「高辖域」功能词 —— 它们对整句语义有全局影响,恰恰最需要高保真的 Softmax 注意力来精确建模。

  图 4:自显著性分数最常选入 SA 分支的 token 词云。高频入选者多为否定、情态、量化等高辖域功能词,说明 SA 优先保留精确的、上下文依赖的交互。

  局部预测到底准不准? 作者还做了一组直接验证:先用局部算出的自显著性分数选出 Top-224 的 token,再运行全注意力,统计后续 query 真正重视的远端 key 与它们的重合率。结果显示,在浅、中、深层的代表性注意力头上,重合率达到 62.9%–84.8%—— 局部窗口内的分数,确实是在线可用的「未来 KV 价值」代理。这正是「用局部窗口预测未来有用的历史信息」的直接证据。

  NP-Map:把预训练范数「原样还回去」

  针对范数失真问题,STILL 提出保范数特征映射(Norm-Preserved Feature Map, NP-Map):不再放任 MLP 自由缩放特征,而是显式地把「方向」与「模长」解耦 ——MLP 只负责精炼方向,预训练的原始范数被重新注入:

  其中 f (⋅) 为可学习 MLP。通过把原始范数 ‖x‖ 重新注回变换后的特征,NP-Map 保证注意力机制严格对齐预训练模型的表征强度,在满足非负约束的同时完整继承预训练 Transformer 的细粒度语义信息。此外,STILL 还引入门控机制提升线性注意力的秩,完整线性分支为:

  实验:短文本反超原版,长文本「复活」

  训练上,STILL 沿用两阶段流水线(Attention Transfer + 低秩线性化),仅在 Cleaned Alpaca 数据上用 20M token、1024 序列长度训练,总成本仅 0.04B token。

  常识与通用推理:

  平均反超原版 Transformer

  在 PIQA、ARC-E/C、HellaSwag、WinoGrande、MMLU 上,以 Llama 3 8B / 3.1 8B 为骨干,STILL 取得全部线性化方法中的最优平均分。Llama 3.1 8B 下,MMLU 从 Liger 的 46.9 大幅提升至 61.3;不含 MMLU 的平均分 74.7,甚至反超原版全注意力模型(73.2)1.5 个点。

  表 1:常识与通用推理任务对比。STILL 在 Llama 3 8B 与 Llama 3.1 8B 上均取得线性化方法中的最优平均表现。

  S-NIAH 大海捞针:

  4K 下 86.2 分,基线几乎全军覆没

  在相同甚至更小的缓存预算下(STILL 缓存总量不超过基线窗口大小),STILL 在 0.5K–4K 各长度全面领先,且差距随长度拉大:4K 时 STILL 达 86.2,而 LoLCATs 仅 8.8、Liger 0.0、SWA 5.8—— 相对最强基线提升 +77.4。

  RULER 全基准:

  47.9 vs 7.2 / 2.1

  在需要长程记忆与状态追踪的 8 项 RULER 子任务上(4K 上下文),STILL 以 128 滑窗 + 896 总缓存上限(基线窗口为 896)取得最高平均分 47.9,大幅超过 LoLCATs(7.2)与 Liger(2.1),也高于 Mamba2(44.7)。

  表 2(左):S-NIAH 任务在 0.5K–4K 上下文、相同缓存预算下的准确率对比。表 3(右):4K 上下文 RULER 全基准对比。

  BABILong 长文推理

  在严格 256 token 总缓存预算下,STILL 在所有长度上全面最优:4K 时 QA1 从 LoLCATs 的 3 提升至 10,QA5 从 6 提升至 24,展现出更强的长程证据聚合能力。

表 4:BABILong 基准在 0K–4K 上下文长度下的对比。

  与层间策略兼容:

  16K 下 S-NIAH-1 满分

  将 STILL 的层内机制嵌入 KL-LS 层间框架(Llama 3.2 1B,25% 层保留全注意力):STILL 在 4K/8K/16K 上 S-NIAH-1 全部满分 100,16K 下 S-NIAH-2 达 74.0(LoLCATs 仅 8.6,+65.4),RULER 总分 27.20,超出 Zebra-Llama(13.37)13.83 分。

  表 5:层内 - 层间混合架构在 RULER 基准 4K–16K 上的结果,所有模型均从 Llama 3.2 线性化而来。

  效率:16K 预填充耗时仅为 Softmax 的 17%

  在单张 H200 上实测 1K–64K:

  •   预填充:Softmax 注意力在 16K 之后遭遇显存瓶颈,STILL 平稳扩展至 64K,16K 时耗时仅为 Softmax 的 17%;延迟选择策略使 STILL 相比 LoLCATs 只增加极小的延迟与显存开销。

  •   解码:Softmax 在 32K 之后 OOM,STILL 显存近乎恒定,16K 时仅为 Softmax 的 20%,吞吐稳定。

图 5:预填充效率。1K 到 64K,STILL 在延迟与显存上全面占优。
图 6:解码效率。STILL 显存近乎恒定,长上下文解码延迟更低。

  消融:每个组件都不可或缺

  •   模块消融(BABILong 1K):去掉门控 −5.8 分,去掉 NP-Map −4.9 分,两者同去 −8.3 分;再去掉自显著性选择,准确率从 27.5 跌至 12.5(−15.0)。

  •   路由指标对比:在 S-NIAH 上,自显著性平均 42.6 分,远高于对角注意力(23.1)、q・k 内积(22.8)、key 范数(12.9)、均匀(12.3)与随机选择(12.8)—— 对角注意力在短上下文尚可近似,长程稀疏保留场景下无法替代自显著性。

  •   chunk 大小:C=64 为最佳平衡点;C 从 16 增至 128 训练时间降 35.4%、显存降 21.7%;逐 token 路由(C=1)会使解码选择延迟增加 3 倍以上、训练速度减半。

  •   骨干泛化:在 Llama 3.2 1B/3B 与 Mistral 7B 上同样最优 ——3B 下 MMLU 52.4(Liger 32.1),Mistral 7B 下含 MMLU 平均 70.9(LoLCATs 69.1,Liger 65.1)。

  结语

  STILL 重新回答了「什么是真正的同层混合注意力」,并给出了完整解法:

  •   首个内容感知的同层混合框架:token 凭内容重要性而非位置进入 Softmax 分支,滑窗只负责近处、远端名额由内容竞争产生,配合固定缓存预算与统一训练 - 推理的分块并行,实现 O (N) 线性复杂度;

  •   用局部窗口预测未来有用的历史:自显著性分数有局部 - 全局一致性的理论保证(单调性 + 保序性),与未来全注意力查询实际重视的远端 key 重合率达 62.9%–84.8%;

  •   NP-Map 保范数特征映射:方向与模长解耦、重注预训练范数,让线性核严格对齐预训练表征强度。

  仅用 0.04B token 的训练成本,STILL 不仅在常识推理上反超原版全注意力模型,更把此前线性化方法在 16K 上下文上「全军覆没」的长文能力救了回来(S-NIAH-1 满分 100),同时带来 80% 显存削减与 83% 预填充加速。代码已开源,欢迎试用与拓展!

特别声明:以上文章内容仅代表作者本人观点,不代表新浪网观点或立场。如有关于作品内容、版权或其它问题请于作品发表后的30日内与新浪网联系。
权利保护声明页/Notice to Right Holders

举报邮箱:jubao@vip.sina.com

Copyright © 1996-2026 SINA Corporation

All Rights Reserved 新浪公司 版权所有