从RNN到Mamba：图解状态空间模型中的‘扫描’到底在扫什么？-开发者社区

从RNN到Mamba：图解状态空间模型中的‘扫描’到底在扫什么？

在序列建模的世界里，我们常常需要处理随时间变化的数据流。想象一下，你正在观看一场网球比赛——每一次击球都依赖于前一次击球的结果，就像我们处理语言或时间序列数据时，每个新词或数据点都建立在之前的信息基础上。传统RNN通过隐状态递归传递信息，而今天我们要探讨的状态空间模型（SSM）则采用了一种被称为"扫描"的机制来完成类似的任务。

1. 序列建模的基本挑战

序列数据的核心特征是时间依赖性。以股票价格预测为例，今天的股价往往与昨天的价格相关。这种依赖关系给计算带来了两个关键挑战：

顺序依赖性：后续计算依赖于先前结果
计算效率：长序列处理需要大量计算资源

传统RNN通过隐状态递归解决第一个问题，但难以应对第二个挑战。LSTM和GRU通过门控机制改善了长程依赖，但本质上仍是顺序计算。状态空间模型引入"扫描"操作，在保持序列建模能力的同时，为并行计算打开了大门。

关键概念：扫描操作本质上是一种序列变换，将输入序列转换为输出序列，同时维护并更新内部状态。

2. 从累加求和理解扫描的本质

让我们从一个简单的累加求和例子开始，这是理解扫描操作最直观的切入点。考虑以下Python代码：

import torch X = torch.tensor([1, 2, 3, 4]) Y = torch.zeros_like(X) Y[0] = X[0] for t in range(1, X.size(0)): Y[t] = Y[t-1] + X[t] # 递归更新

这段代码展示了扫描的核心特征：

状态维护：Y[t-1]保存了到t-1时刻的累积信息
增量更新：每个新时刻t，基于当前输入X[t]更新状态
顺序处理：必须按时间顺序依次计算

这个简单的累加器实际上就是一个最小化的状态空间模型！其中：

X：输入序列
Y：既是输出序列也是隐状态序列
更新规则：Y[t] = Y[t-1] + X[t] 定义了状态转移

2.1 扫描与RNN的对应关系

将上述累加器与RNN对比，可以发现惊人的相似性：

组件	累加求和	RNN	状态空间模型
隐状态	Y[t-1]	h[t-1]	x[t-1]
输入	X[t]	u[t]	u[t]
状态更新	Y[t]=Y[t-1]+X[t]	h[t]=f(h[t-1],u[t])	x[t]=A x[t-1]+B u[t]
输出	Y[t]	y[t]=g(h[t])	y[t]=C x[t]+D u[t]

这种对应关系揭示了扫描操作的本质：它是一类特殊的递归状态更新过程。

3. 并行扫描：当输入序列已知时的优化

顺序扫描虽然直观，但在现代硬件上效率低下。关键突破在于认识到：当整个输入序列已知时，我们可以打破严格的时间顺序。

3.1 并行累加求和的直觉

回到累加求和的例子，假设我们要计算[1,2,3,4]的累加和[1,3,6,10]。顺序计算需要3步：

0+1=1
1+2=3
3+3=6
6+4=10

但如果我们能同时知道所有输入，可以重组计算：

1 2 3 4 ↓ ↓ ↓ ↓ L1: 1 3 3 7 (相邻元素相加) ↓ ↓ L2: 1 10 (跨两元素相加) ↓ L3: 10 (总和)

这种分层计算虽然总操作数相同，但每一层的操作可以并行执行，大大减少实际运行时间。

3.2 Blelloch算法详解

Blelloch算法是并行前缀和计算的经典方法，包含两个阶段：

Up-sweep阶段：自底向上计算部分和
- 将数组视为完全二叉树
- 从叶子开始，逐层向上计算内部节点的和

def up_sweep(X): n = X.size(0) for d in range(int(math.log2(n))): stride = 2**(d+1) for k in range(0, n, stride): X[k+stride-1] += X[k+2**d-1] return X

Down-sweep阶段：自顶向下传播前缀和
- 将根节点置零
- 自上而下传播部分和，构建最终的前缀和

def down_sweep(X): n = X.size(0) X[-1] = 0 # 根节点置零 for d in reversed(range(int(math.log2(n)))): stride = 2**(d+1) for k in range(0, n, stride): t = X[k+2**d-1] X[k+2**d-1] = X[k+stride-1] X[k+stride-1] += t return X

这种算法的优势在于：

工作复杂度：O(n)（与顺序算法相同）
步数复杂度：O(log n)（相比顺序算法的O(n)）

4. Mamba中的选择性扫描机制

Mamba模型将并行扫描思想应用于状态空间模型，实现了高效的序列建模。其核心是选择性扫描（selective scan）操作，动态决定哪些信息需要保留或忽略。

4.1 状态空间模型的扫描方程

Mamba的状态更新方程可以表示为：

x_k = exp(Δ_k A) x_{k-1} + Δ_k B u_k y_k = C x_k + D u_k

其中：

A：状态转移矩阵
B：输入映射矩阵
C：输出映射矩阵
D：直接映射项
Δ：时间步长参数

对应的PyTorch实现核心：

def selective_scan(x, delta, A, B, C, D): deltaA = torch.exp(delta.unsqueeze(-1) * A) # 状态转移因子 deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # 输入映射因子 BX = deltaB * (x.unsqueeze(-1)) # 映射后的输入 hs = pscan(deltaA, BX) # 并行扫描得到隐状态 y = (hs @ C.unsqueeze(-1)).squeeze(3) # 计算输出 return y + D * x

4.2 并行扫描的实际考量

在实际实现中，Mamba面临几个关键挑战：

内存效率：原始Blelloch算法需要O(n)额外空间，但通过优化可以做到原地计算
数值稳定性：指数运算(exp(ΔA))需要特殊处理以避免数值溢出
硬件适配：充分利用GPU的并行计算能力

以下是一个简化的并行扫描实现框架：

def pscan(A, X): # 预处理：确保输入长度为2的幂次 orig_len = A.size(1) padded_len = 2**math.ceil(math.log2(orig_len)) # 填充输入 A_padded = F.pad(A, (0, 0, 0, padded_len - orig_len), value=1) X_padded = F.pad(X, (0, 0, 0, padded_len - orig_len), value=0) # Up-sweep阶段 for d in range(int(math.log2(padded_len))): stride = 2**(d+1) A_padded[:, stride-1::stride] *= A_padded[:, 2**d-1::stride] X_padded[:, stride-1::stride] += A_padded[:, 2**d-1::stride] * X_padded[:, 2**d-1::stride] # Down-sweep阶段 A_padded[:, -1] = 0 X_padded[:, -1] = 0 for d in reversed(range(int(math.log2(padded_len)))): stride = 2**(d+1) temp_A = A_padded[:, 2**d-1::stride] temp_X = X_padded[:, 2**d-1::stride] A_padded[:, 2**d-1::stride] = A_padded[:, stride-1::stride] X_padded[:, 2**d-1::stride] = X_padded[:, stride-1::stride] A_padded[:, stride-1::stride] *= temp_A X_padded[:, stride-1::stride] += temp_A * X_padded[:, stride-1::stride] + temp_X return X_padded[:, :orig_len]