ICLR 2025 Oral｜差分注意力机制引领变革，DIFF Transformer攻克长序列建模难题
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
近年来，Transformer架构在自然语言处理领域取得了巨大成功，从机器翻译到文本生成，其强大的建模能力为语言理解与生成带来了前所未有的突破。
然而，随着模型规模的不断扩大和应用场景的日益复杂，传统Transformer架构逐渐暴露出缺陷，尤其是在处理长文本、关键信息检索以及对抗幻觉等任务时，Transformer常常因过度关注无关上下文而陷入困境，导致模型表现受限。
为攻克这一难题，来自微软和清华的研究团队提出了DIFFTransformer，一种基于差分注意力机制的创新基础模型架构。
论文标题：DifferentialTransformer
论文链接：https ://openreview.net/pdf?id=OvoCm1gGhN
代码链接：https ://aka.ms/Diff-Transformer
其核心思想是通过计算两组Softmax注意力图的差值来放大对关键上下文的关注，同时消除注意力噪声干扰。DIFFTransformer具备以下显著优势：
在语言建模任务中，DIFFTransformer在模型大小、训练token数量等方面展现出了卓越的可扩展性，仅需约65%的模型规模或训练token数量即可达到与传统Transformer相当的性能，大幅提升了语言模型通用表现。
在长文本建模、关键信息检索、数学推理、对抗幻觉、上下文学习、模型激活值量化等一系列任务中，DIFFTransformer展现了独特优势，相比传统Transformer有显著提升。
DIFFTransformer的特性使其在自然语言处理领域具有广阔的应用前景，有望成为推动语言模型发展的新动力。此外，已有跟进研究初步验证方法在视觉、多模态等领域中的有效性，显示出其跨模态通用的潜力。该研究已被ICLR2025接收，并获选为Oral论文（入选比例1.8%）。
方法
本文提出了一种名为DifferentialTransformer（DIFFTransformer）的基础模型架构，旨在解决传统Transformer在长文本建模中对无关上下文过度分配注意力的问题。该方法通过差分注意力机制（DifferentialAttention）放大对关键上下文的关注，同时消除注意力噪声，从而显著提升模型在多种任务中的性能。
差分注意力机制
传统Transformer的注意力机制通过Softmax函数对输入序列中的不同token进行加权，但Softmax的性质导致模型难以完全消除无关上下文的影响。为了克服这一问题，DIFFTransformer引入了差分注意力机制。
具体而言，该机制将查询向量（Query）和键向量（Key）在注意力头（Head）维度分为两组，分别计算两组的Softmax注意力图，然后计算两者的差值作为最终的注意力分数。这一设计类似于电子工程中的差分放大器，以及降噪耳机，通过两组信号相减以消除共有噪声。
差分注意力的数学表达如下：
其中，和分别是两组查询和键向量，是值向量，是一个可学习的标量参数，用于调节两组注意力图的权重。计算过程如图1所示。
图1.差分注意力机制图示与伪代码
为了同步学习速率，将重参数化为：
其中，是可学习的向量，而是用于初始化的常数。
多头差分注意力
为了进一步提升模型的表达能力，DIFFTransformer采用了多头机制。每个注意力头独立计算差分注意力，并将多头输出拼接为最终结果。具体实现如下：
其中是注意力头的数量，是输出投影矩阵。为了保持与Transformer梯度一致，DIFFTransformer在每个头的输出后应用了独立的归一化层，采用RMSNorm实现。
图2.Transformer与DIFFTransformer注意力分数分布可视化
图2展示了DIFFTransformer和传统Transformer在注意力分数分配上的显著差异。作者将一段关键信息插入大段不相关文本的中间位置，并对模型抽取关键信息时的注意力分数分配进行可视化。
传统Transformer的注意力分数被广泛分配到整个上下文中，只有极少分数分配至关键信息；而DIFFTransformer能够将更高的分数集中在目标答案上，并且几乎不向无关上下文分配注意力。
注意力分数分配的稀疏性与精准性也使得DIFFTransformer在处理长文本关键信息检索任务时显著优于Transformer。
实验
作者通过一系列实验验证了DIFFTransformer在多个方面的卓越性能，证明了其在大语言模型中应用的独特潜力与优势。
语言建模
作者研究了DIFFTransformer在扩展模型规模和训练数据量时的性能，如图3所示。实验表明，DIFFTransformer仅需约65%的参数规模或训练数据量即可达到与Transformer相当的语言建模性能。例如，6.8B参数规模的DIFFTransformer在语言建模损失上与11B参数规模的Transformer相当。
图3.语言建模上的模型参数、训练数据量可扩展性实验
长文本建模
作者将模型扩展到64K上下文长度，并在长文本书籍数据上进行了评估。结果显示，考虑累积平均负对数似然（NLL）指标，DIFFTransformer在不同序列位置上均优于Transformer，能够更有效地利用长上下文信息。
图4.长文本书籍数据模型性能评估
关键信息检索
作者通过「多针检索」（Multi-NeedleRetrieval）实验评估了模型从大量上下文中提取关键信息的能力，如图5所示。实验表明，DIFFTransformer在不同上下文长度和答案深度下均表现出更高的准确率，尤其是在文本较长以及答案位于文本更靠前位置时，优势更为明显。例如，在64K上下文中，DIFFTransformer在答案位于25%深度时的准确率比Transformer高出76%。此外，统计信息显示，DIFFTransformer在注意力分数分配上也表现出更高的聚焦能力，能够准确定位关键信息，并展现了更高的信噪比。
图5.多针检索评估
上下文学习
作者从两个角度评估了DIFFTransformer的上下文学习能力：多样本上下文学习和样本顺序鲁棒性测试。如图6所示，在多样本上下文学习任务中，作者使用了4个不同的数据集（TREC、TREC-fine、Banking-77和Clinic-150），并逐步增加示例数量，直到总长度达到64Ktokens。结果显示，DIFFTransformer在不同数据集上均优于Transformer，平均准确率提升显著。
图6.多样本上下文学习
在鲁棒性测试中，作者通过打乱示例顺序的方式评估了模型的性能稳定性。如图7所示，DIFFTransformer在不同示例排列下的性能方差显著低于Transformer，表明其对输入顺序的敏感性更低，具有更强的鲁棒性。
图7.样本顺序鲁棒性测试
幻觉评测
作者利用文本摘要和问答任务作为两个典型的幻觉评测场景，评估了DIFFTransformer在降低大模型幻觉（hallucination）方面的表现。结果如图8所示，DIFFTransformer在生成摘要和回答问题时显著提升了准确率，减少了幻觉现象。这是因为差分注意力机制能够准确定位重要文段，避免无关上下文对模型预测的干扰。
图8.利用文本摘要、问答任务进行幻觉评测
异常激活值分析
图9.注意力激活值的低比特量化
数学推理能力
作者在数学推理任务上进一步验证了DIFFTransformer的性能。作者采用两阶段训练，在3B预训练模型的基础上进行有监督微调，并在MATH等8个数学数据集上评测模型性能。在第一阶段，采用20Btoken合成数学数据对模型进行微调，使模型获得基础数学能力，评测结果如图10所示。从15Btoken开始，DIFFTransformer展现出了显著优于Transformer的数学能力，至20Btoken结束的时候，准确率的差距达到了11%左右。
图10.第一阶段数学合成数据微调
在第二阶段，作者利用Deepseek-R1输出所构造的数据集OpenThoughts-114K-Math对模型进行蒸馏，使模型更强大的深度推理能力。如图11所示，在8个数据集上，DIFFTransformer相较Transformer均有不同程度的提升，平均准确率提升了7.5%，这表明差分注意力机制更强大的上下文建模能力在推理任务中也至关重要。
图11.第二阶段深度推理能力评测
讨论与未来工作
HuggingFace：https ://huggingface.co/papers/2410.05258
alphaXiv：https ://www.alphaxiv.org/abs/2410.05258v1
PetarVeličković：https ://x.com/PetarV_93/status/1874820028975267866
LucasBeyer：https ://x.com/giffmana/status/1873869654252544079
未来工作方面，作者认为可以利用DIFFTransformer的性质设计低比特注意力算子，以及利用差分注意力的稀疏特性进行键值缓存（key-valuecache）的剪枝。此外，将DIFFTransformer应用在除语言以外的其他模态上也值得探索。近期工作DiffCLIP将差分注意力扩展至视觉、多模态领域，揭示了DIFFTransformer在不同模态任务中的更多结构特性与应用潜力。
DiffCLIP：https ://arxiv.org/abs/2503.06626
总结
本文的贡献主要在两个方面：
（1）DIFFTransformer通过创新的差分注意力机制，有效解决了传统Transformer在处理文本时受到噪声干扰、注意力分配不准确的问题；
（2）凭借对关键信息的关注和对噪声的抵御能力，DIFFTransformer在语言建模、长文本建模、关键信息检索、数学推理、对抗幻觉、上下文学习、模型激活值量化等任务中表现出色，有望在自然语言处理、多模态等领域作为基础模型架构。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
