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

Abstract: Weight decay is widely used as a regularizer in LLMs, 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 $L^2$ regularization--by proving it satisfies Villani's criteria for coercive energy functions. Specifically, we show that the regularized loss $\mathcal{F}$ is infinitely differentiable, grows at least quadratically, has Gaussian-integrable tails, and satisfies the differential growth condition $-Δ\mathcal{F} + \tfrac{1}{s}|\nabla\mathcal{F}|^{2} \to \infty$ as $|θ| \to \infty$ for all $s>0$. From this structure, we derive explicit log-Sobolev and Poincaré constants $C_{\mathrm{LS}} \leq λ^{-1} + d/λ^{2}$, linking the regularization strength $λ$ and model dimension $d$ 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(θ) = -Δ\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$, 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.

• The paper establishes that quadratic weight decay induces Villani coercivity in Transformer loss functions, ensuring robust finite-time convergence and generalization guarantees.
• It leverages functional-analytic techniques to derive explicit log-Sobolev and Poincaré constants, validated via spectral analysis and Hutchinson trace estimators.
• The results offer practical hyperparameter guidelines that unify stochastic gradient methods with PAC-Bayesian bounds for large-scale Transformer optimization.

Functional-Analytic Foundations of Optimization and Generalization in Regularized Transformer Loss Landscapes

Overview

The paper "Weight-Decay Turns Transformer Loss Landscapes Villani: Functional-Analytic Foundations for Optimization and Generalization" (2605.06599) rigorously characterizes the mathematical role of $L^2$ regularization (weight decay) in large-scale Transformer training. By connecting regularized cross-entropy loss landscapes to Villani's coercive energy functions—an analytic class from optimal transport and stochastic calculus—this work validates strong convergence and generalization properties for stochastic gradient descent (SGD) and Langevin-based optimizers. The study bridges a theoretical gap by demonstrating that the addition of quadratic regularization is both necessary and sufficient for the loss function to acquire geometric properties (coercivity, integrability, differential growth) that permit explicit log-Sobolev and Poincaré bounds. These results provide finite-time convergence guarantees and principled PAC-Bayesian generalization bounds dependent exclusively on regularization strength $\lambda$ and parameter dimension $d$.

Theoretical Contributions

Villani Criteria and Transformer Objectives

The paper analyzes the standard Transformer loss,

$$\mathcal{F}(\theta) = \frac{1}{N} \sum_{i=1}^{N} [-\log p_\theta(y_i|x_i)] + \frac{\lambda}{2}\|\theta\|^2,$$

where $p_\theta(y|x)$ is the softmax probability produced by the model. The authors prove $\mathcal{F}$ satisfies Villani's three foundational conditions:

• Smoothness: The loss is $C^\infty$ due to the analytic nature of core Transformer operations.
• Coercivity at Infinity: The quadratic term ensures the landscape diverges as $\|\theta\| \rightarrow \infty$, preventing optimization runaway.
• Gaussian-Integrable Tails: The loss induces a probability measure with tails decaying at least as fast as a Gaussian.
• Differential Growth: The field $$-\Delta\mathcal{F} + \frac{1}{s}\|\nabla\mathcal{F}\|^2$$ diverges for all $s > 0$, confirming the confining curvature necessary for log-Sobolev inequalities.

These analytic properties are not satisfied by the unregularized cross-entropy loss, confirming the necessity of weight decay for robust optimization behavior.

Figure 1: A single Transformer block annotated with parameter groups $\lambda$0, highlighting which parameters are encompassed by the regularization term.

Explicit Log-Sobolev and Poincaré Constants

Leveraging Villani theory, the authors derive that the log-Sobolev constant,

$\lambda$1

depends solely on $\lambda$2, $\lambda$3, and user-chosen temperature parameter $\lambda$4, disregarding training data specifics. This constant underpins exponential convergence guarantees for Langevin-type optimizers, including SGD and Adam with noise injection. The paper proves finite-time convergence bounds for noisy gradient descent and quantitatively links regularization strength to mixing speed and generalization capacity.

Figure 2: Comparison of unregularized ("Valley") vs. regularized ("Bowl") loss landscapes, illustrating how weight decay removes flat regions, institutes coercivity, and establishes unique minima.

Geometric Diagnostics and Empirical Validation

Villani Diagnostic $\lambda$5

To operationalize theoretical conditions, a scalable geometric diagnostic,

$\lambda$6

is introduced. The paper demonstrates quadratic divergence of $\lambda$7 in the presence of weight decay and establishes its reliability as a proxy for landscape coercivity in models exceeding $\lambda$8M parameters.

Figure 3: Empirical evaluation of $\lambda$9 as a function of $d$0 along random rays, showing transition from saturation (no weight decay) to quadratic growth (strong regularization).

Spectral Properties

The work provides Hessian spectral analysis, documenting how $d$1 inflates tail eigenvalues and thereby sharpens landscape curvature. This transforms optimization geometry from flat valleys (prone to barren plateaus and poor mixing) to confining bowls (rapid mixing, unique minima).

Figure 4: Spectral radius of $d$2 vs. $d$3, evidencing the linear spectral inflation due to weight decay.

Hutchinson Trace Estimation

Efficient computation of $d$4 via Hutchinson's randomized trace estimator is validated, with theoretical variance guarantees. The diagnostic protocol is shown to be tractable for high-dimensional models and is quantitatively robust across training phases.

Figure 5: Evolution of Hutchinson trace estimator distributions across training, demonstrating stability and convergence to theoretical expectations.

Figure 6: Evolution of variation coefficient in trace estimates, validating probe count and stability for large-scale diagnostics.

Optimization and Generalization Implications

Langevin Mixing and SGD Convergence

The Villani structure implies exponential entropy contraction for Langevin dynamics and noisy SGD. Theoretical bounds predict finite-time convergence rates that improve as $d$5 increases, provided learning rates are tuned accordingly. This finding unifies stochastic optimization with functional inequalities from statistical physics.

Figure 7: Empirical optimization speed versus theoretical exponential bounds, highlighting the influence of weight decay on convergence.

PAC-Bayesian Generalization

The paper leverages the regularized Gibbs measure as a PAC-Bayesian prior, deriving tight generalization bounds,

$d$6

where the KL term depends quadraticly on $d$7. Strong empirical correlation between theoretical bounds and observed validation performance is established.

Figure 8: PAC-Bayesian generalization bounds vs. observed test performance, with tightness increasing for larger weight decay and high $d$8 correlation.

Practitioner Advices

Explicit hyperparameter guidelines are provided, grounded in Villani constants, for learning rate schedules and regularization regimes to optimize mixing speed and generalization trade-off.

Scalability and Extension

The paper analyzes the scalability of Villani diagnostics and log-Sobolev bounds for models up to $d$9 parameters, confirming favorable theoretical scaling and practical feasibility via distributed computing and algorithmic optimizations.

Figure 9: Log-Sobolev constant scaling and computational cost across model sizes, demonstrating tractability for future large Transformer applications.

Broader Implications and Future Directions

This work establishes a principled analytic foundation for Transformer optimization, translating weight decay from an empirical heuristic to a mathematically necessary ingredient for confining landscape geometry and guaranteeing fast, reliable convergence. It substantiates regularization as the primary mechanism for connecting curvature, mixing, and generalization, offering actionable, theoretically-guided hyperparameter recommendations.

Potential future directions include:

• Extending the Villani framework to encoder-decoder and multimodal architectures.
• Tailoring diagnostics to structured or adaptive regularization schemes (e.g., AdamW, LoRA).
• Refining spectral gap estimators for ultra-large models.
• Applying the approach in reinforcement learning, quantum optimization, and physics-informed deep learning.

Conclusion

The paper provides a rigorous, functional-analytic explanation for the role of $$\mathcal{F}(\theta) = \frac{1}{N} \sum_{i=1}^{N} [-\log p_\theta(y_i|x_i)] + \frac{\lambda}{2}\|\theta\|^2,$$0 weight decay in Transformer training. Regularization reshapes loss landscapes to satisfy Villani coercivity, enabling explicit mixing rates and generalization guarantees. By bridging geometric analysis, stochastic dynamics, and learning theory, it offers both mathematical insight and practical guidance for modern deep learning optimization. The methodology and diagnostic tools are broadly extensible, marking a significant step in theoretical interpretability and principled algorithm design for high-capacity predictive models.