Mamba的隐藏代价：逃离二次内存却难逃二次参数需求

引言：注意力机制的二次困境

自Transformer问世以来，自注意力机制（Self-Attention）的二次复杂度（Quadratic Complexity）就成为长序列建模难以绕开的瓶颈。当序列长度为N时，注意力矩阵的计算与存储成本都会以N²的量级增长，直接限制了大模型处理超长上下文的能力，也催生了一大批线性化方案。

要理解这一瓶颈的严重性，需要回到Transformer的核心计算逻辑：自注意力机制需要计算序列中每个token与所有其他token之间的相关性分数，生成一个N×N的注意力矩阵，其中每个元素表示两个位置之间的关联强度。具体而言，给定输入序列的Query(Q)、Key(K)、Value(V)三个矩阵，注意力输出为softmax(QK^T/√d)V，其中QK^T这一步就生成了N×N的注意力矩阵。在多头注意力（Multi-Head Attention）中，虽然每个头的维度d_k较小，但头数h的设置使得总计算量仍为O(N²d)。当N从4096扩展到128K甚至1M时，计算成本会爆炸性增长——128K长度的序列需要的注意力计算量是4K长度的1024倍。对于GPT-4级别的模型处理128K上下文时，单次前向传播中注意力计算就涉及约160亿次浮点运算（仅注意力部分），而存储注意力权重需要约32GB显存（以fp16精度计）。

自注意力机制的二次复杂度不仅体现在计算量上，还体现在内存带宽上。具体而言，对于序列长度N和隐藏维度d，注意力计算需要O(N²d)的浮点运算量和O(N²)的内存来存储注意力权重矩阵。在现代GPU上，由于HBM（高带宽内存）与SRAM（片上缓存）之间的数据搬运成本往往主导实际延迟，FlashAttention通过分块计算将注意力矩阵的分块保持在SRAM中，避免了将完整的N×N矩阵写入HBM，从而在不改变计算复杂度的情况下实现了2-4倍的实际加速。FlashAttention的核心洞察来自对GPU内存层次的分析：现代GPU（如A100）的SRAM容量仅约20MB但带宽极高（19TB/s），而HBM容量为40-80GB但带宽相对有限（2TB/s）。标准注意力实现需要将N×N的注意力矩阵反复在HBM和SRAM之间搬运，导致内存带宽成为瓶颈。FlashAttention采用在线softmax算法（online softmax），将Q、K、V分块加载到SRAM中计算局部注意力，通过维护运行中的最大值和归一化因子来保证数学等价性，从而将HBM访问量从O(N²)降低到O(N²d/M)，其中M是SRAM大小。FlashAttention-2进一步通过优化并行策略和工作分配，在A100上达到了理论峰值FLOPS的72%。但即便如此，计算量本身的O(N²)增长仍然是不可回避的硬限制。这一瓶颈催生了FlashAttention、稀疏注意力（如Longformer的滑动窗口）、以及各类线性注意力方案。

Mamba正是其中最受瞩目的架构之一。凭借状态空间模型（State Space Model, SSM）的选择性机制，Mamba用线性复杂度实现了对长序列的建模，被寄望成为Transformer的有力挑战者。然而，一篇发表在Substack并在Reddit引发热议的分析指出：当你试图逃离二次内存时，二次参数需求可能会在另一个维度上找到你。

Mamba的核心权衡：从内存到参数

状态空间模型的数学基础

在深入讨论Mamba的训练困境之前，有必要理解其底层的状态空间模型原理。SSM源自控制论中的一组连续时间微分方程：dx/dt = Ax + Bu，y = Cx + Du，其中x是隐状态，u是输入，y是输出，A/B/C/D是系统矩阵。将其离散化后，可以得到递推形式：x_k = Āx_{k-1} + B̄u_k。

状态空间模型的数学框架最早来自20世纪60年代的控制论，由Rudolf Kalman等人发展，用于描述动态系统的行为。在深度学习中，SSM的复兴始于2021年的HiPPO（High-order Polynomial Projection Operators）理论，该理论提出了一种最优的在线函数逼近方法，为A矩阵的初始化提供了理论基础——HiPPO矩阵能够以最小误差在有限维状态中压缩连续信号的历史。随后的S4（Structured State Spaces for Sequence Modeling）通过对角化加低秩分解使SSM计算可行，H3引入了门控和短卷积的组合，最终Mamba在2023年12月通过选择性机制和硬件感知的实现取得突破性进展。

Mamba的关键创新在于引入「选择性机制」——让矩阵B和C依赖于输入内容动态变化，使模型能够根据当前输入选择性地记忆或遗忘历史信息。这打破了传统线性时不变SSM的限制，赋予模型内容感知的能力，但也使得模型无法再用卷积形式高效并行训练，转而依赖扫描算法（selective scan）进行序列化计算。

传统的线性时不变（LTI）SSM，如S4模型，其系统矩阵A、B、C在整个序列处理过程中保持不变。这种时不变性带来了一个重要的计算优势：递推关系可以被展开为卷积形式，利用FFT在O(N log N)时间内并行计算整个序列的输出。然而，时不变性也意味着模型对所有输入token一视同仁地进行状态更新，缺乏根据内容选择性过滤信息的能力——这使得LTI-SSM在需要内容感知推理的任务（如选择性复制、归纳推理）上表现不佳。Mamba的选择性机制通过让B和C成为输入的函数，使模型获得了类似于门控机制的内容感知能力，但代价是丧失了卷积形式的并行优势。

序列化梯度带来的优化难题

该分析的核心论点在于Mamba的训练方式。与Transformer可以对整个序列并行计算注意力不同，Mamba的状态更新是**序列化（Sequential）**的——每个token的表示依赖于之前累积的状态。这在推理时是效率优势，但在训练阶段（尤其是反向传播）却带来了微妙的优化困难。

作者用一个四token序列「A、B、C、D」来说明这一过程。在时间反向传播（Backpropagation Through Time, BPTT）中：

• 模型处理A和B时，仅依据A来优化B的表示；
• 接着处理C，此时它面对的是（A, B）的组合状态，需要针对这个累积上下文来优化C；
• 最后处理D，它要基于（A, B, C）的状态进行优化。

BPTT是训练此类循环结构的标准方法，其核心思想是将循环网络沿时间步展开为一个深层前馈网络，然后应用标准的反向传播算法。对于长度为T的序列，梯度需要从最后一个时间步逐步回传到第一个时间步，经过T次矩阵乘法。这带来两个经典挑战：梯度消失（梯度在回传过程中指数衰减，导致早期信息难以被学习）和梯度爆炸（梯度指数增长，导致训练不稳定）。虽然Mamba通过精心设计的参数化——特别是对角化的A矩阵结构——在一定程度上缓解了这些问题，但序列化依赖本身带来的优化耦合效应仍然存在。

在Mamba的实际训练中，虽然前向传播可以通过硬件友好的并行扫描算法（parallel scan / prefix sum）在O(N)时间内完成，但反向传播的梯度计算仍然面临序列依赖的本质挑战。并行扫描（parallel scan）又称前缀和（prefix sum）算法，是并行计算中的经典原语。对于Mamba中的递推关系x_k = A_k·x_{k-1} + B_k·u_k，虽然表面上是纯串行的，但由于矩阵乘法满足结合律，可以将相邻时间步的转移操作两两合并。具体地，定义二元运算(A_i, b_i) ⊕ (A_j, b_j) = (A_j·A_i, A_j·b_i + b_j)，该运算满足结合律，因此可以用经典的Blelloch并行前缀和算法在O(log N)步内计算所有前缀结果。在GPU上，这意味着一个长度为N的序列的前向传播可以在约log₂(N)个同步步骤内完成，相比纯串行的N步有巨大提升。然而，这种并行化需要O(N)的工作量和O(N)的空间来存储中间结果。此外，由于选择性机制使得每个时间步的转移矩阵不同，梯度的精确计算需要存储或重计算每个时间步的局部Jacobian矩阵，这进一步增加了训练时的内存和计算压力。

状态参数的连带优化问题

问题的关键在于：模型在优化每个新token时，并不只是优化这个新token本身，同时也在优化其所依赖的状态参数。这意味着序列中间位置的token承担了更沉重的优化负担——它们既要表达自身信息，又要维系整个累积状态的一致性。

随着序列变长，这种叠加效应愈发明显。位于序列中段的token需要在一个被不断挤压的表示空间中，同时服务于「过去的记忆」和「当前的信息」。这就构成了Mamba隐藏的参数成本。

二次参数需求的理论推论

表示空间为何必须二次增长

作者由此得出一个颇具冲击力的结论：如果不想破坏模型的理解能力，就必须让表示空间（representation space）以二次方式增长。 换句话说，Mamba虽然在内存复杂度上摆脱了二次增长，却可能在参数维度上重新遭遇二次需求。

这是一个「按下葫芦浮起瓢」的经典权衡。Transformer把二次成本花在了注意力矩阵的内存与计算上；而Mamba若要在长序列上维持同等的表达能力，则可能不得不在状态维度（state dimension）上付出二次代价。二者只是把复杂度从一个地方转移到了另一个地方。

从信息瓶颈（Information Bottleneck）理论的视角来看，SSM的固定维度隐状态相当于一个信息瓶颈层。信息瓶颈方法由Tishby等人在1999年提出，是一种在数据压缩与预测能力之间寻找最优平衡的信息论框架。在SSM的语境中，隐状态可以被视为对历史序列的一种有损压缩表示。对于典型的Mamba配置，隐状态维度d通常为16到64（state expansion factor），而模型维度D可能为数千。这意味着在每个时间步，模型需要将D维的输入信息压缩到d×D维的状态空间中。当序列包含需要远距离关联的信息时——例如在128K长度的序列中，第1000个token的信息需要在第100000个token处被精确回忆——这些信息必须在经过99000次状态更新后仍然被保留。每次状态更新都会引入新的信息竞争和潜在的覆写风险，这就是为什么研究者认为状态维度可能需要随序列长度增长。

长序列建模是否存在复杂度守恒律

这个观察触及了一个更深层的问题：长序列建模是否存在某种「守恒律」？直觉上，要在长距离上精确保留信息，模型必须拥有足够的容量来编码这些依赖关系。无论架构如何设计，信息论层面的下限或许难以被彻底规避。

信息论中的数据处理不等式（Data Processing Inequality）为这一直觉提供了形式化支撑：经过任何确定性或随机变换后，信息只会减少不会增加。当SSM用固定维度d的隐状态来压缩长度为N的历史序列时，根据率失真理论（Rate-Distortion Theory），存在一个信息保留能力的理论上限。率失真理论由Shannon在1959年建立，量化了在给定压缩率下可达到的最小失真。具体而言，如果我们将SSM隐状态视为对历史的R比特编码（R ≈ d × log₂(precision)），那么对于具有互信息I(past; future) = Ω(N)的序列（如需要精确回忆N个独立事实），率失真函数告诉我们，要实现零失真（完美回忆），编码率R必须至少等于I(past; future)，即隐状态容量需要线性于N增长。如果序列中包含O(N)比特的相关信息，而隐状态只有O(d)比特的容量，那么当N远大于d时，必然会产生信息丢失。这也解释了为什么有研究者认为，要无损地处理长序列中的全部依赖关系，模型容量必须以某种方式与序列长度相关联。

SSM类模型用固定大小的隐状态压缩历史，这种压缩本身就是有损的——而为了减少信息损失，状态容量就得随之膨胀。这一逻辑在理论上具有说服力，但仍需更多实证支撑。

理性看待：论点的价值与局限

一个待验证的分析视角

需要强调的是，这一论点来自单一来源（Substack作者的分析文章，经Reddit传播），尚未经过同行评审或大规模实验验证。作者提出的「二次参数需求」更多是一种理论推演与直觉论证，而非严格的数学证明或实证结果。

事实上，学界与业界对Mamba的实际表现存在不同看法。许多基准测试显示，Mamba及其变体（如Mamba-2）在长序列任务上确实能以更低的成本达到与Transformer相当的效果。值得一提的是，Mamba-2通过将选择性SSM重新表述为结构化状态空间对偶（SSD）框架，揭示了SSM与注意力机制之间的深层联系——选择性SSM本质上等价于一种特殊的半可分矩阵（semiseparable matrix）形式的注意力。SSD框架的核心发现是：选择性SSM的输入输出关系可以表示为y = M⊙(QK^T)V的形式，其中M是一个下三角的半可分矩阵，Q和K分别由C和B参数化。半可分矩阵的特征是其任意子矩阵的秩都不超过固定常数N（即状态维度），这意味着SSM本质上是一种低秩结构化的注意力。这一对偶性使得Mamba-2可以在SSM模式（适合长序列推理）和类注意力模式（适合并行训练）之间灵活切换，并利用张量核（tensor cores）进行高效的矩阵乘法运算，而非Mamba-1中大量依赖SRAM的元素级操作。这一发现不仅使Mamba-2的训练速度提升2-8倍，还为理解SSM的表达能力边界提供了新的理论视角。这说明所谓的「二次参数需求」即便存在，也未必在实践中构成致命障碍——工程上的选择性机制、混合架构都可能有效缓解这一问题。

对模型架构设计的三点启示

无论该论点最终是否成立，它提出的思考角度都有实际价值：

1. 不存在免费的午餐。任何声称「彻底解决二次复杂度」的架构，都应被审视它把成本转移到了何处。
2. 训练动态与推理效率需分开评估。Mamba的推理优势明显，但训练时的优化难度和参数效率是另一个需要独立考量的维度。
3. 混合架构可能是务实之选。当前许多前沿模型正在尝试将SSM与注意力层结合，取长补短，而非追求纯粹的线性方案。Jamba（AI21）、Zamba（Zyphra）、以及Griffin（DeepMind）等模型都在实践中将SSM层与注意力层交替使用——用SSM层高效处理局部和中程依赖，用少量注意力层捕捉关键的长程依赖，从而在效率与表达能力之间取得务实平衡。

混合架构的设计理念基于一个经验观察：在大多数实际文本中，绝大多数token的依赖关系是局部的（几百到几千token的窗口内），只有少数关键位置需要跨越整个上下文的长程注意力。Jamba模型的设计中，每6个Mamba层搭配1个注意力层，在保持接近纯SSM推理效率的同时，显著提升了需要精确长程检索的任务表现。Zamba则采用了共享注意力层的策略，进一步减少参数开销。实验数据显示，这类混合架构在「大海捞针」（Needle-in-a-Haystack）测试中的表现远优于纯SSM模型，同时在推理吞吐量上仍保持对纯Transformer 2-5倍的优势。这种务实的设计思路表明，架构之争的最终答案可能不是非此即彼的选择，而是精心编排的组合。

结语：Mamba不是银弹，但仍是重要创新

「二次参数需求会在你逃离二次内存时找到你」——这句略带戏剧性的论断，提醒我们对长序列建模的复杂度问题保持清醒。Mamba无疑是近年来最重要的架构创新之一，但它并非银弹。

这场关于Mamba真实成本的讨论，本质上是在追问一个根本命题：在有限的隐状态中压缩无限的历史，究竟需要付出多大的代价？答案或许要等到更多严谨的实验与理论工作来给出。在此之前，保持批判性思考，既不神化线性架构，也不低估它的价值，才是理性的态度。