Boltzmann Attention：用伊辛模型给Transformer加上协同注意力，序列越长效果越好

传统注意力机制softmax/sigmoid本质上都是一个非相互作用的自旋系统——你关注的每一个位置都是独立的彼此之间没有可学习的耦合关系。这就构成了一个明显的表示瓶颈当序列变长位置间的协同依赖关系越来越重要时标准注意力没法直接建模。这篇论文提出Boltzmann Attention把每个注意力决策建模为伊辛模型中的自旋在QK相似度之外引入可学习的成对耦合让模型可以显式学习哪些位置应该一起被关注哪些位置应该互相竞争。实验表明这个改进随着序列长度增加而更明显而且可以通过量子退火进行线性扩展。核心问题J0的表示瓶颈先说问题本身。一篇注意力怎么运作的大家都很熟悉了对于query位置i对每个key位置j计算相似度h_ij然后softmax归一化得到注意力权重。在统计物理的语言里这就是一个局部场local fieldh_ij下的非相互作用自旋系统——每个位置相当于一个自旋它的“行为”是否被关注只受自身场的影响没有自旋之间的耦合J。这个J0的极限对sigmoid attention每个位置独立地通过σ(2h_ij)决定是否被关注完全成立对softmax attention也成立——虽然softmax的分母引入了某种“竞争”但这种竞争是归一化带来的全局约束不是通过学习得到的成对耦合关系。从物理角度看这相当于一个顺磁体paramagnet——每个自旋独立响应外加磁场没有相互作用带来的关联。而现实中序列数据里的相关性往往是结构化的关注一个主语会让关注其谓词更合理关注一个开括号意味着需要关注匹配的闭括号。这些是位置间的协同/对抗关系不是一个独立的场能够表达的。论文用一个形式化的表达阐明了这一点对于标准注意力∂a_ik/∂h_ij 0 当k ≠ j也就是说改变对位置j的关注程度不会直接影响对其他位置的关注——除非通过软最大化分母这种间接方式。这本质上是因为没有可学习的耦合J。随着序列长度增加这个瓶颈会加剧。因为位置对的数量是O(T²)量级的缺失的耦合自由度也随窗口大小平方增长。这也是为什么论文的实验中选择从小窗口开始做可控的验证——他们要排除采样噪音纯粹看耦合J带来的表示增益。方法核心从独立自旋到伊辛模型Boltzmann Attention的做法很直接把注意力决策建模为伊辛模型。每个key位置对应一个二值自旋s_j ∈ {-1, 1}1表示“关注”-1表示“忽略”。能量函数为E_i(s) -∑j h_ij s_j - ∑J_jk s_j s_k其中h_ij来自标准的QK相似度相当于局部场J_jk是可学习的成对耦合参数并且在同一个head内、同一批次中共享。注意力权重从自旋的边际磁化计算a_ij (⟨s_j⟩_i 1) / 2。这里有几个关键点需要理解第一耦合J的作用机制。当J_jk 0铁磁耦合时关注位置j会增加关注位置k的概率也就是学出了“协同关注”。当J_jk 0反铁磁耦合时关注位置j会抑制对位置k的关注产生“竞争性关注”。这意味着一个位置可以因为“和另一个被关注的位置有关联”而获得关注而不仅仅是因为自己的查询-键相似度高。这在结构化的任务里非常自然比如括号匹配中第一个闭括号匹配到第一个开括号后第二个闭括号自然应该匹配第二个开括号——这本质上是一个协同约束。第二模型的结构。论文中使用的是单层、单头的Transformerd_model64d_ff128。使用单头的原因很明确排除多头注意力中“不同头学习不同模式”的解释把所有改进都归因于耦合J。这也让实验设计更加干净——不是多头带来的多样性而是J带来的相关性。第三推理的困难。引入耦合J后能量不再可分解为独立自旋的乘积需要计算配分函数Z_i这意味着对2^T个自旋配置求和。在小窗口T ≤ 16下可以用精确枚举但这显然不能扩展到实际序列长度。论文在后面用DQA非绝热量子退火给出了扩展路径这个后面再细说。实验验证改进来自耦合不是参数化实验设计得很扎实分成两个任务Tiny Shakespeare字符级语言建模从T4到T14对比四个模式 - softmax标准注意力 - hJ完整伊辛模型即Boltzmann Attention - h-onlyJ0等价于sigmoid attention - J-onlyh0纯结构先验没有数据相关的局部场结果很干净 - 在T4时hJ和softmax持平-0.02%可以看作随机变异 - 从T6开始hJ超越softmax优势随T增大0.43% → 0.81% → 0.92% → 1.08%T12→ 1.04%T14 - h-onlysigmoid在T≥8时反而比softmax差−3.08%T8、−7.14%T12 - J-only更差T14时−14.73%这个模式传递的信息很明确 1.不是sigmoid参数化带来的好处——h-only是sigmoid但是表现不如softmax 2.只有hJ的组合才有效——数据相关的局部场和可学习的耦合缺一不可 3.优势随序列长度增大——因为J提供了O(T²)规模的额外自由度窗口越大收益越大论文还分析了学习到的耦合矩阵J的结构在T12时临近位置距离2-4倾向于铁磁耦合正相关较远位置距离≥6倾向于反铁磁耦合负相关。这个距离依赖的模式有多种理解方式可以看作是模型在学习局部的“词组”结构附近单词常常共出现也可以看作是长距离位置间存在某种竞争关系。这个模式在不同随机种子间是稳定的。括号匹配这是一个更直接测试协同注意力的合成任务给定包含Dyck-1词正确嵌套的括号对和随机填充符的序列对于每个闭括号模型必须预测其匹配的开括号位置。为什么这个任务特别适合因为括号匹配天然要求位置间的协同决策第一个闭括号匹配第一个开括号第二个闭括号匹配第二个开括号——每个闭括号的匹配目标取决于其他闭括号已经占用了哪些开括号。如果不引入位置间的耦合仅仅在每个位置独立看QK相似度这个约束是无法直接表达的。结果再次支持上述模式 - 在T16时hJ比softmax高2.89个百分点95.97% vs 93.08% - h-only在T16时下降至90.06%比softmax差 - J-only在T16时下降至87.42%论文还做了一个额外的消融实验去掉FFN只保留注意力层和输出头。这时候差距更大T16时hJ比softmax高4.53pp89.65% vs 85.13%。这说明FFN可以部分补偿J0的瓶颈通过非线性变换在逐位置层面上学习某种隐式的结构但无法完全替代耦合J的作用——因为FFN依然是逐位置操作不会直接创造位置间的相关性。扩展路径量子退火与线性缩放精确枚举的开销是O(2^T)显然不能用于实际应用。论文展示了DQA作为替代方案的可行性。DQA的基本思路是通过一个绝热量子退火过程的非绝热版本在短时间内产生近似的玻尔兹曼样本。论文用Trotter分解模拟了量子电路200个Trotter步退火时间5.0ns并用直通估计器STE通过采样步骤传播梯度。结果是在T4、6、8时DQA训练的性能与精确枚举相当。莎士比亚任务的两个指标差距很小括号匹配的DQA训练甚至达到了100%准确率。这里需要理解DQA的意义它不是要在小窗口下匹配精确枚举而是提供一条线性缩放的路径。在量子硬件上每个query位置只需要一次O(1)时长的退火总开销是O(T)——而精确枚举是O(2^T)。当然现在这个验证还是基于经典模拟的Trotter化量子电路真正的量子硬件实现还需要额外的工作。除了DQA论文也提到经典方法如并行回火parallel tempering可以作为替代路径利用GPU并行进行相关采样。扩展讨论更高阶相互作用和其他观察论文还探索了三体耦合three-body interactions用K表示。结论很有意思 - 在莎士比亚任务上三体耦合没有超过成对耦合hJ的表现 - 但在括号匹配任务上hJK或单独hK在T≥12时显著优于hJ4.93pp at T12这说明高阶耦合对结构化、有组合逻辑的任务如括号匹配是有价值的但对自然语言建模这种统计性任务来说成对耦合可能已经足够捕捉大部分相关结构。这个差异本身就是一个值得注意的发现不同的任务类型可能需要不同阶数的相互作用。论文还做了一个有趣的分析把莎士比亚的困惑度曲线拟合为PPL(T) ae^(-bT) c发现衰减率b在各模型间相近约0.28-0.30但渐近下限c不同。Boltzmann变体的c值5.44-5.45低于softmax的c值5.51更低的是h-only的5.99。这意味着J的作用是降低可达到的性能下限而不是改变衰减速率——也就是在数据更多、窗口更大的情况下耦合依然能带来持续的优势。一些边界和局限需要指出这个工作目前还在小范围可控实验阶段 - 单层单头Transformerd_model64窗口T≤16 - 精确枚举的计算成本限制了实验规模 - DQA验证还是基于经典模拟不是真的量子硬件但从研究价值来看指认J0这个瓶颈并提供一个干净的解决方案本身就是有意义的。不论是通过量子退火、并行回火还是其他采样方法来做扩展核心贡献是识别了“注意力机制中没有位置间可学习耦合”这个结构性问题。论文还提到了一个有趣的边缘情况J-onlyh0的配置虽然性能远低于完整模型但它绕过了QK计算——注意力模式仅由可学习耦合J决定不再需要输入相关的局部场。这虽然听起来像是限制很大但有一种场景可能有用如果任务中的位置关系主要是结构性的比如某种固定的语法模板J-only可以作为快速的先验。而且对于量子退火来说J-only的训练可以做到完全免梯度——梯度估计只需要通过相关性匹配∂L/∂J_jl ∝ ⟨s_j s_l⟩_data - ⟨s_j s_l⟩_model不需要反向传播通过注意力机制。我的看法这篇文章的价值在于把一个常见但未被足够重视的问题明确指出来了我们用了这么多年softmax/sigmoid注意力它们之间有很多数学和实证上的差异但从表示能力的角度看它们都一样——J0。这个角度看问题是有启发的。另一个值得思考的点是论文说“FFN无法完全补偿J0瓶颈”从实验数据看确实如此没有FFN时差距更大。但这更多是一个关于去哪里、如何做扩展的判断如果注意力层本身就能直接建模位置间协同关系就不需要堆更多FFN来间接补偿。对于长序列场景这个设计点可能值得注意。最后耦合矩阵J的可解释性本身也是一个有意思的方向。学习到的J在可视化后呈现出距离依赖的结构这种结构不是被强加的而是从数据中学出的。如果用于更大规模的场景耦合矩阵可能给出关于数据结构的直观理解——比如哪些位置“倾向于一起被关注”这其实等价于学习了一个位置间的共现图。这个方向的解读空间很大。原始链接论文https://arxiv.org/abs/2606.12478 代码HuggingFace上的相关实现/演示https://huggingface.co/huggingface https://huggingface.co/docs/hub/spaces