深入理解 Mamba：从SSM到SSSM的公式推导

Dyanic

已于 2025-06-23 17:17:43 修改

阅读量1.6k

点赞数 46

CC 4.0 BY-SA版权

文章标签： transformer 人工智能

于 2025-06-23 17:15:34 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_43858206/article/details/148840115

一、为什么要用Mamba

在自然语言处理、时间序列分析等领域，序列建模一直面临着两大核心挑战：长距离依赖建模能力与计算效率。传统的循环神经网络（RNN）虽具备理论上的长序列处理能力，但由于梯度消失问题，实际有效建模长度通常局限在数百步；Transformer 通过自注意力机制实现了全局依赖建模，却付出了 $O (n^{2})$ 的时间复杂度代价，难以处理超长序列（如音频、基因组数据）。

Mamba（全称Mamba: A General Purpose Sequence Model）的出现打破了这一困境。作为一种基于状态空间模型（SSM）的新型序列模型，Mamba 兼具以下核心优势：

线性时间复杂度：通过状态空间的递推特性，将序列处理复杂度降至 $O (n)$ ，相比 Transformer 更适合处理数千甚至数万长度的序列。
长距离依赖建模能力：状态空间的动态演化机制可捕获序列中的长期依赖关系，理论上支持无限长度的上下文建模。
硬件加速友好：相比 RNN 的递归结构，Mamba 的矩阵运算更易利用 GPU/TPU 的并行计算能力，实现高效推理。
统一框架灵活性：可无缝兼容 CNN、RNN 等模型的特性，通过参数调整实现不同序列建模场景的适应性。

二、Mamba底层机制

2.1状态空间模型（State Space Model）（SSM）

状态空间模型起源于控制理论，其核心思想是将高阶动态系统分解为一阶状态方程的矩阵形式，便于分析与控制。以经典的弹簧振子系统为例，我们可以直观理解这一抽象过程：
例子引入：

一个弹簧振子系统如图所示：
在这里插入图片描述

系统包含弹簧（弹性系数 $k$ ）、质量块（质量 $m$ ）、阻尼器（阻尼系数 $b$ ），外力 $f (t)$ 作用下，质量块位移 $P (t)$ 满足二阶微分方程：
$m\frac{d^{2}P(t)}{d(t)}+b\frac{dP(t)}{dt}+kP(t)=f(t)$

$令一阶状态变量h_1(t)=P(t)（位移）,h_2(t)=\frac{dP(t)}{dt}（速度）$
原式变为：
$m\frac{h_2(t)}{dt}+bh_2(t)+kh_1(t)=f(t)$
$\frac{h_2(t)}{dt}=\frac{1}{m}(f(t)-bh_2(t)-kh_1(t))$
用矩阵形式表示为：
$\frac{d\begin{bmatrix} h_1(t))\\h_2(t) \end{bmatrix}}{dt}=\begin{bmatrix} \frac{dh_1(t)}{dt}\\ \frac{dh_2(t)}{dt} \end{bmatrix}=\begin{bmatrix} 0 &1 \\ -\frac{k}{m} & -\frac{b}{m} \end{bmatrix}\begin{bmatrix} h_1(t)\\h_2(t) \end{bmatrix}+\begin{bmatrix} 0\\\frac{1}{m} \end{bmatrix}f(t)$
其中输出方程为
$P(t)=\begin{bmatrix} 1&0 \end{bmatrix}\begin{bmatrix} h_1(t)\\h_2(t) \end{bmatrix}$
我们可以观察到，原来的二阶的状态方程，通过转换变成一个一阶的矩阵表达，这个就是SSM的作用，其中 $h_1(t))$ , $h_2(t)$ 就称为状态。定义状态矩阵 $A$ 输入矩阵 $B$ 、输出矩阵 $C$
$A=\begin{bmatrix} 0 &1 \\ -\frac{k}{m} & -\frac{b}{m} \end{bmatrix},B=\begin{bmatrix} 0\\\frac{1}{m} \end{bmatrix},C=\begin{bmatrix} 1&0 \end{bmatrix}$
$f (t)$ 其实就是整个系统的输入， $P (t)$ 其实就是整个系统的输出，我们引入到神经网络中，做一个替换，使得 $x (t) = f (t)$ , $y (t) = P (t)$ ,同时令 $h(t)=[h1(t))h2(t)]h(t)=\begin{bmatrix} h_1(t))\\h_2(t) \end{bmatrix}$ ,则连续时间下的 SSM 可表示为：
$\frac{dh(t)}{dt}=Ah(t)+Bx(t)$
$y (t) = C h (t)$
我们就可以得到了Mamba论文中的第一个公式的推导过程。但是我们可以看到上面的系统是基于一个连续的系统，而我们的神经网络的构建过程是一个离散化的过程，因此需要将连续 SSM 转换为离散形式。Mamba 采用零阶保持法（Zero-Order Hold, ZOH），其核心思想是在离散时间间隔内假设输入值保持不变（如 $x(t_k)$ 在区间 $t_k, t_{k+1}]$ 内恒定）。具体如下图所示：
在这里插入图片描述

为了实现离散化过程，我们在等式两边同时乘 $e^{-At}$ 等式变为：
$e^{-At}\frac{dh(t)}{dt}=e^{-At}Ah(t)+e^{-At}Bx(t)$
$e^{-At}\frac{dh(t)}{dt}-e^{-At}Ah(t)=e^{-At}Bx(t)$
接着我们可以发现其实
$\frac{de^{-At}h(t)}{dt}=e^{-At}\frac{dh(t)}{dt}-e^{-At}Ah(t)$
那么我们的等式变为：
$\frac{de^{-At}h(t)}{dt}=e^{-At}Bx(t)$
我们对两边同时求积分可以得到：
$\int_{t_{k}}^{t_{k+1}} \frac{de^{-At}h(t)}{dt}=\int_{t_{k}}^{t_{k+1}}e^{-At}Bx(t)dt$
由于在零阶保持法中， $x (t)$ 是一个与 $t$ 无关的常量,且用 $t_k$ 的状态来表示，因此在积分中可以直接提取出来因此我们的积分可以进一步简化为：
$e^{-At}h(t)|_{t_k}^{t_{k+1}}=\int_{t_k}^{t_{k+1}}e^{-At}dtBx(t_k)$
$e^{-At_{k+1}}h(t_{k+1})-e^{-At_{k}}h(t_{k})=-\frac{1}{A}e^{-At_{k}}|_{t_k}^{t_{k+1}}Bx(t_k)$
$e^{-At_{k+1}}h(t_{k+1})=e^{-At_{k}}h(t_{k})-\frac{1}{A}e^{-At_{k}}|_{t_k}^{t_{k+1}}Bx(t_k)$
$e^{-At_{k+1}}h(t_{k+1})=e^{-At_{k}}h(t_{k})-\frac{1}{A}e^{-At_{k+1}}Bx(t_k)+\frac{1}{A}e^{-At_{k}}Bx(t_k)$

等式两边同时乘以 $e^{At_{k+1}}$ 可以得到：
$h(t_{k+1})=e^{A(t_{k+1}-t_k)}h(t_{k})-\frac{1}{A}Bx(t_k)+\frac{1}{A}e^{A(t_{k+1}-t_{k})}Bx(t_k)$
令 $tk+1−tk=Δt_{k+1}-t_k=\Delta$
$h(t_{k+1})=e^{A\Delta }h(t_{k})-\frac{1}{A}Bx(t_k)+\frac{1}{A}e^{A\Delta }Bx(t_k)$
进一步化简为：
$h(t_{k+1})=e^{A\Delta }h(t_{k})+A^{-1}(e^{A\Delta }-1)Bx(t_k)$
令 $Aˉ=eAΔ,Bˉ=A−1(eAΔ−1)B\bar{A}=e^{A\Delta },\bar{B}=A^{-1}(e^{A\Delta }-1)B$
代入得：
$h(t_{k+1})=\bar{A}h(t_{k})+\bar{B}x(t_k)$
我们就得到了第二个离散化的公式。

2.2 Mamba对时变系统的核心改进：从固定参数到动态适应

在传统状态空间模型中，离散化后的参数 $Aˉ\bar{A}$ 和 $Bˉ\bar{B}$ 由系统固有属性（如弹簧振子的质量、阻尼系数）决定，属于时不变参数。但实际序列数据（如语音、金融时序）常呈现时变特性（即系统动态特性随时间变化），例如：

语音信号中不同频段的能量分布随发音时刻变化；
金融市场中资产相关性受突发事件影响而波动。

Mamba在保留基础矩阵 $A$ 固定性的前提下，通过动态调节其他参数实现时变适配，将状态转移方程重构为：
$h(t_{k+1}) = A h(t_k) + B_t x(t_k)$
$y(t_k) = C_t h(t_k)$
其中 $A$ 为固定基础矩阵， $B_t$ 、 $C_t$ 为动态参数， $Δt\Delta_t$ 为时变离散化步长，以下是各参数的数学机制解析：

对 $A$ 的解释说明

在传统时不变框架下， $A$ 作为核心状态转移算子，其参数刻画系统固有动态，如物理系统中惯性、阻尼主导的状态演化规律，一旦离散化确定便保持恒定，无法响应序列数据的时变特性。而在适配时变系统的改进逻辑里， $A$ 虽仍承载基础状态传播的 “默认规则”，但需与其他动态模块协同，为 $B_t$ 、 $Δ_t$ 等时变调节提供稳定的状态演化基底。其固定性成为时变适配的 “基准参考系”—— 如同物理定律恒定不变，却能通过外力（时变参数）调节系统输出。这种设计保障状态空间模型在动态调整中仍具备可追溯的核心运算逻辑，避免因过度时变导致状态演化失序。例如，在语音处理中， $A$ 可预设为捕捉声波传播的基础动力学，而时变参数负责适配不同发音时刻的特征变化。

对 $B_t$ 的解释说明

$B_t$ 作为时变输入映射算子，数学表达式为：
$B_t = \text{Project}(\sigma(S_t) \otimes \bar{B})$
其中：

$St=Ws⋅[h(tk);x(tk)]S_t = W_s \cdot [h(t_k); x(t_k)]$ （状态-输入联合特征）；
$σ(⋅)\sigma(\cdot)$ 为激活函数（如Sigmoid），实现输入权重动态筛选；
$⊗\otimes$ 表示元素级乘法， $Project(⋅)\text{Project}(\cdot)$ 为线性投影矩阵。

$B_t$ 作为时变输入映射算子，突破传统固定 $Bˉ\bar{B}$ 的限制，依据时刻 $t$ 的序列数据特征动态调整。面对语音序列，它可感知不同发音时刻频段能量分布差异，适配调制输入与状态的耦合强度；针对金融时序，能捕捉突发事件冲击下资产相关性波动，实时重塑新信息向状态空间的注入方式。通过与选择机制（Selection Mechanism）、投影（Project）模块联动， $B_t$ 实现输入信息的时变筛选与加权投影，让状态更新精准响应数据的瞬时动态，成为模型适配时变特性的 “输入动态调节器” 。

对 $C_t$ 的解释说明

$C_t$ 通过动态解耦状态表征生成输出，数学形式为：
$C_t = \text{MLP}(h(t_k)) \cdot C_{\text{base}}$
其中：

$CbaseC_{\text{base}}$ 为固定基础输出矩阵（如弹簧振子系统中 $C = [1, 0]$ ）；
$MLP(⋅)\text{MLP}(\cdot)$ 为多层感知机，根据当前状态 $h(t_k)$ 生成时变调节权重。

$C_t$ 作为时变输出映射算子，区别于传统固定输出映射，需适配状态空间因时变调整后的内部表征。在时变系统中，随 $A$ 主导的状态演化、 $B_t$ 驱动的输入融入，状态空间的信息分布持续动态重构， $C_t$ 需实时解耦状态特征，生成贴合当前时刻数据特性的输出 $y_t$

对 $Δt\Delta_t$ 的解释说明

$Δt\Delta_t$ 为时变离散化步长，计算逻辑为：
$\Delta_t = \text{Discretize}(v_t, \Delta_{\text{min}}, \Delta_{\text{max}})$
其中：

$vt=∥∇x(tk)∥2v_t = \|\nabla x(t_k)\|_2$ （输入信号变化率）；
$Discretize(⋅)\text{Discretize}(\cdot)$ 为分段函数：
$\Delta_t = \begin{cases} \Delta_{\text{min}} & v_t > \tau \\ \Delta_{\text{max}} & v_t < \epsilon \\ \Delta_{\text{base}} + \beta \cdot v_t & \text{otherwise} \end{cases}$