A Mean-Field Analysis of Multi-Head Self-Attention under Cross-Entropy Training

• URL: http://arxiv.org/abs/2606.10469v1
• Date: Tue, 09 Jun 2026 06:38:27 GMT
• Title: A Mean-Field Analysis of Multi-Head Self-Attention under Cross-Entropy Training
• Abstract summary: 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.
• Score: 0.0
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: This paper develops a mean-field theory for a simplified single-layer causal multi-head self-attention model trained by cross-entropy minimization. Each attention head is treated as a particle in parameter space, and the empirical law of the heads is used as the large-head state variable. 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. Unlike classical mean-field analyses of shallow networks that often focus on square-loss regression, the present model contains the softmax residual from the cross-entropy objective and the query-key-value structure of masked self-attention. We prove a static finite-head approximation bound for the optimal risk, characterize global minimizers through a variational support condition, and establish a quantitative finite-time propagation-of-chaos estimate comparing finite-head stochastic gradient descent with the limiting PDE. We then 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. Finally, we prove local exponential stability under a Wasserstein strong-monotonicity condition and give verifiable stability and instability criteria for Dirac stationary measures. The results provide a rigorous baseline mean-field framework for attention-head training and clarify the additional compactness, landscape, and curvature assumptions needed to pass from stationarity to convergence and stability.

Related papers

• Finite-Sample Metric Non-Collapse for Geometrically Supervised Latent World Models in Control [0.0]
  We establish a finite-sample learning-to-control theory for geometrically supervised latent models of nonlinear deterministic systems.<n>We introduce an encoder-only local--global metric hinge that enforces directional resolution and separated-state discrimination.<n>A modular deterministic transfers the learned certificates to trajectory, finite-horizon cost, learned-cost-head, and corollary guarantees, while a validated finite-net result enables sharper model-specific certification.
  arXiv  Detail & Related papers  (2026-08-07T14:25:33Z)
• Variance Reduction for Stochastic Gradient Generalized Non-reversible Langevin Monte Carlo Algorithms [7.336637417508747]
  We study the leading-order fluctuation of gradient-Maruyama estimators for non-reversible Langevin dynamics.<n>We prove that the empirical average over a horizon of order the inverse squared stepsize satisfies a central limit theorem in the vanishing-stepsize regime.
  arXiv  Detail & Related papers  (2026-06-27T08:25:48Z)
• Stability and Generalization of Push-Sum Based Decentralized Optimization over Directed Graphs [55.77845440440496]
  Push-based decentralized communication enables optimization over communication networks, where information exchange may be asymmetric.<n>We develop a unified uniform-stability framework for the Gradient Push (SGP) algorithm.<n>A key technical ingredient is an imbalance-aware generalization bound through two quantities.
  arXiv  Detail & Related papers  (2026-02-24T05:32:03Z)
• The Procrustean Bed of Time Series: The Optimization Bias of Point-wise Loss [53.542743390809356]
  This paper aims to provide a first-principles analysis of the Expectation of Optimization Bias (EOB)<n>Our analysis reveals a fundamental paradigm paradox: the more deterministic and structured the time series, the more severe the bias by point-wise loss function.<n>We present a concrete solution that simultaneously achieves both principles via DFT or DWT.
  arXiv  Detail & Related papers  (2025-12-21T06:08:22Z)
• Derivative-Free Sequential Quadratic Programming for Equality-Constrained Stochastic Optimization [3.2489082010225485]
  We consider solving nonlinear optimization problems with an objective and deterministic equality constraints.<n>We propose a Derivative-Free Sequential Quadratic Programming (DF-SSQP) method.<n>Under standard assumptions, we establish the global almost-sure convergence of the proposed DF-SSQP method.
  arXiv  Detail & Related papers  (2025-10-25T23:51:20Z)
• 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)
• Asymptotic breakdown point analysis of the minimum density power divergence estimator under independent non-homogeneous setups [2.449909275410287]
  The minimum density power divergence estimator (MDPDE) has gained significant attention in the literature of robust inference.<n>It has been successfully applied in various setups, including the case of independent and non-homogeneous (INH) observations.<n>No general result is known about the global reliability or the breakdown behavior of this estimator under the INH setup.
  arXiv  Detail & Related papers  (2025-08-17T16:33:58Z)
• 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)
• Global Convergence of Over-parameterized Deep Equilibrium Models [52.65330015267245]
  A deep equilibrium model (DEQ) is implicitly defined through an equilibrium point of an infinite-depth weight-tied model with an input-injection. Instead of infinite computations, it solves an equilibrium point directly with root-finding and computes gradients with implicit differentiation. We propose a novel probabilistic framework to overcome the technical difficulty in the non-asymptotic analysis of infinite-depth weight-tied models.
  arXiv  Detail & Related papers  (2022-05-27T08:00:13Z)
• A Local Convergence Theory for the Stochastic Gradient Descent Method in Non-Convex Optimization With Non-isolated Local Minima [0.0]
  Non-isolated minima presents a unique challenge that has remained under-explored. In this paper, we study the local convergence of the gradient descent method to non-isolated global minima.
  arXiv  Detail & Related papers  (2022-03-21T13:33:37Z)
• Nonconvex Stochastic Scaled-Gradient Descent and Generalized Eigenvector Problems [98.34292831923335]
  Motivated by the problem of online correlation analysis, we propose the emphStochastic Scaled-Gradient Descent (SSD) algorithm. We bring these ideas together in an application to online correlation analysis, deriving for the first time an optimal one-time-scale algorithm with an explicit rate of local convergence to normality.
  arXiv  Detail & Related papers  (2021-12-29T18:46:52Z)

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.