Mamba-3 深度剖析 - 02. SSM 基础回顾与演进 (Foundations of SSM)
Mamba-3 深度剖析 - 02. SSM 基础回顾与演进 (Foundations of SSM)
本章摘要
Mamba-3 的核心创新并非凭空而来，而是建立在现代状态空间模型（State Space Models, SSMs）坚实的理论基础之上。为了深刻理解 Mamba-3 的梯形离散化与复数动力学，我们必须首先回顾 SSM 的数学本源，并深入剖析 Mamba-2 确立的“结构化状态空间对偶性”（SSD）架构。本章将推导从连续系统到离散递归的演变过程，解析 Mamba-2 如何统一 RNN 与 Transformer 的视角，并最终揭示其在表达能力上遗留的本质缺陷——这也是 Mamba-3 诞生的直接诱因。
2.1 状态空间模型：从连续到离散
结构化状态空间模型（Structured SSMs）本质上是定义在连续时间上的动力系统。它们通过隐变量 h ( t ) h(t) h(t) 将输入函数 x ( t ) x(t) x(t) 映射为输出函数 y ( t ) y(t) y(t)。
2.1.1 连续时间动力学 (Continuous-Time Dynamics)
一个标准的线性时变（Linear Time-Varying, LTV）系统由以下一阶微分方程描述：
h ˙ ( t ) = A ( t ) h ( t ) + B ( t ) x ( t ) y ( t ) = C ( t ) ⊤ h ( t ) + D ( t ) x ( t ) egin{aligned} dot{h}(t) &= A(t) h(t) + B(t) x(t) y(t) &= C(t)^ op h(t) + D(t) x(t) end{aligned} h˙(t)y(t)=A(t)h(t)+B(t)x(t)=C(t)⊤h(t)+D(t)x(t)
符号定义：
- x ( t ) ∈ R x(t) in mathbb{R} x(t)∈R：输入信号（在语言模型中即为 Token 的 Embedding 投影）。
- h ( t ) ∈ R N h(t) in mathbb{R}^N h(t)∈RN：隐状态（Hidden State），负责记忆历史信息。
- y ( t ) ∈ R y(t) in mathbb{R} y(t)∈R：输出响应。
- A ( t ) ∈ R N × N A(t) in mathbb{R}^{N imes N} A(t)∈RN×N：状态转移矩阵（State Transition Matrix），决定系统的演化轨迹。
- B ( t ) ∈ R N , C ( t ) ∈ R N B(t) in mathbb{R}^N, C(t) in mathbb{R}^N B(t)∈RN,C(t)∈RN：分别控制输入进入状态空间和状态映射到输出的投影。
- D ( t ) ∈ R D(t) in mathbb{R} D(t)∈R：直连/残差连接（Skip Connection），允许输入信号直接传递到输出。
图解：SSM 系统控制流
上图直观展示了状态空间模型的信息流动过程，完美对应了包含跳跃连接的完整公式：
- 输入阶段：左侧输入信号 x ( t ) x(t) x(t) 被分为两路。一路经由矩阵 B B B 投影进入状态方程，另一路经由矩阵 D D D 直接通向输出。
- 状态演化 (Feedback Loop)：
- 求和节点 ( ⨁ igoplus ⨁)：将输入激励 B x ( t ) Bx(t) Bx(t) 与历史状态反馈 A h ( t ) Ah(t) Ah(t) 相加，得到状态的变化率 h ˙ ( t ) dot{h}(t) h˙(t)。
- 积分器 ( ∫ int ∫)：连续时间系统的核心组件（图中方框内的积分符号），它对 h ˙ ( t ) dot{h}(t) h˙(t) 进行积分以更新当前隐状态 h ( t ) h(t) h(t)。
- 反馈矩阵 ( A A A)：状态 h ( t ) h(t) h(t) 经过矩阵 A A A 变换后传回求和节点，形成闭环控制，赋予模型“记忆”能力。
- 输出生成：右侧的隐状态 h ( t ) h(t) h(t) 经由观测矩阵 C C C 投影后，与直连通路 D x ( t ) Dx(t) Dx(t) 在输出端的求和节点 ( ⨁ igoplus ⨁) 汇合，最终生成系统输出 y ( t ) y(t) y(t)。
2.1.2 离散化 (Discretization)
在实际的序列建模（如处理文本 Token）中，我们面对的是离散的时间步 t k t_k tk。为了将连续微分方程应用于离散序列，我们需要引入时间步长 Δ t Delta t Δt 进行离散化。离散化方法的选择直接决定了模型的精度和特性。
欧拉法 (Euler’s Method) —— Mamba-2 的选择
Mamba-2（以及大多数早期 SSM）采用了最简单的欧拉离散化（零阶保持假设）。
其核心假设是在区间
[
t
k
−
1
,
t
k
]
[t_{k-1}, t_k]
[tk−1,tk] 内，
x
(
t
)
x(t)
x(t) 和
B
(
t
)
B(t)
B(t) 保持不变。这导出了如下的离散递归公式：
h t = e Δ t A t h t − 1 + Δ t B t x t h_t = e^{Delta_t A_t} h_{t-1} + Delta_t B_t x_t ht=eΔtAtht−1+ΔtBtxt
离散化参数说明：
- t t t：离散时间步索引（对应连续时间点 t k t_k tk）。
- Δ t Delta_t Δt：在时间步 t t t 的离散化时间步长（Step Size）。注意这里带有下标 t t t，表示步长本身也可以是数据依赖的（Input-Dependent）。
为了简化计算并引入数据依赖性（Selectivity），Mamba-2 进一步将矩阵 A t A_t At 简化为对角实数矩阵（甚至标量），并定义了以下离散参数：
- 遗忘门 (Forget Gate): α t : = e Δ t A t ∈ ( 0 , 1 ) lpha_t := e^{Delta_t A_t} in (0, 1) αt:=eΔtAt∈(0,1)。由于 A t A_t At 通常被参数化为负数， α t lpha_t αt 起到衰减历史信息的作用。
- 输入门 (Input Gate): γ t : = Δ t gamma_t := Delta_t γt:=Δt。
最终，Mamba-2 的核心递归更新规则（Recurrence）可以写为：
h t = α t h t − 1 + γ t B t x t oxed{h_t = lpha_t h_{t-1} + gamma_t B_t x_t} ht=αtht−1+γtBtxt
以及输出方程：
y
t
=
C
t
⊤
h
t
y_t = C_t^ op h_t
yt=Ct⊤ht
这一形式与门控循环单元（GRU/LSTM）极其相似，但完全线性，没有任何非线性激活函数干扰状态更新，这为后续的并行化奠定了基础。
图解：欧拉离散化 (Euler Discretization) 示意图
上图直观展示了 Mamba-2 中采用的欧拉离散化方法及其局限性：
- 零阶保持 (ZOH)：欧拉法（上图左侧部分）假设在每个时间步长 Δ Delta Δ 内，输入函数（蓝色曲线）保持不变。这相当于用一系列矩形（Step Function）来近似原始的光滑曲线。
- 近似误差：图中橙色阴影区域代表了近似值与真实值之间的误差。由于采用了简单的矩形近似，当输入变化较快或步长 Δ Delta Δ 较大时，这种误差会非常显著，导致模型丢失高频信息或动态捕捉不准确。
- 对比梯形法：作为参考，图右侧展示了梯形法则（Trapezoidal Rule，一阶保持），它使用梯形来拟合曲线，能更精确地追踪输入变化，从而大幅减小了阴影部分的误差。这正是 Mamba-3 引入梯形离散化的核心动力。
2.2 Mamba-2 的架构遗产：SSD 与矩阵形式
Mamba-2 的革命性贡献在于提出了结构化状态空间对偶性 (Structured State Space Duality, SSD)。它证明了上述线性递归（RNN 模式）在数学上严格等价于一种受限的线性注意力机制（Attention 模式）。
2.2.1 矩阵形式 (The Matrix Form)
Dao & Gu (2024) 指出，将递归公式在时间维度展开，可以写成一个巨大的矩阵乘法：
Y = ( L ⊙ C B ⊤ ) X Y = (L odot C B^ op) X Y=(L⊙CB⊤)X
其中：
- X ∈ R T × D , Y ∈ R T × D X in mathbb{R}^{T imes D}, Y in mathbb{R}^{T imes D} X∈RT×D,Y∈RT×D 是输入和输出序列（ T T T 为序列长度， D D D 为模型维度）。
- L ∈ R T × T L in mathbb{R}^{T imes T} L∈RT×T 是结构化掩码矩阵 (Structured Mask Matrix)，它编码了 SSM 的因果衰减特性。
- C ∈ R T × N , B ∈ R T × N C in mathbb{R}^{T imes N}, B in mathbb{R}^{T imes N} C∈RT×N,B∈RT×N：此处 B , C B, C B,C 表示将所有时刻的投影参数 B t , C t B_t, C_t Bt,Ct 在时间维度上堆叠而成的矩阵。
- ⊙ odot ⊙：表示哈达玛积 (Hadamard Product)，即矩阵的元素对应相乘（Element-wise Multiplication）。
- C B ⊤ ∈ R T × T C B^ op in mathbb{R}^{T imes T} CB⊤∈RT×T：这是未经掩码的“注意力分数”矩阵（若 N = 1 N=1 N=1 则为标量积），类似于 Attention 中的 Q K ⊤ QK^ op QK⊤。
对于 Mamba-2，这个掩码矩阵 L L L 是一个半可分离矩阵 (Semi-Separable Matrix)，其下三角元素由连续的遗忘门乘积组成：
L i j = { ∏ k = j + 1 i α k if i > j 1 if i = j 0 if i < j L_{ij} = egin{cases} prod_{k=j+1}^i lpha_k & ext{if } i > j 1 & ext{if } i = j 0 & ext{if } i < j end{cases} Lij=⎩ ⎨ ⎧∏k=j+1iαk10if i>jif i=jif i<j
2.2.2 线性注意力视角
如果不考虑掩码
L
L
L 的衰减结构，仅看
C
B
⊤
X
CB^ op X
CB⊤X，这正是线性注意力（Linear Attention）的核心形式。
在 SSD 框架下，Mamba-2 建立了如下映射：
- C → Query (Q) C ightarrow ext{Query (Q)} C→Query (Q)
- B → Key (K) B ightarrow ext{Key (K)} B→Key (K)
- X → Value (V) X ightarrow ext{Value (V)} X→Value (V)
- L → Causal Mask (with decay) L ightarrow ext{Causal Mask (with decay)} L→Causal Mask (with decay)
Mamba-2 Output ≈ Masked ( Q K ⊤ ) V ext{Mamba-2 Output} pprox ext{Masked}(QK^ op)V Mamba-2 Output≈Masked(QK⊤)V
这种对偶性使得 Mamba-2 在训练时可以使用 Tensor Cores 高效计算矩阵乘法（类似 Attention），而在推理时切换回递归模式（RNN）以实现 O ( 1 ) O(1) O(1) 缓存。
图解：SSD 掩码矩阵 ( M mathcal{M} M) 的半可分离结构
上图揭示了 SSD 算法的核心——结构化掩码矩阵 (Structured Mask Matrix)，即公式中的 L L L（图中标记为 M mathcal{M} M）：
- 动态衰减 (Dynamic Decay)：不同于标准 Transformer 中全为 1 的因果掩码（Causal Mask）， M mathcal{M} M 的下三角元素由状态转移系数 A t A_t At（即 α t lpha_t αt）的连乘积构成。例如， M t , s mathcal{M}_{t,s} Mt,s 表示从时间步 s s s 到 t t t 的累积衰减。这种结构使得模型能够根据输入内容动态调整对历史信息的“关注度”。
- 1-半可分离性 (1-Semiseparable)：这是 Mamba-2 算法加速的关键数学属性。矩阵中任何不包含主对角线的矩形子块（如图中绿色方框所示），其秩（Rank）严格为 1。这意味着长距离的相互作用可以被高效分解和压缩。
- 计算复杂度：利用这种半可分离结构，SSD 算法避免了显式构建巨大的 T × T T imes T T×T 矩阵，而是通过分块分解算法将注意力计算的复杂度从 O ( T 2 ) O(T^2) O(T2) 降低到了 O ( T ) O(T) O(T)，实现了线性时间的训练与推理。
2.3 现有模型的局限：实数特征值的桎梏
尽管 Mamba-2 在效率上取得了巨大成功，但其为了硬件亲和性（Hardware-friendliness）所做的简化，在 Mamba-3 的视角看来，成为了限制模型能力的枷锁。
2.3.1 简化的代价
在 Mamba-2 中，状态转移参数 A t A_t At 被严格限制为实数标量（Real Scalar）。这意味着离散化后的状态转移系数 α t = e Δ A lpha_t = e^{Delta A} αt=eΔA 只能是正实数（通常在 0 到 1 之间）。
2.3.2 无法表达旋转动力学
根据线性代数理论，实数特征值对应的动力学行为只能是：
- 收敛（ ∣ λ ∣ < 1 |lambda| < 1 ∣λ∣<1）：遗忘历史。
- 发散（ ∣ λ ∣ > 1 |lambda| > 1 ∣λ∣>1）：不稳定。
- 保持（ ∣ λ ∣ = 1 |lambda| = 1 ∣λ∣=1）：记忆。
它唯独无法表达旋转（Rotation）或震荡（Oscillation）。要表示旋转，特征值必须是复数（ e i θ = cos θ + i sin θ e^{i heta} = cos heta + isin heta eiθ=cosθ+isinθ）。
案例：奇偶校验任务 (The Parity Task)
考虑一个简单的任务：计算输入比特流中 1 的个数是奇数还是偶数（即
∑
x
t
m
o
d
2
sum x_t mod 2
∑xtmod2）。
- 这是一个典型的状态机问题，状态需要在两个状态之间循环翻转。
- 定理 (Grazzi et al., 2024)：限制为实数特征值的 SSM 无法解决奇偶校验问题，甚至无法解决最简单的 2-state 自动机问题。
这就是 Mamba-3 引入复数域 SSM 的根本动机：我们需要找回被 Mamba-2 丢弃的“虚部”，以赋予模型更强大的逻辑推理和状态追踪能力，同时又不能牺牲推理效率。
📚 参考文献与扩展阅读
- Mamba-2: Dao & Gu, “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality”, ICML 2024.
- 理论缺陷: Grazzi et al., “The Expressive Capacity of State Space Models: A Formal Language Perspective”, arXiv 2024.
