力压Transformer，详解Mamba和状态空间模型

大家好，大型语言模型（LLMs）之所以能够在语言理解与生成上取得巨大成功，Transformer架构是其强大的支撑。从开源的Mistral，到OpenAI开发的闭源模型ChatGPT，都采用了这一架构。

然而技术的探索从未止步，为进一步提升LLMs的性能，学界正在研发能够超越Transformer的新架构。其中，Mamba模型以其创新的状态空间模型（State Space Model）成为研究的焦点。

本文将介绍Mamba模型及其在语言建模领域的应用，逐步解析状态空间模型的基本概念，并通过丰富的可视化内容，让大家直观地理解这一技术如何有望挑战现有的Transformer架构。 

1.Transformer架构的挑战

先对Transformer架构做一个快速回顾，并指出其存在的一个主要缺陷。

Transformer架构将文本输入视为由一系列token构成的序列：

其核心优势在于，无论面对何种输入，都能追溯到序列中的早期token，以此来推导出其深层的语义表示。

1.1 Transformer的核心组件

Transformer由架构两部分组成：编码器和解码器。编码器负责解析文本，而解码器则负责生成文本。这种结构的结合，使之能够胜任从文本翻译到内容创作的多种任务。

我们可以仅利用解码器部分来创建生成式模型，这种基于Transformer的模型，即生成预训练Transformer（GPT），通过解码器来续写或补全输入的文本，展现出其在文本生成方面的强大能力。

1.2 自注意力的高效训练

单个解码器块由两个主要部分组成，即“掩蔽自注意力机制（Masked Self-attention）”和“前馈神经网络（Feedforward Neural Network）”。

自注意力是这些模型运行如此良好的主要原因，它提供了整个序列的未压缩视图，并加快了训练速度。具体来说，自注意力机制的工作原理是通过创建一个矩阵，该矩阵对序列中的每个token与之前所有token进行比较，并通过计算它们之间的相关性来确定权重。

在训练过程中，自注意力矩阵是一次性整体构建的，这表示不需要依次等待每个token的注意力计算完成，而是可以同时进行整个序列的注意力计算。例如，在处理“我”和“名字”的关联之前，无需先完成“名字”和“是”的关联计算。

这种设计实现了训练过程的并行化，极大地提升了训练速度，使Transformer架构在处理大规模数据集时更加高效。

1.3 训练与推理的矛盾

然而Transformer架构也有其局限性，每当生成新的token时，必须对整个序列的注意力权重重新进行计算，哪怕此前已经生成了若干token。

生成长度为L的序列需要大约L²次计算，随着序列的延长，计算成本会急剧上升。

这种对序列全面重新计算的需求，是Transformer架构的一个主要瓶颈。

接下来，让我们看看传统的递归神经网络（RNN）是如何克服这一推理过程中的效率问题。

1.4 RNN的潜力

递归神经网络（RNN）是一种处理序列数据的网络结构。在序列的每个时间点，RNN接收两个输入：当前时间点t的输入数据和上一个时间点t-1的隐藏状态，以此来计算下一个隐藏状态并预测输出结果。

RNN具有循环机制，能够将历史信息传递至下一步，类似于将每一步的信息“串联”起来。这种机制可以通过可视化的方式“展开”，以便更清晰地理解其工作原理。

在生成输出时，RNN仅依赖于前一步骤的隐藏状态和当前的输入数据，避免了像Transformer那样需要重新计算整个序列的历史隐藏状态。

正因如此，RNN在进行推理时速度较快，因为它的计算量与序列长度呈线性关系，理论上能够处理无限长的上下文。

举例来说，将RNN应用于之前的输入文本：

每个隐藏状态都是对之前所有状态的压缩汇总。

但这里存在一个问题：随着时间的推移，比如在生成名字"Maarten"时，最后一个隐藏状态可能已经丢失了对"Hello"的记忆，因为RNN在每一步只考虑了前一个状态的信息。

此外，RNN的这种顺序依赖性也导致了另一个问题：它的训练过程无法并行化，必须按顺序逐步进行。

与Transformer相比，RNN在推理速度上具有优势，但在训练并行化方面却存在不足。

这就引出了一个问题：能否找到一种结合了Transformer训练并行化优势和RNN线性推理优势的架构，答案是肯定的，Mamba模型就是。在深入了解Mamba架构之前，先来了解状态空间模型的世界。

2.状态空间模型（SSM）

状态空间模型（SSM），像Transformer和RNN一样，处理信息序列，如文本和信号。在这一部分中，我们将了解SSM的基础知识以及其与文本数据的关系。

2.1 什么是状态空间

状态空间模型是一种通过数学方法全面描述系统状态的方式，包含了描述系统所必需的全部最小变量。

简单来说，就像在迷宫中寻找路径，状态空间就是那张展示所有可能位置（即状态）的地图。在这张地图上，每个点都代表一个独特的位置，并且携带了如距离出口远近等具体信息。

进一步简化这个概念，可以将“状态空间表示”理解为这张地图的提炼，它不仅告诉我们当前所在的位置（即当前状态），还展示了可能的目的地（未来状态），以及如何通过特定的行动（比如右转或左转）达到下一个状态。

虽然状态空间模型依赖方程和矩阵来捕捉系统的行为，但其核心目标是追踪系统的位置、可能的移动方向及其变化路径。

在这个模型中，用以描述状态的变量，如示例中的X和Y坐标或者到出口的距离，统称为“状态向量”。

这听起来有点熟悉，因为在语言模型中，类似的嵌入或向量经常用来描述输入序列的“状态”。例如，你当前的位置状态就可以通过一个向量来表示：

在神经网络的语境下，系统的“状态”通常指的是其隐藏状态，这在生成新token的过程中，尤其是在大型语言模型的背景下，扮演着至关重要的角色。

2.2 状态空间模型（SSM）

状态空间模型（SSM）是一种描述系统状态并预测其未来状态的模型，它能够基于当前状态和特定输入来推断接下来可能的状态变化。

在传统方法中，状态空间模型（SSM）在特定时间点 t 的操作流程如下：

• 映射输入：SSM首先将输入序列 x(t) 映射到潜在的状态表示 h(t) 。以迷宫为例，当你向左下方移动时，这一动作会被转化为与出口的距离以及具体的坐标位置。

• 预测输出：接着，SSM利用这些状态信息来预测输出序列 y(t) 。比如，为了更快接近出口，SSM可能会建议你再次向左移动。

不同于传统的离散序列输入方法，如简单的一步移动，我们采用的是连续序列输入，并预测其输出序列。

在状态空间模型（SSM）中，认为动态系统，例如在三维空间中移动的物体，可以通过其在特定时间点的状态，通过两个核心方程来预测。这两个方程是状态空间模型的基石。

我们的目标是通过解决这些方程，发现统计原理，利用观测到的数据——包括输入序列和之前的状态——来预测系统的未来状态。我们追求的是找到一种状态表示 ℎ(