Attention Sinks and Outliers in Attention Residuals¶
会议: NeurIPS 2026
arXiv: 2605.17887
代码: https://github.com/robinzixuan/OASIS
领域: 模型压缩
关键词: Attention Residuals, Attention Sinks, Activation Outliers, Low-bit Quantization, Softmax1
一句话总结¶
针对 AttnResidual 双重 Softmax 归一化诱发注意力汇聚(Attention Sink)、激活离群值与低比特量化崩塌的病态问题,本文提出 OASIS 框架,通过在 Token 级与深度级引入 Softmax1 显式空路由,并以 Token 级空后验证据自适应调减高空行为分支的深度路由分值,在零训练目标修改下大幅恢复量化精度。
研究背景与动机¶
大语言模型在向极低比特量化和超长上下文外推时,长期面临两大关键数值障碍:一是注意力汇聚(Attention Sink),即模型倾向于将大量注意力权重机械地分配给起始符(如 <|begin_of_text|>)等弱语义甚至无用 Token;二是激活值异常峰值(Activation Outliers),这些极度偏离分布的离群点急剧拉伸动态范围,导致低比特整数量化误差剧增。为了打破标准 Transformer 固定的残差累加模式、提升模型表征容量,Kimi 团队等近期提出了注意力残差(AttnResidual)架构。AttnResidual 在传统 Token 注意力之外新增了跨层残差深度路由,通过第二层 Softmax 学习浅层隐状态向深层的自适应信息流动。
然而,这一额外的灵活性引入了双重单纯形约束(Dual-simplex Constraints)。在标准 Softmax 机制下,所有概率质量必须严格归一化至真实 Token 和真实层分支上,系统缺乏显式的空操作(no-op)路径。一旦某个 Token 或某个前序残差分支在当前计算中无需贡献有效信息,双重归一化强行迫使概率分布向极少数活跃节点坍缩,进而剧烈放大起始 Token 的注意力汇聚,并在深层残差更新中激发严重的激活离群点。实验表明,AttnResidual 的激活峰值范数暴增至原生模型的数倍,导致模型在 W8A8 或 W4A4 低比特量化下困惑度严重爆炸、下游推理能力彻底崩溃。
面对双重归一化引发的集中与离群困境,现有基于剪裁或单层门控的修补手段均治标不治本。本文的核心 idea 是:通过在 Token 级和深度级双重引入 Softmax1 显式空路由以提供数学上的无害衰减通路,并建立 Token 级到深度级的空耦合机制,利用各分支注意力头分配给空状态的置信证据自适应惩罚其深度路由权重,从根本上阻断注意力和残差维度的离群放大与汇聚积累。
方法详解¶
整体框架¶
OASIS 的核心目标是在完全保留 AttnResidual 跨层路由表达能力的前提下,解除双重 Softmax 强行分配概率质量所带来的病态集中。系统整体包含两个协同推进的阶梯:首先是纯激活离群抑制机制(AoS),将 Token 注意力与深度路由中的标准 Softmax 统一替换为具有额外常数偏置的分母扩展形式(Softmax1),赋予网络主动选择“无贡献(Null)”的能力;其次是 Token 到深度的跨层级空耦合机制(Token-to-Depth Null Coupling),动态统计各前序层注意力头分配给空状态的平均质量,作为其信息承载量的度量,直接反向调节深度路由 Logit。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
In["输入前序各层残差隐状态<br/>s_0=h_1 与 s_i=f_i(h_i)"] --> TNull["双层级 Softmax1 显式空路由<br/>引入常数分母提供无损 no-op 通路"]
TNull --> Stat["分支空统计量与中心化偏差<br/>统计各头空概率 psi_i 并减去均值"]
Stat --> Coup["Token 到深度空耦合调制<br/>以 beta 强度调减非活跃分支路由 Logit"]
Coup --> Mix["空感知深度加权聚合更新<br/>输出稳定、低离群的新层隐状态 h_l"]
关键设计¶
1. 双层级 Softmax1 显式空路由:解除归一化强制分配的数学束缚 在因果注意力机制中,早期 Token 暴露给后续所有位置,常被动充当汇聚缓冲池;AttnResidual 的深度残差路由同样会选择性地坍缩至特定层分支。当某一计算步骤实际上不需要信息输入时,传统 Softmax 强制所有输出权重和为 1,迫使网络生成虚假的极大 Logit 差异以压制其他项,直接诱发了激活值极端离群。AoS 在 Token 级别与深度路由级别同时应用 Softmax1 算子,在分母中显式添加常数项 1:
这种形式使得模型分配给真实实体(真实 Token 或真实深度分支)的概率权重总和严格小于 1,未分配的概率质量自然流向为零值的虚拟空状态(Null State)。当某一注意层或残差分支处于近无操作(no-op)状态时,所有真实 Logit 可以同步回落,而无需在真实输出间强行拉开差距,从数学上保证了残差更新的连续衰减,破除了必须依赖极化权重来实现微小更新的内在约束。
2. 分支空统计量与中心化偏差:度量前序层的信息有效性 尽管双重 Softmax1 独立提供了空路由出口,但传统的两级路由在计算上彼此割裂,深度路由器无法感知某个前序残差分支内部是否已退化为无效计算。OASIS 提出利用注意力层自身的空置信度作为分支信息量的评估指标。对于任意前序经过非线性变换的源分支 \(i \in \{1, \dots, \ell-1\}\),提取其包含的全部 \(H\) 个注意力头在当前 Token \(t\) 位置分配给空状态的概率值 \(\tilde{p}_{t,\varnothing}^{(i,h)}\),计算该分支在当前 Token 上的平均空倾向:
为了避免深度路由器产生绝对偏置,并保持初始残差分支(\(i=0\) 代表未经变换的初始输入 \(h_1\),不含注意力头)的相对稳定,OASIS 仅在经过变换的各候选层之间进行统计量中心化:计算均值 \(\bar{\psi}_{\ell,t} = \frac{1}{\ell-1}\sum_{r=1}^{\ell-1} \psi_{r,t}\),并得到中心化相对偏差:
该偏差指标具有严格的物理意义:\(\Delta\psi > 0\) 表明该层相比其他候选层包含了更多不干活的注意力头,信息冗余度高;\(\Delta\psi < 0\) 则说明该层注意力被高度激活,承载了实质性推理表征。
3. Token 到深度空耦合调制:基于信息量反馈的动态路由分值修正 在获得各源分支的中心化空偏差后,深度路由器在将 Logit 输入深度 Softmax1 之前对其进行自适应注入。若原始深度路由 Logit 为 \(g_{i\to\ell,t}^{\text{old}}\),OASIS 根据空偏差反向削弱对应分支的分值:
修正后的 Logit 再次经过深度 Softmax1 归一化得到最终层混合权重 \(\tilde{\alpha}_{i\to\ell,t}\)。当某一层分支表现出高于平均水平的空注意力时,其对应的深度路由 Logit 被平滑惩罚,从而降低其在当前层残差更新中的混合权重;反之,包含高信息量密集注意力的分支获得相对加权。这一机制将 Token 级注意力内部的细粒度行为无缝传递给宏观层间拓扑选择,杜绝了无用层残差在深层被盲目放大的现象。
损失函数 / 训练策略¶
OASIS 严格保持原模型的自回归交叉熵预训练目标不变,不引入任何辅助重构损失或蒸馏约束。耦合强度系数 \(\beta\) 被参数化为一个可学习的非负标量,通过 Softplus 函数实现:\(\beta = \text{softplus}(\theta)\)。在训练初始化时,将底层参数 \(\theta\) 设为 \(-5\)(对应初始 \(\beta \approx 0.007\)),确保在训练初始阶段接近于零扰动,从而使模型平滑过渡并自适应学习最佳层间注意力耦合强度。微调沿用混合精度训练与标准的线性学习率衰减策略。
实验关键数据¶
主实验¶
针对基于 AttnResidual 架构继续预训练的主流开源模型(LLaMA-3.2-1B、Qwen3-0.6B 与 Phi-4),评估激活离群指标(全层平均峰度 Avg. Kurtosis、最大无穷范数 Max Inf. Norm)、WikiText-2 困惑度(FP16 全精度与 W8A8 整数后量化)以及低比特下 GSM8K 复杂推理表现(W4A4 OmniQuant 量化下 Pass@1)。
| 模型 | 方法 | 平均峰度 (Kurtosis) ↓ | 最大范数 (\(\|x\|_\infty\)) ↓ | FP16 困惑度 ↓ | W8A8 困惑度 ↓ | GSM8K FP16 ↑ | GSM8K W4A4 ↑ |
|---|---|---|---|---|---|---|---|
| LLaMA-3.2-1B | 原生 AttnResidual | 91.71 ± 0.08 | 84.43 ± 0.15 | 8.98 ± 0.01 | 30.19 ± 0.06 | 13.64 ± 0.17 | 12.97 ± 0.24 |
| OutEffHop | 3.08 ± 0.02 | 13.78 ± 0.08 | 9.07 ± 0.01 | 10.42 ± 0.01 | 13.58 ± 0.15 | 13.16 ± 0.21 | |
| AoS (无耦合) | 3.03 ± 0.01 | 12.54 ± 0.07 | 9.00 ± 0.01 | 9.36 ± 0.01 | 13.62 ± 0.13 | 13.34 ± 0.18 | |
| OASIS (本文) | 3.02 ± 0.02 | 12.25 ± 0.07 | 8.01 ± 0.01 | 8.16 ± 0.01 | 13.75 ± 0.12 | 13.52 ± 0.17 | |
| Qwen3-0.6B | 原生 AttnResidual | 93.04 ± 0.10 | 74.78 ± 0.13 | 10.57 ± 0.01 | 36.64 ± 0.08 | 62.58 ± 0.24 | 45.67 ± 0.86 |
| OutEffHop | 3.04 ± 0.02 | 14.16 ± 0.06 | 9.99 ± 0.01 | 10.60 ± 0.01 | 62.59 ± 0.19 | 52.36 ± 0.74 | |
| AoS (无耦合) | 3.03 ± 0.02 | 13.90 ± 0.05 | 9.35 ± 0.01 | 9.62 ± 0.01 | 62.65 ± 0.17 | 54.21 ± 0.66 | |
| OASIS (本文) | 2.95 ± 0.02 | 13.00 ± 0.05 | 9.20 ± 0.01 | 9.52 ± 0.01 | 62.68 ± 0.15 | 58.88 ± 0.59 | |
| Phi-4 | 原生 AttnResidual | 38.83 ± 0.02 | 37.96 ± 0.02 | 11.21 ± 0.00 | 1176.73 ± 0.01 | 93.78 ± 0.17 | 43.16 ± 0.25 |
| OutEffHop | 3.95 ± 0.00 | 18.70 ± 0.01 | 10.93 ± 0.00 | 16.98 ± 0.00 | 93.91 ± 0.15 | 71.42 ± 0.23 | |
| AoS (无耦合) | 3.85 ± 0.00 | 9.11 ± 0.01 | 11.09 ± 0.00 | 12.18 ± 0.00 | 93.84 ± 0.16 | 76.85 ± 0.21 | |
| OASIS (本文) | 2.27 ± 0.00 | 8.68 ± 0.01 | 10.94 ± 0.01 | 11.72 ± 0.00 | 94.02 ± 0.15 | 83.37 ± 0.21 |
消融实验¶
在 Qwen3-0.6B 模型上,对 Token 级空路由、深度级空路由以及跨层耦合机制进行系统剥离,评估极端 4-bit 量化(W4A4)条件下的数学推理恢复能力(GSM8K Pass@1 %)。
| 配置 | Token 级空路由 | 深度级空路由 | 跨层级空耦合 | W4A4 GSM8K Pass@1 (%) | 说明 |
|---|---|---|---|---|---|
| 原生 AttnResidual 基线 | × | × | × | 45.67 | 双重归一化未解绑,低比特崩溃严重 |
| 仅 Token 空路由 (OutEffHop) | ✓ | × | × | 52.36 | 仅解绑序列维度,深度仍强制归一化 |
| 仅 深度 空路由 | × | ✓ | × | 48.90 | 仅提供层间空操作,Token 汇聚未除 |
| 独立双空路由 (AoS) | ✓ | ✓ | × | 54.21 | 双层均有空通道,但两层路由决策割裂 |
| 完整 OASIS 模型 | ✓ | ✓ | ✓ | 58.88 | 跨层空信息深度耦合,发挥最大协同增益 (+4.67%) |
关键发现¶
- 离群值与量化灾难被彻底逆转:相较于 Vanilla AttnResidual,OASIS 在三大模型上平均降低激活值最大无穷范数达 81.75%,平均峰度降低 95.90%(恢复至趋近高斯分布的健康区间 \(\approx 3.0\));在 Phi-4 上成功消除 W8A8 下困惑度从 1176.73 崩塌至 11.72 的灾难性发散。
- 跨层耦合带来不可替代的额外增益:消融表明,仅依靠孤立的 Token 级或深度级 Softmax1 只能获得部分改善(52.36% 与 48.90%),唯有建立 Token 到深度的空证据反馈,深度路由器才能有效甄别并下调无效计算层,使 W4A4 性能进一步跃升至 58.88%。
- 长文本能力大幅回血:在 12K 长度的 RULER 与 LongBench 长文本测试中,原生 AttnResidual 性能暴跌超过 30-40%,而 OASIS 挽回了 77.3%~91.8% 的退化损失,证实其消除了假性注意力汇聚后,注意力重归局部语义与真实长程关联。
亮点与洞察¶
- 从数学结构层面诊断双重归一化病态:论文没有停留在“加剪裁”或“调超参”的经验修补,而是深刻剖析了 AttnResidual 架构双重单形曲面导致的条件性浓度坍缩(Lemma 5.1 与 Theorem 5.1),从理论上证明了标准归一化必须靠极端 Logit 差来表达近零更新。
- 极简而优雅的 Token-to-Depth 协同机制:将注意力头在 Token 级的“摸鱼”程度(空路由后验)反哺为深度分支的惩罚项,物理直觉极其自然,且完全在无需任何新监督信号的情况下端到端学习。
- 通用的低精度推理使能范式:不仅适配全量微调与量化,还可无缝兼容 Gated Attention 或 Clipped Softmax 等现有正交方案,为混合残差与自适应路由架构走向边缘端部署提供了可迁移的关键模板。
局限与展望¶
- 空路由比例的潜在过载风险:当前机制在某些模型或层级中可能会过度分配概率给空通道,过大的空权重可能会削弱对长距离微弱语义线索的捕捉敏感度。
- 实验骨干规模受限:主实验主要验证于 0.6B 到 14B 级别的轻量与中型骨干(LLaMA-3.2-1B、Qwen3-0.6B、Phi-4),尚未在 70B+ 等极大规模模型上充分验证极深层路由下的耦合动态。
- 推理 Kernel 的定制化需求:目前 Softmax1 与动态 Logit 偏差注入依赖标准 PyTorch 算子实现,若要在工业端 FlashAttention 等高度融合成熟的算子库中完全释放加速潜力,需要针对性地定制 GPU 底层前向与量化算子。
相关工作与启发¶
- vs OutEffHop (Hu et al., 2024): OutEffHop 率先在注意力层使用 Softmax1 以抑制 Transformer 激活离群值,但它仅关注单层 Token 维度;本文证明在 AttnResidual 的双重归一化下单层治理不足,提出了双层 Softmax1 并首次打通了跨层空耦合。
- vs StreamingLLM (Xiao et al., 2024): StreamingLLM 发现并被动接纳了 Attention Sink 现象,采用保留初始 Token KV 缓存的 engineering 策略维持长文本生成;OASIS 则从归一化根源入手,证明 Sink 是强制概率分配的产物,通过主动提供空流彻底消除不必要的伪汇聚。
- vs Gated Attention / Clipped Softmax (Bondarenko et al., 2023): 前者通过引入显式门控或硬截断 Logit 来允许 head 实施 no-op;OASIS 的 Softmax1 方案更加轻量、平滑,且通过跨层耦合使门控决策与分支拓扑紧密呼应,在量化精度指标上全面优于硬截断方案。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统揭示 AttnResidual 双重归一化的病理机制,创造性提出跨层级空证据耦合。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 大模型骨干、低比特量化(W8A8/W6A6/W4A4)、离群统计、长文本及数学常识全方位评测。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导完备,问题剖析一针见血,方法定义与工程落地严谨自洽。
- 价值: ⭐⭐⭐⭐⭐ 对新型自适应深度模型的大规模部署与高效量化具有极高的实践指导价值。