A First-Order Mean Field Control Analysis of Transformer Layers under Cross-Entropy Training

• URL: http://arxiv.org/abs/2606.23235v1
• Date: Mon, 22 Jun 2026 12:21:37 GMT
• Title: A First-Order Mean Field Control Analysis of Transformer Layers under Cross-Entropy Training
• Abstract summary: We study Transformer-type residual layers under cross-entropy training through a continuous-depth mean field control viewpoint.<n>For fixed controls, we prove an $O(varepsilon)$ pathwise approximation of finite-depth trajectories by the continuous flow.<n>We also give finite-class and metric-entropy uniform estimates, compare optimal values, and discuss existence, stability, continuous-to-discrete recovery, and range estimates for continuous minimizers.
• Score: 1.4180331276028662
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: We study Transformer-type residual layers under cross-entropy training through a continuous-depth mean field control viewpoint. Depth is treated as time, layer parameters as controls, and the residual Transformer recursion as an explicit Euler scheme for a controlled hidden-state flow. For fixed controls, we prove an $O(\varepsilon)$ pathwise approximation of finite-depth trajectories by the continuous flow and combine this with high-probability sampling bounds for the empirical cross-entropy risk. We formulate the limiting population problem as a first-order transport control problem for the law of hidden states and derive a Pontryagin condition whose terminal adjoint contains the softmax residual. We also give finite-class and metric-entropy uniform estimates, compare optimal values, and discuss existence, stability, continuous-to-discrete recovery, initialization, and range estimates for continuous minimizers.

Related papers

• Generalization Bounds on Optimal Control for Transformer Training and Wasserstein Distributional Robustness [8.564319625930894]
  We view data sets as probability laws on pairs of empirical input-output measures, allowing us to interpret the training problem as a finite-horizon Markovian control problem.<n>We show that the same machinery yields a distributionally robust control formulation of the training problem, connecting Transformer generalization to Wasserstein distributionally robust optimization.
  arXiv  Detail & Related papers  (2026-07-30T10:22:38Z)
• Taming Curvature: Architecture Warm-Up for Stable Transformer Training [52.6838951211596]
  We introduce a fast online estimator of the largest (preconditioned) Hessian eigenvalue (i.e., curvature) based on a warm-started variant for power with Hessian-vector products.<n>We find that training instabilities coincide with surges in preconditioned curvature and that curvature grows with depth.<n> Experiments on large Transformers validate that our approach enables efficient curvature tracking and reduces instabilities.
  arXiv  Detail & Related papers  (2026-06-15T14:16:56Z)
• A Mean-Field Analysis of Multi-Head Self-Attention under Cross-Entropy Training [0.0]
  This paper develops a mean-field theory for a single-layer causal multi-head self-attention model trained by cross-entropy minimization.<n>In the infinite-head limit, the averaged attention logits define a risk functional on probability measures, whose first variation generates a nonlinear Wasserstein gradient-flow equation.<n>We study the long-time behavior of the PDE: energy dissipation, convergence to the stationary set under compactness, convergence to a single stationary measure under topological or Kurdyka--ojasiewicz assumptions, and explicit convergence rates under gradient-domination conditions.
  arXiv  Detail & Related papers  (2026-06-09T06:38:27Z)
• ProFlow: Zero-Shot Physics-Consistent Sampling via Proximal Flow Guidance [35.08166384258028]
  ProFlow is a framework for zero-shot physics-consistent sampling.<n>It reconciles strict physical consistency and observational fidelity with the statistical structure of the pre-trained prior.<n>It achieves superior physical and observational consistency, as well as more accurate distributional statistics.
  arXiv  Detail & Related papers  (2026-01-28T03:57:00Z)
• Terminally constrained flow-based generative models from an optimal control perspective [32.87833798690545]
  Terminal Optimal Control with Flow-based models (TOCFlow) is a geometry-aware sampling-time guidance method for pre-trained flows.<n>We show that as the control penalty increases, the controlled process recovers the reference distribution, while as the penalty vanishes, the terminal law converges to a generalized Wasserstein projection onto the constraint manifold.<n>We evaluate TOCFlow on three high-dimensional scientific tasks spanning equality, inequality, and global statistical constraints.
  arXiv  Detail & Related papers  (2026-01-14T13:32:15Z)
• Verifying Closed-Loop Contractivity of Learning-Based Controllers via Partitioning [52.23804865017831]
  We address the problem of verifying closed-loop contraction in nonlinear control systems whose controller and contraction metric are both parameterized by neural networks.<n>We derive a tractable and scalable sufficient condition for closed-loop contractivity that reduces to checking that the dominant eigenvalue of a symmetric Metzler matrix is nonpositive.
  arXiv  Detail & Related papers  (2025-12-01T23:06:56Z)
• Revisiting Zeroth-Order Optimization: Minimum-Variance Two-Point Estimators and Directionally Aligned Perturbations [57.179679246370114]
  We identify the distribution of random perturbations that minimizes the estimator's variance as the perturbation stepsize tends to zero.<n>Our findings reveal that such desired perturbations can align directionally with the true gradient, instead of maintaining a fixed length.
  arXiv  Detail & Related papers  (2025-10-22T19:06:39Z)
• Closed-Form Last Layer Optimization [72.49151473937319]
  Under a squared loss, the optimal solution to the linear last layer weights is known in closed-form.<n>We show this is equivalent to alternating between gradient descent steps on the backbone and closed-form updates on the last layer.
  arXiv  Detail & Related papers  (2025-10-06T09:14:39Z)
• Controlling the Flow: Stability and Convergence for Stochastic Gradient Descent with Decaying Regularization [5.582101184758528]
  We prove strong convergence of reg-SGD to the minimum-norm solution of the original problem without additional boundedness assumptions.<n>Our analysis reveals how vanishing Tikhonov regularization controls the flow of SGD and yields stable learning dynamics.
  arXiv  Detail & Related papers  (2025-05-16T16:53:49Z)
• Beyond the Edge of Stability via Two-step Gradient Updates [49.03389279816152]
  Gradient Descent (GD) is a powerful workhorse of modern machine learning. GD's ability to find local minimisers is only guaranteed for losses with Lipschitz gradients. This work focuses on simple, yet representative, learning problems via analysis of two-step gradient updates.
  arXiv  Detail & Related papers  (2022-06-08T21:32:50Z)
• On the Convergence of Stochastic Extragradient for Bilinear Games with Restarted Iteration Averaging [96.13485146617322]
  We present an analysis of the ExtraGradient (SEG) method with constant step size, and present variations of the method that yield favorable convergence. We prove that when augmented with averaging, SEG provably converges to the Nash equilibrium, and such a rate is provably accelerated by incorporating a scheduled restarting procedure.
  arXiv  Detail & Related papers  (2021-06-30T17:51:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.