通过梯度异质性理解Transformer优化

Transformer难以用随机梯度下降（SGD）优化，很大程度上依赖自适应优化器如Adam。尽管它们在经验上取得了成功，但Adam相对于SGD性能优越的原因仍知之甚少。在本研究中，我们通过梯度异质性的视角分析Transformer模型的优化，梯度异质性定义为参数块间梯度范数的变化。我们提供的理论分析表明，梯度异质性与Hessian异质性共同降低了基于梯度的方法（如SGD）的收敛性，而基于符号的方法对此效应的敏感度显著较低。Adam的逐坐标归一化使其更新方向主要依赖于梯度符号，因此Adam可以被解释为SignSGD的一种软变体。我们的分析利用了SGD和SignSGD在不同范数下遵循最速下降方向的事实，并推导了迭代复杂度的上界，这对SignSGD中的学习率缩放具有启示意义。我们进一步研究了Transformer架构中梯度异质性的起源，并表明它受层归一化放置位置的强烈影响，其中Post-LN架构表现出特别显著的异质性。在NLP和视觉领域微调Transformer的实验结果验证了我们的理论分析。代码可在https://github.com/tom4649/gradient-heterogeneity获取。