AI底层架构的一场革命：Mamba如何用「状态压缩」挑战Transformer的十年统治

Transformer的注意力机制造就了ChatGPT，却也埋下了二次方计算的炸弹。Mamba用选择性状态空间重新定义序列建模，推理内存首次做到恒定不变。2026年3月，Mamba-3以Apache 2.0开源——这场挑战还远未结束。

Transformer的隐藏代价

2017年Google发表《Attention Is All You Need》后，Transformer成为NLP乃至整个深度学习的主干架构。其核心——自注意力机制（Self-Attention）——让模型能够动态计算序列中任意两个位置的相关性，这种"全局感知"能力是GPT、BERT等大模型成功的根基。

但这一机制存在一个结构性缺陷：注意力矩阵的计算复杂度是 O(n²)，其中 n 是序列长度。上下文窗口从8K扩展到128K tokens，注意力层的计算量增加了约256倍。在推理阶段还需维护KV Cache（Key-Value缓存），其内存占用随序列长度线性增长，长文本推理成本居高不下。

工程实况

以GPT-4规模的模型处理200K token上下文为例，仅KV Cache一项就可能占用数十GB显存。这正是为什么长上下文推理在生产环境中依然是高成本操作——Transformer的架构在根本上决定了这一瓶颈难以消除。

状态空间模型：Mamba的数学基础

Mamba的前身是状态空间模型（State Space Model, SSM）。SSM脱胎于控制理论，其核心思路是：不记录完整历史，而是将过去压缩进一个固定维度的"状态向量" h(t)，用两个方程描述系统动态：

// 连续时间SSM核心方程

状态方程（状态演化）

h′(t) = A · h(t) + B · x(t)

输出方程（预测生成）

y(t) = C · h(t) + D · x(t)

x(t)：输入序列｜h(t)：隐藏状态（压缩的历史语境）｜y(t)：输出序列
矩阵A、B、C、D 在深度学习中均为可学习参数，通过预训练优化

• 状态方程， h′(t)=A*h(t)+B*x(t)

处理离散文本序列时，需将连续方程离散化，引入步长参数 Δ（Delta） 控制采样粒度，得到离散递推形式：

// 离散化后（类RNN递推形式）

离散状态更新

h_t = Ā · h_{t-1} + B̄ · x_t

离散输出预测

y_t = C̄ · h_t

Ā、B̄、C̄ 为经零阶保持器（ZOH）离散化后的对应矩阵，包含步长Δ的影响

• 输出方程， y(t)=C*h(t)+D*x(t)

S4：结构化SSM突破长程依赖瓶颈

2021年Albert Gu等人提出S4模型，引入 HiPPO（高阶多项式投影算子）技术初始化矩阵A，通过正交多项式的数学性质使模型天然具备捕捉长距离依赖的能力。S4论文指出：仅将随机初始化的矩阵A替换为HiPPO矩阵，序列MNIST任务的准确率从60%直接跃升至98%。

S4还发现：由于SSM仅含乘法与加法，递推过程可等价展开为一维卷积核，结合快速傅里叶变换（FFT）实现训练阶段的并行加速——训练时以CNN运算，推理时以RNN递推，兼得两者效率。

Mamba的核心创新：选择性状态空间（S6）

传统SSM存在一个根本局限：矩阵A、B、C对所有输入和所有时刻均相同，即"线性时不变性（LTI）"。这意味着模型无法根据当前语境动态决定哪些历史信息重要、哪些可以遗忘——而这恰恰是Transformer注意力机制的核心能力。

Mamba的解法是让关键参数成为当前输入的函数：

选择性机制：三个输入相关参数

Δ_t（步长）：控制当前输入对隐藏状态的影响强度。Δ大：对旧信息快速"遗忘"，强调当前词元；Δ小：历史语境被更多保留，甚至Δ趋零时当前输入对状态几乎无影响。

B_t（输入矩阵）：决定当前词元如何写入隐藏状态。例如输入"yesterday"时，B_t引导状态更新方向为"过去时态语境"。

C_t（输出矩阵）：决定隐藏状态中哪些信息被读取用于预测输出。延续上例：C_t使后续动词输出偏向过去时变位形式。

这三个参数均通过三条并行线性投影层（类比Transformer中Q/K/V的生成方式）从当前词元嵌入向量 x_t 中计算得到。矩阵A保持固定结构（承担全局历史记忆），选择"使用哪部分历史"的职责由B_t和C_t动态完成。作者将这种带选择机制的SSM命名为 S6，以Mamba架构为名发表——"多个S连读形似蛇的嘶鸣"是命名灵感之一。

硬件感知并行扫描：训练效率的工程解法

引入输入相关参数后，训练阶段无法再使用卷积并行化捷径（因为转移核不再恒定）。Dao和Gu的解法是：利用SSM运算的数学结合律，将顺序递推拆解为独立小块，由GPU通过并行前缀和扫描（Parallel Prefix Scan）算法同时处理。

同时借鉴Tri Dao此前开发的FlashAttention技术，将中间计算结果显式分配在GPU的SRAM（片上高速缓存）而非HBM（显存），大幅减少内存读写次数——这正是"硬件感知"的含义所在。

Mamba Block：完整模块结构

在实际网络中，S6选择性SSM作为更大"Mamba模块"的核心子模块运行，配合门控机制共同工作：

// Mamba Block 数据流（简化）

输入 x
├─ 线性扩展（2× 维度） → 拆为 x_proj 和 z_proj
├─ x_proj → 1D卷积（局部特征）→ SiLU激活 → S6（选择性SSM）→ y_ssm
└─ z_proj → SiLU激活 → z_act（门控权重）
y_ssm ⊙ z_act（逐元素相乘：门控调节输出强度）
└─ 线性投影压缩回原维度
└─ + 残差连接（原始输入直接加回）→ 输出

⊙ 表示Hadamard积（逐元素乘法）。门控路径 z_act 类比音量旋钮，控制SSM语境信息的影响强度。

多个Mamba Block顺序堆叠构成完整的Mamba语言模型，最终层接softmax输出词表概率分布。

Mamba-2：状态空间对偶性理论

2024年5月，Dao与Gu发表续作，提出状态空间对偶性（State Space Duality, SSD）框架，从理论层面揭示了SSM与Transformer注意力机制的数学等价关系——两者可被视为同一类算子在不同分解方式下的特例。

这一发现的工程意义是：Mamba得以直接吸收Transformer社区多年积累的优化技术。Mamba-2在此基础上简化了B、C矩阵的生成方式（直接切片而非矩阵乘法），将核心算法压缩至约25行代码，训练吞吐量相比Mamba-1提升显著，并支持类似多头注意力（MHA）的"多头SSM"结构，便于GPU张量并行。

Mamba-3：专为推理重新设计

Mamba-3由相同团队主导，论文发布于2025年底，2026年3月以Apache 2.0许可证正式开源。其设计哲学从"提升训练效率"转向"推理优先（Inference-First Design）"——呼应大规模部署中推理成本远超训练成本的现实。

1. 隐藏状态尺寸 降至Mamba-2一半

2. 语言建模指标 相对Mamba-2提升≈4%

3. 推理延迟 进一步降低

三项核心改进：

MAMBA-3 技术改进

BCNorm：在B、C矩阵投影后引入RMS归一化，类比Transformer中广泛使用的QKNorm，大规模训练稳定性显著提升，同时移除了Mamba-2为稳定性添加的post-gate RMSNorm（在纯Mamba-3模型中）。

复数A矩阵：将矩阵A拆分为实部（送入SSD）与虚部（通过RoPE计算），使模型获得旋转位置编码能力，改善位置理解和长度泛化。

B/C偏置项：为B矩阵添加通道级偏置，增强逐块更新的表达能力，提升模型在复杂序列模式上的跟踪与推理能力。

代码发布于GitHub mamba-ssm仓库，Apache 2.0许可证允许企业直接商用。这是Mamba系列迄今工程实用性最强的版本，开发者可无缝接入现有PyTorch生态。

Mamba vs Transformer：架构能力对比

研究表明，Transformer在需要少样本提示、文本复制、多跳推理等任务中仍有结构性优势——这些任务本质上依赖对历史输入的"精确检索"，而SSM的固定状态压缩会损失细节。业界的选择因此不是"非此即彼"，而是混合架构：Nvidia Nemotron-3、IBM Granite 4.0、AI2 Jamba等主流模型均将Mamba层与Attention层交织，以Mamba层处理长程上下文压缩，以Attention层保留精确检索与推理能力。

演进全景

2021

S4模型：引入HiPPO矩阵初始化，解决SSM长程依赖问题；发现训练卷积化并行加速路径。

2022

H3 / DSS / S4D：S4衍生变体，进一步简化矩阵结构，铺垫Mamba诞生。

2023年12月

Mamba（S6）：引入选择性扫描机制（B_t、C_t、Δ_t均由输入决定）+ 硬件感知并行扫描算法。首次在主流LLM基准上以同等参数量媲美Transformer，推理吞吐最高提升5×。

2024年5月

Mamba-2 + SSD理论：建立SSM与Transformer的数学等价框架（状态空间对偶性）；B/C生成简化为切片操作；支持多头SSM结构；约25行代码实现核心算法。Bamba、Jamba等混合模型基于此构建。

2025年11月 

Mamba-3：推理优先设计；BCNorm稳定大规模训练；复数A矩阵引入RoPE位置编码；状态尺寸减半同时困惑度持平；语言建模效率提升约4%。Apache 2.0，商用可直接使用。

Mamba系列最深远的意义或许不在于某个基准上的几个百分点，而在于它为AI部署提供了一种新的可能性：推理内存恒定意味着理论上支持无限上下文，实际上仅受硬件物理限制。当序列越来越长、模型越来越大，这一特性将从边缘计算、实时推理到超长文档处理等场景中产生实质性的成本优势。

Mamba-3的开源更是将这一优势直接交到了开发者手中。它不是Transformer的终结者——而是那个让AI真正"跑得起来"的架构补丁，混合时代已经到来。

更多transformer，VIT，swin tranformer
参考头条号：人工智能研究所
v号：人工智能研究Suo, 启示AI科技

 动画详解transformer  在线视频教程