The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity

The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity Iosif Lytras1,2,3 , Nikolaos Makras1 , and Sotirios Sabanis1,2,3 1

School of Mathematics, University of Edinburgh, UK 2 National Technical University of Athens, Greece 3 Athena/Archimedes Research Centre, Greece August 7, 2026

arXiv:2608.06283v1 [cs.LG] 6 Aug 2026

Abstract We study the problem of sampling from target distributions whose potentials are simultaneously non-smooth, subject to superlinear gradient growth, and non-convex. We introduce the Subgradient Tamed Unadjusted Langevin Algorithm (SG-TULA), a discretisation of the Langevin diffusion that operates directly on subgradients, without relying on computationally demanding smoothing procedures. To handle the superlinear regime, taming techniques are employed to produce a stable, explicit scheme. We derive non-asymptotic convergence bounds in Wasserstein-2 distance, with all constants tracked explicitly in terms of dimension and inverse temperature, improving upon the currently known rates for subgradient-based Langevin algorithms. We further provide excess risk estimates for the associated optimisation problem. We verify the assumptions, with explicit constants, for the regularized pretraining potential of a LLM in the GPT-2 lineage and the boosted coordinate-wise variant of SG-TULA pretrains the former competitively against finetuned AdamW and Muon, for which no comparable non-asymptotic guarantees are presently available.

1. Introduction, Purpose and Scope The hypothesis under which optimization algorithms are usually analyzed, namely a globally Lipschitz gradient together with convexity or strong convexity, are rarely met by the problems they are applied to. Consider the case of a transformer which is the composition of layers rendering the gradient polynomial in the weights, so that it is not Lipschitz, is rendered discontinuous by ReLU, and the loss is not convex. Each failure has its own practical remedy, clipping for the first and subgradients for the second, and each has been analyzed in isolation. What is missing is a non-asymptotic guarantee when these three occur together. This is the challenge being addressed in this paper. For a potential u : Rd → R, under mild assumptions, we consider the problem (1)

θ⋆ ∈ arg min u(θ), θ∈Rd

which we approach by sampling from the Gibbs measure attached to u, πβ (dx) ∝ e−βu(x) dx on Rd . The connection between the two is the concentration of πβ around the global minimizers of u for large β, a fact going back to Hwang [1980]. In that sense, an algorithm which samples accurately from πβ also solves the associated expected excess risk problem for u. We therefore bound the excess risk E[u(θnλ )] − u(θ⋆ ) of the iterates (θnλ )n∈N by the sampling error together with a temperature gap, the latter quantifying the bias attributable to β. The algorithm producing these iterates, namely the subgradient tamed unadjusted Langevin algorithm (SG-TULA), is an explicit Euler–Maruyama discretisation of the associated Langevin diffusion which operates on a measurable selection h ∈ ∂u and stabilizes the drift by taming, λ θn+1 = θnλ − λhλ (θnλ ) +

p

2λβ −1 ξn+1 ,

hλ (x) =

h(x) − µx + µx, 1 + λ1/2 |h(x) − µx|

with (ξn )n≥1 i.i.d. standard Gaussian, alongside which a coordinate-wise variant is analysed. Taming acts on the subgradient after it has been computed, so the cost of evaluating h is unaffected and the overhead is significantly smaller than one matrix multiplication. Written this way the scheme is visibly an alternative method to gradient 1

clipping, the latter being introduced for recurrent networks in Pascanu et al. [2013] and is now the standard response to exploding gradients in deep learning. However clipping requires a threshold to be fixed in advance and, once active, displaces the drift by an amount that does not diminish as the stepsize is refined. The tamed drift requires no threshold, rescales continuously with the subgradient magnitude, and satisfies hλ → h as λ → 0, so that the scheme remains a discretisation of the intended dynamics. Accordingly, we assume that u is semi-convex A1, strongly convex outside a compact set A2, and every subgradient satisfies |h(x)| ≤ m + K|x|p for some p ≥ 1 A3. Semi-convexity is a weak departure from convexity, in that subdifferential calculus is available and permits gradient discontinuities on a set of measure zero, which is the irregularity encountered in the analysis of neural networks. Additionally, it provides the one-sided Lipschitz bound on which the discretisation estimates rest. Polynomial growth is what taming controls by avoiding exploding gradients. Strong convexity at infinity yields the key property of dissipativity, which together with semi-convexity gives the contraction of the diffusion. It also highlights what the method can address, since a dissipative drift forces πβ to have at most Gaussian tails. Potentials falling outside this regime must be regularized by higher order terms, namely weight decay in the context of ML training. Under these assumptions alone, and with no further regularity imposed on u, our main sampling result takes the following form: Theorem 1 (informal). For every stepsize λ < 1/(4µ) and every n ∈ N,
W2 πβ , L(θnλ ) ≤ 2 CE λ1/4 + e CD e−nλ/T0 . {z } | | {z } discretisation

contraction

All three constants
are given in closed form.
The dimension enters only polynomially through the ratio d/β, with CE = O (d/β)p/2 and CD = O (d/β)1/2 , while the restart horizon T0 , and thus the rate of the contraction term, is independent of the dimension. The dependence on the inverse temperature is exponential in general, since it is inherited from the log-Sobolev constant of πβ . This is to a large extent an artifact of the proof strategy, which routes the optimization guarantee through a sampling one. In particular, the exponential factor is the cost of mixing between the modes of πβ . In the optimization regime πβ has concentrated on the global minimizers since β is large. There the governing constant is a local one and no barrier needs to be crossed, as discussed in Remark 14. Formally, we record in Remark 11 the additional Polyak–Łojasiewicz type condition, under which this dependence becomes polynomial. Given a tolerance ϵ > 0, Corollary 1 provides the complexity of the algorithm, the stepsize being of order λ = O(ϵ4 ) and n = O(ϵ−4 log(1/ϵ)) for iterations respectively. The sampling estimate transfers to the associated optimization problem. Writing θ∗ for a minimiser of u, one obtains the following. Theorem 2 (informal). Under the latter conditions and for every β ≥ 4/µ,
E[u(θnλ )] − u(θ∗ ) ≤ CT1 W2 πβ , L(θnλ ) + | {z } sampling

CT2 |{z}

.

temperature gap

Both constants are explicit, with CT1 = O (d/β)p/2 and CT2 = O (d/β) log(dβ) . This is a non-asymptotic optimization guarantee under hypotheses that the potentials of interest actually satisfy. The complexity O(ϵ−4 ) is inherited entirely from Lemma 9. Since the contraction term carries no power of the stepsize, the order of the strong approximation estimate on the finite restart horizon is the order of the algorithm, and what can be expected of that estimate is a question for the numerical analysis of SDEs with irregular coefficients. Two lines of work there are of particular relevance. The first, developed in Dareiotis and Gerencsér [2020], Dareiotis et al. [2023] and the references therein, obtains order O(λ1/2−ϵ ) for arbitrary ϵ > 0 for the Euler–Maruyama scheme without imposing any regularity on the drift, which is almost optimal in view of the corresponding lower error bounds. However, these results are obtained for globally bounded drift coefficients. Growth is addressed by a second line of work. Müller-Gronbach and Yaroslavtseva [2020] establish order O(λ1/2 ) for piecewise Lipschitz drift in dimension d = 1, and this was extended in Müller-Gronbach et al. [2026] to coefficients of superlinear growth, again in dimension d = 1 and for finitely many points of discontinuity; see also Müller-Gronbach et al. [2025] for the multidimensional case under piecewise Lipschitz drift. By contrast, our setting is multidimensional, superlinear and non-convex, and the set of discontinuities is permitted to be infinite, subject only to being of Lebesgue measure zero. Against this, the

2

benchmark under regularity is O(λ1/2 ) for Lipschitz drift and O(λ) once the Hessian is Lipschitz as well, the latter being an order of Milstein type that the constant diffusion coefficient disguises as Euler’s. Within the sampling literature the same comparison holds. For p = 1 our assumptions reduce to those of Johnston et al. [2025a], where SG-ULA is analyzed under semi-convexity, strong convexity at infinity and subgradients of at most linear growth. Our result recovers that setting and improves the order in W2 from O(λ1/8 ) to O(λ1/4 ). Table 1 places the result among the alternatives. To the best of our knowledge no existing result, other than ours, attains O(λ1/4 ) in W2 under assumptions of non-smoothness, superlinear growth and non-convexity simultaneously and those achieving comparable or better rates do so under smoothness or under strong convexity. Table 1: Hypotheses and W2 rates for related Langevin schemes. Work Durmus and Moulines [2019] Brosse et al. [2019] Lovas et al. [2023] Durmus et al. [2022] Lehec [2023] Fruehwirth and Habring [2024] Johnston et al. [2025a] SG-TULA (this work)

Smooth yes yes yes composite no (subgrad.) no (subgrad.) no (subgrad.) no (subgrad.)

Growth linear superlinear superlinear linear linear linear linear superlinear

Convexity strong strong non-convex at infinity convex strong at infinity at infinity

W2 rate λ1/2 λ1/2 λ1/4 1/2 λ (TV) λ1/4 λ1/2 λ1/8 λ1/4

The algorithm is moreover competitive in practice. In Section I we pretrain nanochat Karpathy [2025] with the boosted variant of the scheme, in which the taming map is composed with the multiplicative factor (boosting) of Lim and Sabanis [2024]. That perturbation is multiplicative and uniformly bounded above and below, so that the growth, closeness and dissipativity estimates of Section B.2 go through with altered constants and the analysis applies unchanged in form. At depth 12 it attains the best validation bits-per-byte and CORE score of the three optimizers tested, ahead of finetuned AdamW and Muon. Moreover, at depth 24 it remains within evaluation noise of both, and does so without the benefit of a scaling law. In Section I we carry out that verification for nanochat, a compact transformer language model in the GPT-2 lineage, and obtain the curvature bound
12L+2 ∇2θ u(θ) ⪰ − c⋆ 1 + |θ| Id¯, with c⋆ explicit in the depth L, the width and the vocabulary size, from which A1–A3 follow for the regularized potential. The exponent 12L + 2 is itself the statement that superlinear growth in deep architectures is a consequence of depth and not an artifact of the analysis. A second verification, for the sparse phase retrieval problem of Section 8.2, exhibits all three failures at once, i.e. the quartic energy is non-convex with p = 3, and the ℓ1 term places the gradient discontinuities on the coordinate hyperplanes. 2. Literature Review Although this is a broad and active research area, including Metropolis-adjusted variants, underdamped Langevin dynamics, Lévy driven dynamics, interacting particles methods, and algorithms designed to handle constraints, we focus on approaches that extend the Unadjusted Langevin Algorithm (ULA) beyond the standard assumptions of linear growth, strong convexity, and smoothness. 2.1. Langevin-based algorithms A systematic study of the ergodicity of the Unadjusted Langevin Algorithm (ULA) was undertaken in Roberts and Tweedie [1996], which also formally introduced gradient clipping under the MALTA (Metropolis-Adjusted Langevin Truncated Algorithm) framework, however, the latter algorithm was not further investigated. Since, theoretical analyses have primarily focused on smooth, convex potentials u. Under these assumptions, i.e. global Lipschitz continuity of the gradient and strong convexity, ULA is known to enjoy geometric ergodicity and admits non-asymptotic convergence guarantees in both Wasserstein and total variation distances. Such results were first established by Dalalyan [2017] and Durmus and Moulines [2017], with refined bounds later provided in Durmus and Moulines [2019]. In the latter work, the authors obtain a global error bound in Wasserstein distance, of order O(λ1/2 ) under the standard smooth and strongly convex setting, which can be improved to O(λ) when additional smoothness of the Hessian is 3

assumed. While this improvement formally corresponds to employing the Milstein discretization scheme, the connection may not be immediately evident to the reader since the diffusion coefficient is constant in this case. More broadly, many results that invoke higher order smoothness assumptions, often phrased in terms of Hessian regularity or beyond, implicitly inherit improved convergence rates that, from the SDE numerical analysis perspective, stem from analogous higher order discretization effects. Considerable effort has since gone into weakening the convexity assumption. A prominent line of work trades convexitytype conditions for a functional inequality on the target measure, while retaining global smoothness of the corresponding gradient. Beginning with Vempala and Wibisono [2019], a substantial literature provides estimates in KL and Rényi divergences under log-Sobolev and Poincaré inequalities, e.g. Erdogdu and Hosseinzadeh [2021], Erdogdu et al. [2022], Chewi et al. [2025]. In Mousavi-Hosseini et al. [2023] the authors work under weak Poincaré inequalities to cover heavy-tailed, non-convex targets, including Cauchy-type distributions for which a Poincaré inequality fails. Under s-Hölder regularity of the gradient, s ∈ (0, 1], they obtain non-asymptotic Rényi divergence guarantees for ULA, providing both iteration and time-horizon complexities. Their results imply a discretization rate of O(λs/2 ) in TV distance.

2.2. Algorithms for superlinear drift It is well established that ULA exhibits numerical instability when the target distribution has a superlinearly growing potential, with moments of the discrete approximations diverging in finite time, Hutzenthaler et al. [2010]. In the deterministic gradient case, SGLD reduces to ULA, and as such can be seen as a standard Euler-Maruyama discretisation of the Langevin diffusion, which inherits the same limitations due to the growth of the drift coefficient. This instability manifests in practical application whenever gradients accumulate too rapidly, as is well documented in the context of recurrent neural Pascanu et al. [2013], motivating the development of stabilization methods. One idea for handling superlinear coefficients, is the taming technique, originally developed for stabilizing explicit Euler schemes of SDEs with exploding coefficients in Hutzenthaler et al. [2012] and subsequently in Sabanis [2013]. The Tamed Unadjusted Langevin Algorithm (TULA), introduced in Brosse et al. [2019], applies this idea to ULA. There, u is assumed to be continuously differentiable and strongly convex but allowed ∇u to grow polynomially. Under these assumptions, they proved convergence of TULA in Wasserstein-2 distance to the target distribution with order O(λ1/2 ). Building on TULA, a line of research has adapted taming to various stochastic and non-convex gradient settings. In the context of optimization, the stochastic gradient setting is understood as mini-batching, whereby only noisy estimates of the full gradient are accessible at each iteration, we refer to Dalalyan and Karagulyan [2019] for an early treatment, and to Barkhagen et al. [2021], Chau et al. [2021] for an extension to the non-i.i.d. case. Within this regime, Lovas et al. [2023] introduced TUSLA for ANN, extending taming to non-convex potentials for which the original density fails to satisfy a dissipativity condition, a higher-order regularization term is added to restore this property, yielding a convergence rate of order O(λ1/4 ) in Wasserstein-2 distance. This direction was further developed in Lim and Sabanis [2024], where taming is applied coordinate-wise and augmented with a boosting factor, which adaptively scales the gradient upward to counteract the vanishing gradient problem. The resulting algorithm, TheoPouLA, retains a convergence rate of order O(λ1/4 ) and continues to require higher-order regularization to ensure dissipativity, yet demonstrates competitive empirical performance against established adaptive gradient-based optimization algorithms in deep learning tasks. Under a different taming function, Neufeld et al. [2025] obtain an O(λ1/2 ) rate in Wasserstein-2 distance under a convexity-at-infinity type condition. However, their constants scale exponentially with the dimension and their analysis requires stronger smoothness assumptions on the Hessian of the potential. In recent work, Yang and Wang [2025] analyze the Unadjusted Langevin Algorithm (ULA) and a projected variant (PLMC) under the log-Sobolev assumption, extending beyond the convex regime. They assume global Lipschitz continuity of ∇u for ULA, and for PLMC they relax this to a polynomial local Lipschitz condition. In principle PLMC can target potentials with super linear growth, however their analysis imposes a growth bound on ∇(∆u) which significantly limits the class of superlinear potentials that can be targeted, see Remark A2 in Li et al. [2022]. Their error bounds in W2 achieve order O(λ), matching the best known rate of the Milstein scheme in the fully smooth and strongly convex case. 4

Further results address the superlinear and non-convex regime in the deterministic gradient case. Under a functional inequality for the target distribution together with a locally Lipschitz gradient, which permits superlinear growth, estimates in KL divergence were obtained in Lytras and Sabanis [2025], Lytras and Mertikopoulos [2025] and later sharpened in Lytras et al. [2025], Lytras and Ntousis [2026] by exploiting additional local regularity of the Hessian. This last requirement was in turn removed through tamed randomized midpoint schemes, which attain Õ(ϵ−1 ) complexity in Wasserstein and total variation distances Wang and Yang [2025], Lytras and Ntousis [2026]. 2.3. Algorithms for non-smooth potentials It is worth distinguishing here between two forms of irregularity. In Sabanis and Zhang [2020] the stochastic gradient is discontinuous in θ, as occurs in quantile and CVaR estimation, where indicators of the form 1{x<θ} appear. The analysis proceeds under a Lipschitz in average condition (Assumption 3 therein), together with a decomposition of the stochastic gradient into a Lipschitz part and a part bounded in θ. Averaging against the law of the data then restores regularity, and the drift h = ∇U of the associated Langevin SDE is globally Lipschitz. The irregularity treated in the present work is of a different nature. The potential u is itself non-smooth, no averaging is available to recover regularity, and the drift is a measurable selection h ∈ ∂u which is genuinely discontinuous. We identify two broad strategies that extend Langevin algorithms to the latter non-smooth setting. (A) Smoothing Proximal Maps One approach to handling non-smoothness is to modify the target potential to recover regularity. A standard technique is Moreau–Yosida regularization, which replaces the potential u with its envelope uγ , a smoothed approximation with a Lipschitz continuous gradient defined via the proximal operator. This leads to the Moreau–Yosida Unadjusted Langevin Algorithm (MYULA) for which non-asymptotic guarantees were first studied in Brosse et al. [2017] for potentials convex on compact bodies, proving convergence in total variation (TV) distance and establishing an O(λ1/2 ) rate in Wasserstein-2 distance, provided the smooth component is strongly convex. In Durmus et al. [2022], the authors provided non-asymptotic TV guarantees for potentials that decompose into smooth and non-smooth convex components, achieving a rate of O(λ1/2 ) up to logarithmic factors. Despite these theoretical strengths, evaluating the Moreau–Yosida envelope remains computationally demanding, as it requires solving a proximal subproblem at each iteration, unless there exists a closed form solution which is uncommon. Furthermore, these methods generally necessitate a composite structure, where the potential is a sum of smooth and non-smooth functions, and assume that the associated proximal maps can be computed efficiently. This requirement limits their applicability to potentials that lack such a decomposition. To mitigate the cost of proximal steps, Ehrhardt et al. [2024] proposed using inexact proximal mappings, though their guarantees are restricted to convex targets. More recently, Habring et al. [2026] introduced Diffusion at Absolute Zero (DAZ), which extends MYULA by incorporating an annealing schedule for the regularization parameter. By assuming the potential is semi-convex and convex outside a ball, DAZ relaxes the global convexity requirement of MYULA and achieves an O(λ1/2 ) rate in TV distance. Under similar composite and convexity assumptions, Renaud et al. [2025] established an O(λ1/(2p) ) rate in Wasserstein-p distances for the Proximal Stochastic Gradient Langevin Algorithm (PSGLA). Very recently, Luu and Wang [2026] replaced the weak convexity requirement of Renaud et al. [2025] to a difference-of-convex structure on the non-smooth term, extending the non-asymptotic guarantees in Wasserstein-p distances with the same convergence rate. Gaussian Smoothing Chatterji et al. [2020] introduce the Perturbed Langevin Monte Carlo (P-LMC) algorithm, which addresses the nonsmoothness of u by injecting noise into the argument of the gradient. They consider a composite potential of the form ū = u+g, where u is convex with an a-Hölder continuous gradient for a ∈ [0, 1], and g is smooth and strongly convex. In Wasserstein-2 distance they recover the optimal convergence rate when a = 1, and obtain a rate of order O(λ1/4 ) in the non-smooth case a = 0 (Remark 3.5 therein). This approach was later generalized to p−Gaussian smoothing in Nguyen et al. [2023], addressing a−mixture weakly smooth potentials under the log-Sobolev inequality. In this setting, the method achieves a convergence rate in Wasserstein-2 distance of order O(λa/2 ) for a ∈ (0, 1]. In a different approach, Laumont et al. [2022] assume a composite structure similar to the proximal ULA literature and smooth the irregular term via Gaussian convolution. The gradient of the smoothed term is approximated through 5

Tweedie’s formula, relating it to a trained denoiser, which leads to the Plug-and-Play ULA (PnP-ULA). The drift is thus determined by a learned network rather than by the potential itself, and their assumptions are correspondingly stated in terms of the denoiser. They establish a bias of order O(λ1/2 ) in V -total variation between the stationary distribution of the scheme and that of the smoothed target, rendering their result not directly comparable to the rates in the present discussion. Moreover, in the absence of any dissipativity yielding condition, their constants depend on the radius of the convex projection set required to recover contraction. The anchored Langevin algorithm of Gurbuzbalaban et al. [2025] addresses a different gap, primarily targeting heavytailed distributions. The authors introduce a reference potential and define a modified SDE whose drift and diffusion are constructed so as to preserve the original invariant measure. For non-smooth targets they consider a composite structure, a smooth and strongly convex component together with an irregular term that is weakly convex and Lipschitz, and hence admits bounded gradients almost everywhere, the reference potential being obtained by replacing the latter with a Gaussian-smoothed approximation. They establish non-asymptotic convergence guarantees in Wasserstein-2 distance. Their final estimate (Theorem 30 therein), however, carries a third term arising from the Monte Carlo approximation of the smoothed reference potential, which is increasing in both the stepsize and the number of iterations and is thus not uniform in time. The allowable stepsize is moreover restricted to O(d−2 ). (B) Direct subgradient methods: Durmus et al. [2019] introduced a subgradient variant of the stochastic unadjusted Langevin algorithm (SSGLD) for sampling from non-smooth, log-concave target distributions. Under the assumption that the potential is Lipschitz and convex, they derive non-asymptotic bounds in KL divergence at rate O(λ). However, their convergence guarantees apply only to the weighted average of the iterates’ distributions, leaving last-iterate convergence unaddressed. In a different direction, Lehec [2023] studies ULA under a subgradient formulation for globally Lipschitz convex potentials and establishes an O(λ1/4 ) rate in Wasserstein-2 distance under a log-Sobolev inequality, with the same rate extending to unconstrained settings when the subgradient has at most linear growth. Fruehwirth and Habring [2024] consider an explicit subgradient Langevin discretization under strong convexity and linear growth of the subgradient. By separately establishing geometric convergence of the iterates to the scheme’s stationary (discrete) measure and bounding the stationary bias W2 (πλ , π) via an energy-entropy functional argument, they recover a rate of order O(λ1/2 ) in W2 , matching the smooth case rate. 2.4. Contraction rates for Langevin dynamics Exponential contraction estimates for the Langevin diffusion beyond convexity have also been explored. In the classical case the contraction Wp (µPt , νPt ) ≤ e−λt Wp (µ, ν), p ≥ 1, holds with prefactor one precisely when the Bakry–Émery curvature is bounded below by λ > 0, which on Rd amounts to ∇2 u ⪰ λI. In the seminal work ofEberle [2016] the reflection coupling introduced there yielded contraction in W1 , assuming only semi-convexity together with strong convexity at infinity. Here semi-convexity is precisely the requirement that the curvature be bounded below, possibly by a negative constant. Luo and Wang [2016] extended this to Lp -Wasserstein distances under the same hypotheses. However their estimate , for p = 2, is of the form W2 (µPt , νPt ) ≤ Ce−λt/2 W1 (µ, ν) + W22 (µ, ν)

1/2

,

where the right-hand side involves W1 under a square root. It is additional term, which halves the rate attainable by algorithms proceeding through such an estimate. Contraction in the form employed here, namely W2 (µPt , νPt ) ≤ ce−λt W2 (µ, ν) with a finite prefactor c > 1, was obtained in Wang [2020] under a log-Sobolev inequality and negatively bounded curvature. The result is established for diffusion semigroups on general Riemannian manifolds and for C 1 drift vector fields. A W2 -contraction was subsequently derived in Monmarché [2023] for elliptic diffusions whose drift contracts distances outside a compact set, though only for a sufficiently large diffusion coefficient, a regime opposite to the one of interest for optimisation. Although recent advances have addressed non-convex settings under superlinear growth, these results are largely confined to highly smooth regimes. Our work aims to bridge this gap by allowing discontinuities in the gradient of u while retaining both superlinearity and non-convexity.

6

3. Notation Pd We introduce some basic notation. For u, v ∈ Rd , define the scalar product ⟨u, v⟩ = i=1 ui vi and the Euclidian 1/2 norm |u| = ⟨u, u⟩ . For all continuously differentiable functions f : Rd → R, ∇f denotes the gradient. For symmetric matrices A and B we write A ⪰ B when A − B is positive semi-definite, and Id denotes the identity on Rd . For semi-convex u we write ∇2 u for the Alexandrov Hessian, which exists Lebesgue-almost everywhere, and any matrix inequality involving it is understood to hold almost everywhere. Since the singular part of the distributional Hessian of a convex function is non-negative, such a bound holds distributionally as well, and the two notions agree with the classical one when u ∈ C 2 . The integer part of a real number x is denoted by ⌊x⌋. For an Rd -valued random variable Z, its law on B(Rd ), i.e. the Borel sigma-algebra of Rd , is denoted by L(Z). We denote by P(Rd ) the set of R d d d probability measures on B(R ) and for any p ∈ N, Pp (R ) = {π ∈ P(R ) : Rd |x|p dπ(x) < ∞} denotes the set of all probability measures over B(Rd ) with finite p-th moment. For any two Borel probability measures µ and ν, we define the Wasserstein distance of order p ≥ 1 as  Wp (µ, ν) =

Z

p

|x − y| dζ(x, y)

inf ζ∈Π(µ,ν)

1/p ,

Rd ×Rd

where Π(µ, ν) is the set of all transference plans of µ and ν. Moreover, for any µ, ν ∈ Pp (Rd ), there exists a transferp ence plan ζ ∗ ∈ Π(µ, ν) such that for any coupling (X, Y ) Rdistributed according to ζ ∗ , Wp (µ, ν) = E1/p [|X − Y | ]. d For µ and ν in P(R ) with µ = f ν, we write H(µ | ν) = f log f dν for the relative entropy of µ with respect to ν, and H(µ | ν) = +∞ whenever µ fails to be absolutely continuous with respect to ν. 4. The Algorithm Consider the overdamped Langevin SDE (Zt )t∈R+ given by dZt = −h(Zt )dt +

p

2β −1 dBt , t ≥ 0,

(2)

with Z0 = θ0 ∈ Rd , where h ∈ ∂u, β is the inverse temperature parameter and (Bt )t≥0 is a standard d-dimensional Brownian motion. To address the superlinear behavior of the subgradient h(x) = ∇u(x), we introduce a parameterized family of drift functions (hλ )λ≥0 , where hλ : Rd → Rd depends on the stepsize λ. These drift functions, which we will refer to as taming functions, serve as close approximations of ∇u. We then define, the tamed subgradient unadjusted Langevin algorithm (SG-TULA), (θnλ )n≥0 , by p λ (SG-TULA): θn+1 = θnλ − λhλ (θnλ ) + 2λβ −1 ξn+1 , θ0λ = θ0 , n ∈ N, (3) where λ > 0 is the stepsize and (ξn )n≥1 is a sequence of i.i.d. standard Gaussian random variables on Rd . We propose two distinct explicit choices for the family (hλ )λ≥0 , which are inspired by previous studies on tamed Langevin based algorithms Johnston et al. [2024], Johnston et al. [2025b]. Define for all x ∈ Rd and λ > 0:   h(x) − µx h(i) (x) − µx(i) (i) hλ,u (x) := + µx and h := + µx . (4) λ,c 1 + λ1/2 |h(x) − µx| 1 + λ1/2 |h(i) (x) − µx(i) | i∈{1,...,d} The Euler scheme 3 with hλ = hλ,u is referred to as the uniform-wise subgradient tamed unadjusted Langevin algorithm (SG-TULAu), and the choice hλ = hλ,c , is referred to as the coordinate-wise subgradient tamed unadjusted Langevin algorithm (SG-TULAc). The corresponding iterates of the algorithms will be denoted by θnλ,u and θnλ,c respectively. For brevity, we shall henceforth suppress this distinction and conduct the analysis under the uniform-wise taming scheme. In the main Theorems, where the difference between the two variants is relevant, it will be explicit.

7

5. Assumptions & Consequences Assumption 1. The potential u is L-semi-convex. There exists L ≥ 0 such that u + L2 | · |2 is convex. Equivalently, every subgradient, h ∈ ∂u, satisfies the one-sided Lipschitz bound ⟨h(x) − h(y), x − y⟩ ≥ −L|x − y|2 ,

∀x, y ∈ Rd .

(5)

Remark 1. Under A1, u is locally Lipschitz, hence by Rademacher’s theorem ∇u exists Lebesgue-almost everywhere, and by Alexandrov’s theorem u is in fact twice differentiable a.e. Consequently the subdifferential ∂u(x) is nonempty and compact for every x ∈ Rd and reduces to {∇u(x)} almost everywhere. In this sense semi-convexity is a natural condition past smooth settings. It admits gradient discontinuities on a set of measure zero, precisely the non-smoothness met in sparsity penalties such as the ℓ1 , total-variation and SCAD priors and in activation functions such as ReLU. Assumption 2. The potential u is strongly convex at infinity (outside a compact set). There exist µ > 0 and R ≥ 0 such that, for every h ∈ ∂u and all x, y ∈ Rd with |x − y| ≥ R, ⟨h(x) − h(y), x − y⟩ ≥ µ|x − y|2 .

(6)

Assumption 3. Each subgradient of u grows at most polynomially. There exist m, K > 0 and p ≥ 1 such that, for every h ∈ ∂u, |h(x)| ≤ m + K|x|p ,

∀x ∈ Rd .

(7)

Remark 2. Let Assumptions A2 and A3 hold, then h is globally dissipative. That is, there exist µ, b > 0, such that ⟨x, h(x)⟩ ≥

µ 2 |x| − b, ∀x ∈ Rd , 2

(8)

 where b = max m2 /(2µ), mR + (µ/2)R2 + KRp+1 . Under the addition of A1, b is further refined to b = max m2 /(2µ), mR + (L + µ/2)R2 . Proof. The proof is postponed to Appendix G. Assumption 4. The initial condition of the algorithm is such that (9)

E|θ0 |2p < ∞.

Assumption 5. The potential is coordinate-wise dissipative: there exist µ, b > 0 such that, for every i ∈ {1, . . . , d}, x(i) h(i) (x) ≥ µ2 |x(i) |2 − b,

∀x ∈ Rd .

(10)

Lemma 1. Let u : Rd → R be L-semi-convex (A1), and suppose there exist µ0 , ρ > 0 such that the Alexandrov Hessian satisfies ∇2 u(z) ⪰ µ0 I a.e. for any z with |z| ≥ ρ. Then, for any selection h ∈ ∂u and all x, y, we have ⟨h(x) − h(y), x − y⟩ ≥ µ20 |x − y|2

whenever

|x − y| ≥ 4ρ(µµ00+L) .

Proof. The proof is postponed to Appendix G. 6. Main Results Theorem 1. Let Assumptions A1-A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for every λ ∈ (0, λ0 ) and n ∈ N0 , the subgradient tamed unadjusted Langevin algorithm (SG-TULAu) given in (3) satisfies W2 (πβ , L(θnλ,u )) ≤ 2CE λ1/4 + eCD e−nλ/T0 , where CE is given in Lemma 9, CD in Lemma 15 and T0 in (37). If additionally Assumption A5 holds, the same bound holds for W2 (πβ , L(θnλ,c )), with the constants CB1 , CB2 , CB4 replaced by CB1 ,c , CB2 ,c , CB4 ,c of Table 2. 8

Proof. The proof is postponed to Appendix G.  4 4 Corollary 1.Let ϵ >  0 and let the assumptions of Theorem 1 hold. Then, for λ ≤ min λ0 , ϵ /(4CE ) and T0 2eCD N≥ log iterations, one achieves λ ϵ λ,u W2 (πβ , L(θN )) ≤ ϵ.

Overall one chooses λ =
O(ϵ4 ) and N = O(ϵ−4
log(1/ϵ)). In terms of the dimension-to-temperature ratio, Table 2 yields λ = O (d/β)−2p and N = O (d/β)2p . The dependence of λ and N on the dimension is a feature of the assumption regime rather than of the scheme. The exponent 2p is inherited from the order λ1/4 of the strong error, which is what the present hypotheses on u permit. Under enough regularity on the Hessian, the strong error would be of order λ and the same argument would give exponent p/2, so that for p = 1 one matches the optimal dimension dependence (d/β)1/2 of a linearly growing gradient, as one has for a Gaussian target. Theorem 2. Let Assumptions A1-A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for any λ ∈ (0, λ0 ), n ∈ N and β ≥ max{4/µ, 1/J}, the following bound holds  
2(b + d/β)β 2 J 2 1 13 d log + log (πd) + , E[u(θnλ )] − u(θ∗ ) ≤ CT1 W2 πβ , L(θnλ ) + 2β µ 2β 6β where J = m + K22p−2 (1 + (2b/µ)p/2 ) + K2p−1 /(p + 1). Proof. Follows directly by invoking Lemmas 10 and 11. 7. Proof strategy Our proof strategy can be summarized in the following steps: 1. We begin our work by ensuring the well-posedness of the Langevin SDE and the uniqueness and existence of the invariant measure under our assumptions. 2. An important element of our proof strategy is to provide contraction rates for the SDE towards the invariant measure. The first step of our work, is to prove a Log-Sobolev inequality for the invariant measure under our assumptions, which relies on classical arguments with some delicate steps to cater for the presence of discontinuities. A combination of the LSI and the semi-convexity assumption yields an exponential contraction from the SDE towards the invariant measure (see Proposition 8), i.e W2 (νPt , πβ ) ≤ Cw e−Cr t W2 (ν, πβ ) , 3. For the discretization analysis we provide uniform in the number of iterations moment bounds for the process generated by the algorithm. The taming design of the scheme and the dissipativity condition are key in this proof. In addition we prove uniform moment bounds for an auxiliary process which is a Langevin SDE started at the law of the algorithm. The details can be found in Section C. 4. By using the growth conditions, the semi-convexity of the potential and the uniform moment bounds we are able to obtain finite time error estimates between the algorithm and the SDE started at a previous point of the algorithm for any n0 < n, denoting P̂ the ones step-Markov kernel of the algorithm (see Lemma 9), i.e. 1

W2 (P̂n−n0 P̂n0 , P̂n−n0 Pλn0 ) ≤ Cecλn0 λ 4 5. Combining the discretization and contraction results

W2 (L(θnλ ), πβ ) ≤ W2 (P̂n−n0 P̂n0 , P̂n−n0 Pλn0 ) + W2 (P̂n−n0 Pλn0 , πβ ) 1

≤ Cecλn0 λ 4 + Cw e−ċλn0 W2 (P̂n−n0 , πβ ) 1

λ = Cecλn0 λ 4 + Cw e−ċλn0 W2 (L(θn−n ), πβ ) 0

9

By picking a suitable n0 such that the prefactor in the second term of the right-hand side becomes smaller than e−1 , one achieves 1 λ ), πβ ) W2 (L(θnλ ), πβ )) ≤ C ∗ λ 4 + e−1 W2 (L(θn−n 0 which by applying a discrete Grownwall’s inequality yields Theorem 1. 6. The proof of Theorem 2 is based on decomposing E[u(θnλ )] − u(θ∗ ) = E[u(θnλ )] − E[u(θ∞ )] + E[u(θ∞ )] − u(θ∗ ) . | {z } | {z } T1

T2

The term T1 refers to the sampling problem and can be bounded by W2 (L(θnλ ), πβ ) in Lemma 10 while the second term refers to the concentration of the target measure around the minimum and can be controlled by a decaying function of β in Lemma 11. 8. Examples 8.1. Nanochat Pretraining We pretrain the nanochat language model of Karpathy [2025] at two depths, driving the matrix-valued parameters with the boosted coordinate-wise map of Lim and Sabanis [2024] evaluated at the subgradient selection h. The architecture, the potential u and the associated bounds are those of Appendix I. We refer to the resulting scheme as SG-TheoPouLA, and it attains the lowest validation bits per byte of the three optimisers at L = 12, where every configuration was tuned directly due to the lack of available scaling laws, and remains competitive at L = 24. ¯

The boosted tamed map. Writing λ for the stepsize and ε > 0 for the boosting threshold, the map acts on v ∈ Rd coordinate by coordinate as
(i)  Tλ,ε (v) = 1+

v (i) λ1/2  , ε + |v (i) | 1 + λ1/2 |v (i) |

¯ i ∈ {1, . . . , d},

(11)

where the right factor is the coordinate-wise tamed drift hλ,c of (4) taken with µ = 0, whereas the left factor lifts those coordinates that taming leaves small. The constant ε acts as a threshold rather than as a strength, since it fixes the magnitude below which a coordinate is treated as small and receives the lift, and it caps that lift at 1 + λ1/2 /ε. The device was introduced in Lim and Sabanis [2024] to address vanishing gradients, which is exactly the regime in which taming alone leaves a coordinate untouched and the iterate stalls, and it is what makes the map competitive on deep architectures. Algorithm 1 SG-TheoPouLA Require: stepsize schedule (λn ), threshold schedule (εn ), averaging coefficient β1 , weight decay w, initial θ0 and m0 = 0 1: for n = 0, 1, 2, . . . do 2: draw a mini-batch and form the subgradient Hn+1 ∈ ∂u(θn ) 3: ψn ← (1 − λn w) θn 4: mn+1 ← β1 mn + (1 − β1 )Hn+1 λn 5: θn+1 ← ψn − Tλn ,εn (mn+1 ) 1 − β1n+1 6: end for From (3) to Algorithm 1. Three choices specialise the analysed scheme to pretraining at this scale. The first is the boosted map itself, adopted because the taming and boosting combination is known from earlier work to be the more competitive of the two on deep architectures, and because the higher-order regularisation of (128) acting on the original gradient field corresponds, under the treatment of equation (11) of Lim and Sabanis [2024], to decoupled weight decay applied on the tamed field. The implementation therefore carries that regularisation in the form the tamed geometry prescribes, and the runs stay within the regime the analysis covers. The second is the temperature, since the 10

inverse temperatures relevant to pretraining are of order β ∼ 1012 and the injected noise term of (3) would contribute fluctuations orders of magnitude below those already supplied by mini-batch sampling. Drawing it would consume computation to no measurable effect, so we work in the zero temperature limit. The third and most consequential is the exponential moving average, which we form on the raw subgradients and to which the map is then applied. We find that ordering substantially stronger at this scale than its converse, in which each subgradient is tamed before entering the average. Two further details complete the description of the implementation. The stepsize entering (11) is the scheduled learning rate and therefore follows a linear decay to a fixed fraction of its peak, and the bias correction rescales the step rather than the buffer, so that the map is evaluated at the uncorrected average. The ε-threshold carries a schedule of its own, descending in log10 scale from its opening value to the working value of Table 4, which holds the multiplier applied to near-zero coordinates in check over the warmup steps while the average is still forming. Protocol. Muon [Jordan et al., 2024] is defined only for matrix-valued parameters and therefore requires a second optimiser for the embeddings, the unembedding and the scalars present in the architecture. We follow Karpathy [2025] in taking AdamW for those groups and hold its hyperparameters fixed throughout, so that every difference recorded below is attributable to the matrix optimiser alone and the three are compared on equal terms. Moreover, the pipeline equips Muon with three devices that are not an integral part of the algorithm but plug-ins, namely a second-moment normalisation of the step, a momentum schedule and a weight decay decaying to zero across training.We disable all three so that the comparison isolates the matrix update rule rather than the engineering accumulated around it. Both models take V = 65 536, context length T = 2048 and full causal attention at every layer, trained on the ClimbMix400B corpus of Diao et al. [2025]. The shallower model has dm = 768 and H = 6, whereas the deeper one has dm = 1536 and H = 12. The shallower runs cover 5000 steps at a batch of 524 288 tokens, whereas the deeper runs cover 7812 steps at a batch of 1 048 576 tokens, and in both cases the learning rate follows a warmup over the tabulated number of steps and a cosine decay thereafter. We report the validation bits per byte recorded by the training loop at the final step, together with CORE, the unweighted mean of the centred accuracies over the 22 tasks of the DCLM suite of Li et al. [2024]. Repeating a configuration unchanged returns validation bits per byte within 0.0008, so differences below that figure are read as level throughout. All runs were carried out on a cluster of four NVIDIA A100 80GB devices. Hyperparameter transfer. Depth L is the single scaling knob in the pipeline, setting the width as dm = 64L rounded up to a multiple of the head dimension 128, from which the number of heads follows as H = dm /128, and fixing the horizon as D = r Nscaling for a tokens-per-parameter ratio r, where Nscaling counts the transformer matrices together with the unembedding. The batch size follows B = Bref (D/Dref )0.383 rounded to the nearest power of two, and every rate but the hardcoded smear gate is then corrected by (B/Bref )1/2 . The table-shaped groups comprising the unembedding and the two embedding tables carry the further width correction (dm /768)−1/2 , from which the matrix and scalar rates are exempt, whereas the weight decay is set by w = wref (B/Bref )1/2 Dref /D so as to hold B/(ηλD) constant. Since doubling the depth doubles both the batch and the width, the two corrections cancel on the table-shaped groups and their rates hold at either scale, whereas the matrix and scalar rates retain the batch correction alone. The rules are calibrated for AdamW and Muon, which therefore move to the deeper model automatically and arrive already close to their best. No comparable study exists for the boosted tamed map, so its configuration at the greater depth was found by manual finetuning and its own scaling law remains to be quantified.

11

Figure 1: Pretraining at L = 12 under the tuned Muon, AdamW and SG-TheoPouLA configurations of Table 4.

Table 2: Pretraining at L = 12, where lower is better for validation bits per byte and higher is better for CORE, with the best of each column in bold. AdamW drives the non-matrix groups in every row. Optimiser Tuning Val bpb ↓ CORE ↑ Muon Karpathy 0.8253 0.1664 Muon ours 0.8199 0.1706 AdamW Karpathy 0.8579 0.1143 AdamW ours 0.8254 0.1523 SG-TheoPouLA (i) ours 0.8163 0.1638 SG-TheoPouLA (ii) ours 0.8199 0.1715 Figure 1 and Table 2 record the shallower comparison, in which SG-TheoPouLA reaches the lowest validation bits per byte of the six configurations under its first setting and the highest CORE under its second, the two differing in warmup, in weight decay and in the final learning rate fraction. Re-tuning improved both baselines over their published defaults, so the comparison is made against Muon and AdamW at their best rather than at their defaults.

Figure 2: Pretraining at L = 24 under the tuned Muon, AdamW and SG-TheoPouLA configurations of Table 4.

12

Table 3: Pretraining at L = 24, where the first two rows are published reference figures rather than runs of ours and the best among our runs is in bold. AdamW drives the non-matrix groups in every row. Model or optimiser Val bpb ↓ CORE ↑ OpenAI GPT-2 [Radford et al., 2019] — 0.2565 Karpathy baseline [Karpathy, 2025] 0.7186 0.2571 Muon 0.7067 0.2637 AdamW 0.7167 0.2577 SG-TheoPouLA 0.7133 0.2574 At the greater depth, recorded in Figure 2 and Table 3, Muon leads and SG-TheoPouLA follows it closely on validation bits per byte while remaining next to AdamW on CORE. All three improve on the published baseline, and they do so from a configuration that was tuned directly rather than transferred. Table 4: Final hyperparameters. The matrix rate and the weight decay are the values supplied to the pipeline, which rescales them with depth, so the rate at L = 24 is not directly comparable with the one above it. The learning rates for the embeddings, the unembedding and the scalars being 0.3, 0.008 and 0.5. Warmup is measured in steps and the final learning rate is a fraction of the peak. Optimiser Muon (Karpathy) Muon (ours) AdamW (Karpathy) AdamW (ours) SG-TheoPouLA (i) SG-TheoPouLA (ii) Muon AdamW SG-TheoPouLA

L Matrix lr Warmup Weight decay Final lr 12 0.02 40 0.15 0.05 12 0.008 40 0.15 0.05 12 3 × 10−4 40 0.15 0.05 12 2 × 10−3 40 0.15 0.05 12 0.01 120 0.10 0.05 12 0.01 80 0.15 0.10 24 0.02 40 0.15 0.05 24 1.8 × 10−3 50 0.15 0.05 24 7 × 10−3 120 0.18 0.10

EMA 0.90 0.90 0.80 0.80 0.82/0.95 0.82/0.95 0.90 0.80 0.82

Boosting — — — — 3 × 10−4 → 10−11 3 × 10−4 → 10−11 — — 3 × 10−4 → 10−11

8.2. Generalized phase retrieval Phase retrieval is the task of reconstructing a signal from the magnitudes of its linear measurements, the phases being lost; it is the standard model for X-ray crystallography, coherent diffraction imaging and astronomical imaging, where detectors record intensities but not phase Sun et al. [2018]. One fixes an unknown signal z⋆ ∈ Cn and known sensing vectors a1 , . . . , aN ∈ Cn , here drawn i.i.d. complex Gaussian, ak = √12 (Xk + iYk ) with Xk , Yk ∼ N (0, In ) independent, and observes only the magnitudes yk = |a∗k z⋆ |, k = 1, . . . , N . Since z⋆ and z⋆ eiϕ produce the same magnitudes, the signal can be recovered only up to a global phase z⋆ 7→ z⋆ eiϕ . Writing bk := yk2 , a standard way to fit the measurements is to minimise the quartic least-squares energy N

f (z) =

2 1 X ∗ 2 |ak z| − bk , 2N

z ∈ Cn ,

k=1

whose associated target is πβ ∝ e−βf . To read f as a potential on Rd we use the canonical identification Cn ∼ = Rd , d = 2n, sending z to x = (ℜz; ℑz). This map is an isometry: the real part of the Hermitian product becomes the Euclidean inner product and |z| = |x|, so lengths are preserved. Each squared magnitude is a Hermitian quadratic form, |a∗k z|2 = z ∗ (ak a∗k )z, and under the identification the Hermitian matrix ak a∗k is represented by the real symmetric positive semi-definite matrix  ⊤ pk pk + qk qk⊤ Ak = ⊤ qk p⊤ k − pk qk

 pk qk⊤ − qk p⊤ k ⊤ , pk p⊤ k + q k qk

ak = pk + iqk , pk , qk ∈ Rn ,

the unique matrix with x⊤ Ak x = |a∗k z|2 for every x; it has rank two and largest eigenvalue |ak |2 . Consequently f ,

13

written in the real variable x, is exactly the potential N

u(x) =

2 1 X ⊤ x Ak x − b k , 2N

(12)

x ∈ Rd ,

k=1

that is, u(x) = f (z) whenever x is the real form of z. Verifying A1–A3 for u is therefore legitimate and concerns the very function we sample. Being a polynomial of degree four, u is smooth, with gradient N
2 X ⊤ x Ak x − bk Ak x, N

h(x) = ∇u(x) =

(13)

k=1

and Hessian (14)

∇2 u(x) = u1 (x) + u2 (x) − u3 , where u1 (x) = N4

2 ⊤ k (Ak x)(Ak x) , u2 (x) = N

P

P

k (x

⊤

Ak x)Ak and u3 = N2

P

k bk Ak .

Assumptions A1 and A3 hold deterministically, for every realisation of the measurements. Remark 3 (A1 holds). For every fixed {(ak , yk )}N k=1 , the potential u is L-semi-convex with L = λmax (u3 ). Proof. In view of (14), the matrices u1 (x) and u2 (x) are positive semi-definite. In particular, u1 (x) is a sum of the rank-one terms (Ak x)(Ak x)⊤ , while u2 (x) combines the positive semi-definite Ak with the nonnegative weights x⊤ Ak x = |a∗k z|2 ≥ 0. Hence ∇2 u(x) ⪰ −u3 ⪰ −λmax (u3 )I = −LI for all x. Since u is smooth, integrating this bound along the segment γ(t) = y + t(x − y) gives Z 1 ⟨h(x) − h(y), x − y⟩ =

(x − y)⊤ ∇2 u(γ(t))(x − y) dt ≥ −L|x − y|2 .

0

Remark 4 (A3 holds). For every fixed {(ak , yk )}N k=1 , h = ∇u satisfies the polynomial growth (7) with p = 3, PN PN 2 2 2 2 4 m = N k=1 yk |ak | and K = N k=1 |ak | + m. Proof. By Cauchy–Schwarz, |a∗k z| ≤ |ak ||z| = |ak ||x|, so that x⊤ Ak x = |a∗k z|2 ≤ |ak |2 |x|2 and |Ak x| = |a∗k z| |ak | ≤ |ak |2 |x|. In view of (13), one calculates N N N 2 X  2 X 
2 X ⊤ 4 3 |h(x)| ≤ x Ak x + bk |Ak x| ≤ |ak | |x| + yk2 |ak |2 |x|. N N N k=1

k=1

k=1

Since |x| ≤ 1 + |x|3 , the linear term is bounded by m(1 + |x|3 ), and one obtains |h(x)| ≤ m + K|x|3 . This example involves two independent probabilistic mechanisms that must not be conflated. (M) The measurement model. The sensing vectors a1 , . . . , aN are drawn once, i.i.d. complex Gaussian, and then fixed. This draw is part of the inverse problem, not of our algorithm; the relevant probabilities are those of Sun et al. [2018]. (A) The algorithm. The injected noise and initialization of SG-TULA. All convergence guarantees of this paper are taken in expectation over (A) alone. The statement “holds with high probability” in Remark 5 below is a probability over (M) only, a property of the measurement ensemble, certified in Sun et al. [2018] and established before our algorithm runs. We fix one such instance, on it u is a deterministic function satisfying A1–A3, and every result of this paper then applies to it deterministically. No probability over (M) ever enters a guarantee over (A).

14

Remark 5 (A2 holds for a generic measurement ensemble). Let a1 , . . . , aN be i.i.d. CN (0, In ), and write L = λmax (u3 ) as in Remark 3. For N large enough, with high probability over the draw of {ak } the fixed potential u satisfies A2: for every µ > 0, the bound (6) holds with radius R=

2 (2µ + L)3/2 . µ

Proof. By [Sun et al., 2018, Lemma 6.4] (with δ = 21 ), for N large enough the event N

E:

1 X ∗ 2 ∗ 2 1 |ak z| |ak w| ≥ |z|2 |w|2 N 2

for all z, w ∈ Cn

k=1

holds with high probability over the draw of {ak }. We fix such a draw. Realifying the bound defining E (using x⊤ Ak x = |a∗k z|2 , w⊤ Ak w = |a∗k w̃|2 and |z| = |x|, |w̃| = |w|) gives, for every w ∈ Rd , N

2 X ⊤ w u2 (x) w = (x Ak x)(w⊤ Ak w) ≥ |x|2 |w|2 , N ⊤

that is

u2 (x) ⪰ |x|2 I.

k=1

Since u1 (x) ⪰ 0 and u3 ⪯ LI, (14) yields ∇2 u(x) ⪰ u2 (x) − u3 ⪰ (|x|2 − L)I. Fix µ > 0 √ √ the decomposition and set µ0 := 2µ and ρ := µ0 + L = 2µ + L; then ∇2 u(x) ⪰ µ0 I whenever |x| ≥ ρ. The potential is Lsemi-convex by Remark 3, so Lemma 1 applies with these µ0 , ρ, L and gives (6) with constant µ0 /2 = µ and radius 4ρ(µ0 + L)/µ0 = 2(2µ + L)3/2 /µ. In practice the energy (12) is rarely minimised on its own. It is common, when the signal is sparse, to be augmented with a sparsity promoting ℓ1 penalty, giving the standard model of sparse phase retrieval Shechtman et al. [2014], N

ũ(x) =

2 1 X ⊤ x Ak x − bk + τ |x|1 , 2N

τ > 0.

k=1

The term τ |x|1 keeps ũ continuous but, being non-differentiable on the coordinate hyperplanes {x(i) = 0}, introduces gradient discontinuities on this set of measure zero, where √ the drift is given by a bounded subgradient selection. As τ |x|1 is convex with subgradient of norm at most τ d, Assumptions A1–A3 are preserved. The target is thus simultaneously superlinear, non-smooth and non-convex. 9. References M. Barkhagen, N.H. Chau, É. Moulines, M. Rásonyi, S. Sabanis and Y. Zhang: On stochastic gradient Langevin dynamics with dependent data streams in the logconcave case, Bernoulli 27(1) (2021) 1–33. S. G. Bobkov and M. Ledoux: From Brunn-Minkowski to Brascamp-Lieb and to logarithmic Sobolev inequalities, Geom. Funct. Anal. 10 (2000) 1028–1052. V.I. Bogachev, M. Röckner and W. Stannat: Uniqueness of invariant measures and essential m-dissipativity of diffusion operators on L1 , Scuola Norm. Superiore (2000). N. Brosse, A. Durmus, É. Moulines and M. Pereyra: Sampling from a log-concave distribution with compact support with proximal Langevin Monte Carlo, Proc. Conf. Learn. Theory, Proc. Mach. Learn. Res. 65 (2017) 319–342. N. Brosse, A. Durmus, É. Moulines and S. Sabanis: The tamed unadjusted Langevin algorithm, Stoch. Proc. App. 129 (2019) 3638–3663. P. Cattiaux, A. Guillin and L.-M. Wu: A note on Talagrand’s transportation inequality and logarithmic Sobolev inequality, Probab. Theory Related Fields 148(1–2) (2010) 285–304. N. Chatterji, J. Diakonikolas, M. I. Jordan and P. Bartlett: Langevin Monte Carlo without smoothness, Proc. Int. Conf. Artif. Intell. Stat., Proc. Mach. Learn. Res. 108 (2020) 1716–1726. 15

N.H. Chau, É. Moulines, M. Rásonyi, S. Sabanis and Y. Zhang: On stochastic gradient Langevin dynamics with dependent data streams: the fully non-convex case, SIAM J. Math. Data Sci. 3(3) (2021) 959–986. A. Y. Chen and K. Sridharan: Optimization, Isoperimetric Inequalities, and Sampling via Lyapunov Potentials, Proc. Mach. Learn. Res. 291 (2025) 1–60. S. Chewi, M. A. Erdogdu, M. Li, R. Shen and M. S. Zhang: Analysis of Langevin Monte Carlo from Poincaré to log-Sobolev, Found. Comput. Math. 25(4) (2025) 1345–1395. S. Chewi and A. J. Stromme: The ballistic limit of the log-Sobolev constant equals the Polyak-Łojasiewicz constant, arXiv preprint arXiv:2411.11415, (2024). A. S. Dalalyan: Theoretical guarantees for approximate sampling from smooth and log-concave densities, J. R. Stat. Soc. Ser. B Stat. Methodol. 79 (2017) 651–676. A.S. Dalalyan and A. Karagulyan: User-friendly guarantees for the Langevin Monte Carlo with inaccurate gradient, Stoch. Proc. App. 129(12) (2019) 5278–5311. K. Dareiotis and M. Gerencsér: On the regularisation of the noise for the Euler–Maruyama scheme with irregular drift, Electron. J. Probab. 25 (2020) Paper No. 82. K. Dareiotis, M. Gerencsér and K. Lê: Quantifying a convergence theorem of Gyöngy and Krylov, Ann. Appl. Probab. 33(3) (2023) 2291–2323. A. Durmus and É. Moulines: Nonasymptotic convergence analysis for the unadjusted Langevin algorithm, Ann. Appl. Probab. 27 (2017) 1551–1587. A. Durmus and É. Moulines: High-dimensional Bayesian inference via the unadjusted Langevin algorithm, Bernoulli 25 (2019) 2854–2882. A. Durmus, S. Majewski and B. Miasojedow: Analysis of Langevin Monte Carlo via convex optimization, J. Mach. Learn. Res. 20 (2019) 73:1–46. A. Durmus, É. Moulines and M. Pereyra: A proximal Markov chain Monte Carlo method for Bayesian inference in imaging inverse problems: When Langevin meets Moreau, SIAM Rev. 64 (2022) 991–1028. A. Eberle: Reflective coupling and its applications to stochastic processes, Probab. Theory Relat. Fields 164 (2016) 73–112. M. J. Ehrhardt, L. Kuger and C.-B. Schönlieb: Proximal Langevin sampling with inexact proximal mapping, SIAM J. Imaging Sci. 17 (2024) 1729–1760. M. A. Erdogdu and R. Hosseinzadeh: On the convergence of Langevin Monte Carlo: the interplay between tail growth and smoothness, Proc. Conf. Learn. Theory, Proc. Mach. Learn. Res. 134 (2021) 1776–1822. M. A. Erdogdu, R. Hosseinzadeh and S. Zhang: Convergence of Langevin Monte Carlo in chi-squared and Rényi divergence, Proc. Int. Conf. Artif. Intell. Stat., Proc. Mach. Learn. Res. 151 (2022) 8151–8175. L. C. Evans and R. F. Gariepy: Measure Theory and Fine Properties of Functions, Revised Edition, Textbooks in Mathematics, CRC Press (2015). L. Fruehwirth and A. Habring: Ergodicity of Langevin dynamics and its discretizations for non-smooth potentials, arXiv preprint arXiv:2411.12051, (2024). M. Gurbuzbalaban, H. M. Nguyen, X. Zhang and L. Zhu: Anchored Langevin algorithms, arXiv preprint arXiv:2509.19455, (2025). I. Gyöngy and N. V. Krylov: Existence of strong solutions for Itô’s stochastic equations via approximations: revisited, Stoch. Partial Differ. Equ. Anal. Comput. 10 (2022) 693–719. A. Habring, A. Falk, M. Zach and T. Pock: Diffusion at absolute zero: Langevin sampling using successive Moreau envelopes, SIAM J. Imaging Sci. 19(1) (2026) 35–77. 16

J.-B. Hiriart-Urruty and C. Lemaréchal: Fundamentals of Convex Analysis, Grundlehren Text Editions, Springer, Berlin (2001). M. Hutzenthaler, A. Jentzen and P. E. Kloeden: Strong and weak divergence in finite time of Euler’s method for stochastic differential equations with non-globally Lipschitz continuous coefficients, Proc. R. Soc. A Math. Phys. Eng. Sci. 467(2130) (2010) 1563–1576. M. Hutzenthaler, A. Jentzen and P. E. Kloeden: Strong convergence of an explicit numerical method for SDEs with nonglobally Lipschitz continuous coefficients, Ann. Appl. Probab. 22 (2012). C.-R. Hwang: Laplace’s method revisited: weak convergence of probability measures, Ann. Probab. 8(6) (1980) 1177– 1182. T. Johnston, I. Lytras and S. Sabanis: Kinetic Langevin MCMC sampling without gradient Lipschitz continuity - the strongly convex case, J. Complexity 85 (2024) 101873. T. Johnston, I. Lytras, N. Makras and S. Sabanis: The performance of the unadjusted Langevin algorithm without smoothness assumptions, Trans. Mach. Learn. Res. (2025). T. Johnston, N. Makras and S. Sabanis: Taming the interacting particle Langevin algorithm: the superlinear case, Appl. Math. Optim. 91(3) (2025) Paper No. 77. A. Karpathy: nanochat: The best ChatGPT that $100 can buy, GitHub repository, (2025). https://github.com/karpathy/nanochat Y. Kinoshita and T. Suzuki: Improved convergence rate of stochastic gradient Langevin dynamics with variance reduction and its application to optimization, Advances in Neural Information Processing Systems (NeurIPS) 35 (2022) 19022–19034. R. Laumont, V. De Bortoli, A. Almansa, J. Delon, A. Durmus and M. Pereyra: Bayesian imaging using Plug & Play priors: when Langevin meets Tweedie, SIAM J. Imaging Sci. 15(2) (2022) 701–737. J. Lehec: The Langevin Monte Carlo algorithm in the non-smooth log-concave case, Ann. Appl. Probab. 33(6A) (2023) 4858–4874. M. B. Li and M. A. Erdogdu: Riemannian Langevin algorithm for solving semidefinite programs, Bernoulli 29(4) (2023) 3093–3113. R. Li, H. Zha and M. Tao: Sqrt(d) dimension dependence of Langevin Monte Carlo, International Conference on Learning Representations (ICLR) (2022). D.-Y. Lim and S. Sabanis: Polygonal unadjusted Langevin algorithms: creating stable and efficient adaptive algorithms for neural networks, J. Mach. Learn. Res. 25(53) (2024) 1–52. A. Lovas, I. Lytras, M. Rásonyi and S. Sabanis: Taming neural networks with TUSLA: nonconvex learning via adaptive stochastic gradient Langevin algorithms, SIAM J. Math. Data Sci. 5(2) (2023) 323–345. D. Luo and J. Wang: Exponential convergence in Lp -Wasserstein distance for diffusion processes without uniformly dissipative drift, Math. Nachr. 289 (2016) 1909–1926. H. P. H. Luu and Z. Wang: DC-LA: Difference-of-Convex Langevin Algorithm, arXiv preprint arXiv:2601.22932, (2026). I. Lytras and P. Mertikopoulos: Tamed Langevin sampling under weaker conditions, Proc. Int. Conf. Artif. Intell. Stat., Proc. Mach. Learn. Res. (2025) 847–855. I. Lytras and A. Ntousis: Error estimates for tamed Euler and randomized Euler schemes for SDEs with locally Lipschitz drift with applications to non-logconcave sampling and optimization, arXiv preprint arXiv:2605.24937, (2026). I. Lytras and S. Sabanis: Taming under isoperimetry, Stoch. Proc. App. 188 (2025) 104684. 17

I. Lytras, S. Sabanis and Y. Zhang: kTULA: A Langevin sampling algorithm with improved KL bounds under superlinear log-gradients, arXiv preprint arXiv:2506.04878, (2025). Y.A. Ma, Y. Chen, C. Jin, N. Flammarion and M.I. Jordan: Sampling can be faster than optimization, Proc. Natl. Acad. Sci. USA (2019). P. A. Markowich and C. Villani: On the trend to equilibrium for the Fokker-Planck equation: an interplay between physics and functional analysis, Mat. Contemp. 19 (2000) 1–29. G. Menz and A. Schlichting: Poincaré and logarithmic Sobolev inequalities by decomposition of the energy landscape, Ann. Probab. 42(5) (2014) 1809–1884. P. Monmarché: Wasserstein contraction and Poincaré inequalities for elliptic diffusions with high temperature, Ann. H. Lebesgue 6 (2023) 941–973. A. Mousavi-Hosseini, T. Farghly, Y. He, K. Balasubramanian and M. A. Erdogdu: Towards a complete analysis of Langevin Monte Carlo: beyond Poincaré inequality, Proc. Conf. Learn. Theory, Proc. Mach. Learn. Res. 195 (2023) 1–35. T. Müller-Gronbach and L. Yaroslavtseva: On the performance of the Euler–Maruyama scheme for SDEs with discontinuous drift coefficient, Ann. Inst. Henri Poincaré Probab. Stat. 56(2) (2020) 1162–1178. T. Müller-Gronbach, C. Rauhögger and L. Yaroslavtseva: On the performance of the Euler–Maruyama scheme for multidimensional SDEs with discontinuous drift coefficient, Numer. Math. (to appear), arXiv:2504.01630, (2025). T. Müller-Gronbach, S. Sabanis and L. Yaroslavtseva: Existence, uniqueness and approximation of solutions of SDEs with superlinear coefficients in the presence of discontinuities of the drift coefficient, Commun. Math. Sci. 23 (2026) 389–429. A. Neufeld, M. C. E. Ng and Y. Zhang: Non-asymptotic convergence bounds for modified tamed unadjusted Langevin algorithm in non-convex setting, J. Math. Anal. Appl. 543(1) (2025) 128892. D. Nguyen, X. Dang and Y. Chen: Unadjusted Langevin algorithm for non-convex weakly smooth potentials, Commun. Math. Stat. (2023). F. Otto and C. Villani: Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality, J. Funct. Anal. 173(2) (2000) 361–400. R. Pascanu, T. Mikolov and Y. Bengio: On the difficulty of training recurrent neural networks, Proc. 30th Int. Conf. Mach. Learn. (ICML), Proc. Mach. Learn. Res. 28(3) (2013) 1310–1318. M. Raginsky, A. Rakhlin and M. Telgarsky: Non-convex learning via stochastic gradient Langevin dynamics: a nonasymptotic analysis, Conf. Learn. Theory (2017) 1674–1703. M. Renaud, V. De Bortoli, A. Leclaire and N. Papadakis: From stability of Langevin diffusion to convergence of proximal MCMC for non-log-concave sampling, Advances in Neural Information Processing Systems (NeurIPS) (2025). G.O. Roberts and R.L. Tweedie: Exponential convergence of Langevin distributions and their discrete approximations, Bernoulli 2(4) (1996) 341–363. R. T. Rockafellar: Convex Analysis, Princeton University Press (1970). R.T. Rockafellar and R.J-B. Wets: Variational Analysis, Springer (1998). S. Sabanis: A note on tamed Euler approximations, Electr. Commun. Probab. 18 (2013) 1–10. S. Sabanis and Y. Zhang: A fully data-driven approach to minimizing CVaR for portfolio of assets via SGLD with discontinuous updating, arXiv preprint arXiv:2007.01672, (2020). J. Shao: Hamilton-Jacobi semi-groups in infinite dimensional spaces, Bull. Sci. Math. 130 (2006) 720–738. 18

Y. Shechtman, A. Beck and Y. C. Eldar: GESPAR: efficient phase retrieval of sparse signals, IEEE Trans. Signal Process. 62(4) (2014) 928–938. J. Sun, Q. Qu and J. Wright: A geometric analysis of phase retrieval, Found. Comput. Math. 18(5) (2018) 1131–1198. S. Vempala and A. Wibisono: Rapid convergence of the unadjusted Langevin algorithm: isoperimetry suffices, Advances in Neural Information Processing Systems (NeurIPS) 32 (2019). C. Villani: Topics in Optimal Transportation, Graduate Studies in Mathematics 58, American Mathematical Society, Providence, RI (2003). C. Villani: Optimal Transport: Old and New, Grundlehren der mathematischen Wissenschaften 338, Springer-Verlag, Berlin (2009). F.-Y. Wang: Logarithmic Sobolev inequalities: conditions and counterexamples, J. Operator Theory 46(1) (2001) 183–197. F.-Y. Wang: Harnack inequalities on manifolds with boundary and applications, J. Math. Pures Appl. 94 (2010) 304– 321. F.-Y. Wang: Exponential contraction in Wasserstein distances for diffusion semigroups with negative curvature, Potential Anal. 53(3) (2020) 1123–1144. X. Wang and B. Yang: When Langevin Monte Carlo meets randomization: non-asymptotic error bounds beyond log-concavity and gradient Lipschitzness, arXiv preprint arXiv:2509.25630, (2025). √ B. Yang and X. Wang: Non-asymptotic error bounds in W2 -distance with d dimension dependence and first order convergence for Langevin Monte Carlo beyond log-concavity, Proc. 42nd Int. Conf. Mach. Learn. (ICML) (2025) 71358–71382. S. Diao, Y. Yang, Y. Fu, X. Dong, D. Su, M. Kliegl, Z. Chen, P. Belcak, Y. Suhara, H. Yin, M. Patwary, Y. Lin, J. Kautz and P. Molchanov: CLIMB: clustering-based iterative data mixture bootstrapping for language model pre-training, arXiv preprint arXiv:2504.13161, (2025). K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse and J. Bernstein: Muon: an optimizer for hidden layers in neural networks, Blog post, (2024). https://kellerjordan.github.io/posts/muon/ J. Li, A. Fang, G. Smyrnis et al.: DataComp-LM: in search of the next generation of training sets for language models, Adv. Neural Inf. Process. Syst. (NeurIPS) 37 (2024) 14200–14282. A. Radford, J. Wu, R. Child, D. Luan, D. Amodei and I. Sutskever: Language models are unsupervised multitask learners, Technical report, OpenAI, (2019). Appendix A. Well posedness Consider the infinitesimal generator L associated with (2) defined for all ϕ ∈ C 2 (Rd ) and x ∈ Rd by Aϕ(x) = −⟨h(x), ∇ϕ(x)⟩ + β −1 ∆ϕ(x). Next define the Lyapunov function Vp (x) = (1 + |x|2 )p/2 for all x ∈ Rd . Note that Vp is twice continuously differentiable, and under Assumptions A2 and A3 so that Remark 2 holds, one gets the following growth condition AVp (x) ≤ C2 Vp (x), ∀x ∈ Rd ,

(15)

where C2 is given in Lemma 12. Moreover under Assumptions A2 and A3, it satisfies the geometric drift condition AVp (x) ≤ −C3 (p)Vp (x) + C4 (p), ∀x ∈ Rd ,

(16)

where C3 (p) and C4 (p) are given in Lemma 13. It follows that lim Vp (x) = +∞, lim LVp (x) = −∞.

|x|→∞

|x|→∞

19

(17)

Proposition 1. Let Assumptions A1-A4 hold. The SDE (2) has a unique strong solution. Proof. Uniqueness is guaranteed under the monotonicity condition A1 and due to the diffusion coefficient being constant. Moreover, all conditions of Theorem 2.8 in Gyöngy and Krylov [2022] are satisfied under our assumptions; therefore, the SDE (2) admits a unique strong solution. In particular, since the drift coefficient is subject to the growth condition of Remark 2 and the diffusion coefficient is constant, in view of (15), they trivially satisfy the conditions (i), (ii) and (iv). Condition (iii) is also satisfied trivially as in our case the domain is D = Rd . Proposition 2. Let Assumptions A1-A3 hold, the Langevin SDE (2) admits a unique invariant measure. Proof. The existence of an invariant measure is established under Assumptions A2 and A3. In particular, the Langevin SDE (2) has a constant diffusion coefficient and Assumption A3 ensures that the drift coefficient is locally integrable. Consequently, in view of (17), all the conditions of Theorem 2.2 in Bogachev et al. [2000] are satisfied, the existence of at least one invariant measure follows. Moreover, with the inclusion of Assumption A1, the contraction results in Section D imply the uniqueness of the invariant measure. This is a direct consequence of Proposition 5, by setting the initial condition Z0 in (2) to be such that L(Z0 ) = L(πβ ). Proposition 3. Let Assumptions A1-A3 hold. The invariant measure πβ of the SDE (2), is characterized by the density Z −1 exp(−βu(x)), with Z being the normalization constant. Proof. Under Assumption A3 one yields that u ∈ H1loc and the rest follow from Theorem 3 in Fruehwirth and Habring [2024]. Appendix B. Technical Prerequisites B.1. Auxiliary Processes We next introduce the auxiliary processes which are used in our analysis. For each λ > 0, the time-scaled process (Ztλ )t∈R+ is defined by Ztλ := Zλt , t ∈ R+ . We note that dZtλ = −λh(Ztλ )dt +

p

2λβ −1 dB̃tλ , Z0λ = θ0 ,

(18)

√ where the Brownian motion (B̃tλ )t≥0 is defined as B̃tλ := Bλt / λ, t ≥ 0. The natural filtration of (B̃tλ )t≥0 is denoted by (Ftλ )t≥0 with Ftλ := Fλt , t ∈ R+ . Then, we define (θ̄tλ )t∈R+ , the continuous-time interpolation of SG-TULA (3), as p λ dθ̄tλ = −λhλ (θ̄⌊t⌋ )dt + 2λβ −1 dB̃tλ , θ̄0λ = θ0 . (19) The law of this process coincides with the law of the algorithm at grid points i.e. L(θ̄nλ ) = L(θnλ ) for every n ∈ N. λ Definition 1. Fix n0 ∈ N. For each n ≥ n0 , let (Wtλ,n )t∈[0,n0 ] denote the Brownian increments Wtλ,n := B̃t+(n−n − 0) λ B̃(n−n . 0)

• The restarted continuous process (ζ̄tλ,n )t∈[0,n0 ] is defined by dζ̄tλ,n = −λh(ζ̄tλ,n )dt +

p 2λβ −1 dWtλ,n ,

λ ζ̄0λ,n = θn−n 0

• The restarted interpolation process (ϑλ,n t )t∈[0,n0 ] is defined by dϑλ,n = −λhλ (ϑλ,n t ⌊t⌋ )dt +

p λ 2λβ −1 dWtλ,n , ϑλ,n = θn−n 0 0

By construction, both processes share the same initial condition and driving noise, and evolve on a common time inλ terval [0, n0 ]. The restarted interpolation process casts n0 steps of SG-TULA, so that L(ϑλ,n n0 ) = L(θ̄n ). Additionally, λ L(ζ̄nλ,n ) corresponds to the same law as the time-scaled Langevin SDE (18), started from the initial condition θn−n 0 0 λ at time n − n0 , and evolves over a fixed time interval of length n0 . That is L(ζ̄nλ,n ) = L( θ̄ )P . n0 λ n−n0 0

20

B.2. Taming Properties Remark 6. For all λ > 0 and x ∈ Rd , it holds that |hλ,u (x)| ≤ µ|x| + λ−1/2 and |hλ,c (x)| ≤

√

√ 2µ|x| +

(20)

2dλ−1/2 .

Proof. The proof is postponed to Appendix G. Remark 7. Let Assumption A3 hold. Then for all λ > 0 and x ∈ Rd , it holds that (21)

|hλ (x) − h(x)| ≤ C1 λ1/2 (1 + |x|2p ),  where either hλ = hλ,u or hλ = hλ,c and C1 = 2µ2 + 4 max m2 , K 2 . Proof. The proof is postponed to Appendix G. Remark 8. Let Assumptions A2-A3 hold. Then for all λ > 0 and x ∈ Rd , it holds that ⟨x, hλ,u (x)⟩ ≥

µ 2 |x| − b, 2

(22)

where b is given in Remark 2. In addition, let Assumption A5 hold. Then for all λ > 0 and x ∈ Rd , it holds that ⟨x, hλ,c (x)⟩ ≥

µ 2 |x| − bd. 2

Proof. The proof is postponed to Appendix G. Appendix C. Moments & One-step Error Estimates Lemma 2. Let Assumptions A2–A3 and A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for every λ ∈ (0, λ0 ) and n ∈ N, one has (23)

E|θnλ |2 ≤ CB1 , where CB1 = E|θ0 |2 + 4(1/µ)(1 + b + dβ −1 ). Proof. The proof is postponed to Appendix G.

Observe that the step-size restriction depends only on µ and not on the growth of the drift. This is a benefit of taming, by (20) the tamed drift grows linearly with constant µ, so µ takes the place of the global Lipschitz constant that the original superlinear drift does not possess. Lemma 3. Let Assumptions A2-A3 and A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for every λ ∈ (0, λ0 ), p ∈ N \ {1} and n ∈ N, one has (24)

E|θnλ |2p ≤ CB2 , where CB2 = E|θ0 |2p + M (p)(1/µ)p (1 + b + dβ −1 )p with M (p) = p(p + 1)26p−2 ((2p − 1)!!)

1/2

.

Proof. The proof is postponed to Appendix G. In particular, the dimension enters only through the ratio d/β. Lemma 4. Let Assumptions A2-A3 and A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for every λ ∈ (0, λ0 ), p ∈ N \ {1}, n0 ∈ N and n ≥ n0 one has sup E[V2p (ζ̄tλ,n )] ≤ CB3 ,

(25)

t∈[0,n0 ]

where CB3 = 2p/2−1 (1 + E|θ0 |p + Cp,µ,b,β,d ), with Cp,µ,b,β,d = O((d/β)p/2 ). Note that CB3 does not depend on n0 . 21

Proof. The proof is postponed to Appendix G. Lemma 5. Let Assumptions A2-A3 and A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for every λ ∈ (0, λ0 ) and t ≥ 0 one has λ E|θ̄⌊t⌋ − θ̄tλ |2 ≤ CB4 λ,

(26)

where CB4 = µCB1 + 4 + 4dβ −1 . The same bound holds for the restarted interpolation (ϑλ,n t )t∈[0,n0 ] of Definition 1. Proof. The proof is postponed to Appendix G. Appendix D. Contraction Estimates D.1. Log-Sobolev Inequality Proposition 4. Let Assumptions A1 and A2 hold. Then the Langevin SDE (2) admits a unique invariant measure −βu π which satisfies a logarithmic Sobolev inequality with some constant ρ(β) > 0, for every ϕ ∈ Cb1 (Rd ) with R β ∝ 2e ϕ dπβ = 1, Rd Z Z ϕ2 log ϕ2 dπβ ≤ ρ(β) |∇ϕ|2 dπβ . (27) Rd

Rd

Proof. Under A1 the potential βu is semi-convex with Hess(βu) ⪰ −βL Id distributionally, and A2 provides a dissipativity at infinity condition. Thus the LSI follows from [Cattiaux et al., 2010, Corollary 2.1], with the HWI step of the underlying Theorem 1.9 justified for our non-smooth βu by [Villani, 2003, Theorem 9.17] as explained in Remark 9. Remark 9. Proposition 4 follows from the curvature of the potential alone. The principle goes back to Wang [2001], where, in the smooth case, a negative uniform lower bound on the curvature (semi-convexity, A1) together with finite exponential moments of πβ suffices for a LSI. However we invoke the result of [Cattiaux et al., 2010, Corollary 2.1] instead, which derives the LSI from a Lyapunov (dissipativity) condition together with the curvature lower bound Hess(V ) ≥ K Id (here V = βu, K = −βL, i.e. A1) and reduces to Theorem 1.9 therein, combining a defective transport inequality with the HWI inequality. Every step of that argument requires no smoothness of u, with a single exception being the HWI inequality, which is quoted from Otto and Villani [2000] in the form proved there for V ∈ C 2 . This is the only place C 2 -smoothness enters. Replacing that reference by [Villani, 2003, Theorem 9.17] establishes the same HWI inequality assuming only a negative lower bound for the distributional hessian. Their argument uses the fact that the Alexandrov Hessian isR dominated by the distributional one. The remaining hypotheses ofR [Villani, 2003, Theorem 9.17] hold here, that is, V e−V < ∞ and πβ ∈ Pac,2 (Rd ) follow from the moment bound Rd V2p dπβ ≤ C4 (2p)/C3 (2p) established in the proof of Lemma 10. Notice that A2 in turn implies dissipativity at infinity, and, together with A3, the global dissipativity of Remark 2. In fact A1 together with dissipativity already suffices for the log-Sobolev inequality, so A2 could be dropped in favour of a dissipativity assumption. We nevertheless retain the strong convexity at infinity A2 for two reasons. In view of Lemma 1 it is straightforward to verify, even for certain neural network architectures, and the dissipativity route yields a log-Sobolev constant that is in general exponential in d, whereas the addition of A2-A3 gives the dimension-free ρ(β) of Remark 10. Remark 10. Let Assumptions A1, A2 and A3 hold. Then the log-Sobolev inequality (27) holds with ρ(β) ≤

 h i 2 exp β 4R (m + KRp ) + µ2 R2 . βµ

Proof. The key idea rests on a perturbation argument that requires no smoothness of the potential. A βµ-strongly convex proxy W0 is constructed, which coincides with βu outside B(0, R) (Lemma 16). The reference measure π0 ∝ e−W0 satisfies a log-Sobolev inequality with constant 2/(βµ) by the Bobkov–Ledoux criterion [Bobkov and Ledoux, 2000, Proposition 3.1], which, in contrast to the Bakry–Émery criterion, applies to non-smooth strongly convex potentials. Since the perturbation ψ = βu − W0 is supported in B(0, R) with osc(ψ) ≤ β[4R(m + KRp ) + µ 2 2 R ] (Lemma 17), the Holley–Stroock perturbation principle (see Theorem 1, Chapter 5 in Markowich and Villani [2000]) yields the stated constant. 22

Remark 10 is a consequence of standard results linking dissipativity, curvature bounds and functional inequalities. Strong convexity at infinity (A2) implies that the drift is dissipative at infinity, see Proof of Remark 2. This yields a Lyapunov condition ensuring that the invariant measure has finite exponential moments, up to a constant in the exponent. Semi-convexity (A1) provides a negative uniform lower bound on the curvature of the potential. Combined, they allow one to establish a LSI for π, see Wang [2001] and references therein. The resulting LSI constant is dimension dependent as the proofs are constructive. However, under the growth condition (A3), the LSI constant can be explicitly controlled using the Holley-Stroock perturbation theorem together with the Bakry-Emery criterion. Similar arguments in the smooth potential setting can be found in Ma et al. [2019]. The significance of Remark 10 is primarily practical. Verifying a LSI directly for a given target measure is often infeasible beyond idealized settings, whereas strong convexity at infinity is typically much easier to establish and can often be checked by algebraic arguments. This includes many nontrivial examples of interest, such as mixture of Gaussians or double-well potentials, which are nonconvex but strongly convex outside a compact set. Remark 11. [Extra assumptions for polynomial dependence] Although in the general case one cannot avoid exponential dependence on β, an additional structure can yield polynomial dependence on β. Earlier works proved this result under smoothness assumptions and strict conditions on the local minima and saddle points in compact Riemannian manifolds (Li and Erdogdu [2023]) and Euclidean spaces (Lytras and Sabanis [2025],Kinoshita and Suzuki [2022]). This line of work was later improved under the assumption of the popular in optimization PL condition i.e for some CP L > 0 u(x) − u∗ ≤ CP L |∇f (x)|2 ∀x ∈ Rd (28) In Chewi and Stromme [2024] it was proved that for u ∈ C 2 such that ∆u ≤ C(1 + ||∇u||2 ) satisfying (28), having a unique global minimizer there holds βρ(β) → CP L as β → ∞. This was further generalized for multiple global minimizers in Chen and Sridharan [2025] where it was shown that under (28) and some dissipativity and growth conditions on the first and second derivatives, that the target measure satisfies a Poincaré inequality (PI) that depends on the local Poincaré constant in a ball around the minimizers set B(W ∗ , r) for some r > 0. Under the additional assumption of 2-dissipativity and semi-convexity it was shown (see Theorem 3.1 in Chen and Sridharan [2025]) that the target measure satisfies LSI with constant ρ(β) ≤ β 2 CP I,B(W ∗ ,r) . This shows that the global LSI constant is connected to the local behavior near the minimizers. More specifically, if the measure admits a Poincaré constant in B(W ∗ , r) that is independent of β (for example if B(W ∗ , r) is convex and u is convex on this set ), the resulting LSI constant is polynomial in β. D.2. Exponential Contraction We follow through the proof of [Wang, 2020, Theorem 2.1(3)], in which A1-A2 are assumed for a smooth drift h ∈ C 2 . The reason for reproducing the argument is twofold: to identify the explicit constants, and to accommodate our setting, in which h is a measurable selection of ∂u and therefore discontinuous. For the latter we use a convolution smoothing argument, which allows us to invoke tools such as the log-Harnack inequality at the level of a smooth potential. Lemma 6. Let φ ∈ Cc∞ (Rd ) satisfy Z φ ≥ 0,

supp φ ⊆ B(0, 1),

φ(w) dw = 1,

(29)

Rd

and for ϵ ∈ (0, 1] define φϵ (w) := ϵ−d φ(w/ϵ), which inherits the properties (29) with supp φϵ ⊆ B(0, ϵ). Set Z hϵ (x) := (h ∗ φϵ )(x) = h(x − w)φϵ (w) dw, uϵ (x) := (u ∗ φϵ )(x). Rd

Let Assumptions A1-A3 hold. Then: (i) uϵ , hϵ ∈ C ∞ (Rd ) and ∇uϵ = hϵ ; (ii) hϵ satisfies A1 with the same constant L and A2 with the same constants µ and R; 23

(30)

(iii) hϵ satisfies A3 in the form p

p

∀x ∈ Rd , ∀ϵ ∈ (0, 1];

|hϵ (x)| ≤ m + K (|x| + ϵ) ≤ m + K (|x| + 1) , (iv) the measure πβϵ ∝ e−βuϵ satisfies (27) with constant ρϵ (β) =

 h 2 µ i exp β 4R (m + K(R + ϵ)p ) + R2 ; βµ 2

(31)

(v) as ϵ → 0, uϵ → u locally uniformly and hϵ → h Lebesgue-almost everywhere on Rd . Proof. The proof is postponed to Appendix G √ Define the time-scaled process Ys := Zβs and B̂s := Bβs / β, then for ϵ ∈ (0, 1] let Y ϵ be the corresponding process with the mollified drift hϵ of Lemma 6. Then √ Y0 = Z0 , (32) dYs = −βh(Ys ) ds + 2 dB̂s , √ ϵ ϵ ϵ dYs = −βhϵ (Ys ) ds + 2 dB̂s , Y0 = Z0 , (33) with semigroups (Qs )s≥0 and (Qϵs )s≥0 , the latter admitting the unique invariant measure πβϵ ∝ e−βuϵ . Lemma 7. Let Assumptions A1-A4 hold and let Y, Y ϵ solve (32) and (33) respectively. Then 2

lim E |Ys − Ysϵ | = 0

ϵ→0

(34)

for every s ≥ 0,

and consequently W2 (νQs , νQϵs ) → 0. Proof. The proof is postponed to Appendix G Lemma 8. Under Assumptions A1-A3 one has W2 (πβϵ , πβ ) → 0 as ϵ → 0. Proof. The proof is postponed to Appendix G Proposition 5. Let Assumptions A1-A4 hold and let ρ(β) be as in Remark 10. Then πβ is the unique invariant measure of (2) and, for every t ≥ 0 and ν ∈ P2 (Rd ), W2 (νPt , πβ ) ≤ Cw e−Cr t W2 (ν, πβ ) ,

Cr =

2 , βρ(β)

(35)

  1 βLρ(β) where, with t⋆ = log 1 + , 2L 2 s

( Cw = max e

2 (L+ βρ(β) )t⋆

, e

2t⋆ βρ(β)

βLρ(β) 2 (1 − e−2Lt⋆ )

s 

) ≤

 βLρ(β) e 1+ . 2

(36)

Proof. The proof is postponed to Appendix G Remark 12. In the convex case, i.e. L = 0, one repeats the calculation so that (35) reads as √ W2 (νPt , πβ ) ≤ e e−2t/(βρ(β)) W2 (ν, πβ ) Appendix E. Discretization Error Estimates For the remainder of the analysis we fix the restart horizon of Definition 1 as   ln Cw + 1 ln Cw + 1 n0 := , so that λn0 ≤ + λ0 =: T0 , Cr λ Cr where Cw , Cr are the contraction constants of Proposition 5. 24

(37)

Lemma 9. Let Assumptions A1-A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for every λ ∈ (0, λ0 ), t ∈ [0, n0 ] and n ≥ n0 , with n0 as in (37), one has   λ,n 1/4 W2 L(ϑλ,n , (38) t ), L(ζ̄t ) ≤ CE λ 
1/2  √ 2 where CE = e5LT0 T0 2LCB4 / µ + L−1 C1 (1 + CB2 ) + 4 max{2m2 , K 2 }(1 + CB2 + CB3 )CB4 . Proof. The proof is postponed to Appendix G. Appendix F. Excess Risk Lemma 10. Let Assumptions A1-A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for any λ ∈ (0, λ0 ) and n ∈ N, the following bound for T1 = E[u(θnλ )] − E[u(θ∞ )], holds
T1 ≤ CT1 W2 πβ , L(θnλ ) , (39) where CT1 = m +

2p−1 K p CB2 + A(2p)p/2 . p+1

Proof. The proof is postponed to Appendix G. Lemma 11. Let Assumptions A1-A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for any λ ∈ (0, λ0 ), n ∈ N and β ≥ max{4/µ, 1/J}, the following bound for T2 = E[u(θ∞ )] − u(θ∗ ), holds   d 2(b + d/β)β 2 J 2 1 13 ∗ T2 ≤ E[u(θ∞ )] − u(θ ) ≤ log + log (πd) + =: CT2 , (40) 2β µ 2β 6β where J = m + K22p−2 (1 + (2b/µ)p/2 ) + K2p−1 /(p + 1). Proof. The proof is postponed to Appendix G. Appendix G. Postponed Proofs Proof of Remark 2 Proof. Let |x| ≥ R, then through A2 one obtains ⟨x, h(x)⟩ = ⟨x − 0, h(x) − h(0)⟩ + ⟨x, h(0)⟩ ≥ µ|x|2 − |x||h(0)| ≥

µ 2 |h(0)|2 |x| − . 2 2µ

(41)

Now let |x| < R, due to the polynomial growth in A3 one writes µ µ ⟨x, h(x)⟩ ≥ −|x||h(x)| ≥ −m|x| − K|x|p+1 ≥ −mR − KRp+1 + R2 − R2 2 2
µ 2 2 p+1 ≥ |x| − mR + (µ/2)R + KR . 2  Combining (41) and (42) yield (8), where b1 = max |h(0)|2 /(2µ), mR + (µ/2)R2 + KRp+1 . Furthermore, if one assumes A1, the dependency on R becomes quadratic. For |x| < R, by A1 for y = 0

(42)

⟨x, h(x)⟩ ≥ ⟨x − 0, h(x) − h(0)⟩ + ⟨x, h(0)⟩ ≥ −L|x|2 − |x||h(0)| ≥ (µ/2)|x|2 − (L + µ/2)R2 + mR ,  and therefore b2 = max m2 /(2µ), mR + (L + µ/2)R2 Proof of Remark 6

25

Proof. One calculates directly from (4) |hλ,u (x)| ≤ µ|x| + λ−1/2

λ1/2 |h(x) − µx| ≤ µ|x| + λ−1/2 . 1 + λ1/2 |h(x) − µx|

Furthermore, |hλ,c

v v u d u d uX (i) uX √ √ t 2 (x)| = |h (x)| ≤ t 2µ2 |x(i) |2 + 2λ−1 ≤ 2µ|x| + 2dλ−1/2 . λ,c

i=1

i=1

Proof of Remark 7 Proof. In view of A1, one calculates h(x) − µx λ1/2 |h(x) − µx| ≤ |h(x) − µx| − (h(x) − µx) 1 + λ1/2 |h(x) − µx| 1 + λ1/2 |h(x) − µx|

≤ λ1/2 |h(x) − µx|2 ≤ λ1/2 2|h(x)|2 + 2µ2 |x|2 ≤ λ1/2 4m2 + 2µ2 |x|2 + 4K 2 |x|2p

|hλ,u (x) − h(x)| ≤

≤ λ1/2 C1 (1 + |x|2p ),  where C1 = 2µ2 + 4 max m2 , K 2 , using |x|2 ≤ 1 + |x|2p for p ≥ 1. Furthermore,

|hλ,c

v v u d u d uX (i) uX 1/2 t t (i) 2 (x) − h(x)| = |h (x) − h (x)| ≤ λ |h(i) (x) − µx(i) |4 λ,c

i=1

≤ λ1/2

d X

i=1

|h(i) (x) − µx(i) |2 = λ1/2 |h(x) − µx|2 ≤ λ1/2 C1 (1 + |x|2p ).

i=1

Proof of Remark 8 Proof. One writes  ⟨hλ,u (x), x⟩ =

h(x) − µx , x + µ|x|2 . 1 + λ1/2 |h(x) − µx|

If ⟨h(x) − µx, x⟩ ≥ 0, then one immediately gets ⟨hλ,u (x), x⟩ ≥ µ|x|2 ≥ (µ/2)|x|2 − b. Otherwise, we calculate     h(x) − µx h(x) − µx ,x ≥ − , x ≥ − |⟨h(x) − µx, x⟩| 1 + λ1/2 |h(x) − µx| 1 + λ1/2 |h(x) − µx| = ⟨h(x) − µx, x⟩ = ⟨h(x), x⟩ − µ|x|2 ≥ −(µ/2)|x|2 − b, where the last inequality holds due to Remark 2. By combining the above bounds, one derives the result in (22). Furthermore, for the coordinate-wise case, we consider the following scenarios: (i)

• If h(i) (x)x(i) − µ|x(i) |2 ≥ 0 for every i ∈ {1, . . . , d}, then we trivially obtain hλ,c (x)x(i) ≥ µ|x(i) |2 ⇒ ⟨hλ,c (x), x⟩ ≥ µ|x|2 ≥ (µ/2)|x|2 − b.

26

• If h(i) (x)x(i) − µ|x(i) |2 ≤ 0 for every i ∈ {1, . . . , d}, then we have d X h(i) (x) − µx(i) h(i) (x) − µx(i) (i) x ≥ − x(i) 1/2 (i) (i) 1/2 |h(i) (x) − µx(i) | 1 + λ |h (x) − µx | 1 + λ i=1 i=1

d X

≥

d X

(h(i) (x) − µx(i) )x(i) ≥ ⟨h(x), x⟩ − µ|x|2

i=1

⇒ ⟨hλ,c (x), x⟩ ≥ (µ/2)|x|2 − b. • If h(i) (x)x(i) − µ|x(i) |2 ≥ 0 for every i ∈ J ⊂ {1, . . . , d}, then we have ⟨hλ,c (x), x⟩ = ⟨hλ,c (x), x⟩J + ⟨hλ,c (x), x⟩J c ≥ µ|x|2J + ⟨h(x), x⟩J c µ µ ≥ µ|x|2J + |x|2J c − b |J c | ≥ |x|2 − bd, 2 2 where the second inequality applies 10 coordinate-wise on J c .

Proof of Lemma 2 Proof. We first observe that E|θnλ |2 < ∞ for every n, by induction from Assumption A4 and the linear growth (20). Consider the SG-TULA iterates (θnλ )n≥0 from (3) p λ |θn+1 |2 = |θnλ − λhλ (θnλ )|2 + 2λβ −1 |ξn+1 |2 + 2 2λβ −1 ⟨θnλ − λhλ (θnλ ), ξn+1 ⟩ Since θnλ is independent of ξn+1 , the last term vanishes under expectation. Thus by taking conditional expectations on both sides and using Remark 6 and 8, we obtain λ E[|θn+1 |2 |θnλ ] = E[|θnλ |2 |θnλ ] − 2λE[⟨θnλ , hλ (θnλ )⟩|θnλ ] + λ2 E[|hλ (θnλ )|2 |θnλ ] + 2λdβ −1

≤ |θnλ |2 − λµ|θnλ |2 + 2λ2 µ2 |θnλ |2 + 2λb + 2λ + 2λdβ −1 (43)

≤ (1 − λµ + 2λ2 µ2 )|θnλ |2 + 2λ(b + 1 + dβ −1 ) Recalling the restriction λ < λ0 < 1/(4µ), iterating inequality (43) yields E|θnλ |2 ≤ (1 − λµ/2)n E|θ0λ |2 +

1 − (1 − λµ/2)n 2λ(b + 1 + dβ −1 ) λµ/2

≤ E|θ0 |2 + 4(b + 1 + dβ −1 )/µ.

Proof of Lemma 3 Proof. Consider the SG-TULA iterates (θnλ )n≥0 from (3) and define the following auxiliary processes p ∆λn := θnλ − λhλ (θnλ ), Gλn := 2λβ −1 ξn+1 . Furthermore, define Aλn := |∆λn |2 and Bnλ := 2⟨∆λn , Gλn ⟩ + |Gλn |2 . For the 2p-th moment, we write λ |θn+1 |2p = (Aλn + Bnλ )p , λ E[|θn+1 |2p | θnλ ] ≤ E[(Aλn )p | θnλ ] + E[2p(Aλn )p−1 Bnλ | θnλ ] + E

" p   X p k=2

27

k

# |Aλn |p−k |Bnλ |k | θnλ

.

(44)

Since Lemma 2 holds, using bound (43), we estimate the first term of (44)
p E[(Aλn )p | θnλ ] = (Aλn )p ≤ (1 − λµ/2)|θnλ |2 + 2λ(b + 1) ≤ (1 + λµ/4)

p−1

p

(1 − λµ/2) |θnλ |2p + 1 + 4λ−1 µ−1

p−1

(2λ(b + 1))p

≤ (1 − λµ/4)p−1 (1 − λµ/2)|θnλ |2p + λ(λ + 4/µ)p−1 (2(1 + b))p (45)

≤ rpλ |θnλ |2p + wpλ ,

where rpλ = (1 − λµ/4)p−1 (1 − λµ/2) and wpλ = λ(λ + 4/µ)p−1 (2(1 + b))p . Note that we applied the elementary inequality (a + b)p ≤ (1 + ϵ)p−1 ap + (1 + 1/ϵ)p−1 bp with ϵ = λµ/4. We similarly estimate the second term of (44) via bound (45) E[2p(Aλn )p−1 Bnλ | θnλ ] ≤ 2p(Aλn )p−1 E[Bnλ | θnλ ] ≤ 4pλdβ −1 (Aλn )p−1 (46)

λ λ ≤ 2p(2λdβ −1 )(rp−1 |θnλ |2(p−1) + wp−1 ).

The third term in inequality (44) expands as follows p   X p k=2

k

p−2  X

 p |Aλn |p−2−m |Bnλ |m+2 m+2 m=0   p−2 X p−2 p(p − 1) |Aλn |p−2−m |Bnλ |m |Bnλ |2 = (m + 2)(m + 1) m m=0  p−2  X p−2 ≤ p(p − 1)|Bnλ |2 |Aλn |p−2−m |Bnλ |m m m=0

|Aλn |p−k |Bnλ |k =

= p(p − 1)(|Aλn | + |Bnλ |)p−2 |Bnλ |2 ≤ p(p − 1)2p−3 |Aλn |p−2 |Bnλ |2 + p(p − 1)2p−3 |Bnλ |p := Dnλ + Fnλ .

(47)

To bound E[Dnλ | θnλ ], we first estimate the factor E[|Bnλ |2 | θnλ ] h i   2 E[|Bnλ |2 | θnλ ] ≤ E 2⟨∆λn , Gλn ⟩ + |Gλn |2 | θnλ ≤ E 8|∆λn |2 |Gλn |2 + 2|Gλn |4 | θnλ     ≤ 8|∆λn |2 E |Gλn |2 | θnλ + 2E |Gλn |4 | θnλ = 8(2λdβ −1 )Aλn + 2(2λdβ −1 )2 . In view of the above result and (45), we derive E[Dnλ | θnλ ] = p(p − 1)2p−3 |Aλn |p−2 E[|Bnλ |2 | θnλ ] ≤ p(p − 1)2p−2 4(2λdβ −1 )|Aλn |p−1 + (2λdβ −1 )2 |Aλn |p−2

λ λ ≤ p(p − 1)2p (2λdβ −1 )(rp−1 |θnλ |2(p−1) + wp−1 ) λ λ + p(p − 1)2p−2 (2λdβ −1 )2 (rp−2 |θnλ |2(p−2) + wp−2 ).

(48)

We now turn to estimating the remainder h i p E[Fnλ | θnλ ] ≤ p(p − 1)2p−3 E 2⟨∆λn , Gλn ⟩ + |Gλn |2 | θnλ ≤ p(p − 1)23p−4 |∆λn |p E[|Gλn |p | θnλ ] + p(p − 1)22p−4 E[|Gλn |2p | θnλ ] ≤ p(p − 1)23p−4 (2λdβ −1 )p/2 (Aλn )p/2 + p(p − 1)22p−4 (2λdβ −1 )p λ λ ≤ p(p − 1)23p−4 (2λdβ −1 )p/2 (rp/2 |θnλ |p + wp/2 ) + p(p − 1)22p−4 (2λdβ −1 )p .

28

(49)

Hence, combining the bounds (45)-(49) and substituting into (44), we obtain λ E[|θn+1 |2p | θnλ ] ≤ rpλ |θnλ |2p + wpλ λ λ + p(p + 1)2p (2λdβ −1 )(rp−1 |θnλ |2(p−1) + wp−1 ) λ λ + p(p − 1)2p−2 (2λdβ −1 )2 (rp−2 |θnλ |2(p−2) + wp−2 ) λ λ + p(p − 1)23p−4 (2λdβ −1 )p/2 (rp/2 |θnλ |p + wp/2 ) + p(p − 1)22p−4 (2λdβ −1 )p .

We establish the desired bound on the event {|θnλ |2 ≥ (4µ)(2dβ −1 )p(p + 1)2p+1 }. Conditioning on this event yields λ λ λ λ |θnλ |2p E[|θn+1 |2p | θnλ ] ≤ rpλ |θnλ |2p + (1/2)(λµ/4)rp−1 |θnλ |2p + (1/2)(λµ/4)2 rp−2 |θnλ |2p + (1/2)(λµ/4)p/2 rp/2 λ λ λ + (2λdβ −1 )p ) + wpλ + p(p − 1)23p−2 ((2λdβ −1 )wp−1 + (2λdβ −1 )2 wp−2 + (2λdβ −1 )p/2 wp/2  p/2 λ ≤ rp/2 (1 − λµ/4) + (1/2)(λµ/4)(1 − λµ/4)p/2−1 + (1/2)(λµ/4)2 (1 − λµ/4)p/2−2  +(1/2)(λµ/4)p/2 |θnλ |2p + λp(p − 1)24p−2 17p−1 (1/4µ)p−1 (1 + b + dβ −1 )p

≤ (1 − λµ/2)|θnλ |2p + λp(p − 1)22p−1 17p−1 (1/µ)p−1 (1 + b + dβ −1 )p ,

(50)

where we have used the restriction λ < 1/(4µ). Subsequently, for the complementary case we have λ E[|θn+1 |2p | θnλ ] ≤ (1 − λµ/2)|θnλ |2p + λ(p(p + 1)2p+1 )p 23p−3 (1/µ)p−1 (dβ −1 )p

+ λp(p − 1)23p−4 (p(p + 1)2p+1 )p−1 (1/µ)p−1 (dβ −1 )p + λp(p − 1)24p−2 (p(p + 1)2p+1 )p/2 (1/µ)p−1 (dβ −1 )p + λp(p − 1)22p−1 17p−1 (1/µ)p−1 (1 + b + dβ −1 )p .

(51)

We conclude by combining (50)-(51), taking expectation on both sides, and iterating the resulting bound as in the proof of Lemma 2 λ E|θn+1 |2p ≤ (1 − λµ/2)E|θnλ |2p + λ24p+2 p(p − 1)(p(p + 1)2p+1 )p/2 (1/µ)p−1 (1 + b + dβ −1 )p

≤ E|θ0λ |2p + 24p+3 p(p − 1)(p(p + 1)2p+1 )p/2 (1/µ)p (1 + b + dβ −1 )p .

Proof of Lemma 4 Proof. Consider the auxiliary process from Definition 1 and fix p ∈ N \ {1}. Using standard arguments involving stopping times, Grönwall’s lemma, Fatou’s lemma and the bound from Lemma 13, we obtain the existence of a constant cn0 , which depends on time, such that supt∈[0,n0 ] E[Vp (ζ̄tλ,n )] ≤ cn0 < ∞. Incorporating the integrating factor eλC3 (p)t and applying Itô’s formula to (t, x) 7→ eλC3 (p)t (1 + |x|2 )p/2 for t ∈ [0, n0 ] yields Z n0   eλC3 (p)n0 Vp (ζ̄tλ,n ) = Vp (ζ̄0λ,n ) + eλC3 (p)s λC3 (p)Vp (ζ̄sλ,n ) + λβ −1 ∆Vp (ζ̄sλ,n ) − λ⟨h(ζ̄sλ,n ), ∇Vp (ζ̄sλ,n )⟩ ds Z n0 p 0 + 2λβ −1 eλC3 (p)s ⟨∇Vp (ζ̄sλ,n ), dWsλ,n ⟩. (52) 0

Since |∇VpR(x)|2 ≤ p2 |Vp (x)|2 = p2 V2p (x) and eλC3 (p)s is bounded on [0, n0 ], the moment bound from Lemma n 13 gives E 0 0 |∇Vp (ζ̄sλ,n )|2 ds < ∞, so the stochastic integral in (52) is a true martingale and vanishes under expectation. Taking expectations and using the drift condition λC3 (p)Vp + λAVp ≤ λC4 (p) from Lemma 13, one

29

writes λC3 (p)n0

e

E[Vp (ζ̄tλ,n )] = E[Vp (ζ̄0λ,n )] + λ

Z n0

  eλC3 (p)s E C3 (p)Vp (ζ̄sλ,n ) + LVp (ζ̄sλ,n ) ds 0 Z n0 λ ≤ E[Vp (θn−n0 )] + λC4 (p) eλC3 (p)s ds 0

C4 (p) λ = E[Vp (θn−n )] + 0 C3 (p)

e

λC3 (p)n0

−1 .

Dividing through by eλC3 (p)n0 , one gets λ E[Vp (ζ̄tλ,n )] ≤ e−λC3 (p)n0 E[Vp (θn−n )] + 0 λ )] + ≤ E[Vp (θn−n 0

C4 (p) 1 − e−λC3 (p)n0 C3 (p)

C4 (p) . C3 (p)

In view of Lemma 3 and Minkowski’s inequality, one gets  p/2 C (p) 4 p/2 λ p 2/p E[Vp (ζ̄tλ,n )] ≤ 1 + E[|θn−n ≤ 2p/2−1 + 2p/2−1 CB2 + (1 + R0 ) | ] + 0 C3 (p)   p/2 ≤ 2p/2−1 + 2p/2−1 E|θ0 |p + M (p/2)(1/µ)p/2 (1 + b + dβ −1 )p/2 + (1 + R0 ) h   
p/2 i ≤ 2p/2−1 1 + E|θ0 |p + 2p/2 + 4p/2 + M (p/2) (1/µ)2p 1 + b + β −1 1 + b + β −1 (d + p − 2) .

Proof of Lemma 5 λ Proof. One considers the difference between (θ̄⌊t⌋ )t≥0 , and (θ̄tλ )t≥0 , to get the one-step error

λ |θ̄⌊t⌋ − θ̄tλ |2 ≤ 2

Z t ⌊t⌋

2 λ λhλ (θ̄⌊s⌋ ) ds

λ 2 + 4λβ −1 |B̃tλ − B̃⌊t⌋ | .

Since ⌊s⌋ = ⌊t⌋ for s ∈ [⌊t⌋, t], the drift integrand is constant on this interval, hence λ λ λ 2 |θ̄⌊t⌋ − θ̄tλ |2 ≤ 2λ2 (t − ⌊t⌋)2 |hλ (θ̄⌊t⌋ )|2 + 4λβ −1 |B̃tλ − B̃⌊t⌋ | .

By (20), |hλ (x)|2 ≤ 2µ2 |x|2 + 2λ−1 . Taking expectations and applying Lemma 2,
λ E|θ̄⌊t⌋ − θ̄tλ |2 ≤ 4λ2 µ2 CB1 + 4λ + 4λdβ −1 ≤ λ µCB1 + 4 + 4dβ −1 ,

(53)

λ,n λ where the last inequality uses 4λµ ≤ 1. The bound for (ϑλ,n t )t∈[0,n0 ] is identical since L(ϑk ) = L(θn−n0 +k ) for k ∈ {0, . . . , n0 }.

Proof of Lemma 9 Proof. Consider the difference between the continuous-time interpolation (19) and the auxiliary process from Definition 1. Applying Itô’s formula to x → |x|2 , then one writes for any t ∈ [0, n0 ] |ϑλ,n − ζ̄tλ,n |2 = −2λ t

Z t 0

λ,n ⟨ϑλ,n − ζ̄sλ,n , hλ (ϑλ,n s ⌊s⌋ ) − h(ζ̄s )⟩ds

Z t Z t λ,n λ,n λ,n λ,n λ,n = −2λ ⟨ϑλ,n − ζ̄ , h(ϑ ) − h( ζ̄ )⟩ds − 2λ ⟨ϑλ,n − ϑλ,n s s s ⌊s⌋ ⌊s⌋ ⌊s⌋ , h(ϑ⌊s⌋ ) − h(ζ̄s )⟩ds 0 0 Z t λ,n − 2λ ⟨ϑλ,n − ζ̄sλ,n , hλ (ϑλ,n (54) s ⌊s⌋ ) − h(ϑ⌊s⌋ )⟩ds := k1 (t) + k2 (t) + k3 (t). 0

30

The first term, k1 (t), is controlled through the one-sided Lipschitz Assumption A1, in particular Z t k1 (t) = −2λ 0

λ,n λ,n λ,n ⟨ϑλ,n ⌊s⌋ − ζ̄s , h(ϑ⌊s⌋ ) − h(ζ̄s )⟩ds ≤ 2Lλ

Z t ≤ 4Lλ

|ϑλ,n − ζ̄sλ,n |2 ds + 4Lλ s

Z t 0

0

Z t 0

λ,n 2 |ϑλ,n ⌊s⌋ − ζ̄s | ds

λ,n 2 |ϑλ,n ⌊s⌋ − ϑs | ds.

(55)

Subsequently, to control the second term, k2 (n0 ), one uses the Cauchy-Schwarz inequality and the polynomial growth Assumption A3, to obtain Z t k2 (t) = −2λ 0

λ,n λ,n ⟨ϑλ,n − ϑλ,n s ⌊s⌋ , h(ϑ⌊s⌋ ) − h(ζ̄s )⟩ds ≤ 2λ

Z t 0

λ,n λ,n |ϑλ,n − ϑλ,n s ⌊s⌋ ||h(ϑ⌊s⌋ ) − h(ζ̄s )|ds.

Under expectation, one writes Z t

λ,n 2 1/2 2 1/2 ds (E[|h(ϑλ,n (E[|ϑλ,n − ϑλ,n s ⌊s⌋ ) − h(ζ̄s )| ]) ⌊s⌋ | ]) 0 Z t √ 1/2 2 λ,n 2 1/2 ds (E[|h(ϑλ,n ≤ 2 2CB4 λ3/2 ⌊s⌋ )| + |h(ζ̄s )| ]) 0 Z n0 1/2 2p 2 λ,n 2p 1/2 ≤ 4CB4 λ3/2 (2m2 + K 2 E[|ϑλ,n ds ⌊s⌋ | ] + K E[|ζ̄s | ])

E[k2 (t)] ≤ 2λ

0

≤ 4(max {2m2 , K 2 }(1 + CB2 + CB3 )CB4 )1/2 λ3/2 n0 ,

(56)

where in the last step, we used the moment estimates from Lemmas 3-4, with 2m2 +K 2 (CB2 +CB3 ) ≤ max{2m2 , K 2 }(1+ CB2 + CB3 ). The last term, k3 (t), is controlled via the properties of the taming function, specifically Remark 7 Z t

λ,n k3 (t) = −2λ ⟨ϑλ,n − ζ̄sλ,n , hλ (ϑλ,n s ⌊s⌋ ) − h(ϑ⌊s⌋ )⟩ds 0 Z n0 λ,n 2 ≤ 2λ (L/2)|ϑλ,n − ζ̄sλ,n |2 + (2L)−1 |hλ (ϑλ,n s ⌊s⌋ ) − h(ϑ⌊s⌋ )| ds Z0 n0 Z n0 λ,n 2 3/2 −1 2p ≤ Lλ |ϑλ,n − ζ̄ | ds + λ L C (1 + |ϑλ,n 1 s s ⌊s⌋ | )ds, 0

(57)

0

where the ϵ−Young inequality was applied with ϵ = L. Taking the expectation of (54) and substituting from (55),(56) and (57), we derive, for every t ∈ [0, n0 ], E[|ϑλ,n − ζ̄tλ,n |2 ] ≤ 5Lλ t

Z t

E[|ϑλ,n − ζ̄sλ,n |2 ]ds + 4LCB4 λ2 n0 s

0

+ 4(max {2m2 , K 2 }(1 + CB2 + CB3 )CB4 )1/2 λ3/2 n0 + λ3/2 L−1 C1 (1 + CB2 )n0 . 1/2

Recalling from (37) that λn0 ≤ T0 , and using λ1/2 ≤ λ0 E[|ϑλ,n − ζ̄tλ,n |2 ] ≤ 5Lλ t

(58)

√ ≤ 1/(2 µ) for the first constant term, one deduces that

Z t

E[|ϑλ,n − ζ̄sλ,n |2 ]ds s 0   √ 1/2 + λ T0 2LCB4 / µ + L−1 C1 (1 + CB2 ) + 4(max {2m2 , K 2 }(1 + CB2 + CB3 )CB4 )1/2 . (59)

Since the second moments of the processes (ζ̄tλ,n )t≥0 and (ϑλ,n t )t≥0 can be controlled, the right-hand side of (59) is finite. Hence, applying Grönwall’s lemma on [0, n0 ] and using 5Lλn0 ≤ 5LT0 leads to   √ E[|ϑλ,n − ζ̄tλ,n |2 ] ≤ λ1/2 e5LT0 T0 2LCB4 / µ + L−1 C1 (1 + CB2 ) + 4(max {2m2 , K 2 }(1 + CB2 + CB3 )CB4 )1/2 . t

31

Proof of Theorem 1 Proof. Let L(θ0 ) = ν and denote by (Pt )t≥0 the Markov semigroup of the SDE (2) and by {P̄k }k∈N0 the discrete-time semigroup of the algorithm (θkλ )k∈N0 . Let n ≥ n0 with n0 as in (37). We triangulate the global error W2 (L(θnλ ), πβ ) = W2 (ν P̄n , πβ ) ≤ W2 (ν P̄n , ν P̄n−n0 Ptn0 ) + W2 (ν P̄n−n0 Ptn0 , π) ≤ W2 (ν P̄n−n0 P̄n0 , ν P̄n−n0 Pn0 λ ) + W2 (ν P̄n−n0 Pn0 λ , πβ ) λ,n λ,n = W2 (L(ϑλ,n n0 ), L(ζ̄n0 )) + W2 (L(ζ̄n0 ), πβ ),

(60)

where the identification of the two laws follows from Definition 1 and writing tn = nλ. The first term is the discretization error over the restart window; by Lemma 9 (at t = n0 ), (61)

λ,n 1/4 W2 (L(ϑλ,n . n0 ), L(ζ̄n0 )) ≤ CE λ

λ The second term is the contraction error. Applying Proposition 5 with initial laws L(θn−n ) and πβ , over the time 0 −Cr n0 λ −1 horizon n0 λ, and using Cw e ≤ e by the choice (37), λ W2 (L(ζ̄nλ,n ), πβ ) ≤ e−1 W2 (L(θn−n ), πβ ), 0 0

∀n ≥ n0 .

(62)

Putting (61) and (62) together yields, for any n ≥ n0 , λ W2 (L(θnλ ), πβ ) ≤ CE λ1/4 + e−1 W2 (L(θn−n ), πβ ). 0

The recursion lemma above (with ϵ = CE λ1/4 , q = e−1 ) then gives, for every n ∈ N0 , W2 (L(θnλ ), πβ ) ≤

CE λ1/4 + e−(n/n0 −1) sup W2 (L(θkλ ), πβ ). 1 − e−1 k∈{0,...,n0 −1}

The supremum is bounded by CD due to Lemma 15. Moreover, by (37), n Cr nλ nλ n ≥ ≥ = . ln Cw + 1 n0 ln Cw + 1 + Cr λ0 T0 +1 Cr λ Since 1/(1 − e−1 ) ≤ 2, one concludes that for any n ∈ N0 , W2 (L(θnλ ), πβ ) ≤ 2CE λ1/4 + eCD e−nλ/T0 .

Proof of Lemma 10 Proof. We notice that the function g(t) = u(tx + (1 − t)y) is locally Lipschitz continuous (since u is semi-convex) so it has a bounded variation in [0, 1]. Then, one can enforce the fundamental theorem of calculus since g ′ (t) = ⟨h(tx + (1 − t)y), x − y⟩ a.e. Let k(p) = 2p−1 K/(p + 1), then one writes Z 1 u(x) − u(y) =

Z 1 ⟨x − y, h((1 − t)y + tx)⟩dt ≤

0

Z 1 ≤

|x − y||h((1 − t)y + tx)|dt 0

|x − y|(m + K|(1 − t)y + tx|p )dt ≤ (m + k(p)|x|p + k(p)|y|p )|x − y|,

(63)

0

where we have used Cauchy-Schwarz and the growth Assumption A3. Now let P be the coupling of µ, ν that achieves W2 (µ, ν). That is, P = L((X, Y )) with µ = L(X), ν = L(Y ) such that W22 = EP |X − Y |2 . Taking expectations

32

in (63), yields Z

Z

gdν = EP [g(X) − g(Y )] q q ≤ EP [(m + k(p)|X|p + k(p)|Y |p )2 ] EP |X − Y |2   q q P 2p P 2p ≤ m + k(p) E |X| + k(p) E |Y | W2 (µ, ν).

gdµ − Rd

Rd

(64)

where the last step follows from Minkowski’s inequality. One concludes by applying inequality (64) for X = θnλ and Y = θ∞   p p E[u(θnλ )] − E[u(θ∞ )] ≤ m + k(p) CB2 + k(p) Cσ,2p W2 (L(θnλ ), πβ ), (65) where Cσ,2p isR the 2pth-moment of πβ and E|θnλ |2p ≤ CB2 by Lemma 3.. Since πβ is the invariant measure of SDE 2, there holds Rd AV2p (x)πβ (dx) = 0. Due to Lemma 13, one estimates the constant by Z Z −1 Cσ,2p ≤ V2p (x)πβ (dx) ≤ −C3 (2p) AV2p (x)πβ (dx) + C4 (2p)/C3 (2p) ≤ C4 (2p)/C3 (2p). Rd

Rd

In view of Lemma 13, Cσ,2p ≤ C4 (2p)/C3 (2p) = Ap (2p), and the stated constant CT1 follows. Proof of Lemma 11 Proof. We follow a similar approach as in Section 3.5 of Raginsky et al. [2017], making necessary adjustments due to the lack of a smoothness condition for the gradient ∇u(x) := h(x) and polynomial growth. According to Raginsky et al. [2017], one obtains the following bound   d 4πe(b + d/β) 1 E[u(θ∞ )] ≤ log − log Z, (66) 2β µd β R where Z := Rd e−βu(x) dx is the normalization constant. One writes Z Z ∗ log Z = log e−βu(x) dx = −βu(θ∗ ) + log eβ(u(θ )−u(x)) dx. (67) Rd

Rd

Now we provide an upper bound for the second term of (67). For the remainder of this analysis, one chooses the version of the subgradient h(x) such that h(θ∗ ) = 0. Evaluating Remark 2 at θ∗ then gives µ2 |θ∗ |2 ≤ b, hence p |θ∗ | ≤ 2b/µ =: R2 . Consequently, one calculates that ∗

∗

Z 1

−(u(θ ) − u(x)) ≤ |u(θ ) − u(x)| ≤

|⟨h(x + t(θ∗ − x)), θ∗ − x⟩| dt

0

Z 1 ≤

|h(x + t(θ∗ − x))||θ∗ − x|dt ≤

0

Z 1 ≤

Z 1

(m + K|x + t(θ∗ − x)|p ) |θ∗ − x|dt

0

m + 2p−1 K|x|p + 2p−1 Ktp |θ∗ − x|p |θ∗ − x|dt

0

≤ (K22p−2 + K2p−1 /(p + 1))|θ∗ − x|p+1 + (m + K22p−2 R2p )|θ∗ − x|. Hence we obtain Z Z ∗ 2p−2 +K2p−1 /(p+1))|θ ∗ −x|p+1 −β(m+K22p−2 R2p )|θ ∗ −x| I= eβ(u(θ )−u(x)) dx ≥ e−β(K2 dx d Rd ZR ∗ p+1 ∗ ≥ e−βJ (|θ −x| +|θ −x|) dx, Rd

33

(68)

where J = m+K22p−2 (1+(2b/µ)p/2 )+K2p−1 /(p+1) and the last inequality uses that both coefficients K22p−2 + K2p−1 /(p + 1) and m + K22p−2 R2p are bounded above by J (recall R2p = (2b/µ)p/2 ). Passing to radial coordinates, with Sd = 2π d/2 Γ−1 (d/2) the surface area of the unit sphere in Rd , one obtains Z ∞ p+1 I ≥ Sd e−βJ (r +r) rd−1 dr. 0

Assume βJ ≥ 1. Restricting the domain of integration to r ∈ [0, (βJ)−1 ], where r ≤ 1 and hence rp+1 ≤ r, gives βJ(rp+1 + r) ≤ 2βJr ≤ 2, so that Z (βJ)−1 I ≥ Sd

−βJ (r p+1 +r ) d−1

e

r

dr ≥ Sd e

−2

Z (βJ)−1

rd−1 dr =

0

0

Sd −2 e (βJ)−d . d

Combining this with (67) yields 1 1 d 2 1 log Z = −u(θ∗ ) + log I ≥ −u(θ∗ ) + log(Sd /d) − log(βJ) − . β β β β β In view of (66), one concludes with d E[u(θ∞ )] − u(θ ) ≤ log 2β ∗



4πe(b + d/β)β 2 J 2 µd



1 2 log(Sd /d) + β β     d 4πe(b + d/β)β 2 J 2 1 2π d/2 2 ≤ log − log + 2β µd β Γ(d/2)d β   2 2 4πe(b + d/β)β J 1 2 d d log log (π) + log (Γ(d/2)(d/2)) + ≤ − 2β µd 2β β β   d 4e(b + d/β)β 2 J 2 1 2 ≤ log + log (Γ(d/2 + 1)) + 2β µd β β   2(b + d/β)β 2 J 2 d 1 13 ≤ log + log (πd) + := CT2 . 2β µ 2β 6β −

Proof of Lemma 1 Proof. Write g(x) = u(x) + (L/2)|x|2 which is convex by A1. Let Lm denote m-dimensional Lebesgue measure, B(z, r) := {w : |zR− w| < r}, Ωρ := {z : |z| > ρ}. Moreover, consider theR mollifier, η ∈ Cc∞ (Rd ) with η ≥ 0, supp η ⊂ B(0, 1), η dLd = 1, and set ηϵ (z) := ϵ−d η(z/ϵ), so that ηϵ ≥ 0, ηϵ dLd = 1, and supp ηϵ ⊂ B(0, ϵ). Henceforth we denote the convolution with ηϵ by g ϵ := ηϵ ∗ g. Let Φ be the set of differentiability points of g. We fix v := x − y,

γ(t) := y + tv (t ∈ [0, 1]),

so that γ(0) = y and γ(1) = x. Since g is convex, ∇g ∈ BVloc (Rd ; Rd ) and its distributional Hessian is a symmetric matrix-valued Radon measure [D2 g] whose directional component µξ := ξ ⊤ [D2 g] ξ is, for each ξ ∈ Rd , a nonnegative Radon measure [Evans and Gariepy, 2015, Theorem. 6.8]. By Aleksandrov’s theorem, the Ld -absolutely continuous density of [D2 g] is the pointwise Hessian [Evans and Gariepy, 2015, Theorem. 6.9], so the density of µξ is ξ ⊤ ∇2 g ξ = ξ ⊤ ∇2 u ξ + L|ξ|2 ≥ (µ0 + L)|ξ|2 a.e. on Ωρ . As µξ ≥ 0, its singular part is nonnegative, hence µξ dominates its absolutely continuous part. Restricting to Ωρ , yields, µξ ≥ (µ0 + L) |ξ|2 Ld . (1) The convolution g ϵ ∈ C ∞ (Rd ) is convex, so ∇2 g ϵ ⪰ 0 on Rd . Moreover ∇2 g ϵ = ηϵ ∗ [D2 g]. Hence for every ξ ∈ Rd

34

and every z such that |z| > ρ + ϵ, we have B(z, ϵ) ⊂ Ωρ . Then, estimate (1) gives Z Z ξ ⊤ ∇2 g ϵ (z) ξ = ηϵ (z − w) dµξ (w) = ηϵ (z − w) dµξ (w) d R B(z,ϵ) Z 2 ≥ (µ0 + L)|ξ| ηϵ (z − w) dLd (w) = (µ0 + L) |ξ|2 .

(2)

B(z,ϵ)

As g ϵ ∈ C ∞ , the fundamental theorem of calculus applied to t 7→ ⟨∇g ϵ (γ(t)), v⟩, that is ⟨∇g ϵ (x) − ∇g ϵ (y), v⟩ =

Z 1

v ⊤ ∇2 g ϵ (γ(t)) v dt.

0

Let Bϵ := {t ∈ [0, 1] : |γ(t)| ≤ ρ + ϵ}. Notice that Bϵ is a sublevel set of the convex map t 7→ |γ(t)|, and thus an interval. Writing Bϵ = [t1 , t2 ], the identity |γ(s) − γ(t)| = |s − t| |v| gives L1 (Bϵ ) = t2 − t1 =

2(ρ + ϵ) |γ(t2 ) − γ(t1 )| ≤ . |v| |v|

Now for every t ∈ [0, 1] \ Bϵ one has |γ(t)| > ρ + ϵ, so (2) applies pointwise and v ⊤ ∇2 g ϵ (γ(t)) v ≥ (µ0 + L)|v|2 . On Bϵ the integrand is positive (due to convexity of g), therefore one writes  
⟨∇g ϵ (x) − ∇g ϵ (y), v⟩ ≥ (µ0 + L) |v|2 L1 [0, 1] \ Bϵ ≥ (µ0 + L) |v|2 1 − 2(ρ+ϵ) . (3) |v| By [Rockafellar, 1970, Theorem. 25.5], Φ is co-null and ∇g is continuous on Φ. In particular, for x ∈ Φ, ∇g(xk ) → ∇g(x) as Φ ∋ xk → x. Let x ∈ Φ be an arbitrary element. Splitting the integral between Φ and its co-null yields: Z Z |∇g ϵ (x) − ∇g(x)| ≤ ηϵ (z) |∇g(x − z) − ∇g(x)| dz + ηϵ (z) |∇g(x − z) − ∇g(x)| dz. {x−z∈Φ}

{x−z ∈Φ} /

The second integral vanishes identically since {z : x − / Φ} is Ld -null and ∇g is locally bounded, while the first R z∈ term is bounded by supw∈Φ∩B(x,ϵ) |∇g(w) − ∇g(x)| ηϵ = supw∈Φ∩B(x,ϵ) |∇g(w) − ∇g(x)| → 0 by continuity of ∇g at x. Hence ∇g ϵ (x) → ∇g(x). Passing to the limit in (3) gives ⟨∇g(x) − ∇g(y), v⟩ ≥ (µ0 + L)|v|2 (1 − 2ρ/|v|). Since ∇g(x) = h(x) + Lx on Φ, the left side is ⟨h(x) − h(y), v⟩ + L|v|2 , so ⟨h(x) − h(y), v⟩ ≥ µ0 |v|2 − 2ρ(µ0 + L) |v|.

(4)

If |v| = |x − y| ≥ 4ρ(µ0 + L)/µ0 , then 2ρ(µ0 + L) ≤ µ20 |v|, hence 2ρ(µ0 + L)|v| ≤ µ20 |v|2 , and (4) gives ⟨h(x) − h(y), x − y⟩ ≥ µ0 |v|2 − µ20 |v|2 = µ20 |x − y|2 .

Proof of Lemma 6 Proof. Both convolutions in (30) are well defined. In particular u is locally Lipschitz by A1 and h is locally bounded by A3, so both belong to L1loc (Rd ). Moreover hϵ is independent of the choice of selection h ∈ ∂u, as any two selections 1,1 agree Lebesgue-almost everywhere. Finally u ∈ Wloc (Rd ) identifies ∇uϵ = ∇ (u ∗ φϵ ) = (∇u) ∗ φϵ = h ∗ φϵ = hϵ , which is (i). R For (ii), let x, y ∈ Rd with |x − y| ≥ R. Trivially |(x − w) − (y − w)| ≥ R, so since φϵ ≥ 0 with φϵ = 1,

35

applying (6) yields Z ⟨hϵ (x) − hϵ (y), x − y⟩ =

⟨h(x − w) − h(y − w), x − y⟩ φϵ (w) dw ≥ µ|x − y|2 .

Rd

Replacing (6) by (5), which carries no restriction on |x − y|, the same computation gives (5). For (iii), A3 together with |x − w| ≤ |x| + ϵ on supp φϵ gives Z p (m + K|x − w|p ) φϵ (w) dw ≤ m + K (|x| + ϵ) , |hϵ (x)| ≤ B(0,ϵ)

the second bound following from ϵ ≤ 1. Since uϵ therefore satisfies A1-A3 by (i)-(iii), part (iv) follows from the proof of Lemma 17, in which the growth constants (m, K, p) enter only through supB(0,R) |h|, here at most m + K(R + ϵ)p . Finally (v) is the standard convergence of mollifications since u is continuous and h ∈ L1loc (Rd ) [Evans and Gariepy, 2015, §4.2.1, Theorem 4.1]. Proof of Lemma 7 Proof. Writing the two equations in integral form, Z s √ Ys = θ0 − β h(Yu ) du + 2B̂s ,

Ysϵ = θ0 − β

0

Z s

hϵ (Yuϵ ) du +

√

2B̂s ,

0

Rs the noise cancels in the difference, so that with ∆s := Ys − Ysϵ we have ∆s = 0 −β (h(Yr ) − hϵ (Yrϵ )) dr, a 2 continuous process of finite variation. Itô’s formula applied to x 7→ |x| thus has no second-order term and, since ∆0 = 0, gives Z s 2 |∆s | = −2β ⟨∆r , h(Yr ) − hϵ (Yrϵ )⟩ dr. 0

Adding and subtracting hϵ (Yr ) splits the right-hand side as Z s Z s −2β ⟨∆r , hϵ (Yr ) − hϵ (Yrϵ )⟩ dr + −2β ⟨∆r , h(Yr ) − hϵ (Yr )⟩ dr . 0 0 | {z } {z } | =:T1

=:T2

By A1 for hϵ , which holds with the same constant L by Lemma 6(ii), we get Z s 2 T1 ≤ 2βL |∆r | dr, 0

while by the Cauchy-Schwarz and Young inequalities Z s Z s Z s 2 2 T2 ≤ 2β |∆r | |h(Yr ) − hϵ (Yr )| dr ≤ β |∆r | dr + β |h(Yr ) − hϵ (Yr )| dr. 0

0

0

Combining the two bounds and taking expectations, Z s h Z s h i i 2 2 2 E |∆s | ≤ β (2L + 1) E |∆r | dr + β E |h(Yr ) − hϵ (Yr )| dr. 0

0 2

Both terms are finite, as Grönwall’s inequality requires: |x|2 ≤ V2p (x) gives |∆r | ≤ 2 (V2p (Yr ) + V2p (Yrϵ )), and
2 2 A3 with Lemma 6(iii) gives |h(x)| + |hϵ (x)| ≤ C 1 + |x|2p ≤ 2C V2p (x) for every ϵ ∈ (0, 1]. Since hϵ inherits A2 by Lemma 6(ii), Lemma 13 at order 2p applies to Y and Y ϵ with ϵ-independent constants and yields, by A4, sup sup E [V2p (Yr ) + V2p (Yrϵ )] < ∞.

ϵ∈(0,1] r≥0

36

(69)

Grönwall’s inequality therefore gives Z s h i 2 2 E |h(Yr ) − hϵ (Yr )| dr. E |∆s | ≤ βeβ(2L+1)s

(70)

0 2

It remains to show that the integral in (70) vanishes as ϵ → 0. Set fϵ := |hϵ − h| . By the above,
|fϵ (x)| ≤ C 1 + |x|2p =: g(x), E [g(Yr )] < ∞, so fϵ (Yr ) is dominated by an integrable function, uniformly in ϵ. Moreover, by Lemma 6(v), fϵ → 0 pointwise on D, whose complement is Lebesgue-null. Since the drift only reweighs the law of the driving Brownian path, one has L(Yr ) ≪ Leb, so Yr almost surely avoids this exceptional set and hϵ (Yr ) → h(Yr ) P-a.s.. Taking any sequence ϵk → 0 and writing fk := fϵk , the dominated convergence theorem yields h i 2 lim E |hϵk (Yr ) − h(Yr )| = 0, k→∞

and since the limit is the same along every such sequence, the convergence holds as ϵ → 0. Hence, by (70), 2 E |Ys − Ysϵ | → 0 for every s ≥ 0. As (Ys , Ysϵ ) is a coupling of νQs and νQϵs , 2

2

W2 (νQs , νQϵs ) ≤ E |Ys − Ysϵ | −→ 0, as ϵ → 0.

Proof of Lemma 8 Proof. Since A1-A3 hold for hϵ with constants independent of ϵ by Lemma 6(ii)-(iii), one shows as in Remark 2 that b > 0, independent of ϵ, with uϵ (x) ≥ µ |x|2 −c for all |x| ≥ R. b dissipativity holds uniformly in ϵ, whence there are c, R 4 Therefore 2

e−βuϵ (x) ≤ eβc e−βµ|x| /4 ∈ L1 (Rd ),

(71)

∀ϵ ∈ (0, 1].

Hence, by the dominated convergence theorem, Zϵ = e−βuϵ → e−βu = Z, and therefore πβϵ → πβ pointwise for every x ∈ Rd . Similarly one uses dominated convergence, bounding via |x|2 times the right-hand side of (71), for the R 2 −βuϵ (x) dx, to conclude that πβϵ → πβ in P2 (Rd ). Hence [Villani, 2009, Theorem 6.9] gives second moment Rd |x| e ϵ W2 (πβ , πβ ) → 0 as ϵ → 0. R For µ ∈ P(Rd ) with µ = f πβϵ we write Hϵ (µ) := f log f dπβϵ , and Hϵ (µ) := +∞ otherwise. R

R

Proof of Proposition 5 Proof. By A1-A2 for uϵ ∈ C 2 , condition (1.8) of Wang [2020] holds with I = 0, since we are in the flat Euclidean case. We argue at the level of the regularized semigroup Qϵs and let ϵ → 0 at the end. Four key estimates are needed, the first three govern the behavior for large s, the fourth for small s. (A) Talagrand inequality. By [Shao, 2006, Theorem 1] the logarithmic Sobolev inequality (27) with constant ρϵ implies
W22 ν, πβϵ ≤ ρϵ Hϵ (ν), ∀ν ∈ P2 (Rd ). (72) (B) Exponential decay of entropy. By [Wang, 2020, Proposition 3.3(2)], 4s

Hϵ (νQϵs ) ≤ e− ρϵ Hϵ (ν),

(C) Log-Harnack inequality and entropy-cost estimate. Let φ be bounded with inf φ > 0 and 2010, Theorem 1.1], Qϵs (log φ) (x) ≤ log Qϵs φ(y) +

βL|x − y|2 , 2 (1 − e−2βLs )

37

(73)

s ≥ 0. R

φ dπβϵ = 1. By [Wang,

s > 0, x, y ∈ Rd .

(74)

Let X ∼ ν and Y ∼ πβϵ be such that E|X − Y |2 = W22 (ν, πβϵ ). Evaluating (74) at (X, Y ) and taking expectations, Z

log φ d (νQϵs ) = E [Qϵs (log φ) (X)] ≤ E [log Qϵs φ(Y )] +

βL W 2 ν, πβϵ , 2 (1 − e−2βLs ) 2

the first equality being the definition of νQϵs . The remaining expectation vanishes, since by the concavity of the logarithm and the invariance of πβϵ , we write E [log Qϵs φ(Y )] =

Z

log Qϵs φ dπβϵ ≤ log

Z

Qϵs φ dπβϵ = log

Z

φ dπβϵ = 0.

Let f be the density of νQϵs with respect to πβϵ , which is admissible after truncation and satisfies R φ = f turns the left-hand side into f log f dπβϵ = Hϵ (νQϵs ), whence Hϵ (νQϵs ) ≤

βL W22 ν, πβϵ , −2βLs 2 (1 − e )

R

f dπβϵ = 1. Taking

(75)

s > 0.

(D) Short-time estimate. Synchronous coupling together with Grönwall’s lemma, as in the proof of Lemma 7, gives

W2 νQϵs , πβϵ ≤ eβLs W2 ν, πβϵ , s ≥ 0. (76) Now we assemble the above estimates. Fix a split, at time s0 > 0. For s ≥ s0 , chaining (72), (73) and (75), 4(s−s0 )

W22 νQϵs , πβϵ ≤ ρϵ Hϵ (νQϵs ) ≤ ρϵ e− ρϵ Hϵ νQϵs0 4(s−s0 )
βL 2 ϵ ≤ ρϵ e− ρϵ W ν, π , 2 β −2βLs 0) 2 (1 − e

that is,
ϵ

W2 νQϵs , πβ ≤

e

2s0 ρϵ

s

ρϵ βL 2 (1 − e−2βLs0 )

!
2s e− ρϵ W2 ν, πβϵ .

For s ≤ s0 , (76) gives instead


2s 2s 2 2 W2 νQϵs , πβϵ ≤ e(βL+ ρϵ )s e− ρϵ W2 ν, πβϵ ≤ e(βL+ ρϵ )s0 e− ρϵ W2 ν, πβϵ . Both regimes carry the same exponential rate, so for every s ≥ 0 s ( 2s0

2s 2 ϵ − ρϵ ϵ W2 νQϵs , πβϵ ≤ Cw e W2 ν, πβϵ , Cw = max e(βL+ ρϵ )s0 , e ρϵ

ρϵ βL 2 (1 − e−2βLs0 )

) .

(77)

Writing s ψ(s) := e

2s ρϵ

ρϵ βL , 2 (1 − e−2βLs )

ψ ′ (s) = ψ(s)



2 βLe−2βLs − ρϵ 1 − e−2βLs

the derivative vanishes precisely when e2βLs = 1 + ρϵ βL/2, that is at   1 ρϵ βL s⋆ = log 1 + . 2βL 2 There 1 − e−2βLs⋆ = ρϵ βL/ (2 + ρϵ βL), so the second expression in (77) equals s   12   ρ 1βL ϵ 2s⋆ ρ βL ρ βL ρ βL ϵ ϵ ϵ e ρϵ = 1+ 1+ , −2βLs ⋆ 2 (1 − e ) 2 2

38

 ,

(78)

Notice that the first expression in (78) agrees with the second for s = s⋆ . Since (1 + r) finally obtains s   ρϵ βL ϵ Cw ≤ e 1 + . 2

1/r

≤ e for every r > 0, one

(79)

The laws converge in W2 , namely νQϵs → νQs by Lemma 7 and πβϵ → πβ by Lemma 8, so by [Villani, 2009, Corollary 6.11] both sides of (77) pass to the limit. Additionally all constants involved depend continuously on ρϵ . Since Pt = Qt/β , substituting s = t/β turns the rate 2/ρ(β) into Cr and t⋆ = βs⋆ , giving (35) with Cw as in (36). Uniqueness follows trivially by contradiction. Remark 13. The split (78) is optimal. The corresponding argument of [Wang, 2020, Theorem 2.1(3)] takes t0 = 1, which suffices to produce a finite prefactor. Appendix H. Supporting Lemmas Lemma 12. Let Assumptions A2-A3 hold. For any p ≥ 1, and x ∈ Rd , we have AVp (x) ≤ C2 Vp (x), where C2 = p(b + β −1 (d + (p − 2)+ )). Proof. Let p ≥ 2, by direct calculation AVp (x) = −p(1 + |x|2 )p/2−1 ⟨x, h(x)⟩ + β −1 p(p − 2)|x|2 (1 + |x|2 )p/2−2 + dβ −1 p(1 + |x|2 )p/2−1 ≤ −p(1 + |x|2 )p/2−1 (µ/2|x|2 − b)) + β −1 p(d + (p − 2))(1 + |x|2 )p/2 ≤ bp(1 + |x|2 )p/2−1 − (µ/2)p(1 + |x|2 )p/2−1 |x|2 + β −1 p(d + (p − 2))Vp (x) ≤ bpVp (x) + β −1 p(d + (p − 2))Vp (x) ≤ C2 Vp (x).

Lemma 13. Let Assumptions A2-A3 hold. For any p ≥ 1, and x ∈ Rd , we have AVp (x) ≤ −C3 (p)Vp (x) + C4 (p), where C4 (p)/C3 (p) ≤ Amax{1,p/2} , with A(p) = 1 + 2(b + β −1 (d + p − 2))/µ. Proof. Let x ∈ Rd and p ≥ 2, one calculates AVp (x) = −p(1 + |x|2 )p/2−1 ⟨x, h(x)⟩ + β −1 p(p − 2)|x|2 (1 + |x|2 )p/2−2 + dβ −1 p(1 + |x|2 )p/2−1 ≤ −(µ/2)p|x|2 (1 + |x|2 )p/2−1 + (b + dβ −1 )p(1 + |x|2 )p/2−1 + β −1 p(p − 2)(1 + |x|2 )p/2−1 ≤ −(µ/2)p|x|2 (1 + |x|2 )p/2−1 + (b + dβ −1 + β −1 (p − 2))p(1 + |x|2 )p/2−1 ≤ −(µ/2)p(1 + |x|2 )(1 + |x|2 )p/2−1 + (µ/2 + b + β −1 (d + p − 2))p(1 + |x|2 )p/2−1   ≤ (µ/2)p A(1 + |x|2 )p/2−1 − (1 + |x|2 )p/2 .

(80)

1 w2 The weighted arithmetic-geometric mean inequality implies aw 1 a2 ≤ w1 a1 + w2 a2 , for weights living on a simplex. Using it with exponents w1 = (p − 2)/p and w2 = 2/p, yields

A(1 + |x|2 )p/2−1 = Ap/2×2/p (1 + |x|2 )p/2×(p−2)/p ≤

2 p/2 p − 2 A + (1 + |x|2 )p/2 p p

In view of (80) one gets AVp (x) ≤ −µVp (x) + µAp/2 39

When p ∈ [1, 2), one immediately gets for any x ∈ Rd AVp (x) ≤ −(pµ/2)Vp (x) + (pµ/2)(1 + 2(b + dβ −1 )/µ).

Lemma 14. Let n0 ∈ N, ϵ > 0, and q ∈ (0, 1). If a non-negative sequence {an }n∈N satisfies an ≤ ϵ + qan−n0 , ∀n ≥ n0 , then an ≤

ϵ + max {ai }q n/n0 −1 , ∀n ∈ N0 . 1 − q i∈{0,...,n0 −1}

Proof. Write n = kn0 + i with k = ⌊n/n0 ⌋ and i ∈ {0, . . . , n0 − 1}. Iterating the hypothesis k times gives an ≤ ϵ(1 + q + · · · + q k−1 ) + q k ai ≤ ϵ/(1 − q) + q k maxi∈{0,...,n0 −1} ai , and q k ≤ q n/n0 −1 since k ≥ n/n0 − 1. Lemma 15. Let Assumptions A2-A4 hold and λ0 ∈ (0, 1/(4µ)). Then, for every λ ∈ (0, λ0 ) and k ∈ N0 ,
1/2 W2 L(θkλ ), πβ ≤ (2CB1 + 2A(2)) . Proof. Let θ∞ ∼ πβ be independent of θkλ . The independent coupling and Lemma 2 give Z Z     W22 (L(θkλ ), πβ ) ≤ E |θkλ − θ∞ |2 ≤ 2E |θkλ |2 + 2 |x|2 πβ (dx) ≤ 2CB1 + 2 Rd

V2 (x) πβ (dx).

Rd

R Since πβ is the invariant measure of the SDE (2), one has Rd AV2 dπβ = 0 (see the proof of Lemma 10), so Lemma 13 at p = 2 yields Z C4 (2) = A(2), V2 dπβ ≤ C3 (2) Rd which concludes the proof. Lemma 16. The strongly convex proxy. Let πβ ∝ e−βu and W := βu. Suppose u satisfies A1, A2 and A3. Then there exists W0 : Rd → R, βµ-strongly βµ 2 2 convex on Rd , such that W0 = W on B c (0, R). Writing W̃ := W − βµ 2 | · | and W0 =: V + 2 | · | , the convex part satisfies inf W̃ (y) ≤ V (x) ≤ sup W̃ (y), ∀x ∈ B(0, R), |y|=R

|y|=R

and ψ := W − W0 is supported in B(0, R). 2 c Proof. As W = βu is βµ-strongly convex on B c (0, R), the function W̃ = W − βµ 2 | · | is convex on B (0, R). Let d V be the convex hull of W̃ |Bc (0,R) , the largest convex minorant on R , nP o P P c V (x) = inf λ W̃ (x ) : λ ≥ 0, λ = 1, λ x = x, x ∈ B (0, R) i i i i i i i i i i

(see Proposition 2.13 in Rockafellar and Wets [1998]), finite by the coercivity implied by Remark 2. Then V is convex on Rd and V = W̃ on B c (0, R). Moreover, Lemma B.2 in Ma et al. [2019], sharpens the bound trivial bounds from βµ 2 2 |y| ≥ R to the sphere |y| = R. Define W0 := V + βµ 2 | · | . Since V = W0 − 2 | · | is convex by construction, W0 βµ d 2 c is βµ-strongly convex on R and we have W0 = W̃ + 2 | · | = W on B (0, R). Hence ψ = W − W0 = W̃ − V vanishes on B c (0, R). Lemma 17. The oscillation bound. Under the hypotheses of Lemma 16, the perturbation ψ = W − W0 satisfies h i osc(ψ) := sup ψ − inf ψ ≤ β 4R (m + KRp ) + µ2 R2 . Rd

Rd

40

Consequently πβ ∝ e−βU satisfies the log-Sobolev inequality (27) with ρ(β) ≤

   2 exp β 4R (m + KRp ) + µ2 R2 . βµ

2 Proof. Fix x, y ∈ B(0, R) and γt = y + t(x − y) ∈ B(0, R). By A1 the function Φ := W + βL 2 | · | is convex, so the convex fundamental theorem of calculus so the convex fundamental theorem of calculus [Hiriart-Urruty and R1 Lemaréchal, 2001, Theorem D.2.3.4] gives Φ(x) − Φ(y) = 0 ⟨gt , x − y⟩ dt with gt ∈ ∂Φ(γt ) = β ∂u(γt ) + βLγt . R1 2 2 Choosing gt = βht + βLγt , ht ∈ ∂u(γt ), and using 0 ⟨βLγt , x − y⟩ dt = βL 2 (|x| − |y| ), the semi-convexity quadratic cancels, hence Z 1 W (x) − W (y) = β ⟨ht , x − y⟩ dt. 0

Since |ht | ≤ m + K|γt | ≤ m + KR by A3 and |x − y| ≤ 2R, p

p

OSCR (W ) :=

|W (x) − W (y)| ≤ 2Rβ (m + KRp ).

sup x,y∈B(0,R)

As ψ = W̃ − V vanishes on B c (0, R), its oscillation is attained on B(0, R). The bound of Lemma 16 gives sup ψ ≤ sup W̃ − inf W̃ , B(0,R)

B(0,R)

inf ψ ≥ inf W̃ − sup W̃ ,

|y|=R

B(0,R)

B(0,R)

|y|=R

hence osc(ψ) ≤ OSCR (W̃ ) + OSC∂B (W̃ ), with OSC∂B the oscillation over {|y| = R}. On that sphere W̃ = 2 W − βµ 2 R differs from W by a constant, so OSC∂B (W̃ ) = OSC∂B (W ) ≤ OSCR (W ), while OSCR (W̃ ) ≤ 2 OSCR (W ) + βµ 2 R . Therefore h i µ 2 2 p osc(ψ) ≤ 2 OSCR (W ) + βµ . 2 R ≤ β 4R (m + KR ) + 2 R The reference π0 ∝ e−W0 has βµ-strongly convex potential. For t, s > 0, t + s = 1, the identity t|x|2 + s|y|2 − |tx + 2 sy|2 = ts|x − y|2 and convexity of V = W0 − βµ 2 | · | give tW0 (x) + sW0 (y) − W0 (tx + sy) ≥

βµ ts |x − y|2 , 2

which is condition (3.1) of Bobkov and Ledoux [2000] with c = βµ. By Proposition 3.1 therein, without requiring 2 . Since additional smoothness of the potential, guarantees that π0 satisfies the log-Sobolev inequality with ρ0 ≤ βµ dπβ dπ0

∝ e−ψ with ψ bounded, the Holley–Stroock perturbation principle (Theorem 1, Chapter 5 in Markowich and Villani [2000]) yields ρ(β) ≤ ρ0 eosc(ψ) ≤

   2 exp β 4R (m + KRp ) + µ2 R2 , βµ

which is (27) with the stated constant. Remark 14. The exponential dependence on β in Lemma [13] is optimal. For a target that is nonconvex on a compact set, such as a mixture of Gaussians or a double-well potential, the logarithmic Sobolev and Poincaré constants grow like eβ∆ , where ∆ is the energy barrier separating the wells. This constant was shown to be sharp in Menz and Schlichting [2014]. Appendix I. Nanochat Architecture & Bounds The potentials arising in the pretraining of language models are simultaneously non-convex, non-smooth and of superlinear gradient growth, which is precisely the regime this work addresses. It is nevertheless one thing to assert that a class of potentials falls within the scope of A1–A3, and quite another to verify it for a fully specified architecture. The purpose of this section is to carry out that verification for one such architecture, and to exhibit the resulting constants

41

explicitly in terms of its depth, width and vocabulary. We work with the architecture of nanochat, a compact version of ChatGPT2 created by Karpathy. The section is organized as follows. Sections I.1 and I.3 lay out the architecture and gather the exact bounds on the scalar and vector nonlinearities the model uses, all of which are uniform in the parameter. Section I.4 writes the forward map and the potential u, and Section I.5 the curvature of u. Sections I.6 and I.7 introduce an exponent tracking notation for polynomially bounded maps, which enables us in Section I.8, to obtain the curvature bound ∇2θ u(θ) ⪰ −c⋆ 1 + |θ|

12L+2

Id¯.

Section I.9 then adds a higher-order regularization term and recovers A1–A3 for the regularized potential. Throughout this section the depth of the network is denoted L, as L already denotes the semi-convexity constant of A1. I.1. Architecture and parameters The architecture introduces the following constants L H V T τ >0

number of blocks (depth), dm number of heads, dh = dm /H vocabulary size, df = 4dm context length, ε>0 logit softcap, cqk > 0

model width (token space), head width, width of the perceptron, normalisation floor, query–key scale.

The parameter vector is the concatenation    L ¯ ℓ θ = E, WQℓ , WK , WVℓ , WOℓ , W1ℓ , W2ℓ ℓ=1 , Wlm ∈ Rd ,

(81)

with E ∈ RV ×dm ,

ℓ WQℓ , WK , WVℓ , WOℓ ∈ Rdm ×dm ,

W1ℓ ∈ Rdf ×dm ,

W2ℓ ∈ Rdm ×df ,

Wlm ∈ RV ×dm ,

so that the parameter space has dimension d¯ = 2V dm + 12Ld2m ,

and

|θ|2 = ∥E∥2F +

L X

∥WQℓ ∥2F + · · · + ∥W2ℓ ∥2F + ∥Wlm ∥2F .

(82)

ℓ=1 ℓ We write WQ,i ∈ Rdh ×dm for the i-th head slice of WQℓ , and likewise for K and V . Notice that nanochat doesn’t use ℓ grouped query attention, in which WK , WVℓ ∈ R(Hkv dh )×dm with Hkv | H, but since the analysis below uses nothing about the weight blocks beyond

∥W ∥ ≤ ∥W ∥F ≤ |θ|,

(83)

the proof’s argument go through for more general LLMs. I.2. Additional Notation For j ∈ {1, . . . , n} we write ej ∈ Rn for the j-th standard basis vector and 1 ∈ Rn for the vector all of whose entries ¯ equal 1. Therefore E⊤ ext as in (94) is understood as the xt -th row of E, and ⟨eyt , ξ⟩ = ξyt . Let F : Rd → RN and ¯ ¯ ¯ ¯ let v, w ∈ Rd directions. For N = 1 we write ∇θ F (θ) ∈ Rd for the gradient and ∇2θ F (θ) ∈ Rd×d for the Hessian matrix, and denote their action on those directions by
∇θ F [v] := v⊤ ∇θ F, ∇2θ F [v, w] := v⊤ ∇2θ F w, (84)

42

that is the directional derivative of F along v and the bilinear form of the Hessian. For N > 1 we abuse the notation in the following sense,
N ∇θ F [v] := ∇θ Fi [v] i=1 ,

N ∇2θ F [v, w] := ∇2θ Fi [v, w] i=1 ,

(85)

both elements of RN . Essentially ∇θ F denotes the N × d¯ Jacobian, whereas ∇2θ F is a bilinear map into RN . The associated operator norms are ∥∇2θ F ∥ :=

∥∇θ F ∥ := sup ∇θ F [v] , |v|=1

sup |v|=|w|=1

(86)

∇2θ F [v, w] ,

so that for N = 1 they reduce to the Euclidean norm of the gradient and the spectral norm of the Hessian, while for N > 1 the second is the corresponding tensor norm. Derivatives in a variable other than θ always carry that variable as a subscript. I.3. Nonlinearities We record exact bounds for the nonlinearities of the model, including their first two derivatives. Within this subsection k denotes a fixed width, which appears either as k = dm or k = dh . Definition 2 (Normalisation). For x ∈ Rk and ε > 0, Nk (x) :=

x , sk (x)

sk (x) :=



2 1 k |x| + ε

1/2

(87)

.

Definition 3 (Softmax, softcap, ReLU-squared, sigmoid). For z ∈ Rn let sm(z)i := ezi / χ(z) := τ tanh(z/τ ),

φ(z) := max(z, 0)2 ,

P

je

zj

, and for z ∈ R

ς(z) := (1 + e−z )−1 ,

each extended to vectors componentwise. Since χ, φ and ς act componentwise, their differentials are diagonal, i.e. ∇χ(z) = diag(χ′ (zj )) and ∇2 χ(z) is the diagonal tensor with entries χ′′ (zj ). Hence ∥∇χ(z)∥ ≤ supt∈R |χ′ (t)| and ∥∇2 χ(z)∥ ≤ supt∈R |χ′′ (t)|, and likewise for φ and ς. It therefore suffices to bound the one dimensional analogue. Lemma 18. Write s = sk (x) as in (87) and σ̄ = sm(z). (i) |Nk (x)| ≤

√

∥∇Nk (x)∥ ≤

k,

1 ≤ ε−1/2 , s

∥∇2 Nk (x)∥ ≤ √

6 6 ≤√ . 2 ks kε

(ii) ∇ sm(z) = diag(σ̄) − σ̄σ̄⊤ ⪰ 0 and ∥∇ sm(z)∥ ≤ 12 , (iii) |χ(z)| ≤ τ ,

′

|χ (z)| ≤ 1,

∥∇2 sm(z)∥ ≤ 1,

|ς ′ | ≤ 14 ,

∥∇2 sm(z)∥∞→1 ≤ 2.

√ 4 3 4 = |χ (z)| ≤ √ . 9τ 3 3τ ′′

(iv) |φ(z)| ≤ |z|2 , ∥∇φ(z)∥ ≤ 2|z|, ∥∇φ1 ∥ ≤ 1 and ∇2 φ1 = 0 a.e. (v) |ς| ≤ 1,

∥∇ sm(z)∥∞→1 ≤ 1,

∥∇2 φ∥ ≤ 2 and for φ1 (z) = max(z, 0) one has |φ1 (z)| ≤ |z|,

1 |ς ′′ | ≤ √ . 6 3

Here the mixed norm is defined as∥M ∥∞→1 := sup|v|∞ ≤1 |M v|1 . Proof. (i) Since ks2 = |x|2 + kε ≥ |x|2 , one obtains |x| ≤

√

43

k s,

(88)

whence |Nk (x)| = |x|/s ≤

√

k. Differentiating (87), ∇Nk (x) =

xx⊤  1 I− 2 . s ks

(89)

The matrix P := xx⊤ /(ks2 ) is symmetric positive semi-definite with ∥P ∥ = |x|2 /(ks2 ) ≤ 1 by (88), so I − P has spectrum in [0, 1] and ∥∇Nk ∥ ≤ 1/s ≤ ε−1/2 . Differentiating (89) once more, ∂l ∂j (Nk )i = −

δij xl + δil xj + δjl xi 3 xi xj xl + , 3 ks k 2 s5

so that, for unit vectors v, w, ∇2 Nk (x)[v, w] = −

v ⟨x, w⟩ + w ⟨x, v⟩ + x ⟨v, w⟩ 3 x ⟨x, v⟩⟨x, w⟩ + . ks3 k 2 s5

In view of the Cauchy–Schwarz inequality and (88), one calculates √ √ 3 k 3k k 6 6 3|x| 3|x|3 2 + 2 2 =√ ≤√ . ∇ Nk (x)[v, w] ≤ 3 + 2 5 ≤ ks k s ks2 k s k s2 kε (ii) From ∂j σ̄i = σ̄i (δij − σ̄j ) one writes ∇ sm(z) = diag(σ̄) − σ̄σ̄⊤ , and therefore, for any v ∈ Rn , v⊤ ∇ sm(z) v =

X

σ̄i vi2 −

X

i

σ̄i vi

2

i

=

X

σ̄i ṽi2 ,

(90)

ṽ := v − ⟨σ̄, v⟩1.

i

P P P √ √ Applying the Cauchy–Schwarz inequality to ( σ̄i )i and ( σ̄i vi )i , and using i σ̄i = 1, gives ( i σ̄i vi )2 ≤ i σ̄i vi2 , so v⊤ ∇ sm(z) v ≥ 0 and thus ∇ sm(z) ⪰ 0. Notice next that ∇ sm(z)1 = σ̄ − σ̄ = 0, so that v⊤ ∇ sm(z) v = (v − c1)⊤ ∇ sm(z) (v − c1) for any c ∈ R. Taking c = 21 (maxi vi + mini vi ) and discarding the (nonnegative) second term in (90), v⊤ ∇ sm(z) v ≤

X

σ̄i (vi − c)2 ≤ max(vi − c)2 =

 max v − min v 2 i

i

i

i

2

i

i

(91)

.

Since (maxi vi − mini vi )2 ≤ 2 maxi vi2 + mini vi2 ≤ 2|v|2 , (91) yields Varσ̄ (v) ≤ 12 |v|2 , that is ∥∇ sm(z)∥ ≤ 12 ; while for |v|∞ ≤ 1 the range is at most 2 and (91) yields v⊤ ∇ sm(z) v ≤ 1. Furthermore, by Cauchy–Schwarz, ∇ sm(z)v 1 =

X i

σ̄i |ṽi | ≤

X

σ̄i ṽi2

1/2

= v⊤ ∇ sm(z) v

1/2

≤1

when |v|∞ ≤ 1,

i

which is the third bound. For the second differential, differentiating ∂j σ̄i once more gives ∂l ∂j σ̄i = σ̄i (δil − σ̄l )(δij − σ̄j ) − σ̄i σ̄j (δjl − σ̄l ), that is, in the notation of (90), ∇2 sm(z)[v, w]i = σ̄i ṽi w̃i −

X

(92)

σ̄i ṽi w̃i .

i

Hence one calculates, ∇2 sm(z)[v, w] ≤ ∇2 sm(z)[v, w] 1 ≤ 2

X

σ̄i |ṽi w̃i | ≤ 2 v⊤ ∇ sm(z) v

1/2

w⊤ ∇ sm(z) w

1/2

.

(93)

i

Plugging in v⊤ ∇ sm(z) v ≤ 12 |u|2 gives ∥∇2 sm∥ ≤ 1, and inserting v⊤ ∇ sm(z) v ≤ 1 for |v|∞ , |w|∞ ≤ 1 gives ∥∇2 sm∥∞→1 ≤ 2. (iii) One calculates χ′ (t) = sech2 (t/τ ) ∈ (0, 1] and χ′′ (t) = − τ2 sech2 (t/τ ) tanh(t/τ ). Writing y = tanh(t/τ ) ∈ 44

√ √ (−1, 1) one obtains |χ′′ (t)| = τ2 |y(1 − y 2 )|, which attains its maximum at y = 1/ 3 with value 2/(3 3).
P P (iv) |φ(z)|2 = j max(zj , 0)4 ≤ ( j zj2 )2 = |z|4 , ∇φ(z) = diag 2 max(zj , 0) has norm at most 2 maxj |zj | ≤ 2 2|z| is diagonal with entries 2·1{zj >0} . For φ1 , one has ∇φ1 (z) = diag(1{zj >0} ), whence |∇φ1 (z)v|2 = P and ∇ φ(z) 2 2 ′′ j 1{zj >0} vj ≤ |v| , and φ1 = 0 a.e. (v) ς ′ = ς(1 − ς) ≤ 41 and ς ′′ = ς(1 − ς)(1 − 2ς). Rewriting it as q = ς − 12 one gets |ς ′′ | = 2|q|( 14 − q 2 ), maximal √ √ at q = 1/(2 3) with value 1/(6 3). I.4. The forward map and the potential Fix a token sequence x = (x0 , . . . , xT −1 ) ∈ {1, . . . , V }T with targets yt = xt+1 , let Rt ∈ Rdh ×dh denote the parameter-free orthogonal rotation applied at position t, and let Sℓ (t) ⊆ {0, . . . , t} be the attention window of layer ℓ at position t. For every t define the following. Embedding. (the xt -th row of E).

ft0 = E⊤ ext ∈ Rdm

(94)

For ℓ ∈ {1, . . . , L}. Block ℓ: Attention.
ḡtℓ = Ndm ftℓ−1 ∈ Rdm

ℓ ℓ ℓ ℓ qt,i = Rt Ndh WQ,i ḡtℓ , ks,i = Rs Ndh WK,i ḡsℓ ∈ Rdh cqk ℓ aℓt,s,i = √ q , kℓ , s ∈ Sℓ (t) dh t,i s,i
pℓt,·,i = sm aℓt,·,i ∈ R|Sℓ (t)| ℓ ℓ vs,i = WV,i ḡsℓ ∈ Rdh

oℓt = WOℓ

H M X

 ℓ pℓt,s,i vs,i ∈ Rdm

(normalise)

(95)

(queries and keys)

(96)

(scores)

(97)

(attention weights)

(98)

(values)

(99)

(output projection)

(100)

(update)

(101)

i=1 s∈Sℓ (t) ℓ− 21

ft

= ftℓ−1 + oℓt

Block ℓ: MLP. ℓ− 21

b̄ℓt = Ndm ft

(normalise)

(102)

mℓt = W2ℓ φ W1ℓ b̄ℓt ∈ Rdm

(perceptron)

(103)

(update)

(104)

ℓ− 12

ftℓ = ft

+ mℓt

Head. (105)

zt = Ndm ftL ∈ Rdm , V

ζt = Wlm zt ∈ R V

ζ̂t = χ(ζt ) ∈ R

(unembedding)

(106)

(softcap)

(107)

Definition 4. (Cross-Entropy) Let D denote the distribution of token sequences. The pretraining potential is u(θ) = Ex∼D

−1 h 1 TX

T t=0

i CEt (θ) ,

CEt (θ) = log

V X j=1

45

eζ̂t,j (θ) − ζ̂t,yt (θ),

(108)

with ζ̂t given by (94)–(107). P Writing CE : RV → R, CE(ξ) = log j eξj − ⟨eyt , ξ⟩, the target yt is a fixed index and (108) reads u = P E[ T1 t CE ◦ζ̂t ], a composition of a fixed convex function on RV with the θ−dependent map ζ̂t . Remark 15. Since D is supported on the finite set {1, . . . , V }T and every derived bound is uniform in the token sequence x, a lower bound of the form ∇2θ ux ⪰ −c I would hold for any data sequence. We therefore fix x and suppress it. I.5. Curvature Fix t and let ξ ∈ RV , so that CEt (ξ) = log

PV

ξj j=1 e − ξyt . Differentiating in ξi ,

∂ CEt e ξi = P ξj − δi,yt = sm(ξ)i − δi,yt , ∂ξi je

(109)

∂ 2 CEt ∂ sm(ξ)i = = sm(ξ)i δij − sm(ξ)j . ∂ξi ∂ξj ∂ξj

(110)

and,

Evaluating (109) and (110) at ξ = ζ̂t and collecting the entries into a vector and a matrix respectively, we set σ̄t := sm(ζ̂t ),

gt := ∇ζ̂t CEt = σ̄t − eyt ,

Λt := ∇2ζ̂ CEt = diag(σ̄t ) − σ̄t σ̄⊤ t. t

(111)

By Lemma 18(ii) the matrix Λt is positive semi-definite, that is CEt is convex in ξ. P Applying the chain rule to u = T1 t CE ◦ζ̂t once gives the gradient of the potential, ∇θ u(θ) =

T −1 T −1 i
⊤ h

⊤
1 X 1 X ∇θ ζ̂t ∇ θ ζt sm(χ(ζt )) − eyt ⊙ χ′ (ζt ) , sm(ζ̂t ) − eyt = T t=0 T t=0

(112)

where ⊙ denotes the Hadamard product and the second equality follows from ∇θ ζ̂t = diag χ′ (ζt ) ∇θ ζt . In the ¯ notation of (84), ∇θ ζ̂t is the V × d¯ Jacobian and ∇θ u(θ) a vector of Rd . Observe that sm(ζ̂t ) − eyt is the predicted distribution less the observed one. Differentiating once more yields the exact identity ∇2θ u(θ) =

T −1 V i X 1 Xh (∇θ ζ̂t )⊤ Λt (∇θ ζ̂t ) + (gt )j ∇2θ ζ̂t,j . {z } T t=0 | j=1 Gauss–Newton term, ⪰ 0 | {z }

(113)

residual term

Here ∇θ ζ̂t is the V × d¯ Jacobian and ∇2θ ζ̂t,j the d¯ × d¯ Hessian of the scalar map θ 7→ ζ̂t,j (θ). One differentiates in θ, while gt and Λt are derivatives of CE in its own argument, evaluated at ζ̂t (θ). Since A⊤ M A inherits the sign of M , the first term of (113) is positive semi-definite and may be discarded in the analysis. Proposition 6. For every θ ∈ / N , where N is the Lebesgue-null set of Remark 18, h 4 i ∇2θ u(θ) ⪰ − √ max∥∇θ ζt ∥2 + 2 max∥∇2θ ζt ∥ Id¯. t 3 3τ t Proof. Discarding the Gauss–Newton term of (113) and writing T −1 V

P :=

1 XX (gt )j ∇2θ ζ̂t,j , T t=0 j=1

46

(114)

one has ∇2θ u ⪰ P ⪰ −∥P ∥Id¯. Observe first that

X σ̄t,j = 2 1 − σ̄t,yt ≤ 2, |gt |1 = 1 − σ̄t,yt +

|gt |∞ ≤ 1.

(115)

j̸=yt

Since χ acts componentwise, the chain rule gives ∇2θ ζ̂t,j = χ′ (ζt,j ) ∇2θ ζt,j + χ′′ (ζt,j ) ∇θ ζt,j ⊗ ∇θ ζt,j , and we bound the two resulting pieces of P , the first gives, by |χ′ | ≤ 1 and (115), X (gt )j χ′ (ζt,j ) ∇2θ ζt,j

≤ |gt |1 max∥∇2θ ζt,j ∥ ≤ 2∥∇2θ ζt ∥, j

j

while the second, being componentwise in χ, reassembles into a quadratic form, X
(gt )j χ′′ (ζt,j ) ∇θ ζt,j ⊗ ∇θ ζt,j = (∇θ ζt )⊤ diag gt ⊙ χ′′ (ζt ) (∇θ ζt ), j

whose norm is at most |gt |∞ ∥χ′′ ∥∞ ∥∇θ ζt ∥2 ≤ 3√43τ ∥∇θ ζt ∥2 by Lemma 18(iii) and (115). Averaging over t and P using T1 t ≤ maxt yields (114). I.6. Tracking the growth We now introduce the device by which we perform the induction through the LLM’s layers. ¯

Definition 5. Let F : Rd → RN be twice differentiable a.e. We write F ∈ T (a0 , a1 , a2 ; C0 , C1 , C2 ) if, for all θ,
a0
a1
a 2 |F (θ)| ≤ C0 1 + |θ| , ∥∇θ F (θ)∥ ≤ C1 1 + |θ| , ∥∇2θ F (θ)∥ ≤ C2 1 + |θ| . (116) When only the exponents are regarded, we abbreviate F ∈ T (a0 , a1 , a2 ). Lemma 19. Let F ∈ T (a0 , a1 , a2 ; C0 , C1 , C2 ) and G ∈ T (b0 , b1 , b2 ; D0 , D1 , D2 ).
(i) (Addition.) F + G ∈ T max{a0 , b0 }, max{a1 , b1 }, max{a2 , b2 }; C0 + D0 , C1 + D1 , C2 + D2 . (ii) (Bilinear maps.) If B is bilinear with |B(ξ, η)| ≤ β|ξ||η|, then B(F, G) ∈ T (c0 , c1 , c2 ; K0 , K1 , K2 ) with c0 = a0 + b0 ,

K0 = βC0 D0 ,

c1 = max{a0 + b1 , a1 + b0 },

K1 = β(C0 D1 + C1 D0 ),

c2 = max{a0 + b2 , a1 + b1 , a2 + b0 },

K2 = β(C0 D2 + 2C1 D1 + C2 D0 ).

(iii) (Composition with a bounded function.) If g : RN → RM satisfies |g| ≤ g0 , ∥∇g∥ ≤ g1 and ∥∇2 g∥ ≤ g2 globally, then
g ◦ F ∈ T 0, a1 , max{2a1 , a2 }; g0 , g1 C1 , g1 C2 + g2 C12 . (iv) (Composition with the activation function.)
φ ◦ F ∈ T 2a0 , a0 + a1 , max{2a1 , a0 + a2 }; C02 , 2C0 C1 , 2C12 + 2C0 C2 , whereas φ1 ◦ F ∈ T (a0 , a1 , a2 ; C0 , C1 , C2 ). (v) (Coordinate blocks.) If F (θ) = W (θ) is a coordinate block of θ, then F ∈ T (1, 0, 0; 1, 1, 0). (vi) (Weighted sum by softmax.) Let S be a finite index set and let the families α = (αs )s∈S and w = (ws )s∈S satisfy the bounds (116) entrywise, that is with maxs in place of the norm, with exponents and constants (a· ; C· )

47

and (b· ; D· ) respectively. Put p = sm(α) and y =

P

s∈S ps ws . Then y ∈ T (c0 , c1 , c2 ; K0 , K1 , K2 ) with

c0 = b0 ,

K0 = D0 ,

c1 = max{a1 + b0 , b1 },

K1 = C1 D0 + D1 ,

c2 = max{a2 + b0 , 2a1 + b0 , a1 + b1 , b2 },

K2 = (C2 + 2C12 )D0 + 2C1 D1 + D2 .

Proof. (i) One trivially has (1 + |θ|)a0 ≤ (1 + |θ|)max{a0 ,b0 } and likewise for b0 , whence |F + G| ≤ |F | + |G| ≤ (C0 + D0 )(1 + |θ|)max{a0 ,b0 } . The same argument applies to ∇θ and ∇2θ . (ii) By the Leibniz rule, ∇θ B(F, G) = B(∇θ F, G)+B(F, ∇θ G) and ∇2θ B(F, G) = B(∇2θ F, G)+2B(∇θ F, ∇θ G)+ B(F, ∇2θ G). One multiplies the corresponding bounds and takes the maximum of the exponents to obtain the claim. (iii) The chain rule gives |g ◦ F | ≤ g0 = g0 (1 + |θ|)0 . Also one writes ∇θ (g ◦ F ) = ∇g(F ) ∇θ F , whence ∥∇θ (g ◦ F )∥ ≤ g1 C1 (1 + |θ|)a1 and     ∇2θ (g ◦ F )[v, w] = ∇2 g(F ) ∇θ F [v], ∇θ F [w] + ∇g(F ) ∇2θ F [v, w] , (117) so that ∥∇2θ (g ◦ F )∥ ≤ g2 C12 (1 + |θ|)2a1 + g1 C2 (1 + |θ|)a2 . (iv) One repeats the computations of Lemma 18(iv) to get : |φ(F )| ≤ |F |2 ≤ C02 (1 + |θ|)2a0 , and ∥∇θ (φ ◦ F )∥ ≤ 2|F |∥∇θ F ∥ ≤ 2C0 C1 (1 + |θ|)a0 +a1 . Finally, by (117), ∥∇2θ (φ ◦ F )∥ ≤ 2C12 (1 + |θ|)2a1 + 2C0 C2 (1 + |θ|)a0 +a2 . For φ1 the three bounds of Lemma 18(iv) are 1, 1 and 0, so that the first term of (117) vanishes and the rule is transparent in all three exponents. (v) |W (θ)| = ∥W ∥F ≤ |θ| ≤ 1 + |θ| by (83) and ∇θ W is defined a.e. with uniformly bounded norm by 1. Moreover, ∇2θ W = 0. P P P (vi) Since p lies in the simplex, |y| ≤ s ps |ws | ≤ maxs |ws |. Next, ∇θ y = s (∇θ ps )ws + s ps ∇θ ws , the second sum being again a convex combination and bounded by maxs ∥∇θ ws ∥. Regarding the first differential, let v be unit vector such that , X ∇θ ps [v] ws ≤ max |ws | ∇ sm(α) ∇θ α[v] 1 ≤ max |ws | ∥∇ sm∥∞→1 ∇θ α[v] ∞ ≤ max |ws | max∥∇θ αs ∥, s

s

s

s

s

by Lemma 18(ii). For the second differential, expand X X X ∇2θ y = (∇2θ ps )ws + 2 (∇θ ps ) ⊗ (∇θ ws ) + ps ∇2θ ws , s

s

s

the last sum being bounded by maxs ∥∇2θ ws ∥ and the middle one by 2 maxs ∥∇θ ws ∥ maxs ∥∇θ αs ∥ as above. For the first, (117) applied to p = sm ◦α gives 2

∇2θ p[v, v] 1 ≤ ∥∇2 sm∥∞→1 ∇θ α[v] ∞ + ∥∇ sm∥∞→1 ∇2θ α[v, v] ∞ ≤ 2 max∥∇θ αs ∥2 + max∥∇2θ αs ∥, s

s

which upon multiplication by maxs |ws | yieldsthe exponents 2a1 + b0 and a2 + b0 and the constants 2C12 D0 and C2 D 0 . Corollary 2. Let F ∈ T (a0 , a1 , a2 ; C0 , C1 , C2 ) take values in Rk . Then   √ Nk ◦ F ∈ T 0, a1 , max{2a1 , a2 }; k, ε−1/2 C1 , ε−1/2 C2 + √6k ε C12 . Proof. Apply Lemma 19(iii) with g = Nk and the constants g0 = all of which are finite because ε > 0.

48

√

√ k, g1 = ε−1/2 and g2 = 6/( kε) of Lemma 18(i),

I.7. The layer recursion We track, for the residual stream at token t,
ftℓ ∈ T Aℓ , Bℓ , Cℓ ,

(118)

suppressing the constants, which are collected later in Section [X]. Lemma 20 (Base). ft0 ∈ T (1, 0, 0), that is A0 = 1, B0 = 0 and C0 = 0. Proof. By (94), ft0 is a coordinate slice of E; apply Lemma 19(v). Lemma 21 (Attention sublayer). Suppose ftℓ−1 ∈ T (A, B, C). Then
oℓt ∈ T 2, 3 + B, max{4 + 2B, 3 + C} ,

(119)

and consequently, by (101) and Lemma 19(i), ℓ− 12

ft

∈ T max{A, 2}, 3 + B, max{4 + 2B, 3 + C} .

(120)

Proof. Write B ⋆ := max{2B, C}. By Corollary 2 applied to (95),
ḡtℓ = Ndm ftℓ−1 ∈ T (0, B, B ⋆ ). ℓ ∈ T (1, 0, 0), so that Lemma 19(ii) applied to the matrix–vector product, Queries and keys. By Lemma 19(v), WQ,i for which β = 1 since |W x| ≤ ∥W ∥F |x|, gives

ℓ WQ,i ḡtℓ ∈ T 1, max{1 + B, 0}, max{1 + B ⋆ , B, 0} = T 1, 1 + B, 1 + B ⋆ ,

the collapse of the two maxima following from B ⋆ ≥ 2B ≥ B. Composing with Ndh through Corollary 2, and with the orthogonal Rt , which is invariant,
ℓ qt,i ∈ T 0, 1 + B, max{2 + 2B, 1 + B ⋆ } , (121) ℓ and identically for ks,i . √ Scores and weights. Equation (97) is a bilinear form with β = cqk / dh , so that Lemma 19(ii) applied to two factors of the class (121) returns that same class,
aℓt,s,i ∈ T 0, 1 + B, α2 , α2 := max{2 + 2B, 1 + B ⋆ }. ℓ ℓ ḡ ℓ ∈ T (1, 1 + B, 1 + B ⋆ ), entrywise in s. Values. Alike the queries, vs,i = WV,i P ℓ s ℓ Contraction. The weighted sum s pt,s,i vs,i is now handled by Lemma 19(vi), with (a0 , a1 , a2 ) = (0, 1 + B, α2 ) and (b0 , b1 , b2 ) = (1, 1 + B, 1 + B ⋆ ):

c0 = 1,

c1 = max{(1 + B) + 1, 1 + B} = 2 + B,  c2 = max α2 + 1, 2(1 + B) + 1, (1 + B) + (1 + B), 1 + B ⋆ = max{3 + 2B, 2 + B ⋆ }, P ℓ where the last equality uses α2 +1 = max{3+2B, 2+B ⋆ }. Thus s pℓt,s,i vs,i ∈ T (1, 2+B, max{3+2B, 2+B ⋆ }) for each head i, and the concatenation over i leaves these exponents unchanged, see Remark 16. Output projection. Applying WOℓ ∈ T (1, 0, 0) through Lemma 19(ii),

oℓt ∈ T 2, max{3 + B, 1}, max{1 + max{3 + 2B, 2 + B ⋆ }, 2 + B, 1} = T 2, 3 + B, max{4 + 2B, 3 + B ⋆ } . Substituting B ⋆ = max{2B, C} gives 3 + B ⋆ = max{3 + 2B, 3 + C} and hence max{4 + 2B, 3 + B ⋆ } = max{4 + 2B, 3 + C}, which is (119). Finally (120) follows from Lemma 19(i) applied to (101), using B ≤ 3 + B and C ≤ max{4 + 2B, 3 + C}.

49

Remark preceding (119) √ are per head, whereas (100) concatenates. Write P 16 (Theℓ head concatenation). The bounds ℓ ℓ yi := s pℓt,s,i vs,i and using |yi | ≤ maxs |WV,i ḡsℓ | ≤ ∥WV,i ∥F dm , then one obtains H M i=1

2

yi

=

H X

|yi |2 ≤ dm

i=1

H X
2 ℓ ∥WV,i ∥2F = dm ∥WVℓ ∥2F ≤ dm 1 + |θ| , i=1

the middle equality holding because the head slices partition the rows of WVℓ . The same argument applies verbatim to the first and second differentials. ℓ− 12

Lemma 22 (Perceptron sublayer). Suppose ft

∈ T (A′ , B ′ , C ′ ). Then (122)

mℓt ∈ T 3, 3 + B ′ , max{3 + 2B ′ , 3 + C ′ } ,
and consequently ftℓ ∈ T max{A′ , 3}, 3 + B ′ , max{3 + 2B ′ , 3 + C ′ } .

Proof. Write B ′⋆ := max{2B ′ , C ′ }. By Corollary 2 applied to (102), b̄ℓt ∈ T (0, B ′ , B ′⋆ ), and by Lemma 19(v) and (ii),
W1ℓ b̄ℓt ∈ T 1, 1 + B ′ , 1 + B ′⋆ . Applying Lemma 19(iv),


φ W1ℓ b̄ℓt ∈ T 2, 1 + (1 + B ′ ), max{2(1 + B ′ ), 1 + (1 + B ′⋆ )} = T 2, 2 + B ′ , max{2 + 2B ′ , 2 + B ′⋆ } , and then W2ℓ ∈ T (1, 0, 0) through Lemma 19(ii),

mℓt ∈ T 3, max{3 + B ′ , 2}, max{1 + max{2 + 2B ′ , 2 + B ′⋆ }, 2 + B ′ , 2} = T 3, 3 + B ′ , max{3 + 2B ′ , 3 + B ′⋆ } . Substituting B ′⋆ = max{2B ′ , C ′ } gives 3 + B ′⋆ = max{3 + 2B ′ , 3 + C ′ }, whence (122). The statement for ftℓ follows from Lemma 19(i) applied to (104), using B ′ ≤ 3 + B ′ and C ′ ≤ max{3 + 2B ′ , 3 + C ′ }. Proposition 7 (Induction). For every ℓ ∈ {0, 1, . . . , L}, ( 1, ℓ = 0 Aℓ = Bℓ = 6ℓ, 3, ℓ ≥ 1,

( Cℓ =

0, 12ℓ − 3,

ℓ=0 ℓ ≥ 1.

(123)

Proof. Composing Lemmas 21 and 22, that is inserting A′ = max{Aℓ−1 , 2}, B ′ = 3 + Bℓ−1 and C ′ = max{4 + 2Bℓ−1 , 3 + Cℓ−1 } into Lemma 22, one obtains the recursions  Aℓ = max{Aℓ−1 , 3}, Bℓ = 6 + Bℓ−1 , Cℓ = max 9 + 2Bℓ−1 , 3 + max{4 + 2Bℓ−1 , 3 + Cℓ−1 } , the first because max{max{Aℓ−1 , 2}, 3} = max{Aℓ−1 , 3}. With A0 = 1 and = B0 = C0 = 0, from Lemma 20, the first two are immediate. For the third, substituting Bℓ−1 = 6ℓ − 6,    Cℓ = max 12ℓ − 3, 3 + max{12ℓ − 8, 3 + Cℓ−1 } = max 12ℓ − 3, 12ℓ − 5, 6 + Cℓ−1 = max 12ℓ − 3, 6 + Cℓ−1 . Thus C1 = max{9, 6} = 9 = 12 − 3, and if Cℓ−1 = 12(ℓ − 1) − 3 then 6 + Cℓ−1 = 12ℓ − 9 < 12ℓ − 3, so that Cℓ = 12ℓ − 3. Remark 17. The exponent Aℓ = 3 does not depend on L since the normalisations (95) and (102) discard the size of the incoming stream and nothing accumulates across the depth. Corollary 3 (The logit map). ftL ∈ T (3, 6L, 12L − 3), and consequently

zt ∈ T 0, 6L, 12L , ζt ∈ T 1, 6L + 1, 12L + 1 .

50

(124)

Proof. By Corollary 2 applied to (105), zt ∈ T (0, 6L, max{12L, 12L − 3}) = T (0, 6L, 12L). By Lemma 19(v) and (ii) applied to (106),
ζt ∈ T 1 + 0, max{1 + 6L, 0}, max{1 + 12L, 6L, 0} = T (1, 6L + 1, 12L + 1).

I.8. The curvature bound (not sharp) Theorem 3. There exists a constant c⋆ = c⋆ (dm , dh , H, ε, τ, cqk , L), independent of θ, of T and of the token sequence, ¯ such that for a.e. θ ∈ Rd ∇2θ u(θ) ⪰ −c⋆ 1 + |θ|

12L+2

(125)

Id¯

The constant is computed in Section [X]. Proof. By Corollary 3 there are constants Γ1 , Γ2 with ∥∇θ ζt ∥ ≤ Γ1 1 + |θ|

6L+1

∥∇2θ ζt ∥ ≤ Γ2 1 + |θ|

,

12L+1

,

uniformly in t. Inserting these into Proposition 6, h 4 Γ2
12L+2
12L+1 i
12L+2 + 2Γ2 1 + |θ| Id¯, Id¯ ⪰ −c⋆ 1 + |θ| ∇2θ u(θ) ⪰ − √ 1 1 + |θ| 3 3τ √ with c⋆ = 4Γ21 /(3 3τ ) + 2Γ2 , the last step using 1 + |θ| ≥ 1 to absorb the lower exponent into the higher. Remark 18. Every function of Lemma 18 is C ∞ , with the exception of Relu-squared, φ, whose second derivative φ′′ = 2 · 1{ · >0} jumps at the origin. The exceptional set is therefore the finite union N =

[n θ:

ℓ− 12

W1ℓ Ndm (ft

o
) j =0 ,

(126)

ℓ,j,t

and each constituent is a C 1 hypersurface, hence Lebesgue-null. Remark 19. The softcap gives |ζ̂t,j | ≤ τ for every t, j and every θ, so that by (108) 0 ≤ u(θ) ≤ log V + 2τ

¯

for all θ ∈ Rd .

(127)

¯

Hence, e−βu is not integrable on Rd , i.e. no Gibbs measure can be associated with u without proper regularization. I.9. Regularization Following Lovas et al. [2023] and Lim and Sabanis [2024], we add a higher-order term to the potential. Fix η > 0 and set û(θ) := u(θ) + η |θ|12L+5 ,

¯

θ ∈ Rd .

(128)

One calculates
∇2θ |θ|12L+5 = (12L + 5)|θ|12L+3 Id¯ + (12L + 5)(12L + 3) |θ|12L+1 θθ⊤ ⪰ (12L + 5)|θ|12L+3 Id¯,

(129)

the discarded term being positive semi-definite. Throughout this subsection we abbreviate c̄⋆ := 212L+1 c⋆ ,

so that

c⋆ 1 + |θ|

51

12L+2

≤ c̄⋆ 1 + |θ|12L+2 .

(130)

¯

Proposition 8 (A1 and A2 hold for û). Let η > 0 and let û be given by (128). Then, for a.e. θ ∈ Rd , i h
∇2θ û(θ) ⪰ η(12L + 5)|θ| − c̄⋆ |θ|12L+2 − c̄⋆ Id¯.

(131)

Consequently: (i) û satisfies A1 with the explicit constant " L = c̄

⋆

1 1+ 12L + 3



(12L + 2) c̄⋆ (12L + 3) η (12L + 5)

12L+2 # ;

(ii) for every µ > 0, ∇2θ û(θ) ⪰ µ Id¯ a.e. on {|θ| ≥ ϱ(µ)}, where   2(c̄⋆ + µ) 1/(12L+3)  2c̄⋆ ; ϱ(µ) = max , η(12L + 5) η(12L + 5)

(132)

(133)

(iii) û satisfies A2 with constant µ/2 and radius R = 4ϱ(µ) µ + L /µ. Proof. Combining Theorem 3, (129) and (130), h
12L+2 i ∇2θ û(θ) ⪰ η(12L + 5)|θ|12L+3 − c⋆ 1 + |θ| Id¯ h i ⪰ η(12L + 5)|θ|12L+3 − c̄⋆ |θ|12L+2 − c̄⋆ Id¯, which is (131) upon factoring |θ|12L+2 . (i) Write q := 12L + 2, κ := η(12L + 5), and ψ(s) := κsq+1 − c̄⋆ sq − c̄⋆ ,

s ≥ 0,

for the scalar appearing in (131), so that ∇2θ û ⪰ −LId¯ a.e. with L = sups≥0 −ψ(s) . One calculates
ψ ′ (s) = −sq−1 c̄⋆ q − κ(q + 1)s ,
which is positive on (0, s⋆ ) and negative on (s⋆ , ∞), with s⋆ := c̄⋆ q/ κ(q + 1) . Hence the supremum is attained at s⋆ , and since c̄⋆ − κs⋆ = c̄⋆ /(q + 1), 
sq⋆  L = −ψ(s⋆ ) = c̄⋆ + sq⋆ c̄⋆ − κs⋆ = c̄⋆ 1 + , q+1 which is (132). Moreover, since −ψ(0) = c̄⋆ , we have that L ≥ c̄⋆ > 0. (132) (ii) Suppose first η(12L + 5)|θ| ≥ 2c̄⋆ , which implies η(12L + 5)|θ|12L+3 − c̄⋆ |θ|12L+2 ≥ 21 η(12L + 5)|θ|12L+3 and hence ψ(|θ|) ≥ 12 η(12L + 5)|θ|12L+3 − c̄⋆ . Suppose in addition |θ| ≥ (2(c̄⋆ + µ)/(η(12L + 5)))1/(12L+3) , then 12 η(12L + 5)|θ|12L+3 ≥ c̄⋆ + µ and therefore ψ(|θ|) ≥ µ. (iii) By (i) the potential û is L-semi-convex, and by (ii) its Alexandrov Hessian satisfies ∇2 û ⪰ µ I a.e. on {|θ| ≥ ϱ(µ)}. Lemma 1 applies with µ0 = µ and ρ = ϱ(µ) and returns A2 with constant µ/2 and radius 4ϱ(µ)(µ+L)/µ.

52

Appendix J. Table of Constants

Table 5: Analytic expressions of constants. The last column reports the order in powers of (d/β), suppressing logarithmic factors and treating m, K, L, µ, R, λ0 and E|θ0 |2p /(d/β)p as O(1). The constants ρ(β), Cr , Cw , T0 are dimension free and are treated as O(1) here, but depend exponentially on β in general, see Remark 11. For the coordinate-wise variants CB1 ,c , CB2 ,c , CB4 ,c the dependence on d is decoupled from β. We write Dp := (2p − 1)!!. No. 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23

Constants  b = max m2 /(2µ),mR + (L + µ/2)R2 C1 = 4 max 2m2 , µ2 , 2K 2 C2 = p(b + β −1 (d + p − 2)) A(p) = 1 + 2(b + β −1 (d + (p − 2)+ ))/µ C4 (p) ≤ A(p)max{1, p/2} C3 (p) −1 CB1 = E|θ0 |2 + 4(1/µ)(1 + pb + dβ ) 3p−3 Λp = p(p + 1)2 Dp M (p) = 23p+1 Λpp CB2 = E|θ0 |2p + M (p)(1/µ)p (1 + b + dβ −1 )p CB3 = 2p−1 (CB2 + 1) + A(2p)p CB4 = µCB1 + 4 + 4dβ −1    2 ρ(β) ≤ exp β 4R(m + KRp ) + µ2 R2 βµ 2 Cr = βρ(β) r 1  12 + βLρ(β)    βLρ(β) Cw = 1 + βLρ(β) ≤ e 1 + 2 2

order O(1) O(1) O(d/β) O(d/β) O((d/β)max{1, p/2} ) O(d/β) O(1) O(1) O((d/β)p ) O((d/β)p ) O(d/β) O(1) O(1) O(1)

ln Cw + 1 O(1) + λ0 Cr  
1/2 2LC C (1+C ) 2 CE = e5LT0 T0 √µB4 + 1 L B2 + 4 max {2m2 , K 2 }(1 + CB2 + CB3 )CB4 O((d/β)p ) T0 =

1/2

CD = (2CB1 + 2A(2))  2p−1 K p CT1 = m + CB2 + A(2p)p/2 p+1 2p−2 p−1 J = m + K2 (1 + (2b/µ)p/2 ) + K2 /(p + 1) 2 2 d 2(b + d/β)β J 1 13 CT2 = log + log(πd) + 2β µ 2β 6β CB1 ,c = E|θ0 |2 + 4(1/µ)d(1 + b + β −1 ) CB2 ,c = E|θ0 |2p + M (p)(1/µ)p dp (1 + b + β −1 )p CB4 ,c = µCB1 ,c + 4d + 4dβ −1

53

O((d/β)1/2 ) O((d/β)p/2 ) O(1) O(d/β log(dβ)) O(d) O(dp ) O(d)

Related documents

• Sampled-Max Subgradient Method for Convex Finite-Max Optimization#1117585
• Online Convex Optimization with Sublinear Noisy Probes#271829
• Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity#1155025
• Localize, Restart, Accelerate: Stochastic Optimization under Generalized Smoothness#974669
• Nonmonotone subgradient methods based on a local descent lemma#1141448