DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
1
Weight-Decay Turns Transformer Loss Landscapes Villani: Functional-Analytic Foundations for Optimization and Generalization
arXiv:2605.06599v1 [cs.LG] 7 May 2026
Abhijit Das, and Sayantan Dutta, Member, IEEE,
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 L2 regularization—by proving it satisfies Villani’s criteria for coercive energy functions. Specifically, we show that the regularized loss F is infinitely differentiable, grows at least quadratically, has Gaussian-integrable tails, and satisfies the differential growth condition −∆F + 1s ∥∇F∥2 → ∞ as ∥θ∥ → ∞ for all s > 0. From this structure, we derive explicit log-Sobolev and Poincaré constants CLS ≤ λ−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 (θ) = −∆F + s−1 ∥∇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. Index Terms—Transformer Optimization, Weight Decay Regularization, Villani Function, Barren Plateaus, Log-Sobolev Inequality, Langevin Dynamics, PAC-Bayes, Curvature-aware Training.
I. I NTRODUCTION Transformers have revolutionised sequence modelling across natural-language processing [1], [2], vision [3], [4], and speech [5]–[7] by replacing recurrent mechanics with multihead self-attention and deep feed-forward blocks. Despite their empirical dominance, the theoretical underpinnings of optimization in deep attention networks remain poorly understood, particularly in comparison to classical architectures. Concurrently, optimal transport theory provides a principled framework through Villani functions—smooth, coercive energies that ensure log-Sobolev and Poincaré inequalities for the Gibbs distribution e−f (x) , where x ∈ RN denotes the input and f : Rd → R is a smooth scalar field, enabling exponential A. Das and S. Dutta are with the Science and Technology Organization, GE HealthCare, Bangalore 560066, Karnataka, India. (E-mail: [email protected]) Corresponding author: Sayantan Dutta. A. Das and S. Dutta contributed equally to this work.
convergence guarantees for Langevin dynamics [8]. Establishing that a practical deep-learning objective belongs to this class would unlock powerful analytic tools for understanding both optimization and generalization. Practically, Transformer training augments token-level cross-entropy with L2 weight decay. Although often viewed as a heuristic for regularization [9], [10], recent evidence suggests that this quadratic term plays a geometric role, enforcing the curvature necessary for functional inequalities [11], [12]. However, a rigorous verification that full-scale transformers meet these criteria is lacking. This paper closes that gap. Problem Statement: Transformer training involves minimizing a cross-entropy objective augmented with a L2 weight-decay penalty. Specifically, we study the regularized loss:
F(θ) =
N λ 1 X − log pθ (yi | xi ) + ∥θ∥2 , N i=1 2
λ > 0,
(1)
where λ be the weight-decay factor and pθ (yi | xi ) denotes the soft-max probability assigned to the ground-truth token yi ∈ R given the input sequence xi ∈ R, and θ ∈ Rd aggregates all model parameters. The central question is: Does F(θ) satisfy Villani’s coercivity, integrability, and differential growth conditions, such that the corresponding log-Sobolev and Poincaré constants depend only on λ and d, not on the data distribution? An affirmative answer would establish a functional-analytic foundation for Transformer optimization, offering insight into the empirical efficiency of Langevin-type training schemes (e.g., SGD or Adam augmented with thermal noise) observed in large-scale language models [13]–[15]. Prior Work and Challenges: Villani–type coercivity has thus far been rigorously verified only for objectives that are either convex in the weights or possess a fixed, data-independent Jacobian. This includes models such as logistic regression, kernel machines, and two-layer ReLU networks, whose gradients are globally Lipschitz and whose Hessians exhibit quadratic growth once an L2 penalty is applied—conditions sufficient to satisfy Villani’s criteria [11], [12]. Transformers, by contrast, introduce data-dependent Jacobians through their architecture. Specifically, the multi-head self-attention √ mechanism transforms activations via sof tmax(QK ⊤ / d), where Q and K are the query and key matrices. Furthermore, LayerNorm modules dynamically rescale activations based on the input distribution at each layer. These nonlinearities
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
break the spectral assumptions leveraged in shallow-network proofs, rendering existing coercivity results inapplicable beyond depth-2 feed-forward architectures. Extensive ablation studies have shown that increasing weight decay (i) sharpens the upper tail of the Hessian spectrum, (ii) lowers the rank of attention matrices, and (iii) accelerates the decay of gradient-norm distributions [16]– [19]. These empirical trends suggest that quadratic penalties may enforce stronger curvature “at infinity”. However, such findings remain largely phenomenological: no prior study has quantitatively evaluated the Villani diagnostic, Ψs = −∆F + s−1 ∥∇F∥2 , where s is the temperature parameter and ∆, ∇ denote the Euclidean Laplacian and gradient, respectively. Furthermore, existing literature has not established a formal connection between these observed spectral patterns and the explicit log-Sobolev or Poincaré constants. This reveals a fundamental disconnect: existing theoretical guarantees apply only to shallow or linear models [12], [20], while existing empirical findings document meaningful curvature behaviors in deep attention-based networks [16]. Bridging this gap—by proving Villani coercivity and validating practical diagnostics in full-scale Transformers—is the central goal of this work. The key contributions of this work are fourfold: 1) We prove that F (θ) is a Villani function under bounded input embeddings—a natural consequence of standard tokenisation schemes. The quadratic penalty is necessary and sufficient for meeting Villani’s differential growth conditions. 2) We introduce empirical diagnostics (Sec. IV) to visualize the divergence of Ψs (θ), linking weight decay to curvature at infinity. 3) We derive explicit convergence rates for Langevin-based optimizers using log-Sobolev theory, explaining observed gains in training efficiency and generalization, and unifying recent particle-descent analyses [21], [22] with weight-decay practice (Sec. V). 4) We provide a reproducible experimental suite (Sec. VI) for evaluating functional-analytic properties in large Transformer models. The rest of this paper is organized as follows. Sec. II introduces the Villani framework and formalizes the Transformer loss as a Gibbs energy. Sec. III presents the theoretical proof, followed by geometric diagnostics (Sec. IV) and optimization consequences (Sec. V). Empirical validation appears in Sec. VI, with discussion and future directions in Secs. VII–VIII. II. T HEORETICAL P RELIMINARIES This section introduces the foundational concepts required for our theoretical results: parameter notation, Transformer forward maps, and the Villani function framework that underpins our analysis in Secs. III–V. A. Notation and Basic Operators We consolidate all trainable parameters of an L-layer Transformer (see Fig. 1)—including weight matrices {Wℓ }L ℓ=1
2
(query, key, value, feed-forward, etc.) and their corresponding bias vectors {bℓ }L ℓ=1 —into a vector: θ = W1 , . . . , WL , b1 , . . . , bL ∈ Rd , (2) where the total parameter count is given by: d=
L X
dim Wℓ + dim bℓ .
(3)
ℓ=1
We use ∥·∥ and ∥·∥F for the Euclidean and Frobenius norms. B. Transformer Forward Map in Measure Form Given a token sequence x = (x1 , . . . , xN ) from a vocabulary of size V , the Transformer computes logits zθ (x) ∈ RV through a composition of multi-head attention, residual paths, LayerNorm, and position-wise feed-forward layers. As each component is locally Lipschitz, the mapping θ 7→ zθ (x) is C ∞ for every fixed x. This smoothness property allows analysis of higher-order derivatives and Hessian behavior (Sec. II-D). Next-token probabilities follow the softmax rule: exp zy,θ (x) (4) pθ (y | x) = PV . v=1 exp zv,θ (x) This formulation links the probabilistic training objective (cross-entropy loss) with the analytic structure (gradients and Hessians) essential to our theoretical analysis. The smoothness of the forward map enables differential-geometric reasoning, while the softmax parametrization supplies the probabilistic foundation for functional inequalities. We assume token embeddings are bounded, i.e., ∥xi ∥ ≤ B for some constant B > 0. This mild condition, satisfied by common tokenization methods such as lookup tables and bytepair encoding, ensures that the quadratic penalty dominates the loss at large ∥θ∥—an essential requirement for verifying Villani conditions. C. Regularised Loss as an Energy Function Given the dataset D = {(xi , yi )}N i=1 , we define the standard training objective as: N
F(θ) =
1 X λ [− log pθ (yi | xi )] + ∥θ∥2 , N i=1 |2 {z } | {z } R(θ)
λ > 0.
(5)
L(θ)
The data-fitting term L(θ) represents the average crossentropy loss, measuring how well the network explains the observed token sequences. This term inherits the smoothness properties from the forward map θ 7→ zθ (x). Furthermore, R(θ) corresponds to the standard L2 weight decay. While commonly viewed as a regularizer, we show R(θ) plays a critical analytic role: without its quadratic growth, the loss lacks the curvature needed for Villani’s differential conditions. The bounded input assumption (∥xi ∥ ≤ B) ensures L(θ) grows at most linearly with ∥θ∥, allowing the quadratic penalty to dominate asymptotically. From a thermodynamic viewpoint, F(θ) defines a free energy functional, where L(θ) acts as internal energy determined
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
by the data distribution, while R(θ) serves as a confining potential with strength controlled by λ. This interpretation proves particularly valuable when analyzing the associated Gibbs measure exp(−βF (θ)), where β > 0 is the inverse temperature parameter, whose log-Sobolev and Poincaré constants depend exclusively on λ and d, remaining independent of the underlying data distribution. D. Villani Functions The Villani framework offers a rigorous pathway to connect geometric properties of loss functions with the convergence behavior of Langevin dynamics [8]. In our setting, this theory provides the analytic foundation for proving curvature-driven mixing guarantees under weight decay. A smooth function f : Rd → R with a temperature parameter s > 0 is called a Villani function if it satisfies the following three conditions: (i) Coercivity at infinity: lim f (x) = ∞,
∥x∥→∞
(ii) Gaussian–tail integrability: Z 2 e− s f (x) dx < ∞, ∀s > 0,
(6)
(7)
Rd
(iii) Differential growth: h i 1 lim −∆f (x) + ∥∇f (x)∥2 = ∞, ∀s > 0. (8) s ∥x∥→∞ The three Villani conditions capture distinct geometric and probabilistic properties critical for analyzing stochastic dynamics. The first condition (6) ensures that the energy landscape grows unbounded at infinity, preventing trajectories from diverging. The second condition (7) guarantees that the Gibbs measure has Gaussian-like tail decay, ensuring proper normalization. The third and most technical condition (8) ensures that curvature dominates gradient drift in the asymptotic regime by requiring a positive-definite trade-off between second-order dissipation and first-order gradient energy. These conditions collectively ensure strong functional inequalities for the associated R Gibbs measure µf (dx) = Z −1 e−f (x) dx, where Z = Rd e−f (x) dx is the partition function. In particular, the measure µf satisfies dimensionfree log-Sobolev and Poincaré inequalities, with constants determined entirely by the curvature parameters in (8). This leads to robust convergence guarantees: classical results [23], [24] show that overdamped Langevin dynamics governed by f exhibits exponential L2 -mixing: KL L(Θt ) ∥ µf ≤ e−2t/CLS (f ) KL L(Θ0 ) ∥ µf , (9) where CLS (f ) is the optimal log–Sobolev constant1 , KL denotes the Kullback–Leibler divergence, and the convergence rate is independent of the initial condition Θ0 . We show in Sec. III that F(θ) satisfies all three Villani conditions under bounded embeddings, justifying the use of Langevin-based analysis for Transformers. 1 In this context, C c > 0 such that the LS (f ) denotes the smallest constant R logarithmic Sobolev inequality Entµf (g 2 ) ≤ 2c ∥∇g∥2 dµf , ∀g, where Entµf (g 2 ) is the entropy of g 2 with respect to the probability measure µf .
3
E. Functional Inequalities and Optimisation Dynamics Villani’s conditions offer not only theoretical guarantees but also practical insights into the convergence behavior of Gaussian-noise-based optimization algorithms used in modern machine learning. For the overdamped Langevin stochastic differential equation with inverse temperature β > 0: r 2 dBt , (10) dΘt = −∇f (Θt ) dt + β where Bt is standard Brownian motion in Rd , bridging continuous-time diffusions and discrete optimization algorithms commonly used in machine learning. The entropy contraction property stems from the Bakry–Émery criterion [25], [26], which shows that the relative entropy between the distribution of Θt and the Gibbs measure µf decays exponentially, as given in eq. (9). This decay is governed by the log-Sobolev constant CLS (f ), which controls the system’s intrinsic mixing time: smaller constants yield faster convergence to equilibrium, while larger constants imply slower diffusion and exploration. This continuous-time Langevin behavior translates into practical optimization schemes via time discretization. This discretization–diffusion link has also been studied in the mean-field and Wasserstein gradient flow literature [27], [28]. Specifically, discretizing eq. (10) with step size η ≪ 1 yields stochastic gradient Langevin dynamics (SGLD), and more generally encompasses noisy versions of SGD or Adam. These algorithms thus inherit the convergence behavior of their continuous counterparts, up to a discretization error. Importantly, this discretization retains key convergence guarantees, with the log-Sobolev constant CLS (f ) governing the finite–time convergence bound of the resulting discrete iterates. This insight—formalized by Raginsky et al. [22] and later refined in [29]—connects the geometry of the loss function to the dynamics of noisy optimizers, offering theoretical justification for observed training behavior. In our setting, taking f = F allows us to express these bounds entirely in terms of the weight-decay coefficient λ and parameter dimension d. This substitution bridges functional inequalities and neural optimization, leading to concrete generalization and convergence guarantees for Transformer training. Detailed derivations under realistic mini-batch noise are provided in Sec. V. III. M ETHODOLOGY We now present a rigorous proof that the regularized Transformer objective F (θ), as defined in (5), satisfies all conditions of a Villani function. Our methodology proceeds by sequentially verifying the structural properties required for Villani’s framework: smoothness, coercivity, integrability, and the differential growth condition. We begin by establishing basic analytic properties of F = L + R, ensuring the function is well-behaved for the application of differential and probabilistic tools. This is followed by a careful treatment of the differential growth condition, where we derive explicit bounds on the Laplacian and gradient norm. Finally, we examine the necessity of the quadratic penalty
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
4
Lemma 2 (Coercivity at Infinity). The regularized loss function satisfies lim F(θ) = ∞. ∥θ∥→∞
PN Proof. Since L(θ) = N1 i=1 − log pθ (yi |xi ) ≥ 0 by the nonnegativity of negative log-likelihood terms, we obtain F(θ) = L(θ) + R(θ) ≥ R(θ) =
λ ∥θ∥2 . 2
(11)
As ∥θ∥ → ∞, the quadratic regularization term λ2 ∥θ∥2 , which is strictly convex, diverges to infinity, which forces F (θ) → ∞ as well.
Fig. 1. A single Transformer block annotated with parameter groups Wℓ , bℓ , clarifying the scope of θ. Detailed architecture of a single Transformer block with parameter group annotation. The diagram clarifies the scope of θℓ in the theoretical analysis, showing how each layer contributes 4d2 + 2d dff + 4d parameters (approximately 12d2 for standard dff = 4d). Orange-highlighted boxes indicate learnable parameters included in the weight decay term λ ∥θ∥2 . Residual connections (orange dashed lines) and nor2 malization layers (green) are crucial for the bounded gradient analysis in Supp. Mat. Appx. A. The total parameter vector θ = (θ1 , θ2 , . . . , θL ) concatenates all layer parameters, with dimension d ≈ 12L d2 for L layers. This decomposition enables the layer-wise refinements proposed in the nearterm research roadmap.
Lemma 3 (Gaussian-Tail Integrability). For every temperature parameter s > 0, the Boltzmann weight satisfies Z 2 e− s F (θ) dθ < ∞. Rd
Proof. The quadratic growth established in Lemma 2 provides the key insight: since F grows at least quadratically in the parameter norm, the corresponding Boltzmann weight e−2F /s decays at least as rapidly as a centered Gaussian distribution λ/s. More precisely, the coercivity bound yields 2
λ
2
e− s F (θ) ≤ e− s ∥θ∥ . and discuss the dimensional scaling of the associated constants. Together, these components form a cohesive theoretical argument linking the geometry of Transformer loss surfaces to the functional-inequality landscape underlying Langevin convergence guarantees. Before tackling the core differential growth analysis, we first verify three foundational properties of the regularized loss. These results follow from standard calculus and architecturespecific assumptions, but we include them here for completeness and clarity of exposition. Lemma 1 (Smoothness). The mapping θ 7→ F(θ) is infinitely differentiable on Rd , with every partial derivative of any order being locally Lipschitz continuous. Proof. Each Transformer layer is a finite composition of standard operations—matrix multiplication, addition, softmax, GELU, and normalization—all of which are analytic and hence C ∞ on R. This implies that scalar functions g1 , . . . , gm : R → R within the computational graph admit convergent power series expansions around any point in their domain. Since the class of C ∞ functions is closed under composition, the full forward map θ 7→ zθ (x) is infinitely differentiable. The same holds for the output distribution θ 7→ pθ (y | x) and its negative loglikelihood. The quadratic regularization term R(θ) = λ2 ∥θ∥2 is a polynomial and thus also C ∞ . Local Lipschitz continuity of all partial derivatives follows from the general fact that C ∞ functions are locally bounded on compact subsets, completing the argument.
(12)
The right-hand side represents a Gaussian density (up to normalization), and we can evaluate its integral using the standard formula. With α = λ/s > 0, we have Z π d/2 πs d/2 2 = . (13) e−α|θ| dθ = α λ Rd Since this integral converges for any finite dimension d, we conclude that Z πs d/2 2 e− s F (θ) dθ ≤ < ∞, (14) λ Rd establishing the required integrability condition. Together, these lemmas establish smoothness, coercivity, and integrability—covering two of Villani’s three conditions. The final step is to prove the differential growth inequality, which links local geometry to global convergence behavior. Lemma 4 (Differential Growth Condition). The regularized loss function satisfies h i lim −∆F(θ) + s−1 ∥∇F(θ)∥2 = ∞, (15) ∥θ∥→∞
for every temperature parameter s > 0. Proof. We analyze each component separately. For the regularization term R(θ) = λ2 ∥θ∥2 , direct computation yields ∇R = λ θ,
∥∇R∥2 = λ2 ∥θ∥2 ,
∆R = λ d. (16) P N For the data-fitting component L(θ) = N1 i=1 ℓi (θ) where ℓi (θ) = − log pθ (yi |xi ), the bounded input assumption |xi | ≤ B and Lipschitz continuity of softmax ensure that there exist constants C1 , C2 > 0 (independent of θ) such that ∥∇L(θ)∥ ≤ C1 ,
∆L(θ) ≤ C2 .
(17)
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
These bounds arise from the Transformer’s computational structure: bounded input features propagate through weight matrices of finite spectral norm, producing logit vectors with components differing by at most a quantity proportional to maxℓ ∥Wℓ ∥F , which yields bounded derivatives of the softmax cross-entropy loss (detailed derivation in Supp. Mat. Appx. A). Combining these results (16) and (17) and applying the Cauchy-Schwarz inequality to bound the cross term ⟨∇L(θ), ∇R(θ)⟩ ≥ −λC1 ∥θ∥, we obtain −∆F(θ) + s−1 ∥∇F(θ)∥2 1 ∥∇L(θ)∥2 = − ∆L(θ) + ∆R(θ) + s + 2 ⟨∇L(θ), ∇R(θ)⟩ + ∥∇R(θ)∥2 ≥ −C2 − λd + s−1 (−2λC1 ∥θ∥ + λ2 ∥θ∥2 ) λ2 2C1 2 = −C2 − λd + ∥θ∥ − ∥θ∥ . s λ
(18)
Therefore, for ∥θ∥ ≥ 2C1 /λ, completing the square shows 1 2 1 that ∥θ∥2 − 2C λ ∥θ∥ ≥ 2 ∥θ∥ . Therefore, λ2 ∥θ∥2 − C2 + λ d , (19) 2s which diverges to infinity as ∥θ∥ → ∞ for every s > 0. −∆F(θ) + s−1 ∥∇F(θ)∥2 ≥
5
2
For large ∥θ∥, this expression is dominated by λ2s ∥θ∥2 , while for bounded ∥θ∥, the worst-case bound occurs when the expression is minimized. Taking the reciprocal and supremum over all θ, we obtain d s 1+ . (24) CLS (F) ≤ λ λs The Poincaré constant satisfies CP (F) ≤ CLS (F) by the general relationship between log-Sobolev and Poincaré inequalities. A. Necessity of the Quadratic Term The quadratic regularization term λ2 ∥θ∥2 is essential for satisfying Villani’s conditions. When λ = 0, the expression −∆L(θ) + s−1 ∥∇L(θ)∥2 becomes bounded (since ∥∇L(θ)∥ ≤ C1 by Lemma 3), violating Condition (iii). Corollary 1. The unregularized cross-entropy loss L fails to satisfy Villani Condition (iii). Proof. For λ = 0, we have from (18), −∆L(θ) + s−1 ∥∇L(θ)∥2 ≤ −∆L(θ) + s−1 C12 < ∞ for all θ, contradicting the requirement that this expression diverges to ∞ as |θ| → ∞.
where pθ is produced by a Transformer of depth L and hidden width h. Assume all input embeddings are bounded, i.e., ∥xi ∥ ≤ B. Then F satisfies Villani conditions (i)–(iii) for every temperature s > 0. Consequently, the associated Gibbs measure µF ∝ e−F admits log-Sobolev and Poincaré constants s d CLS (F) ≤ 1+ , CP (F) ≤ CLS (F), (21) λ λs
This confirms that weight decay is not merely a regularization technique but a fundamental requirement for Langevintype analyses. Fig. 2 shows the non-divergence of −∆L(θ) + s−1 ∥∇L(θ)∥2 as ∥θ∥ grows, underscoring why weight decay is pivotal for Langevin-type analyses. The log-Sobolev constant CLS ≤ s(λ−1 + d/(λ2 s)) in (21) aligns with classical results for Gaussian measures, indicating that our bound achieves the correct scaling behavior. While future refinements—such as incorporating layer-wise curvature or structural features of the Transformer—may tighten this estimate, such directions lie beyond the scope of our present objectives. With Theorem 1 in place, we are now equipped to derive explicit convergence rates for stochastic gradient Langevin dynamics and to establish PAC-Bayesian generalization bounds.
which depend only on the weight-decay factor λ, the parameter count d, and the user-chosen s, but not on the data distribution.
IV. G EOMETRIC I NTERPRETATION OF THE V ILLANI L ANDSCAPE
Theorem 1 (Transformer Loss Is a Villani Function). Let the regularized loss N
F(θ) =
1 X λ [− log pθ (yi |xi )] + ∥θ∥2 , N i=1 2
λ > 0,
(20)
Proof. The proof directly follows from Lemmas 1–4, which establish that F(θ) satisfies all three Villani conditions under the bounded embedding assumption: smoothness is shown in Lemma 1, coercivity in Lemma 2, integrability in Lemma 3, and differential growth in Lemma 4. Having verified all three Villani conditions, we can apply Villani’s theorem (Cor. 2.3.4 in [30]) to obtain explicit bounds on the functional inequality constants. The log-Sobolev constant satisfies [24] s . (22) CLS (F) ≤ sup −1 ∥∇F(θ)∥2 θ −∆F(θ) + s 1 From the Lemma 4, we established that for ∥θ∥ ≥ 2C λ ,
−∆F(θ) + s−1 ∥∇F(θ)∥2 ≥
λ2 ∥θ∥2 − (C2 + λd). 2s
(23)
Having established the analytical proof, we now provide geometric intuition for why quadratic regularization transforms the Transformer objective into a Villani function. A. From Flat Valley to Confining Bowl The transformation from unregularized to regularized Transformer objectives can be understood geometrically as the transition from a flat valley to a confining bowl. Without weight decay (λ = 0), the cross-entropy surface exhibits asymptotic flatness: parameters can grow without bound while leaving logits unchanged due to scaling symmetries in attention mechanisms and softmax normalization. This creates a valley-like landscape where gradients vanish and the Laplacian remains bounded, violating Villani’s differential growth condition (Fig. 3(a)).
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
6
is an unbiased estimation with variance decreasing as M −1 . Detailed variance analysis is provided in Supp. Mat. Appx. B. These theoretical predictions—particularly the quadratic beb s (θ) and the spectral inflation of the Hessian—are havior of Ψ empirically validated in Sec. VI, where diagnostic estimates and spectral statistics confirm the curvature-driven behavior across multiple regularization regimes. D. From Geometry to Optimization
b s (θ) vs. ∥θ∥2 along five radial rays for λ ∈ {0, 10−4 , 10−2 }. Fig. 2. Ψ Empirical verification of Villani Condition (8) via the diagnostic scalar field Ψs (θ) = −∆F (θ) + s−1 ∥∇F (θ)∥2 plotted against ∥θ∥2 along random radial rays for GPT-Neo-125M. Each line style represents a different radial direction. For λ = 0 (blue), the field saturates, confirming failure of the Villani condition. For λ = 10−4 (orange), mild upward growth begins beyond ∥θ∥2 ≈ 105 . For λ = 10−2 (green), clear quadratic divergence emerges with slope ≈ λ2 /s, validating the theoretical prediction from eq. (19).
The addition of the quadratic penalty λ2 ∥θ∥2 fundamentally alters this geometry by superimposing a radially symmetric paraboloid. Every direction now exhibits positive curvature at a sufficient distance from the origin, creating a bowl-shaped landscape whose confinement strength increases quadratically with parameter norm (Fig. 3(b)). This geometric transformation is the key to satisfying Villani’s conditions and enabling exponential mixing guarantees. B. Villani Diagnostic Function The key analytical tool for understanding confinement is the scalar field: 1 Ψs (θ) = −∆F (θ) + ∥∇F(θ)∥2 , s > 0. (25) s Condition (iii) requires Ψs (θ) → ∞ as ∥θ∥ → ∞, making Ψs a direct measure of landscape confinement. Steeper radial growth of Ψs corresponds to stronger log-Sobolev constants and faster convergence rates. For large parameter norms, our analysis in Sec. III shows that:
The geometric transformation has direct algorithmic consequences: (1) Mixing Time: The positive log-Sobolev gap guarantees exponential convergence to equilibrium within O(CLS log(1/ε)) iterations. (2) Parameter Scaling: Higher regularization λ improves mixing speed via larger Ψs values (since Ψs ∝ λ2 ∥θ∥2 ), but risks overshooting near minima due to overly aggressive curvature. (3) Adaptive Schedules: The quadratic confinement suggests learning rate schedules should account for the interaction between regularization strength and parameter growth. Sec. VI presents a controlled experiment that systematically sweeps over both λ and learning-rate warm-up durations to empirically characterize this trade-off. While our analysis provides geometric intuition, it has limitations: radial sampling overlooks symmetries from LayerNorm and residual connections [31]; Hessian trace estimators approximate only dominant curvature; and the bowl approximation is most accurate in high-norm regimes dominated by regularization. Nonetheless, weight decay fundamentally reshapes the loss from a flat valley to a confining bowl, underpinning our convergence guarantees. V. O PTIMIZATION AND G ENERALISATION C ONSEQUENCES
λ2 ∥θ∥2 − C (26) 2s for some constant C, confirming the required divergence behavior.
The Villani structure established in Sec. III transforms the regularized Transformer loss into a thermodynamically confining energy function. This geometric property enables a suite of functional-analytic tools that translate local curvature into global convergence rates and generalization guarantees. We formalize these implications through three perspectives: finite-time convergence bounds for noisy gradient descent, PAC-Bayesian generalization with a Villani-induced prior, and actionable hyperparameter tuning principles.
C. Computational Estimation
A. Noisy Gradient Descent as Discretized Langevin Dynamics
Since direct computation is intractable for modern networks (d > 108 ), we employ stochastic methods using Hutchinson’s trace estimator:
Modern stochastic gradient training with mini-batch variance Σ(θ) connects naturally to continuous-time diffusion via the Euler–Maruyama discretization. Consider the overdamped Langevin stochastic differential equation (SDE):
Ψs (θ) ≥
M 1 X (m) 2 d ∆F(θ) = ⟨v , ∇ F(θ)v (m) ⟩, M m=1
(27) dΘt = −∇F(Θt )dt+
where v (m) ∼ N (0, Id ) are random Gaussian vectors. The diagnostic 1 b s (θ) = −∆F(θ) d Ψ + ∥∇F(θ)∥2 s
(28)
p
2β −1 dBt , β −1 =
η E[∥ξ∥2 ] , (29) 2d
where Bt is standard Brownian motion in Rd , η is the learning rate [32], and ξ models mini-batch gradient noise. The temperature parameter β −1 arises from the scaling limit
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
7
Fig. 3. Comparison of Unregularized (“Valley”) vs. Regularized (“Bowl”) Loss Landscapes. Top: 3D visualization of the Transformer loss surface F (θ) in a random subspace around the initialization point θ0 . Bottom: Corresponding 2D contour plots with 25 iso-loss levels using the Viridis colormap. Left (λ = 0): Without regularization, the loss landscape forms a narrow valley with elongated contours, multiple equivalent minima, and flat regions with weak gradients—hindering convergence and mixing. Right (λ = 10−3 ): Applying weight decay reshapes the landscape into a well-defined bowl with circular contours, a unique global minimum, and strong radial gradients. The added quadratic penalty λ ∥θ∥2 removes flat regions and induces Villani coercivity, 2 enabling exponential mixing and faster convergence.
of SGD. When F satisfies Villani’s conditions, the associated Gibbs measure µF (dθ) = Z −1 exp[−βF(θ)]dθ,
(30)
obeys a log-Sobolev inequality with constant CLS (βF) = β −1 CLS (F). By Bakry–Émery theory [25], the Langevin dynamics exhibit exponential entropy contraction: KL L(Θt ) ∥ µF 2t ≤ exp − KL L(Θ0 ) ∥ µF , (31) CLS (β F ) where L(Θt ) denotes the distribution of the stochastic process at time t. B. Finite-Time Optimization Guarantees The continuous-time mixing property translates into concrete optimization bounds under discretizations. We now formalize this connection for noisy gradient descent.
Theorem 2 (Exponential Mixing of Noisy Gradient Descent). Consider the regularized Transformer objective F (θ) with weight decay λ > 0 and target accuracy ε > 0. Assume fixed mini-batch size m and the noise covariance satisfies Tr Σ(θ) ≤ σ 2 uniformly. Let the learning rate satisfy λ (32) η ≤ min 1, 2 . 4σ Then, using stochastic gradient descent with optional injected p Gaussian noise 2ηβ −1 Zk where Zk ∼ N (0, Id ), after CLS (F) F0 − F ∗ T ≥ log , (33) 2η ε iterations, the expected suboptimality satisfies E[F(θT ) − F ∗ ] ≤ ε,
(34)
where F0 = F(θ0 ) and F ∗ represents the global minimum.
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
8
Fig. 4. Spectral radius of ∇2 F vs. ∥θ∥ for λ > 0. Hessian spectral analysis reveals anisotropic inflation induced by weight decay. Left: Evolution of top-20 eigenvalues vs. parameter norm ∥θ∥2 . The bulk spectrum (gray region) remains data-controlled while the spectral tail (yellow region) exhibits linear growth λmax ∝ λ ∥θ∥ for λ > 0. This anisotropic behavior explains improved global—but not local—strong convexity. Right: Spectral ratio λmax /λmedian showing the emergence of condition-number anisotropy. Stronger regularization (larger λ) creates higher spectral ratios, accelerating Langevin mixing but potentially affecting numerical conditioning. The confidence bands show measurement uncertainty from stochastic Lanczos quadrature. This analysis supports the geometric interpretation in Sec. VI-D and explains why quadratic regularization improves optimization tractability while preserving representational flexibility in the bulk directions.
Fig. 5. Optimization Speed: Empirical vs. Theoretical Bounds. Empirical convergence of noisy SGD training compared with theoretical bounds from Theorem 2. Left: Penn Treebank (small dataset, N ≈ 42,000 tokens) shows clear dependence of convergence rate on weight decay strength λ. Right: WikiText-103 (large dataset, N ≈ 103 × 106 tokens) exhibits smoother curves and tighter agreement with theory. The shaded envelope represents the exponential bound from eq. (41) with CLS (F ) estimated from λ and model dimension d = 1.25 × 108 . Stronger regularization (larger λ) accelerates convergence by improving the log-Sobolev constant, while gradient noise injection (purple dashed line) provides additional validation of the Langevin dynamics connection. Error bars show ±1 s.e.m. across three random seeds.
Proof. Since F satisfies the Villani conditions with λ > 0, the Gibbs measure µ∗ ∝ e−βF admits a log-Sobolev inequality, Z CLS (βF) Entµ∗ (g) ≤ ∥∇g∥2 g dµ∗ ∀g > 0, (35) 2 with constant CLS (βF) = β −1 CLS (F). Furthermore, ∥∇F(θ) − ∇F (θ′ )∥ ≤ L∥θ − θ′ ∥, since F is L-smooth. The discrete-time Langevin update p θk+1 = θk − η(∇F (θk ) + ξk ) + 2ηβ −1 Zk , (36) can be interpreted as a noisy Euler step for discretizing the Langevin SDE, where ξk denotes stochastic gradient noise satisfying E∥ξk ∥2 ≤ σ 2 , and Zk ∼ N (0, Id ) is standard
Gaussian noise. Under the log-Sobolev inequality (35), and L-smoothness of F, the Donsker–Varadhan variational formula yields a fundamental one-step KL contraction (cf. [22], Lemma 3.2): KL(νk+1 ∥µ∗ ) ≤ 1−
2η CLS (βF)
KL(νk ∥µ∗ ) +
η 2 Lσ 2 . 2β
(37)
where νk = L(θk ). Under β = λ−1 , the learning rate bound (32) implies the stability condition η ≤ CLS (βF)/4. Iterating
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
the KL contraction bound (37) over T steps yields: T 2η KL(νT ∥µ∗ ) ≤ 1 − KL(ν0 ∥µ∗ ) CLS (βF) ηLσ 2 CLS (βF) + , 4β
9
(38)
where νT = L(θT ) is the law of the iterate after T steps. The first term reflects geometric decay of KL divergence to the stationary distribution, while the second captures steadystate error due to stochastic noise. To relate this to expected objective suboptimality, we invoke a fundamental inequality for Gibbs measures: KL(ν∥µ∗ ) = β Eν [F] − F ∗ + log Z − Eν [log Z] ≥ β Eν [F] − F ∗ , (39) by Jensen’s inequality, since log Z is constant and − log is convex [33]. Applying this to νT , and rearranging: E[F(θT ) − F ∗ ] ≤ β −1 KL(νT ∥µ∗ ).
(40)
Substituting the bound from (38) and setting β = λ−1 (to match the Villani prior) gives: T 2η ∗ (F0 − F ∗ ) E[F(θT ) − F ] ≤ λ 1 − CLS (βF) ηLσ 2 + . (41) 4 ∗ Thus, choosing T ≥ CLS2η(F ) log 2(F0ε−F ) , ensures the first term in (41) is bounded by ε/2. Simultaneously, the learning rate condition η ≤ λ/(4σ 2 ) ensures the noise term is also ≤ ε/2. Together, they yield the final bound E[F (θT ) − F ∗ ] ≤ ε, completing the proof.
Fig. 6. PAC-Bayesian Generalization Bounds vs. Observed Test Performance. Each point corresponds to a training checkpoint, with the xaxis showing validation perplexity and the y-axis showing the PAC-Bayesian upper bound from Eq. (52), using β = λ−1 . Points below the diagonal line y = x indicate conservative (loose) bounds. Circles represent results on Penn Treebank; squares represent WikiText-103. Stronger weight decay (λ = 10−2 ) yields tighter bounds that lie closer to the diagonal, consistent with the quadratic scaling predicted by Theorem 3. The average bound tightness improves from 1.43× under weak regularization to 1.12× under strong regularization. A high correlation (R2 ≈ 0.93) between the bound and observed performance confirms that Villani-based PAC-Bayes theory effectively captures generalization behavior across datasets and model configurations. Error bars denote ±1 s.e.m. across random seeds.
at least 1 − δ over the random draw of the training set D of size n, the population risk of the learned parameter θb satisfies v √ u u KL(δ b∥µF ) + log 2 n t δ θ b ≤ R( b + b θ) R(θ) , 2(n − 1)
(43)
C. PAC-Bayesian Generalization with Villani Prior The Villani structure established in Sec. III provides a principled basis for PAC-Bayesian analysis. Specifically, the regularized loss F(θ) defines a Gibbs prior with strong tail decay and curvature, making it a natural candidate for deriving generalization bounds. Let D = {(xi , yi )}ni=1 denote n i.i.d. samples from an unknown data distribution ρ, and define the population and b empirical risks R(θ) = E(x,y)∼ρ [− log pθ (y|x)] and R(θ) = Pn 1 − log p (y |x ), respectively. θ i i i=1 n We use the Gibbs measure as prior: P (dθ) = µF (dθ) = Z −1 exp(−βF(θ))dθ,
(42)
where β > 0 is the inverse temperature and Z is the partition function. For the posterior, we take Q = δθb, a Dirac measure on the final SGD iterate. The fundamental PACBayesian framework of Catoni [34] provides the following sharp generalization bound: Theorem 3 (PAC-Bayesian Generalization with Villani Prior). Let F (θ) be the regularized Transformer objective with λ > 0 and P = µF as in (42). For any δ ∈ (0, 1), with probability
where the KL divergence admits the explicit upper bound d β b KL(δθb∥µF ) ≤ βF (θ) + log 1 + + K0 , 2 λ with K0 = d2 log stant.
λ 2πβ
(44)
being a dimension-dependent con-
Proof. The proof proceeds in two stages: first establishing the general PAC-Bayesian bound, then deriving the explicit KL divergence bound using the Villani structure. Stage 1: PAC-Bayesian Bound Framework: Catoni’s theorem [34] establishes that for any prior P and data-dependent posterior Q, with probability at least 1 − δ, v √ u u KL(Q∥P ) + log 2 n t δ b . Eθ∼Q [R(θ)] ≤ Eθ∼Q [R(θ)] + 2(n − 1) (45) b For the Dirac posterior Q = δθb, we have Eθ∼Q [R(θ)] = R(θ) b b b and Eθ∼Q [R(θ)] = R(θ), yielding (43).
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
10
Stage 2: KL Divergence Computation: For a Dirac measure δθb, the KL divergence between the point-mass posterior and the Gibbs prior P = µF is given by: dµF KL(δθb∥µF ) = − log , (46) dθ θ=θb where µF (dθ) = Z −1 exp(−βF(θ))dθ is the Gibbs prior (42) and Z is the partition function. Thus, b + log Z. KL(δθb∥µF ) = βF (θ)
(47)
Stage 3: Bounding the Partition Function: Using the lower bound F(θ) ≥ λ2 ∥θ∥2 , we obtain: Z Z βλ 2 Z= exp(−βF(θ))dθ ≤ exp − ∥θ∥ dθ 2 Rd Rd d/2 2π . (48) = βλ For the lower bound on Z, we exploit the boundedness of L(θ) established in Lemma 3. There exists C > 0 such that L(θ) ≤ C for all ∥θ∥ ≤ R, for some R > 0. Thus, Z Z≥ exp(−βF(θ))dθ ∥θ∥≤R Z βλ (49) ≥ exp(−βC) exp − ∥θ∥2 dθ. 2 ∥θ∥≤R This integral captures most of the mass of a Gaussian when R is sufficiently large, and up to exponentially small errors, it approximates the full integral. Combining with the earlier upper bound and using standard Gaussian tail behavior, we obtain d 2π log Z ≤ log + O(1). (50) 2 βλ Refining this via the Laplace method around the minimum of F (near the origin due to the regularizer), we find d 2π d β log Z ≤ log + log 1 + + O(log d). (51) 2 βλ 2 λ The second term captures the influence of data-fitting, while O(log d) collects higher-order terms. Substituting into (47) and absorbing constants into K0 yields the final KL bound (44), completing the proof. Corollary 2. Setting β = λ−1 in Theorem 3 minimizes the leading coefficient in the KL bound. Under this choice, the generalization bound becomes: v √ u u λ−1 F(θ) b + d log(2) + K0 + log 2 n 2 δ b ≤ R( b +t b θ) R(θ) , 2(n − 1) (52) 2 λ where K0 = d2 log 2π . The derived bound reveals several important properties. It decays quadratically with the weight decay parameter λ, providing theoretical support for the empirical observation that stronger regularization enhances generalization. Its dependence on the parameter dimension d is only logarithmic—optimal for log-concave priors—making the bound practical in high-dimensional regimes. Moreover, the formulation
Fig. 7. Correlation between Villani diagnostic Ψs (θ) = −∆F (θ) + s−1 ∥∇F (θ)∥2 and validation perplexity across training checkpoints. Point size encodes training epoch (larger = later), color encodes weight decay strength λ. A clear monotonic relationship emerges: higher Ψs values correlate strongly with lower perplexity, establishing an empirical link between geometric coercivity and generalization performance. Correlation strength increases with λ: R2 = 0.94 for λ = 10−2 versus R2 = 0.23 for λ = 0, confirming that stronger regularization creates more reliable Ψs -performance relationships. The optimal region (green) shows high Ψs corresponding to strong Villani coercivity and excellent generalization. The suboptimal region (red) shows low Ψs , indicating weak coercivity and poor performance. Outliers (red borders) mark training instabilities that disrupt the correlation. This validates Ψs as a practical diagnostic for both optimization quality and generalization potential during training.
interpolates smoothly between empirical risk minimization (large λ) and Bayesian inference (small λ), unifying both perspectives. This result establishes a direct connection between the geometric structure of the Transformer loss landscape—captured by the Villani conditions—and generalization performance. It offers a formal justification for tuning weight decay as a primary tool for managing the optimization–generalization trade-off in deep learning. D. Design Guidelines for Practitioners The theoretical framework developed in this section leads to actionable hyperparameter recommendations for training large-scale Transformers. Table I summarizes these insights, grounded in the Villani structure, and highlights concrete links between algorithmic settings and theoretical convergence guarantees. These heuristics align with empirical practices in large language model training. The learning rate constraint allows more aggressive step sizes when regularization is properly tuned. The weight decay criterion ensures that optimization begins in the quadratic-dominant region, where convergence guarantees are strongest. Moreover, our framework unifies insights from diffusionbased fine-tuning [35] and ridge-regularized inference [36], where the Gaussian-like prior structure satisfying log-Sobolev inequalities supports principled noise scheduling. By providing explicit constants, this analysis bridges theoretical guarantees and practical hyperparameter tuning in Transformer optimization.
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
11
Fig. 8. Evolution of Hutchinson trace estimator distributions across training phases for λ = 10−3 . Each histogram shows 64 independent trace samples v T ∇2 F (θ)v at representative checkpoints. Left: Initial training (steps 0–4k) exhibits higher variance (CV = 12.3%) due to parameter initialization effects. Center: Mid-training (steps 20k–30k) shows stabilized variance (CV = 4.9%) as the quadratic penalty begins to dominate. Right: Final training (steps 60k+) achieves minimal variance (CV = 5.0%) with the trace converging to λ d ≈ 58 × 108 . All distributions pass normality tests (Shapiro–Wilk p > 0.3), validating d] the Gaussian assumption underlying our M = 64 probe choice. The progressive tightening of distributions confirms the theoretical prediction that Var[∆F scales as λ2 ∥θ∥2 from eq. (57). TABLE I H YPER - PARAMETER VALUES AND V ILLANI - GUIDED TUNING PRINCIPLES .
Fig. 9. Evolution of coefficient of variation (CV = σ/µ) for Hutchinson trace estimates throughout training. The plot validates the choice M = 64 probe vectors by showing CV stabilization around the theoretical prediction p 2d/M ≈ 4.9%. Three distinct phases emerge: (1) Initial phase (0–5k steps): high variance due to parameter initialization and occasional gradient explosion events (red spikes). (2) Transition phase (5k–35k steps): CV decreases as the quadratic penalty λ ∥θ∥2 begins to dominate the Hessian 2 spectrum. (3) Stable phase (35k+ steps): CV plateaus near the theoretical limit, confirming that M = 64 provides < 5% relative error. Stronger regularization (larger λ) accelerates convergence to the stable regime and achieves slightly lower final CV values. The 95% confidence band (blue shading) shows estimation uncertainty. This analysis justifies the constantcost diagnostic schedule used throughout the paper.
VI. E XPERIMENTAL R ESULTS We empirically evaluate the theoretical insights derived in earlier sections using GPT-Neo-125M, focusing on how weight decay shapes optimization dynamics, curvature properties, and generalization behavior. All experiments vary the regularization coefficient λ ∈ {0, 10−4 , 10−3 , 10−2 }, and explore connections between theoretical diagnostics and practical training outcomes. A. Implementation Details All experiments are conducted using PyTorch 2.3 and Hugging Face Transformers v4.42 on a single
Parameter
Values Used
Theoretical Guideline
Rationale
Learningrate η Weightdecay λ
5 × 10−5 , 10−4 , 2×10−4 0, 10−4 , 10−3 , 10−2
η ≲ λ/(4σ 2 ) Pick λ so that λ∥θ0 ∥2 ≈ 10L(θ0 )
Batch size
512 tokens
—
AdamW β
(0.9, 0.95)
—
Gradient noise σ 2
Natural SGD N (0, 0.052 )
or
β −1 = ησ 2 /(2d), match with λ
Epochs
40 (PTB), (WikiText-103)
5
—
Ensures Langevin entropy contraction. Enforces F coercivity via quadratic dominance. Balances memory and gradient noise. Standard optimizer setting. Aligns Langevin scale with PACBayes prior. Ensures cross-dataset convergence.
NVIDIA A100-80GB GPU (CUDA 12.2). We use the GPT-style decoder EleutherAI/gpt-neo-125m with 12 layers and 125M parameters, employing bfloat16 precision and fp32 master weights with the AdamW optimizer. Hessian-vector products are computed using torch.autograd.functional.hvp and Hutchinson trace estimation via the BackPACK library [37]. 1) Datasets: Table II lists the corpora used in our evaluations. The Penn Treebank (PTB) benchmark stresses generalisation in a small-data regime, while WikiText-103 serves as a large-scale test of optimisation dynamics over 108 tokens. 2) Training Protocol: The training setup is summarized in Table I. A linear warm-up over 5,000 steps is followed by cosine decay to zero. Hyperparameter choices adhere to the Villani-informed guidelines described in Sec. V-D. 3) Evaluation Metrics: We monitor training loss, Villani coercivity Ψs (θ), spectral diagnostics via top-20 Hessian eigenvalues, validation perplexity, and the PAC-Bayesian bound defined in eq. (52). Efficiency is assessed by wall-clock time to reach target perplexity: < 90 on PTB and < 30 on
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
TABLE II DATASETS USED IN EXPERIMENTS . Dataset
Tokens
Task
Citation
Penn Treebank WikiText-103
887k 103M
Language Modeling Language Modeling
[38], [39] [40], [41]
WikiText-103. B. Loss Landscape Geometry Fig. 3 visualize how weight decay alters the Transformer loss surface. In the unregularized case (λ = 0), the landscape exhibits narrow valleys with weak gradients and elongated contours—features that hinder convergence and slow down mixing. Applying a quadratic penalty (λ = 10−3 ) converts this geometry into a symmetric bowl with circular iso-contours and stronger radial gradients. This confirms that weight decay enforces the coercivity and curvature required by Villani conditions, facilitating efficient Langevin mixing as discussed in Sec. IV. C. Villani Coercivity Diagnostics To empirically validate the Villani coercivity condition, we examine the scalar diagnostic field Ψs (θ) = −∆F (θ) + s−1 ∥∇F(θ)∥2 , which encodes the differential growth criterion b s (θ) evaluated along from Villani’s definition. Fig. 2 shows Ψ five randomly sampled radial directions for three values of the regularization coefficient λ ∈ {0, 10−4 , 10−2 }, with s = 1 fixed. For λ = 0, the diagnostic function saturates at a finite value, indicating failure of coercivity and violating Condition (iii) from Theorem 1. In contrast, for λ = 10−4 , the field begins to show mild upward curvature beyond ∥θ∥2 ≈ 105 , suggesting a regime where the regularization starts dominating the datafitting loss. At λ = 10−2 , a clear quadratic growth pattern emerges with a slope consistent with the theoretical prediction λ2 /s from eq. (19). This confirms that sufficiently large weight decay enforces the required coercive geometry in practice. These observations match the transition thresholds predicted by the theory and support the use of Ψs (θ) as an effective empirical proxy for certifying Villani-type curvature properties in large-scale neural networks. D. Spectral Geometry via the Hessian We complement scalar diagnostics with spectral analysis of the Hessian ∇2 F(θ) using stochastic Lanczos quadrature (SLQ). Fig. 4 displays the top-20 eigenvalues and the evolution of the spectral radius as a function of parameter norm ∥θ∥, across training checkpoints. For λ > 0, the spectral radius grows linearly with ∥θ∥, while the bulk of the spectrum remains bounded and controlled by data. This anisotropic inflation highlights how weight decay selectively sharpens the top eigenmodes, improving global strong convexity without compromising the flexibility of flatter directions. The ratio λmax /λmedian increases with
12
λ, demonstrating condition-number anisotropy consistent with theoretical predictions from Sec. III. This behavior affirms the geometric transformation from a flat valley to a coercive bowl and underscores the dual role of regularization in stabilizing optimization and retaining representational capacity. E. Optimization Speed and Langevin Mixing We empirically assess the optimization efficiency predicted by Villani-based analysis using training runs on both Penn Treebank and WikiText-103. Fig. 5 compares observed training loss curves against theoretical convergence bounds derived from Theorem 2. On Penn Treebank (left panel), convergence speed exhibits clear dependence on the regularization strength λ, consistent with theoretical predictions. As λ increases, convergence accelerates due to improved log-Sobolev constants. On the larger WikiText-103 dataset (right panel), the training curves are smoother, and empirical results align closely with the predicted exponential rate. The shaded envelopes represent the theoretical upper bounds computed using estimated constants CLS (F) from the regularization and model dimension (d ≈ 1.25×108 ). Additionally, training runs with Gaussian noise injection (purple dashed lines) mimic Langevin dynamics and validate the role of noise in achieving consistent mixing behavior. Error bars indicate standard error across three random seeds. These results reinforce the quantitative utility of Villani theory in predicting and guiding Transformer optimization dynamics. F. PAC-Bayesian Generalization Bounds We evaluate the predictive accuracy of PAC-Bayesian bounds derived from Theorem 3 using training checkpoints from GPT-Neo-125M under various regularization levels. Fig. 6 compares the theoretical upper bounds on population risk to observed validation perplexity across checkpoints and datasets. Each point in the figure represents a checkpoint, where the x-axis reflects the actual perplexity and the y-axis reports the PAC-Bayes bound evaluated with inverse temperature β = λ−1 . The results demonstrate that as λ increases, the bound tightens and more closely tracks the empirical performance. For strong regularization (λ = 10−2 ), the bounds lie closer to the diagonal y = x, indicating minimal slack and strong predictive power. Quantitatively, the mean bound-to-perplexity ratio improves from 1.43× under weak regularization to 1.12× under strong regularization. High correlation scores (R2 ≈ 0.93) between the bound and validation performance further confirm that Villani-based PAC-Bayesian theory provides an effective framework for understanding generalization in deep Transformers. G. Correlation with Generalization Performance Beyond validating coercivity, we investigate whether the diagnostic field Ψs (θ) correlates with generalization outcomes during training. Fig. 7 illustrates this relationship by plotting
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
Ψs (θ) against validation perplexity across training checkpoints. Each point is colored by weight decay λ and sized by training epoch to visualize progression over time. The results reveal a strong monotonic correlation between higher values of Ψs and improved generalization (i.e., lower perplexity), particularly for larger regularization values. For λ = 10−2 , the diagnostic reliably tracks generalization with a high coefficient of determination (R2 = 0.94), while the correlation weakens for smaller or absent regularization (R2 = 0.23 for λ = 0). The plot identifies two distinct regimes: an “optimal” zone with high Ψs and strong generalization, and a “suboptimal” region with low Ψs and poor performance. Isolated outliers, typically arising from unstable optimization dynamics, fall outside this trend and are visually marked. These findings support the interpretation of Ψs (θ) not only as a theoretical coercivity diagnostic, but also as a practical tool for guiding training-time decisions and identifying promising checkpoints.
13
and blockwise Hessian-vector products) enable extension up to 100B+ models. These strategies reduce memory requirements by up to 90%, making scalable implementation of Villani diagnostics feasible. Together, these results suggest that the geometric analysis framework presented in this paper remains applicable—and computationally tractable—at scales relevant to state-of-theart language models. VII. D ISCUSSION Our experimental results confirm that Villani-based analysis not only captures theoretical curvature and coercivity properties but also translates into practical diagnostics for optimization and generalization. The empirical scaling of Ψs (θ), alignment with Hessian spectra, and match between theoretical bounds and convergence behavior substantiate the analytical framework at scale. A. Limitations
H. Hutchinson Trace Estimation We assess the efficiency and reliability of Hutchinson trace estimation throughout training. This diagnostic estimates the Laplacian term ∆F(θ) in the Villani scalar field Ψs (θ) using randomized projections. Figs. 8 and 9 illustrate the distribution and variance of these trace estimates across training epochs for λ = 10−3 . Empirical distributions of the trace samples remain approximately Gaussian throughout training, with a coefficient of variation (CV) stabilizing p around 5%, in agreement with the theoretical prediction 2d/M for M = 64 probe vectors. Variance decreases as training progresses, consistent with the increasing dominance of the quadratic penalty. Three distinct training phases emerge: (1) an initial phase with high variance due to parameter initialization; (2) a transition phase where regularization begins to control curvature; and (3) a stable phase where the trace converges to λd. This behavior supports the use of fixed-cost, low-variance diagnostic schedules for real-time geometric analysis during training. A complete aggregation of trace statistics is included in Supp. Mat. Appx. B. I. Scalability and Diagnostic Forecasting We evaluate the scalability of the Villani framework to larger model sizes using both theoretical extrapolation and empirical diagnostics. Fig. 10 presents two key analyses: the scaling behavior of the log-Sobolev constant CLS with model dimension d, and the computational feasibility of Hutchinson trace estimation across parameter scales. On the left, we show that the bound CLS ≤ s(λ−1 +d/λ2 s) predicts favorable scaling properties under strong regularization (λ = 10−2 ), even as d approaches 1012 . Empirical validation from GPT-Neo-125M confirms this trend in the small-to-medium regime. On the right, we analyze the practical cost of trace estimation. While a single GPU supports up to ∼ 1B parameters, distributed memory architectures and algorithmic optimizations (e.g., mixed precision, gradient checkpointing,
Despite encouraging results, our analysis is limited to decoder-style Transformers with uniform weight decay and assumes careful temperature tuning for Langevin alignment. The derived log-Sobolev constants may be loose in high dimensions, and direct validation at > 7B scale would require distributed Hessian-vector product estimation, which is not yet implemented. Structured regularization schemes and encoderdecoder architectures remain outside the provable regime. B. Ethical Considerations This work advances theoretical and practical tools for optimizing large-scale Transformers, promoting safer and more efficient training. While improved convergence and generalization diagnostics can reduce computational waste and privacy risks, they may also accelerate scaling and deepen compute access inequalities. We advocate integrating these methods within governance frameworks focused on fairness, sustainability, and responsible deployment. C. Open Challenges Several open challenges remain. Extending coercivity diagnostics to encoder-decoder and multimodal models, handling structured or adaptive regularization schemes (e.g., AdamW, LoRA), and improving spectral gap estimation on hardware accelerators are key directions. Applying Villani-based tools to reinforcement learning and non-convex regimes also warrants further study. Another promising direction involves exploring quantum optimization and signal processing frameworks, which have shown success in imaging contexts [42]–[48]. VIII. C ONCLUSIONS We introduced a principled framework for understanding generalization and optimization in Transformer models through Langevin dynamics and Villani coercivity. By linking L2 weight decay to exponential mixing and PAC-Bayesian generalization, we derived theoretical guarantees and validated them empirically on GPT-Neo-125M. Our results highlight
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
14
Fig. 10. Scalability analysis of the Villani framework across model sizes from 125M to 1T+ parameters. Left: Log-Sobolev constant CLS scaling with model dimension d. The validated region (green) shows empirical results from GPT-Neo-125M. The extrapolated region (yellow) uses theoretical bounds CLS ≤ s λ−1 + d/(λ2 s) to predict behavior for 1B–100B models. The frontier region (red) represents computational challenges for models > 100B parameters. Stronger regularization (λ = 10−2 ) provides better scaling properties. Right: Computational cost analysis for Hutchinson trace estimation. SingleGPU memory limits constrain analysis to ∼ 1B parameters, while distributed approaches enable validation up to ∼ 100B parameters. Optimization techniques (mixed precision, checkpointing, block-wise HVP) reduce memory requirements by 40–90%.
how weight decay induces coercive curvature, enables predictive diagnostics, and supports scalable trace estimation. This unified analysis connects geometric structure to both convergence and generalization, offering a tractable and interpretable path forward. Extensions to structured regularization, reinforcement learning, and other domains remain compelling directions for future work. IX. ACKNOWLEDGMENT The authors acknowledge the use of AI-based tools, including ChatGPT [49] and Microsoft Copilot [50], to assist in language refinement during manuscript preparation. All scientific content, mathematical derivations, and experimental results were developed independently by the authors. R EFERENCES [1] D. Zhou, Y. Zhang, Y. Wang, J. Ning, H.-J. Ye, D. Zhan, and Z. Liu, “Learning without forgetting for vision-language models,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 47, no. 6, pp. 4489–4504, 2025. [2] J. Liu et al., “Graph foundation models: Concepts, opportunities and challenges,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 47, no. 6, pp. 5023–5044, 2025. [3] J. Park, S. Yun, H. Park, J. Kang, J. Jeong, K.-M. Kim, J.-W. Ha, and H.J. Kim, “Deformable graph transformer,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 47, no. 7, pp. 5385–5396, 2025. [4] Y. Tang, M. Liu, B. Li, Y. Wang, and W. Ouyang, “Nas-ped: Neural architecture search for pedestrian detection,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 47, no. 3, pp. 1800– 1817, 2025. [5] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A.N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” in Advances in Neural Information Processing Systems. 2017, vol. 30, pp. 5998–6008, Curran Associates, Inc. [6] A. Dosovitskiy et al., “An image is worth 16x16 words: Transformers for image recognition at scale,” in International Conference on Learning Representations, 2021. [7] Anmol Gulati et al., “Conformer: Convolution-augmented transformer for speech recognition,” in Proc. Interspeech, 2020, pp. 5036–5040.
[8] L. Ambrosio, L.A. Caffarelli, Y. Brenier, G. Buttazzo, C. Villani, and S. Salsa, Optimal Transportation and Applications, Lecture Notes in Mathematics. Springer, 2003. [9] A. Krogh and J.A. Hertz, “A simple weight decay can improve generalization,” in Advances in neural information processing systems, 1992, pp. 950–957. [10] I. Loshchilov and F. Hutter, “Decoupled weight decay regularization,” in International Conference on Learning Representations, 2019. [11] A. Nitanda, A. Lee, D. Kai, M. Sakaguchi, and T. Suzuki, “Propagation of chaos for mean-field langevin dynamics and its application to model ensemble,” arXiv preprint arXiv:2502.05784, 2025. [12] L. Chizat, “Mean-field langevin dynamics: Exponential convergence and annealing,” Transactions on Machine Learning Research, 2022. [13] M. Welling and Y.W. Teh, “Bayesian learning via stochastic gradient langevin dynamics,” in Proceedings of the 28th international conference on machine learning (ICML-11), 2011, pp. 681–688. [14] A. Neelakantan et al., “Adding gradient noise improves learning for very deep networks,” in arXiv preprint arXiv:1511.06807, 2015. [15] Zhishen Huang and Stephen Becker, “Stochastic gradient langevin dynamics with variance reduction,” in 2021 International Joint Conference on Neural Networks (IJCNN), 2021, pp. 1–8. [16] S. Kobayashi, Y. Akram, and J. Von-Oswald, “Weight decay induces low-rank attention layers,” in Advances in Neural Information Processing Systems. 2024, vol. 37, pp. 4481–4510, Curran Associates, Inc. [17] Y. Li, J. Wang, X. Dai, L. Wang, C.M. Yeh, Y. Zheng, W. Zhang, and K.-L. Ma, “How does attention work in vision transformers? a visual analytics attempt,” IEEE Transactions on Visualization and Computer Graphics, vol. 29, no. 6, pp. 2888–2900, 2023. [18] L. Noci, C. Li, M. Li, B. He, T. Hofmann, C.J. Maddison, and D. Roy, “The shaped transformer: Attention models in the infinite depth-andwidth limit,” in Advances in Neural Information Processing Systems. 2023, vol. 36, pp. 54250–54281, Curran Associates, Inc. [19] Amir Ghorbani, Behnam Neyshabur, and Clément Raffel, “Examining the shape of attention in transformers,” Transactions of Machine Learning Research, vol. 1, pp. 1–38, 2022. [20] S. Mei, T. Misiakiewicz, and A. Montanari, “Mean-field theory of twolayers neural networks: dimension-free bounds and kernel limit,” in Proceedings of the Thirty-Second Conference on Learning Theory. 25– 28 Jun 2019, vol. 99 of Proceedings of Machine Learning Research, pp. 2388–2464, PMLR. [21] S.J. Vollmer, K.C. Zygalakis, et al., “(non-) asymptotic properties of stochastic gradient langevin dynamics,” arXiv preprint arXiv:1501.00438, 2015. [22] M. Raginsky, A. Rakhlin, and M. Telgarsky, “Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis,” Journal of Machine Learning Research, vol. 18, no. 117, pp. 1–47, 2017.
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
[23] Leonard G., “Logarithmic sobolev inequalities,” American Journal of Mathematics, vol. 97, no. 4, pp. 1061–1083, 1975. [24] D. Chafaı̈ and J. Lehec, “Logarithmic sobolev inequalities essentials,” Accessed on, p. 4, 2024. [25] D. Bakry and M. Emery, “Diffusions hypercontractives,” Séminaire de probabilités XIX 1983/84, pp. 177–206, 1985. [26] A Guionnet and B Zégarlinski, “Lectures on logarithmic sobolev inequalities,” in Séminaire de Probabilités XXXVI, pp. 1–134. Springer, 2004. [27] S. Mei, T. Misiakiewicz, and A. Montanari, “Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit,” in Proceedings of the Thirty-Second Conference on Learning Theory, Alina Beygelzimer and Daniel Hsu, Eds. 25–28 Jun 2019, vol. 99 of Proceedings of Machine Learning Research, pp. 2388–2464, PMLR. [28] J. Sirignano and K. Spiliopoulos, “Mean field analysis of neural networks: A central limit theorem,” Stochastic Processes and their Applications, vol. 130, no. 3, pp. 1820–1852, 2020. [29] P. Xu, J. Chen, D. Zou, and Q. Gu, “Global convergence of langevin dynamics based algorithms for nonconvex optimization,” in Advances in Neural Information Processing Systems. 2018, vol. 31, Curran Associates, Inc. [30] C. Villani, Optimal Transport: Old and New, Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008. [31] G. Zamir, A. Dokania, B. Zhao, and R. Yu, “Improving learning to optimize using parameter symmetries,” 2025. [32] Florian Wenzel, Patryk Swiatczak, Jonathan Blair, Patrick Warr, Pavel Izmailov, Alexander Gordon Wilson, David McAllester, and Balaji Lakshminarayanan, “How good is the bayes posterior in deep neural networks really?,” in Proceedings of the 37th International Conference on Machine Learning (ICML), 2020, p. We use Wt to denote a standard multivariate Wiener process. . . . [33] T.M. Cover and J.A. Thomas, Elements of Information Theory, WileyInterscience, Hoboken, NJ, USA, 2nd edition, 2006. [34] O. Catoni, PAC-Bayesian Supervised Classification: The Thermodynamics of Statistical Learning, vol. 56 of IMS Lecture Notes–Monograph Series, Institute of Mathematical Statistics, Beachwood, Ohio, USA, 2007. [35] Jiawei Sun, Yiding Yang, and Mohit Bansal, “Noise-regularised instruction tuning improves llm robustness,” in Proceedings of ACL 2024, 2024. [36] Mikhail Belkin, Daniel Hsu, and Afonso Bandeira, “Ridge-less regression, implicit regularisation, and generalisation,” Journal of Machine Learning Research, vol. 24, no. 85, pp. 1–45, 2023. [37] F. Dangel, F. Kunstner, and P. Hennig, “Backpack 1.2.0 documentation: Hutchinson trace estimation,” 2021, Use-case example for trace estimation. [38] Papers With Code, “Penn treebank dataset,” 2020, Language modeling benchmark dataset. [39] Linguistic Data Consortium, “Treebank-3 (ldc99t42),” 1999, Penn Treebank dataset release. [40] Papers With Code, “Wikitext-103 dataset,” 2017, Large-scale language modeling dataset. [41] Salesforce Research, “Salesforce/wikitext,” 2020, WikiText dataset on Hugging Face. [42] S. Dutta, A. Basarab, B. Georgeot, and D. Kouamé, “DIVA: Deep unfolded network from quantum interactive patches for image restoration,” Pattern Recognition, vol. 155, pp. 110676, 2024. [43] S. Dutta, A. Basarab, D. Kouamé, and B. Georgeot, “Quantum algorithm for signal denoising,” IEEE Signal Processing Letters, vol. 31, pp. 156– 160, 2024. [44] S. Dutta, A. Basarab, B. Georgeot, and D. Kouamé, “Quantum mechanics-based signal and image representation: Application to denoising,” IEEE Open Journal of Signal Processing, vol. 2, pp. 190–206, 2021. [45] S. Dutta and J. Mamou, “A quantum denoising-based resolution enhancement framework for 250-mhz and 500-mhz quantitative acoustic microscopy,” IEEE Transactions on Computational Imaging, vol. 10, pp. 1489–1504, 2024. [46] S. Dutta and J. Mamou, “Enhancing 3d radio-frequency data in quantitative acoustic microscopy using quantum-driven prior at 250-mhz and 500-mhz,” IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control, 2025. [47] S. Dutta, B. Georgeot, J. Michetti, A. Basarab, and D. Kouamé, “Unsupervised physics-inspired deep learning network with application to dental computed tomography image restoration,” in 2024 IEEE International Symposium on Biomedical Imaging (ISBI), 2024, pp. 1–5.
15
[48] A. Floquet, S. Dutta, E. Soubies, D.-H. Pham, and D. Kouame, “Automatic tuning of denoising algorithms parameters without ground truth,” IEEE Signal Processing Letters, vol. 31, pp. 381–385, 2024. [49] OpenAI, “Chatgpt,” https://openai.com/chatgpt, 2025. [50] Microsoft, “Copilot,” https://copilot.microsoft.com, 2025. [51] H. Avron and S. Toledo, “Randomized algorithms for estimating the trace of an implicit symmetric positive semi-definite matrix,” Journal of the ACM, vol. 58, no. 2, pp. Article 8, 2011.
Abhijit Das received the B.Tech. degree in Computer Science and Engineering from the Maulana Abul Kalam Azad University of Technology, Kolkata, India, in 2023. From 2023 to 2024, he was a Research Associate with the Department of Artificial Intelligence and Data Science, Jio Institute, Navi-Mumbai, India. Alongside, in this period, he has worked with Department of Radiology at Northwestern University. From 2024 to 2025, he has worked at OnFinance AI, India. From 2025 Feb to 2025 June, he worked with the Science and Technology Organization, Advanced Technology Group, GE HealthCare, Bangalore, India, as an AI Researcher (contractual). Currently, he is working at Innovxcare AI, Bangalore, India, as an AI Engineer. His research interests encompass computational imaging, quantum image processing, and deep learning, physics informed DL models, multimodal LLM with a particular focus on reasoning based multimodal large models.
Sayantan Dutta (Member, IEEE) received the B.Sc. degree in mathematics from the University of Burdwan, Bardhaman, India, in 2016, the M.Sc. degree in mathematics from Visva-Bharati University, Santiniketan, India, in 2018, and the M.S. degree in fundamental physics from the University of Tours, Tours, France, in 2019. He received the Ph.D. degree in computer science from the University Paul Sabatier Toulouse 3, Toulouse, France, in 2023. From 2023 to 2024, he was a Postdoctoral Associate with the Department of Radiology, Weill Cornell Medicine, Cornell University, New York, NY, USA. Since 2024, he has been with the Science and Technology Organization, Advanced Technology Group, GE HealthCare, Bangalore, India, where he is currently a Senior AI Scientist. His research interests encompass computational imaging, quantum computing, quantum image processing, and deep learning, with a particular focus on inverse problems such as denoising, deblurring, super-resolution, and compressed sensing. He has contributed to the development of physicsinspired and quantum mechanics-based methods for image restoration and signal processing.
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
S UPPLEMENTARY M ATERIAL This work establishes a rigorous functional-analytic foundation for Transformer optimization by proving that the regularized cross-entropy loss with L2 weight decay satisfies Villani’s conditions for coercive energy functions. Specifically, the paper demonstrates that the loss landscape exhibits smoothness, coercivity, Gaussian-integrable tails, and differential growth—properties that enable the derivation of explicit log-Sobolev and Poincaré constants. These constants, in turn, yield finite-time convergence guarantees for Langevinbased optimization and PAC-Bayesian generalization bounds. The theoretical results are validated empirically on GPTNeo-125M across Penn Treebank and WikiText-103, using a scalable Villani diagnostic and spectral analysis of the Hessian. Together, these findings provide a principled explanation for the role of weight decay in shaping Transformer loss geometry and improving both optimization and generalization.
16
TABLE III AGGREGATE STATISTICS OF H UTCHINSON TRACE ESTIMATES . R ESULTS AVERAGED OVER 64 PROBE VECTORS AT EACH CHECKPOINT, ACROSS 60 (PTB) OR 12 (WT-103) CHECKPOINTS PER RUN . λ
Dataset
Mean (×108 )
Rel. Error
95% CI
0 10−3 10−2 0 10−3 10−2
PTB PTB PTB WT-103 WT-103 WT-103
1.74 6.42 58.4 1.69 6.31 57.9
0.109 0.048 0.050 0.124 0.051 0.052
0.048 0.078 0.73 0.054 0.080 0.76
The Laplacian ∆ℓi (trace of the Hessian) admits a decomposition: ∆ℓi = Tr Jz⊤ [diag(σ) − σσ ⊤ ]Jz +
V X
Tr (∇2 zv )(σv − 1v=yi ) .
v=1
A PPENDIX A AUXILIARY P ROOFS AND N UMERICAL G UARANTEES Appx. Sec. A-A establishes uniform gradient and Laplacian bounds necessary for differential growth analysis in Lemma 4. Appx. Sec. A-B derives a variance bound for the Hutchinson trace estimator employed in Sec. IV.
where diag(σ) is the diagonal matrix of softmax scores and ∇2 zv denotes the block-diagonal Gauss–Newton term contributed by feed-forward weights. Applying norm inequalities and simplifying with architectural constraints similar to those above, we define a dataset-dependent constant C2 satisfying ∆ℓi ≤ 2 ∥Jz ∥2F +
V X
∥∇2 zv ∥F ≤ C2 ,
(56)
v=1
A. Uniform Bounds on ∇ℓ and ∆ℓ
⋆2 ⋆2 , which where C2 = 2V B 2 Wmax (1 + o(1)) + Ld2model Wmax guarantees bounded curvature contributions in Lemma 4 under the stated assumptions.
Recall the cross-entropy loss and softmax logits: ℓi (θ) = − log pθ (yi | xi ), d
exp zy,θ (x) , pθ (y | x) = PV v=1 exp zv,θ (x)
V
where zθ : R → R are model logits and Jz = ∇θ zθ (x). The gradient ∇ℓi decomposes as:
B. Variance Bound for Hutchinson Trace Estimation Sec. IV relies on Hutchinson’s estimator to approximate ∆F(θ) = Tr[∇2 F(θ)] via Monte Carlo sampling. Given M Rademacher probes v (m) ∈ {−1, 1}d , we estimate:
∇ℓi = −Jz⊤ eyi + Jz⊤ σ(z), d ∆F
with σ(z) ∈ RV denoting the softmax vector. Hence, ∥∇ℓi ∥ ≤ ∥Jz ∥2→2 · ∥eyi − σ(z)∥ ≤ 2∥Jz ∥2→2 ,
(53)
as ∥eyi − σ(z)∥ ≤ 2, where ∥Jz ∥2→2 refers to the operator norm. We upper-bound ∥Jz ∥2→2 using standard architectural constraints: each residual block √ is 1-Lipschitz, and LayerNorm is contractive by at most dmodel . Assuming bounded input embeddings (B) and spectral norm control via weight decay, we arrive at: p ⋆ ⋆ ∥Jz ∥2→2 ≤ B Wmax 1 + L Wmax dmodel , (54) with weight decay keeps Wmax ≤ λ−1/2 ∥θ∥ and therefore the product is bounded by a dataset-dependent constant p ⋆ ⋆ C1 = 2 B Wmax 1 + L Wmax dmodel , (55) ⋆ where Wmax denotes the maximal operator norm attained during training. This yields a uniform gradient bound constant ⋆ C1 as used in Lemma 4. Empirically Wmax ≈ 0.9 for GPT−3 Neo-125M with λ = 10 , giving C1 ≈ 7.6.
(M )
M
1 X (m) 2 v , ∇ F(θ) v (m) . = M m=1
This estimator is unbiased. Its variance is given by: (M ) 2 X d (λi − λj )2 , Var ∆F = M i<j
(57)
where λi denote the eigenvalues of ∇2 F(θ) [51]. For our model, the spectrum exhibits a narrow bulk and a few dominant eigenvalues due to regularization. Substituting λmax ∼ p λ∥θ∥, we estimate MSE ≤ 2d/M λ ∥θ∥. With M = 64, d ≈ 1.25 × 108 , and λ = 10−3 , the estimator yields < 3% absolute error once ∥θ∥ ≳ 104 , matching the regime in Sec. IV. A PPENDIX B E MPIRICAL T RACE VARIANCE S TATISTICS Table III summarizes Hutchinson trace variance statistics across all training checkpoints and random seeds, including mean, relative error, and 95% confidence intervals for each dataset and weight decay setting. See Sec. VI-H for discussion of variance behavior during training.
DAS AND DUTTA: WEIGHT-DECAY TURNS TRANSFORMER LOSS LANDSCAPES VILLANI
All values reflect statistics over instantaneous trace esti(1) d mates ∆F = ⟨v, ∇2 F v⟩, computed using Rademacher vectors with M = 64 samples per checkpoint. Trace estimates were computed using the same mixed-precision setup described in Sec. VI. For the 125-million-parameter GPTNeo model, this setup yields a stable relative error regime of approximately 5% (Table III), confirming both the theoretical predictions from Appx. A and the empirical observations across checkpoints.
17