SSM 家族全景图（上）：S5、H3 与 RWKV

前面的文章里，我们深入拆解了 SSM 家族的两座大山：S4 与 Mamba。S4 用 HiPPO、DPLR 和 FFT 把长序列建模做到了当时的最好水平，Mamba 用选择性机制和并行扫描把它推向了语言建模的舞台中央。

本篇跳出单个模型，俯瞰整个家族在 S4 奠基、Mamba 破圈之后长出的其他分支。我们选 5 个最具代表性的变体，每个都按“动机、核心创新、架构与数学”三个部分来拆解。S5、H3、RWKV 放在本篇，Mamba-2 与 Jamba 放在下一篇。

一、S5（Simplified S4）：用最简数学达到同等效果

1. 动机：S4 的工程实现太复杂

S4 虽然强大，但工程实现极其繁琐。要处理复数域、DPLR 结构、Woodbury 恒等式、Cauchy 矩阵和 FFT。研究者自然会问：能不能去掉这些复杂的数学技巧，用最简单的矩阵实现同等的性能？

2. SISO 的含义：$d$ 个独立的小 SSM

要理解 S5 的改进，先要弄清 S4 的 SISO（Single-Input Single-Output，单输入单输出）策略到底指什么。

假设输入为 $\mathbf{X} \in \mathbb{R}^{B \times L \times d}$，其中 $d = 768$ 是特征维度。S4 的处理方式是把 $d$ 维特征向量拆成 $d$ 个独立的标量通道，每个通道独立地跑一个一维的 SISO SSM。

以通道 $i$ 为例，它的小 SSM 结构是这样的：输入是标量 $x_t^{(i)} \in \mathbb{R}$；状态是 $h_t^{(i)} \in \mathbb{R}^N$（$N$ 维，例如 64）；参数是 $\mathbf{A}^{(i)} \in \mathbb{R}^{N \times N}$、$\mathbf{B}^{(i)} \in \mathbb{R}^{N \times 1}$、$\mathbf{C}^{(i)} \in \mathbb{R}^{1 \times N}$；输出是标量 $y_t^{(i)} \in \mathbb{R}$。

768 个这样的标量 SSM 并行运行，各自的输出最后拼接回 $(B, L, 768)$。“SISO”描述的就是每个小 SSM 的输入输出都是标量这一事实。

3. MIMO：S5 的核心改动

S5 作者指出，上述做法有一个根本问题：768 个通道在 SSM 内部完全独立，每个通道只处理自己的时间序列，看不到其他通道的信息。跨通道的信息交流只能靠 SSM 之外的那几层线性投影完成，深度极其有限。

S5 的改法是跑一个大的 MIMO（Multi-Input Multi-Output，多输入多输出）SSM。它的输入是向量 $x_t \in \mathbb{R}^d$，状态是 $h_t \in \mathbb{R}^N$（$N$ 可以设置得更大，例如 256），参数是 $\mathbf{A} \in \mathbb{R}^{N \times N}$、$\mathbf{B} \in \mathbb{R}^{N \times d}$、$\mathbf{C} \in \mathbb{R}^{d \times N}$，状态更新为 $h_t = \bar{\mathbf{A}} h_{t-1} + \bar{\mathbf{B}} x_t$，输出为 $y_t = \mathbf{C} h_t$。

关键区别在于 $\mathbf{B} \in \mathbb{R}^{N \times d}$ 与 $\mathbf{C} \in \mathbb{R}^{d \times N}$ 的形状：它们让 $d$ 维输入在写入状态时被混合，读取时又被重新组合，通道间的信息交互因此天然地发生了。

S5 并没有取消 SSM Block。Block 的整体架构（线性投影、非线性、门控、残差）依然保留，改变的只是 Block 内部“SSM 核心”的计算方式：从 $d$ 个独立的小标量 SSM，变成 1 个大的向量级 SSM。

4. 为什么 MIMO 能简化 $\mathbf{A}$

在 S4 中，因为每个通道是独立的，通道间的耦合必须通过复杂的 $\mathbf{A}$ 矩阵（DPLR 结构）来人为制造。

在 S5 中，通道间的耦合已经由 $\mathbf{B}$ 和 $\mathbf{C}$ 的矩阵乘法天然实现了。所以 $\mathbf{A}$ 只需要承担最简单的状态衰减职责，用纯对角矩阵就够了，甚至不需要 HiPPO 那套复杂初始化，只要把特征值均匀分布在复平面的左半平面即可。

因为 $\mathbf{A}$ 是纯对角的，离散化后的 $\bar{\mathbf{A}}$ 也是纯对角的，并行扫描变得极其简单和快速。

一句话概括：S5 用 MIMO 的架构设计换取了数学上的极大简化，性能媲美 S4，但代码实现容易得多。

5. 对角取值如何确定，实测表现如何

上一小节说 S5 的 $\mathbf{A}$ 只需把特征值均匀分布在复平面的左半平面。这里藏着一个值得想清楚的细节：为什么要均匀分布，而且必须放在左半平面？

对角矩阵承担的是 S4 中 DPLR 结构的职责，它要同时覆盖一批不同的时间尺度——有的模式衰减很快，只留住最近几步；有的模式衰减很慢，能把很久以前的信息带到现在。把特征值均匀铺开，就是让这些时间尺度被均衡地覆盖；把它们放在左半平面，则保证每个模式随时间衰减而不发散。S5 的初始化与参数化正是按这条思路设计的，目的是让简化后的对角系统在行为上贴近 S4 的 DPLR 系统。

计算实现上，S5 的收益很直接。因为状态转移是纯对角的，离散化后的递推相当于逐通道独立的标量乘法，这一步可以完全交给通用的并行扫描（parallel scan）——它是现成的库算子，不需要为对角线之外的结构额外编写定制内核。S5 因此能在保持精度的同时，获得与 S4 相当的计算效率。

实测上，S5 在 Long Range Arena（LRA）上的平均成绩为 87.4%，在最难的 Path-X 任务（序列长度 16K）上达到 98.5%。这两个数字说明了一件重要的事：把 DPLR 换成纯对角、把一堆数学技巧删掉之后，长序列建模能力并没有损失。

二、H3：让 SSM 第一次做好语言建模

1. 动机：S4 为什么做不好语言模型

2022 年，S4 在 LRA 等长序列任务上碾压 Transformer，但在语言建模（Language Modeling，即预测下一个词）上，却始终被 Transformer 压一头。

原因在于，语言模型不仅需要长程记忆，还需要精确记住最近的几个 token。比如“New York”后面大概率跟“City”，这种模式要求模型能精确复现局部的离散组合。S4 的连续时间状态空间倾向于平滑地压缩历史，容易丢失这类离散的、精确的局部模式。

这里可以再往深处追一步：差距究竟出在哪些能力上？研究者用一组合成的语言任务做了诊断，发现当时的 SSM 在两种能力上明显偏弱。第一种是召回序列中较早出现的 token，第二种是跨位置比较两个 token。这两种能力恰好对应语言建模里最常见的两类操作——按内容取回早先的信息，以及判断两处信息是否匹配。这也解释了为什么单纯把状态做大、或者单纯把衰减调慢，都补不上这个缺口：缺的是“精确取出”和“两两对照”这两种操作方式。

2. 核心创新：Shift SSM

H3 的核心思想是：用 SSM 处理长程上下文，再用一个特殊的机制显式地记忆局部 token。

这个特殊机制就是 Shift SSM（移位状态空间）。H3 构造了一个特殊的 $\mathbf{A}$ 矩阵，它承担移位寄存器（Shift Register）的角色，不带任何衰减：

$$\mathbf{A}_{shift} = \begin{pmatrix} 0 & 0 & 0 & \cdots & 0 \ 1 & 0 & 0 & \cdots & 0 \ 0 & 1 & 0 & \cdots & 0 \ \vdots & & \ddots & \ddots & \vdots \ 0 & 0 & \cdots & 1 & 0 \end{pmatrix}$$

配套的输入矩阵设为 $\mathbf{B} = (1, 0, 0, \ldots, 0)^\top \in \mathbb{R}^N$。

它的作用可以直接通过递推看出来。假设 $N = 4$，输入序列为 $x_1, x_2, x_3, x_4, x_5, \ldots$，递推式为 $h_t = \mathbf{A}{shift} h{t-1} + \mathbf{B} x_t$，逐步展开得到：

• $t=1$ 时，$h_1 = [x_1, 0, 0, 0]^\top$。
• $t=2$ 时，$h_2 = [x_2, x_1, 0, 0]^\top$。
• $t=3$ 时，$h_3 = [x_3, x_2, x_1, 0]^\top$。
• $t=4$ 时，$h_4 = [x_4, x_3, x_2, x_1]^\top$。
• $t=5$ 时，$h_5 = [x_5, x_4, x_3, x_2]^\top$。

状态向量 $h_t$ 就是一个精确的滑动窗口缓冲区，原封不动地保存了最近 $N$ 个时间步的输入值，没有任何衰减、混合或压缩。

3. 双路设计

H3 的架构是双路设计：一路是 Shift SSM，输出精确的局部记忆；另一路是标准 SSM（带衰减），负责长程依赖；两路的结果结合（例如逐元素相乘）后作为输出。

两路的分工很明确。Shift SSM 充当精确的短期记忆缓冲区，保证最近 $N$ 个 token 被无损保存，使模型可以精确检索“上上个词是什么”。标准 SSM 则负责压缩和记忆更长的历史。具体到融合方式，H3 把 Shift SSM 的输出与带衰减的对角 SSM 的输出逐元素相乘，用这个乘法门控把两路信息合到一起。这种“一路提供精确的局部值、一路提供逐位置的权重再相乘”的结构与线性注意力同构，H3 的三个组件正是从线性注意力得到的启发。

把三个组件与上一小节诊断出的缺口对应起来看，分工就很清楚了：局部记忆由 Shift SSM 补齐，带衰减的对角 SSM 继续负责长程汇总，而乘法门控提供的是逐位置的交互（形式上与线性注意力一致），承担的正是“跨位置比较”这件事。

4. 局部卷积为什么不够用

第 7 篇讲过，Mamba Block 里有一层一维因果卷积。自然会问：既然有了卷积，为什么还需要 Shift SSM？

一维因果卷积（kernel size 取 4）确实可以捕捉局部模式，但它有两个明显局限。卷积的记忆方式是通过学习的权重做加权求和，这属于有损压缩；它的窗口通常很小，一般在 3 到 5 之间；而且它的卷积核权重需要训练。

Shift SSM 在三方面都不同。它的记忆方式是原封不动地保存；没有信息损失，属于精确复制，因为 $\mathbf{A}_{shift}$ 是固定的、不需要训练；窗口可以很大，例如 64。

核心差异可以用一句话概括：卷积是“学习如何看局部”，Shift SSM 是“把局部完整地存下来”。对于语言建模中需要精确匹配的场景（如“New York”后面接“City”），Shift SSM 的无损记忆更可靠。

结果。 H3 首次在语言建模任务上，让 SSM 模型的性能追平了同等规模的 Transformer，证明了 SSM 完全具备做大语言模型的潜力。

实测结果给了更具体的证据：在 OpenWebText 上，H3 与 Transformer 的困惑度只差 0.4；而一个保留了两层注意力的 125M 参数混合模型，反而比同规模 Transformer 低了 1.0 PPL。把这种混合模型放大到 2.7B 参数并在 The Pile 上训练后，困惑度仍然低于同规模 Transformer，零样本与少样本在 SuperGLUE 的多数任务上也占优。工程侧配套的 FlashConv 把融合的分块 FFT 与状态传递算法结合起来，在 LRA 上取得 2 倍加速，并让混合模型生成文本的速度达到 Transformer 的 2.4 倍。

三、RWKV：社区驱动的线性 RNN

1. 动机与命名

RWKV 严格来说不属于经典控制论意义上的 SSM，但它的设计哲学与 SSM 高度一致，两者殊途同归。它的目标很明确：训练时像 Transformer 一样并行，推理时像 RNN 一样省显存。

RWKV 这个名字来自它的四个核心变量：Reception、Weight、Key、Value。

2. 宏观架构

RWKV 与 Transformer、Mamba 一样，也是 Block 的纵向堆叠。每个 RWKV Block 包含两个子层：Time-Mixing 负责序列维度上的信息混合，Channel-Mixing 负责特征维度上的非线性变换。两个子层外面套着残差连接与 LayerNorm。

3. Time-Mixing：RWKV 的核心

Time-Mixing 负责让序列中的不同时间步交换信息，它分三个步骤。

步骤一：Token Shift（Token 混合）。

$$\hat{x}t = \mu \cdot x_t + (1 - \mu) \cdot x{t-1}$$

其中 $\mu$ 是可学习的标量。这一步把当前 token 和前一个 token 做线性插值，提供极局部的上下文感知。

步骤二：计算 R、K、V、W。 四个变量都由同一个混合后的输入线性变换得到：

$$r_t = W_r \hat{x}_t \quad (\text{Reception})$$

$$k_t = W_k \hat{x}_t \quad (\text{Key})$$

$$v_t = W_v \hat{x}_t \quad (\text{Value})$$

$$w_t = W_w \hat{x}_t \quad (\text{Weight Decay})$$

步骤三：带衰减的加权求和（线性 RNN 视角）。 输出 $o_t$ 的计算方式为：

$$o_t = \frac{\sum_{i=1}^{t} e^{-(t-i)w + k_i} \cdot v_i}{\sum_{i=1}^{t} e^{-(t-i)w + k_i}}$$

（RWKV 的完整形式还包含一个 bonus 项 $u$：分子与分母各多出 $e^{u+k_t}$，让当前 token 单独参与加权。这里为书写简洁先略去。）

这个式子的几个部分各有含义。整体上，它本质上是一个带指数衰减的线性注意力。$k_i$ 决定了历史第 $i$ 个 token 的重要性；$w$ 决定遗忘速度，作用类似 Mamba 中的 $\Delta$；分母是归一化因子，作用类似 softmax 的分母。

推理时的递推形式。 为了在推理时达到每步 $O(1)$，RWKV 维护两个运行状态：加权值累加器 $n_t = e^{-w} n_{t-1} + e^{k_t} v_t$，以及权重累加器 $d_t = e^{-w} d_{t-1} + e^{k_t}$。最终输出为 $o_t = n_t / d_t$。

这个递推在结构上与 SSM 的 $h_t = \bar{A} h_{t-1} + \bar{B} x_t$ 完全同构。这也解释了为什么 RWKV 明明从纯深度学习出发，却和 SSM 走到了一起。

这里还有两个实现层面的细节值得知道。第一，$e^{k_t}$ 与 $e^{-w}$ 都是指数形式，直接算很容易超出半精度的表示范围，所以官方实现为 WKV 的累加与归一化单独写了算子，把指数运算和归一化放在一起处理，兼顾数值范围和显存读写。第二，token shift 在 RWKV 的每个子层都会做一次，Time-Mixing 与 Channel-Mixing 各自有自己的一组混合系数，所以相邻位置的信息在每个子层都有一次交流的机会。

4. Channel-Mixing

这个子层很简单，等价于一个轻量级的 FFN：

$$y_t = \sigma(W_2 \cdot \text{ReLU}(W_1 x_t))$$

其中 $\sigma$ 是 Sigmoid 函数。

5. RWKV 与 Mamba 的对比

理论基础： Mamba 来自控制论的状态空间方程；RWKV 来自纯深度学习设计。

时间混合机制： Mamba 用选择性 SSM（$\Delta$、$B$、$C$ 动态）；RWKV 用带衰减的线性注意力（$w$、$k$ 动态）。

通道混合机制： Mamba 用门控（Gating）；RWKV 用 ReLU 加 Sigmoid 的组合。

状态维度： Mamba 每个通道的状态维度 $N$ 通常取得较小（例如 16，Mamba-2 可以更大）；RWKV 的状态维度与模型维度 $D$ 相同。

6. 为什么 RWKV 值得关注

RWKV 是一个由开源社区主导的项目，没有大厂光环。它凭借极简的数学形式和极高的工程效率，训练出了 7B、14B 甚至更大规模的模型，在真实世界的对话和推理任务中表现优异。

它的意义在于证明了线性序列模型在 LLM 时代绝对有一席之地。

7. 版本演进

RWKV 之后又经历了几个版本的迭代，每次都在扩大状态的表达能力。

RWKV-5（Eagle）与 RWKV-6（Finch）把状态从向量扩展为矩阵值状态（matrix-valued states），并引入多头结构与动态递推，状态容量和表达力随之提高。

RWKV-7（Goose，2025）则把状态更新换成了广义 delta rule：用向量门控与上下文相关的学习率（in-context learning rate）来更新状态，让模型按当前上下文决定写入多少、覆盖多少。

它要解决的问题与 Mamba 的选择性机制是同一个——让状态更新的强度依赖内容。两者的区别在于控制方式：Mamba 用 $\Delta_t$ 一个标量旋钮同时控制衰减与写入，RWKV-7 则用一个向量加一个学习率分别控制写入的方向与幅度。把两个方案并排看，能更清楚地理解“内容感知”在一维递推型模型里到底是怎么落地的。

四、小结

本篇介绍的三个模型分别从不同角度推动了 SSM 家族。

• S5 用 MIMO 架构加纯对角 $\mathbf{A}$ 矩阵，把 S4 身上繁琐的数学技巧（复数域、DPLR、Woodbury、Cauchy 核）替换成最简单的矩阵实现，性能保持一致甚至更好。它的定位是“用最简单的数学，达到同样的效果”。
• H3 通过 Shift SSM 与标准 SSM 的双路设计，给模型补上了精确的局部记忆能力，首次让 SSM 在语言建模上追平同等规模的 Transformer。它的定位是“让 SSM 第一次做好语言建模”。
• RWKV 用带衰减的线性注意力与 Token Shift 组成了一个结构极简的线性 RNN，训练时可并行、推理时每步 $O(1)$，证明了社区驱动的线性序列模型在 LLM 时代有实际竞争力；后续的 RWKV-5/6 把状态扩展为矩阵值，RWKV-7 又引入广义 delta rule 与上下文相关的学习率。

下一篇会继续分析家族中理论意义最大的 Mamba-2，以及把 SSM、Attention 与 MoE 结合起来的工业级混合架构 Jamba。