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

君科程序定做

于 2025-09-22 21:44:29 发布

阅读量1.1k

点赞数 17

CC 4.0 BY-SA版权

文章标签： python

本文链接：https://blog.csdn.net/kangkangluoluo/article/details/151972596

0. 任务与输入

任务举例：图像分类（也可换成分割/检测，最后的“Head”和损失不同）。
输入：I ∈ ℝ^{512×512×3}。

1. Patch 化与嵌入（把 2D 图像变成序列）

切块：以 P=16 切 patch → 32×32=1024 个 patch。
展平+线性映射（等价于 ViT 的 Patch Embedding / Conv16×16, stride 16）：
- 每个 patch 展平后是 16×16×3=768 维；
- 过一层线性映射到模型维度 d（常见 512/768）：
- 得到序列 X ∈ ℝ^{N×d}，N=1024。
位置编码/顺序：选择扫描顺序（行优先、蛇形、Hilbert 曲线等）并加上可学/正余弦位置嵌入（可选）。

小结：到这一步，图像 → token 序列 X(1…N)。

2. Mamba Block 的前向（用 SSM 取代注意力）

一个 Mamba block 的骨架（与 Transformer block 类似：Norm → 主体 → 残差 → MLP）：

2.1 预处理与短卷积（可选）

X ← LayerNorm(X)
X ← Conv1D(kernel=3) （给局部交互一个快捷通道，很多实现会有）

2.2 线性投影，得到选择性参数（这一步“算出 QKV”的类比，但 Mamba 不是 QKV）：
从同一个 X 通过不同线性层得到 4 组量（逐 token）：

U = X W_u（输入通道，像 SSM 的驱动信号）
Δ = Softplus(X W_Δ + b_Δ)（正的步长/时间间隔）
B = X W_B、C = X W_C（输入/输出耦合向量，随输入而变，这就是“Selective”）

还有一组全局可学的状态参数（与层相关而不是与 token 相关）：

Λ（对角稳定矩阵，等价连续系统的 A），通常参数化为负值 Λ=-Softplus(·) 保稳定。

2.3 选择性状态扫描（Selective Scan）
把序列按选定顺序扫描，递推隐藏状态 h_t ∈ ℝ^{m}（m 为状态维度，常小于/等于 d）：

离散稳定更新（简化版写法；实际实现用对角指数/前缀和加速）：
- A_t = exp(Δ_t Λ) （逐 token 的“衰减”系数，元素乘）
- h_t = A_t ⊙ h_{t-1} + B_t ⊙ U_t
- y_t = C_t ⊙ h_t
  其中 ⊙ 为逐元素乘。
并行性：实现上使用“选择性扫描/前缀乘积”把上面递推变为 O(N) 的并行核（不像注意力 O(N²)）。
方向性：很多视觉实现用多方向扫描（→、←、↓、↑ 或多条蛇形/Hilbert），再把各方向输出拼/加，增强 2D 感知。

2.4 投影与残差

Z = y W_o
X ← X + Z（残差）
接一个 Gated-MLP / FFN：X ← X + MLP(LN(X))

堆叠 L 个 Mamba blocks（如 12/24 层），得到最终序列特征 X̂ ∈ ℝ^{N×d}。

3. Head 与损失

分类：g = MeanPool(X̂)（或取第一个 token），logits = g W_cls，
损失 L = CrossEntropy(logits, label)。
语义分割：把 X̂ 还原成 32×32×d，接上解码头/上采样（UPerNet/简单卷积插值），像素级 CE/Dice。
检测/实例分割：作为 backbone，接 FPN/解码头，正常的检测损失。

4. 反向传播 & 优化

核心可学习参数
- Patch/Stem（卷积或线性）
- 各层的 W_u, W_Δ, W_B, W_C, W_o
- 稳定矩阵参数 Λ（对角）以及 MLP 的权重
优化器：AdamW（典型 lr 1e-3/5e-4，wd 0.05），Cosine decay，Warmup。
正则&增强：RandAugment、Mixup、CutMix、Label Smoothing、DropPath（stochastic depth）。
复杂度：时间/显存 O(N)，比 ViT 的 Self-Attention（O(N²)）更省，N=1024 时优势明显；分辨率升高优势更大。

5. 一次前向的“形状流水线”（分类为例）

I:              (512,512,3)
Patchify(16):   (1024, 16*16*3=768)
Linear→d=768:   (1024, 768)
+PosEmbed:      (1024, 768)

Repeat L times:
  LN/Conv1D:    (1024, 768)
  Proj: U,B,C,Δ (1024, 768) each
  SelectiveScan:
      h_t∈ℝ^m   (1024, m)     # m ~ 256/512
      y_t       (1024, m)
  Out proj:     (1024, 768)
  Residual + MLP

GlobalPool:     (768,)
Cls head:       (num_classes,)
Loss: CE

6. 和 ViT/CNN 的“训练核心差异”

与 ViT：不做 QKV 全局两两相似度；改为 选择性状态递推，复杂度从 O(N²) → O(N)，长序列/大图更友好。
与 CNN：不是固定局部卷积核；通过 输入自适应的 B/C/Δ + 稳定的 Λ，在长程上传播与选择信息，同时仍可搭配局部卷积增强细节。
训练核心：学到（1）稳定的长程记忆核 Λ，（2）随内容动态选择/门控的 B,C,Δ，让有用的信息保留并跨很远的 token 传播，最后由任务损失端到端驱动。

小贴士（实践参数）

Patch=16、d=768、层数 L=24、m=256/512，LabelSmoothing=0.1，DropPath=0.1~0.2，224/256/384 输入皆可；512 训练需要更大显存或梯度累积。
遥感/医学等超高分辨率：可用多方向扫描 + 分块滑窗 + 重叠拼接；Mamba 的 O(N) 更适合这类大图。

X指的是啥同一个 X 通过不同线性层得到 4 组量为啥呀

1. X 是什么？

在图像任务里，

经过 Patch Embedding（把 512×512×3 切成 patch → 展平 → 线性映射到 d 维），
得到一个 token 序列：

$[x_1, x_2, ..., x_N], \quad x_i \in \mathbb{R}^d$

这里：

$N$ = patch 数量（比如 32×32=1024），
$d$ = 特征维度（常用 256、512、768）。

你可以把每个 $x_i$ 看作是 一个 patch 的 embedding，类似于句子里的一个词向量。

2. 为什么同一个 X 要过不同的线性层？

这和 Transformer 里「同一个 X 线性投影出 Q、K、V」是类似的思路。
Mamba 的 Selective SSM 里，需要 4 种信息来驱动状态更新：

U = X Wᵤ
- 输入信号，驱动状态转移（相当于“外部刺激”）。
B = X Wᴮ
- 输入耦合向量，决定输入信号如何写入隐藏状态（像门控）。
C = X Wᶜ
- 输出耦合向量，决定隐藏状态如何映射回可见输出。
Δ = Softplus(X WΔ + bΔ)
- 步长参数（必须是正数），相当于「时间尺度/记忆速度」，用来控制指数衰减系数 $At=e−ΔtA_t = e^{-\Delta_t}$ 。

3. 为什么不是只要一个向量？

因为 SSM 的递推公式里需要 多个角色：

状态更新公式（简化版）：

$h_t = A_t \odot h_{t-1} + B_t \odot U_t$

$y_t = C_t \odot h_t$

要计算这个，就必须有：

驱动输入 $U_t$
输入耦合 $B_t$
输出耦合 $C_t$
衰减/步长 $Δt\Delta_t$

所以不能只从 X 得到一个向量，而是要分出不同的功能「通道」。

4. 通俗类比

想象 X 是一行乐谱，乐队要演奏它：

$U$ ：旋律线（音乐的原始内容）
$B$ ：音量控制（音量大/小写入状态）
$C$ ：输出音色（小提琴/钢琴发声方式）
$Δ\Delta$ ：节拍速度（快节奏还是慢节奏，影响“记忆跨度”）

同一份乐谱（X），被不同的“乐器角色”读取出来，才能最终合奏成完整音乐。

✅ 小结：

X = 每个 patch 的特征向量；
同一个 X 投影成 U、B、C、Δ，是因为 SSM 更新公式需要这些不同的量来决定“输入怎么写进状态、状态怎么衰减、输出怎么读出来”。

要不要我帮你画一个 对比图（像 QKV vs U/B/C/Δ），这样你能更直观地看到 Transformer 和 Mamba 在这一步的区别？

「选择性状态扫描」「离散稳定更新」这些听起来很吓人的词，逐一拆开解释

1. 先回顾：Mamba 的核心思想

Mamba 属于 **SSM（状态空间模型）**家族。
它的本质就是：

给定输入序列 $x_1, x_2, ..., x_T$ ，
维护一个「隐藏状态」 $h_t$ ，
让它随时间不断更新，再输出 $y_t$ 。

公式大概像：

$h_t = A_t h_{t-1} + B_t x_t$

$y_t = C_t h_t$

$A_t$ ：控制旧记忆怎么保留/衰减
$B_t$ ：决定输入怎么写进状态
$C_t$ ：决定状态怎么转成输出

2. 什么是选择性状态扫描 (Selective State Space Scanning)

📖 专业解释

这是 Mamba 提出的改进：在每个时间步，不是机械地更新状态，而是根据输入 动态决定更新幅度。

公式变成：

$h_t = e^{-\Delta_t} \odot h_{t-1} + B_t \odot U_t$

这里的 $Δt\Delta_t$ 就是「选择性」的核心，它由输入算出来，能动态调节记忆的快慢。

🧑‍🎓 通俗解释

想象你的大脑记忆：

有些东西（比如「今晚吃什么」）很快就忘掉 → 大的 $Δt\Delta_t$ ，衰减快。
有些东西（比如「女朋友生日」😅）需要长期保存 → 小的 $Δt\Delta_t$ ，衰减慢。

Mamba 就是这样「选择性」地决定：
哪些信息保留长久，哪些信息快速遗忘。

这就是「选择性状态扫描」。

3. 什么是离散稳定更新 (Discretized Stable Update)

📖 专业解释

数学上，SSM 最初是 连续时间系统（微分方程）。
但在计算机里只能处理离散序列（一步一步），所以要「离散化」。

离散化时，如果不小心，可能导致数值不稳定（比如状态发散到无穷大）。
Mamba 采用了 指数衰减形式：

$h_t = e^{-\Delta_t} h_{t-1} + ...$

$e−Δte^{-\Delta_t}$ 永远是正数而且 ≤ 1
保证状态不会爆炸 → 这就是 稳定更新。

🧑‍🎓 通俗解释

想象「热水冷却」：

如果物理模型写错了，可能会算出水越降越热 → 爆炸！
Mamba 的离散稳定更新，就是给每次更新套上一个「冷却系数」 $e−Δte^{-\Delta_t}$ ，保证水温（隐藏状态）只会逐步下降，不会发散。

4. 小总结

选择性状态扫描：动态调节记忆保留/遗忘速度（相当于 LSTM 的遗忘门）。
离散稳定更新：把连续系统转换为离散计算时，保证不会数值爆炸（用指数衰减确保稳定）。

👉 所以，Mamba 的厉害之处在于：

通过「选择性扫描」智能地决定记忆时长；
通过「稳定更新」保证训练不会崩掉。

为啥计算机里只能处理离散序列？

好的，咱们一步步解答你的问题，同时再次简明扼要地阐述 Mamba 的核心思想。

1. 为啥计算机里只能处理离散序列？

📖 专业解释：

计算机本质上是通过数字来进行计算的，它只能处理离散的数据，因为计算机的硬件（尤其是处理器）是基于离散信号（电路的开关状态）来工作的。
连续的数学模型（如微分方程）在现实世界中能够描述很多现象，但这些是连续的值（如温度、速度等）。然而，计算机处理的是有限精度的数值，只能处理离散数据（例如，离散的时间步、离散的空间位置）。

🧑‍🎓 通俗解释：

可以想象你想做 模拟水的流动：

现实世界的水流动：它是连续的，可以随时变化。
计算机模拟水流动：计算机没有办法处理这种无限细的变化，因为计算机是由开关组成的，每一时刻只能给出一个“数字”。因此，必须把连续的水流动分成一个个离散的时间步来模拟。

为什么要用离散化？

离散化让计算机可以通过一个个 有限的步骤 来逼近连续过程。
这种“分成小块来处理”的方式，既能保证计算的可行性，又能近似实际问题的连续动态。

2. Mamba 的核心思想再阐述

Mamba 是一种 选择性状态空间模型（Selective State Space Model，SSM），其核心思想是通过 状态空间更新 来建模输入数据（例如图像）的序列依赖关系。

📖 核心原理：

状态空间模型：Mamba 本质上是一种动态系统模型，用来 逐步更新内部状态，从而捕捉序列数据的长期依赖。
状态更新：Mamba 使用 线性递推（LTI） 更新当前状态，这意味着每一步都会用前一步的状态来更新当前状态，通过一个选择性机制 动态调整记忆的“保留”与“遗忘”。这个选择性机制类似于 LSTM 的遗忘门。
选择性机制：在每一步，模型不仅更新状态，还可以 选择性地“记住”或“遗忘”某些部分信息。通过 自适应的衰减步长 $Δ\Delta$ ，Mamba 动态决定在每个时间步中如何调整状态，这个机制可以让模型学会哪些信息该保留，哪些信息可以丢弃。

🧑‍🎓 简单理解：

输入数据（比如图像的 patch 或者序列的每个元素）依次传入模型。
状态：模型通过内部状态 $h_t$ 来记住每个时间步的历史信息。它不断更新这个状态，来捕捉整个序列的长程依赖。
选择性更新：在每一步，模型通过动态学习的参数，决定应该保留多少过去的信息，哪些信息需要“遗忘”掉。通过这种方式，Mamba 能够避免传统方法中计算复杂的 自注意力机制，以 更高效 的方式处理长期依赖。
输出：最后，模型通过一个输出 $y_t$ 来决定最终的预测结果。

举个例子：

假设我们要用 Mamba 来分析一个图像序列（比如视频的每一帧）：

每一帧图像都对应一个输入 $x_t$ 。
模型在每一帧图像后，更新它的 隐藏状态 $h_t$ ，并通过 选择性扫描机制 决定哪些帧的特征信息需要被保留。
最后，通过 状态空间更新 来生成图像序列的最终输出，如分类或分割。

核心优点：

高效处理长序列：由于 选择性更新 和 线性递推，Mamba 可以高效地处理长序列，避免了像 Transformer 那样的 O(N²) 时间复杂度。
更少的参数：相比于传统的自注意力方法，Mamba 只需要较少的参数来表达序列中的长程依赖，计算成本更低。
选择性记忆机制：通过 动态决定记忆的更新，使得 Mamba 对序列信息的处理更加灵活和智能。

3. 总结：

Mamba 的核心思想就是 动态、选择性地更新状态，它通过 线性递推 和 选择性机制 来高效处理序列数据，从而避免了传统方法（如自注意力）中计算开销大的问题。

在训练过程中，模型学会了 如何根据输入动态决定哪些信息需要保留，哪些需要遗忘，使得模型能够更有效地捕捉长时间序列的依赖关系。

易错点

1. 混淆了 Mamba 和 Transformer/ViT

问题：很多人刚接触 Mamba 时，容易把它和 Transformer 或 ViT 混淆，因为它们看起来结构相似，都是处理序列数据。
错误点：
- Mamba 使用的是 线性递推机制，它通过 状态空间模型（SSM）来处理信息的传递和更新。而 Transformer 和 ViT 使用的是 自注意力机制，通过计算所有元素之间的关系来更新信息。
- Transformer 的 注意力机制 是通过计算 Q/K/V 之间的相似度来捕捉全局依赖，而 Mamba 通过 选择性状态更新 和 衰减机制 来更新信息。
正确理解：
- Mamba 通过状态空间更新（使用离散稳定更新）来高效捕捉序列的长程依赖，它的计算复杂度是 O(N)，比 Transformer 的 O(N²) 要高效得多。

2. 错误理解选择性状态扫描的机制

问题：很多人对 选择性状态扫描（Selective State Space Scanning） 的理解不够清晰，尤其是如何动态选择信息更新。
错误点：
- 选择性状态扫描中的 衰减系数 $Δ\Delta$ （即步长）控制了每一步的 信息保留与遗忘。如果理解成只是对状态做“简单加权”，可能会忽略 状态更新的自适应性。
- 误以为 Mamba 就是“简单的记忆更新”，但它实际上依赖于 动态调节，不同的输入会导致不同的更新策略。
正确理解：
- Mamba 的 选择性状态扫描 通过计算输入的 步长（ $Δt\Delta_t$ ） 来调整每一步的记忆衰减速率。它是 动态的，并不是固定的门控机制，而是根据输入数据来自适应更新模型的状态。

3. 误解 Mamba 的“稳定更新”

问题：许多小白在理解 离散稳定更新（Discretized Stable Update）时，容易把它和传统的递推方法混淆。
错误点：
- 稳定更新 其实是通过 指数衰减函数 $e−Δte^{-\Delta_t}$ 来避免信息在更新过程中“爆炸”。理解为“简单的记忆衰减”时会忽略它在状态空间模型中的作用，导致对 数值稳定性 的错误理解。
正确理解：
- 离散稳定更新 是将连续时间模型（通常为微分方程）转换成离散模型时，确保不会因为数值误差导致状态过度增长或衰减。通过指数衰减系数 $e−Δte^{-\Delta_t}$ ，模型能够稳定地更新状态，而不会出现爆炸或消失的问题。

4. 错误理解 Mamba 与 RNN 的关系

问题：Mamba 和 RNN（递归神经网络） 有相似之处，很多人因此误认为 Mamba 是一种变种的 RNN。
错误点：
- RNN 通过递归地将前一个时间步的状态传递给当前时间步，并依赖 激活函数 来更新状态。而 Mamba 使用的是 线性递推，并且通过 选择性状态扫描 和 稳定更新 来进行信息传递，它不依赖于像 tanh 或 sigmoid 这样的非线性激活函数。
- RNN 是一个 非线性动态系统，它的记忆是通过非线性激活函数进行更新的，而 Mamba 是一个 线性动态系统，通过稳定的线性递推来进行状态更新。
正确理解：
- Mamba 是通过 状态空间模型 实现的，而 RNN 是通过递归函数实现的动态系统。Mamba 更像是 状态空间的高效实现，它避免了 RNN 的梯度消失/爆炸问题，同时具备 更高效的计算复杂度。

5. 误解选择性状态空间的“记忆”

问题：Mamba 强调“选择性记忆”和“状态更新”，但有些人理解成 Mamba 就是一个 “记忆网络”，可以无限记忆所有信息。
错误点：
- Mamba 并不是一个无限记忆的网络，它的“记忆”是 选择性的，每个输入（或者每一步）决定了哪些信息要保留，哪些要丢弃。并不是所有信息都被永远保存。
- 误以为 Mamba 可以解决“所有”的记忆问题，而实际上，它是通过 衰减系数来保证 只保留重要信息。
正确理解：
- Mamba 通过选择性地衰减和更新信息，确保模型在长时间序列中的稳定性和效率，它 不会保留所有的信息，而是依赖 输入特征的自适应性 来动态决定“什么值得记住，什么需要忘记”。

小结：

常见的 Mamba 易错点：

和 Transformer/ViT 混淆：Mamba 用状态空间模型，而 Transformer 用自注意力。
选择性状态扫描：理解为“简单加权”而忽略其动态性和自适应性。
稳定更新：忽略其作为离散化连续时间模型的数值稳定性作用。
与 RNN 混淆：Mamba 是线性递推，不是非线性激活的 RNN。
无限记忆误解：Mamba 只保留重要信息，选择性忘记无关信息。

“Mamba 是线性的” 会觉得奇怪，因为深度学习模型通常都强调非线性能力。

1. 什么叫“线性”

在数学里，线性系统有两个核心性质：

可加性：
$f(x_1 + x_2) = f(x_1) + f(x_2)$
齐次性（比例性）：
$\cdot x) = c \cdot f(x)$ ，其中 $c$ 是常数。

简单说：输入是加法、数乘的形式，输出也能保持相同的规则。

2. 状态空间模型 (SSM) 的线性递推

Mamba 基于 状态空间模型 (SSM)，它的状态更新公式通常是：

$h_{t+1} = A h_t + B x_t$

$y_t = C h_t + D x_t$

其中：

$h_t$ ：隐状态 (hidden state)，是记忆。
$x_t$ ：输入。
$y_t$ ：输出。
$A, B, C, D$ ：都是矩阵（线性变换）。

👉 可以看到，状态更新和输出计算都是 矩阵乘法 + 向量加法，这就是典型的 线性运算。

3. 为什么说 Mamba 是线性的

RNN/LSTM 更新隐状态时，会用到 tanh、sigmoid 这样的 非线性激活函数。
比如：

$h_t = \tanh(Wx_t + Uh_{t-1})$

所以它是 非线性系统。
Mamba 在核心递推里，没有使用 tanh 或 sigmoid，而是用 矩阵乘法和加法，本质是 线性动态系统。
它的非线性主要体现在 输入投影和选择性扫描 上，但状态更新过程本身是 线性的。

4. 通俗解释

想象：

RNN 像一个“会弯曲”的传送带（有非线性弯折）。
Mamba 像一条“直轨道”的传送带（用线性规则传递信息），不过它在轨道上安装了“阀门”（选择性扫描），来决定哪些货物（信息）快点过去，哪些慢点过去。

5. 总结

Mamba 被称为 线性模型，是因为它的 状态更新公式完全由 线性运算（矩阵乘法、加法） 组成，不依赖非线性激活函数。

优势：数值稳定、计算效率高，可以并行化。
不足：如果没有选择性机制，它单靠线性就不够强大，所以才引入了 选择性扫描 (Selective SSM) 来增强表达力。

推一推 Mamba 在图像输入 (512 × 512 × 3) 的情况下，A,B,C,D 矩阵可能有多大。

1. 图像输入是怎么进 SSM 的？

一张彩色图像：

$512 \times 512 \times 3$

$H = 512, W = 512$ ，一共 $262{,}144$ 个像素。
每个像素有 3 个通道 (RGB)。

在 Mamba / ViT 里，通常会先做 切 patch + 线性投影：

例如 patch 大小 = $16 \times 16$ ，那么图像会变成：

$\frac{512}{16} \times \frac{512}{16} = 32 \times 32 = 1024 \, \text{个 patch}$
每个 patch 有 $16 \times 16 \times 3 = 768$ 维特征。
通过线性投影，把 768 维压到一个 模型维度 $d$ ，比如 $d = 128$ 或 $d = 256$ 。