本文一作孟维康是哈尔滨工业大学(深圳)博士研究生,主要研究方向是高效基础模型架构,线性 / 混合注意力机制。通讯作者张正教授是哈尔滨工业大学(深圳)长聘教授、博士生导师,长期从事高效能多模态机器学习研究。
Softmax 注意力的平方复杂度,一直是大模型迈向长上下文的最大拦路虎。线性注意力虽快,从头训练一个大型线性模型却动辄需要成千上万 GPU 时。于是,「把现成的预训练 LLM 直接线性化」成为最具性价比的路线 —— 但现有方法要么按位置粗暴划窗、把远处关键证据挡在高质量注意力之外,要么在特征映射中破坏预训练模型的范数结构,长文能力断崖式下跌。
STILL近日被 NeurIPS 2026 正式接收。STILL 的核心贡献可以概括为两点:
配合保范数的特征映射 NP-Map 与统一训练 - 推理的分块并行实现,STILL 仅用 0.04B 训练 token 就完成线性化,在 16K 上下文下实现 80% 显存削减、83% 预填充加速,并在长上下文基准上相对已有线性化方法最高取得 86.2% 的提升。

引入:层内混合,真的「混」了吗?
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:

图 5:预填充效率。1K 到 64K,STILL 在延迟与显存上全面占优。

图 6:解码效率。STILL 显存近乎恒定,长上下文解码延迟更低。
消融:每个组件都不可或缺
结语
STILL 重新回答了「什么是真正的同层混合注意力」,并给出了完整解法:
仅用 0.04B token 的训练成本,STILL 不仅在常识推理上反超原版全注意力模型,更把此前线性化方法在 16K 上下文上「全军覆没」的长文能力救了回来(S-NIAH-1 满分 100),同时带来 80% 显存削减与 83% 预填充加速。代码已开源,欢迎试用与拓展!
文章来自于微信公众号 “机器之心”,作者 “机器之心”
【开源免费】DeepBI是一款AI原生的数据分析平台。DeepBI充分利用大语言模型的能力来探索、查询、可视化和共享来自任何数据源的数据。用户可以使用DeepBI洞察数据并做出数据驱动的决策。
项目地址:https://github.com/DeepInsight-AI/DeepBI?tab=readme-ov-file
本地安装:https://www.deepbi.com/
【开源免费】airda(Air Data Agent)是面向数据分析的AI智能体,能够理解数据开发和数据分析需求、根据用户需要让数据可视化。
项目地址:https://github.com/hitsz-ids/airda
【开源免费】graphrag是微软推出的RAG项目,与传统的通过 RAG 方法使用向量相似性作为搜索技术不同,GraphRAG是使用知识图谱在推理复杂信息时大幅提高问答性能。
项目地址:https://github.com/microsoft/graphrag
【开源免费】Dify是最早一批实现RAG,Agent,模型管理等一站式AI开发的工具平台,并且项目方一直持续维护。其中在任务编排方面相对领先对手,可以帮助研发实现像字节扣子那样的功能。
项目地址:https://github.com/langgenius/dify
【开源免费】RAGFlow是和Dify类似的开源项目,该项目在大文件解析方面做的更出色,拓展编排方面相对弱一些。
项目地址:https://github.com/infiniflow/ragflow/tree/main
【开源免费】phidata是一个可以实现将数据转化成向量存储,并通过AI实现RAG功能的项目
项目地址:https://github.com/phidatahq/phidata
【开源免费】TaskingAI 是一个提供RAG,Agent,大模型管理等AI项目开发的工具平台,比LangChain更强大的中间件AI平台工具。
项目地址:https://github.com/TaskingAI/TaskingAI
【开源免费】XTuner 是一个高效、灵活、全能的轻量化大模型微调工具库。它帮助开发者提供一个简单易用的平台,可以对大语言模型(LLM)和多模态图文模型(VLM)进行预训练和轻量级微调。XTuner 支持多种微调算法,如 QLoRA、LoRA 和全量参数微调。
项目地址:https://github.com/InternLM/xtuner