返回顶部
热门问答 更多热门问答
技术文章 更多技术文章

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

[复制链接]
admin 显示全部楼层 发表于 2025-12-8 18:24:51 |阅读模式 打印 上一主题 下一主题
大语言模型(LLMs)在处理长上下文时面临全注意力机制带来的二次方计算复杂度瓶颈,限制了其推理效率。稀疏注意力通过限制每个查询仅关注部分历史 token 来缓解这一问题,但无需训练的稀疏策略往往导致性能显著下降。尽管原生稀疏注意力方法(如 NSA、MoBA)通过端到端训练改善了这一状况,却陷入一个关键悖论:其学习到的注意力模式反而比全注意力模型更稠密,削弱了稀疏化的有效性。这源于梯度更新缺陷——未被稀疏机制选中的键值对在前向传播中被跳过,未获得梯度更新,导致无法学会自我抑制。
+ k. u, {9 v, v) k

+ c+ s: [, |/ ]; G1 F
为此,腾讯优图联合伦敦国王学院提出 SSA(Sparse Sparse Attention)训练框架,在每一层同时引入稀疏与全注意力,并强制二者输出双向对齐。该设计保留了所有 token 的梯度流,使模型能主动学习有效稀疏,而非被动剪枝。实验表明,SSA 在多个常识推理基准上达到了稀疏与全注意力推理下的SOTA水平;其模型还能平滑适应不同稀疏预算——随着允许关注的 token 数增加,性能持续提升,支持灵活的计算-性能权衡。尤为突出的是,SSA 在长上下文外推任务中表现最强,通过缓解“汇聚区”(sink areas)中注意力值的过度分配,显著提升了模型对超长序列的泛化能力。
3 `2 X" t5 c/ a) `2 y

7 n5 e0 k, W# |, G8 V1 K5 Q% U

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

; n' S9 I8 S2 V3 s7 {4 T$ z0 |

% a$ I. R% H, U/ W1 {# G1 N& F
论文标题:SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space
论文链接:
https://huggingface.co/papers/2511.20102

# _, r# s5 d. I, k* F0 j% G4 g$ {
1.方法
SSA 设计了两个优化的目标:
标准的下一词预测交叉熵损失,在稀疏与全注意力模式下以相等概率采样计算。
层级别的双向对齐损失,用以约束稀疏注意力与全注意力输出的一致性(详见算法1)。
/ W+ a6 n& h+ c" H/ E

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强
; y& _+ Q4 B) C7 l0 i
其中,Lmode表示在随机采样的注意力模式(全注意力或稀疏注意力)下计算的交叉熵损失,α 为权重系数,Lalignment为双向对齐损失,旨在促进全注意力与稀疏注意力输出的一致性。

: @* K8 |! w2 O5 y- Y+ R/ X

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

8 P# n! I, O! M3 P& |
(1)稀疏与全注意力模式
在训练过程中,以相等概率交替采用全注意力与稀疏注意力模式(如图2所示)。引入双模训练有两个原因:一方面,全注意力能自然形成更具区分度、天然自稀疏的注意力分布;另一方面,稀疏注意力更贴近实际推理时的运行方式。为控制计算开销,并确保模型在训练中处理的 token 总量与基线方法一致,研究者并未同时优化两种模式的损失,而是交替进行更新。

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强
1 o6 N! I# F" g; ?8 Q
(2)对偶注意力对齐机制
$ Y0 l: d& D7 Z, o
为进一步提升注意力稀疏性并增强两种注意力模式之间的一致性,研究者引入了一种对偶注意力对齐机制。在每一层中,除当前主干路径所采用的注意力模式外,额外计算其对应相反模式的辅助注意力输出(例如,若当前流使用全注意力,则同时计算稀疏注意力输出)。该辅助计算仅用于对齐目标,不参与后续层的前向传播。
对齐目标由两个互补的组件组成。第一个是稀疏性损失,它旨在促进全注意力输出模仿稀疏注意力输出,从而促进形成更稀疏和更具选择性的注意力分布:
$ a3 f# W3 |; N# z' X2 B

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

7 V" Y% ~! Y! ~, S! c$ _& Q+ j
其中sg[·]表示梯度截断,afull和asparse分别指全注意力和稀疏注意力输出。
第二个组件为对齐损失,用于对稀疏注意力输出施加正则化,使其与全注意力输出保持一致。
% u3 s: v" T. n5 O+ e8 a

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强
3 Z3 U9 m% W1 Q, N
总对齐损失结合了两个分量:

9 P* W- C% X5 ^; j% j* H

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强
* Y$ }( k/ x! {3 I* L
这种双向对齐机制协同作用,一方面促使全注意力在训练过程中自然趋向更稀疏的分布,另一方面确保稀疏注意力分支在训练中保持稳定,并与其对应的全注意力分支保持一致。从概念上看,该损失以基于值向量(value-aware)的方式对齐两种注意力分布。相较于直接对齐全注意力分布,该方法显著提升了效率:后者需要显式构建稠密的注意力矩阵,不仅与 FlashAttention 等基于在线 softmax 的高效实现不兼容,还会造成较大的内存开销和计算负担。

- X2 U) w+ x/ N& V; i
2.评估
(1)语言建模能力
表1 在全注意力和稀疏注意力推理下,不同训练方法的比较

( ?* g1 }; ~! \* Z  j" @, I% Y

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

0 a) }9 I. Q1 u, f9 c
表6 SSA、MoBA 和 FullAttn 在 KL 散度、注意力稀疏性、困惑度及基准任务准确率方面的对比

1 U! i9 B8 C" q1 g) b: r" E  i

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强
7 W% x6 S* q1 m2 f& Z+ C
SSA 通过引入稀疏训练路径和对齐损失,在保持全注意力性能的同时,显著提升了稀疏推理质量。其核心机制是对齐损失促使全注意力分布变得更稀疏,从而减少与稀疏注意力在推理测试时的表达差异。如表1和表6所示,SSA 的稀疏注意力和全注意力性能差距最小(体现在PPL 和 KL 散度两个指标),验证了“增强内在稀疏性可提升稀疏推理效果”这一假设。
! e+ U: E5 q6 `  p; }8 T
(2)常识推理
如表1所示,在PIQA 、Hellaswag 、ARC‑Easy以及ARC‑Challenge这些常识推理任务中,SSA 不仅优于所有稀疏基线,甚至以仅 256 Token的感受视野(receptive field)超越了全注意力模型。两者在全注意力下的语言建模能力(PPL)相当,但 SSA 的下游任务性能显著更高。
由于二者唯一区别在于 SSA 具有更稀疏的注意力分布,性能提升最可能源于这种内在稀疏性。消融实验进一步验证了这一点:移除对齐损失后,推理性能下降。
表3 消融实验。Train A×B 表示训练时采用感受视野大小为 A、块大小(bl为 B 的配置;FullRatio 指图2中全注意力流(Full Attention Stream)的采样比例。此外,Only Full→Sparse 表示仅将对齐约束从全注意力单向施加至稀疏注意力路径,而 Only Sparse→Full 则表示对齐方向相反,仅从稀疏注意力向全注意力施加约束。
( z& W2 s$ z3 H

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

* L. P* `! F0 N7 p, S
(3)不同稀疏程度下的外推性能
SSA 在不同稀疏程度下展现出良好的外推能力:随着稀疏注意力中 token 数量增加,其在四项任务上的性能基本呈单调提升。

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强
* |+ Q, r6 U8 i5 v: ?9 l3 \
图3 性能与感受视野大小的关系

; U. g8 U& G6 g
(4)长上下文评估
表2 在多种上下文长度下的评估结果
7 @: c% z) {4 m; x4 h9 H

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强

媲美DeepSeek NSA!腾讯优图提出混合注意力机制SSA,长上下文外推更强
  u9 F$ k% W- ^+ w# c1 L$ y
图4 (a) 不同上下文长度下的困惑度。(b) 在 SSA 中提高稀疏注意力训练的比例,可提升模型在长上下文上的外推能力。(c)相较于 MoBA 和 FullAttn,SSA 在局部位置分配了更高的 logits 权重。(d) FullAttn 将大量注意力质量分配给了 8K 以外的 token。其中,图 (a) 使用 1B 参数模型,图 (b–d) 使用 300M 参数模型
如表2所示,在大海捞针任务(Needle-in-a-Haystack)中,SSA 在几乎所有感受视野下(除 1024 外)均为最强的稀疏注意力方法,并在全注意力推理下达到 100% 准确率。
当上下文长度超过训练最大长度(8K)时,FullAttn 性能骤降至 0%,而经稀疏注意力训练的模型仍保持非零检索能力,甚至在全注意力模式下也能恢复可观性能。
在困惑度方面,FullAttn 和 MoBA 在上下文超出预训练窗口后均出现 PPL 剧增;相比之下,SSA 与 NSA 在长达 32K 的上下文中仍保持低且稳定的 PPL。尽管 NSA 的 PPL 略优,但其架构更复杂,且无法外推至全注意力推理;而SSA 在全注意力评估下依然稳定,体现出更强的简洁性与鲁棒性。
在更全面的长上下文理解基准 LongBench 上,SSA 在所有推理模式下均取得最佳结果,进一步验证了其综合优势。
: A; B" T) v, b( k- C
回复

使用道具 举报

您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

AiGoe.com 成立于2014年8月,是目前国内优秀的开源技术社区,拥有超过300万会员,形成了由开源软件库、代码分享、资讯、协作翻译、讨论区和博客等几大频道内容,为IT开发者提供了一个发现、使用、并交流开源技术的平台。
  • 官方手机版

  • 微信公众号

  • 商务合作

  • Powered by Discuz! X3.5 | Copyright © 2001-2020, Tencent Cloud. |
  • | 营业执照 | |星点互联科技有限公司|鲁ICP备19001237号-21|鲁公网安备 4236902302000354号