Mamba：超越Transformer的新一代神经网络架构

news2024/9/23 9:30:15

在过去的七年里，Transformer一直在语言建模领域占据着主导地位。然而，现在有一个新兴的神经网络架构Mamba，正在挑战Transformer的霸主地位。虽然目前Mamba仅在规模较小的模型上进行了测试（参数量达到数十亿），但其结果已经展现出巨大的潜力。此外，与Transformer相比，Mamba所需的计算量更少。对于n个单词的输入序列，Mamba仅需O(nlog(n))的计算量，而Transformer则需要O(n^2)的计算量。因此，基于Mamba的语言模型应该能够处理~~~~更大的上下文信息量。

什么是Mamba？

Mamba通常被认为是状态空间模型（State-Space Model）的一种扩展。状态空间模型是一种序列模型，近年来逐渐受到关注。然而，状态空间模型背后的理论相当复杂，涉及到一些高级数学。幸运的是，Mamba也可以被理解为递归神经网络（Recurrent Neural Networks，简称RNNs）的扩展，而RNNs相对来说更容易理解。因此，我们将通过RNN的路径来理解Mamba。

递归神经网络概述

递归神经网络是一种处理序列数据的神经网络。给定一个输入向量序列，一个卷积层会对连续的向量组应用神经网络。关键在于神经网络每次只处理少量向量，这使得模型易于训练。然而，卷积神经网络（CNNs）的一个缺点是，远处向量的信息不能被立即组合起来，这使得它们难以理解输入中的长距离依赖关系。为了解决这个问题，Transformer架构应运而生，成功地允许单层神经网络组合任意距离的向量信息。

Transformer的局限性

尽管Transformer在许多任务中表现出色，但它的一个显著限制是其计算量随着输入长度呈二次方增长。这对于较小的输入来说不是大问题，但如果输入包含百万个向量，那么计算量将增加到百万的平方，这相当庞大。

RNN的不同方法

递归神经网络采用了一种完全不同的方法来改进卷积层。其基本思想是，将神经网络应用于一个输入向量和先前的输出向量。这看似一个小改动，但却带来了深远的影响：每个输出向量现在包含了所有先前输入向量的信息，而不仅仅是前一个向量的信息。这使得最终的输出向量包含了所有输入向量的信息，而无需额外的计算量。

RNN的实际问题

尽管理论上RNN能够免费地整合长距离信息，但在实际使用中，RNN存在两个主要问题，使其几乎无法使用。首先，尽管递归层的计算量与卷积层相同，但这种计算无法并行化。即使有许多处理器可用，也无法在前一步完成前开始评估下一个输入。相比之下，卷积层的神经网络只需要看到原始输入，可以同时在所有输入上运行。

其次，RNN非常难以训练。尽管理论上单层递归层可以整合任意多的输入信息，但实际中它们只能学习到整合前几步的信息。这使得RNN在过去几十年中逐渐被卷积神经网络和Transformer所取代，几乎没有被使用。

线性RNN的突破

去年发表的一篇论文展示了线性RNN能够避免上述两个问题，因此线性RNN在处理长序列任务时非常有效。那么，什么是线性递归神经网络？简单来说，就是将神经网络替换为线性函数。尽管线性函数只能进行相对简单的变换，但我们可以通过在每个输出向量后应用完整的神经网络来弥补这一缺陷。

线性递归层的并行计算

通过将递归操作纯粹线性化，可以在O(log(n))时间内并行计算。具体来说，线性递归操作可以通过一种称为扫描（scan）的操作来实现，例如累积和（cumulative sum）。这种操作可以在并行计算中显著提高效率。

解决RNN的训练问题

RNN训练困难的根源在于梯度消失和爆炸问题。线性RNN通过在初始化时确保权重接近于1来解决这个问题。具体方法包括使用复数极坐标形式参数化权重，确保在初始化时权重接近于1，并将输入乘以一个可学习参数，使其接近于0，从而在初始化时保持稳定状态。

Mamba的创新

尽管线性RNN在长序列任务中表现出色，但在语言建模中效果不如Transformer。Mamba通过根据输入生成不同的权重来解决这一问题，从而实现信息的选择性遗忘或保留。此外，Mamba将输出向量的大小增加了16倍，使其能够存储更多的输入信息。这些改进使得Mamba在语言建模中表现优于Transformer，同时计算量仅为O(nlog(n))。