2Mamba2Furious：线性复杂度，媲美Softmax注意力的精度

通过简化Mamba-2并引入二阶隐藏状态，实现线性复杂度下接近Softmax注意力精度

前置知识

Softmax注意力

标准Transformer的核心机制，通过计算Query和Key的内积后经softmax归一化得到注意力权重，再与Value相乘。其计算复杂度为序列长度的二次方 $O(N^2)$，内存占用为KV Cache的 $2 \times N \times d_h$。虽然Flash Attention等优化可提升GPU利用效率，但FLOP数仍保持二次复杂度。

理解Softmax注意力的二次复杂度瓶颈是本文研究动机的核心，论文目标就是在保持线性复杂度的前提下逼近其精度

线性注意力

将Softmax注意力中的指数非线性替换为可分解核函数 $\phi(Q)\phi(K)^T$，使得计算可先做Key-Value乘积再与Query相乘，训练复杂度降至 $O(N)$，推理时可实现常数复杂度的RNN形式。但精度通常不如Softmax注意力，因为一阶核函数表达能力有限。

这是本文改进的基础架构，论文的核心贡献就是提升线性注意力的表达能力

Mamba-2

基于状态空间模型(SSM)的高效序列架构，引入了关联扫描算法和衰减掩码(A-mask)。从架构角度看，Mamba-2可视为带有衰减掩码的线性注意力，包含多个组件：输入卷积、A-mask、时间离散化dt、D残差、Z门控等。在效率和精度上显著优于朴素线性注意力。

本文以Mamba-2为起点，通过消融实验拆解其组件，构建简化版本并进一步改进

泰勒展开与高阶隐藏状态

通过泰勒展开，Softmax注意力中的指数函数可分解为内积的所有非负整数幂次之和。一阶线性注意力对应 $p=1$，隐藏状态维度为 $(d_h, d_h)$；高阶RNN的隐藏状态维度约为 $(d_h^p, d_h)$。研究表明，增加线性注意力的阶数可逐步逼近Softmax精度，因为更多高阶项捕获了更丰富的交互信息。

这是本文提出二阶隐藏状态的理论基础，解释了为什么对QK内积平方能大幅提升精度

衰减掩码(A-mask)

在注意力分数上施加的时间衰减机制，使模型倾向于关注近期token而逐渐遗忘远期信息。原Mamba-2中A-mask与时间离散化参数绑定：$A = -\exp(A_{\log}) \odot dt$。本文发现softplus变体 $A = -\text{softplus}(A)$ 更有效，且解耦了与离散化参数的依赖关系。

A-mask是Mamba-2中提升精度最关键的组件之一，本文对其进行了重要改进

遗忘Transformer(Forgetting Transformer)

近期提出的Transformer变体，在Softmax注意力中引入衰减掩码以实现选择性遗忘。其掩码函数使用 $\log(\sigma(x))$，与本文使用的 $-\text{softplus}(-x)$ 在数学上几乎等价（仅差一个符号）。这一发现揭示了SSM模型与Transformer变体之间的深层联系。

本文的指数化变体(2Mamba-E)与遗忘Transformer有深刻联系，这一发现拓展了论文的意义

研究动机

Softmax注意力虽然是现代大语言模型的基石，但其训练时的二次复杂度 $O(N^2)$ 严重限制了长序列应用。即使Flash Attention通过CUDA内核优化提升了GPU利用率，其FLOP数仍保持二次方。线性注意力通过用核函数替代softmax可将复杂度降至线性，但精度始终落后于Softmax注意力。Mamba-2作为当前最强的线性注意力变体之一，其代码库包含众多组件（输入卷积、A-mask、D残差、Z门控、时间离散化等），但这些组件中哪些真正对精度有贡献，哪些是不必要的复杂性，缺乏系统性的消融研究。

本文的目标是本文有三个递进目标：首先系统地消融Mamba-2的各个组件，识别出对精度贡献最大的核心要素；其次基于消融结果构建简化但高效的Mamba-2S基线；最后通过引入二阶隐藏状态（对QK内积平方），在保持线性复杂度的前提下达到甚至超越Softmax注意力的精度水平。

与已有工作不同的是，现有工作如Performer、CosFormer、RetNet、DeltaNet等尝试改进线性注意力，但都是从不同角度出发的独立改进。本文的独特切入点是：以Mamba-2为对象进行系统性拆解和重建，而非从零设计新架构。通过'简化→理解→增强'的研究范式，不仅得到一个更简洁高效的模型，还揭示了隐藏状态阶数与Softmax注意力之间的理论联系（泰勒展开视角），以及指数化变体与遗忘Transformer之间的等价关系。

核心方法

本文采用'消融→简化→增强'三步走的研究路线。直觉上，Mamba-2虽然强大但包含过多未经验证的组件，就像一台满载的赛车，需要先拆解看哪些零件真正提供动力。第一步是系统消融Mamba-2的7个核心组件，分别在隔离环境下测试每个组件对精度的贡献。第二步是基于消融结果，只保留最有价值的组件，构建极简但同样强大的Mamba-2S。第三步是引入高阶隐藏状态理论，将QK内积从一阶提升到二阶（平方），得到最终的2Mamba模型。技术路线的关键洞察是：Softmax注意力可通过泰勒展开理解为所有非负整数阶线性注意力之和，因此增加阶数（从1到2）能显著逼近Softmax精度。

核心创新在于对QK内积进行平方操作 $(QK^T)^2$，将一阶线性注意力提升为二阶。这一操作有两个重要效果：首先，从表达能力角度，二阶隐藏状态捕获了token之间更丰富的交互模式，类似于泰勒展开中高阶项的作用；其次，从数值稳定性角度，平方操作使内积空间严格非负，允许使用在线softmax归一化算法（如Flash Attention风格），这比输出RMSNorm更稳定。与已有方法的本质区别在于：现有工作（如RetNet、DeltaNet）都在一阶框架内做改进，而本文直接提升到二阶，从根本上增加了隐藏状态的表达维度。隐藏状态维度从 $(d_h, d_h)$ 增长到约 $(d_h^2, d_h)$，但通过利用内积的交换律，实际唯一项数为 $d_h(d_h+1)/2$，不到朴素Kronecker积 $d_h^2$ 的一半。

方法步骤详情

Mamba-2S构建过程：从输出归一化的线性注意力出发，依次加入卷积（窗口大小2）和softplus A-mask，这是精度提升最大的两个组件。然后可选地加入时间离散化参数dt。具体前向传播为：(1) 输入经线性投影和1D卷积得到Q、K、V；(2) dt通过softplus激活；(3) A通过 $-\text{softplus}(h \cdot W_A)$ 计算；(4) 对A做累积和得到ACS，构建衰减矩阵 $AM_{ij} = \exp(ACS_i - ACS_j)$；(5) 结合因果掩码M计算注意力输出。2Mamba在此基础上：移除时间离散化（因大模型训练不稳定），移除RMSNorm，将QK内积替换为平方内积 $(QK^T)^2$，加入softmax归一化。2Mamba-E则进一步将平方替换为指数函数 $\exp(QK^T)$，本质上是带softplus A-mask和输入卷积的Softmax注意力。

技术新颖性

技术新颖性体现在多个层面：(1) 首次系统性消融Mamba-2的官方代码库组件，此前这些组件的选择缺乏充分验证；(2) 发现softplus A-mask优于原Mamba-2的指数A-mask，且与时间离散化解耦后更简洁；(3) 将泰勒展开视角的高阶隐藏状态理论应用于Mamba-2架构，实现了从一阶到二阶的跳跃；(4) 发现二阶隐藏状态的记忆效率在序列长度超过约1058时优于Softmax注意力（基于$d_h=64$的计算）；(5) 揭示了指数化2Mamba与遗忘Transformer的等价关系 $\log(\sigma(x)) = -\text{softplus}(-x)$，建立了SSM与Transformer变体之间的理论桥梁；(6) 提供了高效的Triton内核实现二阶内积计算，利用多项式定理减少重复计算。

实验结果

消融实验揭示，A-mask（特别是softplus变体）和输入卷积是Mamba-2中精度贡献最大的两个组件。具体而言，softplus A-mask将测试损失从3.13降至2.88，卷积（窗口2）降至2.93，而其他组件（D残差、Z门控、时间离散化等）的贡献相对微小。构建Mamba-2S时，卷积+softplus A-mask的组合（测试损失2.77）甚至优于单独组件之和，且加入时间离散化后进一步降至2.76。令人意外的是，Z门控反而略微降低了精度。在最终模型比较中，2Mamba（平方内积变体）在300M参数小模型上达到了与Softmax注意力几乎相同的测试损失，且在700M参数的中等模型上同样保持竞争力。2Mamba-E（指数变体）甚至略优于Softmax注意力。在Needle-in-a-Haystack测试中，2Mamba在上下文检索方面略优于Softmax注意力，显著优于Mamba-2，验证了其有效利用上下文的能力。内存分析表明，当序列长度超过约1058个token时，2Mamba的内存占用低于Softmax注意力。

**Figure 8**: Needle-in-a-Haystack基准测试结果

查看结构化数据

任务	指标	本文	基线	提升
语言建模（300M模型，序列长度2048）	测试损失	2Mamba: ~2.7	Softmax Attention: ~2.65, Mamba-2: ~2.85	相比Mamba-2降低约0.15，逼近Softmax注意力
语言建模（700M模型，序列长度8192）	测试损失	2Mamba: ~2.55	Softmax Attention: ~2.5, Mamba-2: ~2.7	在长序列上接近Softmax精度，远超Mamba-2
消融：Softplus A-mask（隔离）	测试损失	2.88	朴素线性注意力: 3.13	单组件最大提升，降低0.25
消融：输入卷积（窗口2，隔离）	测试损失	2.93	朴素线性注意力: 3.13	第二重要组件，降低0.20
Needle-in-a-Haystack（8K上下文）	检索准确率	2Mamba: 略优于Softmax	Mamba-2: 显著低于Softmax	2Mamba在上下文利用上超越Softmax
内存效率（d_h=64）	单头隐藏状态大小	2Mamba: $d_h(d_h+1)/2 + 3d_h = 2272$ 元素	Softmax: $2Nd_h$（N>1058时更大）	序列长度>1058时内存更低

局限与改进

论文存在几个重要局限：(1) 时间离散化参数在大模型中导致数值不稳定，需要强制使用FP32精度（约8倍减速），最终在2Mamba中被移除，损失了一定精度；(2) 所有实验的最大序列长度为8192，未验证在更长上下文（如32K、128K）下的表现；(3) 模型规模仅测试到700M参数，未验证在十亿级或更大模型上的可扩展性；(4) 数据集仅使用FineWeb单一数据集，未在多语言、代码、数学等多样任务上验证；(5) 二阶隐藏状态的计算虽然内存更优，但FLOP数可能高于一阶变体，论文未详细讨论训练速度对比；(6) 消融实验基于300M模型，可能无法完全代表更大模型的组件重要性排序；(7) 论文提到的90K步评估存在bug（实际评估的是每10K步后的第1步），虽然作者认为不影响结论，但降低了结果的严谨性。

独立分析的弱点

论文有几个值得关注的弱点：(1) **时间离散化的稳定性问题未根本解决**：只是简单移除了dt组件，但这是Mamba-2的重要组件之一。改进方向可以是设计数值更稳定的离散化方案，例如使用梯度裁剪或混合精度训练策略；(2) **二阶隐藏状态的计算效率未充分优化**：虽然内存更优，但 $(QK^T)^2$ 的计算本身可能比一阶更耗时。可探索稀疏化或低秩近似来减少二阶项数量；(3) **缺乏与其他最新线性注意力方法的对比**：如Gated DeltaNet（作者在结论中提及但未实验比较），这限制了对2Mamba相对优势的全面理解；(4) **消融实验的统计显著性**：各组件的精度差异较小（测试损失相差0.01-0.05），未报告置信区间或多随机种子的方差；(5) **实际端到端效率验证不足**：论文主要讨论隐藏状态内存，但未提供实际训练吞吐量、推理延迟等端到端性能数据。

未来方向

作者提出了几个有价值的研究方向：(1) 将DeltaNet与2Mamba结合，利用delta rule进一步提升表达能力，这在Gated DeltaNet的成功基础上是有前景的；(2) 探索动态隐藏状态大小，研究是否可以通过自适应调整二阶项数量来优化精度-效率权衡；(3) 研究QK内积平方与其他投影方式（如将Q和K投影到相同维度的低秩空间）的精度对比。此外可延伸的方向包括：(4) 将2Mamba扩展到视觉、多模态等领域，验证其在不同模态下的表现；(5) 探索更高阶（三阶、四阶）隐藏状态的收益递减曲线；(6) 结合稀疏注意力模式，在二阶隐藏状态基础上进一步提升长序列处理能力；(7) 将2Mamba-E（指数变体）与遗忘Transformer进行更深入的理论和实验对比。

复现评估

论文在复现性方面做得很好：代码已在GitHub开源（https://github.com/gmongaras/2Mamba2Furious），HuggingFace上也提供了模型集合。训练使用HuggingFace FineWeb数据集（CC-MAIN-2024-51版本），这是公开可用的大规模数据集。论文提供了完整的算法伪代码（Algorithm 1-5），包含明确的张量形状标注，便于理解实现细节。Triton内核也已开源，这是实现二阶内积计算的关键。主要的复现挑战在于：(1) 训练需要较大算力（300M模型训练100K步，700M模型在多个序列长度上训练）；(2) 2Mamba的Triton内核实现需要GPU编程经验；(3) 论文提到的评估bug可能影响精确数值的复现。总体而言，复现难度中等，适合有一定资源的研究团队。

2Mamba2Furious：线性复杂度，媲美Softmax注意力的精度 2Mamba2Furious: Linear in Complexity, Competitive in Accuracy