PaperlayerView exact revision on arXiv

Source-linked AI summary

On the computational cost of Stochastic Gradient Langevin Dynamics

Mateusz B. Majka, Tigran Nagapetyan, Łukasz Szpruch, Yue Wu, Danqi Zhuang

arXiv:2609.17750v1math.NAcs.NAmath.PRstat.CO

TL;DR

The paper examines whether SGLD’s lower per-step cost from subsampling outweighs the resulting variance when sampling to a prescribed accuracy. It compares SGLD and EM costs across dataset size, subsample size, and accuracy, finding distinct regimes in which either method can be preferable.

  • Problem

    Subsampling lowers computational cost but increases variance, leaving it unclear whether SGLD reduces the total cost of reaching a prescribed mean-square accuracy.

    Source: §1 · p. 1

  • Method

    The paper analyzes computational costs for SGLD and EM in finite-sum-drift settings as functions of dataset size, subsample size, and desired accuracy.

    Source: §1 · p. 2, §2.3 · p. 6

  • Results

    When m ≥ ε^-2, SGLD is no more expensive than EM at leading order across 1 ≤ s ≤ m; EM can be cheaper only in a small-data, aggressive-subsampling regime.

    Source: §2.3 · p. 8

  • Takeaways & Limitations

    There is no uniform cost winner: the leading-order comparison depends on m, ε, and the subsampling factor (m − s)/(ms).

    Source: §2.3 · p. 8

    The analysis focuses on the basic SGLD model rather than more advanced stochastic-gradient schemes.

    Source: §1 · p. 2

Abstract

from arXiv · show

Stochastic Gradient Langevin Dynamics (SGLD) reduces the cost of Langevin-based sampling by replacing full-dataset drift evaluations with mini-batch approximations, but the resulting subsampling error may offset this computational saving. We study this trade-off for stochastic differential equations with finite-sum drifts and compare the computational cost of SGLD with that of the Euler-Maruyama (EM) method. For a prescribed mean-square accuracy $\varepsilon^2$, we derive complexity estimates that explicitly track the dependence on the dataset size $m$, mini-batch size $s$, and accuracy parameter $\varepsilon$. The resulting comparison reveals distinct parameter regimes in which either method is preferable. In particular, EM can have lower leading-order cost only in a small-data, aggressive-subsampling regime, whereas SGLD is favoured over most of the remaining parameter space. In the practically relevant regime $s \ll m$, the transition between the two methods occurs at the scale $m \asymp \varepsilon^{-1}$. We complement the theoretical analysis with numerical experiments based on a Gaussian Bayesian inference model, which examine the predicted cost regimes together with the underlying discretisation error and variance estimates.

1. Introduction

The introduction frames SGLD’s computational savings against the variance introduced by subsampling, making its net cost at a prescribed accuracy uncertain. The paper sharpens the comparison with EM by analyzing how cost depends on dataset size, subsample size, and accuracy, then checks its theoretical regimes numerically.

1. Introduction

  • Subsampling reduces computational cost but increases variance, so whether SGLD lowers the cost of achieving a prescribed mean-square accuracy is problem-dependent.Prior examples report either SGLD or ULA as the better choice [21].Source: §1 · p. 1
  • The paper identifies regimes in m, s, and ε where SGLD or ULA is preferable and confirms its theoretical findings with Bayesian-inference experiments.Its stated goal is a more detailed cost comparison than [21].Source: §1 · p. 2
  • The analysis is limited to the basic SGLD model, although stochastic gradients are also used with kinetic Langevin, Hamiltonian Monte Carlo, tamed ULA, and multilevel Monte Carlo.The authors focus on the basic case because its cost analysis is already challenging.Source: §1 · p. 2
  • The paper studies finite-sum-drift SDEs using standard EM discretisations and SGLD counterparts that evaluate the drift from a subsample rather than all m data points.For Bayesian inference, the finite-sum setting arises naturally from data likelihood contributions.Source: §1 · pp. 2–3
  • Rescaling the drift by 1/m and the noise by 1/√m preserves the invariant distribution while allowing model parameters and step sizes independent of m.The paper says this choice yields cleaner analysis and more accurate cost-regime formulas.Source: Comparison to · p. 3
  • The analysis finds that SGLD cost depends heavily on mini-batch size, correcting [21]’s conclusion that it is independent of that size.The authors attribute the earlier conclusion to bounds that were not sharp enough.Source: Comparison to · p. 3

Paper structure.

  • The paper proceeds from assumptions and cost results in Section 2 to numerical verification in Section 3, theoretical analysis in Section 4, and further variance experiments in Section 5.Proofs for auxiliary results also appear in Appendices A and B.Source: Paper structure · p. 3

2. Main results

For finite-sum drifts, the paper compares the mean-square-accuracy cost of EM and SGLD and identifies parameter regimes where either method is cheaper. EM can be preferable only for small datasets with aggressive subsampling; when s ≪ m, the transition is near m = ε^-1.

2.1. Preliminaries and assumptions.

  • The finite-sum drift is scaled by 1/m so its Lipschitz and dissipativity constants can remain independent of dataset size m.Source: §2.1 · p. 4

2.2. Numerical schemes.

  • EM discretizes the target SDE on a fixed time grid with step size h.Source: §2.2 · p. 5

2.2.1. Subsampling.

  • SGLD uses a subsampled drift estimator without replacement, with s denoting the mini-batch size.Source: §2.2.1 · p. 6

2.3. Results on computational cost.

  • For prescribed mean-square accuracy ε², the analysis selects estimator parameters and derives total-cost dependence on ε, m, and, for SGLD, s.Source: §2.3 · pp. 6–7
  • The leading-order cost comparison depends on m, ε^-2, and (m-s)/(ms), so neither method dominates uniformly across the parameter space.Source: §2.3 · p. 8
  • SGLD has lower leading-order cost over most of the parameter space; EM can be cheaper only for small m with aggressive subsampling, with the s ≪ m transition near m = ε^-1.Source: §2.3 · pp. 9–10

3. Numerical verification of computational cost

The experiments use a Gaussian Bayesian benchmark and an empirical cost protocol to test the predicted regimes. SGLD is generally cheaper for large datasets and small batches, while EM can remain competitive for smaller datasets or near-full batches.

3.1. The numerical set-up.

  • The Gaussian hierarchical model supplies a dataset-dependent Gaussian posterior as the invariant target, with π(f) available by high-accuracy numerical quadrature.Rescaling preserves the invariant distribution while keeping the drift constants independent of m.Source: §3.1 · pp. 11–13

3.2. Cost estimate.

  • The cost protocol makes finite-time bias negligible, calibrates the step size by comparing h with h/2, and estimates path variance before choosing Monte Carlo sample counts.Its per-step cost is m for EM and s for SGLD.Source: §3.2 · p. 13

3.2.1. Non-linear functional: Lipschitz case.

  • For Lipschitz observables f1 and f2, winner maps show SGLD cheaper over most of the fixed-batch test window, though EM wins some lower-m and transition-region configurations.Figure 3 varies m and s jointly and shows SGLD generally favoured when s/m is small; EM becomes increasingly competitive as s approaches m.Source: §3.2.1 · p. 14, Figure 3 · p. 16
  • SGLD is substantially cheaper for large m and small s/m, while EM can compete at smaller m or as s approaches m.The same qualitative pattern appears for the non-Lipschitz observable f3, although the theory does not cover that case.Source: §3.2.1 · p. 14, §3.2.2 · p. 17, Table 1 · p. 17

3.2.2. Non-linear functional: non-Lipschitz case.

  • The non-Lipschitz f3 experiment follows the same qualitative cost pattern, but a general theoretical analysis for non-Lipschitz observables is outside the paper’s scope.The numerical setting matches the one used for the Lipschitz functions.Source: §3.2.2 · pp. 16–17

4. Theoretical analysis and proofs

The analysis builds time-uniform bounds for the exact SDE, EM, and SGLD, then uses bias–variance error decompositions to establish computational cost regimes. Its results rely on stated assumptions and suitable step-size conditions.

4.1. Properties of exact solution.

  • The exact SDE admits an invariant probability measure and has a bounded second moment under Assumptions 2.1 and 2.2.An improved residual-term bound is also obtained under the additional Assumption 2.4.Source: §4.1 · pp. 17–18

4.2. Analysis of Euler-Maruyama method.

  • For EM, the analysis establishes a step-size-conditioned mean-square error bound and a variance bound for Lipschitz functionals.The error estimate compares the discretisation with the exact solution uniformly over time.Source: §4.2 · pp. 19–20

4.3. Analysis of SGLD.

  • For SGLD with subsampling without replacement, the analysis bounds subsampling error, the mean-square difference from EM, and variance for Lipschitz functionals.The section first derives a second-moment identity for the subsampling estimator and a corresponding variance bound.Source: §4.3 · pp. 21, 23

4.4. Proofs for the cost results.

  • The cost proofs decompose mean-square error into bias and variance contributions, balance them against accuracy, and establish four leading-order cost regimes.The resulting comparison is not a single uniform dominance result across parameter settings.Source: §4.4 · pp. 23, 25, 26

5. Additional numerical experiments

Additional experiments test the predicted convergence and variance estimates. They show different strong convergence orders for EM and SGLD, and illustrate how variance varies with dataset size, batch size, and step size.

5.1. Illustration of Theorem 4.5 and Theorem 4.11.

  • EM exhibits strong order one, while SGLD has strong order one-half when s < m in the tested setting.The experiment uses m ∈ {32, 64} and 10^3 sample paths.Source: §5.1 · p. 27, Figure 4 · p. 27

5.2. Illustration of Lemma 4.6 and Theorem 4.12.

  • The single-path variance decreases roughly at order one with m across the tested functions and methods.The Euler variance experiment fixes h = 2^-11 and s = 16.Source: §5.2 · p. 27, Figure 5 · p. 28
  • SGLD variance decreases as s increases; its step-size dependence becomes more visible for larger m, while EM variance is only weakly affected by step size.For SGLD, the h/s contribution can dominate the m^-1 diffusion contribution at large m, and variance approaches an h-independent floor when subsampling contribution becomes negligible.Source: §5.2 · p. 28, Figure 6 · p. 28, Figure 7 · p. 29

Appendix A. Proofs in Section 4.1

The appendix proves contraction and invariant-measure properties for the SDE, together with moment and discretisation-residual bounds used in the main analysis.

Appendix A. Proofs in Section 4.1

  • Solutions driven by the same Brownian motion contract exponentially: their distance is bounded by exp(-Kt) times the initial distance.This contraction supports the invariant-measure argument.Source: Appendix A · pp. 31, 34
  • The contraction argument yields a unique invariant measure, which is invariant for the whole transition semigroup and has finite second moment.The proof uses a fixed-point argument and the moment bound.Source: Appendix A · p. 34
  • The appendix bounds the residual between the exact solution and its discretisation, completing the estimate used in the main error analysis.The residual bound follows from the drift condition and moment bounds.Source: Appendix A · pp. 36, 38

Appendix B. Proofs in Section 4.3

The appendix proof fragments establish squared-error recurrences for a one-step quantity and for the separation between two independent SLGD chains. The latter proof uses stated assumptions and earlier results.

Appendix B. Proofs in Section 4.3

  • The Lemma 4.7 recurrence bounds the next squared error by a contraction term and an additive term involving the step size, dataset size, mini-batch size, and L.Source: Appendix B · p. 40
  • Theorem 4.12 bounds the next-step squared separation of two independent SLGD chains by a contraction term plus noise and sampling contributions.The proof tracks ηk, the difference between the chains.Source: Appendix B · p. 40
  • The cited derivation uses Assumption 2.2 for the second inequality, Eqn. (2.1) for the third, and Corollary 4.10 for the last line.Source: Appendix B · p. 40

Appendix C. The derivation of (21)

The appendix derives posterior distribution (21) using Bayes’ formula, absorbing terms independent of θ into the normalizing constant. It also identifies the posterior variance.

Appendix C. The derivation of (21)

  • The derivation applies Bayes’ formula to y1:m to obtain a posterior density for θ from the prior density.Source: Appendix C · p. 40
  • Terms independent of θ are absorbed into the normalizing constant when deriving the posterior density.Source: Appendix C · p. 41
  • The posterior variance is stated, and the resulting distribution is identified as distribution (21).Source: Appendix C · p. 41

Related papers

Loading 2609.17750v1…