松琳
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 非常火。你觉得它有没有改变在美国的中国留学生对于国内大模型的一些看法?会改变你们毕业之后的职业走向吗?
松琳
人才都会往人才密度高的地方走。以前国内落后很多,所以大家愿意去硅谷的公司,学习更先进的大语言模型。
但现在国内在这方面很容易看出来,已经在追赶,甚至有些地方已经超越了。我觉得国内对海外留学生的吸引力会增加很多。