信号传播的几何动力学预测 Transformer 的可训练性
无序系统与神经网络
2024-03-06 v1 机器学习
摘要
我们研究了深度随机初始化 Transformer 中的前向信号传播和梯度反向传播，得出了关于初始化超参数的简单充要条件，以确保深度 Transformer 的可训练性。我们的方法将个 token 的表示在穿过 Transformer 层时的演化视为个相互作用粒子的离散时间动力系统。我们从置换对称单纯形出发，推导了该粒子系统演化几何的简单更新方程。我们的更新方程表明，在没有 MLP 层的情况下，该系统将坍缩为一条线，这与先前关于 Transformer 秩坍缩的工作一致。然而，与先前工作不同，我们的演化方程可以定量追踪在非线性 MLP 层存在下的粒子几何，并揭示了作为初始化超参数（如注意力和 MLP 残差连接的强度及权重方差）函数的秩序 - 混沌相变。在有序相中，粒子相互吸引并坍缩为一条线；而在混沌相中，粒子相互排斥并收敛于规则的-单纯形。我们解析推导了两个 Lyapunov 指数：一个角度指数，控制该粒子系统偏离混沌边缘的程度；另一个梯度指数，控制反向传播梯度的指数增长或衰减速率。实验表明，令人惊讶的是，训练结束时的最终测试损失仅由训练开始时的这两个指数就能很好地预测，且这两个指数的同时消失构成了实现最小测试损失的简单充要条件。
引用
@article{arxiv.2403.02579,
title = {Geometric Dynamics of Signal Propagation Predict Trainability of Transformers},
author = {Aditya Cowsik and Tamra Nebabu and Xiao-Liang Qi and Surya Ganguli},
journal= {arXiv preprint arXiv:2403.02579},
year = {2024}
}