[BidClub_]
张小珺Jùn|商业访谈录 · · 156 min

逐篇讲解DeepSeek、Kimi、MiniMax注意力机制新论文(含投屏版)——“硬件上的暴力美学”

张小珺松琳

YouTube
TL;DR
  • 三家公司押注的是同一个成本瓶颈:长思维链与长文档正把 attention 的 O(n²) 计算和线性增长的 KV cache 推到 test-time scaling 的中心。DeepSeek、Kimi 选择 Dynamic Sparse Attention,MiniMax-01 则以 7 层 Lightning Attention 配 1 层 Softmax attention;路线不同,目标都是让 prefill、训练与 decoding 更便宜。若性能能与传统 attention 持平,松琳认为高效架构省下的 inference 预算就可以转化为更长的推理时间。
  • DeepSeek 的 Native Sparse Attention(NSA)是三篇工作里最强的技术信号,因为它不只近似 full attention,而是从头预训练后在 loss、下游长文本与 reasoning 评测上都可能更好。其 Triton kernel 在 64K 序列附近相对 Triton 版 full-attention kernel 接近 10 倍加速,短至 8K、16K 已出现明显优势;松琳称其“既要 efficiency,又要 performance”,并给出“八成是会为之后的 V4 来服务”的主观判断。
  • NSA 的护城河不是三个分支本身,而是把算法、GQA 和 GPU 数据流锁成了一套“硬件上的暴力美学”。它用 compressed、selected、sliding-window 三路 attention,并强制同一 GQA group 的 heads 选择相同 KV blocks;再以 head 维度凑 Tensor Core 擅长的矩阵乘法,减少 KV cache 重复读取。代价是结构不够“洁癖”,但 DeepSeek 用增加 head 数量、保留多个 query groups 和额外分支补回表达能力,“像刀尖舔血一样”守住效率原则。
  • Kimi 的 MoBA 走向相反:以 mean pooling 表示 block,只留下 selected-attention 主干,依靠“SGD 的力量”完成路由,换来无新增参数和 full/sparse attention 自由切换。极简的成本也很具体:独立 head 选择引入 indexing/reindexing,因而采用 512-token block、Top 3,而 NSA 是 64-token block、Top 16;MoBA 在 128K 时仍没有明显速度优势,SFT 时还需把最后 3 层切回 full attention 解决梯度稀疏。不过它测到 100 万 token,并据称已经进入 Kimi 产品,属于真正“用脚投票”。
  • MiniMax-01 的核心资产是把 linear attention 首次放大到足以宣称 GPT-4o 级表现的规模,而不是发明 hybrid 本身。80 层模型按“7 层线性、1 层 Softmax”重复 10 次:RNN 式固定状态负责压缩历史、把每步 decoding 做成常数开销,少量 Softmax 层保留完整 KV cache 以补足 retrieval。松琳用一句话概括互补性:“attention 像翻书,RNN 像人的大脑。”
  • 从风险收益看,动态稀疏注意力在松琳看来其实比纯 linear attention 更保守:它保留全部 KV cache,只省计算而不省显存,因此大海捞针能力的下限更容易判断。松琳仍把 MiniMax 评为三者中承担架构风险最多的一家,但也强调 10%—20% Softmax 层的 hybrid 配比此前已在多个较小模型和相关工作中得到验证;三家公司结果无法横向排位,因为训练数据、规模和 infra 均不一致。
  • 资本与人才层面的主线是,数据堆叠接近瓶颈后,架构、开源论文和系统工程正在重新成为可见的公司竞争变量。松琳将 NSA、MoBA、MiniMax-01 排为一、二、三,真正的 aha moment 是 NSA 对 full attention 的“全线压制”;他还认为追赶者开源最能“秀肌肉”、建立投资者信心,并以千问开源与阿里股价表现作例。若国内模型技术继续追赶甚至超越,国内团队对海外留学生的吸引力也会提升,“人才都是会往人才密度高的地方走”。
Digest · the substance, structured for research

1. 长思维链把注意力效率推成下一代模型的前置条件

  • 松琳把论文集中出现的时间点放在 2025 年初理解:DeepSeek 的相关 reasoning 工作与 Kimi 1.5 都依赖更长的 chain of thought 解锁 reasoning,而输出越长,普通 self-attention 的 decoding 成本越难承受。“test-time scaling”因此不只是能力路线,也是架构效率问题。

  • Full attention 要建模所有 token pair,训练与 prefill 的计算复杂度随长度平方增长;autoregressive decoding 又要不断读取随上下文线性增长的 KV cache。长文档和长推理分别放大了这两个瓶颈。

  • 张小珺追问是否在为新模型铺路时,松琳给出明确但保留条件的判断:“DeepSeek 这篇文章有八成是会为之后的 V4 来服务的。”Kimi 既有 Kimi 1.5 的 reasoning 需求,又以长文本产品见长,因此也同时需要降低 prefill 和 decoding 开销。

2. Attention 的本质是为每一对 token 建模交互

  • 松琳从 query、key、value 解释标准 attention:每个 token 用 query 与此前 tokens 的 keys 做内积,得到 attention scores;经 Softmax 归一化后,再以这些权重对 values 做线性组合,形成当前位置的 output。

  • 自回归模型还需 causal mask,防止当前 token 看到未来信息。最简洁的理解是:“attention 就是一个 pairwise 的建模过程”,它刻画序列里任意两个位置的交互程度。

  • Transformer 取代 LSTM 的关键不只在表达能力,也在训练并行度:这种结构能高效使用 GPU,容易堆参数、扩大模型,随后 BERT、GPT 以及主流大语言模型都沿用了这一基础。

3. Dynamic sparsity 让每个 query 自己决定该看哪里

  • BigBird 一类 static sparse attention 预先规定跳跃、窗口等固定 pattern;Dynamic Sparse Attention 则由当前 query 动态选择需要访问的 keys 和 values。同一序列中的不同 tokens,因此可以读取完全不同的历史区域。

  • Kimi 与 DeepSeek 的共同点,是试图把这种 query-aware sparsity 从 inference 加速工具变成可从头预训练的架构。此前的困难并非想法不存在,而是动态、离散的数据访问与现代 GPU 的连续块读取方式“不对齐”。

  • 两篇工作的共同前身是 Quest:它把连续 KV tokens 划成 blocks,为每个 block 建立表示,再让 query 选择最相关的 blocks。连续 block 可以整块读取,避免逐 token 的随机访存;DeepSeek NSA 与 Kimi MoBA 都可看作沿这一框架向预训练扩展。

4. Native 的含义是动态稀疏终于能在硬件上原生训练

  • 松琳把 Native Sparse Attention 称为开创性工作,核心不是第一次提出稀疏或 block selection,而是第一次把大规模 Dynamic Sparse Attention 做到“natively trainable fast attention”:训练、prefill 与 autoregressive decoding 都按硬件约束重新设计。

  • 论文的野心是“既快又好”:不仅降低训练的前向、反向成本和 inference 读取量,还让从头预训练的 NSA 在多个 benchmark 上压过 full attention。松琳反复强调,这使它不再只是对既有模型的有损加速插件。

  • 相比之下,不少既有方法主要是在已经训练好的 full-attention model 上近似其 attention distribution。近似越受 block 粒度和硬件约束,越容易掉点,也很难突破原模型的性能上限;但松琳也指出,Quest 这类模型可以做预训练,只是预训练效率可能不如 NSA。

  • 对 DeepSeek 连续开源项目的内容,松琳拒绝“兼职算命”,只推测可能包括软硬件协同 infra、MoE 和 NSA 相关模型;这部分被他明确标为预测,而不是论文已披露路线。

5. NSA 用三条分支同时保存粗粒度、精细选择与局部先验

  • 第一条 compressed-attention 分支先把每个连续 KV block 压缩成一个 key representation 和 value representation,再让 query 在 block level 上做粗粒度 attention。它既提供全局概览,也产生一个直接进入最终结果的 output。

  • 第二条 selected-attention 分支依据 block scores 取 Top-k,再展开被选 blocks 内的全部 tokens,做更细粒度的 token-level attention。稀疏性正来自这里:不是对所有历史位置做精细计算,只在最相关的少数 blocks 内计算。

  • 第三条 sliding-window 分支直接关注邻近 tokens,利用语言建模中“附近位置通常更重要”的 inductive bias。松琳提到,训练好的 attention map 常把大量权重集中到局部窗口及开头 tokens,后者常被称为 attention sink。

6. 三路门控把全局与局部的配比交给模型自己学习

  • 每条分支的 output 都经过一个由当前 token representation 决定的 gate,随后三者加权相加。张小珺问三路如何协同时,松琳的回答是“不妨把这个选择直接交给模型”,因为人工规定每一路比例反而难以适配不同 query。

  • 这种设计没有宣称 compressed、selected 或 sliding window 各自是新模块;它的价值在于三者组成了可训练、可回传且适合硬件执行的闭环。松琳因此提醒,不应只看架构图便把论文归为“缝合怪”。

7. MHA、MQA 与 GQA 决定了稀疏选择到底能省多少读取

  • Multi-head attention(MHA)为每个 head 保留独立的 query、key、value;Multi-query attention(MQA)让所有 query heads 共享一份 K、V,大幅缩小 inference 时的 KV cache,却可能因共享过度造成 performance degradation。

  • Grouped-query attention(GQA)取折中方案:同一 group 内的 heads 共享 K、V,不同 groups 仍保留不同表示。松琳用“每两个 heads 共享一份 KV”的示意说明,它用温和共享换取 cache 缩减和表达能力之间的平衡。

  • 在 MHA 中,每个 head 本来就必须读取自己的 KV,因此各自选择不同 blocks 不额外扩大读取量;但在 GQA 中,共享 KV 的 heads 若选择不同 blocks,硬件最终必须加载这些集合的并集,原有共享优势便被稀释。

  • 这构成 NSA 与 Quest、MoBA 的关键分岔:后两者保留每个 head 独立选择的自由,NSA 则要求同一 GQA group 内的 heads 读取相同 KV subset,以换取显著更规整的数据流。

8. NSA 主动牺牲 head 级自由度,换取 decoding 的真实吞吐

  • NSA 在 block-selection scores 上对同一 group 的 heads 求和,再统一做 Top-k,从机制上保证它们选中同一组 KV blocks。这样共享 KV 只需加载一次,而不是为各 heads 的不同选择反复读取。

  • 松琳承认这种绑定可能减少 diversity,也正是很多架构研究者会担心“掉点”的地方。DeepSeek 的回应很激进:既然共享对硬件最有效,就增加 heads、保留多个 query groups,再用 compressed、sliding-window 分支补回表达空间。

  • 论文配置中有 4 个 groups、总计 64 个 heads,query head dimension 约 192;相乘得到接近 1.2 万维的投影,远高于约 2560 的 hidden dimension。松琳把这种大幅 up projection 与 DeepSeek V3 的 MLA 风格联系起来:“我反正是从头开始训练嘛”,只要训练和 inference 都高效,结构常规与否不是首要约束。

9. FlashAttention 的成功说明瓶颈不只在 FLOPs,更在 HBM 与 SRAM 之间

  • 松琳回顾 FlashAttention 的 blocking:把一块 queries 和一块 keys 放入 SRAM,用矩阵乘法计算局部 attention block,再读入对应 values;中间 attention scores 不写回 HBM,从而减少高成本的数据搬运。

  • SRAM 容量小但带宽高,HBM 是例如 A100 的 80GB 全局显存,容量大却相对慢。硬件高效算法的共同目标,是尽可能让计算停留在 SRAM,并减少两级 memory hierarchy 间的通信。

  • Full attention 对一个 query block 反正要扫描全部历史 KV blocks,因此连续装载没有浪费;dynamic sparse attention 中,每个 query 选择不同 blocks,若仍把连续 queries 捆在一起,便要加载它们需求集合的并集,最坏情况下几乎重新退化为全量读取。

10. NSA 用 head 维度凑矩阵乘法,绕开稀疏 query 的并集浪费

  • 为避免 query-block 并集,NSA 让每个 token 独立决定 output;但单 token 计算又缺少 Tensor Core 喜欢的矩阵形状。解决办法是利用 GQA 下多份 queries 共享同一 KV selection,把 heads 当作 batch-like 维度,一次凑成 matrix multiplication。

  • 松琳特别欣赏这一 kernel:普通 FlashAttention 用 query block size 凑矩阵维度,NSA 改用 number of heads。“基本上没有什么浪费”,既只读所需 KV,又调用了 Tensor Core,是他第一次看到这种玩法被用于训练阶段。

  • Triton 对相关矩阵尺寸通常要求至少 16,因此每个 group 需要至少 16 个 heads;同时模型又希望保留约 4 个 groups 的选择差异,最终推动了 64-head 设计。看似异常的架构参数,其实是从 kernel 约束反推出来的。

  • 以 A100 为例,松琳称相同 FLOPs 下,半精度矩阵乘法调用 Tensor Core 可比一般 ALU 计算快约 16 倍;到 H100,矩阵计算继续加速,数据 loading 反而更容易成为瓶颈,这就是“memory wall”。

11. 从头训练是 sparse attention 突破 full-attention 上限的唯一出口

  • 对现成 full-attention model 做 sparse inference,本质是逼近既有 attention map。即使先计算完整 score 再做 exact Top-k,它也只是逼近方案的理论上限;Quest 还必须按 block 选取,效率更高却少了一层灵活性,因而通常进一步掉点。

  • 松琳的结论很直接:“如果我们想让 sparse attention 甚至比 full attention 要好,只有一个出路,那就是 train from scratch。”模型从预训练开始适应动态选择,便不必复刻 full attention 的分布,也可能学到更适合稀疏结构的表示。

  • NSA 与 Quest 都保留全部 KV cache,不像 H2O 一类机制直接丢弃 cache。松琳认为“不丢”很重要,因为眼下不相关的位置可能被后续 query 再次调用;动态选择可以暂时跳过信息,却不必永久删除它。

12. 稀疏性可能不是性能税,反而像一种注意力去冗余

  • 松琳提供的直觉解释是,训练好的 full attention 本就“highly sparse”,大量 token pair 的权重近乎冗余。直接以 sparse attention 训练,可能迫使模型更集中于真正相关的位置,减少被无关 tokens “distract”的机会。

  • 这仍是解释而非论文已证明的因果机制;他用“我一个直觉的理解”保留了不确定性。确定的实验现象是,NSA 的训练 loss 基本全程略低于 full attention,多个下游 benchmark 也呈现优势。

  • Needle-in-a-haystack 图几乎“绿油油一片”,说明保留全部 KV cache 后,模型仍能从长上下文找回目标信息。更令松琳意外的是,提升不止出现在长文本,还延伸到 chain-of-thought reasoning 与 R1 风格相关评测。

  • 这也是他全场最大的 aha moment:Dynamic Sparse Attention 不只少算,还可能比 full attention 更强。“它既要 efficiency,又要 performance,这个太厉害了。”

13. NSA 的速度优势在较短上下文已出现,并随长度迅速扩大

  • 论文用 Triton 版本 NSA kernel 对比 Triton 版本 FlashAttention,以减少 CUDA 与 Triton 实现差异造成的不公平。松琳指出,NSA 在较短序列已经更快,到 64K 附近加速接近 10 倍,长度越长,加速比越明显。

  • Decoding 图给出 expected speedup,是因为 autoregressive decoding 通常 memory-bound,效率大体取决于每步读取多少 KV cache。NSA 让 GQA group 共享选择,因此相对 full attention、Quest 和 MoBA 都有天然读取优势。

  • 松琳最看重的不是某一个峰值数字,而是优势同时覆盖训练前向、反向和 decoding:“既快又好”意味着同一架构可同时减少预训练成本和上线 serving 成本,而不是在两阶段之间转移负担。

14. 可回传的 compressed branch 是 NSA 选对 blocks 的训练信号

  • Top-k selection 本身不可导;若架构只有 selected-attention 路径,错误 block selection 很难获得直接梯度纠正。NSA 的 compressed-attention output 参与最终输出,使 block representations 和 scoring 路径可以通过端到端 loss 反向传播。

  • 论文把这一方案与 Quest 等其他 block-selection 方法做 ablation,替代方案表现较差。松琳据此认为 compressed branch 不是装饰,而是让 selection 更 informative、让第一条分支“更加 make sense”的关键。

  • 这也解释了 NSA 为什么没有 MoBA 在 SFT 阶段那么突出的梯度稀疏问题:即使精细 selected path 没覆盖某些 prompt blocks,compressed path 仍把训练信号传回更广泛的历史位置。

15. NSA 的主要创新是系统协同,而不是发明三个新模块

  • 张小珺问它属于“小创新”还是“大创新”时,松琳承认 sliding window、block selection、compression 都能在已有 literature 找到;若只逐模块查 novelty,确实容易得出“没有什么了不起”的判断。

  • 他的反驳是,论文真正的 significance 在于把 query-aware dynamic sparsity 改造成训练与 inference 都 hardware-friendly 的整体。Quest 这类方法也能做预训练,但效率可能不如 NSA;NSA 的 kernel 与 GQA 约束使其更适合大规模使用。

  • 松琳把它归为工程创新更多,但并未用“工程”贬低其价值:“你只有硬件高效,才能降本。”在大模型尺度下,能否连续读取、能否写成矩阵乘法、能否避免无效 cache traffic,本身就是架构可行性的组成部分。

  • 他用“硬件上的暴力美学”和“像刀尖舔血一样”概括这种风格:在 Quest 的大原则——每个 query 动态选 blocks——之内,NSA 找到了一个非常接近当前硬件约束下最优解的处理方式。至于 OpenAI 如何实现类似能力,他以“Closed AI”调侃其技术披露太少,只能说 GPT-3 报告曾提 sparse attention,现有模型无法考证。

16. MoBA 把同一问题压缩成一个近乎无参数的极简方案

  • Kimi 的 MoBA 同样沿 Quest 的 block-selection 框架前进,但每个 key block 只做 mean pooling,query 与这个均值计算 score。它不增加用于 block representation 的 MLP,也没有额外可学习参数。

  • 相比 NSA 的三路输出,MoBA 砍掉 compressed-attention output 和 sliding-window output,只留下 selected attention。其团队表达的是“相信 SGD 的力量”:即使没有额外分支和 loss,优化也能逐渐找到合适 blocks。

  • 每个 attention head 仍可独立选择 KV blocks,因此模型自由度更高,架构图也更“优雅”;但它没有 NSA 的 group-shared selection,无法直接复用后者几乎无浪费的训练与 decoding kernel。

17. MoBA 的简约把复杂度转移到了 indexing 与 block 粒度

  • MoBA 的实现按每个 KV block 找出所有选择它的 query tokens;这些 queries 在原序列中通常不连续,于是先通过 indexing、reindexing 收集成连续 tensor,再调用 FlashAttention kernel。这一步包含大量 indexing 和 reindexing operation。

  • 这一步不是 free:如果 KV block 数量足够多,索引开销可能成为 bottleneck。因此 MoBA 采用 512-token block、Top 3,NSA 则采用 64-token block、Top 16;两者最终读取 token 数量接近,但选择粒度差异很大。

  • Top 3 个大 block 更容易漏掉分散的关键信息,Top 16 个小 block 的容错和细粒度更高。松琳把这视为 MoBA “极简背后付出的代价”,而非单纯的超参数偏好。

  • 速度曲线也体现了这种差别:MoBA 到 128K 左右仍未显示特别明显的速度优势,NSA 在 8K、16K 已有可见加速。松琳因此判断,MoBA 可能更适合极长 prefill,而常见约 8K 的预训练长度未必能获得同等收益。

18. SFT 的稀疏梯度迫使 MoBA 暂时切回 full attention

  • MoBA 的训练 loss 在早期弱于 full attention,随后才逐渐缩小 gap,并没有 NSA 那种稳定压低 full-attention curve 的现象。松琳保留判断:“我不确定是不是 compressed output 被砍掉带来的负面作用。”

  • SFT 常有很长 prompt,但 loss 只计算回答部分;若少数带 loss 的 tokens 没选择某些 prompt blocks,那些 blocks 就收不到梯度。MoBA 报告将此描述为 sub-optimal SFT performance,本质是 sparse routing 叠加 sparse supervision。

  • 解决方案是把最后 3 层切换为 global/full attention,让上层对全部 tokens 产生稠密梯度。由于 MoBA 没新增参数,full 与 sparse attention 可自由切换,这个原本追求极简的设计反过来提供了修复空间。

  • 松琳一度把 90% sparse、10% full 说成训练流程,随后当场更正为 ablation;该实验显示 hybrid 训练可提升表现,但不能据此断言正式模型必然按这一比例训练。这段自我修正也保留了他解读论文时的不确定性。

19. MoBA 的产品落地证明极简路线不是纸面实验

  • Kimi 将评测长度推进到约 100 万 token,长文本结果整体良好;据其公开说法,MoBA 已投入 Kimi 产品。松琳据此认为它确实有效,说明方案至少已经进入真实产品场景。

  • NSA 与 MoBA 的设计哲学因此很清楚:DeepSeek 为硬件效率接受复杂架构与强制共享,Kimi 则尽量不引入模块和参数,再用索引工程与 attention 切换解决问题。“一个是硬件的暴力美学,另外一个是极简的优雅算法。”

  • 两者优化目标并无根本差别,都希望 prefill 和 decoding 越快越好;区别在于 NSA 更早进入加速区间,MoBA 保留更多 head-level freedom,并能无缝切回 full attention。谁更好仍需在统一训练条件下验证。

20. Linear attention 把完整历史压缩成固定大小的矩阵状态

  • 松琳进入自己的研究主场后,从去掉 Softmax 开始推导 linear attention。利用线性运算的结合律,可先累计 keys 与 values 的外积,再让当前 query 读取;历史被编码为一个固定大小的 d×d matrix-valued hidden state。

  • 与 LSTM 的 vector hidden state 相比,这个矩阵状态完成了 state expansion,容量显著更大。每个 decoding step 只需更新并读取固定状态,因此空间复杂度和单步时间复杂度都是常数,总推理复杂度随序列长度线性增长。

  • 代价是它不再保存所有历史 KV cache:全部过去被压进一个有限状态。松琳指出,hidden-state size 与 RNN 的记忆容量直接相关,这既迫使模型学习压缩,也决定了它不可能无损记住任意长历史。

21. Chunkwise 算法让 linear attention 从理论高效变成 GPU 高效

  • 直接用 recurrence 训练会像传统 LSTM 一样逐 token 串行,而且外积与 matrix-vector multiplication 难以充分调用 Tensor Core;若改写成全并行 attention form,又会回到随长度平方增长的训练成本。

  • Chunkwise 算法把序列切成 chunks,只计算每个 chunk 的末端 hidden state;块内用并行 attention 计算 local contribution,块外用同一个历史 state 批量服务整个 chunk 的 queries。相关计算可以利用矩阵乘法,并且是数学等价变换,不是 approximation。

  • 这里的 chunk dimension 再次扮演 batch-like 维度,与 NSA 用 heads 凑矩阵乘法遵循同一硬件原则。选定固定 chunk size 后,训练复杂度可保持 sub-quadratic,同时保留 GPU 并行度。

  • 松琳给出的典型场景是 video generation:几分钟视频转成序列后可轻易达到百万级 tokens,full attention 很快吃不消;chunkwise linear attention 因而可能具有先天优势。

22. 现代 linear attention 已用衰减和选择性补上早期模型的短板

  • 早期 linear attention 的 language-modeling performance 较差,形成了“linear attention 不 work”的行业印象;松琳认为过去两年无论模型效果还是 kernel efficiency 都进步很快,需要更新 prior knowledge。

  • RetNet 是较经典的相关工作;固定遗忘率的机制会使不同位置缺少可选择性。Mamba 2 把 decay 变成 data-dependent,让每个位置决定自己的 decay。松琳的评价是,Mamba 强调的 selectivity 与 LSTM gating “本质 idea 是一样的”,像“新瓶装旧酒”。

  • 从 linear-attention 视角,Mamba 2 的 State Space Duality、block decomposition 和相关训练形式都能得到更直观的等价解释。GLA、Lightning Attention、xLSTM 等现代模型也可放入 attention-with-decay 或更一般的状态更新框架。

23. MiniMax-01 用七比一的混合比例补齐压缩记忆与精确检索

  • MiniMax-01 每 8 层安排 7 层 Lightning Attention、1 层 Softmax attention,整个模块重复 10 次,共 80 层。Linear layers 负责绝大部分顺序建模,少量 full-attention layers 保留访问历史 tokens 的能力。

  • 松琳把两者的互补性讲得很形象:“attention 有点像翻书,RNN 有点像人的大脑。”RNN 的固定容量强迫模型抽取可压缩规律,Softmax attention 则保存全部 KV cache,遇到 needle-in-a-haystack 等 retrieval task 时可以直接查找原始信息。

  • 纯 linear attention 在 retrieval 上可能较弱;hybrid 通过少数 attention layers 修补这一点。RNN 还自带顺序信息,部分 hybrid 架构可以少依赖甚至不用 RoPE,从而避开某些位置编码外推问题。

  • MiniMax 报告称 hybrid 架构的 scaling behavior 优于纯 Softmax attention,并将模型放大到宣称 GPT-4o level performance 的尺寸。令松琳惊艳的是 scale,而不是 hybrid 概念本身:“第一次把这个 hybrid 架构 scale up 到这个尺寸。”

24. 架构创新、开源论文和系统能力正成为同一场公司竞争

  • MiniMax-01 的放大依赖扎实 infra:专家并行、pipeline 并行,以及可能采用的 linear-attention sequence parallelism 等工程优化缺一不可。此前已有多个相关模型验证 hybrid,attention 层占约 10%—20% 也在消融研究中被考察;MiniMax 的 1/8 配比与这一范围相近。

  • 三条路线无法只看最终 benchmark 排名,因为模型规模、数据和训练流程不同。风险判断上,松琳认为 sparse attention 保留全部 KV cache,“大海捞针肯定能做”,下限相对稳;MiniMax 承担的架构风险可能最大,但先前较小规模的相关结果也让 7∶1 hybrid 并非盲赌。

  • 松琳最终把 NSA 排第一、MoBA 第二、MiniMax-01 第三;即使他本人研究 linear attention,hybrid 能 scale 在他看来仍不意外。真正改变其研究兴趣的是 NSA 比 full attention 更好,他表示之后可能会做一些 Dynamic Sparse Attention 方面的研究。

  • 更大的背景是 pretraining data scaling 可能接近瓶颈,行业转向 test-time scaling,也重新寻找底层架构收益。松琳称 DeepSeek 坚持架构创新“非常可贵”,因为这件事“风险与机遇并存”;未来可继续改 RoPE、尝试 contextualized position encoding、用 RNN 承担位置信息,乃至等待 Long Convolution “卷土重来”。

  • 论文与开源也因此成为商业信号。松琳说“天下苦闭源公司久矣”,追赶者通过开源最容易建立技术影响力,并可能支撑投资者信心;他以千问开源与阿里股价上涨作例,将二者联系起来。

  • 对人才流动,他的判断同样带条件:国内大模型过去落后时,留学生自然去硅谷学习;如果如今已追赶、部分技术甚至超越,国内团队对海外留学生的吸引力会提高。“人才都是会往人才密度高的地方走的。”

松琳

Thank you。Timmy 和 DeepSeek 它发的这两篇叫做动态稀疏注意力,Dynamic Sparse Attention,它们这个就是可以用来解决这种嗯长文本的开销的问题。MiniMax 采用的是另外一种方案,叫作混合模型,也就是线性注意力和 Softmax 注意力的结合。

我确实被 Native Sparse Attention 甚至比 full attention 更好这一点吸引住了,可能之后也会做一些 Dynamic Sparse Attention 方面的研究。DeepSeek 可能是比较少见的、非常坚持自己做架构创新的公司,这一点非常可贵,因为很多公司都排斥架构创新,可能觉得堆数据就行了,为什么要冒险做架构创新?但架构创新本身就是风险与机遇并存的。

张小珺

Hello,大家好,欢迎收听《张小俊商业访谈录》,我是小俊,这是一档提供一手高密度信息的商业访谈节目。今天这期节目呢延续了我们的论文系列,我邀请MIT计算机科学与人工智能实验室的在读博士生松林来给大家解读上个星期DeepSeek和Kimmy刚刚发布的全新的技术报告。DeepSeek 和 Kimi 又一次在同一天发布论文,两篇论文都集中在改进注意力机制,以处理长文本任务。两篇论文的署名分别包括梁文锋和杨志林,松林会带领大家来读最新的这两篇论文,同时也会对比对照讲解春节前另外一家大模型公司MiniMax发布的新型注意力的模型。

更早之前,春节前,MiniMax 也发布了一篇注意力机制相关的论文。今天我们会阅读这三篇注意力机制相关的文章,解析不同大模型公司的技术哲学与路线选择,希望让更多人领略 AI 技术平权,体验技术之美。那二零二五我们和AI共同进步。

今天这期节目仍然由他来主讲论文,中间遇到我不懂的地方,我会作为观众来提问,可能在中间,也可能在最后。我们的主要目的是帮助更多人了解最前沿的 AI Researcher 世界到底在发生什么,一起共同学习,做你的“赛博组会”,也希望更多像我这样非技术背景的人能够感受到技术之美。

那开始之前,我想先简单问几个问题,帮大家做一个定位。总体来说,你觉得为什么 DeepSeek 和 Kimi 都选择在这个时间点发布注意力机制相关的论文?注意力机制的改进在 AGI 路线图上有多么重要?

松琳

大家好,我是杨松林,现在是 MIT 计算机科学系博士二年级,我的导师是 Yong King。目前我的研究方向是硬件高效的序列建模,主要集中在高效注意力,尤其是线性注意力。

1. Long Context Changes Everything

因为现在是 2025 年初,之前 DeepSeek 发了一篇阿旺的文章,Kimi 也发布了 Kimi 1.5。它们的技术核心都是长思维链,能够解锁很多 reasoning、推理方面的能力。但是长思维链的开销比较大,因为 Transformer 采用了 self-attention 机制,所以有平方的计算复杂度,以及线性的 KV cache 空间复杂度。

Kimi 和 DeepSeek 这两篇工作叫作动态稀疏注意力(Dynamic Sparse Attention),可以用来解决长文本的开销问题。MiniMax 采用的是另外一种方案,也就是混合模型。混合模型是线性注意力和 Softmax 注意力的结合,如果绝大多数层都换成线性注意力,整体开销,尤其是推理阶段解码的时间开销,就会缩短很多,可能更加适合 test-time scaling。

张小珺

DeepSeek、Kimi、MiniMax 这三篇注意力机制的改进策略,是共性更大,还是差异更大?

松琳

我觉得 Kimi 和 DeepSeek 这两篇论文的共性更多,它们都是动态稀疏注意力。所谓动态,就是它们稀疏的性质不是事先预定好的。之前很多工作叫作静态稀疏注意力(static sparse attention),它们的稀疏模式是事先给定的。

比如 Google 之前有一篇工作叫作 BigBird,它里面有一些固定模式,包括跳跃式的 sparse attention,或者其他事先设定好的 pattern。动态稀疏注意力则是由 query 动态决定稀疏性。对于每个 token 来说,因为它的 query 不同,所以它想要 attend 到的 key 和 value 也可能不同,这就是动态稀疏注意力的核心思想。

Kimi 和 DeepSeek 都想把动态稀疏注意力用在训练上,所以它们非常相似。MiniMax 则是另外一套思路,使用线性注意力。线性注意力可以写成 RNN 的形式,在推理时是常数空间复杂度,每一步也是常数时间复杂度,总体推理复杂度就是线性的。

张小珺

我理解我们今天的节目很大一部分都会围绕注意力机制展开。能不能先给大家介绍一下注意力机制是什么,以及为什么 attention 对大语言模型这么重要?

松琳

2. Attention From First Principles

注意力机制的起源有很多种说法。一般公认的是,之前有一篇工作在 LSTM 的 hidden state 上做了 attention,可能是 2015 年左右的论文,之后这个东西被称为注意力机制。

后来 Google 把它发扬光大,提出了非常知名的《Attention Is All You Need》。他们用这个架构完全取代了之前非常流行的 LSTM,因为 Transformer 的训练并行度很高,在 GPU 上非常高效,也很容易堆参数、扩大规模。后来的 BERT、GPT,以及现在的大模型,基本都是基于 Transformer 架构。

简单来说,对于每两个 token,我们都要建模它们之间的交互,并得到一个 score。每个 token 都要对其他 token 有一个 score,所以我们会计算每两个 token 之间的分数。这个 score 通常叫作 attention score。

在自回归建模中,我们一般会使用 causal mask,防止 token attend 到后面的 token,保证它只能 attend 到前面的 token。具体来说,每个 token 都会有 3 个 vector,分别叫作 query、key 和 value。

以自回归建模为例,每个 token 会用自己的 query 和之前所有 token 的 key 做内积,得到一个标量的 attention score,然后经过 Softmax 得到归一化的注意力分数。最后,再由这些注意力分数对之前位置的 value 向量做线性组合,得到 output,这就是注意力机制的核心思想。

张小珺

怎样才能更简单地理解 attention?

松琳

更简单地说,我觉得 attention 就是对每两个 token 之间的交互进行建模。它其实可以被理解为一种 pairwise 建模过程,用来建模序列中每两个位置之间的交互程度。

张小珺

你觉得今天讲的 DeepSeek、Kimi 和 MiniMax,它们对注意力机制的改进,是在为后续模型训练和产品做铺垫吗?比如 DeepSeek 提出的 NSA,也就是 Native Sparse Attention,是为了给 DeepSeek-V4 的长上下文预训练做铺垫吗?能不能从它们不同的注意力机制改进目标上,看出这几家公司后续策略上的差异?

松琳

3. Test Time Scaling Sets The Agenda

我觉得 DeepSeek 这篇论文有八成会为之后的 V4 服务。大家都知道,DeepSeek 最近出了一个阿旺的工作,这项工作主要是想实现 test-time scaling。因为预训练阶段的 scaling law 似乎已经快要失效了,所以需要把更多计算放到 inference time 上。

在 inference time,我们可能会采用长思维链,也就是 long chain of thought,把输出长度整体拉得很长,这就叫 test-time scaling,也就是推理时间扩展。Test-time scaling 主要是一个 decoding process,所以对 decoding 效率,也就是模型架构在解码阶段的效率,提出了很高的要求。

普通 self-attention 在 decoding 时比较慢,这就是为什么我觉得 DeepSeek 会花精力研究 sparse attention。Kimi 最近也发布了 Kimi 1.5,其中同样有一些 test-time scaling 的内容,思路和 DeepSeek 比较像。后来 OpenAI 也发文说,他们做 reasoning 模型的思路和 DeepSeek、Kimi 两家很相似。

另外,Kimi 还有一个主打功能。以前 Kimi 就是以长文本建模为卖点,用户可以输入一篇 document,让 Kimi 处理整个很长的 document。这涉及类似 pre-filling 的问题,也就是对给定的文本进行建模,有点像训练中的 forward、前向过程。

因为 attention 是平方复杂度,要建模 pairwise 交互,所以当 document 很长时,平方复杂度是无法承受的。这也是 Kimi 想研究 sparse attention 的原因,因为它可以大幅缩短 pre-filling 的时间。

张小珺

接下来连续 5 天,DeepSeek 要开源 5 个项目。能不能对这些新项目做一些预测?我们的节目可能刚好会在他们发布项目的第 1 天或第 2 天播出。

松琳

我觉得他们应该会发布一些 infrastructure 框架。我看到他们已经有一份 report,说是一个软硬件协同设计的框架,所以我觉得更多会是 infra 相关工作,比如可能有 MoE 之类的内容,还有前几天发布的 Native Sparse Attention 开源模型。

至于具体是什么,我也不太能预测,简直是在算命。

张小珺

接下来的时间就交给你,由你来带领大家学习最新的三篇论文。论文讲解时,有些专业名词我们会尽量解释,也请大家见谅,因为有些专业名词确实不好翻译成中文。

松琳

4. Inside Native Sparse Attention

我们先来看《Native Sparse Attention》这篇文章。Native 的意思是,它要实现一种能够在硬件上原生训练的 sparse attention。它可以说是目前第一个把动态稀疏注意力用于预训练的模型。

为什么之前的动态稀疏注意力不能做预训练?后面会讲到,它们在训练时会遇到各种困难,主要是因为它们和当前硬件并不对齐。现在很多 dynamic sparsity 的工作主要是用来加速 inference,而不是从头开始预训练。

DeepSeek 这篇文章是开创性的,它第一次大规模地使用 dynamic sparsity 来做预训练。它采用的机制主要和 Quest 比较相似。Quest 是 ICLR 2024 的一篇论文,由 MIT 韩松老师的团队提出。

Quest 的核心思想是,对于每个 token,它动态决定自己要注意到哪些 key 和 value block。因为每个 token 需要的信息不同,所以它要选择的 KV block 也可能不同。

Quest 具体的做法是,为了更好地适配硬件,对于每个 query,它会在前面寻找 key 和 value 的 block。这个 block 是连续的 chunk,里面包含很多连续的 token,这样硬件就可以连续读取,避免额外开销。

Kimi 和 DeepSeek 这两篇文章都可以看作 Quest 的 follow-up,把 Quest 的思路用于预训练。接下来我们会看到,这篇文章重点讨论了如何针对 Quest 这种 sparse attention 机制进行 hardware-aware design,也就是在算法设计层面让它更加硬件友好,从而提高训练效率。

这篇论文的卖点是,希望在 modern hardware 上直接进行 end-to-end 训练,并且训练效果很好。最后他们可能发现,从头开始预训练的 Native Sparse Attention,甚至比 full attention 还要好,达到了既快又好的效果,野心非常大。

我们先简单看一下最终效果。这里有 3 类 benchmark,包括长文本 benchmark、chain-of-thought reasoning,也就是类似 R1 的推理能力,以及其他能力。可以看到,NSA 在这些 benchmark 上几乎都压过了 full attention。

Full attention 是最传统的 attention 机制,直接建模每两个 token 之间的交互,需要平方级别的时间复杂度。

张小珺

大部分大模型都是用 full attention,对吗?

松琳

对,基本上所有大语言模型都基于 full attention。比如 Qwen、Llama,都是 full attention 架构。不过值得一提的是,它们使用的是 full attention 的一些变种,一般是 GQA,也就是 Grouped-Query Attention,后面会稍微提到。

从右侧的 speed 对比来看,Native Sparse Attention 不仅在 decoding 时更快,在训练的前向和反向传播阶段也比 full attention 快很多。这篇论文的主要思想,就是让算法和硬件进行更多对齐,同时考虑 inference time,也考虑 training 的友好程度,既想加速 pre-training,也想加速 inference。

我们看图 2,这是它架构的核心思想。前面提到,每个 query,也就是每个查询 token,都要 attend 到不同的 key 和 value block,这就是 Quest 的思路。

首先,我们把 key 和 value 分成不同的连续 block。图中一共分成 4 个 block。对于每个 block,会进行一次 compression,也就是把一个 block 里面所有 token 压缩成一个机器学习表示。

之后,每个 chunk 都有一个表示,query 就可以和这些 chunk 的表示进行 attention,得到 attention score。这里会先把 key 压缩成一个表示,同理也会把 value 压缩成一个表示。

这样,我们会得到一个 attention score。给定一个 query,它对每个 chunk 都有一个 attention score。接着 apply causal mask,再对每个 chunk 的 value 做线性组合,得到一个 output。

这个 output 之后还会经过一个 gating,也就是门控。这里有 3 个分支,每个分支都会通过门控决定自己对最终 output 的贡献程度。门控机制是模型自己学习的,通常会把 input representation 经过一个 projection,得到一个 value,这个 value 由当前 token 的表示决定。

第 1 个分支是比较粗粒度的 attention,对每个 chunk 的 value 得到 chunk representation,可以理解成 chunk-level attention。完成粗粒度 attention 后,还要做更加细粒度的 attention。

如果所有细粒度 attention 都做,就会退化成平方复杂度的 full attention。所以 sparse 的来源就在这里:我们不对所有 token 做 attention,只选取 top-k。比如图中 k 等于 2,我们会从 4 个 block 中选出 attention 最高的 2 个 block。

图中选出的是第 2 个和第 4 个 block,然后把这两个 block 的所有 token 放到一起。接着,query 和这些 token 计算 attention score,再对对应位置的 value 进行加权求和,得到第 2 个分支的 output。这个 output 同样会经过 gating,决定它对最终 output 的贡献。

最后是第 3 个分支,也就是 sliding-window attention,滑窗注意力。这个分支主要考虑了语言建模中的一个 inductive bias:相邻 token 往往更加重要。

如果把训练好的 Softmax attention 模型的 attention map 可视化,会发现很多注意力集中在 sliding window,以及最开始的几个 token 上。这个现象一般叫作 attention sink,是 StreamingLLM 中提出的。

这里直接利用了这个 inductive bias,默认 attend 到 sliding window,得到一个 output。同样,这个 output 也会经过 gating。最后,我们对 3 个分支分别做 gating,再把它们相加,得到最终 output。这就是 Native Sparse Attention 的整体思路。

张小珺

在这张图里,它是如何平衡全局和局部信息的?压缩、selected 和 sliding window 这 3 个分支是怎么协同工作的?

松琳

我觉得这是让模型自己学出来的。因为模型有一个门控机制,可以通过每个 query 位置对应的表示,自适应地学出每个分支应该分配多少 attention。

这也是一种自调节机制。与其人为规定每个分支应该选多少,不如把选择直接交给模型,让模型自己学出需要 focus on 哪个分支。

第 2 个分支其实和 Quest 很像。Quest 需要把每个 chunk 得到一个 representation,再通过 query 和 chunk 做 attention,计算 top-k,并完成 block selection。

所以这篇工作的 novelty 并不主要在这些机制本身,后面我会讲。它主要集中在如何解决效率问题。Sliding-window attention,以及粗粒度和细粒度 attention,在已有 literature 中都能找到。它的主要贡献是如何让这些机制变得更加高效。

这篇工作也提到了 Quest,但 Quest 存在一个问题:它不能很好地和更加现代的 attention 机制兼容。

我们先简单回顾一下 GQA。Multi-head attention 中,每个 head 都有自己的 key、value 和 query。Multi-query attention 则是所有 head 共享同一份 key 和 value,因此只需要一份 K 和 V。

这样在 inference 时,整体 KV cache 的大小会从原来的多份缩小到只有 1 份,所以 MQA 的 inference 效率会高很多。但是它共享得过于激进,可能会带来 performance degradation。

为了缓解这个问题,可以做一种折中,不让所有 head 都共享,而是每 2 个 head 共享一份 key 和 value,这就叫 Grouped-Query Attention,缩写是 GQA。

回到 Quest。在 MHA 中,每个 head 的 key 和 value 都不同,所以每个 head 本来就要读取不同的 KV cache。即使每个 head 选择不同的 KV block,整体读取量也不会增加,因为每个 head 都有独立的 KV cache。

但 GQA 不一样。在一个 group 里面,多个 head 共享 key 和 value。如果 group 里的不同 head 选择不同的 KV block,就需要读取不同的 KV cache;如果它们选择相同的 KV block,就只需要读取一份共享的 KV cache,这样能够减少读取开销。

也就是说,在 GQA 中,如果同一个 group 里的不同 head 选择相同的 KV block,inference 的读取量就会减少。因此,Native Sparse Attention 强制要求同一个 group 里的所有 head 选择相同的 KV block,这就是它在 decoding 阶段能够加速的原因。

Quest 和 Kimi 的 MoBA 都允许每个 head 选择不同的 KV subset,而 Native Sparse Attention 为了减少 GQA 下每个 head 选择不同 KV block 带来的开销,强行限制同一个 group 中的 head 选择相同的 KV block。

为了保证同一个 group 里的不同 head 选择相同的 KV cache,在 block selection 时会进行求和。每个 query 对不同 block 都有一个 score,然后把 group 里面所有 head 的 score 加起来,得到唯一的 selection score,再进行 top-k selection。

这样可以保证同一个 group 里面所有 head 的行为一致。Native Sparse Attention 就是通过这种求和,强制 group 中所有 head 选择相同的 KV block,从而获得 GQA 带来的额外读取效率。

这部分还涉及一些 corner case。比如 selection block size 和 compression block size 不一定相同,当两者相同时比较容易处理,不同时就需要额外的机制。但这不影响核心的 high-level idea,所以这里先略过。

5. Hardware Makes Sparsity Work

接下来讲训练效率。我们先回顾一下 FlashAttention 2 的 kernel。它会把一个连续的 query block 放到 SRAM 里,这样 memory access 会非常高效。

对于同一个 query block,其中的不同 token 可能会注意到不同的 KV block。最后,对这个 query block 来说,需要从全局内存中读取的 KV block 就是所有 token 所需 block 的并集。最坏情况下,query block 中每个 token 想要 attend 的 KV block 都不一样,这样就必须把整个并集读进来,会造成低效的 memory access。

我们可以回顾 FlashAttention 的做法。它以 block 的形式读取 query。假设 block size 是 B,每次读取一个 B×D 的 query block。key 从头开始循环读取,每次加载一个 B×D 的小矩阵。

这样,query 的 B×D 矩阵和 key 的 D×B 矩阵做矩阵乘法,就得到 B×B 的 attention matrix。再读取对应的 value,也就是 B×D 的矩阵,先做 online Softmax,再进行矩阵乘法,得到 output。

FlashAttention 采用 blocking,主要是为了利用 tensor core 的加速,同时避免把 attention score 写回 HBM。HBM 可以理解为全局显存,而 SRAM 是 GPU 内部更高效的存储区域,类似 cache。

SRAM 的容量比较小,但 bandwidth 很高。HBM 就是普通 GPU 显存,比如 A100 的 80GB 显存,带宽显著低于 SRAM。硬件高效算法会尽可能把计算保留在 SRAM 中,减少 SRAM 和 HBM 之间的通信。

对于 full attention,每个 query block 反正都要扫描之前所有的 key 和 value block,所以不会有额外浪费。但对于 sparse attention,不同 query 所需的 KV block 不同。如果仍然连续读取 query block,就必须加载每个 query 对应 KV block 的并集,这会造成 memory access inefficiency。

Native Sparse Attention 利用了 GQA 的特性。MQA 或 GQA 在 inference 时,可以通过 attention head 这个维度,把计算凑成矩阵乘法的形式。它有点像 batchify,只不过因为 key 和 value 是共享的,所以存在一个可以 batchify 的维度。

FlashAttention 会把一个 query block,也就是 B×D 的 block,计算成对应的 B×D output。通常 B 可能是 128,所以会并行计算 128 个位置的 output。

但在 DeepSeek 的模型中,为了避免 query block 带来的不必要读取,它不再把多个 token 捆绑在一起计算,而是让每个 token 独立计算自己的 output。FlashAttention 之所以需要捆绑,是因为要凑矩阵乘法;如果 query block size 这一维不存在,就无法使用矩阵乘法。

GPU 非常喜欢矩阵乘法,因为这样可以调用半精度矩阵乘法和 tensor core。Tensor core 是 GPU 上专门用于计算半精度矩阵乘法的计算单元,非常高效。

以 A100 为例,在相同 FLOPs 下,用半精度矩阵乘法计算,比使用其他计算单元,比如 ALU,速度大约快 16 倍。在 H100 等硬件上,tensor core 会越来越快,最后矩阵乘法太快,反而是 memory loading 成为瓶颈,这就是常说的 memory wall。

所以这里的问题是,我们不想把 query block 捆绑起来读取,因为会引入很多不必要的读取,但又想凑成矩阵乘法。前面提到,GQA 存在一个 batchify 的空间。

对于相同的 key 和 value,会有多份 query,也就是有 H 个 query head。Native Sparse Attention 强制这些 head 共享所选择的 KV cache subset,也就是共享选择出的 block。既然它们共享相同的 KV block,就可以把它们打包在一起,同时计算。

这里的 H 可以看成 batch 维度,因此又能使用矩阵乘法。只不过这次矩阵乘法的维度不是 query block size,而是 query head 的数量。

此外,tensor core 并不是所有矩阵乘法都能优化,它有一些尺寸限制。比如在 Triton 中,H、D_K 和 B_K 至少要为 16。Native Sparse Attention 是用 Triton 写的,因此要求 H 至少为 16。

这个要求其实很强,因为通常每个 query group 下面不会有 16 个 head。但同时,模型又希望保留多个 group,以保证不同 token、不同 head 之间仍然有一定 diversity。

一般 GQA 会使用 4 个 query group,比方说像妈妈我记得是用了四个 query group。Native Sparse Attention 也把它强行设成 4 个。但为了满足 tensor core 的限制,它又把整体 head 数量增大。

比如这里 D_Q 是 192,head 数量是 64,两者相乘接近 12,000,而 hidden dimension 只有 2,560。这说明这里存在一个非常大的 up-projection。通常不会这样做,一般可能是从 2,560 投影到 2,560,再分配不同的 head。

Native Sparse Attention 之所以这样设计,完全是为了增加 head 数量,同时保留更多 group,从而保证不同 token 在选择 KV block 时有足够的 diversity。

这非常有 DeepSeek 的特色。DeepSeek-V3 里的 MLA 也是类似做法,会设置很多 attention head,并进行非常大的向上投影。DeepSeek 可能认为,既然是从头开始训练,只要 inference time 和 training time 都足够硬件高效,up-projection 也无所谓。

而且 MoE 本来就有很多参数,所以不会造成 attention 参数和 MLP 参数比例失衡的问题。在 MoE 框架下,这种选择是合理的。

DeepSeek 这里使用的 MoE 是他们之前的 MoE 方案,和这篇文章关系不大,所以就不展开了。

6. Why Sparse Attention Wins

接下来看看训练 loss。和 full attention 相比,它几乎全程都更好一点,两条曲线很接近,但 Native Sparse Attention 一直在下面。

表 1 是一些 benchmark 结果。在这些结果上,NSA 甚至会比 full attention 更好,这证明了预训练的重要性。

如果是在 full attention 的基础上再做 sparse attention,比如 Quest,性能通常会下降。但如果直接把 dynamic sparsity 机制用于预训练,甚至可能得到比 full attention 更好的结果。

很多模型是在 inference 阶段使用 sparse attention 技术来加速。因为它们不是从头训练的,所以需要尽量逼近 full attention 原本学到的 attention 分布。既然是一个逼近过程,就一定会有损失。

比如 Quest 和 full attention 相比会更差。还有一个 exact top-k,它直接算出完整的 attention score 矩阵,再取 top-k,可以看作所有 top-k 方法的上限。它的上限其实并不高,因为它始终是在逼近 full attention,最终上限就是 full attention 本身。

Quest 比 exact top-k 更受限制,因为它需要以 block 为单位选择 KV cache,这样 decoding 时才能一个 block 一个 block 地加载,从而提高效率。也就是说,它获得了硬件效率,但失去了一些灵活性,所以效果比 exact top-k 更差,这是合理的。

如果想让 sparse attention 甚至比 full attention 更好,唯一的出路就是从头训练,也就是 train from scratch。Native Sparse Attention 给出的启示是,不要在 full attention 的基础上再做 sparse attention,而是直接设计一种训练时也表现很好的 sparse attention。

这样就有可能跳过 full attention 的 performance 上限,解锁新的能力,甚至超过 full attention。

这里还有一些解释。首先,Quest 和 NSA 都不丢弃 KV cache。有些模型,比如 H2O,会进行丢弃 KV cache 的操作。近一年有一些工作说明,只要机制会丢 KV cache,就会对 performance 产生很大影响,因为某个 KV cache 可能会在之后被新的 query 使用。

这可能是 Quest 比丢弃 KV cache 的方法效果更好的一个直觉解释。其次,在不丢 KV cache 的基础上,sparse attention 还有一个作用,就是帮助 attention 集中注意力。

Full attention 其实是非常冗余的机制。Attention 自己学出来的分布本来就是 highly sparse 的,里面有很强的 sparse pattern。大家可以把训练好的 attention 模型的 attention map 取出来,很容易发现其中存在大量冗余。

直接用 flat sparse attention 训练,我的直觉是,它在训练阶段就已经去除了一些冗余,能够帮助 attention 更加 focus,不会被无关紧要的 token distract 到,所以效果可能会更好。

最后是 needle in a haystack。图中一片绿色,说明它的 retrieval 能力仍然在线。这也很好理解,因为它没有丢 KV cache,模型只需要动态选择就可以。

更神奇的是,它不仅增强了长文本能力,在 Transformer reasoning 上也有提升。DeepSeek 使用了 DeepSeek-SP 进行 SFT,在这些 SFT data 上训练。结果显示,NSA 的 R1 reward 甚至比 full attention 的 R1 reward 更好,说明 dynamic sparsity 确实很有潜力。

图 6 比较了 Triton-based Native Sparse Attention kernel 和 Triton-based FlashAttention kernel。为了公平比较,都是 Triton 对 Triton。可能他们没有时间写 CUDA kernel,而通常 CUDA kernel 会比 Triton kernel 更快。

公平比较后可以发现,在比较短的序列长度上,Native Sparse Attention 就已经比 global attention,也就是 full attention,更好。序列长度越长,加速比越高。在 64K 长度上,加速接近 10 倍。

最后是 decoding 的 expected speedup。之所以是 expected,是因为 decoding 是 memory-bound 的过程,整体效率基本取决于需要读取多少 KV cache。

前面提到,NSA 对 MQA 更有优势,因为它限制了同一个 group 中的 head attend 到相同的 KV cache,从而减少读取量。相较于 full attention,它会更快;相较于 Quest 或 Kimi 的 MoBA,它在 decoding 上也有天然优势。

这是我非常喜欢 Native Sparse Attention 的一点。它发现让同一个 group 里的 query 共享选择结果,可以带来很多提升。虽然这会有局限,但它通过比较激进的修改,比如增加 head 数量,保证模型 performance 仍然在线。

最终,它既保持了很好的 performance,甚至效果更好,又在 inference time 非常高效。这一点非常难得。

最后看 ablation study。这里研究了 block selection 的机制。前面提到,block selection 中有一个 compressed attention 分支,它基于每个 chunk 压缩后的 key 和 value 计算 output。

这个分支在训练时会有梯度传回去,模型可以通过 end-to-end training 学到哪些 block 更重要。如果没有 compressed attention 分支,block selection 就会存在一个不可导过程。

因为 top-k selection 本身是不可导的,梯度无法通过 selection 传回去。但有了 compressed attention 分支,就能够提供训练信号,使整个 block selection 过程更加 informative。

他们还和其他机制做了 ablation study,比如 Quest。Quest 是一种 training-free 的 importance score computation,具体策略可以参考 Quest 原文。

他们也尝试了其他 block selection 方法,结果都不如当前方案。这说明如何选择 block 非常关键,也证明了 compressed attention output 的重要性。

这篇工作强调 selection 是可回传、可以反向传播的,就是因为有 compressed attention 分支,可以把梯度反传回去。他们把常见的替代方案都拿来比较,证明自己的方案更优。

这就是 Native Sparse Attention 值得讲的部分。

张小珺

我来问你几个问题。这篇工作是小创新多,还是大创新多?

松琳

我觉得大家很有可能会浮于表面。比如看到这张图,就觉得这篇工作没有什么创新:sliding-window attention 已经用得非常多了,block selection 在 Quest 里已经有了,compressed attention branch 也不是罕见的机制。那这篇论文不就是一个缝合怪吗?有什么了不起的?

很多人可能会有这样的观点。但我觉得这篇工作的 significance 不在于这些模块,而在于它对 selected attention 做了改进,使它无论在 autoregressive decoding,还是在训练阶段,都能在硬件上更加高效。

这才是非常重要的地方。只有这样,它才能用于预训练,并且在 decoding 时也有更高效率。所以这篇论文强调的不是 3 个分支的 novelty,而是 natively trainable fast attention,也就是第一个把 query-aware 的 dynamic sparsity 用于预训练的工作。

前面也看到,这套机制在预训练阶段确实很快,而且相比 full attention,表现甚至更好,达到了既快又好的效果。这就不愧是 DeepSeek,非常惊艳。

张小珺

所以它能够在硬件高效实现训练的同时,降本增效,是吗?

松琳

降本其实是硬件高效的体现。只有硬件高效,才能降本,否则训练会比较困难。

Quest 这种模型也能做预训练。我们之后要提到 Kimi 的那个 MoE,它其实就是跟 Quest 的思路差不多,也是能做一些预训练的。但它的预训练效率可能不如 NSA。

Quest 和 MoBA 的 inference efficiency,在 MQA 的情况下也会有一些损失,这一点前面已经反复解释过了。

张小珺

这是算法架构创新,还是工程化创新更多?

松琳

从工程创新角度来看,Native Sparse Attention 可以属于这个类别。它没有提出非常新颖的概念,这些概念都已经存在,它只是把已有概念做得更加适配硬件,从而达到更快的训练和推理速度。

张小珺

也就是说,从这个层面来说,工程创新更多,实现了硬件亲和、hardware friendly。

松琳

对。

张小珺

这篇论文有没有你特别喜欢的部分?

松琳

我非常喜欢 kernel 的设计,因为我自己也很喜欢写 kernel,看到这种 kernel 写法会觉得非常有意思。

很多人可能觉得这个模型架构不够优雅,但我觉得它的 kernel 非常优雅,几乎是一个完美的、带有 sparse attention 的 kernel 设计。它基本没有浪费,又能利用硬件上的计算资源,比如 tensor core。

它通过 head 的维度来凑矩阵乘法,这也是我第一次看到有人在训练阶段这样做。之前已经有人在 inference 阶段这么做了,但这是我第一次看到它被用于训练。

张小珺

它的设计也体现出大胆,对吗?

松琳

确实非常大胆。很多人会担心,这么多 head 共享选择出来的 KV block,会不会掉点。我觉得绝大多数人都会有这样的 concern。

但 DeepSeek 说,我不管,我一定要 share。既然这样硬件高效,掉点的话,就把 head 数量增加,再加一些额外 branch。前面提到它有 3 个分支,这两个额外分支可能会稍微缓解共享带来的负面影响。

最后结果非常大胆地证明,即使 group 中的 head 选择相同的 KV block,performance 仍然很好。从这一点看,确实是非常大胆的举动。

张小珺

你之前跟我说,这是一种“硬件上的暴力美学”。能不能展开讲讲?

松琳

这是非常阴间的暴力美学。首先,它对 GQA 推理的分析非常有见解;同时,它通过强制不同 query 注意到相同的 KV block,再设计出这样的 kernel。

这里每一层都环环相扣,少一点都不行。它在硬件限制下,能够在刀刃上前进,像刀尖舔血一样,同时坚守自己的大原则。

这篇论文的大原则其实来自 Quest:每个 query 想选择不同的 block。Native Sparse Attention 在 Quest 框架下找到了一个非常接近最优解的处理方式,能够高效利用硬件,kernel 速度也非常快。

所以它可以称得上是一种硬件上的暴力美学。

张小珺

OpenAI 是怎么做这项工作的?他们不能叫 OpenAI,应该叫 ClosedAI,因为大家也不知道具体用了什么。

松琳

GPT-3 的 technical report 里提到过 sparse attention,但可能不是这种动态 sparse attention,可能是 BigBird 那种静态 sparse attention。

现在他们的 technical report 提供的 technical detail 越来越少,所以大家也猜不出到底用了什么架构。江湖上可能有很多传言,但没有办法考证。

7. Kimi Simplifies The Sparse Route

接下来我们看第 2 篇工作,也就是 Kimi 的工作。Kimi 最近也在往开源方向发展,MoBA 已经开源了。他们还开源了基于 Muon 新型优化器的方案,以及 KOBE 和其他项目,确实值得称赞。

Kimi 经常和 DeepSeek 对狙,大家会戏称这是“中美对狙”。这也体现出 Kimi 的技术实力。大家都认可 DeepSeek 的实力,但 Kimi 在研发方面同样很有实力。

这篇工作中也能看到一些熟悉的身影,研发团队的水平是在线的,所以可以带来很多和 DeepSeek 对峙的经典作品。

MoBA 和 DeepSeek 的工作非常像,都是在 Quest 这个 framework 下做不仅仅是 inference 加速的工作,还希望用于预训练。Quest 是一种 hardware-aware 的 sparse attention。

MoBA 本身非常简单。前面讲完 Quest 和 Native Sparse Attention 后,MoBA 很短时间内就可以讲完。

在 Quest 框架下,每个 key block 都需要计算一个表示,然后 query 和这个 key block 的表示计算 score。MoBA 直接使用 mean pooling 得到这个表示。DeepSeek 之前使用 MLP 做 projection,MoBA 则直接使用 mean pooling,不引入多余参数。

DeepSeek 有 3 个分支,MoBA 则砍掉了 compressed attention 的 output 分支,也砍掉了 sliding-window attention 的 output 分支,只保留中间的 selected attention 分支。

这样整个模型更加简洁。很多人会觉得 MoBA 比 Native Sparse Attention 更优雅,因为它的模型设计更加简单。

他们的团队相信 SGD 的力量,认为不需要额外的 loss,也可以通过 SGD 选出很好的 top-k block。但这可能和 DeepSeek 的 ablation study 有一些出入,所以具体哪一种更好,还有待进一步研究。

这两篇工作都打开了一个新世界的大门:原来动态 sparse attention 也可以用于预训练。之后可能会有更多研究 focus on dynamic sparsity 的预训练,而不仅仅是 inference。

MoBA 的整体架构非常简单。和 Quest 一样,每个 head 都可以自由选择自己想要 attend 的 KV block。

因此它的痛点就转移到了 implementation。Native Sparse Attention 可以利用 group 共享选择结果,每一步读取都不会浪费太多东西。但在 MoBA 中,不同 query 会选择不同 block,这种浪费无法避免。

它的做法是,对于每个 KV block,找到所有选择这个 block 的 query token。这些 token 可能是不连续的。如果直接读取 block,就会读入多余数据,所以他们会先把这些不连续的 query token 提取出来,存成一个连续 tensor,然后调用 FlashAttention kernel。

也就是说,对于每个 KV block,他们都要把选择这个 block 的 query token 全部提取出来。看他们的 codebase,主要就是各种 indexing 和 reindexing operation。

但这一步有 overhead,并不是免费的。如果 KV block 数量足够多,这一步可能成为 bottleneck。这也是为什么 MoBA 使用了 512 的 block size。

我们对比一下 DeepSeek 的参数。DeepSeek 的 block size 是 64,而 Kimi 这边是 512;DeepSeek 选择 top 16,Kimi 选择 top 3。它们最终覆盖的 token 数量可能差不多。

但因为 MoBA 不能像 Native Sparse Attention 那样写出相应 kernel,所以必须针对每个 KV block,把所有选择该 block 的 query 提取出来。当 block size 较小时,这会引入很多 overhead,可能无法接受。

所以虽然 MoBA 表面上很简约,但也要注意隐藏的问题。第一个就是 KV block size 没有办法设置得太小。

block size 如果太大,选择粒度就会太粗。你只能选择很少的 block,稍不注意就可能漏掉重要信息。DeepSeek 的 block size 是 64,选择 16 个 block,粒度更细,容错率也更高。选择 3 个 block 容易出错,选择 16 个 block,出错概率就会小一些。

这就是 MoBA 简约背后的代价之一:block size 不能太小,无法非常细粒度地选择 KV block。

其次,我们看它的加速比。在 128K 长度时,MoBA 其实没有什么优势。但 DeepSeek 的模型在较短长度,比如 8K、16K 时,加速比就已经非常明显了。这也是 Native Sparse Attention 更加硬件亲和的结果。

再看 training loss。随着训练进行,MoBA 基本上一直压着 full attention。它需要训练到一定阶段,才能把差距缩小,前期还是会比 full attention 差一点,所以我不确定这是不是某种 process 带来的负面作用。

MoBA 砍掉 compressed attention output 分支,相信 SGD 能找到正确 block,但也带来了一个问题:在 SFT 时会出现 sparse 的训练问题。

在 SFT 中,通常有一个很长的 prompt,也就是 prefix,这部分不参与 loss 计算,只计算后面一部分 token 的 loss。如果后面那一小部分 token 没有覆盖到某些 block,那么这些 block 就没有任何梯度信息,造成训练信号稀疏。

他们提到,MoBA 在 SFT 时会有 suboptimal performance。为了解决这个问题,他们采用了 hybrid approach,把最后 3 层切换成 full attention。

最上面的 3 层使用 full attention,会产生非常稠密的梯度信息,对每个 token 都有梯度,从而保证训练信号不会太稀疏,SFT performance 也会提升很多。

他们在知乎 blog 上写得很有意思,把自己的 trial and error 都写出来了,包括在不同阶段遇到的各种问题,以及如何解决这些问题。这可能是非常优质的 research model 训练数据。如果有人想训练一个做 research 的 AI 模型,一定要把他们的博客加入训练数据。

Native Sparse Attention 应该没有这个问题,因为它有 compressed attention 分支,直接参与 output 计算。output 会把 gradient 反传给前面的所有 token,所以不存在梯度稀疏问题。

这是 Kimi 特有的问题,因为它砍掉了 compressed attention,相信 SGD 可以找到正确 block,但 SFT 时信号可能不够,导致它无法找到正确 block。

好在 MoBA 的设计非常简约,没有引入其他东西。因为 mean pooling 不需要参数,整个过程非常简单,所以所有参数都一样,只是 loss 计算方式不同。这样就可以在 full attention 和 block sparse attention 之间自由切换。

他们在 SFT 中利用了这种自由切换的特性。我记得在预训练阶段也做过类似处理:90% 的 token 使用 sparse attention 训练,剩下 10% 的 token 切换到 full attention,做 global attention 训练。

我刚才看错了,这部分是 ablation study。他们发现这种 hybrid 训练方式确实能够提高 performance。也就是说,这种自由切换不但能在 SFT 上带来好处,在预训练阶段可能也有帮助。

这就是极简设计的利与弊。

最后是 evaluation。他们测试了更长的文本,长度比 DeepSeek 更长,测到了 1M 级别,而 DeepSeek 好像只测到 64K。后面的长文本结果都很好。据他们所说,MoBA 也已经实际投入 Kimi 的产品中,可见它确实有效。

这大概就是 Kimi 值得提到的部分。因为第 1 篇里已经讲了很多 preliminary 知识,所以回过头看 Kimi 就很容易理解它在做什么。

张小珺

Kimi 和 DeepSeek 的工作,你觉得有没有相互借鉴的地方?

松琳

我觉得它们可能代表两种设计哲学。DeepSeek 更偏硬件暴力美学,这一点从 MLA 等之前的工作也能看出来。Kimi 则更加极简,想把不必要的地方都砍掉。

具体技术选择路线不同,也会带来实现效率上的差异。如果想让硬件更高效,好像只能走 Native Sparse Attention 这条路。Quest 和 MoBA 也能做预训练,但提速比没有 Native Sparse Attention 那么夸张。

也就是说,MoBA 可能在很长的 pre-filling 上有提速,但一般预训练的文本长度不会很长,可能只有 8K。在这种长度上,MoBA 相较于 full attention 的训练没有明显好处;而 Native Sparse Attention 在 8K 就已经有实打实的加速。

这也和它们的设计哲学有关,因为 DeepSeek 非常强调软硬件协同设计。

张小珺

它对硬件有特殊要求吗?需要特定型号的硬件吗?很多人觉得算法不应该 overfit 某种硬件。

松琳

我觉得还是有很多共通的 principle。无论是什么硬件,有些原则都是共用的。比如矩阵乘法天然就快,不论是什么类似 GPU 的硬件,矩阵乘法都会很快。

连续的数据读取也是一样,在基本所有硬件上都更高效。所以要尽量把运算写成矩阵乘法,读取时尽可能 block by block 连续读取。

这些都是非常常用的 principle,短期内,甚至长期内,也不会发生改变。它告诉我们,设计算法时就要提前考虑 memory access:能不能 block by block 地读取,能不能写成矩阵乘法。

但很多算法从业者会忽略这一点。大家看到 sparse attention 的图,可能觉得它不够漂亮,认为模型架构是在拼凑。但从硬件暴力美学的角度看,只要它足够快,performance 足够好,我觉得它就是美的。

张小珺

所以一个是硬件上的暴力美学,另一个是极简优雅的算法,它们体现了不同的技术哲学。除此之外,它们在优化目标上是一致的吗?

松琳

我觉得优化目标是一致的。大家肯定都希望一个算法在 pre-filling 和 decoding 两个阶段都高效,当然是越快越好。

张小珺

接下来我们简单带大家看一下这周新出的两个动态稀疏注意力工作,然后聚焦一些 hybrid 架构。Hybrid 架构方面,MiniMax 01 是一个很大的模型,据说有 GPT-4o 级别的 performance。

松琳

8. MiniMax Bets On Hybrid Attention

我们来看 MiniMax 的架构图。这里有一个 M× 的标记,说明它有多层线性注意力。在这个例子中,他们使用的是 Lightning Attention。

M 的值是 7,也就是说,模型先有 7 层线性注意力,再有 1 层 Softmax attention,加起来是 8 层。模型总共有 80 层,所以会把这个结构重复 10 次:先 7 层线性注意力,再 1 层 Softmax attention,依次重复。

这个注意力机制和前面两篇完全不同,也是我的专业方向。我主要研究线性注意力。前面是客串讲 sparse attention,因为之前我也和大家一样,从零开始学习 sparse attention,精读之后分享自己的体会。线性注意力则是我的主场,我对它比较熟悉。

什么是线性注意力?我们之前把 attention 叫作 Softmax attention,是因为其中有一个 Softmax operator。线性注意力就是把 Softmax 操作去掉,非常简单。

我们可以看 iterative inference。attention score 是一个标量,我们把它移到后面,再调整两个运算的顺序,因为结果相同。最后会发现,这些都是线性操作,可以利用线性操作的结合律,把它们结合成另一种形式。

这里其实是外积的累加,得到一个恒定大小的矩阵,也就是 D×D 的矩阵。如果把这个矩阵看成一个 RNN 的 hidden state,那么 linear attention 的整体过程就是 linear recurrence,也就是不断累加。

每一步都会把当前时刻的 K、V 和 output 加到 hidden memory 里面。这种 hidden state 是一个 D×D 的矩阵,所以通常叫作 matrix-valued hidden state,用来区别传统 RNN。

比如 LSTM 使用的是 vector-valued hidden state,也就是一维向量,容量大概是低维的。线性注意力则实现了 state expansion,从低维向量变成了更高维的矩阵,因此 hidden state 的容量大很多。

这也是为什么线性注意力比传统 RNN 表现更好:它有一个 state expansion 的过程。State expansion 也是 Mamba 反复强调的概念。

我想强调的是,线性注意力有一种非常硬件高效的 state expansion 操作,因为它本质上是一个 outer product。前面讲 Native Sparse Attention 时反复强调硬件高效,在线性注意力这里又讲到了硬件高效,这足以说明,在这个时代设计算法,最好还是尽量让它硬件高效。

Mamba-2 为了解决 Mamba-1 硬件不高效的问题,进行了改写。后面可以看到,它改写之后和 Linear Attention 非常像。

线性注意力训练的困难在于,如果使用 self-attention 类似的形式,整体复杂度仍然和序列长度成平方关系,所以在很长的 sequence 下,训练没有优势。

如果使用循环形式,第一步是循环计算,缺少并行度,速度会很慢,像传统 LSTM 一样。第二个问题是,这个操作里面没有矩阵乘法,只有外积,以及矩阵向量乘法来完成 memory readout。

前面 Native Sparse Attention 通过比较暴力的方式凑矩阵乘法,在线性注意力中也需要凑矩阵乘法。循环形式做不到,那就把一个 chunk 打包起来,把 chunk size 这个维度凑成矩阵乘法的一部分。

这和 Native Sparse Attention 强行把 head 维度拉进来凑矩阵乘法有点像。很多时候,设计算法时都在思考如何凑矩阵乘法。

比如有一篇工作叫 FlashFFT。FFT,也就是快速傅里叶变换,是非常经典的算法,但在当前硬件上不一定很快,因为它缺少矩阵乘法。FlashFFT 为了硬件加速,把 FFT 写成等价的矩阵乘法形式,从而利用 tensor core,速度提升非常大。

短期来看,还是应该尽可能把算法写成矩阵乘法,这样可以直接利用当前硬件,不需要等待未来某种架构或量子计算机。把算法写成矩阵乘法、做 block-wise 操作,都是不同硬件之间共通的 principle。

Lightning Attention 也提到了 chunk-wise 算法,用来做高效训练。Chunk-wise 的思路不是像循环一样每个 token 计算 output,而是先把序列分成若干个 chunk。

对于每个 chunk,只计算一个 last hidden state。计算 output 时,结合两种算法:历史 context 的贡献可以用循环模式,local chunk 内部的贡献则使用并行模式。

需要强调的是,当 chunk size 取不同值时,这种 chunk-wise 算法可以还原出 linear attention 和 self-attention 的并行形式。它是一种数学上等价的变换,不是 approximation。

简单讲一下 chunk-wise 算法。图中两个灰色 hidden state 表示中间部分被跳过了,直接跳到最后一个 hidden state。我们把窗口中的几个 token 打包,一次性算出中间的 hidden state update。

因为 output 的累加可以写成矩阵乘法,所以这里可以直接用矩阵乘法。在 output computation 之前,用 query attend hidden state,计算 output。

对于整个 chunk 来说,所有 query 都 attend 同一个 memory,这就有点像 batchify。这里的 chunk 维度有点像 batch dimension,因此可以凑出矩阵乘法。

hidden state 已经编码了当前位置之前的所有历史信息,所以可以用循环形式考虑整个历史 context 的贡献。正因为如此,这种 chunk-wise 形式是 exact 的,而不是 approximation。

另外,还有一个 local chunk 的形式,可以看成 local attention,直接计算 chunk 内部的 output,也会使用矩阵乘法。整体复杂度是一个 sub-quadratic 的形式,其中 D 是 hidden dimension,L 是 sequence length,C 是 chunk size。一般会把 C 设成中间的固定值,不随 L 增长,因此总体训练复杂度不是平方关系。

这意味着线性注意力可能用于非常长的预训练。一个很好的应用场景是 video generation,它可以被看成序列建模问题。几分钟的视频序列长度很容易达到百万级,用 self-attention 做 video modeling 会有先天困难。

Chunk-wise 算法非常 general,可以和 attention with decay、attention with delta 等机制结合。Delta 也是很有意思的方向,和最近很火的 test-time training,以及 Google 新出的 Titans,都有很多关联。

最近的现代线性注意力模型包括 Mamba-2、GLA、Gated Linear Attention、Lightning Attention,以及 xLSTM。xLSTM 是 LSTM 作者团队的工作,第一作者就是 LSTM 相关工作的作者之一,它可以看成一种现代的、可并行的 LSTM,也可以理解为线性注意力。

有了 Transformer 之后,我们还要进行 IO 优化,就像 FlashAttention 一样。我们有一个开源库,里面实现了很多硬件高效的线性注意力模型,也支持非常多的 model。

如果大家想训练这些现代线性注意力模型,欢迎使用我们的库。遇到问题可以提 issue,我们一般回复得比较快。

过去大家对线性注意力的印象是 language modeling performance 太差,所以觉得它不 work。但实际上,从前两年到现在,线性注意力进展非常快,performance 已经提升很多。

大家可能需要 update 一下自己的 prior knowledge。它的提升方式有很多,比如增加 D_K。RetNet 是比较经典的工作,来自微软亚洲研究院,当时也很受关注。

Lightning Attention 就是 MiniMax 01 使用的模型。它里面的遗忘率在每个位置都相同,这会带来一个问题:不同位置缺少可选择性。

Mamba 强调 selectivity,LSTM 也强调 selectivity。在我看来,Mamba 的 selectivity 就是 LSTM gating mechanism 的新瓶装旧酒,本质 idea 是一样的,只是在新时代重新被挖掘出来。

Mamba-2 可以写成一种特定形式。和之前的线性注意力相比,它把 gamma 变成了 data-dependent 的 decay,每个位置可以计算自己的 decay。

之前我们提到过 3 种线性注意力形式,同样也可以把这种机制应用到 linear attention with decay。比如有一个 gamma_t,之前是一个 causal mask,而现在需要编码两个位置之间 decay term 的累乘。

这两者的等价性在 Mamba-2 论文中被称为 state space duality。论文是从矩阵分解等角度出发,用 semi-separable matrix 来解释。但这两种形式本质上是等价的。

很多人看 Mamba-2 时会觉得数学内容太多,不容易理解。但从 linear attention 的角度出发,这些东西会更好理解。它也有 Transformer 形式,和 Mamba-2 的 block decomposition 训练方法是等价的。

所以 Mamba-2 完全可以看成一种 linear attention model,只不过它有一个 data-dependent 的 decay。

回到 MiniMax 的工作,它也有 chunk-wise 算法和 IO 高效的训练方式。

9. Hybrid Models Finally Scale

接下来是 scaling law。他们发现 hybrid 架构的 scaling behavior 比纯 Softmax attention 更好。这其实很容易理解,因为 attention 和 RNN 结合,可以利用各自的优势。

RNN 有固定的 hidden state size,会强迫模型学习可压缩的 pattern。因为 hidden state 只有固定容量,所以它必须学习压缩历史信息。这种压缩过程可能和压缩智能的概念有关。

Softmax attention 则没有压缩概念,因为它会把所有 key-value cache 保存下来,但这也使它非常适合做 retrieval。RNN 经常被诟病不能很好完成 retrieval task,因为它只有固定大小的 state,不可能记住所有历史信息。

Attention 则非常适合 retrieval,因为它保留了所有历史信息的 KV cache。一个形象的比喻是,attention 像翻书,RNN 像人的大脑。人的大脑容量是固定的,需要时会去翻书,但一般情况下固定容量的大脑已经够用。

从 performance 对比来看,纯 linear attention 在 retrieval 比较强的任务上可能较弱,比如 needle in a haystack,这很考验模型的 retrieval 能力。

但如果把 attention 和 RNN 结合起来,效果可能更好。值得一提的是,hybrid 架构中的 RNN 还可以提供 positional encoding。很多 hybrid 模型不使用 positional encoding,直接用 RNN 提供位置信息,因为 RNN 自带可学习的位置信息。

这样就不需要 RoPE,也不会遇到 RoPE 的外推问题。所以在外推到更长上下文、做检索任务时,hybrid 架构可能有优势。

速度方面,8K 之后,Softmax attention 会越来越慢。Hybrid 模型也会变慢,因为仍然有 attention 层,但它的 attention 层比较少。比如 MiniMax 01 每 8 层才有 1 层 Softmax attention,所以整体计算量会小很多。

这篇模型在各种 benchmark 上证明了 hybrid 架构的有效性。它非常惊艳的一点,是第一次把这种 hybrid 架构 scale up 到这么大的尺寸,并且获得 GPT-4o 级别的 performance。

它能 scale 到这么大,也离不开硬件优化。前面提到的 chunk-wise 算法,如果放在三四年前、还没有这些算法时,线性注意力可能无法 scale up。

过去大家对线性注意力的印象是,理论上看起来快,但实际跑起来很慢。这个印象需要纠正,因为线性注意力这两年在 performance 和 efficiency 上都发展很快。

MiniMax 这篇文章做了很多工程优化,应该有非常扎实的 infra 团队。这里有各种并行策略,比如 MoE 中常见的 expert parallel,还有 pipeline parallelism。

在线性注意力部分,他们应该做了一些分布式的 chunk-wise attention 扩展,比如 Lightning Attention 团队实现的 sequence parallelism。

可以看出来,工程优化的细节很多,而且非常扎实。这也是他们能够把 hybrid model scale up 到这么大的原因。

这说明,如果一个算法硬件不高效,就很难 scale 到这么大的尺寸;如果不能 scale 到足够大的尺寸,也很难引起关注。现在大家都是 bitter lesson 的信徒:如果一个东西不可 scalable,兴趣就会大幅下降。

最近我们也做了很多线性注意力工作,让它能够硬件高效地训练。

Hybrid 架构并不是凭空出现的,也不是 MiniMax 独自发现它有效。比如 NVIDIA 去年发布过一项工作,作者列表里也有 Mamba 的作者。他们做到了 hybrid 架构,但没有 scale 到 MiniMax 这么大,所以没有引起同样大的关注。

这也说明,现在要让一项工作吸引足够注意力,还是要 scale 到很大的规模,用最终结果说话。在学术界看来是一个非常大的 scale,但可能在业界看来就是一个小模型。

在这种规模上,hybrid 架构已经不输其他架构。它们的 ablation study 也研究了 attention 层应该占多少比例,大概是 10% 到 20%。这和 MiniMax 01 的 attention 配比不谋而合。

社区还有其他工作,比方说 Weida 还有一个叫做 Hamba 的工作,也采用了各种 hybrid 思路。Sliding-window attention 也经常出现在 hybrid model 中。

之前大家其实已经知道混合架构有效,只是缺少一个团队把它 scale 到很大的规模。MiniMax 做到了这一点,所以技术才真正出圈。

我们做线性注意力,也要感谢 MiniMax 01,因为它证明了 hybrid 架构的有效性,让大家对线性注意力的关注度提高了。

张小珺

为什么 MiniMax 没有选择 Kimi 和 DeepSeek 的稀疏注意力机制?

松琳

我觉得可能是因为 sparse attention 的技术路线风险更小。Sparse attention 仍然保留所有 KV cache,而 attention 学出来的分布本来就是 sparse 的。

大家很容易想到,如果预训练一个 sparse attention,同时保证不丢弃 KV cache,那么模型的下限是有保障的。它的风险相对较小。

长文本方面,大家对 RNN 仍然有很多质疑。对于任何 RNN 模型,大家都会担心它的长文本能力,所以不太敢做这方面的研究。

Hybrid 架构可以看作 linear attention 面对这些质疑的一种 compromise,也是一种妥协。像 RWKV 这样的研究者可能不会妥协,他们很有信心,认为 RNN 这条路本身就够了。

但如果更关注长文本能力,比较实际的做法还是加入一些 attention 层。

张小珺

DeepSeek 在改架构上非常激进。这样说的话,MiniMax 是不是更激进?

松琳

MiniMax 有激进的部分,比如采用了这种架构。但从另一个角度看,熟悉这个领域的人会觉得它是比较 safe 的方案,因为它已经在美国工业界被验证了一年左右。

去年年初就有 hybrid 架构的工作,之后还有 NVIDIA 的模型,以及 IBM 的 9B 模型。它们已经反复验证了这条路线是能 work 的。

还有一个,比方说桑把,它也是一个 hybrid 架构,就是 Mamba 加上 sliding window attention,这个是微软他们做的。类似的混合架构还有很多。

如果相信这些 7B 规模的结果能够 transfer 到更大 scale,那么 MiniMax 其实也没有承担特别大的风险。

张小珺

那在这三家公司里,谁承担的风险最大?

松琳

我觉得可能还是 MiniMax。

Sparse attention 尤其是 dynamic sparse attention,主要节省的是计算,不节省显存。大家看到它保留了全部 KV cache,心里就会有底:这个方法首先可以做大海捞针,长文本测试结果也都很好,那就可以大胆使用。

MiniMax 的优化目标和前两个一样吗?

松琳

我觉得是一样的。高效注意力主要有两条路线,一种是 sparse attention,一种是 linear attention。它们的目标都是长文本训练,以及降低 inference 时的复杂度。

张小珺

从结果来看,谁的结果最好?

松琳

大模型很难直接比较。除非是在同一个环境中训练,否则数据、训练方式等都不一样,很难判断结果是架构带来的,还是模型本身带来的。

这个问题需要业界投入更多资源,进行更广泛的验证。

10. Architecture Becomes The New Scaling Law

现在 scaling law 在数据方面遇到了瓶颈,所以它们不约而同地改架构,希望通过架构优化获得更好的结果。

去年六月份的 keynote,伊尔就发表一个演讲,说 training time 的 scaling 可能到头了。最近 test-time scaling 概念很火,高效注意力一方面能够提高 test-time scaling,因为它更高效;如果 performance 能和传统 attention 持平,那么节省下来的 inference 计算就可以直接转化为 test-time scaling 的优势。

另一方面,靠堆数据的 training scaling 可能走不通了,所以确实需要更底层的架构创新。

DeepSeek 是比较少见的、非常坚持自己做架构创新的公司,这一点很可贵。很多公司会排斥架构创新,因为觉得堆数据就行了,为什么要冒险做架构创新?

但架构创新本身就是风险与机遇并存。比如 MLA 的成功,以及可以预见的 Native Sparse Attention 的成功。论文结果显示,它甚至可能比 full attention 更好。

所以在 data scaling 走完之后,大家可能会花更多精力研究架构。既然数据堆不上去了,就只能提升架构,改善整体 performance。

张小珺

今天我们已经分享了 160 分钟。最后问几个小问题。现在发论文已经成为各大模型公司竞争的新维度,你觉得是什么原因?

松琳

天下苦闭源公司久矣。如果开源、发 paper,可以给公司带来更大的技术影响力,客观上也可能给公司股价带来支持。

最近阿里股价涨得很快,就是因为 Qwen 系列模型开源做得非常好,让投资者有信心。公司首先要展示技术、秀肌肉。

对于后来者、追赶者来说,开源是最好的方式。因为你是追赶者,闭源并不会带来技术领先的优势,反而会因为闭源受到不利影响。

所以追赶者经常选择开源,就像 Meta 的 Llama 一样。

张小珺

我们今天讲了 DeepSeek、Kimi、MiniMax 这三篇工作。如果要做一个排名,你相对最喜欢谁的工作?

松琳

我可能会把 Native Sparse Attention 排在第 1。它确实是最近最让我眼前一亮的工作,能够让 Dynamic Sparse Attention 高效训练。

第 2 可能是 Kimi。因为我是做 Linear Attention 的,所以 Hybrid Model 能够 scale up 对我来说并不意外,因此只能无奈把 MiniMax 排在最后。它的东西确实不太意外,但 hybrid 架构值得投入更多资源研究,因为它已经证明了有效性。

张小珺

在 DeepSeek、Kimi、MiniMax 的工作里,有没有所谓的 aha moment?

松琳

不太好说,但我确实被 Native Sparse Attention 甚至比 full attention 更好这一点吸引住了。之后可能会做一些 Dynamic Sparse Attention 方面的研究。

张小珺

Kimi 的报告里没有给出它的架构改进比 full attention 更好的结论,对吧?

松琳

它们各有千秋。在这些 benchmark 上,我们看不到一个非常明确的、相对于 full attention 的优势。

但根据 DeepSeek 报告的结果,Native Sparse Attention 几乎是全线压制:无论是预训练 loss curve,还是下游 evaluation,甚至 reasoning,它都更高效、更好。

这一点让我觉得非常有意思。

张小珺

这个全线压制可能是今天最亮眼的地方。

松琳

对,因为它真的既要 efficiency,又要 performance,太厉害了。

张小珺

你觉得架构接下来还有哪些优化可能性和方向?

松琳

我觉得方向很多。比如 positional encoding,现在普遍使用 RoPE,我觉得 RoPE 也可以继续改进,能不能换成 contextualized positional encoding,或者干脆用 RNN 做 positional encoding?

之前提到,有些 hybrid 工作会直接用 RNN 做 positional encoding。这样的方法能不能直接 scale up,还有很多值得研究的地方。

注意力基本上就是 sparse attention 和 linear attention 两条路。之前还有一个比较火的方向叫 long convolution,很多模型会使用,但后来大家似乎没有继续关注,也许长卷积还会卷土重来。

就像 Linear Attention 和 Sparse Attention 一样,研究方向有时会经历低谷,之后又重新兴起。很多东西都是一阵一阵的。

张小珺

今年春节以来,DeepSeek 非常火。你觉得它有没有改变在美国的中国留学生对于国内大模型的一些看法?会改变你们毕业之后的职业走向吗?

松琳

人才都会往人才密度高的地方走。以前国内落后很多,所以大家愿意去硅谷的公司,学习更先进的大语言模型。

但现在国内在这方面很容易看出来,已经在追赶,甚至有些地方已经超越了。我觉得国内对海外留学生的吸引力会增加很多。

逐篇讲解DeepSeek、Kimi、MiniMax注意力机制新论文(含投屏版)——“硬件上的暴力美学” | BidClub