论文标题中译Mamba基于选择性状态空间的线性时间序列建模COLM 2024Conference on Language Modeling语言建模顶会AbstractFoundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic‑time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers’ computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content‑based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware‑aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end‑to‑end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5× higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million‑length sequences. As a general sequence model backbone, Mamba achieves state‑of‑the‑art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba‑3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.翻译 如今推动深度学习众多前沿应用的基础模型几乎全部建立在 Transformer 架构及其核心注意力模块之上。为了解决Transformer 处理长序列时计算效率低下的问题学界已经提出大量次二次时间复杂度架构包括线性注意力、门控卷积、循环模型以及结构化状态空间模型SSM但在语言这类关键数据模态上这类模型的效果始终不及注意力机制。本文发现这类模型的核心缺陷在于无法执行基于内容的推理并据此提出多项改进。第一令 SSM 参数成为输入的函数以此解决离散模态下的性能短板模型可以依据当前 token 的内容沿着序列维度选择性地传播或者遗忘信息。第二虽然该改动使得模型无法再使用高效卷积实现我们设计了一套面向硬件优化的循环模式并行算法。我们将该选择性状态空间模块Selective SSM整合进一套简化的端到端神经网络架构得到Mamba该架构完全不使用注意力甚至去除了传统 MLP 模块。Mamba 具备极快的推理速度吞吐量是 Transformer 的 5 倍计算量随序列长度呈线性增长在真实数据上即使序列长度达到百万级别模型性能仍可以持续提升。作为通用序列模型骨干Mamba 在语言、音频、基因组学等多种模态上取得 SOTA 结果。在语言建模任务上3B 规模的 Mamba预训练与下游评测效果均优于同等参数量的 Transformer并且可以达到参数量为其两倍的 Transformer 的性能水平。创新点1提出选择性状态空间 S6Selective‑SSM机制打破传统 SSM 的线性时不变LTI约束让 SSM 内部的(Delta,B,C)参数成为输入的函数。传统 S4 的参数是网络静态可学习参数和输入 token 无关而 S6 会根据当前输入内容动态生成状态空间参数使模型拥有基于内容选择记忆或者遗忘信息的能力解决旧 SSM 无法完成选择性复制、归纳头这类需要内容感知推理的任务。2设计硬件感知的并行扫描parallel scan算法。由于 S6 参数随输入变化不再满足线性时不变条件因此无法复用 S4 的快速卷积实现方式只能走循环递推模式。普通串行循环计算速度慢、无法 GPU 并行本文通过内核融合 kernel‑fusion、重计算 recomputation 技术实现硬件友好的并行扫描算子不在高带宽显存 HBM 中完整存储庞大中间状态只在高速 SRAM 完成离散化与递推计算既保留选择性 SSM 的表达能力又实现很高的 GPU 运算速度规避循环模式带来的性能缺陷。3构建极简 Mamba Block 网络架构抛弃自注意力模块并且将传统的卷积、门控操作、SSM 变换、MLP 映射合并为单一的同质网络块。对比 H3 等旧 SSM 架构Mamba 块不再把 SSM 与独立 MLP 块交错堆叠将门控与投影操作融合进同一个模块。整套模型完全不使用 Attention依靠堆叠 Mamba 块完成序列建模推理阶段为纯循环模式不需要 KV 缓存推理吞吐量达到同等规模 Transformer 的 5 倍计算复杂度严格为O(N)线性复杂度支持百万 token 级别的超长序列输入。SSMStructured State‑Space Models结构化状态空间模型SSM即结构化状态空间模型Structured State‑Space Models脱胎于控制理论里的经典状态空间系统是一类用于序列建模的基础模型框架Mamba 是 SSM 体系下的重要改进变体二者不能等同。从控制论本源来看状态空间系统使用一组隐状态用来存储、记忆序列的历史信息建立输入信号到输出信号的映射。连续形式下输入信号 x(t) 通过隐状态 h(t) 完成演化再映射得到输出 y(t)。其中矩阵A控制隐状态自身的演化B控制输入对隐状态的写入C负责从隐状态读出输出信息。由于深度学习处理的是离散的 token 序列不能直接使用连续微分方程SSM 通过零阶保持ZOH离散化规则把连续系统转换为离散时间递推公式得到离散版本的状态更新方程这才是深度学习代码中实际运行的计算形式。传统结构化 SSM 代表工作是S4。S4 有一大突出优势在线性时不变 LTILinear Time‑Invariant的约束之下系统全部参数不随时间步、输入内容变化整套系统既可以写成循环递推形式也等价于一个全局卷积核。训练阶段可以用卷积实现完全并行计算推理阶段切换为循环模式逐 token 生成序列计算复杂度接近线性O(N)相比 Transformer 注意力O(N^2)平方复杂度在长序列场景具备天然效率优势早期 S4 在音频、基因组这类连续信号任务上表现良好。但 LTI 线性时不变同时也是传统 SSM 的核心缺陷。因为所有时间步共用一套固定参数系统行为和输入内容无关模型没有 “内容感知” 的能力。不管输入的是有效信息还是噪声隐状态更新的规则固定不变。Mamba 原论文使用选择性复制、归纳头两个合成任务直观展示该缺陷面对混杂噪声的序列传统 LTI‑SSM 无法做到看到有效内容就记住、看到噪声就遗忘只能执行一套统一的运算逻辑。这直接造成 S4 在语言这类离散、信息稠密的任务上效果远不如 Transformer。简单来说LTI‑SSM 拥有长序列记忆的结构但是没有根据输入内容动态调控记忆的能力。为了解决 LTI 带来的局限Mamba 提出选择性SSMS6依然沿用 SSM 状态空间的基础数学框架但打破线性时不变约束。不再使用全局固定的参数而是通过可学习线性投影让每一个时间步的系统参数都由当前输入 x_t 动态生成。每一步拥有专属的系统参数模型就获得了内容感知能力可以依据输入的内容选择性地记忆或者丢弃信息。打破 LTI 约束带来新的工程难题参数随输入变化后模型不再等价于卷积无法复用卷积做并行训练只能走循环递推。普通循环实现速度很慢于是论文设计硬件感知的并行扫描算法搭配内核融合、重计算技术。不在 GPU 高带宽显存 HBM 保存体积庞大的中间状态将参数直接载入高速 SRAM 完成离散化和递推计算计算结束后再回写输出以此抵消循环模式带来的性能损耗保证 GPU 运行效率。选择性 SSMS6是 Mamba 的底层核心计算单元不等同于完整 Mamba 网络。Mamba 将 S6 模块、一维局部卷积、门控、线性投影、SiLU 激活整合为 Mamba Block大量同构堆叠该模块构建出一套完全不使用注意力的完整序列骨干网络。重要概念区分SSM 是大类S4 是 LTI‑SSMMamba 内部 S6 是选择性 SSMSSM 原生为一维序列模型图像、医学 3D 分割任务U‑Mamba、SegMamba需要把图像 / 体数据转成一维序列才能调用 SSM 模块。SSMS6框架介绍带硬件感知状态扩展的选择性状态空间模型传统状态空间模型SSM为了避免存储庞大的隐状态会固定参数不随输入变化并走高效计算路径但这样缺乏内容推理能力这个新模型给 SSM 加了 “选择机制”—— 让参数随输入动态变化能根据当前内容选信息同时通过硬件优化只在 GPU 高速存储里处理隐状态既解决了传统 SSM 的局限又没牺牲效率是 Mamba 的核心模块结构。引入选择机制的动机我们认为序列建模的核心问题是将上下文压缩到更小的状态中。事实上我们可以从这一角度理解主流序列模型的权衡取舍。例如注意力机制之所以高效且低效是因为它完全不显式压缩上下文。补充上下文压缩可以理解为把一长段历史上下文信息浓缩、存放到一个维度有限的隐状态向量 h 里面。序列处理的时候输入是很长的一串 token上下文 context。Transformer注意力几乎不做压缩。处理的时候要完整保留全部历史 token自回归推理要存完整 KV 缓存。所有历史信息都原封不动保存下来需要的时候再做两两交互。优点信息不丢建模能力强缺点存全部历史存储、计算代价随序列变长爆炸复杂度O(N^2)。RNN / S4 这类循环型模型必须做上下文压缩但是它压缩规则是写死固定的。不管输入序列有多长历史全部信息只能装进一个固定大小的隐状态 h_t。每来一个新输入就更新一遍这个有限大小的状态旧的原始输入 token 不会再保存。优点隐状态尺寸固定推理每一步常数开销内存不随序列长度暴涨痛点有限向量容量有限压缩过程一定会丢失信息。模型性能上限完全取决于这个压缩做得好不好能不能把真正有用的信息留下来扔掉无关噪声。后续希望压缩过程是内容感知的引入选择机制selection mechanism。为理解这一原理我们聚焦于两个合成任务示例1选择性复制任务Selective Copying对经典复制任务进行修改改变需记忆 token 的位置。该任务需要基于内容的推理以记忆相关 token彩色并过滤无关 token白色2归纳头任务Induction Heads一种知名机制被认为是解释大型语言模型上下文学习能力的关键。该任务需要基于上下文的推理以在适当的上下文中生成正确输出黑色。这些任务揭示了线性时间不变模型的失效模式。从循环视角来看其恒定的动态特性无法让它们从上下文中筛选正确信息也无法以输入依赖的方式影响沿序列传递的隐状态。从卷积视角来看已知全局卷积可解决标准复制任务因为该任务仅需时间感知但由于缺乏内容感知它们难以解决选择性复制任务。更具体地说输入与输出之间的间隔是变化的无法通过静态卷积核建模。总之序列模型的效率与有效性权衡取决于状态压缩效果高效模型必须具备小状态而有效模型的状态必须包含上下文中的所有必要信息。由此我们提出构建序列模型的核心原则选择性即基于上下文聚焦或过滤输入到序列状态的能力。具体而言选择机制控制信息沿序列维度的传播或交互方式。整合选择机制到模型将选择机制整合到模型中的一种方法是让影响序列交互的参数例如循环神经网络的循环动态或卷积神经网络的卷积核依赖于输入。算法 1 和算法 2 展示了使用的主要选择机制。核心差异在于将 A、B、C 等多个参数设计为输入的函数并相应调整张量形状。特别需要强调的是这些参数现在包含长度维度 L意味着模型从时间不变变为时变。这一变化导致模型不再等同于卷积。克服现有方法局限的思路总体而言循环神经网络如状态空间模型始终需要在表达能力和速度之间权衡隐状态维度越大的模型理论上效果越好但速度越慢。因此希望在不牺牲速度和内存的前提下最大化隐状态维度循环模式比卷积模式更灵活因为后者是通过展开前者推导而来。然而这需要计算并存储形状为B, L, D, N的隐状态 h其规模乘以状态空间模型的状态维度 N远大于形状为B, L, D的输入 x 和输出 y。因此人们提出了更高效的卷积模式无需计算状态仅需存储大小为B, L, D的卷积核现有线性时间不变状态空间模型利用循环 - 卷积双重形式将有效状态维度提升 N 倍约 10-100远大于传统循环神经网络且无效率损失。选择性扫描概述选择机制的设计旨在克服线性时间不变模型的局限同时我们需要重新审视状态空间模型的计算问题。通过三种经典技术解决这一问题核融合、并行扫描和重计算。朴素循环计算的计算量为 O (BLDN)而卷积计算的计算量为 O (BLD log (L))且前者的常数因子更小。因此对于长序列和不太大的状态维度 N循环模式的计算量实际上更小 面临的两大挑战是循环的顺序性和巨大的内存消耗。为解决后者与卷积模式类似可以尝试不直接存储完整状态 h。核心思路是利用现代加速器GPU的特性仅在更高效的存储层次结构中存储状态 h。具体而言大多数操作除矩阵乘法外都受内存带宽限制。这包括扫描操作通过核融合减少内存 IO从而比标准实现显著提速。具体来说不将大小为B, L, D, N的扫描输入存储在 GPU 的高带宽内存HBM中而是直接将状态空间模型参数Δ、A、B、C从慢速 HBM 加载到快速静态随机存取存储器SRAM在SRAM 中执行离散化和循环计算然后将大小为B, L, D的最终输出写回 HBM。为避免顺序循环发现尽管模型不再是线性的但仍可通过高效并行扫描算法实现并行化。最后还需避免存储反向传播所需的中间状态。巧妙地应用经典的重计算技术以降低内存需求不存储中间状态而是在反向传播时当输入从 HBM 加载到 SRAM 时重新计算中间状态。最终融合的选择性扫描层与优化的 FlashAttention Transformer 实现具有相同的内存需求。Mamba简介Mamba 是一种高效的深度学习序列建模架构于2023年提出基于选择性状态空间模型Selective State Space Model通过引入输入依赖的动态机制使模型能够有选择地处理和保留信息。相比传统的 TransformerMamba 具有线性时间复杂度O(L能高效处理超长序列如数万个 token在语言建模、基因组学和音频处理等任务中表现出色同时显著降低计算和内存开销被视为下一代序列建模的重要方向之一。H3Hungry Hungry Hippos它是此前基于状态空间模型SSM的代表性架构之一核心设计是 “将线性注意力的循环逻辑与 SSM 结合”整体采用 “门控连接 SSM 层 局部卷积” 的组合结构具体来说每个 H3 块会先通过一个 “移位 SSM”本质是简单的局部卷积处理局部序列信息再接入核心的 S4 层一种经典结构化 SSM建模长程依赖同时用门控机制类似 RNN 的门控控制信息的传递与过滤最后还会与线性注意力的思路结合以增强灵活性。不过 H3 的局限在于架构相对复杂且内部的 SSM 层是传统的 “线性时间不变LTI” 类型 —— 参数不随输入变化无法像 Mamba 那样根据内容动态选择信息这也导致它在语言等离散模态上的表现不如 Transformer。Gated MLP门控多层感知机它是 Transformer 架构中 MLP 块的常见增强形式核心是 “用门控机制提升 MLP 的非线性表达能力”传统 MLP 通常是 “线性投影 激活函数” 的简单堆叠而 Gated MLP 会在其中加入 “门控分支”—— 比如先将输入通过两个并行的线性层生成 “信号路” 和 “门控路”再通过逐元素相乘的方式让门控路动态控制信号路的信息权重类似 “开关”常见的 SwiGLU 就是 Gated MLP 的典型变体用 SiLU 激活函数作为门控。在 Transformer 中Gated MLP 的作用是对注意力层输出的 “关联信息” 进行非线性变换补充模型的局部特征建模能力但它本身不具备长程依赖建模能力必须依赖注意力层这也导致 Transformer 整体仍受限于二次时间复杂度。Mamba块的设计正是在 H3 和 Gated MLP 的基础上做了 “整合与简化”。把 H3 的 SSM、一维卷积组件嵌入到 Gated MLP 简洁的双路门控框架内合并了两个模型的优势。Mamba 模块的输入先经过两路线性投影其中一路作为主分支先经过一维卷积做局部时序特征提取再送入选择性 SSMS6之后使用 SiLU 激活函数另一路分支直接经过 SiLU 激活。两条分支的输出进行逐元素相乘最后经过顶层线性投影得到这个 Mamba Block 的输出。顺着数据流完整梳理一遍 Mamba 块内部的计算流程。输入特征首先送入输入线性投影分裂成两个分支。主分支先经过一维卷积用来提取局部邻域信息补充局部上下文紧接着进入核心的选择性 SSM 模块也就是 S6。这里的 SSM 不再是 H3 里传统 LTI 的 S4而是带选择机制的 SSM能够依据输入内容动态生成系统参数实现选择性记忆与遗忘完成长距离上下文压缩。经过 SSM 之后使用 SiLU 激活函数引入非线性。另外一条支路仅经过 SiLU 激活作为门控分支。主分支输出与门控分支逐元素相乘用门控机制控制信息流过滤无用特征相乘后的结果送入输出线性投影层映射回原始通道维度最后配合残差连接完成整个 Mamba Block 的运算。从架构设计思路上可以总结H3 证明 SSM 可以和门控结合用来建模长序列Gated MLP 证明双分支门控结构简单有效。Mamba 没有保留 H3 复杂的三路分支借鉴 Gated MLP 的两路门控范式将 ConvSSM 嵌入其中。这个设计一举两得既保留 SSM 强大的长序列建模能力维持线性 O(N) 的计算复杂度又复用 Gated MLP 简洁的门控结构精简网络分支降低冗余计算。整个完整 Mamba 网络就是大量同构的 Mamba Block 堆叠而成全程不需要注意力机制。