上周,DeepSeek、Kimi 都放出了新的大模型架构改进和优化成果,分别是 NSA、MOBA,它们都聚焦于大模型注意力机制的改进。注意力机制是当前大语言模型的核心机制。回到 2017 年 6 月那篇开启大语言模型革命的论文,标题就是《Attention Is All You Need》——注意力就是你所需要的一切。
优化 Attention 的计算效率和效果,也能帮助解决 AI 学界和工业界都非常关心的一个问题,就是长文本能力。无论是一次输入一整本书,让模型帮我们提炼和理解;还是在现在用强化学习提升推理能力的范式转型后,需要生成很长的思维链;又或者是未来希望模型拥有越来越长的记忆,甚至像人一样拥有一生的记忆,这些都需要长文本能力的支持。
这期节目中,我们邀请了两位做过 Attention 机制改进的 AI 研究者。他们之前的成果,和 NSA、MOBA 一样,都属于大模型稀疏注意力方向的改进。一位是清华大学计算机系自然语言处理实验室的博士生肖朝军,他是 InfLLM 注意力机制改进论文的一作,导师是清华大学计算机系的刘知远;另一位是清华大学电子工程系 NICSEFC 实验室的博士生傅天予,他是 MOA 注意力机制改进论文的共同一作,导师是清华大学电子工程系主任汪玉。汪玉老师的实验室也做了很多与硬件结合的优化工作,两位刚好可以提供一些相互补充的视角。
注意力机制确实是一个涉及较多技术细节的话题,所以我也准备了一些术语和过往注意力优化成果的简单解释,希望可以帮助大家更好地理解。上期我们和美国西北大学的博士王子涵聊过后,有听友反映术语有些太多,听不太懂。最近我密集聊了一些 AI 研究者,有一种逐渐连点成线的奇妙感受。比如这次和朝军、天予聊 NSA、MOBA 这些改进时,他们也聊到了自己做研究的探索过程和一些反思,其中有不少内容和上一期节目形成了呼应。
比如两次我们都聊到了 FlashAttention。上次是在聊 DeepSeek 最新开源的 FlashMLA 时,提到了当年的 FlashAttention,它也是一个从系统层做优化的经典例子。这次聊 Attention 本身,肯定也绕不开 FlashAttention,它是早期对 Full Attention 的重要改进。两次节目我们也都聊到了在工程上做更底层的算子优化的重要性,因为它是把一个好想法真正落地的方法。
如果你虽然不是一个 AI 从业者,但又因为兴趣、好奇或工作需要,非常希望由浅入深地了解更具体的技术机制和进展,那欢迎多听我们的节目。这次我也在 show notes 里贴了一些 AI 大神用相对通俗的方式讲解 AI 的视频节目的地址,让我们一起学习起来。
程曼祺 Manqi Cheng
今天很高兴邀请到两位 AI 研究者,和我们聊一聊最近引起很多关注的大模型 Attention,也就是注意力机制的改进,包括 Moonshot 发布的 MOBA,以及 DeepSeek 发布的 NSA。两位可以先和听友打个招呼,也简单自我介绍一下。
肖朝军
大家好,我是来自清华大学计算机系的博士生肖朝军,一直专注于大模型高效架构方面的研究。
傅天予
大家好,我是来自清华大学电子工程系的博士生傅天予,我们一直关注神经网络模型的高效性。
程曼祺 Manqi Cheng
朝军和天予也可以讲讲,你们之前的工作和注意力机制改进有什么关系。
1. Sparse Attention Starts Here
肖朝军
2024 年的时候,我们发布了一篇文章,叫作 InfLLM,也是做大模型注意力机制的改进。当时的出发点同样是认为,大模型在处理长文本时,需要很多计算资源和存储资源。
在那之前一年左右,大家可以看到,开源模型最多也就能处理 8K,也就是 8,000 左右的文本长度。我们当时就在想,能不能不需要训练,只用稀疏注意力机制,让模型处理更长的文本,比如 128K,甚至 1M 的文本,于是提出了这篇文章。
它整体的 Attention 改进思路,和现在备受关注的 DeepSeek NSA 在 Attention 部分比较类似。NSA 进一步把这种方法做到了预训练阶段,并且写了非常丰富的算子,做了非常精细的设计,使得模型能够在预训练阶段就很好地利用动态稀疏 Attention。这是一个非常大的进步。
傅天予
我这边是在 2024 年左右,发布了一篇名字也和今天聊的 MOBA 很像的文章,叫作 MOA,也就是 Mixture of Attention,或者说 Mixture of Block Attention。
我们和刚才朝军介绍的工作一样,目标都是希望大模型能够更高效地处理特别长的文本。但我们这篇文章更关注的是,我们发现之前一些稀疏注意力的文章,虽然能让模型记住很长的上下文,或者说让模型在很长的上下文上流畅地输出内容,但模型其实并没有真正记住、或者使用很久以前的上下文内容。
所以我们希望分析不同注意力头各自的特性,找到究竟哪些注意力头对于模型真正利用上下文非常关键,再针对性地保留更多这部分注意力头的注意力。这样模型不仅能够顺畅地说话,还能够真正利用很长上下文中的内容。
2. Attention Replaces Fixed Memory
程曼祺 Manqi Cheng
在正式聊 NSA 和 MOBA 的具体改进之前,我们先聊一聊 Attention,也就是注意力机制的一些基础问题。
2017 年提出 Transformer 的那篇论文,标题就是《Attention Is All You Need》,可见 Attention 应该是大语言模型中非常核心的一个模块。两位可以帮我们介绍一下,Attention 是什么,大概用来做什么?
肖朝军
要说清楚 Attention 是什么,可能需要先知道大模型的输入和输出是什么。对于大模型来说,它的输入就是一系列 Token。我们也可以把它理解成一系列词。
Attention 的作用是,对于每一个词,它单独本身有一个含义,但有些词的含义必须结合上下文才能表达。比如代词 “it”,它究竟指什么,就需要结合上下文理解。
所以,注意力机制引入的核心思想,就是希望输入的每一个词能够和之前的词产生关系。在现在的大模型 Attention 机制中,输入的每一个词都会和前面输入的所有词产生关系。它会计算自己和之前所有词的关系,再根据关系的强弱,对之前词的信息进行加权。这样就能帮助大模型不仅理解这个词本身,还能理解它所在的语境。
傅天予
我补充一下。其实在聊 Attention 之前,大家可能需要了解一个更早的概念:神经网络是怎么处理序列的。
当时我们会使用循环神经网络。所谓循环,就是一个词一个词地去理解,同时动态维护一个记忆。每次理解一个词时,再判断这个词和记忆之间有什么关系。
在这个过程中可以预想到,记忆会发生遗忘。技术上也会遇到很多训练问题,比如梯度消失和梯度爆炸。Transformer 这篇论文提出《Attention Is All You Need》,认为传统循环神经网络天然存在缺陷,于是提出了 Attention。
Attention 叫作注意力,可以和人的注意力机制产生一定关联。它仍然是一个 Token 一个 Token 地处理,但问题在于,之前的 Token 需要和记忆做相关性计算,再提取重要内容。Attention 则认为,记忆不必是固定大小的,可以把记忆扩展为全部的过往词元。
这样,在处理当前词元时,它就可以和之前的所有词元逐一匹配,判断哪些词和它最相关,再提取相关信息。因此,它可以改进循环神经网络在处理较长序列时容易遗忘的问题。
程曼祺 Manqi Cheng
循环神经网络就是 RNN,对吗?那 RNN 中的记忆是什么样的?它把什么东西变成了这堆矩阵?
肖朝军
记忆就是一堆数、一堆矩阵。我们假设神经网络需要处理一个句子:“我是肖朝军”。它先处理“我”,再处理“是”,再处理“肖”。循环神经网络每处理一步,就会把当前输入变成记忆。比如处理“肖”的时候,它已经把“我是”这两个词变成了记忆。
在这个过程中,循环神经网络的记忆大小是固定的,只有一个固定大小的记忆。假设我们处理的只是“我是肖朝军”这几个字,遗忘现象并不严重。但如果是今天这样的播客,我在最开始说自己叫肖朝军,后面只说“我”,那这个“我”到底代表什么?模型需要把后面这个“我”和最开始提到的“肖朝军”建立关联。
可是,当输入长度特别长时,记忆大小是有限的。比如播客录了两个小时,模型可能已经忘记最开始出现的“肖朝军”这三个字。这是循环神经网络的缺陷。
对于 Transformer 来说,它会把过往所有词的表示都保存下来。过往词的表示就是一个个向量。两个小时之后,我再次提到“我”,这个“我”会和过往所有词做一次相关性计算,于是它能够很快算出“肖朝军”和“我”之间的高度关联,理解两个小时之后说的“我”仍然指肖朝军。
程曼祺 Manqi Cheng
所以从 RNN 到 Attention 的改进,本身就是为了在输入很长时,更有效地理解长上下文,对吗?
肖朝军
对。当时的“长”其实还没有今天这么长,几百个 Token 就已经被认为很长了。
再往前看,最早提出稀疏 Attention,是为了解决 BERT 这种早期预训练语言模型只能处理 512 个 Token 的限制。稀疏 Attention 出来之后,让模型能够处理 4K,也就是 4,000 个 Token。那是在 2018 年,当时大家已经认为这是非常长的长度了。
而现在,长文本动辄是 128K,甚至 1M,也就是 100 万的长度。所以整个发展非常快。Transformer 当时确实解决了长文本问题,但那个“长”和我们今天所说的长,概念已经不太一样了。
3. Full Attention Hits Its Limit
程曼祺 Manqi Cheng
也就是说,从长文本发展到了长、长、长、长文本。
Attention 在训练阶段和推理阶段都会发挥作用吗?现在很多改进在对比效果时都会对比 Full Attention。Full Attention 是以前标准 Transformer 里的版本吗?为什么后来大家觉得它不够用了?
肖朝军
Full Attention 就是最早 Transformer 里的传统 Attention。刚才提到,Transformer 会把所有词的表示都保存下来。比如我们聊了两个小时,可能已经有上百万个词了。
大家可以算一个数字:在 Transformer 里,一个词会表示成一个向量,一个向量就是一串数字。假设一个数字在计算机里占 2 个字节,当然实际情况还会涉及量化。现在一个词可能就需要几千个数字来表示,一层需要几千个数字,更多层还要更多数字。
假设已经有几百万个词,那么存下来的数据量可能就是几百 GB。电脑的存储显然远远不够,这是存储上的问题。
第二个问题是计算。处理任何一个 Token 时,都需要和前面的所有 Token 做一遍计算。我们聊了两个小时,它就需要和两个小时的词做相关性计算;如果聊四个小时、一天,就要和一天的词做相关性计算。
所以处理长文本时,Attention 是 N² 复杂度,N 就是长度。随着长度增加,时间和存储开销都会变得非常大。现在显卡的显存可能只有几十 GB,因此大家需要解决这个问题,否则就没有办法处理长文本。
傅天予
关于训练和推理阶段的区别,我补充一点。之前很多稀疏注意力工作,都是进行稠密 Attention 训练,在推理阶段再变成稀疏。之所以可以这样做,是因为 Attention 本身就含有稀疏性。
一个主要原因是 Attention 里有 Softmax 机制。可以简单理解为,它是一个更柔和版本的 Max,也就是更柔和的取最大值过程。对于最大值来说,最大的一项是 1,其他项全部是 0。Softmax 则可能让最大的一项是 0.9,第二大的是 0.09,第三大的是 0.009。总之,它是一个会突出最大值的过程。
程曼祺 Manqi Cheng
这个值代表相关性吗?
傅天予
对,这个值就是刚才说的相关性,一般叫 Attention Score,或者 Attention Weight。
第二个原因,从语言学角度看,虽然今天的播客可能有两个小时,但朝军和我之间产生关联的内容是非常稀疏的。前面可能有 100 万个词,但我也许只和“朝军”这一个词有关系,所以关联本身就是稀疏的。
从人脑角度看,大家也发现,人脑中的神经元并不是每个都和所有神经元连接,连接本身也非常稀疏,而且可能随着年龄增长变得越来越稀疏。
所以无论是数学、语言学,还是神经科学方面的启发,大家都认为稀疏 Attention 是存在的,并且可以利用它来提高效率。
肖朝军
总结一下,一个问题是存储会带来很大的需求;另一个问题是计算复杂度会随着文本变长而增加。因为是 N²,也就是平方级增长,所以大家想寻找各种办法优化计算效率。
4. The Attention Optimization Stack
程曼祺 Manqi Cheng
现在比较主流的改进思路有哪些?NSA、MOBA,以及你们做的 MOA,都属于稀疏注意力这个大方向。总体来说,它有哪些分类?
肖朝军
我认为现在的主流分类,主要是稀疏注意力和 RNN。RNN 又复活了,以 Mamba 为例的这些模型,也可以统称为 RNN 模型。
现在最主流的方向,一类是 Transformer 的改进,也就是稀疏 Attention;另一类是 RNN,直接把 Attention 替换掉,换回原来那种计算效率比较高、但带有遗忘机制的层。
傅天予
我个人觉得,现在大家不太愿意说自己是 RNN,更多会说是 Linear Attention,也就是线性注意力。
有些工作确实还是以循环的方式进行;也有一部分工作虽然存在把很长序列压缩成记忆的过程,但这种压缩不是逐个词进行,而是给定一段很长的输入后,一次性把它压缩变小。这部分也可能被归为线性注意力。
所以总体上可以分成稀疏方法和线性注意力方法。线性注意力中,有一部分思路和 RNN 类似。
程曼祺 Manqi Cheng
可以说稀疏注意力更主流吗?因为我看到很多成果都和稀疏注意力有关。
肖朝军
在 Mamba 之前,稀疏注意力更主流。但 Mamba 之后,以 Mamba 为例的线性注意力机制改进也非常多,所以很难说哪个更主流。
程曼祺 Manqi Cheng
Mamba 是 Tri Dao 和 Albert Gu 合作完成的论文,对吗?他们也是 FlashAttention 的作者?
肖朝军
对。Tri Dao 是 FlashAttention 的作者之一。FlashAttention 属于刚才说的哪一类?
它是底层算子的改进,不涉及模型本身的改进,相当于底层算子的精进,和模型是割裂开的。FlashAttention 本质上和 Full Attention 在计算层面完全一致,只是让 Full Attention 计算得更快。
程曼祺 Manqi Cheng
除了 Attention 机制的改进,还有哪些努力可以改进长文本的计算效率?
肖朝军
Attention 机制的改进,处在算法层。对一个用户来说,使用大模型时大概率会接触到三个层次:算法、系统和硬件。
硬件层很好理解,用更快的卡、更快的 GPU,当然可以变快。系统层可能稍微难理解一点。刚才提到的 FlashAttention 就是系统层的改进。
GPU 的峰值算力和峰值显存是固定的,但代码写得好不好,决定了这些算力能发挥出百分之多少。系统层的工作,就是让算力利用率尽量接近 100%,或者让显存带宽的利用率尽量接近 100%。这部分也能提供很大的优化。
比如 FlashAttention 的提出,让 Attention 的计算效率上了一个台阶,同时显存占用至少降低了一个数量级。
再往上是算法层。算法层还可以细分为模型层和更普遍意义上的算法。模型层只考虑模型本身,比如稀疏注意力,或者混合专家系统,也就是 Mixture of Experts,简称 MoE。
更普遍的算法改进,是在模型已经固定之后,去看整体生成范式。现在大模型是一个词一个词往外生成,像打字机一样。有人就会想,能不能两个词两个词地生成,或者一段话一段话地生成;也可以让 10 个大模型各写一段,最后写 10 段所需的总时间就会减少。这些都是不同层次的改进。
程曼祺 Manqi Cheng
因为你在汪玉老师的团队,电子系本来做硬件和系统的会多一些。在这方面,你有什么经验可以分享?
傅天予
系统层现在确实有很多人关注。市面上有很多大模型推理系统,比如 vLLM,FlashAttention 也属于其中一种优化方向。
有些工作专注于优化大模型中的某个组件;还有一些比较新的方向,是针对一个大模型服务很多用户的场景,考虑如何让这些用户的延迟都不要太高,同时让整个系统的吞吐量比较高。
面向不同场景,系统优化的目标和最终结果也会不一样。
肖朝军
算法层有很多改进,是因为长文本问题很多。不仅有计算和存储问题,具体场景不同,优化方向也会不同。
比如 DeepSeek 的 MLA,也就是他们提出的一套新的注意力机制算法,是在 DeepSeek-V2 时提出的。它解决的是什么问题?刚才提到,Transformer 需要把所有词元都存下来,存储量很大;存储量大之后,每次访问它的时间也会很长,需要一个个把它读出来。
可以直观地理解为,从电脑里读取一个实际文件和读取一个 1 MB 的文件,速度肯定有差异。所以 DeepSeek 想把存储维度压低。刚才说一个词可能需要几千个数字来存储,那能不能把它压缩成 500 个数字?MLA 就做了这方面的改进。
所以改进方向很多,要看具体场景是什么。
程曼祺 Manqi Cheng
你们了解之前 Kimi 做的 Mooncake 吗?Mooncake 是什么方向上的改进?
肖朝军
我理解它整体上也是服务长文本的,但更多应该是系统层的改进。它在算法上仍然利用了比较多的 Full Attention,主要是在系统层做优化。
接下来,我们可以以 NSA 和 MOBA 这两个近期受到关注的注意力机制改进为例,更详细地展开聊聊。
在正式聊之前,我有一个比较好奇的问题。我看到 DeepSeek 的 NSA 是先发布在论文平台上,还没有发布 GitHub 代码;Moonshot 的 MOBA 则是先在 GitHub 主页上放出了项目页,也把工程代码放了出来。大家向技术社区释放成果时,是先发论文,还是先发 GitHub?这有什么区别?
傅天予
这在开源层次上是有区别的。比如大家会说 DeepSeek 开源,当然它最近可能也会进一步开源。DeepSeek 开源的是模型,也就是模型参数。
但如果业界想低成本复刻它的训练过程和推理过程,还需要代码,尤其是和硬件适配的代码。这些代码本身不会是一篇论文,而是一份代码,所以可能需要放在 GitHub 上。
现在 DeepSeek 还没有放出来,可能过段时间会发布部署代码。放 GitHub 主要是为了把代码开源。当然放了 GitHub,也可以不开源代码,只把模型放在上面。
论文更多是一种技术摘要和精要,告诉你这个东西怎么做;具体每个细节,还是需要通过代码呈现。所以这是两个层次的问题。
至于 Moonshot 的 MOBA,我没有仔细看它是否把训练代码和底层代码都开源了。
程曼祺 Manqi Cheng
接下来我们具体看一下 NSA 和 MOBA 的改进。
我看了这两篇论文,也看了你们之前的 InfLLM 和 MOA 论文。技术论文大致结构相似:前面的 Introduction 和最后的 Conclusion,会总结论文对注意力机制改进带来的核心特点,以及性能和效率提升;中间的 Method 部分会讲具体的设计架构。
我觉得实验部分也很重要,但之前讨论得比较少:究竟用哪些 Benchmark 来评测改进效果?我们可以按照这些部分展开。
先从 Introduction 和 Conclusion 开始。两位看到 NSA 和 MOBA 的成果后,觉得它们在开头介绍和整体结论里,重点和亮点是什么?
5. Sparse Training Changes The Game
傅天予
我觉得最有意思的是,它们都做了稀疏预训练。这是我们之前很少看到的。就像刚才说的,之前大部分文章都是稠密训练,只有推理时才稀疏。
稠密训练虽然也会存在一定的稀疏性,但如果训练和推理在某些地方不一致,就不可避免会引入误差。所以学术界和业界都非常好奇:如果在训练中,甚至在早期预训练阶段就引入稀疏性,会对模型最终效果产生什么影响?
之前之所以没有这样的探索,一方面可能是大家更追求模型足够聪明,要把数学题做好;效率问题则放到后面解决。因此大部分人更愿意选择保险的稠密训练和稠密注意力,确保模型性能上限不受影响。
但到了现在,模型效率变得越来越重要,尤其是需要服务大量用户时。业界开始投入更多资源尝试稀疏注意力。这其实也在回答大家一直关心的问题:稀疏注意力经过充分训练后,能不能达到和稠密注意力一样的性能上限?
至少从 NSA 论文呈现的效果看,它可以达到,甚至在某些情况下超过稠密注意力。我觉得这是很有意思的观点,也相当于给稀疏注意力领域打了一针强心剂。
以前审稿人可能会质疑:虽然你变快了很多,但也变笨了一点,这一点能不能接受?如果我就是想要一个很聪明的模型怎么办?稀疏训练的工作告诉大家,你可以在保持聪明的同时也很快。
程曼祺 Manqi Cheng
我看论文时注意到,它引用了你们之前的 InfLLM。引用的位置正是在比较它和那些主要在推理阶段做稀疏注意力的工作。它确实有一个区别,就是在预训练阶段就引入了稀疏性。
肖朝军
虽然看起来只是一个在推理阶段、一个在预训练阶段,但这个差距其实很大。一个天然原因是硬件。GPU 本身不太适合做稀疏计算,GPU 很适合做密集计算。
假设我们要把稀疏机制引入预训练阶段,当然可以做,但有可能根本实现不了加速。要在预训练阶段做稀疏计算,整个过程会比较困难。
为什么训练阶段的加速,以及稀疏注意力可训练这件事,现在变得越来越重要?核心还是因为 DeepSeek 的发展路径,以及 R1 的出现。
R1 有很长的思维链,这意味着让思维链变得越来越长,已经成为让模型越来越聪明的必然趋势。之前讨论长文本,大家更多关注的是输入很长,而不是输出很长。
输入很长,我们叫 Prefilling,这是 Moonshot 那篇工作关注的核心,也是 Kimi 做产品时一个很重要的特性:可以输入很长的内容,把一篇长篇小说输入进去,也能很快处理。
但现在 DeepSeek 可能更关注输出很长。NSA 在输入很长时也能加速,但我觉得这篇论文最重要的点,是在输出很长时也能够实现很好的加速。
为什么这件事需要 Trainable?因为 RL 阶段需要迭代训练:先让模型输出很长的内容,再用强化学习更新模型。R1 有一张图显示,随着 RL Training Step 增长,输出长度会越来越长。
可以预想,如果 R1 想继续训练,长度就会不断突破。在这个场景下,把稀疏注意力引入训练阶段,尤其是 RL 阶段,就非常重要。虽然 NSA 论文里还没有做 RL 阶段的训练,但我觉得这一定是未来趋势。
所以,虽然 NSA、MOBA,以及我们之前做的很多工作,都属于稀疏注意力,但到了今天这个时间点,情况已经发生变化。长 CoT 成为非常重要、而之前相对被忽视的问题。
程曼祺 Manqi Cheng
所以长思维链也可以理解为长文本的一种,而且现在相比很长的输入,处理很长的输出也变得非常重要,这和推理模型范式的变化有关。
我看 NSA 和 MOBA 的 Introduction 部分都强调了 Trainable。Trainable 是说可以在训练阶段使用,还是说稀疏性可以学习到?
肖朝军
稀疏性本身一定是天然存在的,但 Trainable 有两个好处。
第一,如果训练阶段就是稀疏的,那么训练时和测试、部署时的注意力使用模式就没有 Gap,也就是没有区别。这样效果一定会更好。
第二,Trainable 的稀疏性能够反过来帮助训练本身,让训练速度提升。这也是非常关键的一点。
6. Dynamic Sparsity Meets Hardware
程曼祺 Manqi Cheng
NSA 和 MOBA 也都强调自己是动态稀疏注意力。可以展开讲讲吗?是不是此前还有静态稀疏注意力?
肖朝军
从大的角度来说,静态通常意味着效率更高,但效果稍差;动态意味着效率稍低,但效果更好。
大模型处理文本时,输入里有两个东西:一个是每个词是什么,另一个是这些词的位置。位置是静态的。如果某个注意力头和位置强绑定,比如总是关注句子开头的几个词,因为开头可能是用户指令或系统设置,那么这个注意力头表现出来的模式就偏静态。
如果注意力头更多和输入的具体内容相关,就会表现出动态性,因为无法预先知道输入内容是什么。
静态注意力可以很好地保留和位置相关的注意力,但如果注意力真正和输入内容相关,就不可避免会损失一些信息。动态注意力则可以同时保留两者,所以效果通常会好一些。
但问题是,在拿到输入的具体内容之前,不知道要看哪里,也不知道要关注哪里。这会给 GPU 运行带来很大的控制开销,导致效率较低。所以这是一种权衡:动态注意力效果好,但慢;静态注意力快,但效果可能差一些。
傅天予
可以补充一个静态的例子。NSA 里面有滑动窗口,滑动窗口就是一种静态注意力机制。
静态稀疏注意力预先规定了,每处理一个词时,要和哪些词做相关性计算。比如滑动窗口只看前面的 512 个 Token,我不需要动态决定哪些 Token 和当前词相关,只要向前看 512 个就可以。
这样做很快,因为不需要知道当前词是什么,就可以提前把前面 512 个词取出来准备好。但如果是“我是肖朝军”这个例子,模型可能需要动态关注很远之前的内容;静态窗口滑不到那里,就关注不到。
动态注意力则必须知道当前词是什么,还要判断它和哪些词相关,整个过程会复杂很多。再加上 GPU 的特性,动态稀疏操作本身就不太友好,所以动态注意力效果好但慢,静态注意力快但效果可能差一些。
MOBA 其实也有静态部分。它固定选择当前词元所在 Block 的 Attention,也就是某种意义上,当前词附近的一段内容必须参与注意力计算。比如临近的 512 个 Token 必须参加 Attention,这也是静态部分。
所以更准确地说,它们都是动态和静态思路的混合。
程曼祺 Manqi Cheng
如果选择动态注意力,就需要针对 GPU 的计算特性做优化或适配。NSA 无论是标题还是 Introduction,都特别强调 Hardware-Aligned,也就是和硬件联动优化。这个可以解释一下吗?
肖朝军
要解释这个问题,可能需要先说 GPU 是怎么计算的。
GPU 的一个特性是 SIMD,也就是 Single Instruction, Multiple Data:给出一条指令,在不同数据上做同样的操作。比如让第一个数字的第一个位置乘以第二个数字的第一个位置。如果只是一个矩阵,GPU 可能计算得比较慢;但如果有 20 个矩阵都做同样的事情,GPU 可以并行处理,效率就很高。
所以并行度有多高,在一定程度上决定了对 GPU 是否友好。
第二个方面是内存访问。内存中的数字通常是连续存储的。如果需要访问的数据也是连续的,GPU 直接顺序访问,速度会很快。但如果数据是碎片化的,GPU 一次会取一大块数据进来,结果只有其中一小部分是需要的,效率就很低。
因此,连续的内存访问对 GPU 更友好。
了解这两点之后再看 NSA,它做了一个对 GPU 非常友好的操作:Block 化。MOBA 标题中的 B 就是 Block。它虽然是稀疏注意力,但粒度是一块一块的,要么整块都要,要么整块都不要。
这样就满足了两个条件:内存访问按块进行,计算时一整块数据做相同的操作,所以对 GPU 更友好。
我想打个广告。2024 年 2 月发表的 InfLLM,核心就是提出 Attention 的稀疏性要在 Block Level 上进行,才能对硬件友好。这是那篇文章的核心观点。
所以 NSA 和 MOBA 都是在 Block 层面做稀疏,对吗?
傅天予
对。
程曼祺 Manqi Cheng
那 DeepSeek 之前给我的解读不太准确。我之前直接问过 DeepSeek NSA 和 MOBA 的区别,它说 MOBA 是在 Block 层面做的,而 NSA 可能有更细的粒度。
肖朝军
NSA 也应该是在 Block Level 做的,只是论文里没有那么强调。它更强调的是后面为硬件做的那些优化。
即使做了 Block,要真正实现硬件上的高效优化,仍然需要很多操作。NSA 后面的硬件优化非常关键。
程曼祺 Manqi Cheng
这些硬件优化具体是靠什么代码实现的?听起来还是软件工作,对吗?
肖朝军
对,是写软件、写算子,也就是系统层的软件。
如果使用英伟达 GPU,通常会用 CUDA,或者更底层的接口。NSA 里使用的是 Triton。Triton 是 OpenAI 开源的、在 CUDA 之上进一步抽象的一套接口。NSA 具体使用 Triton 编写算子。
傅天予
我补充一个有意思的事情。Triton 大概是在 GPT-2 或 GPT-3 那个时期出现的。坊间有一个不一定可靠的传闻,说 GPT 某个版本最开始可能尝试过使用稀疏注意力训练,所以有一部分人尝试做系统层适配,后来就有了 Triton。
不管这个传闻是不是真的,Triton 确实非常适合做块状稀疏计算。现在在学术界,如果想做块状稀疏计算,Triton 已经是大家非常愿意选择的工具。这是 OpenAI 开源后给社区带来的贡献。
肖朝军
不过如果继续往下做,我觉得 DeepSeek 可能还会进一步优化代码。Triton 对编写代码的人比较友好,但如果想进一步提速,可能还要深入到更底层的代码中。
OpenAI 有没有做稀疏注意力,我不知道,但它一定在长文本上做过新的工作。比如最近发布的 Deep Research 等产品,可以想象它的长文本能力一定很强。它可能没有特别强调,也没有公开具体技术,但我倾向于认为,它内部做过一些优化,只是我们不知道具体用了什么技术。
程曼祺 Manqi Cheng
Introduction 最后通常会总结具体提升了多少。关于实际效果,你们觉得哪些比较有亮点?
肖朝军
我觉得 NSA 比较亮眼的还是推理,也就是 Decoding 加速。我认为长 COT 是非常关键的未来趋势。基于这个趋势,无论在训练阶段、RL 阶段,还是直接在推理阶段,能够实现加速都非常重要。
7. The Sparse Attention Blueprint
程曼祺 Manqi Cheng
那接下来进入具体的注意力机制设计。两位看到的重点是什么?
傅天予
刚才提到过,它们都是块状的,也都是静态和动态的混合。
我个人认为,Attention Pattern,也就是注意力层应该如何和过去的 Token 做相关性计算,在这个阶段可能已经没有那么重要了。只要它是 Learnable 的,最终效果差异可能不会特别明显。
所以我更关注效率提升,尤其是在训练阶段能否加速,以及 Decoding 阶段能否实现明显加速。
算法层面的创新当然有,但我个人认为可能没有那么关键。现阶段 Transformer 真正重要的问题,是如何做 Scaling:不断把模型训练得更久、使用更多数据时,长文本效率能不能跟上。
过去我们一直讨论推理阶段的效率,但现在我认为,推理阶段的效率已经没有那么关键。以前的论文测推理效率,通常是一条数据一条数据地测;训练阶段加速则可能是一个 Batch,一次处理很多条数据,两者本质上不太一样。
我对自己之前工作的批判是,当时只能做到一些推理加速,而且没有在算子层面做优化。NSA 出来后对我最大的冲击是,我当时也想把方法扩展到训练阶段,但没有意识到长 COT 会如此重要。
我当时只是觉得预训练开销很大,能不能把这部分开销降下来。现在我感觉,更重要的是在 RL 阶段解决长文本问题,这可能是未来大模型持续 Scaling 的关键。
所以我现在关注的是,论文中的理论加速比能不能真正落地。他们两篇论文做得比较好,理论加速比在 Prefilling 和 Decoding 阶段都和实际结果比较接近。我们之前的 InfLLM 理论加速比可能比较高,但实际加速还不够好。
接下来还需要大家搭建架构、实际测试,验证这两个论文中的理论加速比能不能落到实处。
程曼祺 Manqi Cheng
你们已经开始做类似工作了吗?我看到知乎上有人说 NSA 已经被复现了。
肖朝军
我也看到有人很快就把底层算子搭了一遍。不过最终效果怎么样,还需要验证。理论上的复现和实际使用中的落地,还是两回事。
程曼祺 Manqi Cheng
你们当时没有把方法做到训练阶段,和资源、显卡数量有关吗?从训练阶段就加入稀疏注意力,对投入要求是不是比较高?
肖朝军
肯定有关系,但有两个方面。
一方面是资源投入。学界做预训练本身开销很大,不过我觉得这只是较小的一方面。更重要的是我个人的原因:我一直在做算法,但不太懂硬件,也不太懂底层算子。
当时想把方法落到训练阶段,也和别人讨论过,但我会觉得稀疏计算天然不适合 GPU,最后这个想法就搁置了。看到 NSA 之后,我意识到这其实是自己认知的局限:不懂硬件、不懂底层,所以想法到那里就结束了;但他们真的把它推进了下去。
程曼祺 Manqi Cheng
你当时没有和汪玉老师的实验室讨论吗?
肖朝军
当时没有。InfLLM 是 2024 年 2 月发表的,真正开始做是在 2023 年。那时长文本还没有受到这么多关注,国内开源模型的上下文长度基本还在 8K,比较长的可能是 32K。
现在开源模型已经有 128K,甚至 1M,情况完全不一样。当时长文本没有受到那么多关注,也没有那么多同行可以交流,而且我当时也还不认识天予,否则可以和他讨论。
傅天予
从我的角度来说,学术界还有一个问题:如果要做训练阶段的稀疏注意力,并向大家证明它好用,最重要的是要让它和稠密模型一样强。
但工业界把稠密模型训练得很强所使用的数据,学术界可能没有;训练经验和一些小技巧,学术界可能也不知道。你当然可以用相同的数据和训练方法,重新训练一个相对弱的稠密模型和一个相对弱的稀疏模型,稀疏模型可能稍微聪明一点。
这样好像也能说明问题,但在工业界的认可度可能没有那么高,而且验证成本也会非常高。
所以我们当时做 MOBA 时,最大的出发点是不要训练。我希望它是一个即插即用的方法,任何人训练完一个稠密注意力模型,都可以直接插上我们的模块使用。
还有一点是,我不希望改变原模型的偏好。一旦训练,可能会改变模型的回答风格。比如它本来很礼貌,训练之后突然变得很机灵,这不是我们希望看到的。
所以在做那项工作时,我们是故意拿掉训练部分的。
肖朝军
我们当时最后呈现出来的论文标题里,也有一个非常明显的词:Training-Free,也就是不需要训练。
我们还是受限于当时的认知,没有考虑到在训练阶段引入稀疏性是关键。当时模型训练长文本通常分两个阶段:第一个阶段训练短文本,第二个阶段再用少量数据把上下文长度扩展起来。
那时可能会觉得,第二个阶段的开销没有那么大,先在测试阶段把能力和性能提升上去,也许是更关键的问题。但现在看,这还是受限于当时的认知。训练阶段引入稀疏性不仅效果可能更好,现在也确实变得很有必要。
程曼祺 Manqi Cheng
具体到 NSA 和 MOBA,它们的注意力机制到底是怎么工作的?两者有什么区别?我简单理解,NSA 好像有三个机制,MOBA 好像是一个机制。
傅天予
共同点大概有两个。
第一个是,它们都会在更高层次上选择需要关注的过去信息,也就是过去的上下文。所谓“一块”,就是 Block Level。具体选择方法可能有细微差别,但总的来说,要先选择关注哪一块,再把这块内容拿进来,关注其中的细节。
第二个是,它们都会关注离当前生成词最近的词。这很自然,因为稠密 Attention 天然就会比较关注临近的词。从语言学角度说,当前词要说什么,影响最大的可能是刚才、十分钟前说的话,而不是一小时之前的内容。
还有一个共同特点是,对于每个 Query,它选到的 KV Block 都不一样。也就是说,当前正在生成哪个词,会决定它选择不同的上下文。
肖朝军
可以直观地理解为,先把文本上下文切成 Block,每个 Block 做一次相关性计算,再把最相关的 Block 选出来,之后进行细粒度的 Attention。整体思路就是这样。
具体怎么选,以及如何表示一个 Block、如何让这个 Block 和上文做相关性计算,可能会有细微差别。
所以思想上已经比较固定了。之前学界的很多工作,包括我的 InfLLM,基本都是这个思路。但关键还是把它真正落地,把加速比落实到真实系统里。
程曼祺 Manqi Cheng
NSA 论文里看起来有三个筛选机制,MOBA 是一个筛选机制。是不是一个更简单、一个更好?还是没有这样的推导?
肖朝军
NSA 应该也是两个。虽然它画了三个图,但那三个图实际上对应两个机制:一个是刚才提到的滑动窗口,另一个是选择远距离的 Block,也就是判断哪个 Block 和当前词相关。
MOBA 我理解也是这两个机制,只是图的呈现方式有些区别。
8. Benchmarking Long COT
程曼祺 Manqi Cheng
接下来讨论实验。做这种改进肯定需要实验,这次也可以从 NSA 和 MOBA 选择的 Benchmark 看看,应该怎样评判注意力机制改进得好不好。
你们比较关注哪些评测结果?
肖朝军
我最关注 NSA 在长推理数学题上的效果,反而是那个小表格。Prefilling、“大海捞针”等长文本常用测试集,按预期应该都不会差。
我更关注训练曲线的下降情况,也就是 NSA 的 Figure 4,以及 MOBA 的 Figure 3。训练曲线能告诉我们,稀疏注意力训练到底靠不靠谱,投入足够多资源后,能不能和稠密注意力一样聪明。
推理阶段的能力总是可以提高。如果上不去,把稀疏度降低一点,让它更接近稠密模型,通常就能逐渐靠近。
傅天予
本质上这是一个权衡。如果考虑性能和效率的平衡,我们会关注 Pareto Frontier,也就是帕累托前沿:和我一样快的模型,是不是最聪明的;和我一样聪明的模型,是不是最快的。
如果只看聪明程度,总能往稠密方向靠一点;如果只看效率,也可以调整稀疏程度。所以我们没有那么在意单独的性能指标。
但 Training Curve 会告诉大家,稀疏注意力训练到底靠不靠谱,投入足够多资源后,能不能达到和稠密注意力一样的能力。这是我们之前不知道的,而这两篇工作都给出了比较积极的答案:大概率能够最终趋于一致,甚至稀疏模型会更好。
程曼祺 Manqi Cheng
所以你们更关注效率测试,而不是性能测试?
肖朝军
性能如果和 Full Attention 拉不开差距,能够 Comparable,就已经达到目标了。关键是在不损失性能的前提下,能做得多快。
程曼祺 Manqi Cheng
你们自己做实验时,接下来会重点关注什么?
肖朝军
核心还是看改进了什么。比如 Attention 改进的是访存,那就可能要和 MLA 比,在通用生成、Decoding、Prefilling 等阶段都进行比较。
如果做的是稀疏 Attention,现在自然要和 NSA 比。但如果说 Benchmark,我仍然坚持刚才的观点:重点是长 COT 的生成,以及在这个场景下的效果和速度。
程曼祺 Manqi Cheng
长 COT 生成现在有什么比较好的 Benchmark?
傅天予
如果不关注长 COT 本身的质量,只关注它能不能给模型带来增益,之前的 Reasoning Benchmark 就可以。
程曼祺 Manqi Cheng
只关心质量、不关心是否给模型带来增益,具体是什么意思?
傅天予
只关心质量,就是看最后的能力。最简单的例子是 R1-Zero。论文里也提到,R1-Zero 的思维链不可读,甚至会有很多混合语言。
稀疏注意力可能会改变它的推理过程,让它变得更不可读。当然这只是一种可能。如果你希望长思维链是可读的,Benchmark 就需要调整。
假设你认为长思维链只是一个过程,那可能仍然只需要测最终输出的效果。当思维链不可读时,怎么判断一条长思维链好不好,就取决于你是否关注可读性。
如果关注,就需要测试稀疏注意力对可读性的影响;如果不关注,那就无所谓。
程曼祺 Manqi Cheng
现在还没有这样的 Benchmark 吗?
傅天予
可读性现在应该还没有统一指标。一个 Benchmark 可能由数据和评测指标两部分构成,现在可能有相关数据,但没有成熟的指标。
程曼祺 Manqi Cheng
我想追问一下关于长文本 Benchmark 的问题。我感觉在长文本领域,不同阶段关注的东西不太一样。
我自己喜欢把模型能力分成两点:能不能说人话,能不能说真话。说人话,是指模型不要把不同语言混在一起,不要输出奇怪的字符,至少要输出完整句子。
StreamingLLM 那批早期大模型稀疏注意力工作,最早尝试解决的就是这个问题:当输入长度超出模型训练时见过的长度,怎么保证模型还能顺畅地说人话?
这类问题通常看 Perplexity,或者 Training Loss,也就是看模型生成的内容和人类写的内容之间的匹配程度。当然,不匹配不代表说得不对,模型可能说得比人类更好,但一旦不匹配,分数就会下降。
第二阶段是让模型说真话。比如给它一段很长的上下文:A 的爸爸是 B,B 的爸爸是 C,C 的爸爸是 D,然后问 A 的曾祖父是谁,看模型能不能答对。
这类评测的指标就会变化。不管模型用什么形式说出答案,是大写、小写、名字还是姓,只要能找对人,就算回答正确,通常会用准确率这样的指标。
上一个阶段首先要让模型会说人话,之后才能说真话。现在大家更关注的,可能就是模型的实际能力。再下一个阶段,可能会关注思维链能力:并不是模型思考过程中的每句话都有用,想什么都可以,只要最后得到正确答案就行。
程曼祺 Manqi Cheng
MOBA 做了消融实验,也放了一些消融实验结果;NSA 似乎没有做。做消融实验一般是为了逐个验证机制中的某个部分哪个最有效,对吗?
傅天予
对,消融实验就是一个个把机制去掉之后,看效果会怎么样。
程曼祺 Manqi Cheng
注意力机制一般都需要做消融实验吗?
傅天予
要看实验思路。如果机制比较复杂,可能需要一个个拆开验证;如果机制很简单,就没有必要。
比如把一个简单机制拿掉后,可能就直接变成 Full Attention 了。NSA 的核心其实可以拆成两部分:一部分是底层算子加速,另一部分是稀疏注意力。
稀疏注意力只需要和 Full Attention 比效果,底层算子加速则主要比较速度。这样可以分成两张图来呈现,不一定要用消融实验的形式。
程曼祺 Manqi Cheng
NSA 里还写了很多核心设计,这算写得比较详细吗?能让社区复现吗?
肖朝军
我觉得写得比较详细,基本讲了算法流程,设计思路也写得比较多。
程曼祺 Manqi Cheng
所以你说已经有人复现出来了?
肖朝军
我看到知乎上有人复现了。社区还是很活跃的。
程曼祺 Manqi Cheng
你的意思是,知乎上复现出来的不一定是真的复现了?
肖朝军
不是这个意思。复现是否真正有效,还需要验证。现在无论是谁,包括 DeepSeek 自己,即使发布了 NSA,也需要看它能不能在实际使用中落地,不能盲目信任。
9. The Memory Problem Ahead
程曼祺 Manqi Cheng
最后总结一下注意力机制的更多尝试。再往下看,你们觉得有哪些值得关注的探索方向?
傅天予
刚才提到的长思维链,是一个方向。
另一个问题是,稀疏注意力虽然降低了计算复杂度,但存储复杂度其实没有改变,仍然需要把所有东西存下来。人脑的存储效率可能更高,不需要把过去两个小时的每一个字都存下来。
这种存储效率能不能进一步改进?改进到极致可能就是 RNN,也就是 O(1) 的固定大小存储。但固定大小存储又可能带来能力问题。
现在相当于处在两个极端:稀疏注意力把所有东西都存下来,RNN 只存固定大小的内容。随着上下文继续变长,稀疏注意力迟早会遇到存储问题。有没有可能在两个极端之间找到一种折中?
肖朝军
还可以把它看成另一个问题。假设现在模型的长 COT 只需要解决一道数学题,就已经达到大家的预期;但未来大模型的思维方式肯定要应用到实际任务中。
按照 OpenAI 的规划,下一步可能是创新。作为科学家,或者未来让大模型成为博士生,从选题、做实验到写论文,整个过程可能持续几个月,甚至一两年。
如果这段时间的所有信息都要存下来,肯定会有问题。现在一万、两万的长度可能还能处理,但未来这个问题一定需要解决。
程曼祺 Manqi Cheng
存储这件事可以从硬件上做什么优化?
傅天予
这非常难。计算和存储的发展速度不一样。计算,尤其是英伟达推动的 GPU 计算,发展速度非常快,每年算力可能翻倍,基本能够满足大模型的需求。
但显存的发展非常慢。英伟达几年前的芯片发展到现在,算力可能增加了 10 倍,存储容量可能还没有增加到 2 倍。
快速存储非常占芯片面积。现在芯片很大一部分面积都用来做快速存储。如果进一步扩大存储,就需要增大芯片面积;但芯片面积增大会让良率快速下降,成本也会变得非常高,大家可能无法接受。
程曼祺 Manqi Cheng
而且刚才朝军说的显存限制,并不涉及存储和计算之间的通信和数据搬运,对吗?
肖朝军
这和那个限制没关系。数据搬运当然会影响速度,但我刚才讨论的是存储上限,也就是本身能存多少东西。
存储本身可能比较便宜,不一定非要用显存,也可以使用计算机内存,甚至硬盘。但这样可能会带来速度问题,所以未来怎么解决,值得思考。
更直观地说,人也不需要那么多存储。我们脑子里关于一件事情,可能不会超过 100 GB,甚至不会超过 40 GB。能不能进一步接近人的机制,动态决定什么时候该存、什么东西不该存?
稀疏注意力解决的主要还是计算问题:计算可以稀疏,但存储没有改变。这可能是未来非常重要的方向。
因为如果 AGI 未来要做科研,就需要探索知识边界。我们现在做一篇研究,需要几个月甚至一年,模型显然还做不了这样的事情。
傅天予
我觉得这个问题提得非常好。未来一定会有大量文本进入模型,但是否一定要在注意力机制层面解决,不一定。
还是以做科研为例。假设我读 100 篇论文,这些论文加起来的上下文已经超级长,现有注意力机制可能处理不了。但我可以读一篇论文,让模型写一个 20 字的总结;再读下一篇,再写一个 20 字的总结。100 篇论文最后只剩 2,000 个字,我再看这些总结,决定到底要深入读哪篇论文。
这属于在注意力机制之上更高一层的改进,可能是通用算法层,或者生成范式层的改动。它们都在解决长文本问题,只是放在不同阶段。
10. Multimodal Context Gets Longer
程曼祺 Manqi Cheng
我们现在说的长文本,主要还是语言吗?未来是不是也会变成多模态?
傅天予
感觉这是必然趋势。
以 GPT-4o 为例,它是实时流式模型。假设它的记忆非常有限,只能记 10 秒钟的内容,那流式输入可能就没什么意义;但如果它能记住一天、一年的内容,就可能成为很好的助手。
记忆存储能力的提升,可能会带来一些质的改变。
我们最近也在探索多模态注意力的特点。多模态主要有两个问题。
第一,模态变多后,上下文长度会大幅增加。比如一张图片可能就需要 100 多个 Token,具体取决于分辨率,分辨率越高,Token 越多。
假设视频是每秒 24 帧,每帧 100 个 Token,那么一秒就是 2,400 个 Token,10 秒就是 24,000 个 Token。如果看一分钟、一小时,数量还会继续上升。
现在通常通过抽帧解决,比如一秒一帧,甚至 10 秒一帧。但这样效果会很差,因为 10 秒只看一张图,就没有视频连续的流式信息了。
对标人类来看,人一个小时大概能读 18K 个 Token;听一个小时音频,如果转换成模型输入,大约有 90K;看一个小时视频,即使按一秒一帧,转换成模型输入也可能有 1M Token。
一个小时的视频,在人类看来就是半部电影,并不算长,但对应的 Token 数量可能是 100 万。这个数量大约相当于《哈利·波特》全集的文本量,所以差别非常大。
第二,不同模态会带来不同的稀疏模式。对于文本,最常见的稀疏模式是关注上一个词或临近几个词。
但对于图片和视频,可能需要关注不同帧中的相同位置。比如要看一个球随时间如何移动,就需要关注不同帧中大致相同的位置。这种注意力模式不是一直看相邻 Token,而是每隔 220 个 Token 跳过去看一次。
不同模式也会对注意力设计产生不同影响,对计算特性的需求也不一样。
程曼祺 Manqi Cheng
学术界或工业界有没有相关的新成果?
傅天予
有,比较早就有人研究 KV Compression、KV Eviction 等方向。
比如视频的相邻帧内容基本相同,只发生一些微小变化,因此有些 Token 可以合并,甚至可以丢掉。
我们最近也开源了一个项目,叫作 FrameFusion。它虽然不是直接处理注意力,但同样希望解决长上下文下的视频理解问题。
现在注意力中的丢弃,更多是从重要性角度出发:不重要的地方可以丢掉。但视频里有些内容虽然重要,却在每一帧都出现,不需要重复看这么多次。
所以 FrameFusion 更多是从冗余性出发,只保留重要且独特的部分,把重复部分和不重要的部分都丢掉。
程曼祺 Manqi Cheng
除了长思维链和多模态,还有什么注意力优化方向或趋势?
肖朝军
刚才也提到存储压缩。
另外还有一些比较小的方向,比如稀疏注意力能不能在硬件层面继续优化;或者更多利用计算机内存,做显存和内存之间的数据交换,也就是 Offloading。
还可以把稀疏 Attention 和更多推理加速算法结合起来,比如投机采样。但这些可能更多是工程层面的细节。
傅天予
我觉得可能还是要面向 AGI,思考有哪些本质问题需要改变。
在注意力机制这个问题上,AGI 的大问题可能没有那么多。单纯从注意力机制来说,刚才提到的存储问题可能比较关键。
如果只是做思想实验和实验验证,让模型做科研,它已经有很强的思考和逻辑推理能力。但从注意力机制来说,记忆和存储可能仍然是核心问题。具体怎么做,可能是一个很大的挑战。
程曼祺 Manqi Cheng
从 NSA、MoBA 这样的稀疏注意力,到我们理想中的注意力机制,还有多大差距?
傅天予
我们的理想总是非常理想。
稀疏注意力这个大方向,对注意力的改动相对没有那么激进。相比线性注意力,它的改动比较小,也能和现有很多方法直接结合。
线性注意力描绘了一个更美好的蓝图。从 Scaling Law 的角度看,稀疏注意力只是减少了一个常数比例,序列越来越长时,最终还是会爆掉。
线性注意力则是说,增长速度比稀疏注意力慢。比如在 500 个 Token、长度比较短时,它可能比稀疏注意力慢一点;但长度无限增长时,它的速度增长会更慢。
所以理想状态可能是线性注意力,甚至比线性更低复杂度的方法。但现在学术界尝试较多的仍然是稀疏注意力,因为很多研究证明,纯线性注意力在效果上还不够好,需要和其他方法混合。
肖朝军
我觉得注意力机制的更多尝试,可能还要再往上一层,围绕长文本和长序列记忆来做。
注意力本身是否需要改,可能不是最核心的问题。关键还是从能力出发思考:长文本、长序列和记忆到底需要什么能力。
程曼祺 Manqi Cheng
记忆和长文本有什么区别?它们的相关性很明显。
肖朝军
我认为记忆可能是一种新的架构层面的改进。
长文本是输入:我要处理很长的序列。注意力机制是输入进来之后,模型要怎么处理。
记忆则更多是说,长序列输入进来后,模型不断接受新的输入,需要选择哪些东西要记住、哪些东西不要记住,然后把要记住的内容和当前输入一起交给 Attention 或 Transformer 处理。
所以注意力机制关注的是记忆和输入来了之后怎么处理;记忆本身更关注哪些东西该记、哪些东西不该记;长序列则是整个输入。
傅天予
汪玉老师特别喜欢举一个例子:Y 等于 F(X)。神经网络就是 Y 等于 F(X),X 是输入,Y 是输出,F 是网络本身。
现在有两个趋势:X 越来越长,这是多模态可能带来的;Y 越来越长,这是思维链可能带来的。
F 要解决什么问题?第一,X 和 Y 都变长时,怎么计算得快;第二,在输入和输出都很长的情况下,怎么正确地处理 X,正确地输出想要的 Y。
这就是大家努力的总体目标。至于 F 如何实现,注意力改进只是方法之一,也有其他方法能够达到同样的目的。
程曼祺 Manqi Cheng
这个解释很简单,中学数学也能听懂。
11. Research Becomes AGI's Test
程曼祺 Manqi Cheng
当记忆能力,以及和它相关的长文本推理解决得比较好之后,会解锁哪些新的研发方向或应用方向?
傅天予
记忆本身是架构层面的改进,但架构层面还是要服务于实际能力需求。
我们现在讨论稀疏 Attention,是因为长 COT 未来一定会越来越长,所以长文本问题必须解决。记忆也是一样。
假设有无限大的显存、没有访存时间、无限大的算力,那就不需要考虑记忆,全部存下来就好了。架构层面的问题,本质上是效率问题;但效率问题一定服务于功能问题。
未来如果想让模型做科研,记忆问题肯定要解决。现在的长 COT 够不够让模型做科研?显然还不够。这是两个层面:架构层面解决效率,学习机制层面解决认知能力。
肖朝军
我博士开题的题目和长序列有关,当时整理过长序列能力真正支持之后可能带来的价值,可以归纳为模态、知识、情感和能力。
模态,就是模型可以处理各种需要更长输入的丰富模态。情感,是希望模型像人一样拥有一生的记忆。我和你聊天时,它始终知道很久以前发生过什么。
知识,是让模型类似于搜索引擎:拥有整个互联网的记忆,无论问什么问题,都能给出准确、最新、相关和个性化的回答。
能力,就是刚才说的长 COT 能力,能够深入思考,给出更加合理的答案。
这四个方面都需要长文本和长序列能力。但长序列只是组成部分之一,真正实现这些能力,还需要模型本身进一步精进。
程曼祺 Manqi Cheng
所以大家现在看起来都在卷 Attention,是因为几个事情叠加在一起:一方面需要越来越强的长文本能力,包括多模态,因为多模态转化为序列后也是很长的东西;另一方面,推理模型范式出现后,要处理更长的思维链,也需要更高效的注意力机制。
再往后,新的架构和能力变化,也都可能把 Attention 作为重要模块来改进。
肖朝军
我觉得科研能力可能已经属于 AGI 的终极能力。
科研本身非常综合。如果只讨论智能能力,不讨论情商,那么科研和新知识发现,一定是智力的最终追求。
AGI 最终可能需要做到这件事。拥有科研能力,或者探索新知识的能力,意味着 AI 能够拓展人类知识的边界。
现在 AI 拥有的知识,还是人类已经总结好的语言知识,放在语料里让它学习。但未来它要探索新的范式,探索宇宙的奥秘,就需要科研能力。
掌握科研能力之后,它还会体现出很多认知能力,比如思考能力、记忆能力和高效学习能力。所以我认为,科研能力可能是 AGI 发展到较高水平时的集中体现。
程曼祺 Manqi Cheng
在 OpenAI 的路线图里,创新者是 Agent 之后的第四个阶段,再后面是组织。
肖朝军
我不一定认同 Organization 必须在 Innovation 之后。
组织能力未必需要非常高的智能。比如蚂蚁之间也能组织、协作,这种能力可能不需要非常高的智能。
创新则是颠覆范式、探索新知识的能力。我觉得它反而一定是最高智能的生物才具备的能力。
傅天予
我觉得这可能是两个发展方向:一个是单体智能越来越强,一个是从单体智能向多体合作发展。谁先发展不一定。
我也赞同朝军刚才说的科研能力。它意味着 AI 可能具备自我改进的能力。
现在我们还没有看到 AI 能真正自己改进。大家会想,能不能让 AI 自己生成文本,再自己训练自己,越训练越聪明。但目前这样训练后模型会直接崩溃。
如果真的能做科研,就有希望实现自己迭代、自己升级,甚至可能出现超越人类智能的系统。
程曼祺 Manqi Cheng
所谓自己迭代,是它自己想到新的计算机制,然后自己写代码实现吗?
傅天予
现在还没有看到,但我觉得这应该是 AGI 的未来。
比如 AGI 把我们的工作都取代之后,总会发现自己的缺陷,那它就需要改进自己。问题是,它为什么会有这个动机?动机从哪里来?
肖朝军
现在的 AI 还没有自己的动机,动机来自人给它的指令。
在 R1 之前,AI 的动机是学习人类说话,也就是 Next Token Prediction。它最大的动机是学会人类语言,这是人给它设定的优化目标。
O1、R1 都给了一个新的目标:把题做对。不管思考过程是否和人一样,步骤是否和人一样,只要最后把题做对。
R1-Zero 对学习人类说话的要求更低。R1 会先给一些人类思考的例子进行训练,而 R1-Zero 完全不给这些例子,直接让模型自由思考,所以它的思考过程和人类差得更远,但结果可能仍然正确。
未来要看人类希望它做什么。如果希望它自我迭代,当然可以设计一个损失函数,告诉模型目标是自我迭代、变得越来越聪明。
但问题是,什么叫聪明?这很难定义。人类自己也无法清楚定义智能,只能形成一些共识,而且每个人理解的智能可能不同。
人工智能的定义到现在也有分歧。比如一支话筒可以把声音录下来,帮我完成一个任务,它是不是有智能?一般人的理解是没有。
人脸识别通常被认为是人工智能应用,但它也只是帮我完成识别人脸这个任务。为什么话筒不是智能,人脸识别是智能?人类对智能的定义本身还不完善。
所以我倾向于认为,当我们给 AGI 创造足够多的能力后,未来 AGI 的目标未必由人来决定,而可能由它自己决定。
程曼祺 Manqi Cheng
那它自己决定的目标,是人类想要的目标吗?不一定。它对人类是有好处还是坏处,好处更大还是坏处更大,都很难想象。
肖朝军
我倾向于认为,未来 AGI 可能会形成一种自组织的形式,变得像人一样。
初步阶段可能是取代一部分人的工作,但随着技能水平很高的人逐渐离岗,AGI 可能成为这个领域最专业的那批“人”。之后可能就是 AGI 带领 AGI,一群 AGI 一起完成一群工作。
人类社会会有分工,老板会给大家定目标。未来可能也是类似的过程:AGI 自己设定目标。但这个目标是不是服务于人类社会,是未知的。
傅天予
很多时候,效果上很惊艳的东西,技术上未必复杂;有些效果出现飞跃,技术上也未必真的发生了什么。
比如 O1 或 R1-Zero 的思考过程突然开始胡言乱语。从效果上看,大家会觉得奇怪;但从技术上看,这其实非常合理。
只要目标是让模型最终得到正确答案,那么它说人话还是奇怪的话、说中文还是英文,都无所谓。对这个目标来说,这在技术上很自然。
人类总觉得自己的智能高人一等,大脑的秘密难以捉摸。但我们无法确定,人类智能在机理上是否本质高于其他智能,也无法确定什么样的机理才是最优的。
我们可能不应该因为“越像人类”就认为越好。人脑提供了一个可行解,而且一定是较优解,但不一定是最优解。
肖朝军
当大家还不相信一件事能做成时,提供一个可行解或较优解,可以证明这件事能够实现。但当大家开始质疑怎样做到最好时,只提供一个可行解就不够了。
人脑的效率确实很高。人类进化受到能源限制,每天能摄入的能量有限,过去获取食物的效率也很低。控制能耗可能是生物发展的关键限制。
但对 AI 来说不一定。假设它的智能足够高,我们可以给它提供足够的算力。
程曼祺 Manqi Cheng
地球能源不也是有限的吗?这可能要靠可控核聚变。
肖朝军
那可以交给 AI 研究。让 AI 研究可控核聚变,再用可控核聚变的能源发展 AI,这也很合理。
说到底,我之所以认为科研是最重要的问题,是因为现在所有 AI 的知识和监督信号仍然来自人类。
如果未来 AI 能够真正做科研,就代表它能够自己找到知识边界,拥有自我监督能力。这是比较高级的一种形态。
如果我们在科研过程中需要哪些认知能力,那么这些能力未来也一定是 AGI 所需要的。
程曼祺 Manqi Cheng
今天非常感谢两位做客《晚点聊》。我们从注意力机制一路聊到了 AGI,这也是一个很自然的延伸。
注意力机制是 Transformer 架构中的核心机制,但它的工作原理导致最初的 Full Attention 一直存在显存和计算复杂度问题。当序列或文本变得非常长时,存储和计算就会跟不上。
后来学界做了很多注意力机制改进,都是为了在保证效果的同时,提高注意力机制的存储和计算效率。
我印象很深的一点是,不同于以往主要在模型推理阶段引入稀疏注意力,这次两位都提到,NSA 和 MoBA 的一个亮点,是它们在模型训练阶段就引入了稀疏注意力。
现在的强化学习方法,需要在训练时生成和处理很长的思维链。在这个背景下,把稀疏注意力放到训练阶段就更加有必要。
未来我们还要进行更多多模态理解,而多模态本身也会带来超长序列,因此同样需要注意力机制的改进。
今天我们主要讲了注意力机制改进中的稀疏注意力方向。另一个大的方向,是我们今天也提到过的线性注意力,也可以说是 RNN、循环神经网络的回归。这是一个理论上限很高、也不存在同样存储限制的方向。
顺便预告一下,下期节目我们可能会和 Mini Max 的研究人员详细聊聊线性注意力。Mini Max 在今年一月中旬开源的 Mini Max 零一的核心创新点,就是使用了一种新的线性注意力架构。然后今天的节目就到这里,各位再见,拜拜。本期节目就到这里,欢迎收听。如果你对今天聊的话题有观察、好奇或疑问,欢迎在评论区分享想法,这也会成为我们节目的一部分,让整个讨论更完整。你也可以把我们的节目分享给对这个话题感兴趣的朋友,欢迎推荐更多你想听的主题和嘉宾。你可以从小宇宙、苹果 Podcast 等渠道关注“晚点聊 Late Talk”,也欢迎关注我们的公众号“晚点 Late Post”,下期再见。