Mamba:选择性状态空间的线性时间序列建模

快速答案

Mamba 是一个用选择性状态空间模型(SSM)替代注意力的序列模型:它的循环参数是当前输入的函数,因此能按 token 决定保留还是丢弃信息。收益很具体——随序列长度线性扩展,推理吞吐达到 Transformer 的 5 倍,在真实数据上一直到百万长度序列性能仍在提升,而且 Mamba-3B 在预训练和下游评测中都超过同尺寸 Transformer,并匹敌两倍参数的 Transformer。

为什么以前的高效架构总是输给注意力

为了摆脱注意力的二次方成本,人们尝试过线性注意力、门控卷积、循环网络、经典结构化 SSM,它们在长输入上扩展更好,却在语言这类离散模态上输给注意力。作者把症结归到一点:这些模型做不了基于内容的推理。无论来的是哪个 token,它们的动态都是固定的,因而无法选择性地聚焦某个关键词、忽略某个填充词。注意力能做到,因为每个 token 都与其他所有 token 比较——而这种两两比较恰恰是二次方成本的来源,也是以前的高效模型为了换速度而丢掉的东西。

选择性状态空间

结构化 SSM 通过一个隐藏状态把序列映射出去,而早期工作里这些矩阵在时间上是恒定的。Mamba 的关键一招,是让这些参数(步长以及输入/状态投影)成为输入的函数。这样模型就能在每个位置根据当前 token 决定沿序列传播还是遗忘信息——把基于内容的推理放进了循环里。

这破坏了让以前 SSM 变快的技巧。时不变 SSM 可以写成全局卷积来计算;一旦参数逐 token 变化,卷积就用不了了。作者的解法是面向硬件的并行扫描:以扫描形式计算循环,但把展开后的状态留在快速的 SRAM 里,避免写入较慢的 GPU 显存。正是这一点,让逐 token 变化的循环既有表达力又跑得快。

架构本身也刻意做减法。Mamba 把选择性 SSM 收进一个同质的 block,完全去掉了注意力和标准 MLP block——整个网络就是这种 block 的堆叠。这里的简洁是优点:组件类型更少,只有一个被反复使用的基本单元。

关键结果

推理吞吐是 Transformer 的 5 倍,因为生成是常数显存的循环,而不是不断增长的注意力缓存。
随序列长度线性扩展,且在真实数据上质量一直提升到百万长度序列——这正是注意力变得不切实际的区间。
Mamba-3B 超过同尺寸 Transformer,并匹敌两倍参数的 Transformer,在预训练困惑度和下游评测上都成立。
作为跨语言、音频、基因组的通用 backbone 都有强表现,不是单一领域的技巧。

一句中肯的判断:这些标志性对比都在约 3B 规模,而非前沿规模;“匹敌两倍参数 Transformer”这种说法,随着模型变大会越来越难站住脚。

局限与存疑

Mamba 不是注意力的直接替代品。它的证据上限在 3B 量级附近,而领域里的核心疑问是:在大多数生产模型所处的数百亿参数规模,选择性 SSM 的优势是否还在。还有一个结构性取舍:循环状态是一个固定大小的摘要,因此需要精确回忆任意早前 token 的任务(复制、精确检索)可能更偏好显式注意力,因为它让每个 token 都可寻址。这也是为什么后续大量工作转向混合架构——交替堆叠 Mamba 层与注意力层——而非纯 SSM。工具链成熟度、训练配方、大规模下的稳定性,在论文发表时都未定,如今也只是部分解决。

常见问题

Mamba 和 Transformer 到底有什么不同?

Mamba 没有注意力。它不做 token 两两比较,而是携带一个循环状态,用输入相关的参数决定这个状态保留或遗忘什么。这带来线性时间扩展和常数大小的生成开销,而注意力的成本会随上下文长度增长。

为什么 Mamba 推理比 Transformer 快?

Transformer 生成时要在不断增长的 KV 缓存上重新做注意力,所以每个 token 的成本随上下文上升。Mamba 通过推进一个固定大小的循环状态来生成,因此吞吐约为 5 倍,显存也不会随长度膨胀。

Mamba 在语言建模上能赢过 Transformer 吗?

在测试规模上、同一量级内是的:Mamba-3B 超过同尺寸 Transformer,并在预训练和下游任务上匹敌两倍参数的版本。这个优势在前沿规模是否成立仍是开放问题,所以现在很多系统把 Mamba 层与注意力层结合使用。

什么是选择性状态空间模型?

它是一种结构化 SSM,其循环参数依赖当前输入,而非在时间上固定。这种输入相关性让模型能按 token 选择性地传播或遗忘信息——即早期 SSM 缺失的基于内容的推理——并通过面向硬件的并行扫描高效计算。

Mamba 真正持久的贡献,是把循环从”过时概念”重新定位为注意力的、面向硬件的认真对手——原文见 arXiv:2312.00752。