基础模型现在为深度学习中大多数令人兴奋的应用程序提供支持，几乎普遍基于 Transformer 架构及其核心注意力模块。许多次二次时间架构（例如线性注意力、门控卷积和循环模型以及结构化状态空间模型（SSM））已被开发来解决 Transformers 在长序列上的计算效率低下问题，但它们的表现不如注意力重要的方式，例如语言。我们发现此类模型的一个关键弱点是它们无法执行基于内容的推理，并做出一些改进。首先，简单地让 SSM参数作为输入的函数，可以解决其离散模态的弱点，允许模型根据当前标记选择性地沿序列长度维度传播或忘记信息。其次，尽管这种变化阻碍了高效卷积的使用，但我们在循环模式下设计了一种硬件感知的并行算法。我们将这些选择性 SSM 集成到简化的端到端神经网络架构中，无需注意，甚至不需要 MLP 块 (Mamba)。 Mamba 享有快速推理（比 Transformer 高出 5 倍的吞吐量）和序列长度的线性缩放，并且其性能在高达百万长度序列的实际数据上得到改进。作为通用序列模型骨干，Mamba 在语言、音频和基因组学等多种模式上实现了最先进的性能。在语言建模方面，我们的 Mamba-3B 模型在预训练和下游评估方面都优于相同大小的 Transformer，并且与两倍大小的 Transformer 相匹配。

基础模型现在为深度学习中大多数令人兴奋的应用程序提供支持，几乎普遍基于 Transformer 架构及其核心注意力模块。许多次二次时间架构（例如线性注意力、门控卷积和循环模型以及结构化状态空间模型（SSM））已被开发来解决 Transformers 在长序列上的计算效率低下问题，但它们的表现不如注意力重要的方式，例如语言。我们发现此类模型的一个关键弱点是它们无法执行基于内容的推理，并做出一些改进。首先，简单地让 SSM参数作为输入的函数，可以解决其离散模态的弱点，允许模型根据当前标记选择性地沿序列长度维度传播或忘记信息。其次，尽管这种变化阻碍了高效卷积的使用，但我们在循环模式下设计了一种硬件感知的并行算法。我们将这些选择性 SSM 集成到简化的端到端神经网络架构中，无需注意，甚至不需要 MLP 块 (Mamba)。 Mamba 享有快速推理（比 Transformer 高出 5 倍的吞吐量）和序列长度的线性缩放，并且其性能在高达百万长度序列的实际数据上得到改进。作为通用序列模型骨干，Mamba 在语言、音频和基因组学等多种模式上实现了最先进的性能。在语言建模方面，我们的 Mamba-3B 模型在预训练和下游评估方面都优于相同大小的 Transformer，并且与两倍大小的 Transformer 相匹配。

最近，结构化状态空间序列模型 (SSM)（Gu、Goel 和 Ré，2022 年；Gu、Johnson、Goel 等人，2021 年）已成为一类有前途的序列建模架构。这些模型可以解释为递归神经网络（RNN）和卷积神经网络（CNN）的组合，其灵感来自经典状态空间模型（Kalman 1960）。此类模型可以非常有效地计算为递归或卷积，并且序列长度具有线性或近线性缩放。此外，他们还坚持原则在某些数据模式中对远程依赖关系进行建模的机制（Gu、Dao 等人，2020 年），并且主导了 Long Range Arena 等基准测试（Tay、Dehghani、Abnar 等人，2021 年）。 SSM 的多种风格（Gu, Goel, and Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, andBerant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. al. 2023；Smith、Warrington 和 Linderman 2023）在音频和视觉等涉及连续信号数据的领域取得了成功（Goel 等人，2022；Nguyen、Goel 等人，2022；Saon、Gupta 和 Cui，2023） 。然而，它们在对离散和信息密集的数据（例如文本）进行建模方面效率较低。

我们提出了一类新的选择性状态空间模型，它改进了先前在多个轴上的工作，以实现 Transformer 的建模能力，同时线性缩放序列长度。选择机制。首先，我们确定了先前模型的一个关键限制：以依赖于输入的方式有效选择数据的能力（即关注或忽略特定输入）。基于选择性复制和感应头等重要合成任务的直觉，我们通过根据输入参数化 SSM 参数来设计一种简单的选择机制。这使得模型能够过滤掉不相关的信息并无限地记住相关信息。硬件感知算法。这个简单的改变给模型的计算带来了技术挑战；事实上，所有先前的 SSM 模型都必须具有时间和输入不变性，才能提高计算效率。我们通过硬件感知算法克服了这个问题，该算法通过扫描而不是卷积来循环计算模型，但不会具体化扩展状态，以避免 GPU 内存层次结构的不同级别之间的 IO 访问。由此产生的实现在理论上（与所有基于卷积的 SSM 的伪线性相比，序列长度线性缩放）和现代硬件上（在 A100 GPU 上最多快 3倍）都比以前的方法更快。建筑学。我们通过将先前的 SSM 架构（Dao、Fu、Saab 等人，2023）的设计与 Transformers 的 MLP 块组合成一个块来简化先前的深度序列模型架构，从而形成简单且同质的架构设计（Mamba），选择性状态空间。选择性 SSM 以及扩展的 Mamba 架构是完全循环模型，其关键属性使其适合作为在序列上运行的通用基础模型的骨干。 (i) 高质量：选择性带来了语言和基因组学等密集模式的强大性能。 (ii) 快速训练和推理：在训练期间计算和内存随序列长度线性缩放，并且在推理期间自回归展开模型仅需要每步恒定的时间，因为它不需要缓存先前的元素。 (iii) 长上下文：质量和效率共同提高了序列长度高达 1M 的实际数据的性能。我们根据经验验证了 Mamba 作为通用序列 FM 主干的潜力，在预训练质量和特定领域任务性能方面，在几种类型的模式和设置上：§ 合成材料。在重要的合成任务上，例如复制和归纳头，这些任务被认为是大型语言模型的关键，Mamba 不仅可以轻松解决它们，而且可以无限长地推断解决方案（> 1M 个令牌）。音频和基因组学。在音频波形和 DNA 序列建模方面，Mamba 的性能优于之前最先进的模型，例如 SaShiMi、Hyena 和 Transformers，无论是在预训练质量还是下游指标方面（例如，在具有挑战性的语音生成数据集上将FID 减少更多）超过一半）。在这两种设置中，随着上下文长度达到百万长度的序列，其性能都会得到提高。§ 语言建模。 Mamba 是第一个真正实现 Transformer 质量性能的线性时间序列模型，无论是在预训练困惑度还是下游评估方面。通过高达 1B 参数的缩放法则，我们证明 Mamba 的性能超过了大范围基线，包括基于LLaMa 的非常强大的现代 Transformer 训练方法（Touvron 等人，2023）。与类似大小的 Transformer 相比，我们的 Mamba 语言模型具有 5 代吞吐量，并且 Mamba-3B 的质量与两倍大小的 Transformer 相当（例如，与Pythia-3B 相比，常识推理平均高出 4 分，甚至超过Pythia-7B）。

2 状态空间模型结构化状态空间序列模型 (S4) 是最近一类用于深度学习的序列模型，与 RNN、CNN 和经典状态空间模型广泛相关。它们受到特定的连续系统 (1) 的启发，该系统通过隐式潜在状态映射一维函数或序列 ( ) → ( ) � ( )→ ( ) � ¤ 。具体来说，S4模型由四个参数（α、A、B、C）定义，它们定义了两个阶段的序列到序列的转换。

通常，该模型使用卷积模式 (3) 进行高效的可并行训练（提前查看整个输入序列），并切换到循环模式 (2) 进行高效的自回归推理（其中输入一次看到一个时间步）。线性时不变性 (LTI)。方程 (1) 至 (3) 的一个重要属性是模型的动态随时间变化是恒定的。换句话说，(→,A, B, C) 以及 (A, B) 对于所有时间步都是固定的。这种性质称为线性时不变性（LTI），它与递归和卷积密切相关。非正式地，我们认为 LTI SSM 等同于任何线性递推 (2a) 或卷积 (3b)，并使用 LTI 作为这些模型类别的总称。到目前为止，由于基本效率限制，所有结构化 SSM 都是 LTI（例如计算为卷积），如第 3.3 节中讨论的。然而，这项工作的核心见解是，LTI 模型在对某些类型的数据进行建模时存在根本局限性，而我们的技术贡献包括消除 LTI 约束，同时克服效率瓶颈。结构和尺寸。最后，我们注意到结构化 SSM 之所以如此命名，是因为有效计算它们还需要在 A 矩阵上强加结构。最流行的结构形式是对角线（Gu, Gupta, et al. 2022；Gupta, Gu, and Berant 2022；Smith, Warrington,and Linderman 2023），我们也使用它。此时，A→→→、B→→→1、C→→1→矩阵都可以用数字来表示。为了通过通道对批量大小和长度的输入序列进行操作，SSM 独立应用于每个通道。请注意，在这种情况下，总隐藏状态具有每个输入的维度，并且在序列长度上计算它需要 ( ) 时间和内存；这是 3.3 节中提到的基本效率瓶颈的根源。一般状态空间模型。我们注意到，术语"状态空间模型"具有非常广泛的含义，它仅表示具有潜在状态的任何循环过程的概念。它已被用来指代不同学科中的许多不同概念，包括马尔可夫决策过程（MDP）（强化学习（Hafner et al. 2020））、动态因果模型（DCM）（计算神经科学（Friston、Harrison、和 Penny2003））、卡尔曼滤波器（控制（Kalman 1960））、隐马尔可夫模型 (HMM) 和线性动力系统 (LDS)（机器学习）以及循环（有时是卷积）模型（深度学习） ）。在整篇论文中，我们使用术语"SSM"来专门指代结构化 SSM 或 S4 模型类别（Gu、Goel 和 Ré，2022 年；Gu、Gupta 等人，2022 年；Gupta、Gu 和Berant 2022；Hasani 等人 2023；Ma 等人 2023；Smith、Warrington 和 Linderman 2023），并且这些术语可以互换使用。为了方便起见，我们还可能包括此类模型的衍生模型，例如那些关注线性递归或全局卷积观点的模型（Y. Li et al. 2023; Orvieto et al. 2023; Poli et al. 2023），并澄清必要时的细微差别。SSM 架构。 SSM 是独立的序列转换，可以合并到端到端神经网络架构中。 （我们有时也将 SSM 架构称为SSNN，SSNN 之于 SSM 层，就像 CNN 之于线性卷积层一样。）我们讨论了一些最著名的 SSM 架构，其中许多架构也将作为我们的主要基线。线性注意力（Katharopoulos et al. 2020）是自我注意力的近似，涉及重复，可以被视为退化的线性 SSM。H3 (Dao, Fu, Saab, et al. 2023) 将这种递归推广到使用 S4；它可以被视为一种架构，其中 SSM 夹在两个门控连接之间（图 3）。 H3 还在主 SSM 层之前插入了一个标准局部卷积，将其构建为移位 SSM。Hyena (Poli et al. 2023) 使用与 H3 相同的架构，但用 MLP 参数化全局卷积替换 S4 层 (Romero et al.2021)。RetNet（Y. Sun 等人，2023）在架构中添加了一个额外的门，并使用更简单的 SSM，允许替代的可并行计算路径，使用多头注意力 (MHA) 的变体而不是卷积。

RWKV (B. Peng et al. 2023) 是一种最新的 RNN，设计用于基于另一种线性注意力近似（attention-freeTransformer (S. Zhai et al. 2021)）的语言建模。其主要"WKV"机制涉及 LTI 复发，可以视为两个 SSM 的比率。其他密切相关的 SSM 和架构将在扩展的相关工作（附录 B）中进一步讨论。我们特别强调 S5 (Smith, Warrington, andLinderman 2023)、QRNN (Bradbury et al. 2016) 和 SRU (Lei et al. 2017)，我们认为它们是与我们的核心选择性 SSM 最密切相关的方法。3 选择性状态空间模型我们利用合成任务的直觉来激发我们的选择机制（第 3.1 节），然后解释如何将该机制合并到状态空间模型中（第 3.2 节）。由此产生的时变 SSM 无法使用卷积，这对如何有效地计算它们提出了技术挑战。我们通过利用现代硬件上的内存层次结构的硬件感知算法克服了这个问题（第 3.3 节）。然后，我们描述了一个简单的SSM 架构，没有注意力机制，甚至没有 MLP 模块（第 3.4 节）。最后，我们讨论选择机制的一些附加属性（第3.5 节）。3.1 动机：选择作为压缩手段我们认为序列建模的一个基本问题是将上下文压缩成更小的状态。事实上，我们可以从这个角度来看待流行序列模型的权衡。例如，注意力既有效又低效，因为它明确地根本不压缩上下文。从这一点可以看出，自回归推理需要显式存储整个上下文（即KV缓存），这直接导致Transformers的线性时间推理和二次时间训练缓慢。另一方面，循环模型是高效的，因为它们具有有限状态，这意味着恒定时间推理和线性时间训练。然而，它们的有效性受到这种状态压缩上下文的程度的限制。为了理解这一原理，我们重点关注两个合成任务的运行示例（图 2）。选择性复制任务通过改变要记忆的标记的位置来修改流行的复制任务（Arjovsky、Shah 和 Bengio 2016）。它需要内容感知推理才能记住相关标记（彩色）并过滤掉不相关标记（白色）。▪ Induction Heads 任务是一种众所周知的机制，假设可以解释法学硕士的大多数情境学习能力（Olsson 等人，2022）。它需要上下文感知推理来知道何时在适当的上下文（黑色）中产生正确的输出。这些任务揭示了 LTI 模型的失效模式。从循环的角度来看，它们的恒定动态（例如（2）中的（A，B）转换）不能让它们从上下文中选择正确的信息，或者影响沿输入相关的序列传递的隐藏状态方式。从卷积的角度来看，众所周知，全局卷积可以解决普通复制任务（Romero et al. 2021），因为它只需要时间感知，但由于缺乏内容，它们在选择性复制任务上有困难-意识（图2）。更具体地说，输入到输出之间的间距是变化的，并且不能通过静态卷积核进行建模。总之，序列模型的效率与有效性权衡的特征在于它们压缩状态的程度：有效模型必须有一个小状态，而有效模型必须有一个小状态。包含上下文中所有必要信息的状态。反过来，我们提出构建序列模型的基本原则是选择性：或者是关注或过滤输入到序列状态的上下文感知能力。特别是，选择机制控制信息如何沿着序列维度传播或交互（更多讨论请参见第 3.5 节）。3.2 通过选择改进 SSM将选择机制纳入模型的一种方法是让影响序列交互的参数（例如 RNN 的循环动态或 CNN 的卷积核）依赖于输入

3.3 选择性 SSM 的有效实施诸如卷积（Krizhevsky、Sutskever 和 Hinton 2012）和 Transformer（Vaswani 等人 2017）之类的硬件友好架构得到了广泛的应用。在这里，我们的目标是使选择性 SSM 在现代硬件 (GPU) 上也变得高效。选择机制是非常自然的，早期的作品尝试纳入选择的特殊情况，例如让 α 在循环 SSM 中随时间变化（Gu、Dao 等人，2020）。然而，正如前面提到的，SSM 使用的一个核心限制是它们的计算效率，这就是为什么 S4 和所有衍生品都使用 LTI（非选择性）模型，最常见的是全局卷积的形式。3.3.1 先前模型的动机我们首先重新审视这一动机，并概述我们克服现有方法局限性的方法。▪ 在高层次上，SSM 等循环模型总是在表达性和速度之间进行权衡：如第 3.1 节中所讨论的，具有较大隐藏状态维数的模型应该更有效，但速度更慢。因此我们希望在不付出速度和内存成本的情况下最大化隐藏状态维度。请注意，循环模式比卷积模式更灵活，因为后者 (3) 是从扩展前者 (2) 衍生而来的 (Gu, Goel, and Ré 2022;Gu, Johnson, Goel,等人，2021）。然而，这需要计算和具体化形状为 ( , , , ) 的潜在状态，该形状比形状( , , ) 的输入和输出大得多（SSM 状态维度的一个因子）。因此引入了更高效的卷积模式，它可以绕过状态计算并具体化仅由( , , )组成的卷积核(3a)。先前的 LTI SSM 利用双循环卷积形式将有效状态维度增加 (± 10 – 100) 倍，比传统 RNN 大得多，且没有效率损失。3.3.2 选择性扫描概述：硬件感知状态扩展选择机制旨在克服LTI模型的局限性；与此同时，我们因此需要重新审视SSM的计算问题。我们通过三种经典技术来解决这个问题：内核融合、并行扫描和重新计算。我们提出两个主要观察结果：③ 朴素循环计算使用 ( ) FLOP，而卷积计算使用 ( log( )) FLOP，并且前者的常数因子较低。因此，对于长序列和不太大的状态维度，循环模式实际上可以使用更少的 FLOP。▪ 两个挑战是循环的顺序性和大量的内存使用。为了解决后者，就像卷积模式一样，我们可以尝试不实际实现完整状态 –。主要思想是利用现代加速器 (GPU) 的特性来实现状态——仅在内存层次结构的更有效级别中实现。特别是，大多数运算（矩阵乘法除外）都受到内存带宽的限制（Dao、Fu、Ermon 等人，2022 年；Ivanov 等人，2021年；Williams、Waterman 和 Patterson，2009 年）。这包括我们的扫描操作，并且我们使用内核融合来减少内存 IO 量，从而与标准实现相比显着加速。具体来说，我们不是在 GPU HBM（高带宽内存）中准备大小为 ( , , , ) 的扫描输入 (A, B)，而是直接从慢 HBM加载 SSM 参数 (→, A, B, C)为了快速 SRAM，在 SRAM 中执行离散化和递归，然后将大小为 ( , , ) 的最终输出写回 HBM。为了避免顺序重复，我们观察到尽管不是线性的，但仍然可以使用高效的并行扫描算法进行并行化（Blelloch1990；Martin 和 Cundy 2018；Smith、Warrington 和 Linderman 2023）。最后，我们还必须避免保存反向传播所必需的中间状态。我们小心地应用经典的重新计算技术来减少内存需求：当输入从 HBM 加载到 SRAM 时，中间状态不会被存储，而是在向后传递中重新计算。因此，融合选择性扫描层与使用 FlashAttention 的优化变压器实现具有相同的内存要求。融合内核和重新计算的详细信息参见附录 D。完整的选择性 SSM 层和算法如图 1 所示。

3.4 简化的SSM架构与结构化 SSM 一样，选择性 SSM 是独立的序列转换，可以"灵活地合并到神经网络中"。 H3 架构是最著名的SSM 架构（第 2 节）的基础，该架构通常由受线性注意力启发的块与 MLP（多层感知器）块交织组成。我们通过将这两个组件合并为一个组件来简化这一架构，该组件同质堆叠（图 3）。这是受到门控注意力单元（GAU）（Hua et al. 2022）的启发，它对注意力做了类似的事情。该架构涉及通过可控扩展因子来扩展模型维度。对于每个块，大多数参数 (3 2 ) 位于线性投影中（2 2 表示输入投影，2 表示输出投影），而内部 SSM 贡献较少。 SSM 参数的数量（α、B、C 和

3.5.2 选择机制解读我们详细阐述了选择的两种特殊的机械效应。可变间距。选择性允许过滤掉感兴趣的输入之间可能出现的不相关的噪声标记。这以选择性复制任务为例，但在常见数据模态中普遍存在，特别是对于离散数据，例如"um"等语言填充物的存在。出现此属性的原因是模型可以机械地过滤掉任何特定输入，例如在门控 RNN 情况（定理 1）中，当 ≤ 0 时。过滤上下文。根据经验观察，尽管更多的上下文应该带来严格的更好的性能这一原则，但许多序列模型并没有随着更长的上下文而改进（F. Shi et al. 2023）。一种解释是，许多序列模型无法在必要时有效地忽略不相关的上下文。一个直观的例子是全局卷积（和一般的 LTI 模型）。另一方面，选择性模型可以随时简单地重置其状态以删除无关的历史记录，因此原则上它们的性能随着上下文长度的增加而单调提高（例如第 4.3.2 节）。边界重置。在多个独立序列拼接在一起的设置中，Transformers 可以通过实例化特定的注意力掩模来将它们分开，而 LTI 模型将在序列之间泄露信息。选择性 SSM 还可以在边界处重置其状态（例如，当 ≤ 1 时，≤ ≤≤ ≤ 或定理 1）。这些设置可能是人工发生的（例如将文档打包在一起以提高硬件利用率）或自然发生的（例如强化学习中的情节边界（Lu et al. 2023））。此外，我们还详细阐述了每个选择参数的影响。解释∧。一般来说，∧控制关注或忽略当前输入的程度之间的平衡。它概括了 RNN 门（例如在定理 1 中），机械上，大的 α 会重置状态 α 并关注当前的输入 ，而小的 α 则保持状态并忽略当前的输入。 SSM (1)-(2)可以解释为按时间步长离散化的连续系统 -，在这种情况下，直觉是大 - - 表示系统专注于当前输入更长（因此"选择"它并忘记其当前状态），而较小的"0"表示被忽略的瞬态输入。A 的解释。我们注意到，虽然 A 参数也可以是选择性的，但它最终仅通过 A = exp(−A)（离散化 (4)）与 − 相互作用来影响模型。因此，α中的选择性足以保证(A，B)中的选择性，并且是改进的主要来源。我们假设除了（或代替）- 之外使 A 具有选择性将具有类似的性能，并且为了简单起见将其省略。B 和 C 的解释。正如第 3.1 节中所讨论的，选择性最重要的属性是"过滤掉不相关的信息，以便序列模型的上下文可以被压缩为有效的状态"。在 SSM 中，选择性地修改 B 和 C 可以更细粒度地控制是让输入进入状态 ，还是让状态进入输出 。这些可以解释为允许模型分别根据内容（输入）和上下文（隐藏状态）来调节循环动态。3.6 其他模型详细信息真实与复杂。大多数先前的 SSM 在其状态 α 中使用复数，这对于许多任务的出色性能是必要的（Gu、Goel 和Ré 2022）。然而，根据经验观察，完全实值的 SSM 在某些情况下似乎工作得很好，甚至可能更好（Ma et al.2023）。我们使用真实值作为默认值，这对于除一项任务之外的所有任务都很有效；我们假设复数与实数的权衡与数据模态中的连续离散谱有关，其中复数对于连续模态（例如音频、视频）有帮助，但对离散模态（例如文本、DNA）没有帮助。

初始化。大多数先前的 SSM 还建议特殊的初始化，特别是在复值情况下，这可以在多种设置（例如低数据状态）中有所帮助。我们对复杂情况的默认初始化是 S4D-Lin，对真实情况的默认初始化是 S4D-Real (Gu,Gupta, et al. 2022)，它基于 HIPPO 理论 (Gu, Dao, et al. 2020)。它们分别将 A 的第一个元素定义为−1−2 + 和 −( + 1)。然而，我们期望许多初始化能够正常工作，特别是在大数据和实值 SSM 体系中； 4.6 节考虑了一些消融。§ 参数化。我们将选择性调整定义为 α ( ) = ( 1 ( ))，这是由 α 的机制驱动的（第 3.5 节）。我们观察到它可以从 1 维推广到更大的维度。我们将其设置为 的一小部分，与块中的主要线性投影相比，它使用的参数数量可以忽略不计。我们还注意到，广播操作可以被视为另一个线性投影，初始化为 1 和 0 的特定模式；如果这个投影是可训练的，这会导致替代方案 - ( ) = ( ( ))，它可以被视为低秩投影。在我们的实验中，遵循 SSM 的先前工作（Gu、Johnson、Timalsina），将 α 参数（可以视为偏差项）初始化为 α1 α (([0.001, 0.1]))等人，2023）。备注 3.1。为了简化我们的实验结果，我们有时将选择性 SSM 缩写为 S6 模型，因为它们是具有选择机制并通过扫描计算的 S4 模型。4 实证评价在第 4.1 节中，我们测试 Mamba 解决第 3.1 节中提出的两个综合任务的能力。然后，我们对三个领域进行评估，每个领域都对自回归预训练以及下游任务进行评估。§ 4.3 节：DNA 序列预训练和长序列分类任务的微调。§ 4.4节：音频波形预训练，以及自回归生成的语音片段的质量。最后，第 4.5 节展示了 Mamba 在训练和推理时的计算效率，第 4.6 节消除了架构和选择性 SSM 的各个组件。4.1 综合任务这些任务的完整实验细节，包括任务细节和训练协议，请参见附录 E.1。4.1.1 选择性复制复制任务是序列建模中研究最深入的综合任务之一，最初旨在测试循环模型的记忆能力。正如第 3.1 节中所讨论的，LTI SSM（线性递归和全局卷积）可以通过仅跟踪时间而不是推理数据来轻松解决此任务；例如，通过构造一个长度恰好正确的卷积核（图 2）。这在早期的全局卷积工作中得到了明确的验证（Romero 等人，2021）。选择性复制任务通过随机化标记之间的间距来防止这种快捷方式。请注意，此任务之前已作为去噪任务介绍过（Jing 等人，2019）。请注意，许多先前的工作认为添加架构门控（乘法交互）可以赋予模型"数据依赖性"并解决相关任务（Dao、Fu、Saab 等人，2023 年；Poli 等人，2023 年）。然而，我们直观地发现这种解释是不够的，因为这种门控不会沿着序列轴相互作用，并且不能影响标记之间的间距。特别是架构门控不是选择机制的一个实例（附录 A）。表 1 证实，H3 和 Mamba 等门控架构只能部分提高性能，而选择机制（将 S4 修改为 S6）可以轻松解决此任务，特别是与这些更强大的架构结合使用时。

4.2 语言建模我们根据预训练指标（困惑度）和零样本评估，在标准自回归语言模型上与其他架构评估 Mamba 架构。我们设置模型尺寸（深度和宽度）以反映 GPT3 规范。我们使用 Pile 数据集（L.Gao、Biderman 等人，2020 年），并遵循 Brown 等人中描述的训练方法。 （2020）。所有培训细节均参见附录 E.2。4.2.1 缩放法则对于基线，我们与标准 Transformer 架构（GPT3 架构）以及基于 PaLM 和 LLaMa 的我们所知的最强 Transformer 配方（此处称为 Transformer++）进行比较

4.3.2 缩放：上下文长度在下一个 DNA 实验中，我们研究模型相对于序列长度的缩放特性。我们只比较 HyenaDNA 和 Mamba 模型，因为二次注意力在序列长度较长时变得异常昂贵。我们在序列长度 2 10 = 1024、2 12 = 4096、2 14 = 16384、2 16 =65536、2 18 = 262144、2 20 = 1048576 上预训练模型。我们将模型大小设置为 6 层，宽度为 128（大约1.3M-1.4M参数）。模型接受了 20 个梯度步骤的训练，总共 330 个标记。较长的序列长度使用类似于 (Nguyen,Poli, et al. 2023) 的序列长度预热。结果。图 5（右）显示，Mamba 能够利用更长的上下文，甚至长达 1M 的极长序列，并且其预训练困惑度随着上下文的增加而提高。另一方面，鬣狗 DNA 模型随着序列长度的增加而变得更糟。从 3.5 节关于选择机制属性的讨论中可以直观地看出这一点。特别是，LTI 模型不能选择性地忽略信息；从卷积的角度来看，一个非常长的卷积核正在聚合一个长序列上的所有信息

4.3.3 合成物种分类我们评估下游任务的模型，即通过随机采样连续样本来对 5 个不同物种进行分类。他们的 DNA 的 ous 片段。该任务改编自鬣狗DNA，其中使用了物种{人类、狐猴、小鼠、猪、河马}。我们通过对类人猿物种（人类、黑猩猩、大猩猩、猩猩、倭黑猩猩）进行分类，使任务变得更具挑战性，已知这些类人猿有 99%的 DNA 是相同的。4.4 音频建模和生成对于音频波形模态，我们主要与 SaShiMi 架构和训练协议进行比较（Goel 等人，2022）。该模型包括1. U-Net 主干网络具有两个池化阶段，每个阶段的模型维度加倍，2.在每个阶段交替使用S4和MLP块。我们在 YouTubeMix (DeepSound 2017) 上评估预训练质量（自回归下一个样本预测），这是之前工作使用的标准钢琴音乐数据集，由 4 小时的钢琴独奏音乐组成，采样率为 16000 Hz 预训练细节很大程度上遵循标准语言建模设置（第 4.2 节）。图 7 评估了在保持计算固定的情况下将训练序列长度从 2 13 = 8192 增加到 2 20 – 10 6 的效果。 （数据的整理方式存在一些轻微的边缘情况，这可能会导致缩放曲线出现扭结。例如，只有一分钟长的剪辑可用，因此最大序列长度实际上受到 60 × 16000= 960000 的限制.)Mamba 和 SaShiMi (S4+MLP) 基线都随着上下文长度的增加而持续改进；曼巴在整个比赛中表现更好，并且在更长的距离上差距会扩大。主要指标是每字节位数 (BPB)，它是用于预训练其他模态的标准负对数似然 (NLL) 损失的常数因子 log(2)。我们注意到一个重要的细节：这是本文中唯一一个从真实参数化切换到复杂参数化的实验（第 3.6 节）。我们在附录 E.4 中展示了额外的消融。4.4.2 自回归语音生成SC09 是一个基准语音生成数据集（Donahue、McAuley 和 Puckette 2019；Warden 2018），由以 16000 Hz 采样的 1秒剪辑组成，数字"零"到"九"具有高度可变的特征。我们很大程度上遵循 Goel 等人的自回归训练设置和生成协议。 （2022）。表 4 显示了 Mamba-UNet 模型的自动化指标与 Goel 等人的各种基线的比较。 (2022)：WaveNet（Oord 等人，2016）、SampleRNN（Mehri 等人，2017）、WaveGAN（Donahue、McAuley 和 Puckette，2019）、Diï ̧Wave（Z. Kong 等人，2021）和 SaShiMi。小型 Mamba 模型的性能优于最先进的（以及更大的）基于 GAN 和扩散的模型。与基线匹配的更大模型参数进一步显着提高了保真度指标。表 5 采用小型 Mamba 模型，研究了外部舞台和中心舞台的不同架构的组合。结果表明，Mamba 在外部块中始终优于 S4+MLP，在中心块中 Mamba > S4+MLP > MHA+MLP。

4.6 模型消融我们对模型的组件进行了一系列详细的消融，重点关注 Chinchilla 令牌计数大小为 3.5 亿模型的语言建模设置（与图 4 相同的设置）。4.6.1 架构表 6 研究了架构（块）及其内部 SSM 层的影响（图 3）。我们发现? 在之前的非选择性（LTI）SSM 中，相当于全局卷积，性能非常相似。用实值 S4 变体替换之前工作中的复值 S4 变体不会对性能产生太大影响，这表明（至少对于 LM）在考虑硬件效率时，实值 SSM 可能是更好的选择Ø 效率。用选择性 SSM (S6) 替换其中任何一个都可以显着提高性能，验证第 3 节的动机。Mamba 架构的性能与 H3 架构类似（并且在使用选择性层时似乎稍好一些）。我们还研究了 Mamba 块与其他块的交错，例如附录 E.2.2 中的 MLP（传统架构）MHA（混合注意力架构）。表 7 通过考虑选择性 α、B 和 C 参数的不同组合（算法 2）来消除选择性 SSM 层，表明 α 是最重要的参数，因为它与 RNN 门控相关（定理1）。表 8 考虑了 SSM 的不同初始化，这些初始化已被证明在某些数据模式和设置中产生很大差异（Gu、Goel 和Ré，2022 年；Gu、Gupta 等人，2022 年）。在语言建模方面，我们发现更简单的实值对角线初始化（S4D-Real，第 3 行）而不是更标准的复值参数化（S4D-Lin，第 1 行）表现更好。随机初始化也效果很好，与之前工作的结果一致（Mehta 等人，2023）。表 9 和表 10 分别考虑改变 α 和 (B, C) 投影的尺寸。将它们从静态更改为选择性提供了最大的好处，同时进一步增加维度通常会通过参数数量的小幅增加来适度提高性能。特别值得注意的是，当状态大小增加时，选择性 SSM 的显着改进，仅花费 1% 的额外参数，困惑度就提高了1.0 以上。这验证了我们在 3.1 和 3.3 节中的核心动机。

5 讨论我们讨论相关工作、局限性和一些未来的方向。相关工作。附录 A 讨论了选择机制如何与类似概念相关。附录 B 包含 SSM 和其他相关模型的扩展相关工作。没有免费的午餐：连续-离散频谱。结构化 SSM 最初被定义为连续系统的离散化 (1)，并且对连续时间数据模式（例如感知信号（例如音频、视频））有强烈的归纳偏差。正如第 3.1 节和第 3.5 节所讨论的，选择机制克服了它们在文本和 DNA 等离散模式上的弱点；但这反过来会妨碍他们的表现 LTI SSM 擅长的数据。我们对音频波形的消融更详细地检验了这种权衡。下游A法令。基于 Transformer 的基础模型（特别是 LLM）具有丰富的属性生态系统以及与预训练模型交互的模式，例如微调、适应、提示、上下文学习、指令调整、RLHF、量化等在。我们特别感兴趣的是 Transformer 替代品（例如 SSM）是否具有类似的属性和规定。缩放。我们的实证评估仅限于小型模型，低于大多数强大的开源 LLM（例如 Llama（Touvron 等人，2023））以及其他循环模型，例如 RWKV（B. Peng 等人，2023）和 RetNet （Y. Sun 等人，2023），已在 7B 参数尺度及以上进行了评估。仍有待评估曼巴在这些较大尺寸的情况下是否仍然具有优势。我们还注意到，扩展 SSM 可能会涉及进一步的工程挑战和对模型的调整，这些在本文中没有讨论。六，结论我们向结构化状态空间模型引入了一种选择机制，允许它们执行上下文相关的推理，同时线性缩放序列长度。当合并到简单的免注意力架构中时，Mamba 在多个领域取得了最先进的结果，其性能可匹配或超过强大的Transformer 模型的性能。我们对选择性状态空间模型在为不同领域构建基础模型方面的广泛应用感到兴奋，特别是在需要长上下文的新兴模式中，例如基因组学、音频和视频。我们的结果表明，Mamba 是成为通用序列模型骨干的有力候选者。