Active-Trace Complexity Bounds for Moreau--Yosida Unadjusted Langevin Sampling

This paper studies the Moreau–Yosita unadjusted Langevin algorithm (MYULA) for sampling from a nonsmooth composite target distribution of the form

π(dx) ∝ exp −f(x) − g(x) dx, x ∈ R d,

where f is m-strongly convex with L f-Lipschitz gradient and g is convex and G-Lipschitz. The method replaces g by its Moreau envelope g λ, targets the smoothed measure π λ, and applies the Euler–Maruyama scheme. The standard analysis uses the global smoothness estimate Lip(∇g λ) ≤ λ-1, which leads to a worst-case treatment where the full trace of Moreau curvature is as large as d/λ everywhere. This creates a tension between regularization bias (requiring small λ) and discretization stability (requiring large λ).

The paper introduces the Moreau active trace:

a λ(x):= tr H λ(x) = (1/λ) tr(I − D prox λg(x))

where H λ is the a.e./weak Hessian of g λ. The quantity λa λ(x) is the total local shrinkage of the proximal map, describing how many directions are locally suppressed by the active structure of g.

The central quantity is the reference heat-path active trace:

B ref:= (1/h) ∫0ʰ E[a λ(Y t ref)] dt

where Y t ref = T h(X ref) + √(2)W t ref, with X ref ∼ π λ, T h(x) = x − h∇(f+g λ)(x), and W t ref a Brownian motion. This averages the weak Moreau curvature over the stationary input, the Euler drift, and the Gaussian heat interpolation forming a single discretization step.

Under the standing assumptions, the main result (Theorem 5.2) states:

Φ N ≲ e−mhN/2 Φ 0 + (1/m)[(τ f + G2 + B ref)h + M λ2h2],

where Φ k = W22(μ k, π λ) + 2h KL(μ k∥π λ), τ f:= sup x tr ∇2f(x), and M λ is an a.e. upper bound for a λ.

The fixed-λ iteration complexity is, up to logarithmic factors:

N λ(ε alg) ≲ (1/m)[L f + (τ f + G2 + B ref)/ε alg2 + M λ/ε alg].

The leading ε alg-2 term depends on the reference active trace; the worst-case curvature M λ survives only in a lower-order transfer term.

Proposition 5.9 proves:

√m W2(π λ, π) ≤ G2λ/4.

Thus, choosing λ ≍ ε/G2 gives an end-to-end guarantee for the original target π.

The universal estimate B ref ≤ d/λ yields O(ε-3) accuracy dependence. However, for structured piecewise-linear, lasso-type, group, and total-variation penalties, curvature–tube estimates make B ref independent of λ, yielding O(ε-2) for the same classical MYULA kernel.

The proof separates one MYULA update into a deterministic gradient step and a heat step. The gradient step is controlled by an approximate EVI estimate whose error depends on the Lipschitz size G2 of g, rather than on the Moreau smoothness scale λ-2. The key cancellation is:

∥v∥2 − 2⟨s, v⟩ = ∥r∥2 − ∥s∥2 ≤ G2 − ∥s∥2.

The heat step uses the entropy EVI for the heat semigroup and the heat energy identity for C 1,1 test functions.

The stepwise active trace B k is bounded by the reference active trace via:

B k ≤ B ref + M λ√(K k/2),

using total variation contraction under Markov kernels and Pinsker's inequality.

Proposition 6.3 formalizes: if a component of a λ has height O(λ-1) inside an O(λ) neighborhood of an active stratum Σ, and the reference heat path assigns mass O(r q) to an r-tube around Σ, then this component contributes O(λ q−1) to B ref. Codimension-one layers contribute O(1), while higher-codimension central regions contribute even less as λ ↓ 0.

Lemma 6.5 gives a slice conditional density bound for π λ along a subspace E, governed by directional quantities L E and G E rather than the global Moreau smoothness scale. Lemma 6.6 propagates this through the deterministic Euler map T h and the Gaussian heat step, requiring α E = 1 − h(L E + λ-1) > 0.

For g with finitely many kinks, the Moreau curvature is confined to intervals of total width λΔ (where Δ is the total slope jump). The reference trace satisfies B ref ≤ A 1d:= 2BΔ, independent of λ.

For g(x) = Σi gi(xi), the curvature is confined to coordinate slabs. B ref ≤ A sep:= 2Σi BiΔi, independent of λ. For weighted lasso g(x) = Σi γixi, B ref ≤ A wl:= 4Σi γiBi.

For g(x) = Σ b γ b∥x b∥, the Moreau trace decomposes across blocks. The central ball term and tangential curvature term are controlled separately. B ref ≤ A grp, independent of λ, with explicit block-size dependence discussed in Remark 7.7.

For g(x) = γ∥Dx∥1, the proximal map is affine on polyhedral cells with D prox λg(x) = P ker D A and a λ(x) = rank(D A)/λ. The curvature is localized by thin row slabs. B ref ≤ A D:= 4γΣⱼ RⱼBⱼ/∥dⱼ∥, independent of λ. For one-dimensional anisotropic total variation, B ref ≤ A tv:= 8√2 γB tv(d−1).

The paper compares its results to:

• The original MYULA analysis of Durmus, Moulines, and Pereyra (2018), which gives O(ε-2) fixed-λ bounds but O(ε-4) end-to-end rates when combined with the bias choice λ ≍ ε/G2.

• The Wasserstein–EVI framework of Durmus, Majewski, and Miasojedow (2019), which applied globally to U λ gives O(ε-3) end-to-end.

• The average-smoothness theory of Dalalyan and Karagulyan (2026), which averages across directions but remains worst-case over space, whereas active trace also averages over where the discretized path goes.

The paper develops an active-trace analysis of the classical fixed-λ MYULA kernel. The main bound replaces global Moreau-curvature control by an occupation-weighted trace along a reference heat path. Combined with the Moreau-bias estimate, this gives end-to-end Wasserstein guarantees for the original nonsmooth target and, for the structured penalties considered, O(ε-2) accuracy dependence without changing the MYULA transition.

Improvement 1: Adaptive Step-Size Control for Nonsmooth Sampling

• What: Implement a runtime monitor that estimates the reference heat-path active trace B ref online during sampling, rather than relying on worst-case global curvature bounds.

• What the improved AI system can do: Automatically tune the discretization step h and regularization lambda per problem instance. For structured penalties (lasso, group lasso, TV), it will detect low active trace and use larger steps, achieving O(epsilon-2) convergence without manual parameter tuning. For unstructured penalties, it falls back to conservative steps, preventing instability.

Improvement 2: Curvature-Aware Proximal Operator Selection

• What: Extend the active-trace theory to automatically choose between exact proximal steps and approximate (e.g., iterative) proximal solvers based on the local Moreau curvature a lambda(x).

• What the improved AI system can do: In regions where the active trace is low (most of the space for structured penalties), it uses cheap approximate proximal evaluations. Near active strata (where curvature spikes), it switches to exact solvers. This reduces per-iteration cost by up to an order of magnitude for high-dimensional lasso/TV problems while maintaining the same Wasserstein accuracy guarantee.

Improvement 3: Early-Stopping Criterion with Bias-Variance Trade-off

• What: Use the explicit bias bound sqrt m W 2(pi lambda, pi) at most G 2 lambda/4 and the iteration complexity N lambda(epsilon alg) to design a principled stopping rule that balances discretization error and regularization bias.

• What the improved AI system can do: Given a target total error epsilon, it automatically selects lambda proportional to epsilon/G squared and stops when the algorithmic error epsilon alg is a fixed fraction of epsilon. It outputs samples with a certified end-to-end Wasserstein distance to the true nonsmooth target, eliminating guesswork in hyperparameter selection.

Improvement 4: Dimension-Reduced Sampling for Structured Penalties

• What: Leverage the curvature–tube interface (Proposition 6.3) to identify low-dimensional active subspaces where the Moreau curvature is concentrated, then run the sampler in a reduced coordinate system.

• What the improved AI system can do: For group lasso or generalized lasso with rank-deficient D, it automatically detects the effective dimension r = rank(D A) on each polyhedral cell. It then performs the heat-step diffusion only in the active subspace, reducing computational cost from O(d) to O(r) per step, while preserving the O(epsilon-2) guarantee. This enables scaling to d > 10 6 for sparse signal recovery.

Improvement 5: Robustness Certificates for Nonsmooth Bayesian Inference

• What: Combine the active-trace bounds with the slice-density estimates (Lemma 6.5) to certify that the sampler’s output distribution is within a guaranteed Wasserstein radius of the true posterior, even when the likelihood is nonsmooth (e.g., Laplace noise, Huber loss).

• What the improved AI system can do: For Bayesian inverse problems with sparsity-promoting priors, it provides a rigorous certificate: The posterior mean/quantiles computed from N samples are within epsilon (Wasserstein-2) of the true posterior, with probability 0.95. This is valuable for safety-critical applications (medical imaging, geophysics) where uncertainty quantification must be trustworthy.

Improvement 6: Automatic Detection of Favorable Structure

• What: Build a pre-sampler diagnostic that estimates B ref from a short pilot run (e.g., 100 steps) and classifies the penalty as structured (low B ref, enabling fast rates) or unstructured (high B ref, requiring conservative settings).

• What the improved AI system can do: Given any new nonsmooth target, it automatically determines whether the O(epsilon-2) rate is achievable. If yes, it proceeds with aggressive step sizes; if no, it warns the user and uses the safe O(epsilon-3) configuration. This removes the need for users to manually analyze their penalty structure.