浅谈Transformer

1、简介

在这篇文章中，我们试图从一系列较长的，关于将Transformer理解为物理系统的文章里提出的想法中，提炼出一个统一的观点：

• 深度隐式注意力：注意力机制的平均场理论视角

• 自旋模型的Transformer：近似自由能最小化

我们主张，原型Transformer模块的神经网络架构蓝图可以从经典统计力学中熟悉的物理自旋系统的结构中导出。更具体地说，我们认为Transformer模块的正向传递，可映射为矢量自旋模型中的计算磁化，作为对输入数据的响应。我们将Transformer想象成可微自旋系统的集合，其行为可以通过训练来塑造。

2、Transformer架构从何而来？

纵观自然语言处理和计算视觉中不断增长的 Transformer 架构可以看出，Attention is All You Need（Vaswani 等人，2017）1中引入的设计模式仍然占主导地位。过去四年发布的几乎所有 transformer 模块的架构变体都坚持了一个成功组合：残差连接、类注意力的操作（标记混合）、归一化层和类前馈的操作（通道混合） .

最近的工作，如MetaFormer is Actually What You Need for Vision (Yu et al., 2021) 2适当地将焦点转移到 transformer 模块的高级架构上，并认为其完整结构，而不仅仅是标记混合注意力操作，对Transformer实现具有竞争力的性能至关重要。

那么这个原型设计模式从何而来呢？为什么它似乎一直存在？它的结构背后有什么物理直觉吗？

3. 从能量函数中获得注意力只会让你到此为止

Hopfield Networks is All You Need（Ramsauer 等人，2020 年）3和Large Associative Memory Problem in Neurobiology and Machine Learning（Krotov 和 Hopfield，2020 年）4等近期论文使用基于能量的视角寻找注意力机制背后的物理直觉，用现代连续 Hopfield 网络来表达。主要思路是推导softmax-attention更新规则

(1)

通过使用一些精心选择的能量函数，对输入查询Q的导数进行大的梯度下降更新步骤

(2)

这样，可以采用大梯度步长对普通的softmax attention进行重构。方程式(2)定义的能量形式，实现了一个关联记忆系统，用于存储和检索矢量模式，其中查询流向与其最近的键相关联的谷（请参阅注意力作为能量最小化：可视化能量形式）：

但Transformer模块不仅仅是注意力。在实践中，我们知道残差连接、归一化层和前馈层对于获得良好的实际性能都是必不可少的。

我们能否泛化这种对能量函数求导数以恢复整个Transformer模块的物理直觉？我们可以。但我们必须从能量函数上退后一步，转而关注它们的底层物理系统。

4. 回归本源：物理自旋系统和矢量自旋模型

经典统计力学中的能量函数是对物理系统中的相互作用和约束进行编码的简洁描述。自旋系统是典型的物理系统，通常用作各种现象的玩具模型5。

伊辛模型是一个简单的玩具模型，它描述了一个经典的二元自旋系统，每个位置都具有局部自旋自由度，要么向上要么向下。存在依赖于位置的外部磁场的情况下，N个自旋的二元随机伊辛模型的能量函数由下式给出

(3)

其中Jij矩阵对所有自旋对之间的耦合强度进行编码，而外部磁场hi通过在每个位置提供优先对齐值来充当偏置。(3)定义的模型也称为玻尔兹曼机或Hopfield网络。这个模型的卡通形象看起来像成对耦合的小箭头图6：

         

在热平衡状态下，玻尔兹曼概率分布反映了哪种上下自旋模式或自旋配置是首选。自旋系统的配分函数不仅是一个归一化常数，而且是一个神奇的对象，它通过自由能将波动自旋的微观世界与热力学可观察量联系起来。即使对于简单的自旋系统，在大多数情况下，通过对所有可能配置求和来计算配分函数也是一件非常困难的事情。

二元自旋模型很好，但如今很少能激发机器学习从业者的兴趣。像 Transformer 这样的现代神经网络作用于矢量序列，例如标记嵌入或图像块。我们可以考虑矢量自旋模型，而不是完全放弃自旋模型。将二元自由度替换为d维自由度，我们可以定义一个自旋模型能量函数

(4)

这里标量积变成了向量点积。这种形式的模型首先出现在 1960 年代的统计力学文献中，作为经典d维向量模型。它们也出现在最近关于自旋玻璃模型的高维泛化的研究中7。

现在我们如何关联类似方程式 (4)的矢量自旋系统，到现代神经网络？

5. 为什么我们不直接用数据探测矢量自旋系统？

让我们继续一个直观的想法。想象一下，我们想要揭示方程式(4)表述的矢量自旋系统到矢量数据序列。我们可以通过让序列充当自旋系统的外部磁场来做到这一点(h1,h2,…hN)。然后我们想观察自旋系统如何响应这种特定的模式环境。

如果自旋系统响应计算中的所有步骤都能以可微分的方式实现，我们应该能够通过优化耦合参数来设计其集体行为，以更好地响应未来的输入数据。我们建议根据从自由能计算的磁化来观察自旋系统的响应。

6. 统计力学的一部分：磁化和自由能

为了便于表示，我们称模型参数为

, 自旋和外部磁场. 然后我们可以示意性地将自旋系统的配分函数写成

(5)

和相应的自由能为 

.

磁化是我们的自旋系统对外部磁场

的响应. 从标准热力学中，我们知道可以通过对外部场进行微分来计算自由能的磁化强度8。

(6)

在这种情况下，归结为计算自旋期望值。每个位置的磁化强度取决于耦合，并且通过自旋之间的耦合，也依赖所有位置的外部场的值。磁化揭示了当我们将自旋系统置于模式环境中时，自旋将如何共同趋于对齐。

在我们继续之前，我们必须考虑一个复杂的问题。如果我们想在Transformer模块和矢量自旋系统之间建立对应关系，我们将不得不考虑依赖于外部磁场的耦合。例如，普通transformer 中的注意力矩阵看起来像

(7)

其中矩阵H 是表示外部磁场矢量的堆栈。自旋之间的相互作用是根据输入动态确定的。从物理学的角度来看，这些“摊销”耦合非常奇怪且非常不寻常，但Transformer就是这样。

耦合对外部场的潜在依赖性将等式 (6) 的磁化强度更改为以下形式的表达式

(8)

其中两点相关函数被视为耦合贡献的权重9 。在实践中，我们当然应该让自动微分框架跟踪依赖关系，这样我们就可以摆脱简单的计算

(9)

假定我们有可用的（近似）自由能的可微分表达式。

7. 将可微自旋系统转化为神经网络

现在，让我们使用上面介绍的成分来构建一个神经网络模块，该模块封装了一个矢量自旋系统。给定能量函数方程(4)和自由能

，我们让输入的数据扮演外部磁场的角色，并返回磁化强度作为响应。

好的。但是难道我们之前没有提到配分函数（以及自由能和磁化强度）非常难以计算吗？如果我们无法计算任何东西，为什么还要引入所有这些形式化的表达式？

回顾 1950 年代至 70 年代的统计力学论文，事实证明物理学家已经开发出几种可用于处理矢量自旋系统的技巧和近似方法。可以在深度隐式注意力（下图左）和近似自由能最小化（下图右）中找到上述配分函数方法适用于矢量自旋系统的计算证据。

在这些例子中，配分函数方程(5)的近似值分别根据平均场理论和最速下降法获得。我们对这两种方法的数值实现，在内部都依赖于深度隐式层，以确保固定点计算和求根步骤是有效可微的。

8. 眯眼练习：认识transformer模块

根据方程式(9) 从深度隐式注意和近似自由能最小化中获得的（近似）自由能计算磁化，揭示了一个令人惊讶的熟悉的高级结构：残差连接、标记混合、归一化和通道混合的模式。从另一个方向接近症结所在，我们认为transformer模块通过对一般磁化响应 (8)实施特定近似来对输入做出反应。

残差连接与输入成正比，由外部磁场的存在引起。标记混合贡献来自能量函数中的耦合项，并且混合输入而不作用于局部向量自旋维度。归一化源于要求自旋系统的能量保持与晶格位置的数量成线性比例以及归一化外部磁场矢量。通道混合贡献包括可以在局部应用的磁化项，例如平均场方法中的 Onsager 自校正项或（近似）来自方程式(8)中依赖于输入的耦合的贡献。

综上所述，这些观察表明我们可以将Transformer模块的正向传递描绘成矢量自旋系统的封装器：模块的输入被路由到外部磁场（并且，可选地，到参数化耦合函数），之后磁化作为输出返回。Transformer模块与可微物理系统有着惊人的相似之处，我们可以通过训练来控制其集体行为。

9. 训练 transformer 模块塑造集体行为

既然我们可以将 transformer 模块想象成响应数据探测的物理自旋系统，让我们想象一下它们的训练是什么样的。

在我们的自旋系统方程的能量函数(4)的水平上，我们可以通过引入（离散）时间维度并使外部磁场随时间变化来对Transformer模块的训练过程进行建模，从而得出

(10)

在每个训练步骤t, 一系列输入数据

承担外磁场的作用。在正向传播过程中，磁化mi根据当前模型参数和外部磁场的存在以可微分的方式计算。在物理上，我们在每个训练步骤都考虑带有“冻结”耦合的“淬火”系统。在向后传递期间，模块的耦合参数Jij得到更新，推动自旋系统中的相互作用，从而影响其在未来迭代中对类似数据的磁化响应。

我们可以将此训练过程视为逐渐塑造由数据驱动的可微矢量自旋系统的集体行为。如果耦合取决于输入，如方程式 (7)，我们应该使耦合在方程式中也与时间相关 (10). 在这种情况下，外部磁场以及参数化耦合在每个训练步骤都会立即发生变化。

10. 训练深度Transformer编排自旋系统集

训练一个深度transformer模型，相当于通过建立一个可微的关联结构来编排一堆transformer模块，其中一个自旋系统的磁化驱动下一个自旋系统。训练过程中的摆动(数十亿)参数，推动自旋系统集合的级联反应行为，以更好地适应由数据和损失函数指定的集合(元)任务。

11.结论

在这篇文章中，我们认为transformer模块的正向传递映射到响应数据的矢量自旋模型中的计算磁化。通过对一个精心选择的能量函数求导，我们推广了之前在现代连续Hopfield网络中理解softmax注意力模块的工作，我们建议对一个一般矢量自旋系统的自由能求导，以得到一个完整transformer模块的架构蓝图。         

通过从不相干的、统计力学的角度缩小和接近transformer，我们获得了transformer的物理直觉，当我们把自己局限在在纷繁变化的显式神经网络架构时，这种直觉似乎很难获得。将transformer模块视为伪装的自旋模型不仅可以作为近似计算磁化的不同方法，统一架构变化，而且还可以阐释transformer在深度学习中的经验性质的成功。