Training Infinitely Deep and Wide Transformers

Transformer 已成为现代机器学习中占主导地位的模型架构，然而对其训练动力学的理论理解仍十分有限。本文构建了一个严格的数学框架，用于在平均场（mean-field）极限下分析基于梯度的 Transformer 训练过程；在此极限下，网络深度（即层数）与宽度（即注意力头数）均趋于无穷大。与 ResNet 的训练可被理解为对一个神经常微分方程（neural ODE）的控制不同，Transformer 的训练则对应于对一个神经偏微分方程（neural PDE）的控制——其根源在于注意力机制将多个词元（token）分布耦合在了一起。我们的平均场模型采用两类测度（measure）表示：一类是随网络层数演化而变化的词元分布；另一类是每一层所对应的注意力参数。我们严格证明了无限深 Transformer 前向传播过程的适定性（well-posedness），并借助函数空间中的常微分方程所刻画的流映射（flow maps），完整表征了词元分布的演化规律。通过伴随灵敏度分析（adjoint sensitivity analysis），我们推导出训练风险函数关于条件 Wasserstein 梯度（conditional Wasserstein gradient）的显式表达式，其中伴随变量满足一组反向常微分方程。我们进一步证明：在条件 Wasserstein 度量空间中，梯度流曲线存在且唯一，从而为基于梯度的 Transformer 训练奠定了坚实的理论基础。本文一项关键的技术贡献在于：首次给出了注意力机制下神经正切核（Neural Tangent Kernel, NTK）可逆性的充要条件——我们证明，NTK 可逆当且仅当若干 log-sum-exp 函数在仿射函数类模意义下线性无关；而该条件对多种实际常见的词元分布均成立，包括离散分布、均匀分布以及高斯混合分布。在此 NTK 可逆性假设下，我们进一步证明：只要初始损失足够小，梯度流必将收敛至全局最小值，从而彻底消除了优化景观中出现虚假局部极小值（spurious local minima）的可能性。

论文旨在解决Transformer模型梯度训练动力学的理论基础缺失问题，特别是在深度（层数）与宽度（注意力头数）同时趋于无穷的mean-field极限下，如何严格刻画其参数更新、前向传播与损失最小化过程。该问题新颖：尽管ResNet已有神经ODE等成熟mean-field理论，但Transformer因自注意力机制引入跨token耦合，其连续极限本质上是PDE而非ODE，此前尚无严格数学框架。

提出首个Transformer mean-field理论框架：将token序列建模为层间演化的概率测度流，注意力参数建模为层依赖的条件测度；前向传播由函数空间中的流映射ODE刻画；通过条件Wasserstein梯度与伴随敏感性分析导出可证伪的梯度流方程；关键创新在于将注意力机制的NTK可逆性归约为log-sum-exp函数系对仿射函数的线性无关性，并证明该条件在常见数据分布（离散、均匀、高斯混合）下自然成立。

理论贡献为主，无实验验证（属纯数学分析）；核心亮点包括：（1）首次建立Transformer无限深宽极限下的well-posed神经PDE模型；（2）给出NTK单射性的充要解析判据并验证其广泛适用性；（3）在小初值损失下证明全局收敛，排除spurious local minima；（4）为后续研究提供新工具——条件Wasserstein梯度流、注意力流形几何、以及log-sum-exp函数论在ML中的应用。未使用数据集，无代码开源（纯理论工作）。值得深入的方向：有限宽度/深度的误差界估计、实际训练中NTK条件的实证检验、以及该框架对稀疏注意力或线性注意力的拓展。

Neural ODEs for ResNets (Chen et al., NeurIPS 2018); Mean-field limit of two-layer neural nets (Mei, Montanari & Nguyen, PNAS 2018); Wasserstein gradient flows in ML (Arjovsky et al., ICML 2017; Sanjabi et al., NeurIPS 2018); NTK theory for MLPs/CNNs (Jacot et al., NeurIPS 2018); Recent attempts on Transformer linearization (Yang et al., ICLR 2022; HaoChen et al., NeurIPS 2023), but none establish PDE limits or conditional Wasserstein analysis.