Mamba 状态空间模型原理与 Transformer 对比

Figure 3:(Architecture.) Our simplified block design combines the H3 block, which is the basis of most SSM architectures, with the ubiquitous MLP block of modern neural networks. Instead of interleaving these two blocks, we simply repeat the Mamba block homogenously. Compared to the H3 block, Mamba replaces the first multiplicative gate with an activation function. Compared to the MLP block, Mamba adds an SSM to the main branch. Forσ\sigmawe use the SiLU / Swish activation[hendrycks2016gaussian,ramachandran2017swish].

文章分类：大模型与架构
文章子分类：模型架构

---

1. 引言与动机

当前，基础模型（Foundation Models, FMs）几乎由 Transformer 主导。其核心 self-attention 机制能够在上下文窗口内密集地传递信息，从而建模复杂数据。然而，这一特性也带来了两个根本缺陷：

1. 二次复杂度：attention 的计算量和内存随窗口长度呈二次增长；
2. 固定上下文窗口：无法直接建模窗口外的信息。

为克服上述问题，研究者提出了大量次二次时间复杂度的架构，包括 linear attention、gated convolution、RNN 变体以及结构化状态空间模型（Structured State Space Models, SSMs，如 S4）。这些模型可视为 RNN 与 CNN 的结合，能够以线性或近线性复杂度进行序列建模，并在音频、视觉等连续信号模态上表现优异。

然而，在文本、DNA 等离散、信息密集的模态上，这些 SSMs 的表现仍不及 Transformer。作者指出，其核心瓶颈在于缺乏基于内容的选择能力（content-based selection）：传统 SSM 的参数对时间/输入保持不变，无法根据当前 token 选择性地保留或遗忘信息。

针对这一问题，Mamba 提出了选择性状态空间模型（Selective State Space Models）。通过让 SSM 参数成为输入的函数，模型获得类似 attention 的内容选择能力；同时，借助硬件感知的并行 scan 算法，Mamba 在保持线性序列复杂度的同时，达到了 Transformer 级别的建模能力。

---

2. 状态空间模型基础

2.1 连续与离散状态空间

结构化状态空间模型 S4 受经典状态空间模型启发，通过一个隐式潜状态 $h(t)$ 将输入 $x(t)$ 映射到输出 $y(t)$。连续系统定义为：

$$ h’(t) = \mathbf{A}h(t) + \mathbf{B}x(t) \tag{1a} $$

$$ y(t) = \mathbf{C}h(t) \tag{1b} $$

为了在离散序列上计算，需要通过离散化规则（如 zero-order hold, ZOH）将连续参数 $(\Delta, \mathbf{A}, \mathbf{B})$ 映射为离散参数 $(\overline{\mathbf{A}}, \overline{\mathbf{B}})$：

$$ \overline{\mathbf{A}} = \exp(\Delta \mathbf{A}), \quad \overline{\mathbf{B}} = (\Delta \mathbf{A})^{-1}(\exp(\Delta \mathbf{A}) - \mathbf{I}) \cdot \Delta \mathbf{B} \tag{2} $$

离散化后，模型可以按两种方式计算：

• Recurrent 形式：

$$ h_t = \overline{\mathbf{A}} h_{t-1} + \overline{\mathbf{B}} x_t \tag{3a} $$

$$ y_t = \mathbf{C} h_t \tag{3b} $$

• Convolution 形式：

$$ \overline{\mathbf{K}} = (\mathbf{C}\overline{\mathbf{B}}, \mathbf{C}\overline{\mathbf{A}}\overline{\mathbf{B}}, \dots, \mathbf{C}\overline{\mathbf{A}}^k\overline{\mathbf{B}}, \dots) \tag{4a} $$

$$ y = x * \overline{\mathbf{K}} \tag{4b} $$

2.2 线性时不变性（LTI）的局限

传统 SSM 要求参数 $(\Delta, \mathbf{A}, \mathbf{B}, \mathbf{C})$ 在时间/输入上保持不变，这一性质称为线性时不变性（Linear Time Invariance, LTI）。LTI 使得模型可以高效地通过卷积并行计算，但也意味着模型无法根据输入内容动态调整状态转移，从而在文本等离散模态上表现受限。

---

3. 选择性状态空间机制

3.1 核心思想：选择即压缩

序列建模的一个根本问题是如何将上下文压缩到有限状态中。Attention 不压缩上下文，因此效果好但效率低；RNN/SSM 压缩状态，因此效率高但效果受限于压缩质量。

Mamba 从两个合成任务出发说明问题：

• Selective Copying：在 Copying 任务基础上随机化待记忆 token 的位置，要求模型根据内容选择性地记住相关 token、忽略噪声；
• Induction Heads：要求模型根据上下文进行关联回忆。

LTI 模型在这两个任务上表现不佳，因为其动态固定，无法根据输入选择性地传播或过滤信息。

3.2 让 SSM 参数依赖于输入

Mamba 的选择机制核心在于：让影响序列交互的参数成为输入的函数。具体地，参数 $\Delta, \mathbf{B}, \mathbf{C}$ 被设为输入 $x$ 的函数：

$$ s_B(x) = \text{Linear}_N(x), \quad s_C(x) = \text{Linear}N(x), \quad s\Delta(x) = \text{Broadcast}_D(\text{Linear}_1(x)) $$

并通过 softplus 等变换得到最终的 $\Delta$。这一改动使模型从时间不变变为时间/输入可变，赋予其内容选择能力，但也打破了与卷积形式的等价性，对计算效率提出挑战。

3.3 选择机制的解释

• $\Delta$ 的作用：控制对当前输入的关注或忽略程度，可视为 RNN 门控机制的推广；
• $\mathbf{B}$ 与 $\mathbf{C}$ 的作用：分别控制输入是否进入状态，以及状态是否参与输出，实现基于内容和上下文的状态调制；
• 边界重置与上下文过滤：选择性使模型能够在任意位置重置状态、过滤无关历史，因此随着上下文增长性能可持续提升。

---

4. 硬件感知的并行算法

由于参数随输入变化，Mamba 无法使用高效的卷积计算，必须回到 recurrent scan。然而，直接物化形状为 $(B, L, D, N)$ 的潜状态会占用大量内存并造成 IO 瓶颈。

4.1 核心优化策略

为应对上述挑战，Mamba 采用三种经典技术：

1. Kernel fusion：将离散化、recurrence 等操作融合为一个 CUDA kernel，减少 HBM 与 SRAM 之间的数据搬运；
2. Parallel scan（scan）：利用 work-efficient 的并行 scan 算法打破 recurrence 的串行性；
3. Recomputation：前向传播时不保存中间状态，反向传播时从 HBM 重新加载输入并在 SRAM 中重新计算。

具体流程为：将 SSM 参数 $(\Delta, \mathbf{A}, \mathbf{B}, \mathbf{C})$ 从慢速 HBM 加载到快速 SRAM，在 SRAM 中完成离散化和 recurrence，最终只将输出 $(B, L, D)$ 写回 HBM。

4.2 复杂度与速度

• 理论上，selective scan 的 FLOPs 为 $O(BLDN)$，序列长度严格线性扩展；
• 在 A100 等硬件上，优化后的 scan 比标准 PyTorch 实现快约 20–40 倍，比卷积式 SSM 快约 3 倍；
• 端到端推理吞吐比同规模 Transformer 高约 5 倍，因为无需维护历史 KV cache。

---

5. Mamba 整体架构

Figure 9:(Audio Pretraining.) Mamba improves performance over prior state-of-the-art (Sashimi) in autoregressive audio modeling, while improving up to minute-long context or million-length sequences (controlling for computation).

Figure 1:(Overview.) Structured SSMs independently map each channel (e.g.D=5D=5) of an inputxxto outputyythrough a higher dimensional latent statehh(e.g.N=4N=4). Prior SSMs avoid materializing this large effective state (D​NDN, times batch sizeBBand sequence lengthLL) through clever alternate computation paths requiring time-invariance: the(Δ,𝑨,𝑩,𝑪)(\Delta,\bm{A},\bm{B},\bm{C})parameters are constant across time. Our selection mechanism adds back input-dependent dynamics, which also requires a careful hardware-aware algorithm to only materialize the expanded states in more efficient levels of the GPU memory hierarchy.

Figure 2:(Left) The standard version of the Copying task involves constant spacing between input and output elements and is easily solved by time-invariant models such as linear recurrences and global convolutions. (Right Top) The Selective Copying task has random spacing in between inputs and requires time-varying models that canselectivelyremember or ignore inputs depending on their content. (Right Bottom) The Induction Heads task is an example of associative recall that requires retrieving an answer based on context, a key ability for LLMs.

Figure 5:(Induction Heads.) Models are trained on sequence length28=2562^{8}=256, and tested on increasing sequence lengths of26=642^{6}=64up to220=10485762^{20}=1048576. Full numbers inLABEL:tab:induction.

Mamba 将 SSM 模块与 Transformer 的 MLP 块融合为单一、同质的 Mamba 块，去除了独立的 attention 或 MLP 块。每个 Mamba 块通过可控的扩展因子 $E$ 扩展模型维度 $D$，其中大部分参数集中在线性投影，内部 SSM 参数占比较小。

整体架构特点如下：

• 端到端、纯 recurrent 序列模型；
• 训练时计算和内存随序列长度线性增长；
• 自回归推理时每步仅需常数时间，无需 KV cache；
• 支持百万级上下文长度。

---

6. 实验验证与性能

6.1 合成任务

在 Selective Copying 和 Induction Heads 任务上，Mamba 不仅能轻松解决，还能外推到超过 100 万 token 的序列长度，而 LTI 模型和 attention 在更长序列上表现受限。

Selective Copying 结果：

Model	Arch.	Layer	Acc.
S4	No gate	S4	18.3
-	No gate	S6	97.0
H3	H3	S4	57.0
Hyena	H3	Hyena	30.1
-	H3	S6	99.7
-	Mamba	S4	56.4
-	Mamba	Hyena	28.4
Mamba	Mamba	S6	99.8

6.2 语言建模

Mamba-3B 在预训练困惑度和下游任务上匹敌两倍参数 Transformer，生成吞吐提升约 5 倍。

下游零样本评估（部分结果）：

Model	Token.	Pile ppl↓	LAMBADA ppl↓	LAMBADA acc↑	HellaSwag	PIQA	Arc-E	Arc-C	WinoGrande	Average
Hybrid H3-130M	GPT2	—	89.48	25.77	31.7	64.2	44.4	24.2	50.6	40.1
Pythia-160M	NeoX	29.64	38.10	33.0	30.2	61.4	43.2	24.1	51.9	40.6
Mamba-130M	NeoX	10.56	16.07	44.3	35.3	64.5	48.0	24.3	51.9	44.7
Hybrid H3-360M	GPT2	—	12.58	48.0	41.5	68.1	51.4	24.7	54.1	48.0
Pythia-410M	NeoX	9.95	10.84	51.4	40.6	66.9	52.1	24.6	53.8	48.2
Mamba-370M	NeoX	8.28	8.14	55.6	46.5	69.5	55.1	28.0	55.3	50.0
Pythia-1B	NeoX	7.82	7.92	56.1	47.2	70.7	57.0	27.1	53.5	51.9
Mamba-790M	NeoX	7.33	6.02	62.7	55.1	72.1	61.2	29.5	56.1	57.1
GPT-Neo 1.3B	GPT2	—	7.50	57.2	48.9	71.1	56.2	25.9	54.9	52.4
Hybrid H3-1.3B	GPT2	—	11.25	49.6	52.6	71.3	59.2	28.1	56.9	53.0
OPT-1.3B	OPT	—	6.64	58.0	53.7	72.4	56.7	29.6	59.5	55.0
Pythia-1.4B	NeoX	7.51	6.08	61.7	52.1	71.0	60.5	28.5	57.2	55.2
RWKV-1.5B	NeoX	7.70	7.04	56.4	52.5	72.4	60.5	29.4	54.6	54.3
Mamba-1.4B	NeoX	6.80	5.04	64.9	59.1	74.2	65.5	32.8	61.5	59.7
GPT-Neo 2.7B	GPT2	—	5.63	62.2	55.8	72.1	61.1	30.2	57.6	56.5
Hybrid H3-2.7B	GPT2	—	7.92	55.7	59.7	73.3	65.6	32.3	61.4	58.0
OPT-2.7B	OPT	—	5.12	63.6	60.6	74.8	60.8	31.3	61.0	58.7
Pythia-2.8B	NeoX	6.73	5.04	64.7	59.3	74.0	64.1	32.9	59.7	59.1
RWKV-3B	NeoX	7.00	5.24	63.9	59.6	73.7	67.8	33.1	59.6	59.6
Mamba-2.8B	NeoX	6.22	4.23	69.2	66.1	75.2	69.7	36.3	63.5	63.3
GPT-J-6B	GPT2	–	4.10	68.3	66.3	75.4	67.0	36.6	64.1	63.0
OPT-6.7B	OPT	–	4.25	67.7	67.2	76.3	65.6	34.9	65.5	62.9
Pythia-6.9B	NeoX	6.51	4.45	67.1	64.0	75.2	67.3	35.5	61.3	61.7
RWKV-7.4B	NeoX	6.31	4.38	67.2	65.5	76.1	67.8	37.5	61.0	62.5

6.3 基因组学（DNA）

在 HG38 人类基因组数据上，Mamba 随模型规模和上下文长度增长均表现更好。当上下文长度从 1K 增加到 1M 时，Mamba 的预训练困惑度持续改善，而 HyenaDNA 等 LTI 模型性能反而下降。在区分大猩猩、黑猩猩等近缘物种的分类任务上，Mamba 也能利用长达百万的上下文。

6.4 音频建模与生成

在 YouTubeMix 钢琴音乐和 SC09 语音生成任务上，Mamba 超越了 SaShiMi、Hyena 和 Transformer。Mamba-UNet 在 SC09 上的 FID 显著低于更大规模的 GAN/diffusion 模型。

Model	Params (M)	NLL↓	FID↓	IS↑	mIS↑	AM↓
SampleRNN	35.0	2.042	8.96	1.71	3.02	1.76
WaveNet	4.2	1.925	5.08	2.27	5.80	1.47
SaShiMi	5.8	1.873	1.99	5.13	42.57	0.74
WaveGAN	19.1	-	2.03	4.90	36.10	0.80
DiffWave	24.1	-	1.92	5.26	51.21	0.68
+ SaShiMi	23.0	-	1.42	5.94	69.17	0.59
Mamba	6.1	1.852	0.94	6.26	88.54	0.52
Mamba	24.3	1.860	0.67	7.33	144.9	0.36

Outer	Center	NLL↓	FID↓	IS↑	mIS↑	AM↓
S4+MLP	MHA+MLP	1.859	1.45	5.06	47.03	0.70
S4+MLP	S4+MLP	1.867	1.43	5.42	53.54	0.65
S4+MLP	Mamba	1.859	1.42	5.71	56.51	0.64
Mamba	MHA+MLP	1.850	1.37	5.63	58.23	0.62
Mamba	S4+MLP	1.853	1.07	6.05	73.34	0.55
Mamba	Mamba	1.852	0.94	6.26	88.54	0.52

6.5 消融实验

架构与 SSM 层消融：

Model	Arch.	SSM Layer	Perplexity
Hyena	H3	Hyena	10.24
H3	H3	S4 (complex)	10.30
-	H3	S4 (real)	10.34
-	H3	S6	8.95
-	Mamba	Hyena	10.75
-	Mamba	S4 (complex)	10.54
-	Mamba	S4 (real)	10.56
Mamba	Mamba	S6	8.69

选择性参数消融：

Selective $\Delta$	Selective $\mathbf{B}$	Selective $\mathbf{C}$	Perplexity
✗	✗	✗	10.93
✗	✓	✗	10.15
✗	✗	✓	9.98
✓	✗	✗	9.81
✓	✓	✓	8.71

$\mathbf{A}$ 初始化消融：

$\mathbf{A}_n$ Initialization	Field	Perplexity
$-1/2 + ni$	Complex	9.16
$-1/2$	Real	8.85
$-(n+1)$	Real	8.71
$\exp(\mathcal{N}(0,1))$	Real	8.71

$\Delta$ 投影维度消融：

Size of $\Delta$ proj.	Params (M)	Perplexity
-	358.9	9.12
1	359.1	8.97
2	359.3	8.97
4	359.7	8.91
8	360.5	8.83
16	362.1	8.84
32	365.2	8.80
64	371.5	8.71

状态维度 $N$ 消融（注：原文表格未明确标注上下两组的完整条件，此处按原文列出）：

State dimension $N$	Params (M)	Perplexity
1	367.1	9.88
2	367.4	9.86
4	368.0	9.82
8	369.1	9.82
16	371.5	9.81
1	367.1	9.73
2	367.4	9.40
4	368.0	9.09
8	369.1	8.84
16	371.5	8.71

6.6 总结

Mamba 通过选择性状态空间机制，使 SSM 获得内容选择能力，同时借助硬件感知 scan 算法保持线性复杂度。实验表明，Mamba 在语言、音频、基因组学等多个模态上达到或超越 Transformer 水平，尤其在长序列场景下优势显著，并具备更高的推理吞吐。

---

原文来源

Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
论文链接：https://arxiv.org/abs/2312.00752
代码与预训练模型：https://github.com/state-spaces/mamba

原文链接

https://arxiv.org/html/2312.00752