STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs¶

会议: NeurIPS 2026
arXiv: 2602.02180
OpenReview: https://openreview.net/forum?id=rgEJDjkYDK
论文: https://arxiv.org/abs/2602.02180
代码: https://github.com/SMULL-Group/STILL
领域: LLM 其他
关键词: 大模型线性化, 层内混合注意力, 自显著性评分, 模长保持特征映射, 块级延迟选择

一句话总结¶

针对预训练大语言模型线性化中位置路由丢失全局长程依赖与可学习映射扭曲特征模长的问题，STILL 提出局部自显著性评分筛选显著 token 维持 Softmax 注意力、模长保持特征映射（NP-Map）消除表征强度失真，并借助块级延迟选择实现了常数级显存占用与长上下文无损保留。

研究背景与动机¶

基于 Transformer 的自回归大语言模型虽然具备出色的长程上下文建模能力，但标准 Softmax 注意力在时间与显存开销上随序列长度呈二次方增长，成为长文本扩展的关键瓶颈。从头预训练混合架构（如 MiniMax-Text-01、Kimi-linear 或 DeepSeek-V4）需要消耗庞大的算力与海量 GPU 时数，使得学界普遍难以承受。为此，近期以 LoLCATs、Liger 和 KL-LS 为代表的层内混合注意力（Intra-Layer Hybrid Attention）范式备受关注：它们直接继承预训练权重，在单个 Transformer 层内将注意力拆分为滑动窗口局部 Softmax 注意力（SA）与全局线性注意力（LA），仅用少量数据微调即可实现线性复杂度。

然而，现有层内混合范式存在两大根本缺陷。首先，现存的 Token 路由机制严重依赖滑动窗口等静态绝对位置启发式，隐式假设关键依赖仅存在于局部邻域；对于远距离上下文，现有方法要么采用均匀线性注意力进行粗糙压缩，要么直接截断舍弃，导致窗口外至关重要的长程证据丢失高保真 Softmax 建模机会。其次，预训练 Softmax 注意力具有“模长感知（norm-aware）”特性，即向量的幅度调节着表征强度和注意力分布的尖锐程度；而现有方案引入的可学习特征映射（如基于 MLP 的 Hedgehog）在微调过程中可以任意缩放特征范数，破坏了预训练模型的几何分布，使得后续线性核函数在失真的模长统计量上计算。

本文的切入角度是：如果能仅通过因果局部计算准确估计 Token 的全局重要性，并从根本上解耦特征方向与模长，就能用极低的预算将预训练模型无损线性化。核心 idea：STILL 提出具有局部-全局一致性的自显著性评分（Self-Saliency Score）在局部滑动窗口内筛选全局显著 token 予以保留在固定预算的 Softmax 缓存中，配合解耦方向与幅度、重注预训练范数的模长保持特征映射（NP-Map）来汇聚其余背景上下文，并在块级粒度上延迟决策以兼顾硬件并行效率。

方法详解¶

整体框架¶

STILL 将预训练 Transformer 解码器的每个自注意力层转换为层内混合线性注意力模块。整个流水线分为三大协同阶段：首先在局部滑动窗口内计算各 Token 的自显著性评分；随后在块（Chunk）级别执行延迟路由决策，将得分最高的少量显著 Token 分配并永久保留在固定容量的 Softmax 注意力（SA）缓存中，其余 Token 则被投影并累积进循环线性注意力（LA）状态；最后，线性注意力分支通过解耦重构的 NP-Map 计算状态累加，与局部及历史显著 SA 分支的输出进行统一归一化融合，得到最终的注意力输出。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    IN["输入序列与投影<br/>{q, k, v, g}"] --> S3["自显著性评分与局部全局一致性<br/>滑动窗口内评估自身注意力敏感度"]
    S3 --> CHUNK["块级延迟选择与固定预算路由<br/>按 Chunk 批量排序选出 Top-λ 进入 SA"]
    CHUNK -->|未选中上下文| LA["模长保持特征映射 NP-Map<br/>解耦方向与模长 + 门控线性累积"]
    CHUNK -->|局部与历史显著 Token| SA["固定预算 Softmax 缓存<br/>局部滑动窗口 + 显著关键证据"]
    LA --> COMB["双分支统一归一化融合<br/>分子累加 / 分母累加合并输出"]
    SA --> COMB
    COMB --> OUT["最终层注意力输出 y"]

关键设计¶

1. 自显著性评分与局部全局一致性：摆脱基于位置的启发式路由

现有层内混合方案全凭相对位置进行滑动窗口截断，无法感知不同 Token 的实际语义贡献。STILL 引入自显著性评分（Self-Saliency Score, $\mathrm{Score}_t$），核心是度量当前 Token 本身在局部注意力分布中的支配地位。在半径为 $w$ 的局部滑动窗口 $\mathcal{W}_t$ 内，分别计算包含对角元（自身）与移除对角元后的两条注意力概率分布 $\mathbf{a}^{\mathrm{diag}}_{t,j}$ 与 $\mathbf{a}^{\mathrm{nodiag}}_{t,j}$： $$ \mathbf{a}^{\mathrm{diag}}_{t,j} = \frac{\exp(\mathbf{q}_t \mathbf{k}j^\top)}{\sum}_t} \exp(\mathbf{q}_t \mathbf{km^\top)}, \quad \mathbf{a}^{\mathrm{nodiag}}} = \frac{\exp(\mathbf{q}_t \mathbf{kj^\top)}{\sum}_t \setminus {t}} \exp(\mathbf{qt \mathbf{k}m^\top)} \ (j \neq t) $$ 通过两者之间的 KL 散度形式得到自显著性得分： $$ \mathrm{Score}t = \sumt} \mathbf{a}^{\mathrm{diag}}} \log \frac{\mathbf{a}^{\mathrm{diag}} + \epsilon}{\mathbf{a}^{\mathrm{nodiag} $$ 论文从理论上严格证明了局部-全局一致性（Local-Global Consistency）：在概率期望意义下，局部自显著性评分关于局部自注意力概率严格单调递增，并且保序地反映全局自注意力概率的排序（} + \epsilon$\mathrm{Score}_i > \mathrm{Score}_j \implies \mathbb{E}[p_{\mathrm{global}}^{(i)}] > \mathbb{E}[p_{\mathrm{global}}^{(j)}]$）。可视化分析表明，该指标高度优先捕获否定词（not, no, never）、情态动词（will, would）和量词（all, any）等对全局语义至关重要的高阶功能算子，确保真正决定上下文逻辑的 Token 获得高保真 Softmax 建模。

2. 模长保持特征映射 NP-Map：对齐预训练表征强度

预训练 Softmax 注意力本质是模长感知的，向量模长直接控制注意力权重的尖锐程度。此前类似 Hedgehog 的方法使用轻量 MLP 变换特征后再做 Softmax 映射，但在微调时 MLP 会自由缩放特征范数，破坏预训练模型的模长一致性。STILL 提出 Norm-Preserved Feature Map (NP-Map)，在经由 MLP $f(\cdot)$ 进行变换时，将特征的方向与模长显式解耦，仅利用 MLP 调整特征方向，而强制注入预训练原始键值向量的欧氏范数 $\|\mathbf{x}\|$： $$ \mathbf{u} = \frac{f(\mathbf{x})}{|f(\mathbf{x})|} |\mathbf{x}|, \quad \phi_{\mathrm{NP}}(\mathbf{x}) = [\operatorname{softmax}(\mathbf{u}) \oplus \operatorname{softmax}(-\mathbf{u})] $$ 通过将非负性特征映射与重注模长结合，NP-Map 避免了分布漂移并完美适配预训练几何。此外，线性分支进一步引入门控机制 $\mathbf{g}_t$，不仅提升了线性注意力状态矩阵的有效秩，而且仅引入极微小的计算开销。

3. 块级延迟选择与固定预算路由：兼顾硬件并行度与恒定内存

如果对每个生成步骤逐 Token 进行即时打分与路由，会导致严重的 GPU 分支发散并破坏核函数并行。STILL 采用块级延迟选择（Delayed Selection Strategy）：以块（Chunk 大小为 $C$）为计算单元，当最新完成的一个块 $a$ 离开局部窗口时，在进入下一个块首位置时对其内部全部 $C$ 个 Token 集中排序，仅选取 Top-$\lambda$ 个 Token 的键值对永久放入显著 SA 缓存，其余 $C-\lambda$ 个 Token 则以批量矩阵乘法一次性折叠进线性注意力的循环隐状态 $(\mathbf{s}_t, \mathbf{z}_t)$。

为了在理论与工程上确立严格的线性复杂度与恒定显存，STILL 为每个注意力头预先分配固定上限为 $m$ 的 KV 缓存槽位，其中固定预留 $W$ 个槽位给局部滑动窗口，剩余 $m-W$ 个槽位平均分配给长上下文中的历史显著块（即每块配额 $\lambda = \lfloor(m-W)/\lceil N_{\max}/C \rceil\rfloor$）。一旦选入 SA 缓存便不再被剔除或重新打分，使得整个注意力计算的理论复杂度为 $\mathcal{O}(Nmd) + \mathcal{O}((N-m)dd') + \mathcal{O}(T \log C) = \mathcal{O}(N)$，既满足了预填阶段的块级高度并行，又满足了自回归解码阶段的固定常数显存占用。

损失函数 / 训练策略¶

STILL 严格沿用两阶段超轻量训练流水线，总共仅消耗不到 0.04B（约 20M）训练 Token： - 第一阶段（Attention Transfer）：冻结预训练 Teacher 模型主干参数，仅在 Alpaca 数据集（序列长度 1024）上通过均方误差（MSE）损失优化 NP-Map 中的轻量 MLP 参数与门控投影矩阵 $\mathbf{W}_g$，使线性注意力核函数输出尽可能拟合原始 Softmax 注意力图。 - 第二阶段（Low-rank Linearization）：对查询、键、值、输出投影矩阵及门控层（$\mathbf{W}_q, \mathbf{W}_k, \mathbf{W}_v, \mathbf{W}_o, \mathbf{W}_g$）启用低秩自适应（LoRA）微调，端到端对齐混合算子与原模型的长短程表征分布。

实验关键数据¶

主实验¶

在常识与通用推理基准（PIQA, ARC-e, ARC-c, HellaSwag, WinoGrande, MMLU）以及具有挑战性的长程检索基准（RULER S-NIAH-1 任务，4K/8K 长度）上，STILL 显著超越了此前的线性化方法，甚至在推理任务上超越了全注意力原始底座。

模型	底座	训练 Token	常识推理平均 (w. MMLU)	MMLU (5-shot)	S-NIAH-1 (4K)	S-NIAH-1 (8K)
Llama 3 8B (全注意力)	-	15T+	71.8	65.3	100.0	100.0
LoLCATs	Llama 3 8B	0.04B	67.3	43.9	8.8	0.0
Liger-GLA	Llama 3 8B	0.02B	67.6	43.4	0.0	0.0
STILL (本文)	Llama 3 8B	0.04B	71.5	59.7	86.2	36.2
Llama 3.1 8B (全注意力)	-	15T+	74.2	68.0	100.0	100.0
LoLCATs	Llama 3.1 8B	0.04B	69.9	50.8	8.8	0.0
Liger-GLA	Llama 3.1 8B	0.02B	68.3	46.9	0.0	0.0
Llamba	Llama 3.1 8B	12B	71.5	60.0	-	-
STILL (本文)	Llama 3.1 8B	0.04B	72.5	61.3	86.2	36.2

在全套 4K 长度 RULER 复杂状态追踪测试（涵盖多键检索、变量跟踪、长文档问答等 8 项子任务）中，STILL 取得 47.9 的平均分，大幅领先 LoLCATs（7.2）、Liger-GLA（2.1）甚至优于从头训练的 Mamba2（44.7）。而在层内-层间联合混合（Intra-Inter Hybrid）架构扩展实验中，STILL 更是完美解决了 16K 长度下的 S-NIAH-1 检索（得分 100.0），16K RULER 综合得分达到 27.20，大幅领先 Zebra-Llama（13.37）与 LoLCATs（7.59）。

消融实验¶

消融实验深入验证了各个关键模块在长程问答基准 BABILong（1K 长度，QA 任务平均准确率）及不同路由策略在长程检索 S-NIAH-1 中的不可替代性。

配置	变体 / 路由指标	评估基准与指标	相比完整模型变化
完整模型 (STILL)	Self-Saliency + NP-Map + Gate	BABILong (1K) 准确率: 27.5	基准
w/o Gate	移除门控机制	BABILong (1K) 准确率: 21.7	-5.8
w/o NP-Map	移除模长保持，退化为传统 MLP	BABILong (1K) 准确率: 22.4	-4.9
w/o Gate & NP-Map	传统线性映射且无门控	BABILong (1K) 准确率: 19.2	-8.3
w/o Self-Saliency	移除自显著性，仅靠窗口与传统压缩	BABILong (1K) 准确率: 12.5	-15.0
路由指标消融	自显著性评分 (Self-Saliency)	S-NIAH-1 (0.5K/1K 平均): 42.6	基准
对比指标 1	仅对角自注意力 ($a^{\mathrm{diag}}_{t,t}$)	S-NIAH-1 平均: 23.1	-19.5 (0.5K 下跌 27.6)
对比指标 2	键值内积 ($\mathbf{q}_t^\top \mathbf{k}_t$)	S-NIAH-1 平均: 22.8	-19.8
对比指标 3	键向量范数 ($\\|\mathbf{k}_t\\|$)	S-NIAH-1 平均: 12.9	-29.7
对比指标 4	随机选择 / 块内均匀抽取	S-NIAH-1 平均: 12.8 / 12.3	-29.8 / -30.3

关键发现¶

局部显著性是对全局注意力的强保序代理：在未来 KV 效用实测中，因果局部计算出的 Top-224 显著 Token，能够精确覆盖未来真实全注意力查询实际重点关注的远程键的 62.9%～84.8%（表 12），从根本上解释了为何不需要二次扫描全局即可做出精准路由。
纯对角注意力无法替代相对变化分布：虽然对角注意力在短上下文推理上仅比自显著性微跌 0.5 分，但在长程证据检索（S-NIAH-1）中准确率直接腰斩（从 42.6 暴跌至 23.1）。这证明在深层稀疏筛选中，剔除自身后引起邻域注意力分布重构的扰动程度（KL 散度式结构）才是判别关键承载词的真正信号。
硬件显存与延迟优势：在单张 NVIDIA H200 GPU 上，当上下文达到 16K 时，STILL 的预填时间仅为标准 Softmax 注意力的 17%；在自回归生成解码中，STILL 的显存占用随长度从 1K 到 64K 完全保持恒定平滑（平直曲线），显存仅为 16K 时 Softmax 的 20%，彻底消除了长文本 OOM 现象。

亮点与洞察¶

解耦模长与方向的即插即用修复：NP-Map 明确指出了此前线性化模型中被普遍忽略的“注意力模长敏感性”问题，通过简单的单位球投影与原模长重注，无需引入新权重就从几何底层修复了微调带来的特征漂移。
硬件友好的延迟块级路由设计：避免了在自回归每一步进行单 Token 排序带来的分支开销，利用微小的时间延迟在块边界进行批量调度，以 $C=64$ 实现了最优的吞吐与长文本召回平衡。
极高的数据与算力性价比：仅需 0.04B（2000 万）Token 的微调训练，便成功将 8B 规模的 Llama 3/3.1 和 Mistral 7B 无损线性化，且有效恢复了 16K 以上极其脆弱的长距离检索与联想能力。

局限与展望¶

模型规模尚未外推至更大参数量：论文受限于计算资源，实证验证主要停留在 1B 到 8B 参数级别（Llama 3/3.1、Mistral 7B），尚未在 70B 及以上超大规模模型上充分验证极长文本下的路由鲁棒性。
静态超参配置尚未自适应：当前的块大小 $C$ 和各头配额 $\lambda$ 采用全层一致的静态预设，未来可结合层级/注意力头功能分工实现注意力头自适应的预算动态分配。
模态局限于纯文本：目前 STILL 仅在自回归纯文本模型上进行验证，扩展到多模态大型模型（VLM）处理长视频流和高分辨率特征序列是极具前景的后续方向。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 提出结合局部-全局一致性的因果自显著性评分与模长保持特征映射，理论推导扎实，优雅地解决了线性化中的长距离信息与表征几何丢失问题。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 Llama 3、3.1、3.2、Mistral 多种底座，在常识推理、RULER、BABILong 及 64K 真实硬件显存延迟上均有详实对比与消融。
写作质量: ⭐⭐⭐⭐⭐ 叙述逻辑严密，图表清晰直观，理论证明与工程实现细节完备。
价值: ⭐⭐⭐⭐⭐ 为开源大型语言模型低成本、高保真线性化提供了兼具理论自洽性与极高工程可用性的优秀范式。