Mamba学习笔记(1)——原理基础

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Author：Albert Gu（Carnegie Mellon University）& Tri Dao（Princeton University） [github] [arXiv]

0 Abstract

基础模型（Foundation models），如今已成为深度学习中许多令人兴奋的应用的动力源，几乎普遍基于Transformer架构及其核心的注意力模块。为了应对Transformer在长序列上计算效率低的问题，已经开发了许多亚二次时间复杂度的架构，例如线性注意力、门控卷积和递归模型，以及结构化状态空间模型（SSM）。然而，在重要模态（如语言）上，这些架构的表现不如注意力机制。我们发现，这些模型的关键弱点在于它们无法进行基于内容的推理，并为此做出了几项改进。首先，让SSM参数作为输入的函数可以解决它们在离散模态中的弱点，使模型能够根据当前的输入选择性地沿序列长度传播或忘记信息。其次，尽管这一变化阻止了高效卷积的使用，我们设计了一个面向硬件的并行算法，以递归模式进行计算。我们将这些选择性SSM集成到一个无需注意力甚至MLP块的简化端到端神经网络架构（Mamba）中。Mamba享有比Transformer高5倍的推理吞吐量，并且在序列长度上呈线性扩展，其性能在真实数据上（最长达到百万长度的序列）得到改善。作为一种通用的序列模型骨干，Mamba在语言、音频和基因组等多个模态上实现了最先进的性能。在语言建模中，我们的Mamba-3B模型在预训练和下游评估中均超越了与之同等大小的Transformer，并且与其两倍大小的Transformer表现相当。
1、Transformer的局限性

现有的Transformer架构在长序列上计算效率低下，无法高效处理离散模态（如语言数据）

2、SSM改进

• 参数动态性：将SSM的参数设计为随输入变化，使得模型能够选择性地保留或忘记信息，从而提升基于内容的推理能力。
• 硬件优化：设计了一个并行的递归算法，以避免高效卷积的限制并提高模型的计算效率。

3、Mamba架构

• 无需注意力机制：通过选择性SSM设计出简化的神经网络架构，不依赖Transformer的注意力模块。
• 高效推理与训练：Mamba推理速度是Transformer的5倍，能够在线性时间内处理序列，最长可达百万长度。

4、多模态应用

Mamba在多个数据模态中达到了最先进的性能，特别是在语言建模中，Mamba-3B模型超越了同等大小的Transformer模型，并能与更大模型相匹敌。

1 Introduction

基础模型（Foundation models，FMs），即在大量数据上进行预训练然后适用于下游任务的大型模型，已经成为现代机器学习中的一种有效范式。这些FMs的核心通常是序列模型，它可以处理来自各种领域（如语言、图像、语音、音频、时间序列和基因组）的任意输入序列。虽然这个概念并不依赖于特定的模型架构，但现代FMs主要基于一种单一的序列模型：Transformer及其核心的自注意力层。自注意力机制的有效性归因于它能够在上下文窗口内密集路由信息，使其能够建模复杂数据。然而，这一特性也带来了根本性的缺陷：无法建模上下文窗口外的内容，并且随着窗口长度的增加，计算复杂度呈二次增长。

为了解决这些问题，已经出现了大量关于更高效注意力机制变体的研究，但这些变体往往以牺牲自注意力的有效性为代价。目前为止，这些变体尚未在跨领域的大规模应用中表现出良好的效果。

最近，结构化状态空间序列模型（SSMs）作为一种序列建模的新型架构引起了关注。这些模型可以被解释为递归神经网络（RNNs）和卷积神经网络（CNNs）的结合，并受到经典状态空间模型的启发。SSMs可以通过递归或卷积进行非常高效的计算，序列长度的扩展通常是线性或接近线性的。此外，它们在特定数据模态（如音频）中具有建模长距离依赖关系的机制，并在长序列数据集的基准测试中表现出色。然而，在建模离散和信息密集的数据（如文本）时，它们的表现较差。

我们提出了一种新的选择性状态空间模型类别，在多个方面改进了现有工作，以实现与Transformer相当的建模能力，同时保持线性扩展的计算效率。
Selection Mechanism 首先，我们发现了之前模型的一个关键限制：无法高效地基于输入选择数据（即关注或忽略特定输入）。基于在选择性复制和诱导头等合成任务中的直觉，我们设计了一种简单的选择机制，通过基于输入参数化SSM参数。这使模型能够过滤掉无关信息，并无限期地记住相关信息。
Hardware-aware Algorithm 这个简单的改变对模型计算提出了技术挑战；事实上，所有先前的SSM模型必须在时间和输入上保持不变，才能保证计算效率。我们通过硬件感知算法解决了这一问题，该算法通过扫描而不是卷积递归计算模型，并避免扩展状态物化，从而减少GPU内存层次间的IO访问。最终的实现无论在理论上（序列长度线性扩展）还是在现代硬件上（在A100 GPU上加速至3倍）都比之前的方法更快。
Architecture 我们通过将之前SSM架构的设计与Transformer的MLP模块结合成一个单一模块，简化了之前的深度序列模型架构，形成了一个简单且同质化的架构设计（Mamba），并引入了selective state spaces。
选择性SSM，以及其扩展架构Mamba，都是完全递归的模型，具有一些关键特性，使它们非常适合作为通用基础模型的主干，用于处理序列数据。(i) 高质量：选择性机制在语言和基因组等密集模态上表现出色。(ii) 快速训练和推理：在训练期间，计算和内存使用随着序列长度线性扩展，而在推理过程中，由于不需要缓存之前的元素，模型的自回归展开只需常数时间每步。(iii) 长上下文：高质量和高效率相结合，在真实数据上提供了序列长度达100万的性能改进。
我们通过不同模态和设置的实验证明了Mamba作为通用序列FM主干的潜力，无论是在预训练质量还是特定任务表现上。
Synthetics 在一些被认为对大型语言模型至关重要的合成任务中，例如复制任务和诱导头任务，Mamba不仅能够轻松解决这些问题，还可以无限延长解决方案（超过100万个tokens）
Audio and Genomics Mamba在建模音频波形和DNA序列时，表现优于之前的先进模型如SaShiMi、Hyena和Transformers，无论是在预训练质量还是下游指标上（例如在一个具有挑战性的语音生成数据集上，FID降低了一半以上）。在这两种场景中，随着上下文长度的增加，其性能在百万级长度的序列上得到了改善。
Language Modeling Mamba是第一个真正实现Transformer质量表现的线性时间序列模型，无论在预训练困难度还是下游评估中均表现出色。通过扩展到10亿参数的规模，我们展示了Mamba超越了大量基线模型的表现，包括基于LLaMa的非常强大的现代Transformer训练方案。我们的Mamba语言模型的生成速度比同等规模的Transformers快5倍，且Mamba-3B的质量与其两倍规模的Transformers相当（例如，在常识推理上比Pythia-3B平均高4分，甚至超过了Pythia-7B）。

1、基础模型（FMs）

基础模型是预训练的大型模型，广泛应用于各个领域的序列数据（如语言、图像、音频等）。当前，Transformer和自注意力层是FMs的主要架构。

2、Transformer的局限性

虽然自注意力机制有效，但随着上下文窗口长度的增加，计算复杂度呈二次增长。 Transformer无法处理超出上下文窗口的内容，限制了其应用。

3、结构化状态空间模型（SSMs）

SSMs结合了RNNs和CNNs的特性，计算效率高，尤其适用于处理长序列数据，且具有线性或近线性扩展的能力。
然而，SSMs在处理离散数据（如文本）时表现欠佳。

4、本文贡献

提出了一种选择性状态空间模型（Selective SSM），通过改进现有方法，在保持线性扩展的前提下，实现了与Transformer相当的建模能力。

2 State Space Models

结构化状态空间序列模型（S4）是一类用于深度学习的序列模型，受到连续系统的启发。该系统将一维函数或序列 x(t)（属于实数集R）通过隐式的潜在状态 h(t)（属于RN）映射到 y(t)（属于R）。S4 模型通过四个参数（Δ, 𝐀, 𝐁, 𝐂）定义，并分两个阶段实现序列到序列的转换。
具体来说，S4模型由四个参数（Δ, 𝐀, 𝐁, 𝐂）定义，这些参数在两阶段中定义了一个序列到序列的转换过程。

公式1a到3b描述了状态空间模型（SSM）如何通过递归和卷积操作对序列进行处理，涉及状态更新和输出生成的具体机制公式。公式 (1a) 和 (1b) 是状态更新和输出生成的连续时间形式，(2a) 和 (2b) 则是它们在离散时间下的对应形式。公式 (3a) 和 (3b) 则将状态更新的递归形式转换为了全局卷积的形式，可以提高并行化处理序列的效率。我们仔细看下这六个公式：

公式（1a） 表示状态更新。其中h(t)表示在时间t时刻的隐状态；A表示状态转移矩阵，控制状态h(t)是如何在不同时间中演变的；B表示输入矩阵，控制输入x(t)对状态的影响。h'(t)是更新后的状态，他结合了之前的状态h(t)和当前输入x(t)

公式（1b） 表示输出生成。其中h(t)表示在时间t时刻的隐状态；C表示输出矩阵，用于将隐状态h(t)转换为输出y(t)。

公式（2a） 表示状态更新的离散形式，即将连续时间系统转换为离散时间系统的版本。它表示在离散时间t的状态ht，是由前一时间t-1的状态ht-1和当前输入xt更新得到的。

公式（2b） 表示输出生成的离散形式，表示在离散时间t的输出yt，是通过当前时刻的隐状态ht和输出矩阵C生成的。

公式（3a） 定义了一个卷积核K，他描述了系统对输入序列的响应。这个卷积和由C、A和B组合形成，表示了系统如何随着时间k积累信息。

公式（3b） 表示卷积形式K的状态更新。表示输出y是输入x与卷积核K的卷积结果，示了通过全局卷积处理整个序列的能力。卷积操作可以看作是对输入序列进行整体性处理的方式，从而避免了逐步计算隐状态的递归方式。

Discretization 第一阶段通过A = fA(Δ, A)和B = fB(Δ, A, B)将“连续参数”（Δ, 𝐀, 𝐁）转化为“离散参数”（𝐀, 𝐁），其中，fA和fB是离散化规则的函数组合。可以使用各种规则，比如在公式（4）中定义的zero-order hold（ZOH）：

在状态空间模型中，连续时间系统通过离散化转化为离散时间系统，以便在数字计算中高效处理。其中，fA用于将连续的转台转移矩阵A离散化，fB适用于将输入矩阵B从连续时间转为离散时间。

离散化规则 零阶保持（Zero-Order Hold, ZOH）便是文中离散化规则之一。ZOH是一种常见的离散化方法，适用于保持系统在当前采样点之间的输入不变，从而简化计算。A=exp(ΔA)表示离散化后的状态转移矩阵A，他通过指数映射将连续时间矩阵A转换为离散时间矩阵，保持了系统的动态特征。B=(ΔA) ^−1(exp(ΔA)−I)⋅ΔB定义了离散化后的输入矩阵B，用于确保系统在离散时间下的输入响应与连续时间下的一致性。

离散化与连续时间系统有着深刻的联系，这使得模型可以获得额外的性质，例如分辨率不变性，并自动确保模型正确归一化。此外，它还与RNNs的门控机制相关联，我们将在第3.5节重新讨论这一点。然而，从机械角度来看，离散化可以简单地被视为SSM前向传递中计算图的第一步。不同类型的SSM可以绕过离散化步骤，直接参数化A和B，这可能更容易理解。
Computation （Δ, 𝐀, 𝐁, C）转化为“（𝐀, 𝐁, C）后，模型可以通过两种该方式进行计算：线性递归（linear recurrence ）或者全局卷积（global convolution）。
通常情况下，模型在训练时使用卷积模式以实现高效的并行处理（即可以提前看到整个输出序列），而在自回归推理时切换为递归模式（即每次只能看到一个时间的输入）以提高推理效率。

linear recurrence 对应（2a和2b），在递归模式下，模型会逐步处理输入序列中的每一个时间步。递归模式适合自回归推理（autoregressive inference），即每次仅看到一个时间步的输入，利用当前的输入和之前的状态ht-1来计算当前的状态ht和输出yt。这种模式效率高，适合生成任务，如逐步生成序列（文本或音频）

Global Convolution 对应（3a和3b），在卷积模式下，模型会通过整个输入序列进行卷积计算输出。卷积核K由参数𝐀,B,C组成，模型可以一次性看到整个序列的所有输入。这种模式非常适合并行化训练（parallelizable training），因为所有时间的输入都可以同时处理，不需要逐步递归。并行化卷积训练能够显著提高大规模数据集上的训练效率。

两种模式的应用场景
卷积模式, 通常用于高效的并行训练，因为模型在训练时可以同时处理整个输入序列，从而加速训练过程。
递归模式, 用于高效的自回归推理，因为在推理过程中，模型需要根据先前的状态和当前输入逐步计算输出。

Linear Time Invariance (LTI) 方程 (1) 到 (3) 的一个重要特性是模型的动态在时间上保持不变。换句话说，（Δ, 𝐀, 𝐁, C），以及由此得到的（𝐀, 𝐁）在所有时间步长上都是固定的。这个特性被称为线性时不变性 (LTI)，它与递归和卷积有着深刻的联系。我们将 LTI SSMs 视为等同于任何线性递归 (2a) 或卷积 (3b)，并使用 LTI 作为这些模型类别的统称。
到目前为止，所有的结构化 SSMs 都是 LTI 的（例如通过卷积计算），这是由于基本的效率限制，具体讨论见第3.3节。然而，本研究的核心观点是，LTI 模型在建模某些类型的数据时存在基本的局限性，而我们的技术贡献在于去除 LTI 约束，同时克服效率瓶颈。

LTI 是状态空间模型的一个重要特性，它意味着模型的动力学（即状态更新和输出生成的方式）在整个时间序列中保持不变。

LTI的关键特性

• 动力学恒定 模型的状态转移矩阵𝐀、输入矩阵B和输出矩阵C不会随时间的变化而变化，这就是所谓的LTI。换句话说，公式（1）-（3）在时间序列的每个时间步都是一样的。
• LTI与递归和卷积的关系 首先，LTI意味着递归公式和卷积公式在时间序列的每个时刻都采用相同的权重和转换方式；其次，在递归计算中状态通过相同的矩阵 𝐀和输入矩阵B从前一时刻传递到下一时刻；最后，在卷积计算中，输入序列与固定的卷积核 K进行卷积操作，因此整个过程保持不变。

LTI的局限性
虽然LTI模型具有高效的计算特性，尤其是在并行计算（如卷积）中，然而，它们在处理某些类型的数据时存在局限性，尤其是在需要动态调整模型参数来适应不同输入的情况下。LTI模型无法适应输入的变化，因为其参数在所有时间步都是固定的，这限制了其对复杂数据（例如离散和信息密集型数据，如文本）的建模能力。

Structure and Dimensions 结构化SSMs之所以命名为“Structure”，是因为为了高效计算，它们需要对矩阵𝐀施加一定的结构。最常见的结构形式是对角形式，我们也采用了这种结构。在这种情况下，矩阵𝑨 ∈ R𝑁×𝑁, 𝑩 ∈ R𝑁×1, 𝑪 ∈ R1×𝑁都可以用𝑁表示。为了对具有批量大小 𝐵、长度 𝐿 和 𝐷 个通道的输入序列 𝑥 进行操作，SSM 独立地应用于每个通道。需要注意的是，在这种情况下，每个输入的总隐藏状态维度为 𝐷𝑁，而在整个序列长度上进行计算需要 𝑂(𝐵𝐿𝐷𝑁) 的时间和内存；这正是第 3.3 节中解决的效率瓶颈的根源。

该部分讨论了 结构化状态空间模型（SSM） 在计算上的效率问题，以及矩阵A结构化的必要性。

Structure SSMs特点

• 对矩阵A施加一定的结构约束（对角矩阵结构），其中A是一个方阵，但在结构化模型中将其限制为对角阵，B是一个列向量，C是一个行向量。
• 计算瓶颈 复杂度𝑂(𝐵𝐿𝐷𝑁)表明，随着batch_size(B)，序列长度L，通道数D和维度N的增加，计算和内存需求也会快速增加。这种复杂度是SSM在高效计算时面临的根本瓶颈，尤其当序列很长时，内存消耗和计算负担变得非常大。

General State Space Models 我们指出，“State Space Models”这一术语具有非常广泛的含义，简单来说，它表示任何具有潜在状态的递归过程的概念。它在不同领域中被用于描述许多不同的概念，包括马尔可夫决策过程（MDP）（强化学习）、动态因果模型（DCM）（计算神经科学）、卡尔曼滤波器（控制理论）、隐马尔可夫模型（HMM）和线性动态系统（LDS）（机器学习）、以及广义的递归模型（有时包括卷积模型）（深度学习）。
在本文中，我们使用“SSM”一词专指结构化SSM或S4模型的类别，并交替使用这些术语。为了方便起见，我们还包括此类模型的派生模型，例如那些侧重于线性递归或全局卷积观点的模型，并在必要时澄清细微差别。

State Space Models 指的是任何具有潜在状态的递归过程，它是一个通用概念，在多个领域有不同的形式和应用。

马尔可夫决策过程（MDP） 强化学习中的一个重要概念，用来建模具有潜在状态和随机性决策的系统。

动态因果模型（DCM） 在计算神经科学中，用来建模大脑活动和因果关系的动态过程。

卡尔曼滤波器 控制系统中的经典方法，用来估计动态系统中的隐状态。

隐马尔可夫模型（HMM） 和 线性动态系统（LDS） 在机器学习中，用来处理具有潜在状态和噪声的时间序列数据。

递归模型和卷积模型 在深度学习中，递归神经网络（RNNs）和卷积神经网络（CNNs）也可以看作是特定形式的状态空间模型，因为它们通过潜在状态或局部特征提取对序列或图像数据进行处理。

SSM Architectures SSM是独立的序列转换模块，可以集成到端到端的神经网络架构中。（我们有时也称SSM架构为SSNNs，SSM层与神经网络的关系类似于线性卷积层和CNN的关系。）以下是一些最著名的SSM架构，许多将作为我们的主要基准模型。
Linear attention 这是自注意力的近似模型，涉及一种递归机制，可视为简化的线性SSM。
H3 此架构使用S4进行递归的一般化，可以视为SSM层前后夹带两个门控连接的模型（见图3）。H3还在主要SSM层之前插入了一个标准的局部卷积层，称为“shift-SSM”。
Hyena Hyena架构与H3相似，但将S4层替换为MLP参数化的全局卷积层。
RetNet 此架构增加了一个额外的门控机制，并使用更简单的SSM，允许替代的并行计算路径，使用多头注意力（MHA）的变体而非卷积。
RWKV 这是一个用于语言建模的RNN，基于线性注意力的另一种近似机制，即无需注意力的Transformer。其核心机制“WKV”涉及线性时间不变（LTI）的递归，可视为两个SSM的比率。
其他相关的SSM架构在扩展的相关工作中进一步讨论（见附录B）。特别值得一提的是S5、QRNN和SRU，它们与我们的核心选择性SSM方法最为接近。

SSM架构是可以替代传统注意力机制和卷积操作的强大工具，在不同的变体下展示了极大的灵活性和计算效率。

3 Selective State Space Models

我们通过合成任务的直观理解引出选择机制（第3.1节），接着解释如何将该机制融入状态空间模型（第3.2节）。由此产生的时变SSM不能使用卷积，带来了如何高效计算它们的技术挑战。我们通过一个硬件感知的算法解决了这一问题，该算法利用了现代硬件的内存层次结构（第3.3节）。然后，我们介绍了一个不使用注意力甚至MLP模块的简单SSM架构（第3.4节）。最后，我们讨论了选择机制的一些附加特性（第3.5节）。

3.1 Motivation: Selection as a Means of Compression

我们认为序列建模的一个根本问题是如何将上下文压缩到一个较小的状态中。事实上，可以从这个角度来看各种流行的序列模型之间的权衡。例如，注意力机制既有效又低效，因为它根本没有压缩上下文。这可以从自回归推理需要明确存储整个上下文（即KV缓存）这一事实看出，这直接导致了Transformer在推理时线性时间和训练时平方时间的瓶颈。另一方面，循环模型因为其状态是有限的，具备常量时间推理和线性时间训练的效率。然而，它们的有效性受限于这个状态能在多大程度上压缩上下文。
为了理解这一原理，我们聚焦于两个常用的合成任务示例（图2）：

• 选择性复制任务（Selective Copying）对流行的复制任务（Copying task）进行了修改，通过改变需要记忆的标记的位置。这项任务要求具备基于内容的推理能力，能够记忆相关的标记（有色）并过滤掉无关的标记（白色）。
• 引导任务（Induction Heads）是一种被认为能解释大规模语言模型（LLMs）大部分上下文学习能力的机制。该任务要求具备基于上下文的推理能力，以便在适当的上下文中产生正确的输出（黑色）。
  这些任务揭示了线性时不变（LTI）模型的失败模式。从循环模型的视角看，其恒定的动态（例如方程(2)中的𝑨、𝑩转换）不能根据输入选择正确的上下文信息，或者影响沿序列传播的隐藏状态。从卷积模型的视角来看，已知全局卷积能够解决基础的复制任务，因为它只需要时间感知，但在选择性复制任务中表现不佳，因为缺乏内容感知。更具体地说，输入到输出的间距是变化的，无法用静态卷积核建模。
  总结来说，序列模型的效率与有效性权衡在于它们如何压缩状态：高效的模型必须有较小的状态，而有效的模型必须能够在状态中包含所有必要的上下文信息。因此，我们提出构建序列模型的一个基本原则是选择性：即基于上下文的能力，专注于或过滤输入到序列状态中的信息。特别是，选择机制控制信息如何沿着序列维度传播或相互作用（参见第3.5节）。

注意力机制的有效性和低效率
自注意力（如Transformer中的注意力机制）虽然在建模复杂依赖关系时非常有效，但它不对上下文进行压缩。推理时，Transformer需要显式存储整个上下文（通过KV缓存），这导致了推理时的线性时间复杂度和训练时的二次时间复杂度。

递归模型的效率和局限性
递归模型（如RNN）因为具有有限状态，所以推理时间为常数，训练时间为线性，具有较高的效率。然而，递归模型的有效性取决于它们如何压缩上下文信息到有限的状态中。如果压缩不充分，模型的效果会大打折扣。

示例任务解释压缩问题

• 选择性复制任务（Selective Copying Task）：这是经典复制任务的一种变体，要求模型记住某些标记位置的令牌，同时忽略其他无关的令牌。这个任务需要模型具备“基于内容的推理能力”来选择性记忆相关信息，而忽略无关信息。
• 归纳头任务（Induction Heads Task）：该任务与语言模型的“上下文学习”能力相关，模型需要知道何时在合适的上下文中输出正确的内容。这也是对模型基于上下文的推理能力的测试。

LTI的限制

• 从递归模型的角度来看，LTI（线性时不变）模型的动态是固定的，因此无法根据输入选择正确的信息。它们的状态更新过程无法根据输入数据进行调整，从而无法有效过滤或选择有用信息。
• 从卷积模型的角度来看，全局卷积可以解决基本的复制任务，因为它只需要时间上的感知能力，但在处理选择性复制任务时会出现问题，因为卷积核是固定的，无法基于输入内容进行调整。

序列模型的效率与有效性权衡

• 效率：为了实现高效率，模型的状态必须足够小，这样可以减少计算量和内存消耗。
• 有效性：为了实现高效性，模型的状态必须包含所有必要的上下文信息。如果上下文压缩不足，模型的性能会受到影响。

Selection的核心原则
作者提出，构建序列模型的一个核心原则是选择性。选择性意味着模型能够基于上下文进行推理，选择性地聚焦或过滤输入信息，并将有用的信息压缩到一个较小的状态中。选择机制控制了信息如何在序列维度上传播或交互。

3.2 Improving SSMs with Selection

将选择机制融入模型的一种方法是让影响序列交互的参数（例如RNN的递归动态或CNN的卷积核）依赖于输入。算法1和算法2展示了我们使用的主要选择机制。主要区别在于使多个参数 Δ，B，C成为输入的函数，同时对张量形状进行相关的变化。特别强调的是，这些参数现在具有长度维度L.这意味着模型已从时间不变变为时间变化。（注意，形状注释在第2节中进行了描述。）这使得与卷积的等价性（3）失效，从而影响其效率，后面将对此进行讨论。我们特别选择：𝑠𝐵(𝑥) = Linear𝑁(𝑥), 𝑠𝐶(𝑥) = Linear𝑁(𝑥), 𝑠Δ(𝑥) = Broadcast𝐷(Linear1(𝑥)), 𝜏Δ = softplus。其中，Linear𝑑是投影到维度d的参数化过程。选择𝑠Δ和𝜏Δ的原因与第3.5节中解释的RNN门控机制相关。
这种输入依赖的选择机制允许模型根据输入的不同动态调整其行为，从而提高了对不同输入上下文的适应性。尽管这可能会影响模型的效率，但通过精确控制信息的流动，它也为更复杂的序列建模任务提供了更多灵活性。

作者提出了一种通过选择机制将模型的参数与输入数据动态关联的方法。这种方法的核心思想是使与序列交互相关的参数依赖于输入数据，从而实现基于内容的动态调整。

• SB(x)=LinearN(x)通过输入x进行线性投影，将输入映射到维度N
• SC(x)=LinearN(x)与B类似，C也通过线性投影映射到维度N
• 𝑠Δ(𝑥) = Broadcast𝐷(Linear1(𝑥)) 中Δ的选择采用先通过线性投影映射到1维，然后广播到维度 D
• 𝜏Δ = softplus 中𝜏Δ是一个激活函数，选择softplus 作为激活的原因与RNN中的门控机制有关，相关的细节在第3.5节中解释。

3.3 Efficient Implementation of Selective SSMs

卷积和注意力机制等硬件友好型操作因其高效性而被广泛应用。在此，我们的目标是使选择性SSM在现代硬件（如GPU）上同样具备高效性。选择机制本质上是自然而然的，早期的研究已经尝试融入某些特殊选择机制，例如让Δ在递归SSM中随时间变化。然而，如前所述，SSM应用中的核心限制是其计算效率，这也是S4及其所有衍生模型使用LTI（非选择性）模型的原因，最常见的形式是全局卷积。

3.3.1 Motivation of Prior Models

我们首先回顾这一动机，并概述我们克服以往方法限制的方式：

• 从总体上看，像SSM这样的递归模型始终在表达能力与速度之间权衡：如第3.1节所讨论，隐藏状态维度越大，模型的表现应越好，但速度会变慢。因此，我们希望在不增加计算和内存成本的情况下，最大化隐藏状态的维度。
• 递归模式比卷积模式更灵活，因为后者是从前者扩展出来的。然而，这需要计算并存储形状为(B, L, D, N)的潜在状态ℎ，这比输入𝑥和输出𝑦的形状(B, L, D)大了一个因子𝑁（即SSM的状态维度）。为此，引入了更高效的卷积模式，可以绕过状态计算，并仅生成形状为(B, L, D)的卷积核。
• 先前的LTI（线性时不变）状态空间模型利用递归和卷积的双重形式，将有效状态维度扩大了𝑁倍（通常约为10到100），远远超过了传统RNN，但没有影响效率。

递归模型的权衡
递归模型（如SSM）总是在表达能力和速度之间寻找平衡。模型的隐藏状态维度越大，其表达能力（即处理复杂序列的能力）就越强，但速度和内存消耗也会增加。因此，关键目标是在不增加速度和内存成本的情况下，最大化隐藏状态维度。

递归模式 vs 卷积模式
递归模式（公式2）比卷积模式（公式3）更加灵活，因为卷积模式是通过展开递归模式得到的。因此，从理论上讲，递归模式可以处理更多的复杂性。然而，使用递归模式时，必须计算和存储隐藏状态h,其形状为(B,L,D,N)，这里的N是SSM的状态维度。这个形状比输入x和输出y的形状(B,L,D)要大很多。因此，处理递归模式时需要更多的内存和计算。

卷积模式的效率优势
为了解决上述问题，提出了更高效的卷积模式，通过跳过状态 ℎ的计算，直接计算卷积核（公式3a），其形状为 (B,L,D)，与输入和输出的大小相同。这显著降低了内存和计算成本。

LTI状态空间模型的创新
之前的线性时不变（LTI）状态空间模型，通过结合递归和卷积的形式，能够有效增加状态维度 N（约为10到100），远远超过传统的RNN模型。通过这种方式，模型在不牺牲计算效率的情况下，获得了更高的表达能力。

3.3.2 Overview of Selective Scan: Hardware-Aware State Expansion

Selection Mechanism旨在克服LTI模型的局限性，同时我们需要重新审视SSM的计算问题。我们使用了三种经典技术来解决这一问题：内核融合、并行扫描和重计算。主要有以下两点观察：
朴素的递归计算使用 O(BLDN) 的FLOPs（浮点运算次数），而卷积计算使用 O(BLDlog(L)) 的FLOPs，且递归计算的常数因子更低。因此，对于长序列和不太大的状态维度 N，递归模式实际上可以使用更少的FLOPs。
两个主要挑战是递归的顺序性和大量的内存使用。为了解决后者，我们可以效仿卷积模式，不实际存储完整的状态 h。主要思想是利用现代加速器（如GPU）的特性，只在更高效的内存层次中存储状态 h。大多数操作（除了矩阵乘法）受到内存带宽的限制。通过内核融合减少内存IOs，大大加速了计算过程。
具体而言，我们直接从慢速HBM（高带宽内存）加载SSM参数到更快的SRAM，进行离散化和递归操作，然后将最终输出写回HBM。为了避免递归的顺序性，我们采用了工作效率高的并行扫描算法来实现并行化。为了减少内存需求，在反向传播时我们通过重计算来避免保存中间状态，从而将内存需求降低到与FlashAttention的优化Transformer实现相当的水平。
完整的选择SSM层和算法如图1所示，详细的内核融合和重计算技术见附录D。

递归计算 vs 卷积计算的 FLOPs

• 递归计算的浮点运算（FLOPs）复杂度为 O(BLDN)，即批量大小 B、序列长度 L、通道数 D 和状态维度 N 的乘积。
• 卷积计算的浮点运算复杂度为 O(BLDlogL)，比递归计算更高效，尤其在长序列情况下。然而，递归计算的常数因子较低，因此对于较长序列和适中的状态维度 N，递归模式的 FLOPs 实际上可能更少。

递归计算的挑战

• 序列性：递归计算的一个主要挑战是它的顺序依赖性，不能像卷积那样并行处理。
• 内存使用：递归计算需要存储大量的中间状态 h，导致内存需求巨大。

内核融合（Kernel Fusion）

• 为了减少内存输入/输出（IO）的操作，作者使用了内核融合技术，特别是在现代加速器（如GPU）的高带宽内存（HBM）和快速SRAM之间进行数据传输。
• 具体做法是将SSM参数 (Δ,A,B,C) 从慢速的HBM传输到快速SRAM，在SRAM中进行离散化和递归计算，然后再将最终的输出 (B,L,D) 写回HBM。这减少了中间状态 h 的存储开销。

并行扫描（Parallel Scan）

• 尽管递归计算具有序列依赖性，作者指出，这种递归尽管不是线性的，但仍然可以通过并行扫描算法来实现并行化处理。通过该方法，模型可以在较短时间内完成递归计算。
• 作者引用了并行扫描算法的经典研究（如Blelloch, Martin和Smith等人的工作）来证明该技术的有效性。

重新计算（Recomputation）

• 为了减少内存需求，作者使用了重新计算技术。通过在反向传播时不存储中间状态，而是重新计算这些状态，节省了内存开销。
• 具体来说，反向传播时，当输入数据从HBM加载到SRAM时，重新计算中间状态 h，避免了在前向传播过程中保存这些状态。这样，内存需求降低到了与经过优化的Transformer（使用FlashAttention）实现相同的水平。

3.4 A Simplified SSM Architecture

选择性SSM与结构化SSM一样，可以作为独立的序列变换灵活地集成到神经网络中。H3架构是最著名的SSM架构的基础，它通常由一个受到线性注意力启发的模块和MLP（多层感知机）模块交错组成。为了简化这一架构，我们将这两个组件合并为一个同质堆叠的模块，灵感来源于门控注意力单元（GAU），类似于GAU对注意力机制的简化。
在这个架构中，模型的维度D被可控的扩展因子E扩展。对于每个模块，大部分参数3ED^2位于线性投影中（输入投影为2ED^2，输出投影为ED^2），而SSM的贡献相对较小（投影矩阵如Δ、B、C、A的参数较少）。该模块通过标准的归一化和残差连接重复堆叠，形成Mamba架构。
在实验中，我们固定扩展因子 (E = 2)，并使用两个模块堆栈，使其参数量12D^2与Transformer的多头注意力和MLP模块匹配。使用了SiLU/Swish激活函数，目的是使门控MLP变为流行的“SwiGLU”变体。为了进一步优化，我们还选择了LayerNorm进行归一化，与RetNet在类似位置使用归一化层的动机一致。
这种选择性SSM的设计通过简化架构、结合不同的现代优化方法，同时保持了与Transformer相当的性能和参数效率。

3.5 Properties of Selection Mechanisms

选择机制是一个更广泛的概念，可以以不同的方式应用，比如在传统的RNN或CNN中，对不同的参数（例如算法2中的矩阵𝑨）进行调整，或者使用不同的变换𝑠(𝑥) 来处理输入数据。

3.5.1 Connection to Gating Mechanisms

我们强调了最重要的联系：经典的RNN门控机制是我们在SSM中选择机制的一个实例。我们注意到RNN门控与连续时间系统离散化之间的联系已被广泛认可。事实上，定理1是对Gu, Johnson, Goel等人（2021年，第3.1条引理）的改进，推广到了ZOH离散化和依赖输入的门控（证明见附录C）。更广泛地看，SSM中的Δ可以视为RNN门控机制的广义形式。与之前的工作一致，我们认为SSM的离散化是启发式门控机制的理论基础。
定理1：当𝑁=1，𝑨=-1，𝑩=1，𝑠Δ=Linear(𝑥)，且𝜏Δ=softplus时，选择性SSM的递归形式（算法2）为

正如在第3.2节中提到的，我们选择𝑠Δ和𝜏Δ的原因来自于这一联系。特别需要注意的是，如果某个输入𝑥𝑡应当被完全忽略（如在某些合成任务中需要的那样），所有𝐷个通道都应该忽略它，因此我们将输入投影到1个维度，然后通过Δ进行重复/广播。

RNN门控与SSM选择机制的联系

• RNN中的门控机制，如LSTM和GRU中的输入门、遗忘门等，控制了信息如何在递归模型中传播。本文指出，RNN的门控机制实际上是SSM选择机制的一种实例。
• 作者引用了Funahashi和Nakamura（1993）以及Tallec和Ollivier（2018）的工作，这些研究已经明确了RNN门控与连续时间系统离散化之间的联系。

定理1的改进

• 定理1 是对Gu, Johnson, Goel等（2021年）提出的引理（Lemma 3.1）的改进，本文将其推广到ZOH（零阶保持）离散化和输入相关的门控机制。具体证明在附录C中。
• 定理1的核心思想是将SSM中的选择机制和RNN中的门控机制视为类似的过程，特别是SSM中的参数Δ扮演了RNN门控机制的广义角色。

Δ 的广义角色

• 在SSM模型中，参数 Δ 通过选择机制控制信息如何在序列中传播。它的作用类似于RNN中的门控机制，如控制信息流动和记忆信息的更新。
• 作者将SSM的离散化视为RNN启发式门控机制的一个原则性基础，这一观点与现有文献一致。

3.5.2 Interpretation of Selection Mechanisms

我们进一步详细说明选择机制的三个具体效应：

1. 可变间距(Variable Spacing)：选择性允许滤除在关注的输入之间出现的无关噪声标记。选择性复制任务是这一效应的例子，而在常见的数据模式中，尤其是离散数据，这一效应普遍存在，例如语言填充词如“呃”的出现。这一特性产生的原因是模型可以机制性地滤除任何特定的输入𝑥𝑡，例如在带门控的RNN案例中（定理1）当𝑔𝑡趋向于0时。
2. 上下文过滤(Filtering Context)：经验表明，许多序列模型在较长的上下文下并没有显著改善，尽管更多的上下文应该理论上提高性能。原因在于许多序列模型无法有效地在必要时忽略无关的上下文；一个直观的例子是全局卷积（和一般的LTI模型）。而选择性模型则可以随时重置其状态，从而删除多余的历史信息，因此它们的性能理论上随着上下文长度的增加而单调提高（如第4.3.2节所述）。
3. 边界重置(Boundary Resetting)：在多条独立的序列被拼接在一起的情况下，Transformers可以通过指定一个特定的注意力掩码将这些序列分开，而LTI模型则会在这些序列之间传播信息。选择性SSMs同样可以在边界处重置其状态（例如当Δ𝑡 → ∞，或定理1中𝑔𝑡 → 1时）。这些场景可能是人为创建的（例如将文档打包在一起以提高硬件利用率），也可能是自然发生的（例如强化学习中的情节边界）。

此外，我们还进一步讨论了每个选择性参数的效应：

• Δ的解释：通常，Δ控制了模型关注或忽略当前输入𝑥𝑡的程度。它推广了RNN的门控机制（例如定理1中的𝑔𝑡）：从机制上来说，大的Δ会重置状态ℎ并关注当前输入𝑥，而小的Δ则保持状态并忽略当前输入。在SSM中，Δ可以解释为时间步的离散化参数，大的Δ → ∞表示系统长时间关注当前输入，而小的Δ → 0则表示忽略瞬时输入。
• 𝑨的解释：我们指出，虽然𝑨参数也可以是选择性的，但它最终通过𝑨 = exp(Δ𝑨)（离散化公式(4)）与Δ相互作用来影响模型。因此，Δ的选择性足以确保(𝑨, 𝑩)的选择性，Δ是主要的性能提升来源。我们假设将𝑨也设为选择性的（或者替代Δ）会产生类似的性能，但为简化处理暂不讨论这一点。
• 𝑩和𝑪的解释：正如第3.1节中所讨论的，选择性最重要的特性是滤除无关信息，使得序列模型的上下文可以压缩为一个高效的状态。在SSM中，将𝑩和𝑪设为选择性可以更精细地控制是否让输入𝑥𝑡进入状态ℎ𝑡，或者让状态进入输出𝑦𝑡。这可以解释为允许模型基于内容（输入）和上下文（隐藏状态）分别调节其循环动态。

3.6 Additional Model Details

• 实数 vs. 复数：大多数先前的SSM（状态空间模型）使用复数作为其状态ℎ，这在感知模式任务中对于实现强大性能是必要的。然而，经验观察表明，在某些场景下，完全实数的SSM似乎也能很好地工作，甚至可能表现更好。我们默认使用实数，这在除了一个任务之外的所有任务中都表现良好。我们假设复数与实数之间的权衡与数据模式的连续-离散光谱相关，复数对于连续数据模式（如音频、视频）更有帮助，而对离散数据（如文本、DNA）则帮助不大。

实数和复数的选择取决于数据的模态，复数更适合连续模态，实数更适合离散模态。

• 初始化：大多数先前的SSM也建议使用特殊的初始化，尤其是在复数情况下，这可以在低数据场景中提供帮助。我们的复数情况下默认初始化为S4D-Lin，实数情况下为S4D-Real，这基于HIPPO理论。这些初始化方式分别将𝑨的第n个元素定义为−1/2 + 𝑛𝑖和−(𝑛+1)。然而，在大量数据和实数SSM情况下，我们预计多种初始化方法都能正常工作，相关消融实验在第4.6节中讨论。

初始化策略依赖于S4D模型，实数和复数的初始化各有其优势，特别是在不同的数据环境下。

• Δ的参数化：我们将对Δ的选择性调整定义为𝑠Δ(𝑥) = Broadcast𝐷(Linear1(𝑥))，这源于Δ的机制（第3.5节）。我们观察到它可以从1维推广到更大的维度𝑅。我们设置𝑅为𝑑的一个小分数，其使用的参数数量相对于模块中的主线性投影可以忽略不计。此外，广播操作可以被视为另一个线性投影，初始化为特定的1和0模式；如果这个投影是可训练的，它会导致另一种形式𝑠Δ(𝑥) = Linear𝐷(Linear𝑅(𝑥))，这可以视为低秩投影。在我们的实验中，Δ参数（可视为一个偏置项）被初始化为𝜏−1 Δ(Uniform([0.001, 0.1]))，遵循SSM的先前研究。

Δ 参数化可以灵活调整，并通过低秩投影增强模型的选择性机制。

备注3.1：为了简化实验结果的表示，我们有时将选择性SSM简称为S6模型，因为它们是带有选择机制并通过扫描计算的S4模型。