理解Transformer的训练动态

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

核心发现

当前突破性AI模型大多基于Transformer架构，其核心特征在于注意力机制的应用。例如在大型语言模型（LLM）中，Transformer通过注意力机制确定生成下一个词时需要重点关注文本中的哪些词汇；在多模态模型中，则可能决定在计算像素值时需要关注指令中的哪些词语。

随着Transformer模型的重要性日益提升，自然需要更深入理解其训练动态——例如训练过程是否会收敛到有效模型、收敛速度如何，以及哪种架构变体最适合特定任务。然而注意力机制的复杂性使得传统分析工具难以适用。

通过理论分析，首先确定了能够概率性保证收敛到全局最优解的超参数和初始化条件。消融研究表明，注意力核（用于计算注意力权重的函数）的选择会影响收敛速度：高斯核在某些情况下能够实现收敛，而常用的Softmax核则无法实现。

实证研究进一步表明，在特定设置下，由于更平滑的优化 landscape，使用高斯核训练的模型比使用Softmax核的模型收敛更快。

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

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

在模型训练期间，通常同时更新所有三个矩阵。但通过分析仅更新部分矩阵而固定其他矩阵的结果，能够识别哪些矩阵和核函数对收敛速度影响最大：

这表明在某些情况下，常用Softmax核可能存在缺陷。通过在文本分类和图像解释分割两个不同数据集上训练Transformer模型对（一对使用高斯核，一对使用Softmax核）验证了这一发现：在两个任务中，高斯核都能实现更快的收敛速度和更高的模型精度。

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

研究同时提出了一组保证收敛的必要初始化条件，包括：核操作矩阵需满秩（即列线性独立），查询矩阵和键矩阵的特征值与值矩阵特征值之比需超过指定阈值。

这些发现为理解Transformer的优化机制提供了新视角，随着Transformer在日常生活中的作用越来越大，希望AI社区能够在此基础上进一步扩展相关研究。