摘要本文解读 ICML 2026 论文《Attention Sink Forges Native MoE in Attention Layers: Sink-Aware Training to Address Head Collapse》北京大学集成电路学院与人工智能研究院。该论文提出注意力 Sink 即隐式路由器这一核心洞察通过融合隐式门控等价性证明、头部坍缩的 MoE 归因与Sink 感知负载均衡损失证明 Vanilla / Sink / Gated 三种注意力机制在数学上统一为注意力层内的原生 MoE其特别之处在于把要消除的伪影重新表述为要训练的路由器。实验表明辅助损失在 0.6B/1B/2B 全部九组配置上一致提升最高 1.40pp微调 LongBench 最高 4.71pp训练开销 2%为高效大模型训练与注意力机制设计提供了重要借鉴。视频讲解点击观看 B 站视频摘要论文基本信息背景与动机研究主线从问题到结论基准/方法设计分类全景方法细节实验设计与结果结果对比总结关键发现局限性常见问题FAQ注意力 Sink 是什么Sink 为什么是 MoE 路由器什么是头部坍缩head collapseSink 感知辅助损失如何工作训练开销大吗为什么消除 Sink 能提升性能参考链接论文基本信息项目内容标题英文Attention Sink Forges Native MoE in Attention Layers: Sink-Aware Training to Address Head Collapse标题中文注意力 Sink 锻造注意力层的原生 MoE Sink 感知训练解决头部坍缩作者Zizhuo Fu (付子卓), Wenxuan Zeng, Runsheng Wang (王润声), Meng Li (李萌)机构Institute for Artificial Intelligence School of Integrated Circuits, Peking University会议ICML 2026arXivhttps://arxiv.org/abs/2602.01203项目网站无公开项目页详见 arXiv 与视频讲解背景与动机大语言模型LLM普遍存在注意力 Sinkattention sink现象模型对序列第一个 token 分配不成比例的注意力权重与该 token 的语义无关。这一现象自StreamingLLM2023命名以来被广泛应用于 KV cache 优化、长上下文推理、注意力头剪枝与量化。围绕如何消除 Sink业界形成两条路线Sink AttentionGPT-OSS2025在 softmax 分母中加入可学习 sink 参数吸收冗余注意力Gated AttentionQwen3-Next2025为每个注意力头引入 sigmoid 显式门控因子。但现有工作存在两个关键缺口其一消除 Sink 为何能提升性能——机制层面没有统一解释其二DuoAttention2025等观测发现只有固定子集注意力头真正参与生成零化超过 25% 的头几乎不损失精度——这一头部坍缩现象同样没有成因解释。本文的工作恰好填补这两个缺口Sink 不是有害冗余而是隐式路由器头部坍缩就是 MoE 专家坍缩。从历史视角看Sink 研究经历现象2023→ 机理2024value drain、Active-Dormant→ 可利用结构2025–2026的演进本文是把 Sink 变成训练信号的转折点。研究主线从问题到结论图 9Mermaid 流程图研究主线——问题 → 动机 → 设计 → 方法 → 实验 → 结论。基准/方法设计本文方法的根基是一个统一门控视角注意力层天然就是混合专家MoE——每个注意力头是独立专家sink 上的注意力权重充当路由器。对 Vanilla / Sink Attention冗余注意力被吸收到零贡献组件首 token 的 value drain或可学习 sink 参数因此输出可改写为$O_t^{l,h} (1 - A_{t,\text{sink}}) \cdot \sum_{j \neq \text{sink}} \tilde{A}_{t,j} \mathbf{v}_j$其中 $\tilde{A}$ 是非 sink token 的重归一化 softmax 权重因子 $G_{t}^{l,h} 1 - A_{t,\text{sink}}$ 就是隐式门控与 Gated Attention 的显式门控 $G_{t}^{l,h} \sigma(\mathbf{x}t^l W\theta^{l,h})$ 直接对应。Sink Attention 的门控还可写成 $\sigma(\mathrm{LSE} - \mathbf{sink})$ 的 sigmoid 等价形式——逐头独立 sigmoid 激活恰好复现 Gated Attention 的两大设计选择附录 A 有完整推导。图 1Vanilla / Sink / Gated 三种注意力机制对比——前两者的 sink 权重等价于第三者的显式门控。分类全景图 10Mermaid 流程图三种注意力机制的分类全景——Vanilla / Sink / Gated 统一到隐式门控。方法细节基于门控视角本文提出Sink 感知训练三要素Head importance头部重要性$\mathrm{Imp}^{l,h}$ 定义为门控因子 $G^{l,h}$ 在整批 token 上的均值——三种机制用同一把尺子度量头部激活。Head imbalance头部不均衡度每层内 importance 的变异系数 $\mathrm{CV} \mathrm{std}/\mathrm{mean}$取所有层的平均——量化坍缩严重程度。辅助负载均衡损失$\mathcal{L}_{\text{aux}} \lambda \sum_l N_h [\mathrm{CV}]^2$借鉴 MoE 负载均衡思想Switch Transformer / DeepSeek-MoE从零训练时直接鼓励各头均匀利用。微调适配是关键工程点直接把损失套到预训练模型上会产生头部钉死效应——主导头门控抵抗变化其余头被迫对齐。本文借鉴 DeepSeek-MoE 的共享专家 路由专家设计把每层 top-$m$ 重要头保留为共享头不动仅对剩余路由头施加负载均衡。Flash Attention 集成则利用前向已有的 log-sum-expLSE值计算门控无需退回 eager 模式训练延迟开销小于 2%。图 2Value drain 现象——模型训练出零贡献吸收器承接冗余注意力这是隐式门控成立的物理基础。图 3消除 Sink 后查询-键几何解绑——softmax 检索精度更高这是长上下文受益的几何解释。机制上本文提出坍缩链条假设value drain → attention sink → 隐式门控 → head collapse。Gated Attention 的门控从训练开始就显式存在因此坍缩最早、最剧烈Vanilla 需要先完成 value drain 才能建立门控坍缩最晚——门控形成是坍缩的先决条件。实验设计与结果实验设置从零训练覆盖 3 机制 × 3 规模0.6B/1B/2B× {base, baseaux}数据为 FineWeb-Edutoken 预算按 Chinchilla 法则取参数量的 20 倍评测 10 项基准MMLU、GSM8K、HumanEval、ARC-E/C、OpenBookQA、HellaSwag、PIQA、BoolQ、WinoGrande平均准确率。微调实验在 Qwen3-4B/8B、LLaMA3.1-8B 上以 LoRArank 16在 AceReason-Nemotron 长序列推理数据上训练评测 6 项推理基准 6 项 LongBench 任务。主结果10 项基准平均准确率模型规模注意力机制basebaseauxΔ0.6BVanilla39.5440.821.280.6BSink40.3141.110.800.6BGated40.4341.350.921BVanilla42.5943.350.761BSink43.2743.750.481BGated42.6943.981.292BVanilla45.4646.831.372BSink46.3546.500.152BGated46.3047.701.40图 4三种代表模型的头部坍缩热图——固定子集 head 高激活、其余低激活跨机制普遍存在。图 5不均衡度随训练上升并平台化坍缩自然涌现baseaux 曲线被压低并稳定在低位。图 6无辅助损失时激活集中于少数 head 且跨数据集雷同有损失时分布均匀且随任务分化。图 7Sink 与 Gated 从训练早期即持续领先 Vanilla支持消除 Sink 提升表达力的分析。微调结果推理平均 Qwen3-4B 74.73→75.24、Qwen3-8B 74.31→76.01、LLaMA3.1-8B 56.50→58.97LongBench 平均 Qwen3-4B 51.46→52.75、Qwen3-8B 50.67→55.384.71、LLaMA3.1-8B 40.51→43.98。亮点LLaMA3.1-8B 的 ProcessBench 从 23.63 跃升至 31.147.51。结果对比总结图 11Mermaid 流程图结果对比总结——辅助损失全配置提升微调推理与长上下文双涨。关键发现Sink 权重就是隐式门控$G 1 - A_{\text{sink}}$ 与显式 sigmoid 门控数学等价附录严格推导三种注意力机制统一为注意力层原生 MoE——最强双模式 Oral 组合P9 等价性证明 P1 假设审计扭转的典型画像。头部坍缩 专家坍缩固定子集 head 持续主导是门控建立后的自然结果LLaMA-3.1-8B、GPT-OSS-20B、Qwen3-Next-80B 均呈现同样格局与 DuoAttention 等零化 25% 头几乎无损的观测一致。辅助损失全配置一致有效3 机制 × 3 规模 9 组配置全部提升最高 1.40pp2B Gated收益集中在训练后期且模型越大收益越大与 MoE 负载均衡文献趋势一致。坍缩时序揭示机制差异Gated 早期骤升后平台化显式门控立即专家化Vanilla/Sink 渐进需先建立门控——门控是坍缩的先决条件。微调双涨Qwen3-8B LongBench 4.71pp、LLaMA3.1-8B 推理 2.47pp、ProcessBench 7.51pp共享头 路由头策略避免破坏已学行为。工程代价极低Flash Attention 经 LSE 集成训练延迟 2%——免费的负载均衡。图 8机制假设附录——value drain → attention sink → 隐式门控 → head collapse门控形成是坍缩的先决条件解释了三种机制的坍缩时序差异。局限性单 Sink 假设隐式门控假设只有一个 sink 组件但标点等低语义 token 也可能吸收冗余注意力一般化形式应为 $G 1 - \sum_{j \in \mathcal{S}} A_{t,j}$。DeepSeek-V4 反例sink 参数未吸收全部冗余时首 token 仍获高注意力——单组件度量会低估真实头部激活。多门控交互显式门控与隐式 sink 并存时的完整刻画留作未来工作。规模边界预训练验证限于 2B 以下微调限于 8B 以下更大规模的坍缩动力学未验证。常见问题FAQ注意力 Sink 是什么LLM 对序列第一个 token 分配不成比例的注意力权重与该 token 语义无关这一现象由 StreamingLLM2023命名为 attention sink。它源于 softmax 归一化当头部不需要满额分配时冗余权重被倒进一个零贡献组件。Sink 为什么是 MoE 路由器冗余注意力被吸收到零贡献 sink 上后输出可分解为门控因子 $(1 - A_{\text{sink}})$ 乘重归一化加权和。这个因子逐头独立、呈 sigmoid 形式与 Gated Attention 的显式门控完全同构——head 是专家门控是路由注意力层天然满足 MoE 定义。什么是头部坍缩head collapse只有固定子集注意力头持续高激活、其余头低激活的现象等价于 MoE 的专家坍缩。它由门控建立后的负载不均衡自然涌现主导头获得不成比例的梯度更新其余头被锁死在低利用率。Sink 感知辅助损失如何工作用门控因子定义每头重要性最小化每层层内重要性变异系数的平方借鉴 MoE 负载均衡。从零训练直接全头均衡微调时把 top-m 重要头保留为共享头只均衡剩余路由头。训练开销大吗不大。利用 Flash Attention 前向已有的 LSE 值计算门控无需显式物化注意力权重额外训练延迟小于 2%。为什么消除 Sink 能提升性能Vanilla 中首 token 的键向量约束了查询向量的分布双极结构softmax 区分 token 的精度受限Sink/Gated 解绑几何约束后查询-键匹配更灵活长上下文检索更精准。参考链接arXiv 论文页https://arxiv.org/abs/2602.01203视频讲解B 站https://www.bilibili.com/video/BV1HPMS6AEuSStreamingLLMSink 现象发现https://arxiv.org/abs/2309.17453GPT-OSSSink Attentionhttps://arxiv.org/abs/2506.01516Gated Attention / Qwen3-Nexthttps://arxiv.org/abs/2510.13232DuoAttention固定重要头观测https://arxiv.org/abs/2410.10819DeepSeek-MoE共享/路由专家设计https://arxiv.org/abs/2401.06066给大家推荐一款自用写文献综述、无虚构文献的 AI复旦大学 FudanNLP 团队自研 切问学术官网qiewenpaper.com覆盖3.6 亿篇可溯源真实中英文文献能自动整合文献观点生成规范综述还能挖掘研究创新点、复现实验配合视频教学新手快速上手文献综述写作后记博客的关键词集中在编程、算法、机器人、人工智能、数学等等持续高质量输出中。讨论QQ群白拾的小屋 (750365700)⭐B站账号白拾的物理AI组会活跃于知识区和动画区✨GitHub主页YhbCode000工程文件
