晚点
：两位可以先自我介绍一下，以及讲讲你们过去和注意力机制改进相关的研究。
肖朝军
：我是来自清华大学计算机系的博士生肖朝军，一直专注于大模型高效架构层面的研究。2024 年，我们发布了 InfLLM，也是做大模型注意力机制的改进。
一年多前，开源模型最多也就处理 8K 就是 8000 token 左右的文本长度。我们就想，能不能在无需训练的情况下，在推理阶段（可以理解为大模型的使用阶段）用稀疏注意力机制处理更长的 128K 甚至 1000K 的文本。这个注意力机制的思路和现在的 DeepSeek NSA 类似。而 NSA 是进一步把稀疏注意力做到了预训练阶段，写了非常丰富的算子，做了很精细的设计，是一个很大的进步。
（注：大模型里，稀疏指仅部分连接，效率更高；对应的稠密，也就是 dense，指数据 / 计算全连接。）
傅天予
：我是来自清华大学电子系的博士生傅天予，我们也一直关注神经网络模型的高效性优化。
也是在 24 年，我们发布了一篇和今天要聊的 MoBA 标题很像的论文，叫 MoA——Mixture of Sparse Attention（混合稀疏注意力机制），而 MoBA 是 Mixture of Block Attention（混合分块注意力机制）。我们发现，之前的一些稀疏注意力机制，虽然能使模型流畅地输出内容，但模型其实没有有效利用远距离的上下文信息，我们就希望能分析不同注意力头的各自特性，找到究竟其中哪些对模型能真正利用长上下文很关键，然后更多保留这部分注意力头。
晚点
：我们可以先聊关于 Attention 的一些基础问题。2017 年提出目前大语言模型主流架构 Transformer 的那篇论文，标题就是 Attention Is All You Need。简单来说，Attention 作为大模型的核心模块，它是用来做什么的？
傅天予
：要解释这个，我们要先知道一个大模型的输入和输出是什么，它的输入就是一系列 Token，也就是 “词元”，简单理解就是一个词。每个词本身有一个含义，但也有必须联系上下文才能表达的含义，比如代词 “it”，要放到上下文里才知道指什么。所以 Attention 机制的核心思想就是希望能知道输入的每一个词和之前的词的关系。
在标准 Attention 里，每新输入一个词都要和前面输入的所有词去计算相关性，它是一个关系强弱加权。这样不仅能帮大模型理解这个词本身，还能理解它的语境
肖朝军
：补充一下，大家也可能需要了解，再往前，神经网络是怎么处理序列（输入模型的一系列 Token 整体上是一个序列）的。
之前的主流方法是循环神经网络（RNN）。循环是什么意思呢？就是一个词、一个词地去理解，去动态维护一个 “记忆”。
比如当 RNN 来处理 “我是肖朝军” 这句话时，会先处理 “我”，再是 “是”，再是 “肖”，每处理一步，就把之前所有输入都变成 “记忆”，也就是系统里的一堆矩阵。在 “我是肖朝军” 的例子里，RNN 处理到 “肖” 时，已经把 “我” 和 “是” 转化成了一个 “记忆”。
RNN 里的 “记忆” 的存储大小是不变的，但这会使它 “遗忘” 很久以前的词元。当只是一句话时，这个问题不严重。但如果像我们今天这样录 2 个小时的播客，我最开始说完我是谁后，之后只会说 “我”，这样到后面，模型就会忘掉 “我” 和 “肖朝军” 的关系。这是 RNN 的缺陷。
而 Transformer，它所使用的 Attention 机制，就是把过往所有词都以一个个向量的形式存储下来。这样在 2 小时的聊天里，我每次提到 “我”，它都会去和过往所有词做一次相关性计算，就能知道 “我” 和 “肖朝军” 高度相关。
这也导致，Attention 的 “记忆” 不是固定的存储大小，而是把 “记忆” 扩展为过往全部词元，这就能改进 RNN 长序列的遗忘问题。
晚点
：所以当年从 RNN 到 Transformer 提出 Attention，本身就是为了解决长文本问题。
肖朝军
：对，不过当时的长，也没多长，就是几百 Token。
在 Attention 之后，很快就提出了稀疏 Attention，最开始是为了 BERT（注：Google 最早的 Transformer 架构的大语言模型，但它的优化目标和同为 Transformer 架构的 OpenAI 的 GPT 系列模型略微不同，BERT 一度是比 GTP 更主流的大模型），它只能处理 512 Token。
然后稀疏 Attention 出来后，让模型能处理 4K token 了，这是 18、19 年的事，那时大家认为这已经非常长了。
而到了现在，动辄是 128k 甚至 1M（100 万）的长度，发展非常快。所以 Transformer 确实是为了长文本，只是这个长和现在的长还不太一样。
原初 Attention 两大瓶颈：“显存开销” 和 “计算复杂度”
晚点
：从长文本到了长长长长长文本。后来原初 Transfomer 里标准的 Full Attention 又遇到了什么瓶颈，为什么不够用呢？
肖朝军
：首先就是存储问题，其实就是刚才讲的，它的 “记忆” 大小是不固定的，序列越长，记忆就越大。比如有几百万个词，那存下来要好几百 G，现在 GPU 的显存（显存是 GPU 的 “专属内存”，最初主要存储图像信息，在大语言模型里，它用来存储模型参数和训练数据）就几十 G，存储远远不够。
第二是计算复杂度的问题，因为 Full Attention，是每处理一个 Token 时，都需要和前面所有 Token 做一遍相关性计算，如果序列长度是 N，它的计算复杂度就是 N 的平方，当长度越长，计算时间就会越来越高。
傅天予
：我补充一点关于训练和推理阶段的 Attention 的区别。之前很多稀疏注意力工作都是一个稠密训练的 Attention，但在推理阶段把它变成稀疏的。这是因为我们发现 Attention 本身含有稀疏性。
这可能有几个主要原因：一是 Attension 里有 Softmax（软最大值）机制。简单理解，它是一个 “更软” 版本的取最大值的过程。一般取最大值时，会把最大的数取为 1，其他全部取 0。而 Softmax 是把最大的那个取为 0.9，第二大的可能是 0.09，第三大的是 0.009……这导致在这么多相关性数值里，我们一般叫 Attention Score 或 Attention Weights，本来就只有很少的数很大，这就为稀疏性提供了一个潜在的数学保证。
（注：如前文所说，Transfomer 里的 Attention 是每输入一个新 Token，要计算它与前面所有 Token 的相关性，所以会得到一组值。常规在一组值里取最大值时，是将最大的设为 1，其他都设为 0；而 Softmax 是对所有值进行一种函数处理，使它们都分布在 0 到 1 之间，且总和为 1，最大值会得到最高的权重。这些权重之间的数值悬殊很大，是 “稀疏” 的，一些很小的值，后续处理时可能被舍弃。）
二是从语言学的角度，比如我们今天播客可能有 2 小时，但在 “朝军” 和 “我” 之间产生的关联非常稀疏，就是这 2 小时里，前面有几万个 Token，但和 “朝军” 有强相关性的只有 “我” 这个词，关联本身就很稀疏。
最后是从人脑角度，也不是所有神经元都和所有神经元连接，它的连接也非常稀疏，并且随着年龄增大，会越来越稀疏。
所以无论从数学上，还是语言学上，还是神经科学上的启发，大家就觉得稀疏 Attention 本来就存在，并且是可以被利用起来提高效率。
晚点
：总结而言，在 Full Attention 的基础上，后来有哪些改进思路？
肖朝军
：我认为主流分类就是稀疏注意力和刚才提到的 RNN（循环神经网络），它又复活了，比如 Mamba 这类模型，RNN 计算很高效，但有遗忘机制。
（注：Mamaba 是由卡内基梅隆大学在 2023 年提出的基于状态空间模型的高效长序列建模架构，主要作者是 Tri Dao 和 Albert Gu 等，后文提及的 FlashAttention 的作者也是 Tri Dao。）
傅天予
：不过现在大家好像不喜欢说自己是 RNN，会说是线性注意力。我觉得线性注意力里的一些工作确实还是以循环的方式在进行，但也会有少部分工作，虽然也是把很长的序列压缩到一个 “记忆” 里，但不像传统 RNN 是逐词迭代进行的。
晚点
：哪种更主流呢？最近的 NSA 和 MoBA 都是稀疏注意力方向的改进。
肖朝军
：在 Mamba 之前，稀疏更主流。但 Mamba 之后，线性注意力机制的改进也非常多。
傅天予
：稀疏注意力整体对原本注意力的改动没那么激进，与现有很多方法能直接搭上。而线性注意力可能描绘了一个更美好的蓝图。
因为它的存储大小是固定的（注：线性注意力可理解为 RNN，前文有提及 RNN 的存储大小是固定的），计算复杂度随序列变长，不是平方增长，而是线性增长（这和它的计算方式有关）。
所以从 Scaling Laws 角度看，当输入比较短时，线性注意力的提效可能不如稀疏注意力，但输入特别长时，线性注意力的效率优势会上来。
不过目前大家更多尝试稀疏注意力，因为很多学界探索表明，纯线性注意力的效果（效率衡量的是速度和计算资源消耗，效果衡量的是模型 “聪明” 程度）欠佳，还需要混合其它注意力机制。
晚点
：Attention 的提出和后来的改进都是为了处理长文本。除了改进 Attention 机制本身，还有哪些努力能帮助处理长文本？
傅天予
：优化一个大模型，会有 3 个层次：算法、系统和硬件。硬件层很好理解，你用更快的 GPU，当然可以变快。
系统层稍难理解，比如 FlashAttention 就是一个系统层改进。虽然 GPU 的峰值算力和峰值显存是固定的，但代码写得好不好，决定了你能用到百分之多少。系统层的工作就是追求提高这个百分比，充分利用带宽、显存、算力。FlashAttention 让 Attention 的计算效率提升了一个台阶，显存则至少降了一个量级。
再上面是算法层，细分又有模型层和更普遍意义的算法。稀疏注意力、MoE（混合专家系统）都是对模型层的改进。而改变生成范式——比如从一个、一个词往外蹦，变成两个词、两个词往外蹦，则是更普遍的算法改进。
（注：MoE 现在是一个主流的大模型架构，比如 DeepSeek-V3 就是一个 MoE 模型，它也是 DeepSeek-R1 的基座模型。V3 也有生成方式上的改进，使用了 Meta 提出的 Multi-token Prediciton，一次生成两个词。）
肖朝军
：算法层有很多改进，因为长文本的问题很多，不仅仅是计算复杂度和显存。那针对不同场景，就有不同的优化方向。
比如 DeepSeek 在 DeepSeek-V2 里提出的 MLA 也是一个算法改进，它解决什么问题呢？就是刚才提到的，Transformer 需要把所有词元都存下来，存储量很大，这样每次去访问它的时间会很长，要一个一个读出来。所以 DeepSeek 就想把这个维度压低、存储变小，就做了这方面的一个改进。
（注：MLA 是 Multi-head Latent Attention，多头潜在注意力。它通过一种联合压缩技术，把算注意力分数时的显存占用减少了超 90%，同时保持了模型性能。）
所以我特别关注 NSA 和 MoBA 的理论加速比能不能落到实际，从论文看，他们都做得比较好。我们之前的 InfLLM 就是理论加速比高，但实际加速比不够好。
晚点
：为什么之前加速比很难落下去？
肖朝军
：我一直在搞算法，不是很懂底层系统，当时我也和别人讨论过，就觉得稀疏天然不适合 GPU，就做罢了。
但看到 NSA 之后，我觉得还是自己的认知有局限，就是我不懂硬件，没有在算子层做优化。而 DeepSeek 真能把这个想法推下去，这是 NSA 给我的最大冲击。
改进稀疏注意力的思路已比较成熟，但 DeepSeek 和 Kimi 真的把加速落到了真实系统
晚点
：这次 NSA 也特别强调了它是 hardware-aligned（硬件协同）的，就是结合硬件层做了优化，这个可以解释一下吗？
傅天予
：这需要理解 GPU 到底在怎么做计算。GPU 的一个特性就是 “single instruction, multiple data”（单指令、多数据流），即给出一条相同指令，对不同数据做同样的操作。
比如只计算一个矩阵时，GPU 比较慢，但如果对 20 个矩阵做同样操作，它就很高效。所以 GPU 是高度并行计算，并行度越高，对 GPU 越友好。这是计算方面。
在内存方面，计算机的内存是存一串连续数字。GPU 顺序访问连续数字的效率很高，但如果要访问的数据比较碎片化，GPU 一次取一大块数，却只有其中一小部分有用，效率就会很低，反之，如果取一大块数，这些数都有用，效率就很高。
理解这两点后再来看 NSA，它做了一个对 GPU 很友好的操作，就是它是 Block（块，可以简单为理解为是一组数）的，这又进到了 Kimi 的 MoBA 的标题里了，MoBA 的 B 就指 Block。
所以 NSA 虽然是一个稀疏注意力，但它的粒度是到 “块” 的，要么这一块都有用，要么都没用。这就满足了刚才说的两个期待：内存上是连续访问，计算上是对同一块数据做相同计算。它对硬件层的主要优化，就是提升 Block 的计算。
肖朝军
：对，打个广告。我 24 年 2 月的 InfLLM 的核心观点也是，稀疏注意力，就应该在块层级去做，才会对硬件更友好。
晚点
：其实我也问过 R1，NSA 和 MoBA 的异同，它说：“MoBA 更侧重 Block 级别的调控，而 NSA 更侧重 Token 级别的精细化操作和硬件优化。” 实际上二者都是 Block 级别的吗？
肖朝军
：NSA 也是在 Block level 做的，但文章里没特别强调。因为即使是 Block，真要做到非常好的硬件优化，还是需要很多操作，所以它后面更多强调了它的硬件优化。
晚点
：硬件优化靠什么具体方法实现呢？就是前面提到的写算子？
肖朝军
：对，就是写一些系统层的算子（注：算子是神经网络中执行具体数学运算的模块、它在模型较底层，直接与芯片硬件交互）。DeepSeek 在 NSA 上写算子时用的是 Triton，这是 OpenAI 开源的、在英伟达的 CUDA 上再抽象了一层的接口。
傅天予
：我想补充一个非常有意思的事，因为 Triton 是 GPT 3 那会儿出来的（注：Triton 开源于 2021 年，GPT-3 发布于 2020 年 6 月）。一个不靠谱的坊间传闻是，GPT 曾尝试用稀疏注意力来训练，所以去做了系统适配，这就有了 Triton。
不管传闻真假，Triton 确实非常适合做块状稀疏计算。虽然 OpenAI 最后自己可能没有用 Triton，但它已经成为学界要做块状稀疏计算的很好的工具。
晚点
：所以这是 OpenAI 还 open 的时候，给社区带来的一些贡献。
肖朝军
：补充一下，我觉得 DeepSeek 未来可能会继续深入到 CUDA 层（CUDA 是英伟达自己提供的，可以帮开发者更好使用 GPU 的一套软件平台）去写优化。因为 Triton 是对编码友好，CUDA 可以做更深。
还想补充一点， OpenAI 有没有做稀疏注意力我不知道，但它一定在长文本上做了新的工作。最近它发了 DeepResearch，可以想象，长文本能力一定很强，只不过它自己没有强调。
晚点
：具体到 MoBA 和 NSA 的注意力机制的设计，它们的异同是什么呢？
傅天予
：我觉得共同点主要有三点。一是，它们都是先从更高层次选一块需要关注的上下文，这个所谓的 “一块” 就是 “Block Level”，具体选法，二者有细微差别，但总之要先选一块，把这一块内容拿进来，再去关注块内部的细节。
二是，它们选块时，都倾向关注离当前这个词最近的词，这也非常自然，因为即使是稠密注意力，天然也会表现出对临近的词更关注。这从语言学上也能理解。
