PaperlayerView exact revision on arXiv

Source-linked AI summary

Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning

Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, Andrew Gordon Wilson

arXiv:1902.03932v2cs.LGcs.AIcs.CVstat.MEstat.ML

TL;DR

Modern Bayesian deep learning needs scalable inference for high-dimensional, multimodal neural-network posteriors, while traditional SG-MCMC can mix poorly in limited time. The paper develops cyclical SG-MCMC, alternating large-step exploration with small-step sampling, and reports improved mode exploration, uncertainty estimation, theoretical convergence support, and large-scale results. Its main scope boundary is that the theoretical guarantee relies on standard assumptions, with stricter convex-case improvements not pursued because of their mismatch with real applications.

  • Problem

    MCMC methods are underused for modern Bayesian deep learning, where multimodal posteriors and limited-time estimation make effective exploration difficult.

  • Method

    Cyclical SG-MCMC alternates large-stepsize exploration with small-stepsize sampling, using warm restarts to explore modes and characterize local densities.

  • Results

    Cyclical SG-MCMC explores complex multimodal and unimodal distributions more efficiently, improves uncertainty estimation, and is supported by non-asymptotic convergence analysis.

  • Takeaways & Limitations

    Cyclical SG-MCMC offers a practical route toward MCMC-based inference in modern Bayesian deep learning, including scalable experiments on challenging datasets.

  • Takeaways & Limitations

    The theoretical bound uses standard assumptions, while stricter convex-case improvements are not considered because of their discrepancy from real applications.

Abstract

from arXiv · show

The posteriors over neural network weights are high dimensional and multimodal. Each mode typically characterizes a meaningfully different representation of the data. We develop Cyclical Stochastic Gradient MCMC (SG-MCMC) to automatically explore such distributions. In particular, we propose a cyclical stepsize schedule, where larger steps discover new modes, and smaller steps characterize each mode. We also prove non-asymptotic convergence of our proposed algorithm. Moreover, we provide extensive experimental results, including ImageNet, to demonstrate the scalability and effectiveness of cyclical SG-MCMC in learning complex multimodal distributions, especially for fully Bayesian inference with modern deep neural networks.

1 INTRODUCTION

Bayesian inference offers uncertainty-aware alternatives for deep-network training, but MCMC methods remain underused because modern neural networks create challenging multimodal posteriors. The paper proposes cyclical SG-MCMC to improve exploration, supports it theoretically, and evaluates it on large-scale Bayesian deep-learning problems.

  • Bayesian methods account for uncertainty in neural-network weight space while balancing model complexity and data fitting.
  • MCMC methods remain largely unused for modern deep neural networks despite their earlier strong performance with smaller networks.
  • Existing SG-MCMC methods face bounded estimation error under decreasing stepsizes and practical demands for effective inference in short training times.
  • The proposed cSG-MCMC uses warm restarts to combine local-density characterization with exploration of multimodal posterior distributions.
  • cSG-MCMC is reported to provide better limited-time mixing than traditional SG-MCMC and approximate parallel MCMC at a fraction of its cost.
  • The paper proves non-asymptotic convergence and reports experiments on Bayesian neural networks, uncertainty estimation, and challenging datasets including ImageNet.

2 PRELIMINARIES: SG-MCMC WITH A DECREASING STEPSIZE SCHEDULE

SG-MCMC scales Bayesian inference through minibatch-based stochastic gradients, but asymptotic consistency traditionally requires decreasing stepsizes. Standard schedules reduce numerical bias over time, while limited-time inference remains subject to approximation error.

  • SG-MCMC methods use minibatches to approximate the full-data potential energy when evaluating all datapoints is too expensive.
  • SGLD & SGHMC: SGLD updates posterior samples with stochastic gradients and Gaussian noise controlled by the stepsize.
  • SGLD & SGHMC: SGHMC adds an auxiliary momentum variable and friction term to improve mixing while counteracting minibatch noise.
  • Asymptotic consistency requires decreasing stepsizes satisfying the stated summability conditions.
  • A typical schedule is α_k = a(b + k)^−γ with γ ∈ (0.5, 1] and positive constants a and b.

3 CYCLICAL SG-MCMC

cSG-MCMC alternates large-step exploration with small-step sampling in repeated cycles, using warm restarts to traverse multimodal posteriors while preserving local characterization. The method is motivated as a low-cost approximation to parallel MCMC and is paired with theoretical and practical safeguards.

  • cSG-MCMC consists of exploration and sampling stages repeated across cycles.
  • The cyclical schedule starts each period with a large stepsize for aggressive exploration, then decreases it to examine local posterior regions.
  • Each new cycle restarts with a large stepsize, encouraging escape from the current mode and discovery of another posterior region.
  • 3.1 EXPLORATION: The method avoids collecting samples at overly large stepsizes because minibatch noise can distort the stationary distribution and make stochastic MH correction costly.
  • 3.1 EXPLORATION: Exploration uses optimization-like behavior, while sampling begins at T = 1 when the cycle enters its later phase.
  • The schedule is presented as combining large-step perturbation, faster within-cycle decay, and optimization’s super-convergence advantage.
  • 3.2 SAMPLING: Compared with parallel MCMC, cSG-MCMC uses one training pass and has the cost of a single chain while estimating different posterior regions across cycles.
  • 3.2 SAMPLING: Repeated cycles collect samples from multiple modes, with a separate scheme provided for combining them during testing.

4 THEORETICAL ANALYSIS

The analysis establishes non-asymptotic convergence guarantees for cyclical SG-MCMC using weak convergence and 2-Wasserstein distance, while separating exact-process convergence from numerical error.

  • Convergence analysis: The analysis studies cSG-MCMC convergence through both test-function averages and Wasserstein distance.These provide complementary views of convergence behavior.
  • Weak convergence: For cSGLD, Theorem 1 bounds the bias and MSE of posterior-average estimates for smooth test functions.The estimates use generated samples to approximate the posterior average.
  • Convergence under the Wasserstein distance: The 2-Wasserstein analysis compares the algorithm’s sample distribution µK with the stationary distribution ν∞, which matches the target distribution.The Wasserstein distance is defined using joint distributions whose marginals are the compared distributions.
  • Convergence under the Wasserstein distance: Theorem 2 provides a convergence bound for cSGLD with a cyclical stepsize sequence at iterations satisfying K mod M = 0.The bound uses constants independent of the stepsizes.
  • Theoretical interpretation: Under α0 = O(K^-β) for every β > 1, the analysis gives a stated specialized Wasserstein bound.The paper notes that the overall dependency on K has the same order as standard SGLD.
  • Theoretical interpretation: The convergence bound separates exact-process convergence to stationarity from numerical discretization error.The first term concerns convergence of the exact solution, while the second concerns the gap between the algorithm and the discretized process.

5 EXPERIMENTS

Experiments evaluate cSG-MCMC on synthetic multimodal sampling, Bayesian neural-network classification, ImageNet, and uncertainty estimation. Across these settings, cyclical schedules improve mode exploration, sample diversity, predictive performance, and uncertainty estimates relative to traditional approaches.

  • Synthetic Multimodal Data: cSGLD fully explored all 25 Gaussian modes, whereas traditional SGLD discovered only 4 using 200k samples.The cyclical schedule uses large steps to discover modes and small steps to characterize local modes.
  • Bayesian Neural Networks: cSG-MCMC significantly outperformed Snapshot ensembles and traditional SG-MCMC on image-classification test error.The comparison attributes the advantage to combining diverse-mode coverage with fine-scale characterization within modes.
  • Bayesian Neural Networks: On CIFAR-100, cSG-MCMC samples formed 4 weight-space clusters, while traditional SG-MCMC samples formed one cluster.Interpolation spikes between cSG-MCMC samples further indicate that neighboring samples come from different modes.
  • Bayesian Neural Networks: A cyclical single chain performed best with half the computational cost of the compared parallel decreasing-schedule setup.The experiments identify warm restarts from the cyclical schedule as important for this result.
  • Bayesian Neural Networks: At a fixed number of cycles, cSG-MCMC achieved substantially lower test errors than Snapshot optimization.Snapshot provides one minimum per cycle, whereas cSG-MCMC collects more samples to characterize each mode.
  • Uncertainty Evaluation: cSGHMC obtained the lowest ImageNet testing NLL, while cSGLD and cSGHMC assigned lower probability to low-entropy predictions on notMNIST.The notMNIST result is linked to exploring multiple modes that provide differing predictions for out-of-distribution inputs.

6 DISCUSSION

The discussion presents cyclical SG-MCMC as an approach for exploring multimodal and unimodal posteriors, with empirical evidence of broader mode coverage, improved mixing, and more accurate uncertainty estimation.

  • Discussion: Cyclical SG-MCMC explores complex multimodal distributions and captures diverse hypotheses corresponding to different parameter settings.The paper also reports more efficient exploration of unimodal distributions and improved uncertainty estimation.
  • Discussion: 50K samples let cSGLD explore significantly more of a 25-Gaussian mixture, whereas traditional SGLD discovered only one mode.The comparison uses the non-parallel setting shown in Figure 5.
  • Discussion: Mode coverage is measured by counting a mode when at least 100 samples fall within radius 0.25 of its center, using 10 runs.This metric is used to quantify multimodal exploration.
  • Discussion: In Bayesian logistic regression, cSG-MCMC dramatically outperforms SG-MCMC overall, while SGHMC and cSGHMC match on the simple Heart posterior.Higher-dimensional Australian and German settings show significantly higher ESS for cSG-MCMC.
  • Discussion: The analysis treats ψ as a Poisson-equation solution functional and assumes bounded derivatives together with dissipativity and gradient-noise conditions.These assumptions support the theoretical convergence analysis.

C PROOF OF THEOREM 1

The proof of Theorem 1 specializes a general decreasing-step-size SG-MCMC result to cSGLD. It also clarifies that asymptotic unbiasedness requires the decreasing-step-size assumption, which cSGLD does not satisfy and is therefore biased.

  • Proof of Theorem 1: The general decreasing-step-size analysis assumes positive decreasing stepsizes whose sum diverges and whose limiting behavior satisfies the stated convergence conditions.These are the assumptions used for the general SG-MCMC theorem.
  • Proof of Theorem 1: The proof bounds the bias and mean-squared error by decomposing discretization and stochastic-gradient contributions.The supplied derivation states that one higher-order gradient-noise term is omitted from big-O notation.
  • Proof of Theorem 1: Assumption 4 is needed only for asymptotic unbiasedness; without it, the resulting algorithm can remain biased, as in cSGLD.The lemma still applies without Assumption 4.
  • Proof of Theorem 1: The proof derives Theorem 1 as a special case of a general SG-MCMC lemma for a first-order cSGLD integrator.The specialization sets the integrator order to N = 1.

D PROOF OF THEOREM 2

The proof of Theorem 2 analyzes cSGLD through continuous-time interpolation and Wasserstein-2 distance to the stationary posterior. It establishes a convergence bound under standard regularity, dissipativity, and gradient-noise assumptions.

  • Proof of Theorem 2: The proof introduces an SDE whose stationary distribution is p(θ|D), then relates the discrete iterates to continuous-time distributions.The argument uses µk and νt for the corresponding discrete and continuous distributions.
  • Proof of Theorem 2: The analysis assumes unbiased stochastic gradients, bounded update-noise behavior, independence of updates, and additional regularity conditions.These assumptions are stated alongside the continuous-time interpolation and SDE analysis.
  • Proof of Theorem 2: The proof bounds intermediate relative-entropy and moment terms before applying transport and auxiliary lemmas to obtain the Wasserstein convergence result.The final bound follows after combining these components and using monotonicity in the accumulated stepsize.
  • Proof of Theorem 2: The final derivation concludes that the relevant Wasserstein bound decreases with K under the stated assumptions.This is the convergence behavior used in the theorem.

E RELATION WITH SGLD

The comparison argues that cyclical SGLD converges faster than polynomially decayed SGLD under matched diffusion-time considerations. The broader method combines samples from cycles by estimating their relative weights for posterior expectations.

  • Relation with SGLD: Standard SGLD uses a polynomially decaying stepsize, while cSGLD uses cyclical schedules with repeated exploration and sampling phases.The theoretical comparison contrasts their accumulated diffusion time and approximation errors.
  • Relation with SGLD: When α0 ≥ h0, cSGLD accumulates diffusion time faster than standard SGLD and therefore incurs less error for K iterations.The paper states that this yields faster convergence for cSGLD.
  • Relation with SGLD: Increasing standard SGLD’s initial stepsize to match cSGLD’s speed makes its approximation error much larger, again favoring cSGLD.The comparison reports h0 scaling as O(α0K/log K) and a larger error term for standard SGLD.
  • Relation with SGLD: Cyclical SG-MCMC combines samples from cycle-specific sub-posteriors to estimate the full posterior expectation.Each cycle characterizes a region Θm, and cycle weights are estimated with a harmonic-mean method.
  • Relation with SGLD: The cycle-combination estimator can be adapted to overlapping regions by replacing them with disjoint sub-regions.This preserves the estimator when the original cycle regions overlap.
  • Relation with SGLD: The analysis assumes unbiased noisy gradients and independence between fresh update noise and previous gradient-noise variables.These assumptions support the stochastic-gradient convergence derivation.

G.1 THEOREM

The section derives non-asymptotic convergence bounds for the cyclical stepsize schedule under two cases determined by mαmin + Mα0. It also notes practical hyperparameter trade-offs and a possible future application to Wasserstein gradient flows.

  • Theorem statement: The convergence bounds use positive constants M, m, δ, and σ defined in Assumption 5.The supplied theorem excerpt identifies these constants but does not reproduce the full bounds.
  • Stepsize conditions: Under the stated stepsize construction, αk varies from α0 toward αmin while remaining within the required stability range.The proof states 0 < αk < min(2/M, 1/m) for every k > 0.
  • Convergence proof: The convergence proof applies Lemma 2 after bounding the contraction factors ρk across the cyclical stepsize sequence.The proof first establishes 0 < ρk < 1, then uses the lemma to complete the convergence argument.
  • Case analysis: When mαmin + Mα0 ≤ 2, the proof bounds ρk by 1 − mαmin for every k > 0.This case substitutes the bound into the preceding convergence inequality.
  • Case analysis: When mαmin + Mα0 > 2, the proof instead bounds ρk by Mα0 − 1 for every k > 0.This alternative bound yields the second convergence case.
  • Practical settings: The training budget creates a trade-off between the number of cycles M and cycle length, with β tuned so sampling begins after reaching a good region.For CIFAR-10, the text gives a 50-epoch cycle as an example and reports β should be selected by cross-validation.

J TEMPERING IN BAYESIAN NEURAL NETWORKS

The section studies temperature selection for Bayesian neural networks, motivated by tempering as a response to overparameterized models. With cSGLD, the best NLL and error occur at different temperatures on CIFAR-10 and CIFAR-100, while tempered sampling improves predictive performance.

  • Motivation: Tempering reflects the belief that model capacity is misspecified and shapes the posterior used for prediction.The paper frames tempering as combining model-capacity beliefs with data to form a predictive distribution.
  • Experimental setup: The main neural-network experiments use the prescribed pSGLD temperature T ≈ 0.0045, while the sensitivity study evaluates seven temperatures.The evaluated values are T ∈ [1, 0.5, 0.1, 0.05, 0.01, 0.005, 0].
  • Results: Sampling from the tempered posterior outperforms optimization in both NLL and accuracy.The paper interprets this as finding a better predictive distribution for overparameterized neural networks.
  • Temperature sensitivity: On CIFAR-10 with cSGLD, the best performance for both NLL and error is achieved at T = 0.1.Figure 6 reports NLL and error (%) as functions of temperature.
  • Temperature sensitivity: On CIFAR-100 with cSGLD, the best performance for both NLL and error is achieved at T = 0.01.Figure 7 reports NLL and error (%) as functions of temperature.
  • Experimental setup: The reported cSGLD and cSGHMC settings include M = 100 and β = 0.01 for the listed datasets.The passage gives separate initial-step parameters for the two cyclical samplers and dataset-specific settings.

K.3 UNCERTAINTY EVALUATION

The uncertainty-evaluation settings use four cycles with an exploration proportion of 0.8 for cyclical SG-MCMC.

  • Hyperparameter settings: For both cSG-MCMC and Snapshot, M = 4; cSG-MCMC uses β = 0.8.The passage also specifies α0N = 0.01 for cSGLD and 0.008 for cSGHMC.
Loading 1902.03932v2…