理解Transformer的训练动态

理论分析揭示了模型训练过程中的优化机制，并发现对于某些优化任务，高斯注意力核可能比softmax表现更优。

核心发现

在近期举办的2024年神经信息处理系统大会（NeurIPS）上，一项针对Transformer架构的新分析表明：

1. 确定了能够概率性保证收敛到全局最优解的超参数和初始化条件
2. 通过消融研究证明注意力核的选择影响收敛速度
3. 高斯核在某些情况下能够实现收敛，而常用的softmax核则无法实现
4. 在特定设置下，使用高斯核训练的模型因优化景观更平滑而收敛更快

三矩阵机制

Transformer中的注意力权重计算涉及三个矩阵：

• 查询矩阵（query matrix）
• 键矩阵（key matrix）
• 值矩阵（value matrix）

在自注意力机制中，查询矩阵和键矩阵应用于相同输入；在交叉注意力机制中，它们应用于不同输入。注意力核定义了在查询和键编码上执行的操作，其结果表示输入集之间的相关性。

矩阵更新策略研究

通过分析仅更新部分矩阵而固定其他矩阵的结果，发现：

• 当所有三个矩阵都可更新时，普通梯度下降（GD）可实现全局最优性
• 仅更新值矩阵时，GD仍然是最优的
• 仅更新查询矩阵时，仅高斯核能保证GD收敛

实验验证

在两个不同数据集（文本分类任务和图像解释分割任务）上的实验表明：

• 高斯核能够实现更快的收敛速度和更高的模型精度
• 理论分析表明收敛性主要依赖于值矩阵的更新
• 值矩阵与核操作结果的乘法是线性操作，而核操作是非线性的

初始化条件要求

为保证收敛性，需要满足一组初始化条件：

• 核操作矩阵必须满秩（列线性独立）
• 查询矩阵和键矩阵的特征值与值矩阵特征值的比率需超过指定阈值

这些发现为理解Transformer训练动态提供了重要见解，随着Transformer在日常生活中的作用越来越大，希望AI社区的其他成员能够在此基础上进一步扩展相关分析。

相关研究成果详见NeurIPS 2024会议论文《Unraveling the gradient descent dynamics of transformers》