Mamba 作者谈 LLM 未来架构

本片文章来自【机器之心】对Mamba作者进行采访所进行的编译整理。

前言

在大模型领域，一直稳站C位的 Transformer 最近似乎有被超越的趋势。

这个挑战者就是一项名为【Mamba】的研究，其在语言、音频和基因组学等多种模态中都达到了SOTA性能。在语言建模方面，无论是预训练还是下游评估，Mamba-3B模型都优于同等规模的Transformer模型，并能与两倍于其规模的 Transformer 模型相媲美。

论文一经发表，引起了不小的轰动。惊叹之余，大家发现论文的作者只有两位，一位是卡内基梅隆大学机器学习系助理教授 Albert Gu，另一位是 Together.AI 首席科学家、普林斯顿大学计算机科学助理教授（即将上任）Tri Dao。

这项研究的一个重要创新是引入了一个名为选择性SSM（selective state space model）的架构，相比于Transformer中的自注意力机制的计算量会随着上下文长度的增加呈平方级增长，比如上下文增加32倍时，计算量可能会增长1000倍，Mamba可以随着上下文长度的增加实现线性扩展，其性能在实际数据中可提高到百万token长度序列，并实现5倍的推理吞吐量提升。而这些，都离不开选择性SSM。

看到Mamba如此优异的成绩后，很多研究者对 SSM 相关研究产生了好奇。

在最近的一次采访中，来自艾伦人工智能研究所（ AI2）的机器学习研究员 Nathan Lambert 与 Mamba 论文作者之一 Tri Dao，以及同样来自 Together.AI 的科学家 Michael Poli 进行了深入交流。

他们主要讨论了 LLM 架构的未来，此外，这三位研究者还探讨了状态空间模型（SSM）在新兴的 LLM 市场中的应用前景。对话中涉及的知识点也比较密集，比如为什么 Transformer 中的注意力机制有效、其扩展限制是什么、Mamba 介绍及其硬件优化以及对未来架构预测的讨论等。

1、为什么注意力机制有效？

Nathan Lambert：我们首先来讨论一下，为什么注意力机制有效，以及注意力机制的局限性在哪里。Transformer有多少部分是基于注意力机制构建的，还有没有其他的机制在起作用，在这方面可能会遇到哪些挑战？

Tri Dao：没错，所谓Transformer，就是目前推动我们看到的大多数激动人心的应用的架构。如你所说，注意力机制是其中的核心层。其实，注意力机制早在2014年到2015年就受到了关注，随后 Transformer 的概念出现，融合了注意力机制，注重多层感知机（MLP）和注意力机制的交织使用。

我认为它的成功很大程度上在于，这些模型似乎能够很好地进行扩展，你可以通过增加更多的参数和数据来扩大模型规模。这就是成功的秘诀。虽然现在看起来显而易见，但我认为五年前这还不是一个明确的概念。

Transformer成功的几个原因包括：首先，它足够通用，能够从大量数据中学到很多东西。其次，它对硬件非常友好。与之前的循环神经网络（RNN）不同，它没有顺序依赖性。

因此，它可以在GPU、TPU上运行得很好，可以进行扩展，同时非常高效地利用硬件。我个人也致力于让它更加高效地利用硬件。所以，这就是成功的秘诀——作一个既通用又能很好地进行扩展的架构。如果你是研究NLP的，也许你会考虑加入一些归纳偏置来增强模型。就我个人而言，我认为Transformer是一个非常通用的架构，可扩展性好，而且对硬件很友好。

Nathan Lambert：是的，是的。现在回想起来，这一切似乎都显得很明显。现在，在研究它的替代方案时，一个有趣的维度是上下文长度。Michael，你怎么看？

Michael Poil：是的，我有几点想说。首先，目前仍然有很多优秀的研究试图从第一性原理出发解释 Transformer，为什么它能学习这些有趣的回路？人们会拆解计算过程，比如不同的transformer中的 head 组合等。

有一些工作将 Transformer 作为一种编码的编程语言来理解。但我认为，正如 Trey提到的， Transformer中有一些非常有趣的设计选择，注意力和MLP的交织使用是相当重要的。而且， Transformer中一开始之所以成功，是因为它采用了一些已经为RNN和其他传统NLP模型所开发的技术，比如使用门控机制来调节模型吸收哪些信息，以及决定某些内容在这种并行形式中被忘记的速度。这就像是一些可以在GPU上进行优化的宝石，虽然不是很容易，但确实可以优化。

2、注意力计算量呈平方级增长

Nathan Lambert：是的，这些都很棒。我想说的比较具体一点是，注意力机制最终呈现出一种随输入序列长度呈二次方增长的计算成本。假设你有一个长度为L的输入序列，并且你想输出一个长度也为L的序列。如果你深入到数学细节中去，观察大多数库在推理时发生的情况，你会发现有这样一个上三角注意力矩阵，其中，你只能考虑文本的过去部分。随着处理的进行，你会发现它形成了一个L平方的关系，其中第一个token只考虑一个元素，然后每个后续的token需要考虑的过去token逐渐增多。刚才我们已经讨论了RNN，以及一些非注意力的方法是如何做到不去查看序列中所有文本历史的。当你给聊天机器人GPT写一个长长的提示时，你真的希望所有那些信息都被编码进去吗？除了这种密集的注意力矩阵，我们还有什么其他的选择？

Tri Dao：循环神经网络可以追溯到80年代，也许其中一些更著名的是长短期记忆网络（LSTM）、门控循环单元（GRU）。大约在2012年到2016年左右，它们在翻译、语音识别等方面非常流行，当时它们是NLP领域的SOTA技术。

它们以一种序列方式处理文本：逐个观察token，然后改变隐藏状态，每次看到一个新的token时都会更新隐藏状态。我认为从某种意义上来说，这模仿了人脑处理信息的方式，就像你阅读一句话或一段话，就像你在大脑中存储一些信息。当你读完一个文档时，你可能能够回答关于那个文档的问题，而无需再次参考该文档。所以，RNN就是这样工作的。它们处理文本，然后改变隐藏状态，隐藏状态是可以用来生成新token或对文档进行分类的表示。

大约在2016年左右，这些方法效果很好。但是，它们逐渐不受青睐，从实验结果来看，它们的表现并不像Transformer那么好。正如你提到的，由于Transformer存在这种二次扩展特性，每个token都会与之前的各个token进行比较，这为信息传播提供了一种非常简单的方式。我认为这是Transformer和注意力机制表现出色的部分原因。

但最近，一些新的RNN架构似乎表现得相当不错，RWKV是其中较早的一个。我非常钦佩那个项目，那是一位名叫Bo Peng 的研究者主导研发的。它似乎在以一己之力对抗Transformer统治地位，证明了RNN可以非常强大。

Nathan Lambert：是的。我之前也读过这篇论文。在技术层面上，他们尝试通过两个线性RNN来复制类似于注意力机制中的查询键值查找的东西，本质上是为了消除像特定注意力扩展这样的潜在问题。在这两个RNN拥有更好的长上下文行为和不同的实现规则。他们还训练了参数量高达140亿的模型。这也引出了我接下来想问的一些问题，包括Mamba和Striped Hyena。我们可以一个一个来聊。

3、Striped Hyena 是个什么模型？

Nathan Lambert：我进入了Together API，做了一次Mistral 和 Striped Hyena的对比测试。结果显示，Striped Hyena是一个不错的语言模型。它能够回答大多数问题，没有明显的失败模式。Michael，你怎么看待这个模型？

Michael Poli：首先我想说，这些新方法之间有一个有趣的联系。存在一种凸集，它有一个中心点，线性注意力（即没有softmax注意力）、线性RNN以及基于状态的模型（SSM）之间的关联都在这个凸集里。在某种程度上，这种基础模型的数学表述是相同的，我这里不是指基础架构，而是指基础模型。

然后你可以朝不同的方向发展，每个方向都有自己的权衡，比如特征映射方向、kernel方向。所以当你分解或去除掉softmax时，你可以处理查询（queries）和键（keys）时采取不同的方法。这些查询和键是构成你的注意力矩阵的基本实体。在去除softmax之后，你可以构建其他类似kernel的函数，或者其他你希望能够近似实现注意力机制功能的函数。

你可以做一些类似泰勒近似（Taylor approximation）或泰勒展开（Taylor expansion）这样的事情。你会得到一个略显不同的视角，但你得到的东西又非常相似。你可以转向Time variance。这意味着你修改RNN，使其计算更多地依赖于输入序列。也就是说，在线性RNN中的计算量由输入序列所决定的。你可以使用类似门控（gates）的东西，我们已经看到了跟多工作，例如，用额外的门控来更新内部tension，让你更好地利用你的固定状态维度。第三个方向——至少在我看来——是使用卷积形式的，更多地使用其他类型的线性运算符，这些运算符仍然是可结合的，仍然允许你并行训练。

所以这里的内容包括时间不变系统。我可以详细解释这些点，但是存在一些可以在卷积和循环之间切换的模型，这些模型还配备了额外的门控机制。我参与的一个项目，它诞生于我刚才提到的这第三类架构。我们真正试图做到的是，打造出每浮点运算性能最佳的架构。我们反复验证的一个原则是，似乎将不同层、不同类别的模块甚至是全注意力层结合起来，得到的东西比单独的组件要好。

因此，我们正在试图更深入地理解这些模型的组合方面。这种理解有助于我们创造每浮点运算性能更优的预训练模型。使用这种模型，我们运行了一整套的规模扩展法则。混合化也给了我们一些优势，因为我们想要的是一种即开即用的东西，它使得过程变得更加简单。

在为更长的上下文进行微调时，我们可以应用一些为 Transformer 开发的技术，而且有点出人意料的是，这些技术对于混合体也同样有效。比如，用于旋转嵌入的线性缩放等等，你可以深入了解细节。所以这主要是一个尝试性项目，试图弄清楚，在当前的环境下，我们最好能做到什么程度。

4、什么是 Mamba?

Nathan Lambert：Striped Hyena 使用了一组新的模型移植技术（model grafting）进行了优化，使我们能够在训练期间更改模型架构，对我来说，感觉好像有很多事情正在发生，比如一些你可能不能谈论太多的内容如数据。

所以，我不认为所有数据都得到了很好的解释，比如较长的上下文数据是什么，你能从模型的角度解释一下吗？我认为即使只是总结一下，对人们来说也是很好的体验。

该领域有很多很酷的工作，因此，AI领域有很多新的项目在展开，例如，有些人试图将 Lama模型拆解并继续训练。事实上，这种做法有点狂野，人们尝试采用强大的模型并试图让它们变得更小，同时还想获得媲美大模型的性能优势。

虽然这是个题外话，但我没有预料到的是当你关注社交媒体的时候，你会看到人们会说，哦，原来还是状态非注意力模型赢了。在我看来，这种说法掩盖了很多有趣的细节。

好吧，让我们回到Mamba。如果没记错的话，我认为Mamba 套件中最大的模型是2.8亿参数，在NLP基准测试种，包括GPT J、Pythia模型套件给出的基准分数都非常强。

Tria Dao：Mamba 是我与 Albert Gu合作完成的，当时他是斯坦福大学的博士生，我们就是在那里认识的，他现在是 CMU的助理教授。所以这是一次美妙的合作，Mamba的成功要归功于他。Albert一直致力于状态空间模型的研究，从某种意义上说，正如前面提到的，线性张量、线性RNN、卷积、神经网络等领域他都有涉及。

在我过去参与的几个项目中，我还致力于空间、状态空间的研究，我研究的角度是如何使状态空间更加硬件高效，并提高其性能。所以和Albert Gu一起工作真是太棒了。我认为参与Mamba的研究过程更多的是一个概念证明，状态空间实际上可以像NLP领域的 transformer 一样好吗？因此，才有了 Mamba，这项研究表明状态空间在音频方面可能会更好。但是，对于状态空间模型来说，语言模型一直是最难获得，最难做好的。

而且，语言也是人们现在最关心的事情，所以我做的更多的是一个概念验证，也就是说，我们想表明状态空间模型也可以具有竞争力，甚至可以和transformer 一较高下。我们实验中验证的 token 数从 3B 到 300B 不等。

所以从绝对意义上讲，这些都不是非常强大的模型，这些还不是我们真正想要的模型。我认为我们正在做的更多的是进行学术比较。比如训练相同数量的 token，状态空间模型可能比 transformer 稍微好一些。

这件事对我们来说尤其令人兴奋，我认为 Albert 已经推动这件事有一段时间了。

结果是，我们的研究在推理上可能会更快，也许我们会有不同的方式来理解上下文学习是如何发生的。我很期待未来的工作。

5、Mamba 硬件优化

Nathan Lambert：您能否谈一谈实现这些新 CUDA 内核实际上需要做什么？

Tri Dao：关于状态空间的研究，它在某种意义上是循环神经网络。状态大小是你在遍历或处理序列时，用来存储信息的缓冲区。

从某种意义上讲，Transformer也可以这样理解，它保存的整个历史记录通常称为 KV 缓存，Transformer 保留历史并不断引用它。对于RNN，它们具有固定大小的状态；对于 transformer，你可以认为状态大小正在增加。而且，我们的直觉是，状态规模越大，模型就表现越好。

所以基本上，你需要有更多的空间来存储你需要记住的东西。因此，之前的模型（例如S4等）具有相当大的隐式状态大小，它们使用卷积试图来避免将状态具体化。

我们希望在循环中加入更多的输入依赖性，然而，这样做阻止了我们使用可以提高效率的卷积视图。

因此，我们必须找到一种不同的方法来提高效率，所以我们把重点专注于在GPU上提高效率。我们的想法是，我们要有一个大的状态大小，但我们不需要用到实际的GPU内存，比如HBM，我们可以把大的状态保存在一个更快的内存中，称之为 SRAM，你可以把它想象成一个缓存。如果你更熟悉CPU，这通常是一个缓存和RAM。

所以，如果你有较大的状态，你可以将它保存在缓存中，这样一来，你不会遭受太大的损失。

6、2024年架构预测

Nathan Lambert：我现在对GPU 与 TPU最有力的洞察是，MoE在TPU中效果不佳，因为你必须在基本层上加入一些MoE。

在分布式训练中，前馈层可能最终分布在不同的TPU节点上，而TPU之间是通过邻近节点进行通信的。因此，与GPU相比，TPU在这方面会受到较大的影响。2024年这个领域会发生什么？

Tri Dao：我认为Transformer 仍然是一个非常强大架构，现在可以扩展到一万亿级别的参数，人们往往想要性能最好的模型，这些模型在硬件上运行效率最高，并且在软件方面拥有最多的支持。

我最近有些新想法，比如状态空间。我们已经看到，正如Michael提到的，混合这些组件似乎可以提高性能，我认为在7B大小的模型上得到了验证，也许状态空间模型可以在更大规模的模型上起作用。

当前大多数人关注的是基于 Lime 架构的数据和基础设施建设，尽管现有的Transformer架构在生产环境中仍然有非常强大且得到广泛支持，但也存在一些边缘领域，比如长上下文、音频、基因组学等，研究这些领域中的替代架构会非常有趣。这些领域提出了一些有意义的科学问题，例如模型是否像人类一样理解指令和直觉、是否能与量化方法协同工作等。

此外，即使现在人们还在使用Transformer架构，未来可能会融入更多的新思想和组件，例如增加更多的层和注意力机制，尽管它们可能仍然被称为Transformer。

简而言之，尽管目前的人工智能领域倾向于保守，专注于现代架构，但新的架构和思想正在逐渐涌现，这些新颖的观点和方法可能会为人工智能的发展带来新的动力和方向。

Michael Poli：是的，我百分之百同意 Tri Dao的看法，注意力机制作为计算原理依然很重要。注意力机制作为一种高效便捷的方式，可以有效增加序列处理器的状态容量。

状态维度与序列长度之间存在权衡关系，当模型尺寸变大，也就是模型变得更宽时，将有效地引入更多的状态和序列长度。这样一来，一些边际效应可能会消失，而某些权衡将会发生改变，尤其是对于那些非常大的模型，如 14B、30B这类模型。

未来，架构设计将变得更加有趣和复杂，将会有更多的创新发生。无论是混合模型还是引入新的模块，我们都将看到更多激动人心的创新。

7、对 AI 更多的预测

Nathan Lambert：混合专家模型（Mixture of Experts，MoE）和状态空间模型在近期内呈现出一种流行的趋势。

然而，在开源和学术界，目前还没有人真正尝试在混合专家模型上做早期的尝试和改进。模型嫁接（Model Grafting）现在正在变得更加实用。

跟进这些进展非常有趣，希望这些发展能够为学术界和科学家们提供更多影响行业对话的方法，特别是在目前行业更多关注于扩大模型规模的情况下。我建议开源公司应该在他们的语言模型中进行特定的改进，以获得商业优势。

你们在机器学习方面还关注其他什么吗？它不一定是关于状态空间模型的。明年最让你们兴奋的是什么？

Tri Dao：我个人认为数据仍然是最重要的因素。我们正在深入研究数据如何影响模型性能，例如通过一些与模型性能高度相关的合成任务来进行研究。这种方法一直是我们论文和研究工作中的主要动力和示例。未来一段时间我们将专注于数据领域。

虽然所有的架构工作都很有趣，但其在硬件上高效运行也很有趣，但最终还是关于数据的。如果你了解扩展率（scaling law），你就会知道不同的模型架构通常会有相同的斜率，只是偏移量不同，似乎唯一能改变斜率的是数据质量。

Michael Poil：是的，我们添加了数据。数据确实很有趣，比如在架构设计上的微型化，找出并分解语言建模等任务所涉及的各个方面，我们试图将它们打包成可以用来迭代的东西，这非常令人兴奋。

我个人对新应用，尤其是基因组学工作感到非常兴奋，但更多地是从工程角度出发，我们正看到一个转变。目前，语言仍是获得最多点击和最大兴趣的领域，但我认为这将随着时间的推移而改变。

Nathan Lambert：是的，每个人都在谈论语言，但我觉得图像、视频将会是产生巨大价值的东西。我不知道语言的上限在哪里。我很兴奋，我已经开始尝试这些了，比如我会拿博客的文本，让模型转换成图片，然后转换成带有音频的视频，所有这些都是用一个Python脚本完成的，这真的很容易做到。所以我同意你的观点，超越语言的东西很有趣。

Tri Dao：在你的经验中，当你将所有这些东西拼凑在一起时，它们确实能够合理地工作吗？

Nathan Lambert：它还不是那么完美，DALL·E生成的图片比较相似，但我的做法非常简单，就是直接把文本拿过来，然后用一个系统提示，让模型生成各种各样的图片，我认为我可以做的更好。据我知道，大概在一年内，就会有一个文本到视频的API，到时我会切换API，那种体验会很棒。

Tri Dao：是的，我认为这些进步确实产生了大量的经济价值，我们已经看到了这一点。现在很多公司都开始转向使用这些技术。我认为这将改变我们的工作方式，正如你提到的，改变我们的工作方式和娱乐方式。所以这是一个非常令人兴奋的未来。

原文链接：https://www.interconnects.ai/p/interviewing-tri-dao-and-michael?continueFlag=5d10d34c97637bebcfeba6470c0f0d9b