理解Transformer的训练动态

理论分析为模型训练过程中的优化提供了深刻见解，并揭示对于某些优化场景，高斯注意力核可能比softmax表现更佳。

核心发现

当前大多数突破性AI模型都基于Transformer架构，其显著特点是采用了注意力机制。例如，在大型语言模型中，Transformer决定在生成下一个词时需要特别关注文本字符串中的哪些词；在视觉语言模型中，它可能决定在计算像素值时需要关注指令中的哪些词语。

鉴于Transformer模型日益重要，我们自然希望更好地理解它们的动态特性——例如训练过程是否会收敛到一个有用的模型，收敛速度如何，或者哪些架构变体最适合什么目的。然而，注意力机制的复杂性使得传统的分析工具难以应用。

在2024年神经信息处理系统大会(NeurIPS)上，我们提出了对Transformer架构的新分析。首先，我们确定了能够概率性保证收敛到全局最优解的超参数和初始化条件。

通过消融研究，我们还展示了注意力核的选择——用于计算注意力权重的函数——会影响收敛速度。具体来说，高斯注意力核有时能够在更常见的softmax核无法收敛的情况下实现收敛。最后，我们进行了一项实证研究，表明在某些特定设置下，由于更平滑的优化景观，使用高斯核训练的模型比使用softmax核训练的模型收敛更快。

高斯核和softmax核在两种不同机器学习任务中的优化景观。由于优化景观更平滑，使用高斯核训练的模型比使用softmax核训练的模型收敛更快。

三个矩阵的故事

在Transformer中，注意力权重计算涉及三个矩阵：查询矩阵、键矩阵和值矩阵。这三个矩阵都用于生成输入数据的编码。在自注意力机制中（如LLMs的情况），查询矩阵和键矩阵应用于相同的输入。在交叉注意力机制中，它们应用于不同的输入：例如，在多模态模型中，一个矩阵可能用于编码文本，而另一个用于编码图像。

注意力核定义了在查询和键编码上执行的操作；该操作的结果表示一组输入与另一组（或自身）的相关性。值矩阵产生的编码表示数据的语义属性。核操作的结果与值矩阵产生的编码相乘，强调某些语义特征并减弱其他特征。其结果本质上是模型下一个输出的语义内容配方。

通常在模型训练期间，所有三个矩阵会一起更新。但我们分析了仅更新矩阵子集而其他矩阵保持固定的结果。这使我们能够确定哪些矩阵和核函数对收敛速度影响最大。结果如下：

• 如果所有三个矩阵都可以更新，普通梯度下降(GD)可以实现全局最优性，无论使用高斯还是softmax注意力核；
• 如果只有值矩阵可以更新，GD仍然是最优的，无论使用哪种核；
• 如果只有查询矩阵可以更新，GD收敛仅在使用高斯核时得到保证。

这表明在某些情况下，常用的softmax核可能存在缺点，我们进行了一组实验来验证这一直觉。在两个不同的数据集上——一个用于文本分类任务，一个用于图像解释和分割任务——我们训练了成对的Transformer模型，一个使用高斯核，一个使用softmax核。在这两个任务中，高斯核都实现了更快的收敛速度和更高的模型准确率。

文本分类任务（上）和图像解释与分割任务（下）的实验结果。高斯核的结果为蓝色，softmax核的结果为红色。在准确率测量中（左），分数越高越好；在训练损失测量中（右），分数越低越好。

我们的分析还表明，理论上，收敛主要取决于值矩阵的更新，因为值矩阵与核操作结果的乘法是线性操作，而核操作是非线性的。

最后，我们的论文还提出了一组保证收敛所必需的初始化条件。这些条件包括核操作矩阵必须具有满秩——即其列线性无关——以及查询矩阵和键矩阵的特征值与值矩阵的特征值之比必须高于指定阈值。

更多细节可以在我们的论文中找到。我们希望AI社区的其他成员能够扩展我们的分析，随着Transformer在我们日常生活中扮演越来越重要的角色，进一步扩展我们对它们的理解。