文献阅读-Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Abstract

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

1. Introduction

基础模型(FMs)，或在海量数据上进行预训练，然后适应下游任务的大型模型，已经成为现代机器学习的一种有效范例。这些FMs的骨干通常是序列模型，操作来自各种领域的任意序列输入，如语言、图像、语音、音频、时间序列和基因组学(Brown et al . 2020;Dosovitskiy等2020;Ismail Fawaz等2019;Oord等人2016;Poli等人2023;Sutskever, Vinyals和Quoc V Le 2014)。虽然这一概念与模型架构的特定选择无关，但现代FMs主要基于单一类型的序列模型:Transformer (Vaswani等人2017年)及其核心关注层(Bahdanau, Cho和Bengio 2015年)。自注意力的有效性归因于其在上下文窗口内密集路由信息的能力，使其能够对复杂数据进行建模。然而，这种特性带来了根本性的缺点:无法在有限窗口之外对任何东西进行建模，并且相对于窗口长度进行二次缩放。为了克服这些缺点，大量的研究已经出现了更有效的注意力变体(Tay, Dehghani, Bahri等人，2022)，但往往是以牺牲使其有效的特性为代价的。到目前为止，这些变体都没有被证明在跨领域的规模上是有效。

最近，结构化状态空间序列模型(SSMs) (Gu, Goel, and r20122;Gu, Johnson, Goel等人(2021)已经成为序列建模的一个有前途的体系结构类别。这些模型可以被解释为循环神经网络(rnn)和卷积神经网络(cnn)的组合，灵感来自经典的状态空间模型(Kalman 1960)。这类模型可以非常有效地计算为递归或卷积，在序列长度上具有线性或近线性缩放。此外，他们在特定的数据模态下对于建模长程依赖有原则性准则(Gu, Dao, et al. 2020) ，并主导了 Long Range Arena等基准(Tay, Dehghani, Abnar等2021)。许多种类的ssm (Gu, Goel, and r 2022;顾，古普塔，等2022;Gupta, Gu, and Berant 2022;李阳等2023;马等人 2023;Orvieto等人2023;Smith, Warrington, and Linderman 2023)在涉及音频和视觉等连续信号数据的领域取得了成功(Goel et al 2022;Nguyen, Goel, et al . 2022;Saon, Gupta, and Cui, 2023)。然而，它们在对离散和信息密集的数据(如文本)进行建模时效果较差。

我们提出了一类新的选择性状态空间模型，它改进了先前在几个轴上的工作，以实现transformer在序列长度线性缩放的情况下的建模能力。

选择机制 首先，我们确定了先前模型的一个关键限制:以输入依赖的方式有效选择数据的能力(即关注或忽略特定的输入)。基于基于重要合成任务(如选择性复制和感应头)的直觉，我们通过基于输入参数化SSM参数来设计一个简单的选择机制。这使得模型可以过滤掉不相关的信息，并无限期地记住相关的信息。

硬件感知算法 这种简单的改变对模型的计算提出了技术挑战;事实上，为了提高计算效率，所有先前的ssm模型都必须是时间和输入不变的。我们使用硬件感知算法克服了这一点，该算法使用扫描而不是卷积来循环计算模型，但不实现扩展状态，以避免在GPU内存层次结构的不同级别之间进行IO访问。由此产生的实现在理论上(序列长度线性缩放，与所有基于卷积的ssm的伪线性相比)和现代硬件上(在A100 gpu上快3倍)都比以前的方法快。

结构 我们通过将先前的SSM架构(Dao, Fu, Saab, et al . 2023)的设计与transformer的MLP块合并为单个块来简化先前的深度序列模型架构，从而获得包含选择性状态空间的简单且同质的架构设计(Mamba)。

选择性ssm，以及扩展的Mamba架构，是完全循环的模型，具有关键属性，使它们适合作为在序列上操作的一般基础模型的骨干。(i)高质量:选择性使语言和基因组学等密集模式的表现强劲。(ii)快速训练和推理:在训练期间，计算和内存在序列长度上呈线性扩展，并且在推理期间自回归展开模型只需要每步恒定的时间，因为它不需要先前元素的缓存。(iii)长上下文:质量和效率共同提高了实际数据的性能，最长可达1M序列长度。

在几种类型的模式和设置下，我们通过经验验证了Mamba在预训练质量和特定领域任务性能方面作为一般序列FM骨干的潜力:

•人工合成物。对于重要的合成任务，如复制和感应头，已经被提议作为大型语言模型的关键，Mamba不仅可以轻松解决它们，而且可以无限长地推断解决方案(>1M标记)。

•音频和基因组学。在模拟音频波形和DNA序列方面，Mamba在预训练质量和下游指标(例如:DNA序列)方面优于SaShiMi、Hyena和Transformers等先前最先进的模型。

将具有挑战性的语音生成数据集上的FID减少一半以上)。在这两种情况下，它的性能随着更长的上下文而提高，最长可达百万长度的序列。

•语言建模。Mamba是第一个真正实现transformer质量性能的线性时间序列模型，无论是在预训练复杂度上还是在下游评价上。通过缩放定律高达1B个参数，我们表明Mamba超过了大范围基线的性能，包括基于LLaMa的非常强大的现代变形金刚训练食谱(Touvron等人2023)。我们的Mamba语言模型与同等大小的transformer相比具有5倍的生成吞吐量，Mamba- 3b的质量与其两倍大小的transformer相匹配(例如，在常识推理方面，与Pythia-3B相比平均高出4分，甚至超过了Pythia-7B)。

2. 状态空间模型

结构化状态空间序列模型(S4)是最近一类用于深度学习的序列模型，与rnn、cnn和经典状态空间模型广泛相关。它们的灵感来自于一个特定的连续系统(1)，它映射了一个一维函数或序列通过内在的隐藏状态。

具体来说，S4模型定义了四个参数(∆，A, B, C)，定义了两个阶段的序列到序列的转换。

离散化 第一阶段通过固定公式将“连续参数”(∆，A, B)转化为“离散参数”(A, B)，其中配对的被称为离散化准则，可以使用各种规则，例如式(4)中定义的零阶保持器(ZOH)。

离散化与连续时间系统有着深刻的联系，这可以赋予它们额外的属性，如分辨率不变性(Nguyen, Goel, et al . 2022)，并自动确保模型正确归一化(Gu, Johnson, Timalsina, et al . 2023;Orvieto等人，2023)。它还与rnn的门控机制有关(Gu, Gulcehre, et al . 2020;Tallec和Ollivier 2018)，我们将在3.5节中重新讨论。然而，从力学的角度来看，离散化可以简单地看作是SSM向前传递计算图的第一步。ssm的替代风格可以绕过离散化步骤，直接参数化(A, B) (Zhang et al . 2023)，这可能更容易推理。

计算 在参数从之后，模型可以用两种方式计算，要么是线性递归(2)，要么是全局卷积(3)。通常，该模型使用卷积模式(3)进行有效的并行训练(提前看到整个输入序列)，并切换到循环模式(2)进行有效的自回归推理(一次看到一个时间步)。

线性时间不变性 方程(1)到(3)的一个重要性质是模型的动力学随时间是恒定的。换句话说，(∆，A, B, C)以及对于所有时间步长都是固定的。这种性质被称为线性时不变性(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)，我们也使用。

在这种情况下，矩阵都可以用数字N表示。为了操作具有D通道的批量大小B和长度L的输入序列x，SSM独立地应用于每个通道。注意，在这种情况下，每个输入的总隐藏状态的维度为DN，并且在序列长度上计算它需要O(BLDN)时间和内存;这就是3.3节讨论的基本效率瓶颈的根源。

一般状态空间模型。我们注意到，状态空间模型一词具有非常广泛的含义，它简单地表示具有潜在状态的任何循环过程的概念。它已被用于指代不同学科中的许多不同概念，包括马尔可夫决策过程(MDP)(强化学习(Hafner et al 2020))，动态因果建模(DCM)(计算神经科学(Friston, Harrison, and Penny 2003))，卡尔曼滤波器(控制(Kalman 1960))，隐马尔可夫模型(HMM)和线性动力系统(LDS)(机器学习)，以及循环(有时是卷积)模型(深度学习)。

在整篇论文中，我们使用术语“SSM”专门指结构化SSM或S4模型(Gu, Goel, and r 2022;顾，古普塔，等2022;Gupta, Gu, and Berant 2022;Hasani et al . 2023;马等人2023;Smith, Warrington, and Linderman 2023)，并且可以互换使用这些术语。为了方便起见，我们还可以包括这些模型的导数，例如那些专注于线性递归或全局卷积观点的模型(Y. Li et al . 2023;Orvieto等人2023;Poli et al . 2023)，并在必要时澄清细微差别。

SSM架构。ssm是独立的序列转换，可以合并到端到端的神经网络架构中。(我们有时也称SSM架构为ssn, ssn之于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 et al . 2023)在架构中添加了一个额外的门，并使用更简单的SSM，允许另一种可并行计算路径，使用多头注意(MHA)的变体而不是卷积。
•RWKV (B. Peng et al . 2023)是一种基于另一种线性注意力近似(无注意力变压器(S. Zhai et al . 2021))的语言建模的最新RNN。其主要的“WKV”机制涉及LTI复发，可以看作是两个ssm的比率。

其他密切相关的SSM和架构将在扩展的相关工作中进一步讨论(附录B)。我们特别强调了S5 (Smith, Warrington, and Linderman 2023)， QRNN (Bradbury et al 2016)和SRU (Lei et al 2017)，我们认为它们是与我们的核心选择性SSM最密切相关的方法。

3. Selective SSM

我们使用来自合成任务的直觉来激发我们的选择机制(第3.1节)，然后解释如何将该机制合并到状态空间模型中(第3.2节)。由此产生的时变ssm不能使用卷积，这对如何有效地计算它们提出了技术挑战。我们使用硬件感知算法来克服这个问题，该算法利用现代硬件上的内存层次结构(第3.3节)。然后，我们描述了一个简单的SSM架构，没有注意甚至MLP块(第3.4节)。最后，我们将讨论选择机制的一些附加属性(第3.5节)。

3.1 Motivation: Selection as a Means of Compression

我们认为序列建模的一个基本问题是将上下文压缩成更小的状态。事实上，我们可以从这个角度来看待流行序列模型的权衡。例如，注意力既有效又低效，因为它根本不压缩上下文。这可以从自回归推理需要显式存储整个上下文(即KV缓存)这一事实中看出，这直接导致了transformer的线性时间推理和二次时间训练缓慢。另一方面，循环模型是有效的，因为它们有一个有限的状态，意味着常数时间推理和线性时间训练。然而，它们的有效性受到这种状态对上下文的压缩程度的限制。

为了理解这一原则，我们关注两个正在运行的合成任务示例(图2)。

•选择性复制任务修改了流行的复制任务(Arjovsky, Shah, and Bengio 2016)，通过改变要记住的令牌的位置。它需要内容感知推理，以便能够记住相关的标记(彩色)并过滤掉不相关的标记(白色)。

•诱导头任务是一个众所周知的机制，假设可以解释LLMs的大部分上下文学习能力(Olsson et al 2022)。它需要上下文感知推理来知道何时在适当的上下文中产生正确的输出(黑色)。

这些任务揭示了LTI模型的失效模式。从循环的角度来看，它们的恒定动态(例如(2)中的转换)不能让它们从它们的上下文中选择正确的信息，或者影响沿序列传递的隐藏状态，并且不依赖于输入。从卷积的角度来看，我们知道全局卷积可以解决普通复制任务(Romero et al 2021)，因为它只需要时间感知，但由于缺乏内容感知，它们在选择性复制任务中存在困难(图2)。更具体地说，输入到输出之间的间隔是变化的，不能用静态卷积核来建模。

总之，序列模型的效率与有效性的权衡是由它们如何很好地压缩它们的状态来表征的:高效的模型必须有一个小的状态，而有效的模型必须有一个包含来自上下文的所有必要信息的状态。反过来，我们提出构建序列模型的基本原则是选择性:或上下文感知能力，以关注或过滤掉输入到序列状态。特别是，选择机制控制信息如何沿着序列维度传播或交互(更多讨论请参见3.5节)。

3.2 Improving SSMs with Selection

将选择机制纳入模型的一种方法是让影响序列交互的参数(例如RNN的循环动力学或CNN的卷积核)依赖于输入。

算法1和算法2说明了我们使用的主要选择机制。主要的区别是简单地将几个参数∆，B, C作为输入函数，以及整个张量形状的相关变化。特别地，我们强调这些参数现在有一个长度维度L，这意味着模型已经从时不变变为时变。(请注意，形状注释在第2节中进行了描述)。这失去了与卷积(3)的等价性，并影响了其效率，下面将讨论。

我们特别地选择，其中是对于维度d的参数化映射，选择是由于与第3.5节中解释的RNN门通机制的连接。

3.3 Efficient Implementation of Selective SSMs

卷积(Krizhevsky, Sutskever, and Hinton 2012)和Transformers(Vaswani et al 2017)等硬件友好架构得到了广泛的应用。在这里，我们的目标是使选择性ssm在现代硬件(GPU)上高效。选择机制是非常自然的，早期的研究试图纳入特殊的选择情况，例如在周期性ssm中让∆随时间变化(Gu, Dao, et al 2020)。然而，正如前面提到的，ssm使用的一个核心限制是它们的计算效率，这就是为什么S4和所有导数使用LTI(非选择性)模型，最常见的形式是全局卷积。

3.3.1 Motivation of Prior Models

我们首先重新审视这一动机，并概述我们克服先前方法局限性的方法。

•在高层次上，像ssm这样的循环模型总是在表达性和速度之间取得平衡:正如3.1节所讨论的，具有更大隐藏状态维度的模型应该更有效，但速度更慢。因此我们希望在不支付速度和内存成本的情况下最大化隐藏状态维度。

•请注意，循环模式比卷积模式更灵活，因为后者(3)是通过扩展前者(2)推导出来的(Gu, Goel, and r 2022;Gu, Johnson, Goel等人2021)。然而,这将要求的计算和落实隐藏状态h用(B,L,D,N)形状，远远大于输入和输出(B,L,D)形状。因此，引入了一种更高效的卷积模式，它可以绕过状态计算，实现一个只有(B,L,D)的卷积核(3a)。

•先前的LTI ssm利用双重递归卷积形式将有效状态维度增加N(≈10−100)，比传统rnn大得多，而且没有效率损失。

3.3.2 Overview of Selective Scan: Hardware-Aware State Expansion

该选择机制旨在克服LTI模型的局限性;同时，我们需要重新审视ssm的计算问题。我们用三种经典技术来解决这个问题:核融合、并行扫描和重计算。我们做了两个主要观察:
•朴素循环计算使用O(BLDN)FLOPs，而卷积计算使用O(BLDlog(L))FLOPs，前者具有较低的常数因子。因此，对于长序列和不太大的状态维푁，循环模式实际上可以使用更少的FLOPs。

•两个挑战是递归的顺序性，以及大量的内存使用。为了解决后者，就像卷积模式一样，我们可以尝试不实际实现完整状态。

其主要思想是利用现代加速器(gpu)的特性，只在内存层次结构中更有效的级别实现状态。特别是，大多数操作(除了矩阵乘法)都受到内存带宽的限制(Dao, Fu, Ermon, et al 2022;Ivanov等2021;Williams, Waterman, and Patterson, 2009)。这包括我们的扫描操作，我们使用内核融合来减少内存io的数量，与标准实现相比，这导致了显著的加速。

具体来说，我们不是在GPU HBM(高带宽内存)中准备size(B，L，D，N)的扫描输入，而是直接将SSM参数(∆，A, B, C)从慢速HBM加载到快速SRAM中，在SRAM中进行离散化和递归，然后将size(B，L，D)的最终输出写回HBM。

为了避免顺序递归，我们观察到，尽管不是线性的，它仍然可以用工作效率高的并行扫描算法并行化(Blelloch 1990;Martin and Cundy 2018;史密斯，沃灵顿和林德曼2023)。

最后，我们还必须避免保存中间状态，这是反向传播所必需的。我们仔细地应用了经典的重新计算技术来减少内存需求:当输入从HBM加载到SRAM时，中间状态不被存储，而是在反向传递中重新计算。因此，融合选择扫描层具有与FlashAttention优化变压器实现相同的内存要求。

融合核和重新计算的细节见附录d。完整的选择性SSM层和算法如图1所示。

3.4 A Simplified SSM Architecture

与结构化ssm一样，选择性ssm是独立的序列变换，可以灵活地纳入神经网络。H3架构是最著名的SSM架构的基础(第2节)，它通常由一个受线性注意力启发的块与一个MLP(多层感知器)块交织组成。我们通过将这两个组件组合成一个组件来简化该架构，该组件均匀堆叠(图3)。这是受到门控注意力单元(GAU) (Hua et al 2022)的启发，GAU对注意力做了类似的事情。

该体系结构包括通过一个可控的扩展因子E扩展模型维度D。对于每个块，大多数参数(3ED^{2)都在线性投影中(2ED}2为输入投影，ED^2为输出投影)，而内部SSM的贡献较小。SSM参数的数量(∆、B、C和矩阵A)相比之下要小得多。我们重复这个块，与标准的规范化和剩余的连接交织，形成曼巴建筑。在我们的实验中，我们总是固定为퐸= 2，并使用两个块堆栈来匹配Transformer的交错MHA(多头注意)和MLP块的12个퐷2参数。我们使用了SiLU / Swish激活函数(Hendrycks and Gimpel 2016;Ramachandran, Zoph, and Quoc V Le 2017)，促使门控MLP成为流行的“SwiGLU”变体(Chowdhery等人2023;Shazeer 2020;Touvron等人2023)。最后，我们还使用了一个可选的规范化层(我们选择LayerNorm (J. L. Ba, Kiros, and Hinton 2016))，其动机是RetNet在类似位置使用了规范化层(Y. Sun et al . 2023)。

3.5 Properties of Selection Mechanisms

选择机制是一个更广泛的概念，可以以不同的方式应用，例如更传统的rnn或cnn，不同的参数(例如算法2中的a)，或使用不同的转换s(x)。

3.5.1 Connection to Gating Mechanisms

我们强调了最重要的联系:rnn的经典门控机制是ssm选择机制的一个实例。我们注意到RNN门控与连续时间系统离散化之间的联系是很好的建立的(Funahashi和Nakamura 1993;Tallec and Ollivier 2018)。事实上，定理1是对Gu, Johnson, Goel等人(2021，引理3.1)的改进，将其推广到ZOH离散化和输入依赖门(证明见附录C)。更广泛地说，ssm中的∆可以看作是RNN门通机制的广义作用。与先前的工作一致，我们认为ssm的离散化是启发式门控机制的原则基础。

如3.2节所述，我们对S∆的具体选择，τ∆即来自于此连接。特别要注意的是，如果一个给定的输入xt应该被完全忽略(在合成任务中是必要的)，那么所有的D通道都应该忽略它，所以我们在用∆重复/广播之前将输入投影到一维。

3.5.2 Interpretation of Selection Mechanisms

我们详细阐述了选择的两种特殊的机制效应。
可变间距。选择性允许过滤掉在感兴趣的输入之间可能出现的不相关的噪声标记。选择性复制任务就是一个例子，但是在常见的数据模式中无处不在，特别是对于离散数据——例如语言填充符(如“um”)的存在。这个属性的出现是因为模型可以机械地过滤掉任何特定的输入xt，例如在门通RNN的情况下(定理1)gt→0。

过滤背景。根据经验观察，许多序列模型在更长的上下文下并没有得到改善(F. Shi et al . 2023)，尽管有更多上下文应该导致严格更好的性能的原则。一种解释是，许多序列模型在必要时不能有效地忽略无关的上下文;一个直观的例子是全局卷积(和一般LTI模型)。另一方面，选择性模型可以在任何时候简单地重置它们的状态以删除无关的历史，因此它们的性能原则上随着上下文长度单调地提高(例如第4.3.2节)。

边界重置 在多个独立序列拼接在一起的情况下，Transformers可以通过实例化一个特定的注意掩码来保持它们的分离，而LTI模型将在序列之间传递信息。选择性ssm也可以在边界处重置它们的状态(例如，gt→1时∆t→∞或定理1)。这些设置可以人为地发生(例如，将文档打包在一起以提高硬件利用率)或自然地发生(例如，强化学习中的情节边界(Lu et al 2023))。

此外，我们详细说明了每个选择参数的影响。

∆的解释。一般来说，∆控制着力于或忽略当前输入xt之间的平衡。它概括了RNN门(例如定理1中的gt)，机械地，一个大的∆重置状态并关注当前输入x，而一个小的∆保持状态并忽略当前输入。ssm(1)-(2)可以被解释为一个被时间步长∆离散的连续系统，在这种情况下，直观的感觉是大∆→∞代表系统更长时间地关注当前输入(因此“选择”它并忘记它的当前状态)，而小∆→0代表被忽略的瞬态输入。

A的解释。我们注意到，虽然A参数也可以是选择性的，但它最终仅通过(离散化(4))与∆的相互作用来影响模型。因此，∆的选择性足以保证(A, B)的选择性，并且是改进的主要来源。我们假设将A选择性地添加(或代替)∆也会有类似的效果，为了简单起见，我们将其省略。

B和C的解释。正如3.1节所讨论的，选择性最重要的特性是过滤掉不相关的信息，以便将序列模型的上下文压缩成有效的状态。在SSM中，将B和C修改为可选的，可以更细粒度地控制是让输入푥푡进入状态푡还是让状态进入输出푦푡。这些可以解释为允许模型分别基于内容(输入)和上下文(隐藏状态)调制循环动态。

3.6 Additional Model Details

Real vs. Complex 大多数先前的ssm在其状态中使用复数，这对于在许多任务上的强大性能是必要的(Gu, Goel, and r 2022)。然而，根据经验观察，在某些情况下，完全实值ssm似乎工作得很好，甚至可能更好(Ma et al 2023)。我们使用真实值作为默认值，这对于我们的所有任务都很有效;我们假设复实权衡与数据模式中的连续-离散频谱有关，其中复数对连续模式(例如音频，视频)有帮助，但对离散模式(例如文本，DNA)没有帮助。

Initialization 大多数先前的ssm还建议进行特殊的初始化，特别是在复杂值的情况下，这可以在一些设置(如低数据制度)中有所帮助。我们对复杂情况的默认初始化是S4D-Lin，对真实情况的默认初始化是S4D-Real (Gu, Gupta, et al . 2022)，这是基于HIPPO理论的(Gu, Dao, et al . 2020)。它们分别定义A的푛-th元素为−1∕2 +ni和−(n+ 1)。然而，我们期望许多初始化工作良好，特别是在大数据和实值SSM制度;第4.6节讨论了一些消融。

参数化∆。我们将∆的选择性调整定义为s∆(x)=BroadcastD(Linear1(x))，其动机是∆的机制(第3.5节)。我们观察到它可以从1维推广到更大的维度R。我们将其设置为D的一小部分，与块中的主要线性投影相比，它使用的参数数量可以忽略不计。我们还注意到，广播操作可以被视为另一个线性投影，初始化为1和0的特定模式;如果这个投影是可训练的，这将导致另一个s∆(x)=LinearD(LinearR(x))，它可以被视为一个低秩投影。

在我们的实验中，∆参数(可视为偏差项)初始化为，遵循先前对ssm的研究(Gu, Johnson, Timalsina, et al . 2023)。

评论3.1 为了实验结果的简洁，我们有时将选择性ssm缩写为S6模型，因为它们是具有选择机制的S4模型，并通过扫描计算。

4. 实验评估

在4.1节我们测试曼巴的能力，以解决在3.1节动机的两个合成任务。然后我们在三个领域进行评估，每个领域都在自回归预训练和下游任务上进行评估。

•第4.2节:语言模型预训练(缩放定律)和零射击下游评估。

•第4.3节:DNA序列预训练和长序列分类任务的微调。

•第4.4节:音频波形预训练，以及自回归生成的语音片段的质量。

最后，第4.5节展示了Mamba在训练和推理时间的计算效率，第4.6节详细介绍了体系结构和选择性ssm的各个组成部分。

4.1 合成任务

这些任务的完整实验细节，包括任务细节和训练方案见附录E.1。

4.1.1 选择拷贝

拷贝任务是序列建模中研究得最多的综合任务之一，最初设计用于测试循环模型的记忆能力。正如3.1节所讨论的，LTI ssm(线性递归和全局卷积)可以通过只跟踪时间而不是对数据进行推理来轻松解决这个任务;例如，通过构建一个长度正好合适的卷积核(图2)。这在早期关于全局卷积的工作中得到了明确的验证(Romero et al 2021)。选择性复制任务通过随机化令牌之间的间距来防止这种快捷方式。请注意，该任务之前已经作为去噪任务引入(Jing et al . 2019)。

请注意，许多先前的工作认为，添加架构门控制(乘法交互)可以赋予模型“数据依赖性”并解决相关任务(Dao, Fu, Saab, et al . 2023;Poli等人，2023)。然而，我们发现这种解释在直观上是不够的，因为这种门控不会沿着序列轴相互作用，也不会影响令牌之间的间距。特别是体系结构门控并不是选择机制的一个实例(附录a)。

表1证实，H3和Mamba等门控制架构只能部分提高性能，而选择机制(将S4修改为S6)很容易解决这个问题，特别是在与这些更强大的架构结合使用时。

4.1.2 Induction Heads

诱导头(Olsson et al . 2022)是来自机制可解释性视角(Elhage et al . 2021)的一个简单任务，它令人惊讶地预测了LLMs的上下文学习能力。它要求模型执行联想回忆和复制:例如，如果模型在序列中看到了“哈利波特”这样的双字母组合，那么下一次“哈利”出现在同一序列中时，模型应该能够通过从历史中复制来预测“波特”。

数据集。我们在序列长度为256的诱导头任务上训练了一个2层模型，词汇大小为16，这与之前在该任务上的工作相当(Dao, Fu, Saab, et al . 2023)，但序列更长。我们还通过在测试时评估从2^{6 = 64到2}20 = 1048576的序列长度范围来研究泛化和外推能力。

模型。在感应头的既定工作之后，我们使用了2层模型，这允许人们注意机械地解决感应头任务(Olsson et al 2022)。我们测试了多头注意(8个头，不同的位置编码)和SSM变体。我们对Mamba使用64的模型维度퐷，对其他模型使用128。

结果。表2显示了mamba(或者更准确地说，是它的选择性SSM层)能够完美地解决这个任务，因为它能够选择性地记住相关的令牌，同时忽略中间的所有其他内容。它可以完美地泛化到百万长度的序列，或者比训练时看到的长4000倍，而其他方法都不能超过2倍。

在注意模型的位置编码变体中，xPos(设计用于长度外推)略好于其他;还要注意的是，由于内存限制，所有的注意力模型只测试到序列长度2^14 = 16384。在其他SSMs中，H3和Hyena是相似的，这与Poli等人(2023)的发现相反。

4.2 Language Modeling

我们在标准自回归语言建模上对Mamba架构与其他架构进行了评估，包括预训练指标(困惑度)和零射击评估。我们设置模型大小(深度和宽度)以反映GPT3规范。我们使用Pile数据集(L. Gao, Biderman, et al . 2020)，并遵循Brown et al .(2020)中描述的训练配方。所有培训细节见附录E.2。

4.2.1 Scaling Laws

对于baseline，我们比较标准的Transformer体系结构(GPT3体系结构)，以及我们所知道的基于PaLM和LLaMa架构(例如旋转嵌入，SwiGLU MLP, RMSNorm代替LayerNorm，无线性偏差，以及更高的学习率)的最强Transformer配方(这里称为Transformer++)。我们还比较了其他最近的次二次结构(图4)。所有的模型细节都在附录E.2中。

图4显示了标准Chinchilla (Hoffmann et al 2022)协议下，在≈125M到≈1.3B参数的模型下的标度规律。Mamba是第一个与非常强大的Transformer配方(Transformer++)的性能相匹配的无需关注的模型，现在已经成为标准，特别是随着序列长度的增长。我们注意到，由于缺乏有效的实现导致内存不足或不切实际的计算需求，RWKV和RetNet基线缺乏上下文长度8k的完整结果，之前的强循环模型也可以解释为ssmM。

4.2.2 Downstream Evaluations

表3显示了Mamba在一系列流行的下游零射击评估任务上的表现。我们比较了这些规模上最著名的开源模型，最重要的是Pythia (Biderman等)
2023)和RWKV (B. Peng et al . 2023)，它们与我们的模型使用相同的标记器、数据集和训练长度(300B个标记)进行训练。(注意，曼巴和皮提亚是用上下文长度2048训练的，而RWKV是用上下文长度1024训练的。)

4.3 DNA Modeling

受大型语言模型成功的激励，最近有一些探索将基础模型范式用于基因组学。DNA被比作语言，因为它由具有有限词汇的离散符号序列组成。它还以需要远程依赖关系来建模而闻名(Avsec等人2021)。我们将曼巴作为调频主干进行预训练和微调，与最近对DNA长序列模型的研究相同(Nguyen, Poli, et al . 2023)。我们特别关注模型大小和序列长度之间缩放规律的两个探索(图5)，以及需要长上下文的困难下游合成分类任务(图6)。

对于预训练，我们主要遵循标准的因果语言建模(下一个令牌预测)设置来进行训练和模型细节(参见附录E.2)。对于数据集，我们在很大程度上遵循HyenaDNA (Nguyen, Poli, et al . 2023)的设置，它使用HG38数据集进行预训练，该数据集由单个人类基因组组成，在训练分割中具有约45亿个标记(DNA碱基对)。

4.3.1 Scaling: Model Size

在本实验中，我们研究了具有不同模型主干的基因组学基础模型的缩放特性(图5左)。
训练。为了利用基线，我们在1024的短序列长度上进行训练;如4.3.2节所示，我们预计在更长的序列长度上，结果更有利于曼巴。我们对于每个batch总共2^20=1M的token将全局批处理大小设置为1024，模型训练了10K个梯度步骤，总共有10B个令牌。

结果。图5(左)显示，Mamba的预训练困惑度随着模型大小的增加而平稳提高，并且Mamba的可伸缩性优于HyenaDNA和Transformer++。例如，在≈40M参数的最大模型尺寸下，曲线显示Mamba可以用大约3到4倍的参数匹配Transformer++和HyenaDNA模型。

4.3.2 Scaling: Context Length

在接下来的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)的序列长度warmup。

结果。图5(右)显示，曼巴能够利用更长的上下文，甚至长达1M的极长序列，其预训练困惑度随着上下文的增加而提高。另一方面，HyenaDNA模型随着序列长度的增加而变差。从第3.5节关于选择机制属性的讨论中可以直观地看出这一点。特别是，LTI模型不能选择性地忽略信息;从卷积的角度来看，一个很长的卷积核是在一个很长的序列中聚集所有的信息可能会包含很多噪声。请注意，虽然HyenaDNA声称可以在更长的上下文环境下改善，但他们的结果并没有控制计算时间。

4.3.3 Synthetic Species Classication

我们通过随机取样连续的DNA片段来评估5个不同物种的下游分类任务模型。这个任务改编自HyenaDNA，它使用了物种{人类，狐猴，老鼠，猪，河马}。我们通过将五种类人猿(人类，黑猩猩，大猩猩，猩猩，倭黑猩猩)进行分类来修改任务，使其更具挑战性，这五种类人猿已知有99%的DNA是相同的。

4.4 音频建模和生成

对于音频波形模式，我们主要与SaShiMi架构和训练协议进行比较(Goel et al 2022)。该模型包括1;一个U-Net骨干网，有两个阶段的池化，每个阶段的模型维度加倍D。在每个阶段交替使用S4和MLP块。

我们考虑用Mamba块代替S4+MLP块。实验细节见附录E.4。

4.4.1 长背景自回归预训练

我们在YouTubeMix (DeepSound 2017)上评估预训练质量(自回归下一个样本预测)，这是一个标准钢琴音乐数据集，由先前的工作组成，包括4小时的独奏钢琴音乐，采样率为16000 Hz预训练细节主要遵循标准语言建模设置(第4.2节)。图7评价了在保持计算量不变的情况下，将训练序列长度从2^{13 = 8192增加到2}20≈106的效果。

Mamba和SaShiMi (S4+MLP)基线随着上下文长度的增加而持续改善;曼巴在整个过程中表现更好，距离越长，差距就越大。主要度量是比特每字节(BPB)，它是标准负对数似然(NLL)损失的常数因子log(2)，用于预训练其他模式。

我们注意到一个重要的细节:这是本文中唯一一个从真实参数化切换到复杂参数化的实验(第3.6节)。我们在附录E.4中显示了其他消融。

4.4.2自回归语音生成

SC09是一个基准语音生成数据集(Donahue, McAuley, and Puckette 2019;Warden 2018)，由以16000 Hz采样的数字“0”到“9”的1秒片段组成，具有高度可变的特征。我们在很大程度上遵循Goel等人(2022)的自回归训练设置和生成协议。

表4显示了与Goel等人的各种基线相比，Mamba-UNet模型的自动度量

(2022): WaveNet (Oord等人2016)，SampleRNN (Mehri等人2017)，WaveGAN (Donahue, McAuley, and Puckette 2019)， DiffWave (Z. Kong等人2021)和SaShiMi。一个小曼巴模型优于国家的最先进的(和更大的)氮化镓和扩散为基础的模型。更大的模型参数——与基线匹配——进一步显著提高了保真度指标。

表5以小型Mamba模型为例，研究了外部阶段和中心阶段不同架构的组合。结果表明，Mamba在外围区块的表现优于S4+MLP，在中心区块的表现优于MHA+MLP。

4.5 速度和内存测试

我们对SSM扫描操作的速度(状态扩展푁= 16)以及Mamba的端到端推理吞吐量进行了基准测试，见图8。我们高效的SSM扫描比我们所知道的最好的注意力实现(FlashAttention-2 (Dao 2023))在序列长度超过2K的情况下要快，比PyTorch中的标准扫描实现快20-40倍。Mamba实现了4-5倍比类似大小的变压器更高的推理吞吐量，因为没有KV缓存，它可以使用更高的批处理大小。例如，Mamba-6.9B(未经训练)将比5倍小的Transformer-1.3B具有更高的推理吞吐量。详细信息见附录E.5，其中还包括内存消耗的基准测试。

4.6 模型消融

我们对模型的组件执行了一系列详细的精简，重点关注在Chinchilla令牌计数下大小≈350M模型的语言建模设置(与图4设置相同)。

4.6.1 结构

表6研究了架构(块)及其内部SSM层的影响(图3)。我们发现•在以前的非选择性(LTI) SSM中，其相当于全局卷积，性能非常相似。

•用实值ssm替换之前工作中的复值S4变体不会对性能产生太大影响，这表明(至少对于LM)在考虑硬件效率时，实值ssm可能是更好的选择。

•用选择性SSM (S6)替换其中任何一个都能显著提高性能，验证了第3节的动机。

•Mamba架构的性能与H3架构相似(在使用选择性层时似乎略好)。

我们还在附录E.2.2中研究了将Mamba块与其他块(如MLP(传统架构)MHA(混合注意力架构))交织在一起的问题。

4.6.2 Selective SSM

表7通过考虑选择性∆、B、C参数的不同组合(算法2)，对选择性SSM层进行了消融，可以看出，由于∆与RNN门通的关系，它是最重要的参数(定理1)。

表8考虑了SSM的不同初始化，这些初始化已被证明会在某些数据模式和设置中产生很大差异(Gu, Goel, and r 2022;Gu, Gupta, et al . 2022)。在语言建模方面，我们发现更简单的实值对角初始化(S4D-Real，第3行)比更标准的复值参数化(S4D-Lin，第1行)表现更好。随机初始化也很有效，这与之前的研究结果一致(Mehta et al . 2023)。

表9和表10分别考虑改变∆和(B, C)投影的尺寸。将它们从静态更改为选择性提供了最大的好处，而进一步增加维度通常会略微提高性能，并增加少量参数计数。

特别值得注意的是，当状态大小푁增加时，选择性SSM的显著改进，仅以1%的额外参数为代价，困惑度提高了1.0以上。这证实了我们在3.1节和3.3节中的核心动机。

5. Discussion

我们讨论了相关的工作，局限性和一些未来的方向。

相关的工作。附录A讨论了选择机制与类似概念的关系。附录B扩展了ssm和其他相关模型的相关工作。

**没有免费的午餐:连续-离散频谱。**结构化ssm最初被定义为连续系统的离散化(1)，并且对连续时间数据模式(如感知信号(如音频，视频))具有强烈的归纳偏见。正如3.1节和3.5节所讨论的，选择机制克服了它们在离散模式(如文本和DNA)上的弱点;但是这反过来又会阻碍他们的表现在LTI ssm擅长的数据上。我们对音频波形的消融更详细地考察了这种权衡。

下游可供性 基于transformer的基础模型(特别是llm)具有丰富的属性生态系统，以及与预训练模型交互的模式，例如微调、适应、提示、上下文学习、指令调整、RLHF、量化等等。我们特别感兴趣的是，诸如ssm之类的Transformer替代品是否具有类似的属性和功能。

缩放 我们的经验评估仅限于较小的模型尺寸，低于大多数强大的开源llm(例如Llama (Touvron等人，2023))以及其他循环模型(如RWKV (B. Peng等人，2023)和RetNet (Y. Sun等人，2023)的阈值，这些模型已在7B参数尺度及以上进行了评估。曼巴在这些更大的体型中是否仍然具有优势还有待评估。我们还注意到，扩展ssm可能涉及进一步的工程挑战和对模型的调整，这些在本文中没有讨论。

6. Conclusion

我们在结构化状态空间模型中引入了一种选择机制，允许它们在序列长度线性缩放的同时执行上下文相关的推理。当合并到一个简单的无需关注的架构中时，Mamba在不同的领域集合上实现了最先进的结果，在这些领域中，它匹配或超过了强大的Transformer模型的性能。我们对选择性状态空间模型在不同领域建立基础模型的广泛应用感到兴奋，特别是在需要长上下文的新兴模式中，如基因组学、音频和视频。我们的结果表明，Mamba作为一般序列模型骨干是一个强有力的候选人。