技术前沿

Mamba：2 状态空间模型

2024 年 5 月 6 日

文｜庞德公

编辑｜郭嘉

若在阅读过程中有些知识点存在盲区，可以回到如何优雅的谈论大模型重新阅读。另外斯坦福2024人工智能报告解读为通识性读物。若对于如果构建生成级别的AI架构则可以关注AI架构设计。技术宅麻烦死磕LLM背后的基础模型。当然最重要的是订阅跟随“鲁班模锤”。

Mamba自从出道就一直被拿来和Transformer对比，文章巨多，带着标题的爆炸性字样“颠覆Transfomer”，“全面包围”等等。文章深浅不一，小编认为其背后的设计思路以及演化过程更加重要。为此本专题将从山上有座庙开始，沿着RNN、SSM、S4、HIPPO一路走来。期间抛弃冗杂的文辞，以更接地气的模式协助读者打开新的逻辑思维。

状态空间

先来看一个重要的概念，状态空间。它包含了一组最少变量，但是又能完整的描述系统。它通过数学方式用来表示系统所有可能存在的状态。

打个比方街机三国志的游戏，游戏地图的所有位置组合而成“状态空间”，而“状态空间表示”是对其简化且抽象表示。例如现在什么状态，可以做哪些动作，未来会导致什么变化。当然单个状态可以包括位置、目前关卡、血量、武器装备和待操作的动作选项等等。要是把这些状态进行编码，是不是可以数字化表示。听起来是不是似曾相识，其实就是Embeddings（链接可以帮助回忆下Embeddings技术）。

大模型中的神经网络其实也有隐藏状态，分散在不同的神经单元。当预测下一个字符（Token）的时候也要根据当时的状态，同时结合上一个输出和语境。因此状态空间能在大模型技术发挥作用，而且它也成为和Transformer不一样的架构。

状态空间模型

状态空间模型（SSM）是用于描述这些状态表示并根据某些输入预测一个状态的模型。在时间t（请注意t是时间的意思，它不使用离散序列，而是将连续序列作为输入并预测输出序列）：

模型映射输入序列x(t) —（例如，游戏中的操作，直冲还是挥拳）
根据输入映射到潜在状态h(t) —（例如，快要挂了的状态）
并预测输出序列y(t) —（例如，要放大招）

上图为wiki中关于状态空间（State-space Representation）的示意图。有了上一章节的铺垫，相信读者对该系统的理解较为轻松。换句大白话：输出伴随着输入变化而变化，模型会存储每次处理的内部状态（上下文状态），同时会利用四个关键的矩阵完成内部状态的切换和控制输出。这四个矩阵就是该模型的参数。

上图输入为连续信号u，输出为连续信号y。中间x为状态变量（上一章节的h），A为状态矩阵，主要用于将上个状态xt-1转换为xt，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)系统，时不变系统是输出不会直接随着时间变化的系统。任意时间延迟的输入将得到相同时间延迟的输出（ABCD其实可以带上时间参数，那就是时变系统，延时输入会导致不同的输出）。

这里带出一个概念LTI(线性时不变)。时不变已经解释了，线性的意思大概就是f(a+b)= f(a)+ f(b)。该模型是一个连续(continuous)的，若要运用到实际的场景需要做一些变化，例如将输入离散转连续，得到连续输出在将结果离散化。

离散到连续

如果连续信号的输入，要找到状态来表示h(t)在分析上的确很有挑战。而且要是想将状态空间模型运用到大模型，首先还是要解决离散与连续互转的问题。毕竟大模型输入是离散的，经过Embedding之后的向量序列。

zero-order hold (零阶保持) 是传统数模转换器(DAC)的实际信号重建的数学模型。它通过将每个样本值保留一个样本间隔Δ来将离散时间信号转换为连续时间信号。因此包括大模型在内的任何离散信号可以通过ZOH转化为连续的信号供SSM使用。SSM输出为连续的信号，则可以通过采样转为离散信号继续后面的流程的使用。