arXiv:2605.17660math.OCcs.AI2026-05被引 2

从数学上证明了无限深宽Transformer可稳定训练,为优化提供理论支撑。

Training Infinitely Deep and Wide Transformers

  • 用平均场理论建模无限深度宽的Transformer,将训练看作控制神经偏微分方程。
  • 在初始损失足够小时,梯度流可收敛到全局最优,消除局部极小陷阱。
  • 揭示注意力机制下NTK可逆的充要条件,适用于多种分布数据。

Transformer已成为现代机器学习的主流架构,但其训练动态的理论理解仍有限。本文建立了一个严格的数学框架,分析均值场条件下基于梯度的Transformer训练,其中深度(层数)和宽度(注意力头数)均趋于无穷。由于注意力机制耦合多个标记分布,ResNet训练可视为控制神经常微分方程,而Transformer训练则对应控制神经偏微分方程。我们的均值场模型包含两类测度表示:逐层演化的标记分布与各层的注意力参数。我们建立了无限深Transformer前向传播的适定性,通过函数空间中的常微分方程描述标记演化。利用伴随敏感性分析,推导出训练风险的条件Wasserstein梯度显式公式,其中伴随变量由反向常微分方程控制。我们证明了条件Wasserstein度量空间中梯度流曲线的存在性与唯一性,为基于梯度的Transformer训练提供了严格基础。关键技术贡献是给出注意力机制下神经正切核(NTK)可逆的充要条件:我们证明其等价于对数求和指数函数在仿射函数模下的线性无关性,该条件在离散分布、均匀分布和高斯混合分布等多种情况下均成立。在此NTK可逆假设下,当初始损失足够小时,梯度流收敛至全局最小值,从而消除了优化景观中的伪局部极小。

原文摘要 · Abstract (English)

Transformers have become the dominant architecture in modern machine learning, yet the theoretical understanding of their training dynamics remains limited. This paper develops a rigorous mathematical framework for analyzing gradient-based training of transformers in the mean-field regime, where both the depth (number of layers) and width (number of attention heads) tend to infinity. While ResNet training can be understood as controlling a neural ODE, transformer training corresponds to controlling a neural PDE, due to the coupling of multiple token distributions through the attention mechanism. Our mean-field model features two types of measure representations: token distributions evolving through layers and attention parameters at each layer. We establish well-posedness of the forward pass through infinitely deep transformers, characterizing token evolution via flow maps that satisfy ODEs in function spaces. Using adjoint sensitivity analysis, we derive an explicit formula for the conditional Wasserstein gradient of the training risk, involving adjoint variables governed by backward ODEs. We prove the existence and uniqueness of gradient flow curves in the conditional Wasserstein metric space, establishing a rigorous foundation for gradient-based transformer training. A key technical contribution is providing necessary and sufficient conditions for injectivity of the Neural Tangent Kernel (NTK) for attention mechanisms: we show that NTK injectivity is equivalent to linear independence of log-sum-exp functions modulo affine functions, a condition satisfied by diverse token distributions, including discrete distributions, uniform distributions, and Gaussian mixtures. Under this NTK injectivity assumption, we prove that gradient flow converges to global minima when the initial loss is sufficiently small, eliminating spurious local minima from the optimization landscape.

Transformer优化理论均值场梯度流

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。