Mamba选择性状态空间模型

Mamba 是一种以选择性状态空间层为核心的序列模型架构。它让状态更新随当前输入变化，使模型能在有限的递归状态中保留有用信息、淡化无关内容；再用硬件感知的选择性扫描并行计算这些递归更新。本文只介绍 Gu 与 Dao 的原始 Mamba 论文，讨论其机制、网络块、论文实验及其适用边界。

状态空间模型（SSM）用一个隐状态沿序列传递信息。离散形式可写作 h_t = Ā h_(t−1) + B̄ x_t、y_t = C h_t：当前输入 x_t 更新状态，状态再生成输出 y_t。传统结构化 SSM 的递推参数在各位置保持不变，因此同一套动态规则会处理每个 token。固定参数让递推可改写为卷积，便于整段并行计算；代价是模型不容易依据内容决定该记住什么、忽略什么。

这个限制在“选择性复制”任务中很直观：模型读到一串夹杂噪声的符号，只需记住其中指定的数据，等到序列末尾再按顺序写出。若要依赖固定间距或时间位置，间隔一旦随机变化，规则就不够用了。语言也有类似难点：一个词是否重要，通常取决于它的内容和前后上下文，而不是它在固定的第几步出现。

Mamba 的选择性状态空间层（论文中称 S6）把 Δ、B 和 C 改成由当前输入计算的量；结构化矩阵 A 仍是学习得到、沿位置共享的参数。每步的 Δ_t 参与离散化，调节本步的状态转移和输入写入幅度；B_t 控制当前内容怎样进入状态，C_t 控制从状态读出什么。这样，同一个模型可以根据正在读到的 token 改变信息传播方式，而不是对全序列使用固定滤波器。

用递推式概括，h_t = Ā_t h_(t−1) + B̄_t x_t，y_t = C_t h_t。若本步内容不值得保留，更新可少写入新信息或淡化已有状态；若内容相关，状态便能保留它供后续位置使用。这里的“选择”是训练出的连续状态更新，不是先列出关键词再硬性删除，也不是像自注意力那样在当前步显式比较并读取所有历史 token。有限状态需要压缩上下文，具体能保留多少取决于模型参数、训练数据和任务。

论文用一个简化的一维情形解释这种门控直觉：当状态维度取 1 且参数固定为特定值时，递推可化成 h_t = (1 − g_t)h_(t−1) + g_t x_t，其中 g_t 由当前输入计算。它说明输入可以控制新内容写入与旧状态保留的比例。真实 Mamba 块使用更高维的结构化状态和多个通道；这个一维式子只是解释参数作用的特例，不是实际网络的完整公式。

让参数随输入变化后，每个位置的更新规则都不同，传统的单个全局卷积核就不再适用。逐 token 执行递推虽然自然，却会形成顺序依赖。Mamba 将状态更新写成可组合的递推变换，再对这些变换做并行前缀扫描（scan）；扫描保留递推语义，同时能把整段训练序列分块并行处理。它不是对 token 两两计算注意力矩阵。

并行扫描的中间状态如果全部展开，会形成大小近似为“批次 × 序列长度 × 通道 × 状态维度”的张量。反复把这类张量写入和读回 GPU 高带宽显存，会抵消递推的计算优势。论文的硬件感知算法把参数离散化、扫描和输出投影融合到内核中，在更快的片上 SRAM 中完成较大的中间状态计算，只把较小的最终输出写回高带宽显存；训练反向传播时还会重算部分中间量，减少存储与读取。它针对的是数据搬运和中间状态展开成本，因此“线性时间”仍需结合状态维度、内核实现和硬件来理解。

训练和生成使用同一组选择性更新，但输入可见方式不同：训练时整段样本已知，扫描可以在序列维度上并行；自回归生成时每次只到达下一个 token，模型以递归方式更新并携带状态。并行扫描解决整段计算的顺序依赖，递归状态则让逐步生成无需把整段历史重新展开。它们是同一状态方程的两种执行场景，不是两套不同的模型定义。

在自回归推理时，模型可以保留每层递归状态，处理下一个 token 时更新这份状态，不必像全注意力 Transformer 那样为所有历史 token 保存完整的逐层键值缓存。因此，序列越长，Mamba 的历史状态占用不会按 token 数持续增长；这不表示模型完全不占用状态内存，也不代表提示词处理、投影、批次并发和所有 GPU 操作都没有成本。

原论文把选择性 SSM 放进一个重复堆叠的网络块中，并把以往常分开的序列变换和 MLP 计算合并到较统一的结构。输入投影扩展通道后，一条路径经短卷积和 SiLU 激活，另一条路径提供门控信号；选择性 SSM 处理主路径，随后与门控路径结合，再经输出投影。多个块之间使用归一化与残差连接。该架构不包含自注意力层，也没有单独交替堆叠传统 MLP 块；论文实验使用的扩展系数为 2。

合成任务用于隔离机制本身。选择性复制实验以长度 4,096 的序列训练，词表含 16 种符号（包括噪声），要求模型记住其中 16 个数据符号。附录中的设置是两层、模型维度 64，并训练 400,000 步；图表显示，Mamba 架构与选择性 SSM 组合在该任务上的准确率为 99.8%。这项结果针对随机间隔下筛选和保留符号的受控任务。

归纳头实验测试模型能否根据序列中先前出现的符号关联来预测后续内容。论文训练两层模型，训练序列长度为 256、词表大小为 16；Mamba 的模型维度为 64，其他模型为 128。测试序列长度从 64 增至 1,048,576，论文报告 Mamba 的选择性 SSM 在该任务上可泛化到百万长度，是训练长度的 4,000 倍；注意力模型因内存限制只测试到 16,384。这个明确的百万长度测试属于归纳头实验，不是上段 4,096 长度的选择性复制设置。两种任务都属于受控合成测试，不能保证对任意长文本进行精确检索。

语言模型部分报告了预训练和零样本下游评测。缩放实验覆盖约 1.25 亿到 13 亿参数，并在 The Pile 数据集上对比多种架构；下游评测还比较了参数规模相近的公开模型。论文作者报告，Mamba 在所列规模和任务上具竞争力，部分对比接近或超过参数量约两倍的模型。结果依赖论文中的分词器、训练 token 数、上下文长度、训练配方和对照模型，不能概括成“任意 Mamba 都胜过任意 Transformer”。

语言缩放实验采用 Chinchilla 训练量配比，模型在 The Pile 上训练；零样本评测表则汇总 Pile 验证困惑度、LAMBADA 和常识推理等指标。论文说明 Mamba 与 Pythia 使用相同 tokenizer、数据和 3000 亿 token 训练长度，二者上下文长度为 2,048；其他模型的 tokenizer 或训练条件并不都完全相同。因此，比较必须看具体行和匹配条件，表格中的平均分不能脱离训练设置解释成架构的普遍排名。

基因组实验使用 DNA 序列建模并评估长序列分类。论文将上下文长度从 1,024 扩展到 1,048,576 个字符；其中一组上下文比较固定模型结构，并以约 3,300 亿训练 token 对比 Mamba 与 HyenaDNA。另一项下游测试从人类、狐猴、老鼠、猪、河马等早期任务变体调整为五种大型猿类的物种分类。作者报告 Mamba 的预训练困惑度随上下文增加而改善。它反映的是指定序列建模和物种分类设置，不是临床变异解释、基因功能发现或全部生物任务的能力证明。

音频评估包含 YouTubeMix 钢琴波形自回归预训练，以及 SC09 一秒数字语音片段生成；SC09 样本以 16 kHz 采样，内容是数字“零”至“九”。前者最长使用约 96 万采样点的片段，并在固定计算量下比较不同上下文长度；后者对语音生成使用自动评价指标。论文报告 Mamba 在所测设置中提升了部分指标，但也注明 SC09 数据集较小，较大的模型会过拟合。这些结果不是对一般语音识别、自然语音对话或商业语音系统的测试。

因此，“线性时间”指的是在隐藏维度等其他条件固定时，沿序列计算的主要成本随序列长度近似线性增长；它不是对任意长度、硬件和实现都更快的保证。注意力内核、批次形状、状态维度和数据搬运都会影响实际时延；短序列上启动和调度成本也可能占比较高。Mamba 的有限状态有助于节省随历史长度增长的缓存，但同时要求模型把历史压缩进状态，不能由复杂度结论推导出对任意过去 token 的逐字精确回忆。

论文把选择性扫描内核与卷积、FlashAttention-2 和未融合的 PyTorch 扫描实现进行比较。核心扫描微基准使用单批次、通道维度 1,024、状态维度 16、BF16 输入，并在 A100 80GB PCIe 上测量；其结论只涵盖所比较的核心操作，不含其他架构运算。论文报告，超过约 2,000 长度后该扫描在这些测试中快于所选注意力实现，并相对未融合的标准扫描快约 20 至 40 倍。

端到端推理测试使用 2,048 长度提示词、生成 128 个 token，在 batch size 1 到 128 间测量，并在相同 A100 上把 Mamba 1.4B 与 Transformer 1.3B 对比；作者还比较了未经训练的 Mamba 6.9B 与 Transformer 6.7B。论文报告在这套基准中 Mamba 的推理吞吐量约为相近规模 Transformer 的 4 至 5 倍，原因之一是递归状态不需要随序列增长的完整 KV 缓存，因而可提高批次规模。数字受 GPU、批次、提示词长度、模型实现和指标影响；不能当作当前部署中的速度承诺。论文的 125M 训练内存表格也显示，在其设置下 Mamba 与使用 FlashAttention-2 的 Transformer 处于相近量级，但 Mamba 的测得数值略高。

Mamba 描述的是模型怎样沿序列更新表示和状态。PagedAttention 讨论的是 Transformer 服务如何分配、寻址和共享注意力 KV 缓存；它属于推理系统的内存管理层，不是 Mamba 的选择性递推。推理时计算扩展则是在回答阶段增加采样、搜索或修订等测试时计算，属于另一种推理策略。三者可以处在同一系统的不同层面，但概念和实验指标不能互换。

Albert Gu、Tri Dao，Mamba: Linear-Time Sequence Modeling with Selective State Spaces，arXiv:2312.00752，v2，2024。全文：https://arxiv.org/html/2312.00752v2

Woosuk Kwon 等，Efficient Memory Management for Large Language Model Serving with PagedAttention，arXiv:2309.06180，2023。全文：https://arxiv.org/html/2309.06180

Charlie Snell 等，Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters，arXiv:2408.03314，2024。全文：https://arxiv.org/html/2408.03314