Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Mamba：基于选择状态空间的线性时间序列建模
论文两位作者Albert Gu和Tri Dao，博士都毕业于斯坦福大学，导师为Christopher Ré。
Albert Gu现在是CMU助理教授，多年来一直推动SSM架构发展。他曾在DeepMind 工作，目前是Cartesia AI的联合创始人及首席科学家。
Tri Dao，以FlashAttention、FlashDecoding系列工作闻名，现在是普林斯顿助理教授，和Together AI首席科学家，也在Cartesia AI担任顾问。
Code：https://github.com/state-spaces/mamba
Background
基础模型（FM），即在海量数据上进行预训练，然后针对下游任务进行调整的大型模型，已成为现代机器学习的有效范式。这些模型几乎都基于 Transformer 架构及其核心注意力模块。
自注意力的效果主要归功于它能够在上下文窗口中密集地传递信息的能力，使其能够对复杂数据进行建模。然而，这一特性也带来了一些根本性的缺陷：
**无法对有限窗口外的任何内容进行建模：**自注意力机制通常在一个固定的上下文窗口（如一个句子或一个图像的一部分）内操作，这意味着它只能考虑并处理这个窗口内的信息。当需要理解或分析的信息分散在更大的上下文中，或者与当前窗口相隔较远时，自注意力可能无法有效捕捉这些信息，因为自注意力缺乏直接的机制来处理或引入窗口之外的信息。这在处理长篇文本或大规模图像时尤其成问题，因为重要信息可能分布在整个文档或图像中。**窗口长度的二次缩放问题：**自注意力的一个关键特点是它在计算时需要比较和加权窗口内的所有元素对。这意味着如果窗口长度增加，所需的比较（和因此的计算资源）会以二次方的方式增加。例如，如果窗口长度加倍，那么所需的计算量将增加四倍。这种二次缩放性质使得自注意力在处理大窗口或长序列时变得效率低下。
为了解决 Transformer 在长序列上的计算效率低下问题，人们开发了许多亚二次时间架构，如线性注意力、门控卷积和递归模型以及结构化状态空间模型（SSM）。此类模型的一个关键弱点是无法进行基于内容的推理。
其中SSMs最有潜力，它可以被解释为递归神经网络（RNN）和卷积神经网络（CNN）的结合，其灵感来自经典的状态空间模型（卡尔曼，1960）。但它们在离散和信息密集数据（如文本）的建模方面不太有效。
Contribution
改进
作者提出了一类新的选择性状态空间模型，该模型在多个维度上改进了先前的工作，从而在序列长度线性缩放的同时，实现了Transformer的建模能力。主要的改进有以下几点：
选择机制 先前的模型，特别是那些用于处理序列数据的模型，可能在有效选择数据方面存在局限，它们可能不够有效地关注重要的输入信息或忽略不相关的输入信息。选择性复制和归纳头等合成任务在理解和改进模型的选择机制方面提供了重要的直觉。于是作者设计了一种简单的选择机制，根据输入对 SSM 参数进行参数化（让 SSM 参数成为输入的函数）。这样，模型就能过滤掉无关信息，并无限期地记住相关信息。硬件感知算法 先前的SSM模型为了计算效率，必须是时间不变和输入不变的，这意味着它们的计算方式不随时间或输入数据的变化而改变。而上一点提到的改进，虽然提高了模型的灵活性和效果，但也带来了技术上的挑战。 于是作者开发了一种硬件感知算法。这种算法考虑了硬件的特性和限制，特别是在使用GPU进行计算时。算法使用递归的方式通过扫描来计算模型，而不是使用传统的卷积。 算法设计考虑到了GPU内存层次结构，为了避免不同级别之间的IO访问，它不会实体化扩展的状态，这一步骤有助于减少内存访问和相关的时间延迟。 在理论上，这种实现方式使得处理速度随序列长度线性增长，与基于卷积的所有SSM模型的伪线性增长相比有显著优势。在A100 GPU上，实现的性能比之前的方法快3倍。模型架构 作者将先前的 SSM 架构设计（Dao、Fu、Saab 等人，2023 ）与 Transformers 的 MLP 模块合并为一个模块，集成到一个简化的端到端神经网络架构中，该架构没有注意力，甚至没有 MLP 块，简化了先前的深度序列模型架构，形成了一种包含选择性状态空间的简单、同质的架构设计（Mamba）。
优势
选择性 SSM 以及 Mamba 架构是完全递归模型，其关键特性使其适合作为在序列上运行的通用基础模型的骨干，它具有以下几个特点：
**高质量：**选择性为语言和基因组学等密集模式带来了强大的性能。**快速训练和推理：**在训练过程中，计算量和内存与序列长度成线性关系，而在推理过程中，由于不需要缓存以前的元素，因此自回归展开模型每一步只需要恒定的时间。**上下文长：**质量和效率共同提高了实际数据的性能，序列长度可达 100 万。
应用范围
合成 在一些重要的合成任务上，如被认为是大型语言模型关键的复制和归纳头，Mamba 不仅能轻松解决，还能推导出无限长（>100 万个词组）的解决方案。 音频和基因组学 在音频波形和 DNA 序列建模方面，Mamba 在预训练质量和下游指标（例如，在具有挑战性的语音生成数据集上，FID 降低了一半以上）方面都优于 SaShiMi、Hyena 和 Transformers 等先前的一流模型。在这两种情况下，它的性能都随着上下文长度的增加而提高，最高可达百万长度的序列。 语言建模 Mamba 是第一个线性时间序列模型，无论是在预训练复杂度还是下游评估方面，都真正达到了 Transformer 质量的性能。通过多达 1B 个参数的缩放规律，作者发现 Mamba 的性能超过了大量基线模型，包括基于 LLaMa（Touvron 等人，2023 年）的非常强大的现代Transfomer训练方法。Mamba 语言模型与类似规模的 Transformer 相比，具有 5 倍的生成吞吐量，而且 Mamba-3B 的质量与两倍于其规模的 Transformer 相当（例如，与 Pythia-3B 相比，常识推理的平均值高出 4 分，甚至超过 Pythia-7B）。
State Space Models（状态空间模型）
“状态空间模型”一词广泛涵盖涉及潜在状态的任何循环过程，并已用于描述跨多个学科的各种概念。
可以先从物理中的状态空间模型来理解：
由常规物理规律可以研究系统的三个维度：系统输入、系统输出和状态量，上述模型中，u(t)为系统输入即拉力，y(t)为系统输出即位移量，该系统的状态可以有位移、速度、加速度等等更深层的潜在特征。
以弹簧-质量-阻尼系统 (SMD)为例，我们想要知道，在给定一个力
f
(
t
)
f(t)
f(t)作为系统输入量
u
(
t
)
u(t)
u(t)的时候， 以及质量块 M 随时间变化的位移
x
x
x ，速度
v
v
v 的方程。其中我们设置位移
x
x
x 作为系统的输出量
y
(
t
)
y(t)
y(t)。
将
x
x
x 和
v
v
v 作为向量
X
X
X 的两个基。这里的
X
X
X 就是系统的状态向量，而
x
x
x 和
v
v
v 就是这个系统两个状态分量，用向量的形式写出来就是：
X
=
[
x
1
x
2
]
=
[
x
v
]
X = \left[ \begin{matrix} x_1\\x_2 \end{matrix} \right] = \left[ \begin{matrix} x\\ v \end{matrix} \right]
X=[x1x2]=[xv]
X与y的关系可表示为：
X
=
[
x
1
(
t
)
x
2
(
t
)
]
=
[
y
(
t
)
d
y
(
t
)
d
t
]
X = \left[ \begin{matrix} x_1(t)\\x_2(t) \end{matrix} \right] = \left[ \begin{matrix} y(t)\\ \frac{dy(t)}{dt} \end{matrix} \right]
X=[x1(t)x2(t)]=[y(t)dtdy(t)]
根据牛顿第二定律
F
=
M
a
F=Ma
F=Ma，弹簧提供的力为
−
k
x
-kx
−kx，阻尼器提供的力为
−
b
v
-bv
−bv（负号表示与位移方向相反），外部力为
u
(
t
)
u(t)
u(t)。可以得到如下微分方程：
M
a
=
−
k
x
−
b
v
+
u
(
t
)
Ma = -kx-bv+u(t)
Ma=−kx−bv+u(t)
其中
a
=
d
2
x
d
t
2
a = \frac{d^2x}{dt^2}
a=dt2d2x为加速度，
x
x
x是位移，
v
=
d
x
d
t
v=\frac{dx}{dt}
v=dtdx是速度，于是可得到微分方程：
M
d
2
x
d
t
2
+
b
d
x
d
t
+
k
x
=
u
(
t
)
M\frac{d^2x}{dt^2} +b\frac{dx}{dt}+ kx = u(t)
Mdt2d2x+bdtdx+kx=u(t)
可表示为：
M
x
˙
2
+
b
x
2
+
k
x
1
=
u
(
t
)
M\dot{x}_2+bx_2+kx_1=u(t)
Mx˙2+bx2+kx1=u(t)
求解该方程，可得：
x
˙
1
=
x
2
\dot{x}_1 = x_2
x˙1=x2
x
˙
2
=
u
(
t
)
−
b
x
2
−
k
x
1
M
\dot{x}_2 = \frac{u(t)-bx_2-kx_1}{M}
x˙2=Mu(t)−bx2−kx1
即：
X
˙
=
[
0
1
−
k
M
−
b
M
]
[
x
1
x
2
]
+
[
0
1
M
]
u
(
t
)
\dot{X} = \left[ \begin{matrix} 0&1\\-\frac{k}{M} & -\frac{b}{M}\end{matrix} \right] \left[ \begin{matrix} x_1\\x_2\end{matrix} \right] + \left[ \begin{matrix} 0\\ \frac{1}{M} \end{matrix} \right]u(t)
X˙=[0−Mk1−Mb][x1x2]+[0M1]u(t)
而系统的输出量
y
(
t
)
y(t)
y(t)也可以写成关于当前状态的方程：
y
=
[
1
0
]
[
x
1
x
2
]
+
[
0
]
u
(
t
)
y=\left[ \begin{matrix} 1&0 \end{matrix} \right]\left[ \begin{matrix} x_1\\x_2 \end{matrix} \right] + [0]u(t)
y=[10][x1x2]+[0]u(t)
与深度学习领域 SSM 模型对应：
通过该种方法抽象而来的深度网络，在序列预测任务中，可以解决传统的 Transformer 模型对离散深层潜在特征不敏感的问题，以对时序性数据建立更强的逻辑。
下面开始正式介绍SSM：
本文中的SSM指的是结构化状态空间序列模型（Structured state space sequence models，S4）是最近一类用于深度学习的序列模型，与 RNN、CNN 和经典状态空间模型有广泛联系。它们受到一个特殊连续系统的启发，该系统通过隐含的潜在状态
h
(
t
)
∈
R
N
h(t)\in \mathbb{R}^N
h(t)∈RN，映射一个一维函数或序列
x
(
t
)
∈
R
↦
y
(
t
)
∈
R
x(t)\in\mathbb{R}\mapsto y(t)\in \mathbb{R}
x(t)∈R↦y(t)∈R。
S4模型可由四个参数
(
Δ
,
A
,
B
,
C
)
(\Delta,A,B,C)
(Δ,A,B,C) 定义，
Δ
\Delta
Δ为选择性因子，形状与输入相同，用于控制每个时间步对应的状态更新程度，A 表示状态转移矩阵，B 和 C 分别表示输入到状态和状态到输出的映射矩阵。他们分两个阶段定义序列到序列的转换。
离散化
第一阶段通过公式
A
ˉ
=
f
A
(
Δ
.
A
)
\bar{A}=f_A(\Delta.A)
Aˉ=fA(Δ.A) 和
B
ˉ
=
f
B
(
Δ
,
A
,
B
)
\bar{B}=f_B(\Delta,A,B)
Bˉ=fB(Δ,A,B) 将“连续参数”
(
Δ
,
A
,
B
)
(\Delta,A,B)
(Δ,A,B) 转换到“离散参数”
(
A
ˉ
,
B
ˉ
)
(\bar{A},\bar{B})
(Aˉ,Bˉ)。其中
(
f
A
,
f
B
)
(f_A,f_B)
(fA,fB)被称为离散化规则，可以使用各种规则，如零阶保持 ( zero-order hold，ZOH )：
离散化可以赋予连续时间系统更多的特性，如分辨率不变性和自动确保模型正确归一化。它还与 RNN 的门控机制有关。不过，离散化可以简单地看作是 SSM 前向传递中计算图的第一步。另一种 SSM 可以绕过离散化步骤，直接参数化 (A,B)（Zhang 等，2023 年），这可能更容易推理。
计算
在参数从
(
∆
,
A
,
B
,
C
)
↦
(
A
,
B
,
C
)
(∆,A,B,C) ↦ (A,B,C)
(∆,A,B,C)↦(A,B,C) 转化之后，模型可以通过两种方式计算，一种是线性递归 (式2) （使用RNN），另一种是全局卷积 (3) （使用CNN）。
通常在训练时使用卷积模式 (式3) （提前看到整个输入序列），推理时使用递归模式 (式2) （每次看到一个时间步的输入）。
线性时不变性（Linear Time Invariance，LTI）
S4 的状态模型参数 ($ \Delta, A, B, C $ ) 在所有时间步中都是固定不变的，这一特性被称为线性时不变性。LTI 是递归和卷积的基石，为构建序列模型提供了一个简化但功能强大的框架。
迄今为止，所有结构化 SSM 都是 LTI 模型，因为存在基本的效率限制。然而，本文工作的一个核心观点是，LTI 模型在对某些类型的数据进行建模时具有根本性的局限性，本文的技术贡献在于消除 LTI 限制，同时克服效率瓶颈。
结构和维度
结构化 SSM 之所以被命名为结构化 SSM，是因为高效计算 SSM 还需要对 A 矩阵施加结构（结构化即给矩阵 A 设定一个特定的格式或模式，使得它更易于处理和优化），最常用的结构形式是对角线结构，本文也使用这种结构。
在此情况下，
A
∈
R
N
×
N
,
B
∈
R
N
×
1
,
C
∈
R
1
×
N
A\in\mathbb{R}^{N\times N},B\in\mathbb{R}^{N\times 1},C\in\mathbb{R}^{1\times N}
A∈RN×N,B∈RN×1,C∈R1×N都可以用N个数字来表示。若输入大小为
x
∈
R
B
×
L
×
D
x\in \mathbb{R}^{B \times L \times D}
x∈RB×L×D，其中B为batch size，L为输入长度，D为输入通道数，将SSM独立应用于每个通道，每个输入的总隐藏状态维数为 DN，在序列长度上它需要O(BLDN)的时间和内存，这是基本效率瓶颈的来源。
SSM 架构
SSM 是独立的序列变换，可被纳入端到端神经网络架构。(我们有时也称 SSM 架构为 SSNN，它与 SSM 层的关系就像 CNN 与线性卷积层的关系）。下面将介绍一些最著名的 SSM 架构，其中许多也将作为主要 baseline。
Linear attention: 线性注意力，它是自注意力的近似，涉及一个递归，可以看作是一个退化的线性 SSM。 H3: 它在S4的基础上进行了扩展；可以被看作是一种由两个门控连接夹着一个 SSM 的架构（如下图）。H3还在主SSM层之前插入了一个标准的局部卷积，将这部分定义为一个shift-SSM。
Hyena: 使用与 H3 相同的架构，但用 MLP 参数化全局卷积取代了 S4 层。 RetNet: 在架构中增加了一个额外的门，并使用更简单的 SSM，允许另一种可并行计算的路径，使用多头注意力（MHA）的变体来代替卷积。 RWKV: 是最近基于另一种线性注意近似（attention-free Transformer）设计的用于语言建模的 RNN。它的主要 "WKV "机制涉及 LTI 递归，可视为两个 SSM 的比值。 其他的方法还有S5、QRNN、SRU等。
选择性状态空间模型（S6）
分为以下几部分介绍：
利用合成任务的直觉来引发选择机制。解释如何将这一机制纳入状态空间模型。由此产生的时变 SSM 不能使用卷积，给如何高效计算卷积带来了技术挑战。于是采用了一种硬件感知算法，利用现代硬件的内存层次结构来克服这一难题。描述了一个简单的 SSM 架构，它不需要注意力，甚至不需要 MLP 块。讨论选择机制的一些其他特性。
动机：选择是一种压缩手段
序列建模的一个基本问题是将上下文压缩到更小的状态中。注意力机制虽然有效，但由于不压缩上下文信息，因此在推断和训练过程中效率低下（Transformer就像人类每写一个字之前，都把前面的所有字+输入都复习一遍，所以写的慢）。而递归模型由于具有有限状态，推断和训练效率较高，但其有效性取决于它们的有效性受限于状态对上下文压缩的质量（RNN就像每次只参考前面固定的字数，写的快，但容易忘掉更前面的内容）。
作者提出了两个合成任务作为示例：
选择性复制任务：改编自流行的复制任务，通过变化需要记忆的标记的位置，要求模型能够识别相关标记并过滤掉不相关的标记。归纳头任务：用于解释大型语言模型（LLMs）的上下文学习能力，要求模型在适当的上下文中产生正确的输出。
（左）复制任务的标准版本涉及输入和输出元素之间的恒定间距，可通过线性递归和全局卷积等时变模型轻松解决。(右上）选择性复制任务的输入间距是随机的，需要时变模型能够根据输入内容选择性地记住或忽略输入。(右下）"归纳头 "任务是联想记忆的一个例子，要求根据上下文检索答案，这是 LLM 的一项关键能力。
这些任务揭示了线性时间不变（LTI）模型的局限性。从递归模型的角度来看，它们恒定的动态（例如
(
A
ˉ
,
B
ˉ
)
(\bar{A},\bar{B})
(Aˉ,Bˉ)）不能使它们从上下文中选择正确的信息，或以输入依赖的方式影响沿序列传递的隐藏状态。从卷积模型的角度来看，虽然全局卷积可以解决标准复制任务，但在处理需要内容意识的选择性复制任务时则存在困难，因为输入到输出之间的间距是变化的，无法用静态卷积核建模。
序列模型的效率与有效性之间的权衡由它们压缩状态的能力决定：高效的模型必须有一个小的状态，而有效的模型必须包含所有必要的上下文信息。作者提出，构建序列模型的一个基本原则是选择性，或者说具有上下文意识的能力，专注于或过滤掉输入到序列状态的信息。特别是，选择机制控制信息如何在序列维度上传播或相互作用。（Mamba就像是每次参考前面所有内容的一个概括，越往后写对前面内容概括得越狠，丢掉细节、保留大意）
通过选择改进 SSM
将选择机制纳入模型的一种方法是，让影响序列交互的参数（如 RNN 的递归动力或 CNN 的卷积核）取决于输入。
主要做的改变就是将
(
Δ
,
A
,
B
,
C
)
(\Delta,A,B,C)
(Δ,A,B,C) 从固定的参数变为输入的函数，同时改变张量形状。使用线性层进行，其中参数现在有了一个长度维度L，意味着模型从时不变变成了时可变。这就失去了与卷积 (式3) 的等价性，对效率产生了影响。
高效实施 SSM 的选择性
目标是让选择性 SSM 在现代硬件（GPU）上也能高效运行。
先前模型的动机
在高层次上，诸如 SSM 这样的递归模型总是要在表现力和速度之间做出权衡：隐藏状态维度越大的模型应该越有效，但速度也越慢。因此，需要在不付出速度和内存代价的情况下最大化隐藏状态维度。递归模式比卷积模式更灵活，因为（式3）是由（式2）扩展来的。但是递归模式需要实体化形状为 (B, L, D, N) 的阴寒状态 h ，要比形状为 (B, L, D) 的输出 x 和 y 大得多。因此作者引入了更高效的卷积模式（式3a），可以绕过状态计算，仅使用形状为 (B, L, D) 的卷积核。传统的 LTI SSMs 利用双递归-卷积形式 (Dual Recurrent-Convolutional)，将有效状态维度提高到了 N(≈ 10 - 100)，远大于传统的RNN，且不会降低效率。
选择性扫描概述：硬件感知状态扩展
扫描（scan）： 指SSM采用了“扫描”操作来处理序列数据。扫描操作是一种算法，它按顺序处理序列中的元素，并累积结果。在计算机科学中，扫描通常指的是累积或逐步应用函数到列表或数组的元素上，以生成一个新的序列，其中每个元素是原序列到该点的函数应用结果。在SSM处理序列数据时，每个时间步的状态更新是基于前一时间步的状态和当前输入的累积效应。这种方法允许模型在每个步骤中选择性地集成新信息，同时考虑之前的状态，从而在整个序列上有效地传播信息。
这里借鉴了Flash Attention
选择机制旨在克服 LTI 模型的局限性，因此需要重新审视 SSM 的计算问题。用三种经典技术来解决这个问题：内核融合、并行扫描和重新计算。作者提出了两个主要观点：
原始的递归运算的复杂度为O(BLDN)，卷积的复杂度为O(BLDlog(L))，因此对于长序列且不太大的状态维度N，递归模式的复杂度更小。递归的顺序性和大量的内存使用是两个挑战。为了解决后者，我们可以尝试像卷积模式一样，不实体化完整的状态 h。
主要思路就是仅在更高效的内存层次结构中实体化状态 h，但大多数的操作会收到内存带宽的限制，包括扫描操作，于是作者使用核融合 (kernel fusion) 来减少内存 IO 的数量，从而比标准实现显著提速。
核融合（Kernel Fusion）: 核融合是一种编程技术，用于减少GPU上内存输入输出操作的次数。在GPU编程中，一个“核”是一个在GPU上执行的函数。通常，每个核函数的执行都需要从全局内存读取数据，处理后再写回内存。通过将多个核函数合并为一个（即“核融合”），可以减少内存访问次数，因为合并后的核可以在内部共享数据而无需每次都访问全局内存。这种技术可以显著提高性能，尤其是在需要多个步骤或操作才能完成的任务中。
传统方法需要在GPU HBM（高带宽内存，即显存） 上准备形状为 (B, L, D, N) 的扫描输入 (A, B) ，会占用大量内存空间。而作者直接将SSM参数 (Δ, A, B, C) 从慢速HBM加载到快速的SRAM（静态随机存取内存，即缓存）中，在SRAM中进行离散化和递归运算（SRAM比HBM有更快的访问速度，这可以提高计算效率）。然后将形状为 (B, L, D) 的最终输出写回HBM。
为了避免传统顺序递归的性能瓶颈，尽管SSM的递归不是线性的，但仍然可以通过高效的并行扫描算法实现并行化。
为了反向传播，通常需要存储中间状态。但在这种方法中，为了减少内存需求，作者采用了“重计算”技术。在反向传播过程中，当输入从HBM重新加载到SRAM时，中间状态被重新计算，而不是存储。因此这种融合的选择性扫描层的内存需求与使用 FlashAttention 的优化 Transformer 实现相当。
完整的选择性SSM层和算法如图所示：
简化的 SSM 架构
与结构化 SSM 一样，选择性 SSM 也是独立的序列变换，可以灵活地融入神经网络。H3架构通常由一个受线性注意力启发的块与MLP块交错组成。受门控注意力单元（GAU）的启发，在Mamba架构中，这两个组件被合并为一个单一的重复块，形成了一个均匀堆叠的结构。
本文的简化区块设计结合了 H3 区块（大多数 SSM 架构的基础）和现代神经网络中无处不在的 MLP 区块。这只是简单地重复 Mamba 模块，而不是交错使用这两个模块。与 H3 模块相比，Mamba 用激活函数取代了第一个乘法门。与 MLP 模块相比，Mamba 在主分支上增加了一个 SSM。对于$ \sigma $，使用 SiLU / Swish 激活。
Mamba架构通过一个可控的扩展因子 E 来扩大模型维度 D。在每个块中，大部分参数（
3
E
D
2
3ED^2
3ED2）用于线性投影（
2
E
D
2
2ED^2
2ED2用于输入投影，
E
D
2
ED^2
ED2用于输出投影）。与线性投影相比，SSM的参数（Δ, B, C 的投影和矩阵 A）数量要少得多。通过重复这个块，并与标准的归一化和残差连接交错，构成了Mamba架构。
在实验中，扩展因子 E 总是固定为2，使用两层堆叠的块来匹配Transformer的交错多头注意力（Multi-Head Attention, MHA）和MLP块的
12
D
2
12D^2
12D2 参数。采用SiLU/Swish激活函数，使得Gated MLP成为流行的“SwiGLU”变体。最后，受到RetNet在类似位置使用归一化层的启发，还使用了一个可选的归一化层（选择了LayerNorm）。
选择机制的特性
与门控机制的联系
RNN中的门控机制可以被视为SSM选择机制的一个实例，且这种联系在之前的研究中已经得到了建立。这种关系通过将连续时间系统离散化与RNN门控机制联系起来。定理1是对 ZOH 离散化和输入依赖门的改进，SSM 中的 Δ 可以看作是 RNN 门控机制的一种广义作用。作者认为 SSM 的离散化是启发式门控机制的原则基础。
其中
g
t
g_t
gt 是在时间步 t 的门控信号，
h
t
h_t
ht 是隐藏状态，
x
t
x_t
xt 是输入，
σ
σ
σ是激活函数。如果某个特定的输入
x
t
x_t
xt 应该被完全忽略，那么所有 D 通道都应该忽略它。为了实现这一点，作者将输入压缩到一维，然后与 Δ 结合（通过重复或广播）。
这部分提出的观点是，SSM的离散化是启发式门控机制的有原理的基础，这意味着通过离散化处理，SSM能以一种系统性的方式来决定如何从输入选择信息，并将其融入到状态更新中。这种方法能够在保持SSM模型的连续性的同时，允许模型动态地响应不同的输入，这在合成任务中尤其重要。这种选择性SSM递归公式提供了一种灵活的框架，能够在各种神经网络结构中灵活地集成SSM。
选择机制的解释
选择的特殊机制效应：
可变间距（Variable Spacing）: 选择性允许模型过滤掉可能出现在感兴趣输入之间的无关噪声标记。这在选择性复制任务中得到了体现，但在常见数据模态中无处不在，尤其是在处理离散数据时，例如语言中的填充词“um”。这个属性的产生是因为模型可以机械性地过滤掉任何特定的输入
x
t
x_t
xt，例如在门控RNN的情况下（定理1），当
g
t
→
0
g_t→0
gt→0 时。 过滤上下文（Filtering Context）：实证观察表明，许多序列模型在更长的上下文下并没有显示出改善，尽管原则上更多的上下文应该导致严格更好的性能。一个解释是许多序列模型在必要时不能有效地忽略无关的上下文，如全局卷积，以及一般的LTI模型。而选择性模型可以随时重置它们的状态以移除多余的历史，因此它们的性能原则上随着上下文长度的增加而单调改善。 边界重置（Boundary Resetting）：在将多个独立序列拼接在一起的设置中，Transformer可以通过实例化特定的注意力掩码来保持它们的独立，而LTI模型将会在序列之间混淆信息。选择性SSM也可以在边界处重置它们的状态（例如当
Δ
t
→
∞
Δ_t→∞
Δt→∞ 或定理1中的
g
t
→
1
g_t→1
gt→1）。
每个选择性参数的影响：
参数Δ的作用： 它控制着模型对当前输入
x
t
x_t
xt 的关注程度，以及应该保留多少历史状态信息。调节 Δ 的大小，可以模拟不同的系统行为，从完全关注当前输入到完全保留历史状态，实现对输入的选择性关注。Δ 在SSMs中起着类似于RNN门控信号的作用，如在定理1中提到的
g
t
g_t
gt，但在SSMs的框架下提供了一种更一般化的形式。 当 Δ 很大时，模型会重置状态
h
h
h，这相当于让模型更多地关注当前的输入
x
x
x，而非之前的状态。当 Δ 很小时，模型保持现有状态
h
h
h 的持久性，对当前的输入
x
x
x 给予较少的关注，从而忽略它。 SSMs 可以被看作是一个连续系统通过时间步长 Δ 离散化后的结果。在这个离散化的连续系统中，大的 Δ（趋向于无穷）意味着系统在较长时间内专注于当前输入，相当于“选择”了当前输入并忘记了当前状态。相反，小的 Δ（趋向于零）意味着当前输入是短暂的，可以被忽略。 参数A的作用： 虽然理论上 A 也可以具备选择性，但 A 对模型的主要影响是通过它与 Δ 的相互作用来实现的（通过离散化公式 $A=exp(ΔA) $）。作者认为，只要 Δ 具有选择性，就可以保证整个模型的选择性，而且 Δ 是提高模型性能的关键。为了保持模型的简洁性，作者选择不让 A 参数具备选择性。 参数 B 和 C 的作用： 选择性最重要的特性是过滤掉无关信息，从而将序列模型的上下文压缩成有效的状态。在 SSM 中，修改 B 和 C 使其具有选择性，可以使模型能够更精细地控制输入和状态的流动。
B 控制着输入
x
t
x_t
xt 是否被引入到状态 $h_t $ 中，即选择性地决定哪些输入对状态的更新至关重要。C 影响着状态
h
t
h_t
ht 如何转化为输出
y
t
y_t
yt，即如何基于当前的状态信息生成最终的输出。 这意味着，SSM可以基于当前处理的数据内容以及已经编码在隐藏状态中的上下文信息，来动态地调整其内部状态的更新方式和输出。这增加了模型处理序列数据时的灵活性和有效性，因为它可以根据数据的具体特征和任务需求来优化信息流。
其他模型细节
实数与复数
之前的大多数SSM使用复数来表示状态 h，这对于许多任务来说是必要的，因为它能提高性能。然而，近期的经验观察发现，完全基于实数的SSM在某些情况下也能表现良好，甚至可能更好。作者提出，实数与复数的选择可能与数据模态的连续性-离散性谱有关，复数对于连续模态（如音频、视频）有帮助，而对于离散模态（如文本、DNA）则可能不是必要的。
初始化
大多数之前的SSM建议使用特殊的初始化，尤其是在复数值情况下，这可以帮助在低数据情况下提高性能。对于复数情况的默认初始化是S4D-Lin，对于实数情况是S4D-Real，这是基于HIPPO理论的。这些初始化定义了矩阵 A 的第 n 个元素为
−
1
2
+
n
i
-\frac{1}{2}+ni
−21+ni和
−
(
n
+
1
)
-(n+1)
−(n+1)。作者认为，在大数据和实数SSM情况下，许多初始化方式都可能有效。
Δ 的参数化
选择性调整
Δ
\Delta
Δ 被定义为
S
Δ
(
x
)
=
Broadcast
D
(
Linear
1
(
x
)
)
S_{\Delta}(x)=\text{Broadcast}_D(\text{Linear}_1(x))
SΔ(x)=BroadcastD(Linear1(x))。这里，
B
r
o
a
d
c
a
s
t
D
Broadcast_D
BroadcastD是一个广播操作，它将一维线性变换
L
i
n
e
a
r
1
(
x
)
Linear_1(x)
Linear1(x)的结果扩展到
D
D
D 维空间。作者提到，这个调整可以从一维推广到更大的维度
R
R
R。这意味着可以将输入
x
x
x 映射到一个更大的空间，然后应用在这个扩展空间的操作。
R
R
R 相较于
D
D
D 很小，这意味着相比于模型的主要线性投影部分，这个调整使用的参数数量很少。推广到更高维度允许模型在更复杂的输入空间中进行选择性调整。这样可以更灵活地处理含有多个特征或维度的数据，从而提升模型对复杂数据结构的处理能力。 广播操作也可以被视为另一种线性投影，其初始化为特定的1和0模式。如果这个投影是可训练的，这将导致一个替代形式
S
Δ
(
x
)
=
Linear
D
(
Linear
R
(
x
)
)
S_{\Delta}(x) = \text{Linear}_D(\text{Linear}_R(x))
SΔ(x)=LinearD(LinearR(x))，这可以被视为一种低秩（即维度较低）的投影。
在实验中，
Δ
\Delta
Δ 参数的初始化遵循之前关于SSM的工作，通常初始化为
τ
Δ
−
1
(
Uniform
(
[
0.001
,
0.1
]
)
)
\tau^{-1}_{\Delta}(\text{Uniform}([0.001, 0.1]))
τΔ−1(Uniform([0.001,0.1]))。
这部分描述了SSM中
Δ
\Delta
Δ 的参数化方式和思路，这种参数化方式允许模型更灵活地处理输入数据，并对模型的行为产生显著影响，尤其是在处理不同维度空间的数据时。通过合理的初始化和维度设置，这种方法旨在优化模型的性能，同时保持参数数量的经济性。
总结
Mamba模型是一个创新的线性时间序列建模方法，巧妙地结合了递归神经网络（RNN）和卷积神经网络（CNN）的特点，解决了处理长序列时的计算效率问题。它通过状态空间模型（SSM）的框架，实现了RNN的逐步处理能力和CNN的全局信息处理能力的融合。在训练阶段，Mamba使用卷积模式来一次性处理整个输入序列，而在推理阶段则采用递归模式，逐步处理输入，这样的设计使得Mamba既能充分利用CNN的高效并行处理能力，又能保持RNN在序列数据处理上的灵活性。
Mamba模型的核心创新在于引入了选择性机制，这一机制通过使SSM参数依赖于输入，使得模型可以根据不同的输入动态调整其行为。这类似于RNN中的门控机制，但在SSM的框架下提供了更广泛的应用可能性。通过这种方式，Mamba能够有效地过滤掉无关信息，同时保留和加强与任务相关的信息，提高了对长序列数据的处理能力。
为了优化计算效率，Mamba采用了硬件感知算法，特别是利用GPU的内存层次结构来提高扫描操作的计算速度和降低内存需求。这种方法结合了RNN的递归计算效率和CNN的并行处理优势，使得Mamba在处理长序列数据时更加高效。
在架构设计上，Mamba简化了传统的SSM架构，通过合并H3和MLP块，形成了一个均匀堆叠的结构。这不仅简化了模型的结构，还提高了模型的灵活性和效率。
Mamba模型通过结合RNN和CNN的优点，并引入创新的选择性机制和硬件感知算法，成功地解决了传统Transformer模型在处理长序列数据时的计算效率问题，展现了在处理各种序列数据，尤其是语言、音频和基因组学数据时的出色性能和效率。
文章链接
您阅读本篇文章共花了：
发表评论