Mamba：2 状态空间模型

在正式进入今天的内容之前，先简单说明一下。这篇文章将从一个非常基础的概念——状态空间——开始讲起。如果你对某些前置知识点还不太熟悉，可以随时翻翻我们之前发过的系列文章补补课。当然这不是硬性要求，直接往下看也没问题。

Mamba从出道开始就一直被拿来和Transformer做对比，相关文章多到看不过来。标题一个比一个生猛，“碘伏Transformer”、“全面包围”之类的话术层出不穷。说得夸张一点，这些文章深浅不一，有的讲得浅尝辄止，有的又过于硬核。其实更重要的是背后的设计思路以及它整个演化的过程。为此，这个专题将从“山上有座庙”讲起，沿着RNN、SSM、S4、HIPPO一路梳理。全程会尽量丢掉那些冗长的术语，用更接地气的方式来帮大家打开一个新的思维框架。

状态空间

先来看一个非常重要的概念：状态空间。它由一组最少的变量组成，但又能完整地描述一个系统。用数学的方式来说，它就是用来表示系统所有可能状态的集合。 举个例子：你玩过街机版的《三国志》吗？游戏地图上所有可能的位置组合起来，就构成了一个“状态空间”。而所谓的“状态空间表示”，是对这个空间的一种简化和抽象。比如当前是什么状态，能执行哪些动作，这些动作未来会导致什么变化。单个状态可以包括角色位置、当前关卡、血量、武器装备、以及下一步待操作的选项等等。如果能把这些状态进行编码，是不是就能用数字来表示它们了？听起来是不是有点耳熟？其实这就是Embedding的思路。 大模型里的神经网络也有隐藏状态，它们分散在不同的神经单元中。当模型要预测下一个字符（Token）的时候，也需要根据当前的状态，同时结合上一个输出结果和上下文语境。所以，状态空间不仅能在大模型技术里发挥作用，而且它也构成了一个和Transformer完全不同的架构。

状态空间模型

状态空间模型（SSM），就是用来描述这些状态表示，并根据某些输入预测下一个状态的模型。在时间t（注意这里的t表示时间，它不是用离散序列，而是把连续序列作为输入并预测输出序列）这个过程中，模型主要做三件事： 1. 映射输入序列x(t)——比如游戏里的操作，是直冲还是挥拳。 2. 根据输入映射到潜在状态h(t)——比如“快要挂了”的状态。 3. 预测输出序列y(t)——比如“该放大招了”。 上图为wiki中关于状态空间（State-space Representation）的示意图。有了前面的铺垫，相信大家理解这张图会相对轻松。用大白话来说：输出会随着输入的变化而变化，模型会把每次处理过程中的内部状态（也就是上下文状态）存储起来，同时利用四个关键的矩阵来完成内部状态的切换并把控输出。这四个矩阵，就是这个模型的全部参数。

上图输入为连续信号u，输出为连续信号y。中间x为状态变量（就是上一段的h）。A是状态矩阵，主要负责把上一个状态x_t-1转换到x_t；B是输入矩阵；C是输出矩阵；D是前馈矩阵。

这个模型是Mamba的基石，它有几个很突出的特点。第一，模型不仅能够反映系统内部的状态，还能揭示系统内部状态和外部输入、输出变量之间的联系。第二，它可以把多个变量的时间序列统一处理成向量时间序列。从变量到向量的这种转变，更适合解决多输入、多输出变量的建模问题。第三，这个模型能用最少的信息来描述系统状态，在运算效率上天然有优势。 接下来这张图值得仔细看看。这张图的原图在网上被引用了很多次，但我认真推敲后发现原图其实存在一些问题，所以干脆用PS重新画了一版。这张图展示的是状态空间具体是如何运算的。

假设输入向量为x（d维），输出向量为y（q维），状态向量为h（n维）。需要强调一下，在原始模型中，输入和输出的维度不必一致。那么矩阵A的维度是n×n，矩阵B的维度是n×d，矩阵C的维度是q×n，矩阵D的维度是q×d。上图的例子中，输入n和输出q均为3维，中间状态向量为4维。

运算公式如下图。也可以很直观地看出来：h的状态更新，用的是新的输入和上一轮的中间状态，这一点和RNN非常相似。而输出，则由状态h和输入分别乘以两个矩阵再累加来完成。这里留一个有趣的作业：可以带着下面的公式去推敲一下上图的矩阵运算，看看能不能对得上。 聊到这里，技术门槛其实并没有想象中那么高。但大家还是要完全消化这些内容，比如状态空间模型的参数有哪些、运算模式是什么、运行的原理是什么。否则在后面深入的部分，很容易迷失方向。 最后值得提一句，这个模型是一个时不变（time-invariant）系统。所谓时不变，就是说系统的输出不会直接随着时间的变化而改变。简单理解就是：任意时间延迟的输入，都会得到相同时间延迟的输出。当然，如果A、B、C、D这四个矩阵带上时间参数，就变成时变系统了，那时延时输入就会导致不同的输出。 这里引出一个概念——LTI（线性时不变）。时不变刚才已经解释过了，线性的意思大致可以用公式f(a+b)= f(a)+ f(b)来表达。而这个模型本身是连续的，要把它真正应用到实际场景中，还需要做一些变化，比如将输入从离散转成连续，得到连续输出后再把结果离散化。

离散到连续

如果是连续信号的输入，想找到状态来表示h(t)在分析上是很有挑战的。而且要把状态空间模型应用到大模型上，首先得解决离散和连续互相转换的问题。毕竟大模型的输入本身是离散的，经过Embedding之后才变成向量序列。 Zero-order hold（零阶保持）是传统数模转换器（DAC）在实际信号重建中常用的数学模型。它通过把每个样本值保持一个样本间隔Δ，将离散时间信号转换为连续时间信号。所以，包括大模型在内的任何离散信号，都可以通过ZOH转化为连续的信号供SSM使用。SSM输出的连续信号，再通过采样转回离散信号，继续后续流程。 上图展示的是离散转连续，下图展示的是连续转离散。当然，离散化的方法不止一种，Mamba采用的正是ZOH。 在Mamba的原始论文中，它对连续SSM的改造思路，主要是围绕将连续的矩阵参数（Δ，A，B）调整为离散化的参数矩阵（A_new，B_new）。具体的算法如下，能看懂公式的同学可以直接看公式。 同时，Δ、A、B按照ZOH的思路进行调整和设计，具体的实现方式如下： 到了这一步，已经完成了从连续SSM到离散SSM的转化。接下来就可以开始进一步的使用了。注意这里下标从t变成了k（k代表离散的每一步），参数也从A、B变成了A_new、B_new。当然，在训练过程中保存的还是参数矩阵A，只是在真正使用的时候才进行离散化。 拥有了离散的SSM之后，就可以稍微远离微积分，用纯粹的矩阵运算来解决一切问题。下面这张图展示的是每一步离散迭代计算的流程，看到这里，是不是让你想起了上一章节里讲过的RNN？还是那句话，透过现象看本质——现象会随着时间不断产生事件x_t，而人类，正是从这些事件中不断捕获内在的本质h_k。