点击下方卡片，关注「AI视界引擎」公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
序列模型如Transformer和RNN常对无关上下文过度分配注意力，导致中间表示存在噪声。这通过助长幻觉、削弱长距离和检索能力以及降低鲁棒性来损害大语言模型（LLM）的性能。
近期研究表明，差异设计可以缓解Transformer中的这一问题，提升其在各种应用中的有效性。
本文探讨了最初为Transformer开发的技术是否可以应用于Mamba——一种基于选择性状态空间层的最新架构，该架构以更高的效率实现了Transformer Level 的性能。作者证明，将差异设计直接应用于Mamba是不够的，需要仔细的架构修改。
为此，作者为Mamba引入了一种新型差异机制，并在语言建模基准上进行了经验验证，展示了改进的检索能力和超越原版Mamba的优越性能。
最后，作者进行了广泛的消融研究和实证分析，以论证作者的设计选择，并提供证据表明2507有效缓解了基于Mamba的模型中的过度分配问题。
unsetunset1 引言unsetunset
为深度序列建模设计增强型架构是机器学习领域的一项关键任务，因为序列模型推动了自然语言处理中的ChatGPT[1]和计算机视觉中的Stable diffusion[32]等突破性进展，这些进展是现代生成模型的基础。然而，这些模型面临两大主要挑战：效率，受限于Transformer架构在序列长度上的二次时间复杂度[36]，以及鲁棒性，由于不一致性、可靠性问题和幻觉现象，导致性能欠佳。作者的工作通过增强具有亚二次复杂度的近期高效架构（如Mamba[16]）的鲁棒性，使其更加可靠和鲁棒，从而应对这些挑战。
为解决这一问题，作者旨在减少隐藏层中对无关上下文的过度注意力分配，这通常会导致噪声表示。2507受到Ye等人[43]的启发，他们在Transformer中通过差异设计缓解了这一问题。这种架构修改是Diff-Transformer的核心，Diff-Transformer是一种Transformer变体，它相比原始模型在实证性能上有所提升，包括更强的鲁棒性、增强的检索能力和长上下文处理能力。
作者的重点是改进最近引入的Mamba架构，该架构基于选择性状态空间层（S6）。该架构以其效率而闻名，在序列长度上引入了次二次复杂度，同时还能实现自回归解码，其复杂度保持独立于序列长度。除了这些效率优势之外，最近的研究表明，基于Mamba的架构即使在相对较大的规模下也能匹配甚至超越Transformer的SoTA性能。例如，纯Mamba模型Falcon-Mamba 7B [46]在多个语言任务上达到了与LLaMA-3-8B相当的性能。Mamba在大语言模型（LLM）中的其他值得注意的整合案例包括Jamba [23]、Zamba [14]、Hymba [11]以及其他模型[39, 37, 31]。更重要的是，最近的基于Mamba的模型在推理模型方面表现出色[28, 40]，突显了它们在持续的测试时扩展革命中的核心作用，例如OpenAI的O1 [20]和DeepSeek的R1 [19]等模型。
虽然过度分配问题是一个普遍存在的、并非特定于任何架构的问题，作者假设基于Mamba的大语言模型（LLMs）比Transformer模型更容易出现过度分配。这主要归因于两个因素：
（i）Mamba是一种无softmax的架构，这意味着它缺乏有助于抑制无关注意力权重的指数缩放效应，以及（ii）作为基于状态的模型，Mamba进行局部操作，无法在不考虑所有中间token的情况下直接处理远距离的token，导致重要token分散在无关token之间。
这引出了作者的核心研究问题：
差分设计能否用于提升Mamba模型的鲁棒性？作者假设通过缓解过度分配问题来提升鲁棒性，可以增强模型的整体能力，改善检索和长上下文处理，并可能减少大语言模型中与一致性和幻觉相关的问题。
作者通过展示，尽管差分机制的天真实现并未提升Mamba架构的性能，但更精心设计的机制则可以，从而对这个问题给出了肯定的答案。通过在语言任务上的系统评估、消融研究和实证分析，作者得出结论，作者的变体相较于原版Mamba更具优势。
作者的主要贡献如下：
（i）作者提出了Diff-Mamba，这是基于Diff-Transformer灵感对Mamba架构的一种改进，它缓解了将过多的注意力分数分配给无关上下文的问题，并提升了模型的通用语言建模能力。
（ii）通过一系列消融研究和利用机制可解释性工具进行的实证分析，作者论证了作者的设计选择，并证明了2507获得的中间表示确实更少嘈杂。
（iii）最后，作者展示了Diff-Mamba相较于Mamba在检索和长上下文能力方面有所提升。这一点尤为重要，因为像Mamba这样的循环LLM主要设计用于解决Transformer二次复杂度带来的低效问题，这在长上下文场景下尤为关键。
因此，Diff-Mamba在架构最需要的领域实现了对Mamba的改进[6]。
unsetunset2 背景unsetunset
在此，作者描述了科学背景并介绍了讨论2507的关键术语。
2.1 差分Transformer
自注意力机制 自注意力机制是Transformer架构的基本组成部分[36]，对自然语言处理和计算机视觉领域的最新进展产生了深远影响。该机制通过捕捉成对token之间的依赖关系，实现动态的焦点分配，使模型能够确定序列中每个token的相对重要性。数学上，其表达式为：
注
意
力
α
α
在此公式中，
、
和
分别代表 Query 、 Key和Value ，而
表示键的维度。Transformer通过采用
个并行注意力头扩展了这一机制，使模型能够捕捉更广泛的依赖关系。
差分注意力为了解决对无关token过度分配注意力的问题。Ye等人[43]引入了Diff Transformer，这是一种通过将每个注意力头分成两个并从其中一个注意力图中减去另一个注意力图来减少注意力噪声的机制。该机制可以描述如下：
注
意
力
其中
是一个可学习的标量。为了更好地改善训练动态，对
进行重新参数化，并在每个头的末尾（减法之后）应用 Group Normalization [42]。
2.2 状态空间层
状态空间层最初由Gu等人[17]引入，后来由S4模型[18]进行了显著改进。此后，它们在各种领域展现出强大的性能，包括自然语言处理[12, 26]、音频生成[15]、图像建模[3, 27]、长时程视频理解[38]、强化学习和语音识别[33]。这些模型实现了基于时不变状态空间公式的线性递归更新规则，这些规则可以通过卷积并行高效计算，且具有亚二次复杂度。
2.3 Mamba和选择性状态空间层
一个Mamba模块处理一个信号
，其中
是隐藏维度，
是token的数量。其核心机制是S6层，其前向路径形式化为：
σ
σ
Ŷ
其中
是S6层的输入，也是S6层的输入，
。函数
表示SiLU激活函数，
表示门控分支的逐元素乘法。每个Mamba模块主要由线性层和卷积层参数化，以及下面描述的S6层内部组件。
S6层是SSM中最流行的变体，它采用实值、对角线选择型SSM。标准实值和对角线型SSM由一个对角线转移矩阵
、6个输入输出矩阵
参数化，其中
是状态大小，以及一个时间尺度
。此类SSM的每个通道可视为通过以下递归规则将输入标量序列
映射到输出标量序列
的映射：
其中
是离散化函数，离散系统矩阵为
和
。公式4中的递归规则可以使用高效的并行扫描 [35] 或通过快速傅里叶变换（FFT）的简单标量卷积 [17] 在现代硬件加速器上并行高效计算。注意，公式4是从
到
的映射，为了处理
个通道，使用了多个独立的实例。
S6层与标准SSM的不同之处在于采用了选择性机制，其中系统矩阵是输入相关的。因此，系统变为时不变系统，每一步的系统矩阵由所有通道集确定，然后独立应用于处理每个通道。整个机制可以通过以下方式计算：
和
来定义时变矩阵如下：
以及时变递归规则：
作者研究了“多对一”的设置，其中模型处理整个输入序列以生成单个输出。这种模式在自然语言处理（NLP）中广泛使用，涵盖了自回归下一个词预测和序列分类任务。
2.3.1 Mamba作为隐式注意力
Mamba与线性注意力层之间的联系在[2, 8, 34]中已有明确阐述。具体而言，研究表明S6的单通道时变递归更新规则（见公式6）在A为隐式注意力矩阵时，可以明确展开为线性注意力公式
。
这种观点进一步被Zimerman等人[45]扩展，他们将S6的隐式注意力解释从公式7推广到包含Mamba模块的大多数组件，包括激活函数、归一化层、卷积层和门控分支，将其整合到一个统一的隐式注意力公式中。
数据控制线性算子 Data-Controlled Linear Operators 的概念首先由Poli等人 [30] 提出，他们证明了自注意力机制可以被视为此类算子的一种表达形式。这一原则指导了作者设计Hyena层。随后，Ali等人 [2] 展示了S6层也可以在隐式变体下统一于这一表述，Zimerman等人 [45] 进一步将此扩展到其他架构，包括整个Mamba模块、RWKV [29] 和Griffin [10]。这种扩展视角启发了2507，使作者将微分设计解释为一种隐式参数化低噪声数据控制线性算子的方法。这一见解促使作者决定在Mamba Level 应用微分设计。
unsetunset3 方法unsetunset
作者从将微分机制最简单地整合到Mamba模块开始。2507受到微分 Transformer [43]的启发，后者在注意力 Level 而非 Transformer Level 应用减法。该机制建立在自注意力之上，可以描述为：
对
于
所
有
属
于
：
，
，
其中
是softmax函数，
被唯一参数化以确保稳定性并改善训练动态：
为了更好地将此技术应用于Mamba模型，作者通过数据控制的线性算子[30]的视角重新诠释了差分注意力，从而得到以下公式：
这里，矩阵A定义了一个数据控制的线性算子。
3.1 差分S6
公式11定义了一种实现差分Mamba的直截了当的方法，即通过减去S6层得到的值而不是注意力层得到的值。这种方法建立在S6层和注意力层之间的两个关键相似性之上：（i）在Mamba中，S6层的作用与Transformer中的注意力层相同——捕获 Token 之间的交互，（ii）S6层已被证明是一种隐式的因果线性注意力形式。
因此，将微分机制整合到Mamba模块中可以通过以下方式实现：
其中
的定义与公式 10 类似。
类似于公式11，该公式可以重写为数据控制线性算子的形式：
在任意时间步
和模型索引
下，
和
是由系统矩阵
和
控制的S6的隐式注意力矩阵，定义如下：
Diff S6与Diff Attention（参见公式12和9）之间一个关键的区别在于，Diff Attention在同一尺度上对元素进行相减，因为softmax函数产生的值位于[0,1]范围内。相比之下，S6产生未归一化且无界的输出。为了解决这种差异，作者引入一个额外的归一化步骤，记为
。
为简化起见，作者定义
为：
其中
是一个可学习的参数
，用于将
参数化为一个正值且更稳定的权重。
3.2Diff-Mamba
然而，正如结果部分详细所述，Diff S6的表现不佳，在实际应用中未能达到标准Mamba层的水平。作者推测这源于S6过于简单，未能作为通用、具有表达力的注意力层替代方案。因此，它无法充分发挥微分技术的全部潜力。为解决这一局限性，作者借鉴了Zimerman等人[45]的研究成果，他们证明Mamba模块可以作为注意力层的替代混合机制，通过将其构建为数据控制的线性算子来实现。
特别是，作者们展示了Mamba可以被重新表述为隐式注意力，通过：
其中M是一个表示卷积层的矩阵，而
是S6对应的线性算子，如Ali等人[2]所正式定义。
这种表述将Mamba定义为一个数据控制线性算子，具有更丰富和更具表达力的隐式注意力矩阵。基于这一见解，作者引入了Diff-Mamba，这是一种将微分方法扩展到完整Mamba模块的机制：
λ
类似于公式11和13，这可以重写为一个数据控制的线性算子，其定义为：
Diff Attention与Diff-Mamba之间的一个关键区别在于后者在更广泛的组件集上应用了减法操作，如图1所示。
类似于公式15中归一化的Diff S6变体，作者添加一个归一化项：
由此得到归一化的Diff-Mamba机制，这是作者主要的贡献。
最后，遵循Ye等人[43]的方法，作者将所有变体的输出乘以
3.3 整体架构
遵循Mamba架构，2507避免了将Mamba或自注意力等序列混合器与MLP交错混合。相反，它完全由Diff-Mamba块组成，并辅以标准的深度学习组件，如 Token Embedding 层、 Token 解嵌入层、 Shortcut 和归一化层。
高效实现 直接计算归一化Diff-Mamba（公式20）的朴素方法几乎将推理延迟增加一倍。为缓解这一问题，作者通过将模型划分为两条独立路径，在Mamba的单次前向传递中并行计算被减数和减数。为保持原始Mamba架构的计算复杂度、参数数量和内存占用，作者避免了Mamba块中标准的通道扩展做法，传统上通道数量会翻倍。相反，作者在通道维度上复制输入表示，有效保留模型规模。该并行实现的可视化展示在图2中。
unsetunset4 实验unsetunset
在本节中，作者通过实证方法评估了Diff-Mamba架构的有效性。作者首先在4.1节中证明，Diff-Mamba在多个数据集上的小规模语言建模任务中表现优于原始Mamba架构。在4.2节中，作者通过一系列全面的消融研究来论证作者的关键设计决策。随后，在4.3节中，作者展示了Diff-Mamba相对于Mamba的优越检索性能。最后，在4.4节中，作者利用机制可解释性领域的工具，如Tuned-lens [4]，来比较Diff-Mamba和Mamba的内部表示，通过实证验证作者的差异方法能够有效减少表示中的噪声。
中介表示。作者的实验设置完整描述在附录A中。
4.1 语言建模
为评估Diff-Mamba相对于Mamba在通用NLP任务上的性能，作者使用相同规模的模型和完全一致的训练设置，从头开始训练这两个模型，包括相同的代码库[18]、数据集和超参数。作者关注三个广泛使用的基准：WikiText-103、Text8和Enwik8，并针对不同深度的模型进行实验。最终结果报告在表1中。为更全面地展示优化过程，作者在图3中包含了训练过程中的测试曲线。
如表1所示，Diff-Mamba在所有评估环境中均优于Mamba，实现了持续更低的损失。具体而言，对于12层模型，Diff-Mamba在WikiText-103上比Mamba提高了0.4困惑度，在Text8上提高了0.046比特每字节（bpb），在Enwik8上提高了0.041 bpb。对于6层模型，Diff-Mamba在WikiText-103上比Mamba提高了0.201，在Text8上提高了0.02，在Enwik8上提高了0.007。有趣的是，作者观察到随着层数的增加，Mamba中的差异设计变得越来越有效。对此的一个可能解释是，在较低层中，由隐式注意力矩阵捕获的依赖关系较短且简单。在这种情况下，Mamba能够有效地管理过度分配，而无需差异机制。然而，在较高层中，依赖关系变得更加复杂，跨越更长的范围[5]并表现出更多样化的模式。这加剧了过度分配的影响，从而使差异设计的优势更加显著。此外，图3中的训练曲线揭示了Diff-Mamba的优化特性，显示其始终优于Mamba并实现更快的收敛。作者假设这种现象源于差异设计减少了噪声量，而噪声似乎对提高收敛性至关重要[21, 44]。表2：消融研究比较（i）差异设计的作用范围
4.2 消融实验分析
为验证作者关于以下设计决策的实验结果：(i) 在S6层与Mamba层应用差分操作，(ii) 在相减操作前加入额外的归一化子层，以及(iii) 将重新参数化的
转换为标量，作者在Text8基准数据集上进行了专门的消融实验。表2总结了实验结果。所有模型具有相同的参数数量，并使用针对 Baseline 模型优化的相同超参数进行训练。实验验证了以下方面：(i) 差分机制的应用（Diff-S6与DiffMamba），(ii) 包含归一化操作的效果（"w.
"）与不包含归一化操作的效果（"w.o Nrm"），以及(iii) 对
进行重新参数化的重要性（"re.
"）。所有模型均在完整的Text8数据集上训练，且具有相同的参数数量。报告的值为第10个epoch上的测试困惑度（PPL）。越低越好。
可以观察到，所有三种设计选择都是合理的，具体来说，Diff-S6在无归一化情况下表现优于Mamba，而Diff-Mamba则优于Diff-S6。加入额外的归一化能够提升结果，分别为Diff-Mamba和Diff-S6带来了0.015和0.008的困惑度提升。最后，
重新参数化并未带来更好的性能，这一点在Diff-S6和Diff-Mamba模型中均有体现。
4.3 检索
Diff-Transformer相较于原始Transformer架构展现出显著提升的检索能力。因此，作者在零样本和微调设置下，针对BABILong基准测试[22]中的五个检索任务，开展专项实验，以评估这些增强的检索能力是否能够迁移至Diff-Mamba。BABILong包含多种推理任务，如事实链、简单归纳和演绎，其中相关事实嵌入在长度不一、上下文各异的自然语言段落中。这种设置提供了严格的基准，用于评估模型在长上下文中检索和推理的能力。在此实验中，Diff-Mamba通过用Diff-Mamba层替换预训练Mamba 370M模型的最后12层而开发。为确保公平比较，Diff-Mamba和Mamba均在PG19数据集的2B token上进行微调，随后在BABILong的五个任务上进行零样本评估。此外，两者均在BABILong任务上进行微调，最多使用1k token，以促进更具针对性的比较，确保其学习过程与相同目标保持一致。
结果（图4）仅报告在测试集上的表现。每个单元格展示了其中一个模型优于另一个模型的比率。绿色表示Diff-Mamba获胜，红色表示Mamba获胜。从结果中可以看出，Diff-Mamba在较长的上下文长度下始终优于Mamba，表现出更慢的性能退化以及随着上下文长度增加而更大的比率分数。尽管其在零样本设置下的性能在PG19微调过程中受到次优对齐的部分影响，但Diff-Mamba在零样本场景下仍保持强劲表现，突显其有效关注相关上下文信息的能力。在BABILong微调测试中，Diff-Mamba持续获胜，比率分数高达2.11，并在零样本评估中赢得了长达3.5的上下文任务。总体而言，作者的研究结果表明Diff-Mamba实现了卓越的检索性能，特别是在BABI-long微调和零样本设置下的长上下文场景中。
4.4 中间表示中的噪声分析
作者在第4.3节和4.1节中的实证分析，以及Diff Transformer的结果表明，差异化设计可以缓解过度分配问题并提升整体性能。为了进一步探究这种现象的深层原因，作者利用机制可解释性领域的工具分析模型的内部表征。具体而言，作者借助调谐透镜[4]——一种通过训练仿射 Prob 将每一层的激活映射到模型最终预测来检查中间表征的方法，从而实现逐层可解释性并洞察模型的内部计算。基于这一工具，作者对比了Diff-Mamba与标准Mamba模型在隐藏表征中的信噪比。调谐透镜工具将每一层的表征投影到检索任务中下一个token的预测上。通过测量针token的预测概率，作者可以估计不同层的信噪比。值得注意的是，如图5所示，在大多数层中，Diff-Mamba表现出比Mamba更高的信噪比。这种差异在早期层尤为明显，其中Diff-Mamba中针token的预测概率高出几个数量级。这项分析从经验上证明了Diff-Mamba产生了更少噪声的表征，直接符合差异化机制背后的核心原理，该机制旨在缓解过度分配问题。调谐透镜在Diff-Mamba和Mamba上的线性投影均使用了与原始论文相同的配置进行训练。
4.5 中等规模下的初步结果
为了全面评估Diff-Mamba，作者从零开始训练了Mamba和Diff-Mamba模型，每个模型参数量为370M，使用The Pile数据集[13]中5000亿token的子集进行训练。经过消融研究（附录中的图10），发现最有效的Diff-Mamba变体是网络中反复交替Mamba和Diff-Mamba层，以实现可扩展性能。这种变体旨在结合两者的优势：其性能优于Mamba和Diff-Mamba，并受益于Diff-Mamba改进的鲁棒性、长上下文和检索能力。随后，作者在LongCraw164数据集[7]上进行了零样本测试，该数据集是RedPajama-v2[41]的一个长序列子集，专门用于长上下文建模研究。在图6(a)中，作者按照[24]的方法计算了数据集上的每token损失。此外，在图6(b)中，作者计算了The-Pile测试子集上的PPL。Diff-Mamba展现出令人印象深刻的长期上下文能力，在不同上下文长度下每token损失均维持在约9.98，而Mamba随着上下文增长显著增加。此外，Diff-Mamba在The Pile和PG19测试集上的PPL得分分别优于Mamba 0.131和1.445（图6）。作者认为中等规模下的积极结果具有前景，表明在Mamba和Diff-Mamba层之间交替能够产生更有效、可扩展和鲁棒的架构。
(b) Pile和PG19测试集上的PPL结果
(a) x轴为token索引，y轴为相应的每个token的损失。
unsetunset5 讨论：差分设计的意义unsetunset
从本文和Diff-Transformer的实证研究结果中引出的一个关键问题是，为什么差异设计如此有效，以及其强大能力的根本原因是什么。作者的研究通过首先表明过度分配注意力到无关上下文的问题并非Transformer独有的现象，而是架构研究中的普遍挑战，从而揭示了这一问题。
此外，虽然以往的研究主要通过降噪耳机和差分放大器的类比来推动差异设计，并以其强大的实证性能作为支持，但作者在此基础上更进一步。
在图5中，作者利用模型理解领域的工具进行了定量分析，表明Diff-Mamba模型中的中间表示相较于其非差异Mamba对应模型具有更高的信噪比，为文献中先前提出的动机提供了实证支持。
然而，差异设计实证成功的根本原因在很大程度上仍未得到探索，需要从理论和实证角度进行进一步研究，以推进这一重要方向的发展。
unsetunset参考unsetunset
[1]. Differential Mamba
点击上方卡片，关注「AI视界引擎」公众号