Transformer是状态空间模型: 通过结构化状态空间对偶实现通用模型和高效算法。

1. 状态空间对偶

结构化状态空间模型(Structured State Space Model, SSM)通过引入隐藏状态hh定义了从输入xx到输出yy的映射：

ht=Atht−1+Btxtyt=C⊤tht

对于输入x的P个特征可以独立地使用SSM，并且矩阵A可以简化为标量形式（即所有元素相同的对角矩阵），此时SSM记为：

Y(T,P)=SSM(A(T),B(T,N),C(T,N))(X(T,P))

或写成递推的形式：

y0=C⊤0h0=C⊤0B0x0y1=C⊤1h1=C⊤1A1B0x0+C⊤1B1x1y2=C⊤2h2=C⊤2A2A1B0x0+C⊤2A2B1x1+C⊤2B2x2⋯yT=C⊤TAT⋯A1B0x0+C⊤TAT⋯A2B1x1+⋯+C⊤TATBT−1+C⊤TBTxT

基于此可以写出Y=MX对应的矩阵M：

[C⊤0B0C⊤1A1B0C⊤1B1C⊤2A2A1B0C⊤2A2B1C⊤2B2⋮⋮⋱⋱C⊤TAT…A1B0C⊤TAT…A2B1…C⊤TATBT−1C⊤TBT]

注意到：

C⊤iA×i:jBj=A×i:j⋅(C⊤iBj)

上式与因果线性注意力Y=(L∘QK⊤)V是等价的。作者把它定义为状态空间对偶(State Space Duality, SSD)，揭示了SSM和自注意力机制的相关性；与标准的softmax注意力相比，SSD有两个主要区别:

  1. （线性化）去掉了softmax运算
  2. （因果化）引入了因果掩码矩阵：
L=[1⋮⋱1⋯1]

2. SSD算法

通过对上述矩阵进行块分解，可以高效、快速地计算矩阵乘法。具体来说，可以分成以下四个部分：

上述四种颜色的运算可以分别并行实现，从而实现加速计算。

3. Mamba-2 模型

Mamba-2通过移除了序列线性映射简化了Mamba结构。主要改动包括：

实验结果表明，Mamba-2性能比Mamba更好，并且训练速度更快。