Mamba的发展,是从SSM->HiPPO->S4->Mamba 演化过来。所以,了解Mamba,得从SSM开始。
SSM,状态空间模型
SSM,就是状态空间模型。
为什么需要SSM?查看三十年前的教科书,控制论的发展,让人们对控制系统不再只满足于研究输出量的变化,对于系统内部的状态变量同时感到兴趣,以便设计和控制这些参数达到最佳控制目的。所以 ,需要状态变量分析。
然后有大段的篇幅进行状态分析的名词定义和状态方程的公式推导,暂且放下,略去将来再看。
SSM的方程如下。
状态方程就是:状态的导数是关于状态和输入的函数。这里的定义是线性函数。
对应的结构图如下:
为什么状态的导数的可以表示为这种反馈的结构?不是很懂,教科书中还有很多课需要补,教科书里既然这么写,先放一下,就先接受这个结论。
由于矩阵D类似于跳跃连接,因此在没有跳跃连接的情况下,SSM可以只关注矩阵A、B、C:
上一篇《Mamba学习(三):离散化SSM的矩阵计算》提到,除了连续的输入之外,还会通常碰到离散的输入(如文本序列)。所以SSM需要离散化形式。
另外,为了加速训练过程,还需要卷积结构表示。
于是,整个SSM就是下面几个主要公式,分别是连续状态方程、离散状态方程和卷积结构表示:
连续状态方程如下:
连续和离散状态方程如下:
卷积结构表示
那上面的公式(4)如何得来?按照时间步骤展开就可以得到:
见《一文通透想颠覆Transformer的Mamba:从SSM、HiPPO、S4到Mamba》的例子,具体展开可以看到,和上面的公式(4)是一样的。
这样,SSM在训练的时候,就可以类似图形处理的CNN网络中的卷积核:
不过文本而不是图像,因此转换为一维结构:
根据上面的公式,其中的卷积核是:
与RNN的关系
RNN可以表示成如下,分别是循环和展开的表示:
SSM的表示对应如下:
可以看到,RNN和SSM在结构上非常相似。
前面文章《Mamba学习(一):总体架构》提到,Mamba可以简化为一个带门控的 RNN。
总结
SSM在推理时,可以类似RNN的循环结构,获得比较快的推理速度;而在训练的时候,可以通过采用卷积核实现并行化,获得较快的训练速读。
- Transformer:训练快,推理慢
- RNN:训练慢,推理快
- Mamba:并行化训练,循环结构推理,一举两得
RNN为啥没法写成卷积形式?因为RNN多了一个非线性的激活函数。
RNN一般表示为ht=f(Wht-1 + Uxt),其中f是非线性的激活函数,比如tanh。
RNN的并行化
作为备注,需要继续前面的话题,那么,RNN就真的不能并行化处理吗?不然!
如果有兴趣,请继续观看苏神的两篇大作:
《Google新作试图“复活”RNN:RNN能否再次辉煌?》
文章链接:https://mp.weixin.qq.com/s/sTlUBXB-PVKI2l0HqRStSA
《脑洞大开:非线性RNN居然也可以并行计算?》
文章链接:https://mp.weixin.qq.com/s/Yj8XFpIgBxTASUeifeb6XQ
另有一篇《将注意力视为RNN!Bengio等人新作媲美Transformer,但超级省内存》也可以看看。