下一词预测中 Transformer 的非渐近收敛性

Transformer 在现代机器学习中因其处理序列数据的卓越能力而取得了非凡的成功，尤其是在下一词预测 (NTP) 任务中。然而，关于其在 NTP 中的性能的理论理解仍有限，现有研究主要关注渐近性能。本文对单层由自注意力模块和前馈层组成的 Transformer 的训练动力学进行细粒度非渐近分析。我们首先使用基于部分序的数学框架刻画了 NTP 训练数据集的关键结构属性。随后，我们设计了两阶段训练算法，其中训练前馈层的预处理阶段和训练注意力层的主阶段均表现出快速收敛性能。具体而言，两层均以其相应最大边际解的方向收敛呈亚线性收敛。我们还表明交叉熵损失具有线性收敛速率。进一步，我们表明训练后的 Transformer 在数据分布迁移下仍具非平凡预测能力，这为 Transformer 异常的泛化性能提供了解释。我们的分析技术涉及发展注意力梯度的新性质，并进一步深入分析这些性质如何促成训练过程的收敛。我们的实验进一步验证了我们的理论发现。