cs.LGMay 7, 2026

Weight-Decay Turns Transformer Loss Landscapes Villani: Functional-Analytic Foundations for Optimization and Generalization

Authors: Abhijit DasSayantan Dutta

Organizations: GE HealthCare, Bangalore 560066, Karnataka, India

Abstract

Weight decay is widely used as a regularizer in large language models, yet its precise role in shaping Transformer loss landscapes remains theoretically underexplored. This paper provides the first rigorous functional-analytic characterization of the standard Transformer objective--cross-entropy loss with L2L^2 regularization--by proving it satisfies Villani's criteria for coercive energy functions. Specifically, we show that the regularized loss F\mathcal{F} is infinitely differentiable, grows at least quadratically, has Gaussian-integrable tails, and satisfies the differential growth condition ΔF+1sF2-Δ\mathcal{F} + \tfrac{1}{s}\|\nabla\mathcal{F}\|^{2} \to \infty as θ\|θ\| \to \infty for all s>0s>0. From this structure, we derive explicit log-Sobolev and Poincaré constants CLSλ1+d/λ2C_{\mathrm{LS}} \leq λ^{-1} + d/λ^{2}, linking the regularization strength λλ and model dimension dd to finite-time convergence guarantees for noisy stochastic gradient descent and PAC-Bayesian generalization bounds that tighten with increasing λλ. To validate our theory, we introduce a scalable Villani diagnostic Ψs(θ)=ΔF+s1F2Ψ_s(θ) = -Δ\mathcal{F} + s^{-1}\|\nabla \mathcal{F}\|^2 and estimate it efficiently using Hutchinson trace probes in models with over 100M parameters. Experiments on GPT-Neo-125M across Penn Treebank and WikiText-103 confirm the predicted quadratic growth of ΨsΨ_s, spectral inflation of the Hessian, and exponential convergence behavior consistent with our log-Sobolev analysis. These results demonstrate that weight decay not only improves generalization empirically but also establishes the mathematical conditions required for fast Langevin mixing and theoretically grounded curvature-aware optimization in deep learning.

Explore similar work

CardsList