Mamba：我们所知的ChatGPT时代的终结者？

作者：精选海外教程postcode发布时间：2024-04-08

自2017年问世以来，Transformer架构已俨然成为自然语言建模（生成文本的模型）领域的“不二之选”。

ChatGPT、Gemini、Claude等，无一不是基于这一开创性架构。其影响力之深远，以至于“ChatGPT”中的“T”即代表“Transformer”。

作为一种序列到序列模型（接收一个序列作为输入，可以是文段或图像中像素的序列，然后输出另一个序列，通常是新文本），Transformer的独门秘籍在于其注意力机制。

简单来说，它对序列中所有词元进行两两间的乘法运算，使它们“对话”，以揭示词与词之间的关系。也就是说，该机制指示每个词应关注序列中的哪些词。例如，对于代词的注意力机制将使其关注相应的名词。

此机制效果显著，但代价高昂。

⭐️⭐️⭐️更多海外高质量教程： https://www.postcode.vip/ ⭐️⭐️⭐️

掌握低效之术
简言之，Transformer虽表现出令人瞩目的性能，却付出了极度低效的代价。
要理解像ChatGPT这样的模型如何处理语言，不妨想象您正在阅读哈利·波特系列小说的中途。
然而，为了成功理解下一页的内容，您必须在脑海中存储之前读过的所有页，即完整的上下文。简单来说，要读下一个单词，您必须重新阅读至此为止的所有内容。
如此反复。
实际上，Transformer并不会每次都真正地重新阅读一切，而是存储先前计算结果的缓存。
然而，这迫使它们在内存中存储所有这些计算，相比之下，这就像您为了成功阅读下一个单词，必须记住之前读过的每一个单词一样。
这就是ChatGPT等Transformer处理语言的方式。听起来非常低效，对吧？
没错，因为它确实如此。

那么人类又是如何阅读哈利·波特书籍的呢？
诚然，我们会记住此前阅读过许多内容，但会遗忘一些无关紧要的信息，比如赫敏的暑期活动。这些内容与主线剧情无关，因此我们的大脑会选择忽略。
换句话说，我们并非保留到现在为止阅读的所有信息，而是构建故事的压缩表示，保留相关数据并抹去无关细节。
用通俗的话说，我想表达的是：Transformer并不压缩语境。由于没有对语境进行压缩，随着文本序列长度增加，Transformer的计算需求会显著增大。
具体而言，Transformer在训练和推理阶段分别具有二次方和线性的时间复杂度。
也就是说，在训练阶段，序列长度翻倍会导致成本增加四倍；而在推理（执行）阶段，成本则会翻倍。

那么研究实验室又该如何应对呢？

为了避免成本失控，他们会对模型的“工作空间”设置一个上下文窗口的限制，这就是为什么ChatGPT和其他模型只能处理一定大小的文本。

近年来，人们提出了许多具有更“高效”注意力机制的不同架构，但性能损失阻碍了它们真正取代原始Transformer。

那么Mamba的研究人员做了什么呢？

简单来说，他们摒弃了过去几年来最重要的算法突破——注意力机制。进入状态化世界

Mamba架构源于一个关键问题：我们能否在保持高效率的同时，像Transformer那样有效地建模语言？答案是肯定的，这得益于我们所定义的“状态”。

回到哈利·波特的例子，当我们阅读一本书时，会不断更新对书的理解状态；我们逐步构建对书中发生事件的压缩、近似理解，保留关键要素并舍弃其余部分。

本质上，这正是Mamba所做的。

给我三个“s”

正如注意力模块位于Transformer的核心位置，选择性状态空间模型（Selective State Space Model，简称Selective SSM）位于Mamba的核心。

SSM是一种相对较新的语言建模架构，灵感来源于上世纪60年代的状态空间模型。

简单来说，模型保持一种“状态”或记忆，作为上下文使用。换句话说，下一个输出将是当前输入和到目前为止我的当前状态的函数。

如果当前输入是“Harry”，SSM将利用其状态推断下一个词很可能是“Potter”。因为之前很可能已经看到这两个词一起出现，状态会记住这一点。

然而，ChatGPT也能预测出“Potter”，对吧？

是的，但这里的关键在于，选择性SSM只保留内存中重要的上下文，这意味着它们在训练和推理阶段分别实现了线性和常数时间复杂度。

换句话说，如果序列长度翻倍，训练成本仅翻倍（Transformer则翻四倍），而无论序列多长，推理成本始终保持不变！

之所以能实现常数推理复杂度，是因为状态具有固定大小，成本不会随序列长度增加而增加，模型只需存储关键信息，遗忘其余部分。

另一方面，这类模型的预测质量取决于其压缩后的上下文质量。

那么，选择性SSM为何能选择保留哪些上下文，而Transformer做不到呢？

答案就在于“选择性”。

选或不选，这是一个问题

Mamba的独特之处在于其SSM模块具有选择性，能够选择保留哪些上下文、丢弃哪些上下文。换句话说，它通过选择实现我们所定义的压缩。

与任何其他状态空间模型一样，模型由两个方程驱动：

借助参数A、B，我们首先获得状态h（即记忆），然后通过C得到新的输出y。

您可以将这些参数视为“门”。

由于A乘以前一状态，它决定了前一状态是否对下一个预测相关。

B决定当前输入是否重要，从而计算新的状态。

最终，C决定如何将这些信息转化为输出。

关键在于，与所有先前架构不同，Mamba依赖于输入和时间。

在这种情况下，研究人员决定让A保持恒定，尽管也可以对其进行参数化。

选择性SSM为了让模型能够压缩上下文，它必须具备选择哪些内容应作为上下文、哪些不应作为上下文的能力。例如，如果下一个词是“um”，这是英国人常用的口头禅，那么它可能对上下文不太重要，因此会被忽略。为实现这一现象，如前所述，Mamba引入了一种新范式，其中模型权重取决于输入，并随时间变化。在Transformer中，权重是固定的，意味着它无法选择保留什么或不保留什么。通过将权重设为输入的变量，它们可以适应任何输入。参数∆的作用是什么？它充当主要的选择门，最终决定模型对输入x赋予多大重要性。从模型层面来看，这个SSM模块被插入到Mamba块中，并同等地堆叠起来，构建实际的模型。总之，选择性SSM是Mamba架构的核心组件，它通过动态调整权重，实现了对输入信息的筛选与压缩，使得模型能够在处理长文本时保持高效且准确的上下文理解能力，克服了传统Transformer架构在处理长序列时面临的复杂度瓶颈。这种创新的设计使得Mamba能够在保持与Transformer相当的语言建模效果的同时，显著提升了计算效率，尤其在推理阶段实现了常数时间复杂度，为大规模文本处理任务提供了更为实用且经济的解决方案。

Mamba块但如果关键仅仅是使权重依赖于输入，为什么不在早期的状态空间模型中这样做呢？状态空间模型是循环模型，意味着输入是按顺序处理的。或者换种说法，它们的计算不可并行化……除非将模型变为线性时不变（Linear-time invariant，LTI）。由于LTI状态空间模型是线性的且不随时间变化，我们可以对前面看到的两个方程应用卷积，无需显式构建隐藏状态，从而大幅度降低成本并实现计算的并行化。

借助卷积，我们无需计算h即可计算y然而，由于Mamba使模型权重依赖于输入，就失去了启用并行化的卷积能力，这对于在GPU上运行这些模型至关重要。为解决这个问题，研究人员使Mamba具备硬件感知能力，这意味着它仅将隐藏状态（由于模型无法应用卷积，必须这样做）快速地存储到SRAM级别内存中。他们实际上所做的是执行内核融合，即并行扫描而非卷积，并在训练期间反向传播期间重新计算中间状态。不深入细节，他们本质上是在应用多种优化措施，以减少I/O内存事件的数量，这反过来意味着，尽管模型是循环的，但仍然像Transformer一样高效。但真正的问题是：与Transformer相比，Mamba表现如何？令人兴奋，但仍存疑问在不同的小型规模（最多70亿个参数）上进行测试，Mamba击败了包括同类规模GPT在内的所有其他模型，无论是在困惑度（衡量模型预测下一个标记的准确性指标）还是在准确率方面。它甚至能够匹配比其大两倍以上的Transformer的结果！此外，随着长度增加，其准确率并未下降，这在此前闻所未闻：