Transformer模型在各种领域展现出非凡的性能，成为解决序列机器学习问题的最先进方案。尽管我们对Transformer架构的基本组件有了大致的了解，但对其运作机制和预期动态知之甚少。近年来，人们越来越关注注意力机制和Hopfield网络之间的关系，这有望揭示Transformer网络的统计物理学。然而，迄今为止，类似Transformer模型的动态机制尚未得到深入研究。
本文通过使用非平衡态下不对称Hopfield网络的研究方法，填补了这一空白。具体来说，我们采用了生成函数上的路径积分方法，得出了由并发平均场变量控制的动力学。假设1比特的token和权重，我们推导出大型自注意力神经网络与softmax输出耦合的解析近似，在网络规模趋于无穷大时，该近似将变得精确。我们的发现揭示了非平凡的动态现象，包括与混沌分岔相关的非平衡相变，即使对于具有少量编码特征和非常短的上下文窗口的简单配置也是如此。最后，我们讨论了我们的分析方法在提高我们对Transformer模型内部运作的理解方面的潜力，这可能有助于降低计算训练成本并提高模型的可解释性。
1. Transformer与Hopfield网络的奇妙联系
Hopfield网络描述了一个系统，其中一组自旋xi（i∈{1, .., N}）的概率p(x)由一个能量函数定义，该函数考虑了系统中存储的M个记忆ξa = {ξa,0, .., ξi,a, .., ξa,N}（a∈{1, .., M}，ξi,a∈{−1, +1}）。一个状态的概率定义为：
p(x) = Z−1 exp (β/N ∑a ∑i<j xiξi,aξj,axj)
其中，Z = ∑x exp(β∑a ∑ij xiξi,aξj,axj) 是配分函数，耦合通常是对称的（神经元i和j之间的耦合值为∑a ξi,aξj,a），描述了一个能量景观E，在这个能量景观的最小值处，我们可以恢复系统的记忆。参数β是一个常数，定义了逆温度。最后，1/N的归一化确保了系统的能量是广泛的。
类似于受限玻尔兹曼机，我们可以通过定义两组变量x = {x1,t, . . . , xN,t}（t∈{1, 2, . . . , T}）和一组由矩阵Wk, Wq（k, q∈{1, 2, . . . , M}）编码的模式来描述一个二分Hopfield网络。该系统的概率定义为：
p(k, q) = Z−1 exp (β/N ∑a ∑ij kiWk i,aWq j,aqj)
请注意，现在k, q之间的连接是不对称的（即，可能Wk a ̸= Wq a）。对于给定的q，系统的条件分布可以表示为：
p(k|q) = exp(β/N ∑a ∑ij kiWk i,aWq j,aqj) / ∑k′ exp(β/N ∑a ∑ij k′ iWk i,aWq j,aqj)
当输出k再次被馈送到输入q时，对称随机耦合会导致自旋玻璃行为，而不对称随机耦合会导致非平衡稳态，表现出有序-无序相变和混沌动力学。
如果我们仔细观察[16, 31]中定义的单个头的自注意力函数At = softmaxτ (1/U ∑τ ∑i (qtWq)⊺W kkt−τ )τ W vvt−τ ，我们可以观察到查询键交互矩阵与上面定义的Hopfield网络具有类似的形式。引入一个缩放参数γ（类似于上面的逆温度β，但我们将其保留用于后面的输出），我们可以描述：
pA(kτ |qt) = softmaxτ (γ/U ∑a ∑ij kj,t−τ Wk i,aqi,t)
其中，L是上下文窗口的长度，a∈{1, .., M}定义了编码的特征，如公式(1)所示。重要的是，在自注意力中，每个键、查询和值token都等于不同时间步xt的输入，因此kt = qt = vt = xt，如公式(8)所示（相反，在交叉注意力中，查询token可以取不同输入的值）。请注意，我们选择用配分函数来表示softmax，得到与公式(3)中不对称Hopfield网络的条件分布等效的形式。我们注意到，即使在向量x不限于二进制值的情况下，等效性仍然有效。重要的是，注意力与二分Hopfield网络相比有两个重要的区别。首先，配分函数的求和不是对所有查询q进行，而是只对之前时间步t−τ′出现的查询进行，如公式(5)所示。其次，点积被归一化，以便softmax函数始终处于梯度相对较大的区域，假设∑i ki,tWk i,a项在a上的方差约为1阶[31]。通常，单位方差是通过专门的归一化层和初始化权重来实现的，以保持方差。对于大量的特征，softmax使用归一化项U = √M。在我们的例子中，为了简单起见，我们省略了归一化层，并将权重归一化包含在我们的归一化常数中，得到U = N^2√M。
2. 简化的注意力-输出层
通常，Transformer网络中的输出被定义为一个softmax。我们的目标主要是设计自注意力层能够产生的各种序列行为。因此，我们选择将一个高度简化的Transformer网络表示为注意力层和softmax输出的组合。类似的具有非线性输出的单层注意力网络可以完全记忆有限样本，并且是连续函数的通用逼近器。为了专注于注意力的动态，我们忽略了加法、归一化和前馈网络块，并将softmax直接应用于定义为注意力值线性组合的对数，如p(ot|At) = softmaxot(A t ⊺W oot)。为了进一步简化模型，我们假设输出对数是通过注意力值的线性变换Wo获得的，即o = WoAt。那么：
p(ot|At) = softmaxot(β/N ∑a ∑i i,aAa t) = exp(β/N ∑i oi,tWo i,aAa t) / ∑o′ exp(β/N ∑i o′ i,tWo i,aAa t)
其中，β = T−1是逆温度，1/N的归一化确保了对数的单位方差。
为了生成Transformer的动态，我们定义了一系列状态x0:T = {x1,t, . . . , xN,t}（t∈{0, 1, . . . , T}），其中xt = {x1,t, . . . , xN,t}。然后我们定义模型的动态，从公式((4), (5))生成注意力值，其中vτ = kτ，然后使用公式(6)以p(xt+1|At)的形式自回归地生成下一个token xt+1 = ot，得到：
p(x0:T ) = ∏t=0T−1 ∑τ pA(xτ |xt) p(xt+1|At)
其中：
Aa t = ∑i Wv i,axi,t−τ
3. 1比特token编码的注意力层的动态平均场理论
在本节中，我们使用标准方法来研究简化Transformer的循环动态，这些方法用于检查循环网络的非平衡统计物理学[5]。为了简单起见，并且为了保持对Hopfield模型的描述，我们假设权重和token的二进制编码。二进制权重Transformer正迅速成为一种有吸引力的替代方案，可以实现更具成本效益的模型。二进制token不太常见，但已被探索以降低计算成本（尽管有效的缩放似乎是保持性能的关键因素）。无论如何，我们的框架可以很容易地扩展到离散的非二进制token集（例如，用softmax替换我们结果中的tanh函数）。
由于模式之间的相互作用不像公式(1)那样是对称的，我们没有一个解析表达式来计算系统的矩。相反，我们定义了一个生成函数，它充当矩生成函数，在非平衡统计力学中扮演与配分函数等效的角色，但适用于非平衡设置。生成函数定义为：
Z(g) = ∑x0:T p(x0:T ) exp(∑a,t ∑α gα a,t 1/N ∑i Wα i,axi,t)
其中，α∈{q, k, v, o}是查询、键、值和输出变量特征的索引，路径概率定义在公式(7)中。请注意：
∂Z(g)/∂gα a,t |g=0 = 1/N ∑i Wα i,a ⟨xi,t⟩
恢复了Transformer中编码的特征的统计信息。因此，生成函数充当描述非平衡统计物理学中过程的一种动态配分函数。
我们使用路径积分方法[5]来求解系统，引入平均场变量：
mα a,t = 1/N ∑i Wα i,a ⟨xi,t⟩
在网络规模趋于无穷大时，使用最速下降法计算公式(9)（附录A）。在1比特token编码的情况下，这将导致一个生成函数：
Z(g) = exp(∑t=0T ∑i log 2 cosh(∑a Wo i,aβ ˆAa,t−1 + 1/N ∑a,α Wα i,agα a,t−1))
由以下平均场变量方程描述：
mα a,t = 1/N ∑i Wα i,a tanh(β ∑b Wo i,b ˆAb,t−1)
ˆAa t = ∑τ mv b,t−τ exp(γ ∑a mq a,tmk a,t−τ) / ∑τ ′ exp(γ ∑a mq a,tmk a,t−τ ′)
其中，ˆAa t是使用平均场变量评估的归一化注意力值（即除以N）。这些方程在网络规模趋于无穷大时是精确的，尽管我们应该期望在较小的规模下出现热涨落。
我们可以观察到，对于所有平均场mα，解都采用相同的形式。在公式(13)中，我们看到，对于所有模式b，自旋xi,t+1的平均行为是在tanh函数内计算的，然后投影到与每个与模式a相关的矩阵Wα a相对应的平均场。所有自旋i的信息被收集起来，以获得模式a的总行为。然而，我们可能会注意到，这样的网络不包含关于token排序的任何信息。Transformer模型通过添加一个位置编码的形式来缓解这个问题，该位置编码是一个外部信号[31]。为此，我们将携带位置编码的NP个单元添加到我们的token中（即语义嵌入）。在这个嵌入中，我们简单地将关于时间步t选择的token位置的信息编码为一个比特数组。我们将位置token单元定义为pi,t = (1)⌈t/i⌉（其中⌈x⌉是一个向上取整运算符），是时间序列值t的二进制编码中的第i个比特。
mα a,t+1 =(1−ϵ) 1/N ∑i=1N Wα i,a tanh(β 1/U ∑b Wα i,b ˆAb,t) + ϵ 1/NP ∑i=1NP Wα i,api,t+1
其中，ϵ决定了位置编码的相对权重。上面的方程中的位置编码只是添加到由softmax输出生成的token（或平均场）中。请注意，公式(13)中的解仍然是正确的，但我们只是通过添加一个外部信号pt来重新计算它的值。与Transformer网络中通常的做法一样，位置嵌入和语义嵌入通过同一个矩阵Wα进行投影。
公式(15)中的平均场方程不能直接在网络规模无限大的情况下计算。然而，对于在W中编码的有限数量的模式，其值可以取±1，我们可以用模式对之间的相关值Wα a , Wα′ b来代替对i的求和，如附录B中的公式(29)所示，得到：
mα a,t = 1/2M (∑b σb ⟨W o i,bWα i,a⟩i + ∑b<c<d σbσcσd tanh(β ∑b σb ˆAb,t−1) ⟨W o i,bW o i,cW o i,dWα i,a⟩i + · · ·)
其中，σ是一个包含M个二进制变量σa = ±1的数组。省略号包括sigma分量的奇数乘积乘以Wα i,a与W o i分量的偶数乘积在i上的平均值。⟨…⟩i表示对i索引的平均值。
4. 结果
为了观察简化Transformer网络的动态行为机制，我们模拟了具有随机二进制权重和随机相关值的网络（附录C），并使用了一个长度为L = 4个token的上下文窗口。我们对网络进行了1.2×10^6步的模拟，其中每一步都将轨迹的最后4个token作为输入（如上下文窗口定义），并生成下一个token。为了避免瞬态轨迹，我们丢弃了前10^6步。在接下来的部分，我们将展示其中一个网络初始化的行为。即使是参数的一种配置也允许观察到非常不同的行为，并且类似的结果可以在其他组合中复制。下一节中获得的结果是在手动探索后使用γ = 220和ϵ = 0.02获得的。对于其他参数也观察到了类似的动态，但我们选择这种组合，因为它在不同的β下产生了更大的动态多样性。随机生成相关性的种子是手动选择的，设置相关性的过程在附录C中解释。
实验是在几个2个英特尔至强E5-2683 @ 2.10GHz节点上进行的。每个模拟过程的执行（具有上面提到的设置）是在数十个节点上并行进行的。每个计算大约需要5分钟才能完成，每个过程使用2个专用核心。我们为每个过程分配了8GB的内存。
4.1 非平衡相变
在本节中，我们探索了模型在不同逆温度β = T−1下的行为。在Transformer中，下一个token的概率取决于一个温度值，该温度值会修改softmax概率，如公式(6)所示。我们使用4001个值将温度值从0修改到3。这类似于在标准Transformer架构中探索softmax输出的不同温度。
我们以两种形式生成系统的分岔图。通常，分岔图[27]描绘了某个变量在特定参数（在本例中为β）下的所有取值。在分岔图的第一种形式中，我们简单地从一个平均场变量mo 1,t（其他变量的行为类似）的稳态中采样20000个点，并将每个β的取值表示为图3中的黑色和黄色点，具体取决于吸引子是周期性的还是非周期性的。在分岔图的第二种形式中，如果吸引子是非周期性的，我们绘制轨迹中与平面mo 2 = 0（在0.001的小误差值内）相交的点，这些点在图3中用橙色和紫色点表示，具体取决于吸引子是准周期性的还是混沌的。虽然在未来的研究中我们将系统地研究吸引子，但在这里我们只通过目视检查和对每个分岔图中点的数量进行计数来区分它们。在图3.a和图3.b中，我们可以看到β值在范围[0, 3]和[1.24, 1.28]内的分岔图。在较低的β处，系统陷入周期性吸引子，然后发展到一个具有准周期性行为的区域，最终突然过渡到混沌状态。我们可以看到，尽管如此，不同的状态并不稳定，β的微小变化会对所表现出的行为产生很大的影响。
为了可视化这一点，我们从分岔图中选择了β值，并在同一个平面上绘制了mo 2(t)的轨迹。在这里，我们可以区分：1) 在一组固定点之间跳跃的周期性轨迹（例如，图4中的β = 1.27）；2) 在一组连续曲线上的点之间跳跃的准周期性循环（例如，图4中的β = 1.255和β = 1.26405）；或者3) 混沌轨迹（例如，图4中的β = 1.266，β = 1.28和β = 1.4）。
4.2 动态和时间记忆
在图5（顶部）中，我们观察到一个准周期性和一个混沌轨迹的示例。观察它们的傅里叶频谱和自相关函数（图5，中间和底部），我们看到所有信号都具有与低频分量相关的长相关性，特别是在混沌状态的情况下。在准周期性轨迹的情况下，我们观察到的低频分量比与上下文长度相关的低频分量（f = 1/L = 0.25）更低。这揭示了注意力和Transformer模型中经常被忽视的一方面，即它们的记忆容量并不局限于上下文窗口的跨度，而是系统可以显示出丰富的动态机制，从而诱导低频分量，显著提高模型“记忆”先前token的能力。
此外，准周期性和特别是混沌轨迹——对于语义表示，其中相关的token在嵌入空间中彼此靠近——可以被解释为吸引子，允许系统以不同的方式表达类似的结构。此外，图5（右上角）所示的不同吸引子的共存表明，模型可以由混沌动态驱动，在不同的结构之间切换。总的来说，这里展示的简单示例表明，类似Transformer的架构可以轻松地产生具有非平凡记忆效应的丰富动态结构。
5. 总结
本文提出了一种针对注意力网络的动态平均场理论，并展示了它在一个简单的模型中的应用，该模型结合了注意力网络和softmax输出token生成器。
我们使用非平衡统计力学技术来研究注意力机制作为不对称Hopfield网络。我们的计算在网络规模趋于无穷大时，得到了路径概率统计的精确结果。虽然为了简单起见，我们使用了权重和token的1比特描述，但结果可以轻松地扩展到其他配置。
即使对于相对简单的配置，包括每个级别（键、查询、值和输出）只有三个特征，以及非常有限的上下文（4个token），我们也发现了丰富的动态行为景观，包括平均场变量之间的周期性、准周期性和混沌吸引子之间的多个相变。此外，在所有情况下，我们都观察到复杂的动态，其记忆效应远远超出了上下文窗口中存储的内容。虽然我们没有考虑Transformer网络中的其他块，例如前馈层和残差连接，但我们预计额外的非线性和反馈回路会增加观察到的动态的复杂性。
尽管我们的示例仅限于简化的模型配置，但我们的方法有可能为更现实的设置中Transformer模型的动态提供见解。理解Transformer网络一般配置中预期的相图，可以更深入地理解架构，以及表征参数空间中的理想配置和相关的相变。此外，我们的框架可以扩展到实现技术，以理解记忆容量和动态之间的关系，即非平衡态下自旋玻璃的著名平均场理论的非平衡等效物。
此外，我们假设我们的平均场计算可以为计算Transformer模型中注意力层的统计信息提供一种经济高效的替代方案，这可能通过减轻学习过程中计算损失函数梯度的计算成本，来加速模型训练过程。尽管在平均场变量的计算中假设了无限大的系统，但我们的方法可以针对有限大小的网络进行改进，以获得更准确的近似，例如使用非平衡自旋模型中的理论。我们希望在未来的工作中探索这个想法。
我们渴望为提高Transformer模型的可解释性而开发的方法做出贡献。大型语言模型中特征的可解释性以及识别系统关键特征中涉及的回路，引起了极大的兴趣。通过非平衡统计物理学的视角，我们可以通过描述一组低维平均场变量来阐明大型异构网络的动态。理解这些序参量会导致识别模型预测中的关键特征和模式。我们推测，理解这些关键特征中的相变可以提供与对齐问题等挑战相关的关键见解。
总而言之，这些结果有助于发展注意力机制及其相变的非平衡统计物理学的精确分析理论。我们设想，沿着这条线的贡献有可能促进将统计物理学的理论概念更系统地整合到生成式AI模型领域。
致谢和资金披露
感谢Ivan Garashchuk和Iñigo Urteaga对本文稿的宝贵意见。APL和MA获得了“la Caixa”基金会（ID 100010434，代码LCF/BQ/PI23/11970024）的MA青年领袖奖学金资助。APL和MA感谢巴斯克政府的BERC 2022-2025计划和西班牙科学与创新部的塞维罗·奥乔亚认证（CEX2021-001142-S / MICIN/AEI/10.13039/501100011033）的支持。MA部分获得了约翰·坦普尔顿基金会（资助ID 62828）、巴斯克政府（ELKARTEK 2023计划，项目KK-2023/00085）和西班牙科学、创新和大学部的资助MICIU/AEI /10.13039/501100011033的支持。
参考文献
[1] https://arxiv.org/pdf/2406.07247