Mamba 深度解读：用状态空间模型挑战 Transformer 霸主地位
用状态空间模型挑战 Transformer
如今，AI 正在吞噬世界。
而这里的 AI，我特指 Transformer。过去几年 AI 领域的几乎所有重大突破都归功于 Transformer。
然而，Mamba 是另一类被称为状态空间模型（State Space Models, SSM）的替代模型之一。重要的是，Mamba 首次承诺在实现长序列（比如 100 万个 token）时，能够达到与 Transformer 相似的性能（以及至关重要的相似缩放定律）。为了实现这一长上下文，Mamba 的作者移除了注意力机制中的"二次瓶颈"。Mamba 还运行得很快——"比 Transformer 快 5 倍"。
Gu 和 Dao（Mamba 的作者）写道：
Mamba 享受快速推理和序列长度的线性缩放，其性能在真实数据上可延伸到百万级序列。作为通用序列模型骨干，Mamba 在语言、音频和基因组学等多个模态上达到了最先进的性能。在语言建模方面，我们的 Mamba-3B 模型优于相同规模的 Transformer，并匹配其两倍大小的 Transformer，无论是在预训练还是下游评估中。
本文将讨论：
- Mamba 与 Transformer 相比的优势（和劣势）
- 理解 Mamba 的类比和直觉
- Mamba 对可解释性、AI 安全性和应用的意义
Transformer 的问题——也许注意力并非你所需要的一切
我们正处于 Transformer 时代。机器学习曾经只是检测猫和狗。而现在，借助 Transformer，我们能生成如诗般优美的文字、写出超越中位数竞赛程序员的代码、解决蛋白质折叠问题。
但 Transformer 有一个核心问题。在 Transformer 中，每个 token 在预测时都可以回头看所有之前的 token。为了实现这种回溯，我们在所谓的 KV 缓存中缓存了每个 token 的详细信息。
这种两两之间的通信意味着在训练中，前向传播的时间复杂度是 O(n²)（可怕的二次瓶颈），而自回归生成每个新 token 需要 O(n) 时间。换句话说，随着上下文规模的增加，模型变得越来越慢。
雪上加霜的是，存储这个键值（KV）缓存需要 O(n) 的空间。因此，随着内存占用的增长，可怕的 CUDA 内存不足（OOM）错误变成了一个重大威胁。如果只是空间问题，我们也许可以考虑增加更多 GPU；然而，随着延迟呈二次方增长，简单地增加算力可能不是可行的解决方案。
在边际上，我们可以通过滑动窗口注意力或巧妙的 CUDA 优化（如FlashAttention）来缓解二次瓶颈。但归根结底，对于超长上下文窗口（比如一个记得每次对话的聊天机器人），我们需要不同的方法。
基础模型骨干
从根本上说，所有优秀的 ML 架构骨干都有两个重要操作的组件：
- token 之间的通信
- token 内部的计算
在 Transformer 中，这是注意力（通信）和 MLP（计算）。我们通过优化这两种操作来改进 Transformer。
我们希望用另一种促进 token 间通信的机制来替代注意力组件。具体来说，Mamba 采用了一种受控制理论启发的状态空间模型（SSM）用于通信，同时保留 MLP 风格的计算投影。
就像 Transformer 由堆叠的 Transformer 块组成一样，Mamba 也由堆叠的 Mamba 块组成。
我们希望理解并说明选择 SSM 进行序列转换的原因。
Mamba 的动机——回到 Temple Run
假设我们在构建一个 Temple Run 智能体。它需要在每个时刻决定角色是向左还是向右移动。
要成功选择正确的方向，我们需要了解周围环境的信息。我们把相关信息的集合称为状态（state）。这里的状态可能包括你的当前位置和速度、最近障碍物的位置、天气状况等。
请注意，你不需要一直看整个屏幕。通过观察障碍物随着奔跑而向下移动，你可以推断出屏幕大部分区域的变化。你只需要看屏幕顶部来了解新信息，然后模拟其余部分。
这引出了一个自然的数学表述。设 h 为隐藏状态（关于世界的相关知识），x 为输入（每次获得的观测值）。h' 表示隐藏状态的导数，即状态如何演化。我们试图预测 y——最优的下一个动作（左或右）。
更具体地说，状态 h 可以表示为微分方程（公式 1a）：
h'(t) = Ah(t) + Bx(t)知道 h 就能确定你的下一个动作 y（公式 1b）：
y(t) = Ch(t) + Dx(t)系统演化由其当前状态和新获取的观测决定。少量的新观测就足够了，因为大部分状态可以通过将已知的状态动力学应用于先前状态来推断。也就是说，屏幕的大部分内容并不是新的，它只是前一状态自然向下轨迹的延续。
通过观察屏幕顶部，你可以学到很多关于系统动力学的知识。例如，上部的速度增加可能意味着屏幕其余部分也在加速，所以我们可以推断游戏正在加速。这样一来，即使我们一开始对游戏一无所知，只有有限的观测，也能很快获得对屏幕动力学的整体理解。
什么是状态？
这里的状态指的是那些与输入变量结合后能完全决定未来系统行为的变量。理论上，一旦我们拥有了状态，就无需了解过去的任何其他信息来预测未来。有了这种状态选择，系统就转变为一个马尔可夫决策过程。理想情况下，状态是一小部分信息，它捕捉了系统的基本属性。也就是说，状态是过去的压缩。
离散化——如何在量化世界中生存
太好了！那么，给定一些状态和输入观测，我们有了一个自回归风格的系统来确定下一个动作。棒极了！
但实际上，这里有一个小问题。我们在将时间建模为连续的。但在现实生活中，我们在离散的时间步长中获取新输入并采取新动作。
我们需要将这个连续时间微分方程转换为离散时间差分方程。这个转换过程称为离散化。离散化在文献中是一个被充分研究的问题。Mamba 使用零阶保持（ZOH）离散化。
从公式 1a 出发：
h'(t) = Ah(t) + Bx(t)对于小的 Δ：
h'(t) ≈ (h(t+Δ) - h(t)) / Δ令 h_t = h(t) 和 h_{t+1} = h(t + Δ)，代入公式 1a：
h_{t+1} - h_t ≈ Δ(Ah_t + Bx_t)
→ h_{t+1} ≈ (I + ΔA)h_t + (ΔB)x_t
如果你曾经看过 RNN，并且觉得这很眼熟——相信你的直觉：我们有某个输入 x，它通过某种变换与先前的隐藏状态结合，生成新的隐藏状态。然后我们在每个时间步使用隐藏状态计算输出。
理解 SSM 矩阵
现在我们可以更直观地解释 A、B、C、D 矩阵：
- A：状态转移矩阵——决定状态如何演化（遗忘）
- B：输入投影——决定新输入如何进入状态（记忆）
- C：输出投影——决定如何从状态中提取输出
- D：跳跃连接——直接从输入到输出的路径
此外，Δ 有一个很好的解释——它是步长，或者我们可以称之为停留时间。对于大的 Δ，你更关注该 token；对于小的 Δ，你立即跳过该 token，不会将它太多地纳入下一个状态。
就是这样！这就是 SSM——Mamba 块中注意力（通信）的替代品。Mamba 架构中的计算来自常规的线性投影、非线性变换和局部卷积。
好的，理论讲完了——但这真的有效吗？
效果与效率：注意力是聚焦，选择性是优先级排序
在 1997 年的 WWDC 上，乔布斯有一句名言："专注就是说不"。专注是无情的优先级排序。我们通常从正面看待注意力——选择注意到什么。而在乔布斯的意义上，我们反而可以从负面来理解注意力——选择丢弃什么。
机器学习中有一个经典的直觉案例，称为鸡尾酒会问题。想象一个派对，有几十个同时进行的嘈杂对话：我们如何识别一个人在说什么，而其他人同时在说话？大脑通过将"注意力"集中在某个特定刺激上来解决这个问题，从而尽可能压制所有其他声音。
Transformer 使用点积注意力来聚焦于最相关的 token。注意力之所以如此出色，一个很大的原因是你有机会回顾上下文中发生的所有事情。如果运用得当，这就像照相记忆。
Transformer 非常有效。但它们并不非常高效。它们存储过去的一切，以便理论上可以完美地回顾 token。
传统的 RNN 则相反——它们遗忘很多，只在隐藏状态中回忆少量信息，并丢弃其余部分。它们非常高效——状态很小。但它们效率较低，因为丢弃的信息无法恢复。
我们想要更接近效果/效率权衡的帕累托前沿——比传统 RNN 更有效，比 Transformer 更高效。
Mamba 架构似乎提供了一种推动效果/效率帕累托前沿的解决方案。
SSM 和 RNN 一样高效，但我们可能会怀疑它们的效果如何。毕竟，它们似乎很难只丢弃不必要的信息而保留所有相关内容。如果每个 token 都按照相同的方式处理，对 token 应用相同的 A 和 B 矩阵（就像在 token 工厂流水线上一样），那就没有上下文依赖性。我们希望遗忘和记忆矩阵（分别是 A 和 B）能够变化并动态适应输入。
选择机制
选择性允许每个 token 以独特于其自身需求的方式转换到状态中。选择性正是将我们从普通的 SSM 模型（对每个输入应用相同的 A（遗忘）和 B（记忆）矩阵）带到了 Mamba——选择性状态空间模型。
在常规 SSM 中，A、B、C 和 D 是学习到的矩阵。而在 Mamba 的选择机制中，A、B、C 和 D 也是 x 的函数——矩阵是上下文相关的，而不是静态的。
使 A 和 B 成为 x 的函数使我们能够两全其美。
Mamba 论文作者写道：
序列模型的效果与效率权衡取决于它们压缩状态的程度：高效模型必须拥有小状态，而有效模型必须拥有包含上下文所有必要信息的状态。因此，我们提出构建序列模型的一个基本原则是选择性：或者说，将输入过滤或聚焦到序列状态中的上下文感知能力。特别是，选择机制控制信息如何沿着序列维度传播或交互。
人类（大部分情况下）对一生——甚至一天中经历的所有事情都没有照相记忆！信息太多了，不可能全部保留。我们在潜意识中通过选择遗忘来决定记住什么，在遇到信息时就抛弃大部分。Transformer 在回忆时决定关注什么。人类在形成记忆时也决定扔掉什么。人类早早地频繁过滤信息。
如果我们有无限的记忆容量，显然 Transformer 方法比人类方法更好——它确实更有效。但它效率较低——Transformer 必须存储大量可能无关的过去信息。Transformer 只在回忆时才决定什么相关。Mamba 的创新在于允许模型更早地找到更好的遗忘方式——它通过选择性来选择丢弃什么，在形成记忆时扔掉不太相关的信息。
选择性的问题
应用选择机制也有些陷阱。非选择性 SSM（即 A、B 不依赖于 x）在训练中可以快速计算。这是因为 Yt 中依赖于 xi 的部分可以表示为线性映射，即一个可以预计算的矩阵。
例如（忽略 D 组件，即跳跃连接）：
y₂ = CBx₂ + CABx₁ + CAABx₀
如果我们仔细观察，可能会发现更好的东西——这个表达式可以写成一个卷积。因此我们可以应用快速傅里叶变换和卷积定理来非常高效地在硬件上计算。
不幸的是，有了选择机制，我们就失去了卷积形式。很多注意力都集中在使用与 Tri Dao 的 FlashAttention 类似的硬件优化技巧来让 Mamba 在现代 GPU 硬件上保持高效。借助硬件优化，Mamba 的运行速度可以超过相当规模的 Transformer。
为政治经济学家准备的机器学习——状态应该有多大？
Mamba 作者写道："序列模型的效果与效率权衡取决于它们压缩状态的程度"。换句话说，就像在政治经济学中一样，根本问题是如何管理状态。
- 传统 RNN 是无政府主义的：它们有一个小的、最小的状态。状态大小有界限。状态压缩很差。
- Transformer 是共产主义的：它们有最大化的状态。"状态"只是整个历史的缓存，没有压缩。每个上下文 token 在回忆前都被平等对待。
- Mamba 拥有压缩状态：但它有选择地决定什么进入状态。Mamba 说，如果状态聚焦良好且高效，我们可以用小状态来达成目的。
结论是状态表示至关重要。较小的状态更高效；较大的状态更有效。关键在于选择性地、动态地将数据压缩到状态中。Mamba 的选择机制允许上下文相关的推理、聚焦和忽略。对于性能和可解释性来说，理解状态似乎都非常有用。
Transformer 与 Mamba 中的信息流
Transformer 是怎么"知道"东西的？在初始化时，transformer 并不智能。它通过以下方式学习：
训练数据：模型从训练数据中学习。这相当于将输入数据有损压缩到权重中。
上下文数据：Transformer 将上下文作为短期记忆，可以以近乎完美的保真度回忆。
检索：注意，Transformer 在回忆前完全不过滤上下文。因此，如果我们有一堆可能有用的信息，我们在 Transformer 外部过滤它（使用信息检索策略），然后将结果塞进提示中。这个过程称为检索增强生成（RAG）。
Mamba 的信息流：训练数据对 Mamba 的影响类似。然而，上下文数据和检索的界限稍微模糊。Mamba 的上下文数据类似于 Transformer 的检索数据，被压缩/过滤。这些上下文数据也可以像 Transformer 一样被查找（尽管保真度稍低）。
Transformer 的上下文之于 Mamba 的状态，就像短期记忆之于长期记忆。Mamba 不仅有"RAM"，它还有一块硬盘。
状态交换作为一种新的提示范式
目前，我们经常使用 RAG 为 Transformer 提供上下文信息。
对于类 Mamba 模型，你可以想象拥有一个状态库，通过让模型运行在专门数据上创建而成。状态可以像图像模型的 LoRA 一样共享。
例如，我可以在 20 本物理教科书和 100 道物理题及其解答上做推理。然后我得到一个状态，我可以把它给你。现在你不需要添加任何 few-shot 示例；你只需要简单地提问。上下文学习就在状态中。
换句话说，你可以将下载的状态拖放到你的模型中，就像插入卡带一样。注意，"训练"一个状态不需要任何反向传播。它更像是一种高度专业化的单次固定大小压缩算法。这是在推理时以零计算或零延迟应用无限上下文学习。
这比 few-shot 提示更便宜、更快速（因为状态可以无限重用，没有推理成本）。它也比微调便宜得多，并且不需要任何梯度更新。我们可以想象在上下文中额外检索状态。
Mamba 与机械可解释性
Transformer 的可解释性通常涉及分析注意力模式、神经元激活和消融研究。对于 Mamba，大多数我们想做的消融研究仍然有效，但理解 token 通信现在更加微妙。所有信息都通过隐藏状态在 token 之间移动，而不是通过注意力机制——后者可以将信息从序列中的一个位置"传送"到另一个位置。
为了理解 Mamba 的上下文学习（ICL）任务，我们将对 SSM 状态进行干预。一个经典的上下文学习任务是间接宾语识别。由于有人假设 Transformer 中大部分的上下文学习源于更底层的序列位置操作（如归纳头），Mamba 能够完成这个任务表明它具有更通用的上下文学习能力。
Mamba 和 SSM 的未来
类 Mamba 模型很可能在需要极长上下文和长期记忆的场景中表现出色。例如：
一个说明性的例子是拥有长期目标的智能体。假设你有一个与世界交互的智能体。最终，其体验超出了 Transformer 上下文窗口的承载能力。智能体不得不将其经验压缩或总结成更紧凑的表示。
但是如何确定哪些信息作为总结最有价值呢？对于语言任务，LLM 在总结方面相当不错。然而，对于其他领域，如何总结可能并不明确。例如，总结一部 2 小时电影的最佳方式是什么？模型本身能否自然学会这样做，而不是采用像用文字描述电影美学这样的权宜之计？
这就是 Mamba 所允许的：真正的长期记忆。一个真实的状态，模型在其中学会保留重要的东西。预测即压缩——学会什么对预测下一步有用，必然导致构建有用的先前 token 压缩。
对智能助手的影响是明确的：你的聊天机器人与你共同进化。它记得。
智能体与 AI 安全
对 AGI 存在风险持乐观态度的原因之一是语言模型。以前，通过自我对弈训练的深度强化学习智能体看起来最可能成为第一个 AGI。语言模型本质上更安全，因为它们不是以长期目标训练的。
这里长期序列推理的潜力重新凸显了基于智能体的 AI 安全的重要性。对于只有 8K 上下文窗口的 Transformer，很少有智能体相关的担忧。但对于拥有令人印象深刻的长期记忆和可能的工具性目标的系统，很多担忧都是相关的。
Transformer × Mamba：最佳合作
Mamba 作者表明，将 Mamba 的长上下文与 Transformer 在短序列上的高保真度结合是有价值的。例如，如果你在制作长视频，你可能无法将整部电影放入 Transformer 的注意力上下文中。你可以想象让注意力查看最近的帧以获得短期流畅性，同时用 SSM 保持长期叙事一致性。
这不是 Transformer 的终结。它们的高效性正是许多任务所需要的。但现在 Transformer 不再是唯一的选择。其他架构也是真正可行的。
所以我们不是在后 Transformer 时代。但这是第一次，我们生活在后-只有-Transformer的时代。而这为具有极端上下文长度和原生长期记忆的序列建模打开了无限可能。
两位 ML 研究员——Sasha Rush（HuggingFace、Annotated Transformer、康奈尔教授）和 Jonathan Frankle（彩票假说、MosaicML、哈佛教授）——目前正在就这个问题进行打赌。
目前 Transformer 遥遥领先。还有 3 年时间，现在有一个研究方向有了真正的机会。
最后剩下要问的是：注意力真的是我们所需要的全部吗？
zhirenhun
一个热爱技术的程序员，喜欢分享前沿AI知识和开发经验。
