mamba 学习

Posted on 2025-09-24 Edited on 2025-12-03

对Mamba的一点学习

pre Knowledge

时不变性

如果在输入信号上有一个时移，那么输出信号上会产生相同的时移，那么这个系统就是时不变的。也就是说输入信号为 $x(t - t_0)$ 时，输出信号为 $y(t - t_0)$。

如何去理解呢？系统的特性和行为不随时间改变，信号的当前状态与时间无关，这个系统就是时不变的。举例来说，一个电路你今天做一次实验，和明天做一次实验，得到的结果是一样的，与时间无关。你的输入不随时间发生变化，第一次计算之后掌握的特征储存下来，将来使用的时候就不用重新计算特征了。

例：如果我们考虑一个连续时间系统，定义如下：

\[ y(t) = \sin[x(t)] \]

为了确认这个系统是时不变的，就必须判定对于任何输入和任何时移 $t_0$，时不变性是否成立。令 $x_1(t)$ 是系统的任意输入，并令

\[ y_1(t) = \sin[x_1(t)] \]

是其相应的输出。

然后，考虑将 $x_1(t)$ 时移作为第二个输入：

\[ x_2(t) = x_1(t - t_0) \]

对于这个系统的输出是：

\[ \begin{aligned} y_2(t) &= \sin[x_2(t)] = \sin[x_1(t - t_0)] \\ y_1(t - t_0) &= \sin[x_1(t - t_0)] \\ y_2(t) &= y_1(t - t_0) \end{aligned} \]

即证时不变性。

要证明一个系统是时变性的话，仍然采用的是上面的方法，观察时移前后的输出并对比即可，或者找一个条件，使时不变性不成立(找反例)。例如 $y(t) = x(2t)$ 就是一个时变系统，尺度变换下不同时间的输入会导致输出的不同。

SSM 状态空间模型

状态空间模型是一种用不可直接观测的隐含状态（hidden state）来描述系统动态行为的数学模型。系统的观测输出由该隐含状态演化而来。基本结构

SSM 通常由两个核心方程组成： 1. 状态方程（State Equation）描述隐含状态如何随时间演化（以离散时间为例）： [ t = {t-1} + x_t ] 2. 观测方程（Observation/Output Equation）描述如何从隐含状态生成可观测输出： [ y_t = _t + D x_t ]

其中： (x_t)：时刻 (t) 的输入 (_t)：时刻 (t) 的隐含状态（通常维度远高于输入） (y_t)：时刻 (t) 的输出 (, , , D)：模型参数（在线性 SSM 中为固定矩阵/标量）

MAMBA

首先我们对于之前的状态空间模型应用到实际的计算机领域，我们注意到，在SSM的过程中有 \[ \begin{cases} \dfrac{dh(t)}{dt} = A h(t) + B x(t) \\ y(t) = C h(t) \end{cases} \] 这是连续情况的SSM。所以为了将其离散化，我们需要对其进行采样。设采样间隔为$ $

\[ \dfrac{dh(t)}{dt} = A h(t) + B x(t) \] 的解为 \[ h(t) = e^{At} h(0) + \int_0^t e^{A(t-\tau)} B x(\tau) d\tau \] 其中$e^{At}$是矩阵指数函数 \[ e^{At} = \sum_{n=0}^{\infty} \frac{(At)^n}{n!} \]

假设在每个采样时刻 $ t_k = k$，输入 $ x(t) $ 在区间 $ [t_k, t_{k+1}] $ 内保持恒定，即：

\[ x(t) = x_k \quad \text{for } t \in [k\Delta, (k+1)\Delta] \]

那么从 $ t = k$ 到 $ t = (k+1)$ 的状态演化为：

\[ h((k+1)\Delta) = e^{A\Delta} h(k\Delta) + \int_0^\Delta e^{A(\Delta - \tau)} B x_k d\tau \]

令 $ ' = - $，则积分变为：

\[ \int_0^\Delta e^{A(\Delta - \tau)} B x_k d\tau = \left( \int_0^\Delta e^{A\tau'} d\tau' \right) B x_k \]

而：

\[ \int_0^\Delta e^{A\tau} d\tau = A^{-1}(e^{A\Delta} - I) \quad \text{（若 } A \text{ 可逆）} \]

所以：

\[ h_{k+1} = e^{A\Delta} h_k + A^{-1}(e^{A\Delta} - I) B x_k \]

令$ {A} = e^{A} $ $ {B} = _0^e{A} dB = A^{-1}(e{A} - I) B $ （若 $ A $ 可逆）于是得到离散系统： \[ \begin{cases} h_{k+1} = \bar{A} h_k + \bar{B} x_k \\ y_k = C h_k \end{cases} \] ### 离散 SSM 的卷积形式与时不变性

在上一小节中，我们已经得到了离散化后的状态空间模型（假设采样间隔为常数 $\Delta$）：

\[ \begin{cases} h_{k+1} = \bar{A} h_k + \bar{B} x_k \quad y_k = C h_k \end{cases} \]

其中 $\bar{A} = e^{A\Delta},\quad \bar{B} = \left(\displaystyle\int_0^\Delta e^{A\tau} d\tau\right) B = A^{-1}(e^{A\Delta}-I)B$（若 (A) 可逆）。

从递推式出发，逐步展开前几步状态： \[ \begin{aligned} h_1 &= \bar{A} h_0 + \bar{B} x_0 \\ h_2 &= \bar{A} h_1 + \bar{B} x_1 = \bar{A}^2 h_0 + \bar{A}\bar{B} x_0 + \bar{B} x_1 \\ h_3 &= \bar{A} h_2 + \bar{B} x_2 = \bar{A}^3 h_0 + \bar{A}^2 \bar{B} x_0 + \bar{A} \bar{B} x_1 + \bar{B} x_2 \end{aligned} \]

可以归纳得到一般形式：

\[ h_k = \bar{A}^k h_0 + \sum_{j=0}^{k-1} \bar{A}^{k-1-j} \bar{B}x_j \]

令初始状态 (h_0 = 0)，则进一步简化为：

\[ h_k = \sum_{j=0}^{k-1} \bar{A}^{k-1-j} \bar{B}x_j \]

可以看到，$h_k$ 是对过去所有输入 ${x_j}_{j=0}^{k-1}$ 的“线性加权和”，权重矩阵由 $\bar{A}^{k-1-j}\bar{B}$ 决定。越久远的输入，前面会乘上更多次 $\bar{A}$，在一定条件下会产生“记忆衰减”的效果。

2. 输出写成卷积形式

根据观测方程： \[ y_k = C h_k \]

将上式中的 $h_k$ 代入(仍令 ($h_0$ = 0)) \[ y_k=Ch_k=C \sum_{j=0}^{k-1} \bar{A}^{k-1-j} \bar{B},x_j \\ = \sum_{j=0}^{k-1} C \bar{A}^{k-1-j} \bar{B},x_j \]

定义“卷积核矩阵”（只与时间间隔 $\ell$ 有关）： \[ K_\ell := C \bar{A}^{\ell} \bar{B}, \quad \ell = 0,1,2,\dots \]

令 $\ell = k-1-j$，则有：

\[ y_k = \sum_{\ell=0}^{k-1} K_\ell x_{k-1-\ell} \]

这样就转换成了卷积的形式，只不过这里的卷积核 $K_\ell$ 是矩阵，输入 $x_k$ 是向量，输出 $y_k$ 也是向量。

现在我们这里的卷积核$K_\ell$就只依赖于时间差$\ell$,不依赖于绝对时刻(k). 系统满足输入整体时移，输出整体同样时移的性质，是一个线性时不变(LTI)系统。接下来将引入可变时间步长 ### 引入可变步长 $\Delta_k$

上面假设采样间隔为常数 $\Delta$。现在我们考虑更一般的情况：现在，每一步的采样间隔可以不同，记为 $\Delta_k$。

1. 连续系统仍为常系数

连续时间的状态空间模型仍然是：

\[ \dfrac{d h(t)}{dt} = A h(t) + B x(t), \quad y(t) = C h(t) \]

这里的 (A,B,C) 仍是常系数矩阵。

不同的是，我们现在选取一列不等间隔的采样时刻：

\[ t_0,\ t_1,\ t_2,\ \dots,\quad t_{k+1} = t_k + \Delta_k \]

在每个区间 $[t_k,t_{k+1}]$ 内，仍然采用零阶保持（Zero-Order Hold）假设：

\[ x(t) \approx x_k,\quad t\in[t_k,t_{k+1}] \]

2. 在每个区间内做离散化

从 $t=t_k$ 到 $t=t_{k+1} = t_k + \Delta_k$ 的状态演化可以写为（类似之前的推导，只是用 $\Delta_k$ 替换 $\Delta$）：

\[ h(t_{k+1}) = e^{A\Delta_k} h(t_k)* \left(\int_0^{\Delta_k} e^{A(\Delta_k - \tau)} d\tau\right) B x_k \] 于是定义 \[ \begin{aligned} \bar{A}_k &= e^{A\Delta_k} \ \bar{B}_k &= \left(\int_0^{\Delta_k} e^{A\tau} d\tau\right) B = A^{-1}(e^{A\Delta_k}-I)B \quad (\text{若 } A \text{ 可逆}) \end{aligned} \]

离散系统变为： \[ \begin{cases} h_{k+1} = \bar{A}_k h_k + \bar{B}_k x_k \\ y_k = C h_k \end{cases} \] 与固定步长的情形相比，这里多了一个下标 $k$： * $\bar{A}_k,\bar{B}_k$ 不再是常数矩阵，而是随步长 $\Delta_k$ 变化的矩阵； * 因为每一步的有效“离散参数”不同，所以该离散系统已经不再是时不变系统，而是一个线性时变系统（LTV）。

$\Delta_k$ 作为输入函数：Mamba 的 Selective SSM

只要让采样步长 $\Delta_k$变化，离散系统的矩阵 $\bar{A}_k,\bar{B}_k$ 就会随$k$变化，从而得到一个线性时变系统。而Mamba让这种变化不只是随“时间步数”改变，而是随当前输入$x_k$改变，实现“选择性”。 #### 1. 让 $\Delta_k$ 依赖于输入 $x_k$ 定义一个小的神经网络（例如线性层 + 激活函数）：

$ k = (x_k) $

例如可以取 $ k = (Wx_k + b_) $

用$\phi_ \Delta(\cdot)$保证$\Delta_k > 0$。

然后，将 $\Delta_k$ 代入前一节中 $\bar{A}_k,\bar{B}_k$ 的定义：

\[ \begin{aligned} \bar{A}*k &= e^{A\Delta_k} = e^{A\phi*\Delta(x_k)} \\ \bar{B}_k &= \left(\int_0^{\Delta_k} e^{A\tau} d\tau\right) B = A^{-1}(e^{A\Delta_k}-I)B \end{aligned} \]

此时状态更新为： \[ h_{k+1} = \bar{A}_k(x_k), h_k + \bar{B}_k(x_k)x_k \]

可以看到，这里的“有效离散参数” ({A}_k,{B}_k) 已经是输入驱动的函数：

\[ \bar{A}_k = \bar{A}(\phi*\Delta(x_k)),\quad \bar{B}_k = \bar{B}(\phi*\Delta(x_k)) \]

从“连续时间常系数系统”出发，通过引入输入相关的步长 $\Delta_k$，我们得到了一个随输入变化的离散矩阵系统，这正是 Mamba 论文中所谓 Selective State Space Model 的重要一环。

2. 更一般的 Selective SSM 形式

在 Mamba 中，不仅 $\Delta_k$ 可以依赖于 $x_k$，还可以让 $A_k,B_k,C_k$ 也依赖于 $x_k$。抽象地，可以写成：

\[ \begin{cases} (A_k,B_k,C_k,\Delta_k) = g(x_k) \\ \bar{A}_k = \bar{A}(A_k,\Delta_k),\quad \bar{B}*k = \bar{B}(A_k,\Delta_k,B_k) h*{k+1} = \bar{A}_k h_k + \bar{B}_k x_k \\ y_k = C_k h_k \end{cases} \]

其中 $g(\cdot)$ 是由神经网络（如线性层、MLP）构成的可学习函数。

这时，系统的状态演化与输出不仅取决于过去的状态，也显式地取决于当前输入内容所决定的参数 $(A_k,B_k,C_k,\Delta_k)$——这就是 “Selective” 一词的来源：模型可以根据输入内容，选择性地记忆或遗忘信息。