跳到主要内容

绕开平方瓶颈 — SSM、RWKV、Mamba 与 TTT

这一章讲四件事: 为什么长文本逼人离开 Transformer;状态空间(记下系统此刻处境的一组数)模型——简称 SSM,用一组方程描述「状态怎么随时间变」——怎么用一套方程同时拿到「能并行训练」和「省算推理」;RWKV 与 Mamba 各自在这个骨架上加了什么;以及更激进的测试时训练(TTT)。 链条位置:第 02 章末尾留了一句「注意力随序列长度平方增长」——这一章就是那句话的账单和分期付款方案。书把它放在 2024 版教材里,是想说:Transformer 的统治并非终点。

1. 顶层全景​

  Transformer:注意力每步要跟所有历史词打分 → 计算量 ∝ 长度²,长文本撞墙
  传统 RNN(GRU/LSTM):只看隐状态,理论可无限长 → 但梯度消失/爆炸(第 01 章),长依赖学不动
      ↓ 两头都要
  SSM 家族:训练走卷积形式(并行) + 推理走递归形式(线性)     ← RWKV / Mamba
  TTT:       隐藏状态不够用 → 拿模型参数本身当记忆,边推理边更新
  图说:两条路线共享同一个目标——把「记住全部历史」的代价从平方压到线性。

2. SSM:一套方程,两种形态​

这一节是全章的地基,后面两个模型都是它的变体。

2.1 从控制理论借来的骨架​

SSM 的思想源自控制理论中的动力系统:用一组状态变量(概括系统此刻处境的数)描述系统随时间的连续变化——这种连续表示天然适合描述长时间范围的依赖1。骨架是三个随时间变化的量(状态 x、输入 u、输出 y)加四个可学习矩阵(状态矩阵 A、控制矩阵 B、输出矩阵 C、命令矩阵 D),写成两条方程:状态方程描述状态如何随输入和上一状态变化,输出方程描述状态如何变成输出2。

2.2 离散化:关键的一步​

连续形式直接算非常慢,所以要做离散化——书称之为「SSM 中最为关键的步骤」。用梯形法把积分近似掉之后,方程变成两种形态3。

  • 递归(每步用到上一步的结果)形式——展开:当前状态=前一步状态×矩阵+当前输入×矩阵。

    「迭代」即反复套用同一式子。

    像 RNN,推理时每个词只算一步,复杂度随长度线性(直着涨,不是平方涨);但和 RNN 一样没法并行训练4。

    像 RNN,推理时每个词只算一步,复杂度随长度线性(直着涨,不是平方涨);但和 RNN 一样没法并行训练4。

  • 卷积(拿一个小模式贴着序列滑一遍、每个位置做一次加权求和的数学运算)形式——把递归方程展开后,输出恰好等于输入与一个固定小模式核的卷积。因为参数在整个序列上固定(时不变性),它能高效并行训练。

  • 代价是上下文长度固定,自回归推理的延迟(每一步都要干等上一步)大5。

书给的用法一句话记住:训练时用卷积形式,推理时用递归形式——同一套参数,两种形态按需切换6。

3. RWKV:第一个扩到百亿的非 Transformer​

RWKV(Receptance Weighted Key Value)把 SSM 范式做成了可用的语言模型。核心机制 WKV 在数学上可以看作两个 SSM 之比;它的工程性格是两头占:训练时可当作 Transformer 并行,推理时可当作 RNN,计算与内存开销恒定7。

结构上有两个模块:时间混合模块处理不同时间步之间的关系(注意力管的事),通道混合模块处理同一时间步内不同特征通道的交互(FFN 管的事)——书 footnote 解释,通道就是向量里的每一个特征维度8。时间混合里有四个基本元素:接收向量 R(整合历史)、键 K、值 V(与注意力同义)、权重 W——W 是随时间衰减的权重,让每个通道的记忆随距离打折,越久远的信息权重越低9。

成绩单:参数扩到 140 亿,是第一个扩展到数百亿参数的非 Transformer 架构;NLP——自然语言处理——任务上与同规模 Transformer 相当,长上下文基准仅次于 S410。

顺带备好一个词:算法(一套明确的计算步骤),下文会反复出现。

4. Mamba:让参数「看菜下饭」​

这一节是本章主走查的落点。 SSM 的时不变性是把双刃剑:换来并行训练,但处理文本这类信息密集数据的能力较弱——同样的参数处理每个词,不能根据内容调整11。Mamba 的两件武器:

① 选择机制:把离散化 SSM 中的参数 B、C、Δ 变成输入的函数(比如过一层线性投影),模型参数由此带上时间维度——每个时间步参数值不同,从时不变变成时变。这相当于让模型每个位置自己决定「这条信息记多少、忘多少」12。但参数随输入变了,卷积形式的高效并行就没了。

② 硬件感知算法(一套明确的计算步骤):把丢掉的效率从 GPU(显卡)的内存层级里抢回来。

三招:内核融合(减少内存读写)、并行扫描、重计算(反向传播——训练时把误差往回传的步骤——时重算中间状态换内存)。

具体实现是把 SSM 参数从慢的大容量显存 HBM——High Bandwidth Memory,显存的一种——搬进更快的就近小存储 SRAM 里算完再写回。

那种就近的小存储,行话叫缓存(离计算单元更近的小快存)。GPU(显卡)的这两种内存速度差好几倍13。

书给的三个结果(都是可核对的主张):训练时计算与内存随长度线性增长、推理每步只需常数时间;在 A100 GPU 上推理吞吐(单位时间能处理的量)比同规模 Transformer 高 5 倍;Mamba-3B 在语言建模上超过两倍参数量的 Transformer;能处理百万级序列14。

判断(我们的,不是书里的): Mamba 的选择机制其实是把「注意力的按内容分配」偷渡回了 RNN 世界——时不变换来并行,时变换来内容感知,两样不可全得,只能用工程补齐。这个取舍结构在第 02 章的解码策略(稳与浪)、第 14 章的「何时检索」里会反复出现。 如果错,会错在: 如果后续硬件或算法让「时变」也能高效并行,这个取舍就不再是本质约束,而只是 2023-2024 年的工程现状。

5. TTT:拿参数当记忆,边推理边训练​

SSM 家族还有一个共同天花板:所有历史被压进固定长度的隐藏状态,压得太狠就会饱和——书给出实锤:Mamba 在上下文长度超过 16k 后,困惑度基本不再下降15。

TTT(测试时训练,Test-Time Training)的思路是换一个存储介质:用模型自身的参数来当隐藏状态、记住上文;推理时每读一个词就对参数做一次梯度更新——一边循环训练一边推理,这与「推理阶段模型保持静止」的传统范式完全相反16。

预训练时它有两层循环:外部循环照旧做下一词预测、优化全局权重;内部循环用自监督(不用人工答案,自己给自己出题)方式更新「作为记忆的隐藏状态」。

具体做法:给定当前输入与历史,算一个重构损失,再用梯度下降(沿最陡的下坡方向小步调参)更新记忆参数;推理时只跑内部循环17。结果:相对 RWKV/Mamba 能更有效地捕捉超长上下文,复杂度仍是线性18。

6. 作者的判断与证据​

  • 给了证据的:Mamba 16k 饱和(引文献);A100 吞吐 5 倍、Mamba-3B 超两倍参数 Transformer、百万级序列(引 Mamba 论文);RWKV 14B 与长上下文仅次于 S4。

  • 书自己的判断:「Transformer 并非完美」的开场与「TTT 有望在超长序列任务中发挥重要作用」的收尾——都是方向性判断,书没有给 TTT 的大规模实验数据,这部分读作展望而非结论。

  • 结构性的坦白:书承认 Mamba「在硬件依赖性和模型复杂度上存在一定的局限性」——它吃特定 GPU 特性,换环境未必划算19。

7. 边界与局限​

  • 本章三个模型都是 2023-2024 年的早期形态:RWKV 书里注明讨论的是 v4;Mamba 后续的 2 代、混合架构(注意力与 SSM 分层混用)不在书内。

  • 书没有给三者与 Transformer 在中文任务上的对比数据;所有基准均为英文为主。

  • TTT 的「推理时训练」意味着每个请求都要做梯度更新——书的效率分析停在线性复杂度,没算这条额外开销的账。

  • 这些架构的生态(推理框架、微调工具、量化——把数压短省内存——的支持)远不如 Transformer 成熟,选型时不只看论文数字。

8. 可带走的​

  1. 注意力的账单是平方;RNN 免单但学不动长依赖——现代架构要两头占。
  2. SSM 一套状态方程,离散化后两种形态:训练用卷积(并行)、推理用递归(线性)。
  3. 时不变性=可并行,但处理信息密集数据弱;这是本章所有取舍的总开关。
  4. RWKV:WKV 可看作两个 SSM 之比,权重随时间衰减;第一个到 14B 的非 Transformer。
  5. Mamba:选择机制让参数随内容变(时变),硬件感知算法(HBM→SRAM)抢回效率;吞吐高 5 倍。
  6. 固定长度隐藏状态会饱和:Mamba 超 16k 困惑度不再降。
  7. TTT:模型参数即记忆,推理时边测边训;超长上下文的新路线,但每步都有训练开销。

9. 原文地图​

主题原书章原文位置
平方瓶颈与两类现代变体2.6 非 Transformer 架构text/07-ch02-06-2-6-transformer.txt:29(搜「平方增长」) · text/07-ch02-06-2-6-transformer.txt:39(搜「状态空间模型」) · text/07-ch02-06-2-6-transformer.txt:42(搜「线性时间复杂度」)
SSM 控制理论来源与四矩阵2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:58(搜「控制理论」) · text/07-ch02-06-2-6-transformer.txt:73(搜「状态矩阵」)
离散化关键步骤2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:95(搜「离散化」)
递归与卷积形式2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:117(搜「递归形式」) · text/07-ch02-06-2-6-transformer.txt:141(搜「时不变性」)
训练卷积、推理递归2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:147(搜「训练时使用卷积形式」)
S4 与 HiPPO2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:161(搜「HiPPO」)
RWKV 结构与 WKV2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:168(搜「RWKV」) · text/07-ch02-06-2-6-transformer.txt:176(搜「时间混合模块」) · text/07-ch02-06-2-6-transformer.txt:180(搜「特征通道」)
14B 与恒定开销2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:263(搜「14B」) · text/07-ch02-06-2-6-transformer.txt:267(搜「恒定的计算和内存复杂度」)
Mamba 选择机制2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:284(搜「信息密集」) · text/07-ch02-06-2-6-transformer.txt:286(搜「选择机制」) · text/07-ch02-06-2-6-transformer.txt:306(搜「从时间不变」)
硬件感知算法2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:338(搜「并行扫描」) · text/07-ch02-06-2-6-transformer.txt:344(搜「HBM」)
5 倍吞吐与百万序列2.6.1 状态空间模型 SSMtext/07-ch02-06-2-6-transformer.txt:367(搜「5 倍」) · text/07-ch02-06-2-6-transformer.txt:371(搜「两倍参数」) · text/07-ch02-06-2-6-transformer.txt:374(搜「百万级别」)
16k 饱和2.6.2 训练时更新 TTTtext/07-ch02-06-2-6-transformer.txt:396(搜「16k」)
TTT 机制与双循环2.6.2 训练时更新 TTTtext/07-ch02-06-2-6-transformer.txt:400(搜「TTT」) · text/07-ch02-06-2-6-transformer.txt:402(搜「存储隐藏状态」) · text/07-ch02-06-2-6-transformer.txt:415(搜「内部循环」)

Footnotes​

  1. 出处:「2.6.1 状态空间模型 SSM」第 58 段(text/07-ch02-06-2-6-transformer.txt:58,搜「控制理论」)。 ↩

  2. 出处:「2.6.1 状态空间模型 SSM」第 73 段(text/07-ch02-06-2-6-transformer.txt:73,搜「状态矩阵」)。四个矩阵与两条系统方程(式 2.5)在该段;Du(t) 可视为残差连接、可忽略在同段末(text/07-ch02-06-2-6-transformer.txt:89,搜「残差连接」)。 ↩

  3. 出处:「2.6.1 状态空间模型 SSM」第 95 段(text/07-ch02-06-2-6-transformer.txt:95,搜「离散化」)。梯形法原理同段(text/07-ch02-06-2-6-transformer.txt:99,搜「梯形」)。 ↩

  4. 出处:「2.6.1 状态空间模型 SSM」第 117 段(text/07-ch02-06-2-6-transformer.txt:117,搜「递归形式」)。线性推理与不可并行训练同段。 ↩ ↩2

  5. 出处:「2.6.1 状态空间模型 SSM」第 141 段(text/07-ch02-06-2-6-transformer.txt:141,搜「时不变性」)。卷积核由矩阵参数决定、上下文长度固定导致自回归延迟大在同段(text/07-ch02-06-2-6-transformer.txt:143,搜「上下文长度固定」)。 ↩

  6. 出处:「2.6.1 状态空间模型 SSM」第 147 段(text/07-ch02-06-2-6-transformer.txt:147,搜「训练时使用卷积形式」)。 ↩

  7. 出处:「2.6.1 状态空间模型 SSM」第 168 段(text/07-ch02-06-2-6-transformer.txt:168,搜「RWKV」)。WKV 是两个 SSM 之比同段;训练可并行、推理恒定在第 267 段(text/07-ch02-06-2-6-transformer.txt:267,搜「恒定的计算和内存复杂度」)。 ↩

  8. 出处:「2.6.1 状态空间模型 SSM」第 176 段(text/07-ch02-06-2-6-transformer.txt:176,搜「时间混合模块」)与脚注正文(text/07-ch02-06-2-6-transformer.txt:180,搜「特征通道」)。 ↩

  9. 出处:「2.6.1 状态空间模型 SSM」第 182 段(text/07-ch02-06-2-6-transformer.txt:182,搜「接收向量」)与第 226 段(text/07-ch02-06-2-6-transformer.txt:226,搜「时间衰减」)。 ↩

  10. 出处:「2.6.1 状态空间模型 SSM」第 263 段(text/07-ch02-06-2-6-transformer.txt:263,搜「14B」)与第 270 段(text/07-ch02-06-2-6-transformer.txt:270,搜「仅次于 S4」)。 ↩

  11. 出处:「2.6.1 状态空间模型 SSM」第 282 段(text/07-ch02-06-2-6-transformer.txt:282,搜「时不变性使得」)与第 284 段(text/07-ch02-06-2-6-transformer.txt:284,搜「信息密集」)。 ↩

  12. 出处:「2.6.1 状态空间模型 SSM」第 286 段(text/07-ch02-06-2-6-transformer.txt:286,搜「选择机制」)与第 306 段(text/07-ch02-06-2-6-transformer.txt:306,搜「从时间不变」)。参数函数化(sB/sC/Δ 的具体形式)在第 298 段(text/07-ch02-06-2-6-transformer.txt:298,搜「Linear」)。 ↩

  13. 出处:「2.6.1 状态空间模型 SSM」第 338 段(text/07-ch02-06-2-6-transformer.txt:338,搜「并行扫描」)与第 344 段(text/07-ch02-06-2-6-transformer.txt:344,搜「HBM」)。 ↩

  14. 出处:「2.6.1 状态空间模型 SSM」第 362 段(text/07-ch02-06-2-6-transformer.txt:362,搜「快速训练和推理」)、第 367 段(text/07-ch02-06-2-6-transformer.txt:367,搜「5 倍」)、第 371 段(text/07-ch02-06-2-6-transformer.txt:371,搜「两倍参数」)与第 374 段(text/07-ch02-06-2-6-transformer.txt:374,搜「百万级别」)。 ↩

  15. 出处:「2.6.2 训练时更新 TTT」第 396 段(text/07-ch02-06-2-6-transformer.txt:396,搜「16k」)。 ↩

  16. 出处:「2.6.2 训练时更新 TTT」第 400 段(text/07-ch02-06-2-6-transformer.txt:400,搜「TTT」)与第 402 段(text/07-ch02-06-2-6-transformer.txt:402,搜「存储隐藏状态」)。 ↩

  17. 出处:「2.6.2 训练时更新 TTT」第 415 段(text/07-ch02-06-2-6-transformer.txt:415,搜「内部循环」)。重构损失与梯度更新(式 2.11–2.13)在同段(text/07-ch02-06-2-6-transformer.txt:427,搜「重构损失」)。 ↩

  18. 出处:「2.6.2 训练时更新 TTT」第 445 段(text/07-ch02-06-2-6-transformer.txt:445,搜「线性时间复杂度」)。 ↩

  19. 出处:「2.6.1 状态空间模型 SSM」第 378 段(text/07-ch02-06-2-6-transformer.txt:378,搜「硬件依赖性」)。 ↩