Mamba-2 状态空间对偶性理论

概述

状态空间对偶性（State Space Duality, SSD）框架 是 Mamba-2 的核心理论基础，它揭示了状态空间模型（SSM）与结构化注意力之间的深层数学联系。这一框架将 SSM 视为**半可分矩阵（Semi-separable Matrix）的矩阵变换，同时将线性注意力视为结构化掩码注意力（Structured Masked Attention, SMA）**的张量收缩运算，而两者的交集正是 SSD 模型。

核心洞见：SSD 模型可以通过两种完全独立的方式推导——一种从 SSM 视角出发，另一种从注意力视角出发——最终证明它们在数学上是等价的。¹

SSD 框架图景

┌─────────────────────────────────────────────────────────────┐
│                    SSD Framework                            │
│                                                             │
│        ┌──────────────────┐    ┌──────────────────┐        │
│        │  Structured SSMs │    │ Structured       │        │
│        │  (半可分矩阵)     │    │ Attention        │        │
│        │                  │    │ (结构化掩码注意力) │        │
│        └────────┬─────────┘    └────────┬─────────┘        │
│                 │                       │                  │
│                 └───────────┬───────────┘                  │
│                             │                              │
│                      ┌──────┴──────┐                       │
│                      │  SSD Model  │                       │
│                      │  (Mamba-2)  │                       │
│                      └─────────────┘                       │
└─────────────────────────────────────────────────────────────┘

1. SSM 视角：半可分矩阵变换

1.1 序列变换的矩阵形式

许多序列模型可以表示为矩阵序列变换（或称矩阵混合器）：

$$Y=M\cdot X$$

其中 $M$ 是一个矩阵，$X\in {\mathbb{R}}^{(T,P)}$ 是输入序列，$Y\in {\mathbb{R}}^{(T,P)}$ 是输出序列。

自注意力的例子是 $M=\mathsf{softmax}(Q{K}^{\top })$，而 SSM 同样可以写成这种形式。

1.2 选择性 SSM 的定义

Mamba-2 中的选择性状态空间模型定义为¹：

$$\begin{array}{rl}{h}_t& =A_t{h}_{t-1}+B_t{x}_t\\ y_t& =C_t^{\top }{h}_t\end{array}$$

其中 $A_t$ 是标量-单位阵结构（scalar-identity structure）的矩阵。

1.3 SSM 的矩阵展开

将 SSM 递推展开为矩阵形式，得到：

$$Y=\mathsf{SSM}(A,B,C)(X)=MX$$

其中矩阵 $M$ 的元素定义为：

$$M_{ij}=C_i^{\top }{A}_{i:j}^{\times }{B}_j:=C_i^{\top }{A}_i\cdots A_{j+1}{B}_j$$

且 $M_{ij}=0$ 当 $i<j$（下三角）。

展开后的矩阵形式如下：

$$M=\left[\begin{array}{ccccc}{C}_0^{\top }{B}_0& & & & \\ C_1^{\top }{A}_1{B}_0& C_1^{\top }{B}_1& & & \\ C_2^{\top }{A}_2{A}_1{B}_0& C_2^{\top }{A}_2{B}_1& C_2^{\top }{B}_2& & \\ ⋮& ⋮& \ddots & \ddots & \\ C_T^{\top }{A}_{T-1}\cdots A_1{B}_0& C_T^{\top }{A}_{T-1}\cdots A_2{B}_1& \cdots & C_T^{\top }{A}_{T-1}{B}_{T-2}& C_T^{\top }{B}_{T-1}\end{array}\right]$$

1.4 半可分矩阵的定义与性质

定义：上述矩阵称为**（三角）半可分矩阵（Semi-separable Matrix）**。¹

半可分矩阵具有结构化秩性质：对角线及其下方的任意子矩阵都是低秩的。

┌─────────────────────────────┐
│  ★ ★ ★ ★ ★ ★ ★ ★ ★ ★ ★ ★  │  ← 对角线及以下区域
│    ★ ★ ★ ★ ★ ★ ★ ★ ★ ★ ★  │    所有子矩阵
│      ★ ★ ★ ★ ★ ★ ★ ★ ★ ★  │    都是低秩的
│        ★ ★ ★ ★ ★ ★ ★ ★ ★  │
│          ★ ★ ★ ★ ★ ★ ★ ★  │
│            ★ ★ ★ ★ ★ ★ ★  │
│              ★ ★ ★ ★ ★ ★  │
│                ★ ★ ★ ★ ★  │
│                  ★ ★ ★ ★  │
│                    ★ ★ ★  │
│                      ★ ★  │
│                        ★  │
└─────────────────────────────┘

1.5 从 SSM 到注意力的对偶性推导

对于标量 $A_t$ 的情况（即 SSD 模型），由于 $A_t$ 是标量，可以提取出来：

$$C_i^{\top }{A}_{i:j}^{\times }{B}_j=A_{i:j}^{\times }\cdot (C_i^{\top }{B}_j)$$

这直接导出 SSD 的注意力形式：

$$M=L\circ C{B}^{\top }\in {\mathbb{R}}^{(T,T)}$$

其中 $\circ$ 表示逐元素乘法，$L$ 是因果掩码矩阵。

2. 注意力视角：张量收缩与结构化掩码

2.1 核注意力定义

将注意力定义为函数：

$$(Q^{(T,N)},K^{(S,N)},V^{(S,P)})\mapsto Y^{(T,P)}$$

标准形式为：

$$Y=(Q{K}^{\top })\cdot V$$

2.2 线性注意力的张量收缩证明

因果线性注意力的二次形式为：

$$Y=(L\circ Q{K}^{\top })\cdot V$$

其中 $L$ 是因果掩码矩阵：

$$L=\left[\begin{array}{cccc}1& & & \\ 1& 1& & \\ 1& 1& 1& \\ ⋮& ⋮& ⋮& \ddots \end{array}\right]$$

使用**爱因斯坦求和约定（einsum）**将上述运算写成张量收缩形式：

$$\begin{array}{rl}G& =\mathsf{contract}(\text{TN},\text{SN}\to \text{TS})(Q,K)\\ M& =\mathsf{contract}(\text{TS},\text{TS}\to \text{TS})(G,L)\\ Y& =\mathsf{contract}(\text{TS},\text{SP}\to \text{TP})(M,V)\end{array}$$

简化为单次四路收缩：

$$y=\mathsf{contract}(\text{TN},\text{SN},\text{SP},\text{TS}\to \text{TP})(Q,K,V,L)$$

2.3 线性形式的推导

选择不同的收缩顺序可以避免二次复杂度。重新排序为 $V,K,L,Q$：

$$\begin{array}{rl}Z& =\mathsf{contract}(\text{SP},\text{SN}\to \text{SPN})(V,K)\\ H& =\mathsf{contract}(\text{TS},\text{SPN}\to \text{TPN})(L,Z)\\ Y& =\mathsf{contract}(\text{TN},\text{TPN}\to \text{TP})(Q,H)\end{array}$$

关键观察：第二行的 $H=L\cdot Z$ 恰好是矩阵乘法，可以通过**累积求和（cumsum）**线性时间计算。

2.4 结构化掩码注意力（SMA）

定义：结构化掩码注意力（SMA）定义为四路张量收缩，其中注意力掩码 $L$ 是任意结构化矩阵。¹

这推广了原始的线性注意力框架。常见的结构化矩阵包括：

• 因果掩码 → 线性注意力（Katharopoulos et al.）
• 指数衰减掩码 → Retentive Network
• 半可分矩阵 → SSD 模型（Mamba-2）

3. 两种视角的等价性证明

3.1 1-半可分矩阵

SSD 模型对应的掩码矩阵 $L$ 是 1-半可分矩阵（1-SS Matrix）：

$$L=\left[\begin{array}{ccccc}1& & & & \\ a_1& 1& & & \\ a_2{a}_1& a_2& 1& & \\ ⋮& ⋮& \ddots & \ddots & \\ a_{T-1}\cdots a_1& a_{T-1}\cdots a_2& \cdots & a_{T-1}& 1\end{array}\right]$$

因此，SSD 模型本质上是 1-半可分掩码注意力 或 1-SS SMA。

3.2 从 SMA 到 SSM 的等价性

给定 1-SS 矩阵 $L$，矩阵-向量乘法 $y=Lx$ 可以通过标量递推线性时间计算：

$$\begin{array}{rl}{y}_0& =x_0\\ y_1& =a_1{x}_0+x_1\\ y_2& =a_2{a}_1{x}_0+a_2{x}_1+x_2=a_2{y}_1+x_2\\ ⋮& =⋮\end{array}$$

这恰好对应原始的 SSM 递推形式！

3.3 对偶性的本质

对偶表示	SSM 视角	SMA 视角
核心表示	半可分矩阵变换	四路张量收缩
广义化	状态空间模型	线性注意力
SSD 实例化	标量 SSM（$A_t$ 为标量）	1-SS 掩码注意力
线性-二次对偶	结构化矩阵乘法算法	张量收缩约简顺序

核心洞见：虽然普遍认为引入注意力掩码 $L$ 后无法重排序矩阵乘法，但张量收缩约简顺序的灵活性使得这仍然可行。矩阵乘法的结合律是张量收缩约简顺序的特殊情况。¹

4. 与 Transformer 的理论联系

4.1 SSM 作为通用序列变换

Transformer 的核心是自注意力机制：

$$Y=\mathsf{softmax}(Q{K}^{\top })\cdot V$$

而 SSM 通过半可分矩阵提供了一种不同的序列混合方式。SSD 框架证明了：

Transformer 是 SSM 的一种特例²

具体来说，当 SSM 的状态维度 $N\to \infty$ 且参数结构特殊化时，可以逼近 Transformer 的表达能力。

4.2 线性注意力与 SSM 的联系

SSD 框架揭示了以下等价链条：

Transformer (softmax 注意力)
        ↓
Kernel Attention (特征映射)
        ↓
Linear Attention (累积求和)
        ↓
SSD (半可分掩码注意力)
        ↓
SSM (状态空间递推)

4.3 混合算法的启示

SSD 框架启发了混合算法的设计：通过半可分矩阵的块分解，可以结合线性形式和二次形式的优点，在计算效率和表达能力之间取得平衡。¹

5. 关键数学公式汇总

5.1 选择性 SSM

$$\begin{array}{rl}{h}_t& =A_t{h}_{t-1}+B_t{x}_t\\ y_t& =C_t^{\top }{h}_t\end{array}$$

5.2 半可分矩阵元素

$$M_{ij}=C_i^{\top }{A}_{i:j}^{\times }{B}_j=C_i^{\top }{A}_i{A}_{i+1}\cdots A_{j+1}{B}_j$$

5.3 SSD 的注意力形式

$$M=L\circ C{B}^{\top }$$

5.4 结构化掩码注意力的张量收缩

$$y=\mathsf{contract}(\text{TN},\text{SN},\text{SP},\text{TS}\to \text{TP})(Q,K,V,L)$$

5.5 1-半可分矩阵

$$L_{ij}=\{\begin{array}{ll}{\prod }_{k=i}^{j-1}{a}_k& \text{if }i<j\\ 1& \text{if }i=j\\ 0& \text{if }i>j\end{array}$$

6. 算法意义

6.1 两种计算方式

计算方式	算法	复杂度
线性形式	结构化矩阵乘法（利用半可分结构）	$O(T)$
二次形式	标准矩阵乘法（展开掩码矩阵）	$O(T^2)$

6.2 系统优化

SSD 框架使得 Mamba-2 能够借鉴 Transformer 的系统优化技术²：

• 张量并行（Tensor Parallelism）
• 序列并行（Sequence Parallelism）
• 可变序列长度支持

6.3 状态维度的影响

模型	状态维度 $N$	表达能力	计算效率
Mamba-1	$O(N)$	高	中
Mamba-2 (SSD)	$O(1)$ 标量	中	高
Transformer	$N\to \infty$	最高	低

7. 未来研究方向

基于 SSD 框架的理论基础，以下是潜在的研究方向：

1. 新的序列模型：寻找具有特定性质的结构化矩阵类
2. 双向扩展：自然地将 Mamba 扩展为双向模型
3. 知识蒸馏：从 Transformer 蒸馏到 SSM 的理论工具
4. 可解释性：利用矩阵形式分析 SSM 的内部表示
5. 新型结构：探索 Toeplitz 或 Fourier 结构化注意力

---

参考资料

相关链接

• Mamba-2 架构解析
• Mamba-2 算法实现
• 状态空间模型入门
• Transformer 与注意力机制

Footnotes

1. Tri Dao, “State Space Duality (Mamba-2) Part II - The Theory”, 2024. https://tridao.me/blog/2024/mamba2-part2-theory/ ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶

2. Tri Dao & Albert Gu, “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality”, arXiv:2405.21060, 2024. https://arxiv.org/abs/2405.21060 ↩ ↩²