Mamba

Note

from https://www.youtube.com/watch?v=N6Piou4oYx8

没错，经过 7 年的统治，Transformer 模型终于被取代了。不过，也许目前为止，Mamba 模型仅在几十亿参数规模的小型模型上进行了测试，但迄今为止的结果非常有希望！此外，Mamba 的计算需求比 Transformer 模型要低。对于一个由 n 个词组成的输入序列，Mamba 仅使用 O(nlog(n)) 的计算量，而 Transformer 模型需要使用 \(O(n^2)\) 的计算量。因此，基于 Mamba 的语言模型应该能够使用更大的上下文大小。在这个视频中，我们将深入探讨 Mamba 架构，包括它是什么，为什么效果这么好，以及你如何能设计出这样的架构。

通常，Mamba 被描述为某种叫做状态空间模型的扩展。状态空间模型是另一种序列模型，近几年逐渐流行起来，但说实话，状态空间模型背后的理论极其复杂，涉及一些相当高级的数学。幸运的是，Mamba 也可以被理解为循环神经网络（简称 RNN）的一种扩展，这种网络更容易理解。因此，在这个视频中，我们将通过 RNN 的路径来理解 Mamba。现在让我们开始吧：什么是循环神经网络？给定一个输入向量序列，卷积层会将神经网络应用于连续的向量组。关键在于，神经网络一次只能看到少量的向量，这使得模型易于训练。

这种方法的缺点是，距离较远的向量之间的信息无法组合，除非应用了多个卷积层。这使得卷积神经网络难以理解输入中的长期依赖性，而这种长期依赖性在自然语言文本中经常出现。为了解决这一缺陷，发明了 Transformer 架构，它成功地允许单个层组合来自任意远距离的向量信息。我之前制作了一个视频，详细解释了 Transformer 如何以及为什么有效，你可以在这里找到。

尽管 Transformer 效果很好，但它们有一个重大的局限性，那就是它们使用的计算量与输入长度呈二次方关系。对于小型输入，这不是什么大问题，但如果你想在输入中有一百万个向量，这意味着你需要进行一百万乘一百万的操作，这是一大笔计算量。

循环神经网络采取了完全不同的方法来改进卷积层。其思想非常简单：不是将神经网络应用于两个连续的输入向量，而是将其应用于一个输入向量和该神经网络的前一个输出。这看似一个小变化，但其后果深远：每个输出向量现在包含了它之前所有输入向量的信息，而不仅仅是前一个向量的信息。这最终的输出向量包含了输入中每个向量的信息，无论有多少个向量。而且我们没有使用比卷积层更多的计算。我们已经免费融入了长距离信息。这正是我们想要的。

或者至少，如果不是因为 RNN 有两个小问题使得它们在实践中几乎无法使用的话。第一个问题是，虽然循环层使用的计算量与卷积层相同，但这种计算无法在多个处理器之间并行。即使你有很多处理器，也不能开始评估神经网络的输入，直到所有之前的步骤完成，因为你需要将前一步的输出送入神经网络。与此相比，卷积层只需要看到原始输入。只要有足够的处理器，你可以同时在所有输入上运行神经网络。

由于现代硬件，如 GPU，专门用于并行计算，拥有成千上万的处理器，RNN 在实践中实际上比 CNN 慢得多。事实上，尽管 RNN 的计算量较少，但它们比 Transformer 还要慢。第二个问题是，RNN 极其难以训练。理论上，单个循环层可以融合任意多的输入信息，但实际上，它们不行。相反，它们最多只能学习融合前几十个输入的信息。RNN 的想法自 1980 年代以来已经存在，但由于这两个问题，RNN 已经不受青睐，卷积神经网络和 Transformer 在实践中更为成功。

实际上，在过去十年中，RNN 几乎根本没有被使用。直到现在。去年，一篇新论文发布，显示线性 RNN 可以避免这两个问题，因此线性 RNN 是高效的长序列模型。那么什么是线性循环神经网络呢？你只需用线性函数替换神经网络。这似乎是个坏主意，因为线性函数只能对其输入执行相对简单的变换，但我们可以通过在每个输出向量之后应用一个完整的神经网络来弥补它。

这类似于在 Transformer 中，你可以用简单的线性函数替换值神经网络，然后在自注意层之间添加神经网络来弥补缺乏非线性处理能力的问题。所以就像在 Transformer 中一样，我们将交替线性循环层和按元素的神经网络。但重要的是，通过使循环操作纯粹线性，解决了 RNN 的两个问题变得可能。

首先，我将解释如何在并行中仅用 \(O(log(n))\) 的时间计算应用于 \(n\) 个向量的线性递归。然后，我将解释如何解决困扰常规 RNN 的训练问题。

线性递归运算符由以下公式给出：要获得第 \(i\) 个输出向量，你需要用矩阵 \(W_y\) 乘以前一个 \((i-1)\) 输出向量，并加上第 \(i\) 个输入向量乘以另一个不同的矩阵 \(W_x\)。

\(W\) 矩阵中的条目是模型将学习的参数，因此它们最初是以 \(0\) 为中心的正态分布随机样本，然后通过梯度下降进行更新。由于 \(W_x\) 矩阵仅独立应用于每个输入，我们实际上可以认为它是前一层的一部分，因此我们可以简化我们的递归运算符，只需添加输入 \(x\)，假设在前一层已经应用了线性函数。

线性递归实际上是一种更通用操作称为扫描的特例，所以让我们从扫描的最简单示例开始：累积和。给定 \(n\) 个数字的输入列表，目标是计算每个术语的部分和列表。因此，输出列表中的第 \(i\) 项应该是输入列表中前 \(i\) 项的总和。虽然通过一次加一个数字简单地计算这个是很简单的，但我们希望并行执行它。结果显示我们可以这样做：首先将每一对连续的数字加在一起。然后，从结果列表中，将相隔 2 步的数字对加在一起。然后是 4 步。然后是 8 步……依此类推，每次迭代步长加倍，直到步长与整个输入列表一样大，这将在 \(log(n)\) 步后完成。这个算法之所以有效，是因为在每次迭代中，第 \(i\) 个输出元素包含了前一步长数字的总和。例如，在第一次迭代中，每个输出数字是前两项的总和。在下一次迭代中，每个项包含前 2 项的总和加上从 2 开始的前 2 项的总和，即前 4 项的总和。依此类推。当步长是输入大小时，每个输出包含所有前期项的总和，这是所需的。很容易看出，每次迭代都可以并行计算，然而不同的迭代仍然需要顺序计算，共有 \(log(n)\) 次迭代。所以，如果你有 n 个处理器，这个算法的总运行时间是 \(O(log(n))\)，从简单的顺序版本的 \(O(n)\) 减少。

这个相同的算法适用于计算任何二元运算符的累积应用列表，只要二元运算符是关联的。关联意味着你可以改变应用顺序，仍然得到相同的结果。这对加法是成立的，这就是我们的并行累积和算法有效的原因。它也适用于许多其他操作。特别是，这个二元运算符是关联的：\(f((W_1, x_1), (W_2, x_2)) = (W_1W_2, W_1x_1+x_2)\)。请注意，这个运算符使用一对矩阵和向量作为输入和输出，而不仅仅是一个数字，如加法。值得注意的是，使用这个运算符执行扫描等同于线性递归。我们首先需要将输入向量列表替换为一对列表，其中第一个元素是循环权重矩阵，第二个元素是输入向量，但然后我们像往常一样执行扫描。你可以自己检查这个运算符实际上是关联的，通过扩展其他顺序的几个术语。总之，我们只需要用这个运算符代替加法执行我们的并行累积和算法，我们就可以在仅用 O(log(n)) 时间得到线性循环层的结果。

除了一个小问题外。如果你仔细观察这个操作，其工作方式是使用元组的第一个元素作为累积矩阵，其中包含到目前为止所见所有矩阵的乘积。这就是为什么输出元组的第一个元素是两个输入矩阵的乘积。但这意味着我们在每一步都要执行一个 [d, d] 乘以 [d, d] 的矩阵乘法，其中 d 是向量的维度。这个过程非常慢。

请注意，在原始的顺序 RNN 中，我们不需要跟踪这个累积矩阵，因此我们只在每一步将权重矩阵与长度为 [d] 的输入向量相乘，这是一个 \(O(d^2)\) 操作。但现在我们每一步都必须进行一个 \(O(d^3)\) 操作。对于标准模型大小来说，这无疑是计算量增加了一千倍。这是不利的。

幸运的是，有一个解决办法：矩阵对角化。你看，几乎每个方阵都可以分解成一个可逆矩阵 \(P\)、一个对角矩阵 \(D\) 和 \(P^{-1}\) 的乘积，只要允许矩阵元素为复数即可。这里有一个例子。请注意这个中间矩阵是对角的，也就是说除了主对角线外的所有元素都是 0。

关于这一点有趣的是，当你在这种形式下将矩阵自乘时，内部的 \(P\) 逆和 \(P\) 项会相消，两个对角矩阵的乘积只是对角矩阵中元素的乘积。也就是说，为了计算 \(D^2\)，你只需要对 \(D\) 的主对角线上的元素进行平方，这只需要 \(O(m)\) 操作，而不是 \(O(m^3)\)，这要好得多。

因此，我们可以做的是用对角化形式表示循环权重矩阵，这意味着我们只需要使用一个包含 D 的主对角线元素的复数向量。也就是说，我们首先将复数矩阵 \(P\) 应用于输入向量，然后使用元素逐个相乘的方式进行线性递归，最后将 \(P^{-1}\) 应用于输出。这样计算的结果将等同于某个实数值权重矩阵 \(W\) 的线性递归。但以这种方式计算时，递归运算符只需要在两个向量之间进行元素逐个相乘来更新累积权重，而不是矩阵乘法。

当我们将这个运算符嵌入到我们的并行扫描算法中时，总计算量现在只是 \(O(dn\log(n))\)，并行运行时间是 \(O(\log(n))\)。这要好得多。

注意，这一层的参数是复数向量 \(w\) 和矩阵 \(P\) 中的复数项。实际上，你只需要使用两个单独的实数来表示每个参数的实部和虚部，这些参数通过从以 0 为中心的正态分布抽样初始化，并像往常一样用梯度下降法更新。

最后，计算矩阵逆是非常慢的，所以在实践中我们通常不会这么做，而是在线性递归前后使用两个独立的复数矩阵。这实际上使模型比实值线性 RNN 更具表现力，并节省了计算。但这确实意味着模型不再等价于实值递归，而且输出现在可能是一个复数，因此我们需要取输出的实数部分再传递到下一层。

好的，我们已经看到如何使线性 RNN 在现代硬件上运行得更快，但是关于 RNN 非常难以训练的另一个问题呢？在我们解决这个问题之前，这里简要回顾一下为什么训练 RNN 如此棘手：神经网络是通过从模型中每个权重减去损失函数的梯度来训练的。

梯度是什么？想象一下评估神经网络，然后将权重的值增加一点非常小的量，再次评估它。这两次评分的差异（成比例地）是那个权重的梯度，它告诉你如何改变权重以使神经网络变得更好。那么让我们评估一下线性循环层的梯度。为了简化这个问题，假设除了第一个输入外，之后的输入都是 0，这样我们就可以忽略它们了。

当我们评估循环层时，每一步中前一个输出都会被权重向量乘，所以经过 \(n\) 步后，输出向量等于循环权重向量的 \(n\) 次幂乘以第一个向量 \(x_1\)。当我们将权重稍微增加一点并再次评估它时，我们得到这个结果。取差异，我们得到，直到一个常数缩放因子，\(w^{(n-1) }x_1\)。

这里的问题是，随着 \(n\) 变大，这个项 \(w^{(n-1)}\)，要么变得非常小，要么变得非常大，这取决于 w 中的值是小于还是大于 1。无论哪种情况都是一个问题：如果梯度非常大，那么神经网络的权重会改变太多，神经网络已经学到的现有功能就会被破坏。如果梯度非常小，那么权重变化不足，神经网络根本学不到任何东西。

这就是训练 RNN 困难的原因，虽然理论上 RNN 可以使用无限长的上下文，但实际上，使用基于梯度的训练技术，RNN 只会学习使用与梯度保持适当大小的步数的上下文。这被称为梯度消失和梯度爆炸的问题。

当我们重新加入非零输入时，这个问题只会变得更糟，因为额外的输入使得梯度变得更不稳定。清楚地说，这对常规神经网络不是问题的原因是它们在每层使用不同的权重。一些层的权重可以小于 1，一些层的权重可以大于 1，只要梯度保持大致相同的大小，神经网络就能学习。

有很多不同的权重配置可以产生稳定的梯度，而且在整个训练过程中保持稳定的配置是很容易的。但对于 RNN 来说，你在每一步都使用相同的权重，所以只有一个稳定的配置，那就是权重为 1。任何偏离 1 的情况都会导致梯度指数级增长或衰减。注意当权重是复数时，同样的论点适用，只是使用权重的绝对值。

那么我们如何修复消失和爆炸的梯度问题呢？我们看到，只要循环权重为 1 且输入为 0，RNN 的梯度就是稳定的，所以在线性 RNN 论文中，作者提出在这个稳定状态下初始化他们的线性 RNN。具体来说，他们将权重参数化为复极坐标形式 \(ae^{ib}\)，其中 a 是幅度，b 是角度。然后他们通过运行 a 通过这个 \(e^{-e^{()}}\) 函数来限制幅度小于 1，这个函数总是输出一个 0 到 1 之间的数字，而不是像我们通常那样从以 0 为中心的正态分布中随机抽样 a，他们初始化 a 使得幅度 \(e^{-e^{()}}\) 在 0.999 到 1 之间均匀分布。他们在 0 到π/10 弧度之间均匀初始化角度 b。这确保了在初始化时，权重都非常接近 1。最后他们将输入乘以 \(\sqrt{e^{-e^{a}}}\)，这是另一个可学习的参数，初始化时由于 \(e^{-e^{a}}\) 接近于 1，这是一个非常小的数字。这确保了在初始化时，输入都接近于 0，因此它们不会干扰递归。

所以，在初始化时，这个模型大致相同于我之前向你展示的稳定的 RNN。当模型开始训练且权重发生变化后，不能保证它将保持稳定，但实际上看来，像这样初始化模型足以让它学会记住长达数万步的上下文。有了这些修改，我们现在拥有一个计算快速且能学习使用极长上下文的线性 RNN。

在线性 RNN 论文中，他们在长距离竞技场基准测试上评估了这个模型，这是一系列 6 个合成任务的集合，用于评估模型执行长期推理的能力。例如，在 PathX 任务中，模型必须将图像分类为是否包含两个圆之间的完整点状路径。只是图像被展平成一长串 16000 像素的序列。线性 RNN 在长距离竞技场上实现了最先进的性能，平均在各任务中比 Transformer 高出约 33%。

既然我们已经理解了线性 RNN，那么关于状态空间模型的所有讨论是怎么回事呢？事实证明，状态空间模型只是线性 RNN。状态空间模型受控制理论的启发，源自一个完全不同的尝试将连续动态系统离散化的想法，但最终结果仅仅是一个线性 RNN，只是初始化方案略有不同。最常见的状态空间模型形式将每个循环权重参数化为 \(w=e^{(a+bi)}\)，其中是一个可学习的参数，通常初始化为一个非常小的数字，通常在 0.0001 到 0.1 之间。权重乘以一个小数字使其接近 0，当你对接近 0 的东西取幂时，你会得到接近 1 的东西。这再次确保在初始化时循环权重都大约为 1，所以训练是稳定的。

状态空间模型还将输入乘以 \(((a+bi))^{-1(w-1)}\)，因为这是控制理论所规定的，但从经验上看，当你只按照线性 RNN 设置缩放输入时，你会得到相同的性能。在长距离竞技场上，受控制理论启发的状态空间初始化与线性 RNN 初始化性能大致相同。无论何时，当你听到状态空间模型时，就想到线性 RNN。

最后我们可以谈谈 Mamba 了。你看，虽然线性 RNN 在长距离竞技场基准测试上表现得确实很好，但这并不意味着它们是好的语言模型。对于语言建模来说，线性 RNN 的性能远不如 Transformer。这里是各种最先进语言模型的性能，图上显示越低越好。如你所见，包括状态空间模型在内的所有模型的性能都明显不如 Transformer。Mamba 论文中指出的原因是，线性 RNN 无法选择性地忘记输出向量中的信息。如果权重接近 0，则每次输入后输出向量将被设置为 0，实际上模型将立即忘记当前输入之前的任何内容。如果循环权重接近 1，则输出向量在与权重相乘时不会改变，因此输出向量将积累所有观察到的输入的信息。

你希望模型能够根据所看到的输入决定何时存储信息以及何时忘记信息。Mamba 提出了一个优雅的解决方案：不是在每一步使用相同的权重，而是使用取决于输入的不同权重。Mamba 对每个输入向量应用一个线性函数以生成该输入的单独权重向量。然后使用这些生成的权重执行循环扫描。这样，某些输入可以生成接近 0 的权重，从而从输出向量中擦除信息，但其他输入可以生成接近 1 的权重，从而使输出向量保持不变。

我还怀疑，在每一步使用不同的权重有助于消除梯度消失和爆炸，因为现在应该有许多不同的稳定配置，就像前馈网络一样，尽管这在 Mamba 论文中没有提及。Mamba 还使用了另一个技巧，即增加输出向量的大小。在标准 RNN 中，输出向量的大小与输入向量的大小相同。Mamba 将输出向量的大小扩大了 16 倍。这使它能够存储更多以前输入的信息。然后，输出向量在传递到下一层之前被投影回原始大小。

通常，这将使计算时间增加 16 倍，但事实证明，在现代 GPU 上，Mamba 层的主要瓶颈是读写高性能内存中的数据所需的时间。你看，现代 GPU 实际上有两种不同类型的内存。数据存储在主内存中，但为了进行计算，数据首先需要转移到高性能内存中。对于 mamba 递归操作，事实证明，传输数据所需的时间实际上比进行计算本身所需的时间要长得多。因此，Mamba 将输入向量和模型参数转移到高性能内存中，然后在一个单独的块中计算整个 mamba 操作，包括将输出投影回较小的原始大小，然后将结果写回主内存。这样，你只需要在高性能内存中传输原始大小的向量，所以传输时间不变。实际的计算时间慢了 16 倍，但与传输时间相比，计算时间如此之小，以至于它并不真正影响所花的总时间。你基本上可以免费使用 16 倍大的向量。

根据元评审，Mamba 并未在长距离竞技场基准测试上进行测试。记得我之前谈到的那个基准测试吗，线性 RNN 在那里表现得比 Transformer 好得多？这位评审员希望看到 Mamba 在这个任务上的表现如何。现在，这是拒绝一篇论文的一个非常愚蠢的理由，因为长距离竞技场是一项完全不同的任务，而 Mamba 特别是一个语言模型。请记住，Transformer 在长距离竞技场上的表现远不如线性 RNN，但 Transformer 仍然是更好的语言模型。所以在长距离竞技场上的表现绝不表明语言建模能力。Mamba 为语言建模树立了新的艺术标准，它不应该因为没有同时解决一些其他无关的任务而被拒绝。

元评审中的另一个主要批评是 Mamba 只在语言建模上进行了评估，即在预测文本中下一个单词时的准确性。评审者认为这一指标并不能指示语言模型的实用性，而应该评估 Mamba 在衡量模型推理能力的下游任务上的表现。但实际上，这正是他们在 Mamba 论文中所做的，他们将 Mamba 预训练为语言模型，然后对一系列标准的下游基准任务进行了零次提示。而且令人惊讶的是，Mamba 在所有其他语言模型中表现最佳。

作为额外的信息，另一位评审者说，并且我引用，“Mamba 在训练期间有一个二次方的内存需求，就像 Transformer 一样”。这...完全不正确。Mamba 和 Transformer 都没有二次方的内存成本。Transformer 有一个二次方的计算成本，但它们的内存成本是线性的。Mamba 的也是如此。我不确定如果你真的了解它的工作原理，怎么可能得出 Mamba 有二次方内存成本的结论…

因此，你可以想象，这个不够理想的同行评审在机器学习社区中引发了一些关于同行评审实践的辩论，以及 Mamba 是否应该被拒绝的问题。你可能已经猜到了我站在辩论的哪一边，但我想知道你对学术同行评审有多破碎的看法，请在下面的评论中告诉我。或者关于 Mamba 架构本身的想法，我想那也可以。