在正式进入今天的内容之前,先简单说明一下。这篇文章将从一个非常基础的概念——状态空间——开始讲起。如果你对某些前置知识点还不太熟悉,可以随时翻翻我们之前发过的系列文章补补课。当然这不是硬性要求,直接往下看也没问题。
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是状态矩阵,主要负责把上一个状态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)系统。所谓时不变,就是说系统的输出不会直接随着时间的变化而改变。简单理解就是:任意时间延迟的输入,都会得到相同时间延迟的输出。当然,如果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)调整为离散化的参数矩阵(Anew,Bnew)。具体的算法如下,能看懂公式的同学可以直接看公式。
同时,Δ、A、B按照ZOH的思路进行调整和设计,具体的实现方式如下:
到了这一步,已经完成了从连续SSM到离散SSM的转化。接下来就可以开始进一步的使用了。注意这里下标从t变成了k(k代表离散的每一步),参数也从A、B变成了Anew、Bnew。当然,在训练过程中保存的还是参数矩阵A,只是在真正使用的时候才进行离散化。
拥有了离散的SSM之后,就可以稍微远离微积分,用纯粹的矩阵运算来解决一切问题。下面这张图展示的是每一步离散迭代计算的流程,看到这里,是不是让你想起了上一章节里讲过的RNN?还是那句话,透过现象看本质——现象会随着时间不断产生事件xt,而人类,正是从这些事件中不断捕获内在的本质hk。