Transformer是状态空间模型: 通过结构化状态空间对偶实现通用模型和高效算法。
1. 状态空间对偶
结构化状态空间模型(Structured State Space Model, SSM)通过引入隐藏状态$h$定义了从输入$x$到输出$y$的映射：
\[\begin{aligned} h_{t} &= A_t h_{t-1} + B_t x_t \\ y_t &= C_t^{\top} h_t \end{aligned}\]对于输入$x$的$P$个特征可以独立地使用SSM，并且矩阵$A$可以简化为标量形式（即所有元素相同的对角矩阵），此时SSM记为：
\[Y^\mathtt{(T,P)} = \mathsf{SSM}(A^\mathtt{(T)}, B^\mathtt{(T,N)}, C^\mathtt{(T,N)})(X^\mathtt{(T,P)})\]或写成递推的形式：
\[\begin{aligned} \mathbf{y}_0 &= C_0^\top \mathbf{h}_0= C_0^\top B_0 \mathbf{x}_0 \\ \mathbf{y}_1 &= C_1^\top \mathbf{h}_1 = C_1^\top A_1B_0 \mathbf{x}_0 + C_1^\top B_1 \mathbf{x}_1 \\ \mathbf{y}_2 &= C_2^\top \mathbf{h}_2 = C_2^\top A_2A_1B_0 \mathbf{x}_0 +C_2^\top A_2B_1 \mathbf{x}_1 + C_2^\top B_2 \mathbf{x}_2 \\ & \cdots \\ \mathbf{y}_{T} &= C_T^\top A_T\cdots A_1B_0 \mathbf{x}_0 +C_T^\top A_T\cdots A_2B_1 \mathbf{x}_1 + \cdots + C_T^\top A_TB_{T-1} + C_T^\top B_T \mathbf{x}_T \\ \end{aligned}\]基于此可以写出$Y=MX$对应的矩阵$M$：
\[\begin{bmatrix} C_0^\top B_0 & \\ C_1^\top A_1 B_0 & C_1^\top B_1 & \\ C_2^\top A_2A_1 B_0 & C_2^\top A_2 B_1 & C_2^\top B_2 \\ \vdots & \vdots & \ddots & \ddots \\ C_\mathtt{T}^\top A_{\mathtt{T}}\dots A_1 B_0 & C_\mathtt{T}^\top A_{\mathtt{T}}\dots A_2 B_1 & \dots & C_\mathtt{T}^\top A_{\mathtt{T}} B_{\mathtt{T}-1} & C_\mathtt{T}^\top B_{\mathtt{T}} \\ \end{bmatrix}\]
注意到：
\[C_i^\top A_{i:j}^\times B_j = A_{i:j}^\times \cdot (C_i^\top B_j)\]上式与因果线性注意力$Y = (L \circ QK^\top) V$是等价的。作者把它定义为状态空间对偶(State Space Duality, SSD)，揭示了SSM和自注意力机制的相关性；与标准的softmax注意力相比，SSD有两个主要区别:
- （线性化）去掉了softmax运算
- （因果化）引入了因果掩码矩阵：
2. SSD算法
通过对上述矩阵进行块分解，可以高效、快速地计算矩阵乘法。具体来说，可以分成以下四个部分：
- （橙色）对角块，输入到输出，可以直接计算
- （绿色）低秩块：输入到状态
- （黄色）低秩块：状态到状态
- （蓝色）低秩块：状态到输出
上述四种颜色的运算可以分别并行实现，从而实现加速计算。
3. Mamba-2 模型
Mamba-2通过移除了序列线性映射简化了Mamba结构。主要改动包括：
- 通过单个投影并行生成$A,X,B,C$
- 在最终输出投影之前添加一个归一化层
实验结果表明，Mamba-2性能比Mamba更好，并且训练速度更快。
