CYCLICAL STOCHASTIC GRADIENT MCMC FOR BAYESIAN DEEP LEARNING
Year: 2020

Abstract:
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 prove non-asymptotic convergence of our proposed algorithm. Moreover, we provide extensive experimental results, including ImageNet, to demonstrate the effectiveness of cyclical SG-MCMC in learning complex multimodal distributions, especially for fully Bayesian inference with modern deep neural networks.

heading: 1 INTRODUCTION
text: Deep neural networks are often trained with stochastic optimization methods such as stochastic gradient decent (SGD) and its variants. Bayesian methods provide a principled alternative, which account for model uncertainty in weight space (MacKay, 1992; Neal, 1996), and achieve an automatic balance between model complexity and data fitting. Indeed, Bayesian methods have been shown to improve the generalization performance of DNNs (Hernández-Lobato & Adams, 2015; Blundell et al., 2015; Li et al., 2016a; Maddox et al., 2019), while providing a principled representation of uncertainty on predictions which is crucial for decision making.
Approximate inference for Bayesian deep learning has typically focused on deterministic approaches, such as variational methods (Hernández-Lobato & Adams, 2015; Blundell et al., 2015). By contrast, MCMC methods are now essentially unused for inference with modern deep neural networks, despite previously providing the gold standard of performance with smaller neural networks (Neal, 1996). Stochastic gradient Markov Chain Monte Carlo (SG-MCMC) methods (Welling & Teh, 2011; Chen et al., 2014; Ding et al., 2014; Li et al., 2016a) provide a promising direction for a sampling based approach to inference in Bayesian deep learning. Indeed, it has been shown that stochastic methods, which use mini-batches of data, are crucial for finding weight parameters that provide good generalization in modern deep neural networks (Keskar et al., 2016).
However, SG-MCMC algorithms for inference with modern neural networks face several challenges: (i) In theory, SG-MCMC asymptotically converges to target distributions via a decreasing stepsize scheme, but suffers from a bounded estimation error in limited time (Teh et al., 2016; Chen et al., 2015). (ii) In practice, empirical successes have been reported by training DNNs in relatively short time (Li et al., 2016b; Chen et al., 2014; Gan et al., 2016; Neelakantan et al., 2016; Saatchi & Wilson, 2017). For example, Saatchi & Wilson (2017) apply SG-MCMC to generative adversarial networks (GANs) to solve the mode collapse problem and capture diverse generation styles. However, the loss surface for DNNs is highly multimodal (Auer et al., 1996; Choromanska et al., 2015). In order for MCMC to be effective for posterior inference in modern neural networks, a crucial question remains: how do we make SG-MCMC efficiently explore a highly multimodal parameter space given a practical computational budget?
Several attempts have been made to improve the sampling efficiency of SG-MCMC. Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) (Chen et al., 2014) introduces momentum to Langevin dynamics. Preconditioned stochastic gradient Langevin dynamics (pSGLD) (Li et al., 2016a) adaptively adjusts the sampler’s step size according to the local geometry of parameter space. Though simple and promising, these methods are still inefficient at exploring multimodal distributions in practice. It is our contention that this limitation arises from difficulties escaping local modes when using the small stepsizes that SG-MCMC methods typically require. Note that the stepsize in SGMCMC controls the sampler’s behavior in two ways: the magnitude to deterministically drift towards high density regions wrt. the current stochastic gradient, and the level of injecting noise to randomly explore the parameter space. Therefore, a small stepsize reduces both abilities, resulting in a large numbers of iterations for the sampler to move across the modes.
In this paper, we propose to replace the traditional decreasing stepsize schedule in SG-MCMC with a cyclical variant. To note the distinction from traditional SG-MCMC, we refer to this method as Cyclical SG-MCMC (cSG-MCMC). The comparison is illustrated in Figure 1. The blue curve is the traditional decay, while the red curve shows the proposed cyclical schedule. Cyclical SG-MCMC operates in two stages: (i) Exploration: when the stepsize is large (dashed red curves), we consider this stage as an effective burn-in mechanism, encouraging the sampler to take large
moves and leave the local mode using the stochastic gradient. (ii) Sampling: when the stepsize is small (solid red curves), the sampler explores one local mode. We collect samples for local distribution estimation during this stage. Further, we propose two practical techniques to improve estimation efficiency: (1) a system temperature for exploration and exploitation; (2) A weighted combination scheme for samples collected in different cycles to reflect their relative importance.
This procedure can be viewed as SG-MCMC with warm restarts: the exploration stage provides the warm restarts for its following sampling stage. cSG-MCMC combines the advantages from (1) the traditional SG-MCMC to characterize the fine-scale local density of a distribution and (2) the cyclical schedule in optimization to efficiently explore multimodal posterior distributions of the parameter space. In limited time, cSG-MCMC is a practical tool to provide significantly better mixing than the traditional SG-MCMC for complex distributions. cSG-MCMC can also be considered as an efficient approximation to parallel MCMC; cSG-MCMC can achieve similar performance to parallel MCMC with only a fraction of cost (reciprocal to the number of chains) that parallel MCMC requires.
To support our proposal, we also prove the non-asymptotic convergence for the cyclical schedule. We note that this is the first convergence analysis of a cyclical stepsize algorithm (including work in optimization). Moreover, we provide extensive experimental results to demonstrate the advantages of cSG-MCMC in sampling from multimodal distributions, including Bayesian neural networks and uncertainty estimation on several large and challenging datasets such as ImageNet.
In short, cSG-MCMC provides a simple and automatic approach to inference in modern Bayesian deep learning, with promising results, and theoretical support. This work is a step towards enabling MCMC approaches in Bayesian deep learning. We release code at https://github.com/ruqizhang/csgmcmc.

heading: 2 PRELIMINARIES: SG-MCMC WITH A DECREASING STEPSIZE SCHEDULE
text: SG-MCMC is a family of scalable sampling methods that enables inference with mini-batches of data. For a dataset D = {di}Ni=1 and a θ-parameterized model, we have the likelihood p(D|θ) and prior p(θ). The posterior distribution is p(θ|D) ∝ exp(−U(θ)) , where U(θ) is the potential energy given by U(θ) = − log p(D|θ)− log p(θ) .
When D is too large, it is expensive to evaluate U(θ) for all the data points at each iteration. Instead, SG-MCMC methods use a minibatch to approximate U(θ): Ũ(θ) = −N ′
N ∑N ′ i=1 log p(xi|θ)−
log p(θ) , where N ′ N is the size of minibatch. We recommend Ma et al. (2015) for a general review of SG-MCMC algorithms. We describe two SG-MCMC algorithms considered in this paper.
SGLD & SGHMC Welling & Teh (2011) proposed Stochastic Gradient Langevin Dynamics (SGLD), which uses stochastic gradients with Gaussian noise. Posterior samples are updated at the k-th step as: θk = θk−1−αk∇Ũ(θk) + √ 2αk k, where αk is the stepsize and k has a standard Gaussian distribution.
To improve mixing over SGLD, Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) (Chen et al., 2014) introduces an auxiliary momentum variable v. SGHMC is built upon HMC, with an additional friction term to counteract the noise introduced by a mini-batch. The update rule for posterior samples is: θk = θk−1 + vk−1, and vk = vk−1−αk∇Ũ(θk)− ηvk−1 + √ 2(η − γ̂)αk k, where 1− η is the momentum term and γ̂ is the estimate of the noise. To guarantee asymptotic consistency with the true distribution, SG-MCMC requires that the step sizes satisfy the following assumption:
Assumption 1. The step sizes {αk} are decreasing, i.e., 0 < αk+1 < αk, with 1) ∑∞ k=1 αk = ∞;
and 2) ∑∞ k=1 α 2 k <∞.
Without a decreasing step-size, the estimation error from numerical approximations is asymptotically biased. One typical decaying step-size schedule is αk = a(b + k)−γ , with γ ∈ (0.5, 1] and (a, b) some positive constants (Welling & Teh, 2011).

heading: 3 CYCLICAL SG-MCMC
text: We now introduce our cyclical SG-MCMC (cSG-MCMC) algorithm. cSG-MCMC consists of two stages: exploration and sampling. In the following, we first introduce the cyclical step-size schedule, and then describe the exploration stage in Section 3.1 and the sampling stage in Section 3.2. We propose an approach to combining samples for testing in Section F.
Assumption 1 guarantees the consistency of our estimation with the true distribution in the asymptotic time. The approximation error in limited time is characterized as the risk of an estimator R = B2 +V , whereB is the bias and V is the variance. In the case of infinite computation time, the traditional SG-MCMC setting can reduce the bias and variance to zero. However, the time budget is often limited in practice, and there is always a trade-off between bias and variance. We therefore decrease the overall approximation error R by reducing the variance through obtaining more effective samples. The effective sample size can be increased if fewer correlated samples from different distribution modes are collected.
For deep neural networks, the parameter space is highly multimodal. In practice, SG-MCMC with the traditional decreasing stepsize schedule becomes trapped in a local mode, though injecting noise may help the sampler to escape in the asymptotic regime (Zhang et al., 2017). Inspired to improve the exploration of the multimodal posteriors for deep neural networks, with a simple and automatic approach, we propose the cyclical cosine stepsize schedule for SG-MCMC. The stepsize at iteration k is defined as:
αk = α0 2
[ cos ( π mod(k − 1, dK/Me)
dK/Me
) + 1 ] , (1)
where α0 is the initial stepsize, M is the number of cycles and K is the number of total iterations (Loshchilov & Hutter, 2016; Huang et al., 2017).
The stepsize αk varies periodically with k. In each period, αk starts at α0, and gradually decreases to 0. Within one period, SG-MCMC starts with a large stepsize, resulting in aggressive exploration in the parameter space; as the stepsize is decreasing, SG-MCMC explores local regions. In the next period, the Markov chain restarts with a large stepsize, encouraging the sampler to escape from the current mode and explore a new area of the posterior.
Related work in optimization. In optimization, the cyclical cosine annealing stepsize schedule has been demonstrated to be able to find diverse solutions in multimodal objectives, though not
specifically different modes, using stochastic gradient methods (Loshchilov & Hutter, 2016; Huang et al., 2017; Garipov et al., 2018; Fu et al., 2019). Alternatively, we adopt the technique to SGMCMC as an effective scheme for sampling from multimodal distributions.

heading: 3.1 EXPLORATION
text: The first stage of cyclical SG-MCMC, exploration, discovers parameters near local modes of an objective function. Unfortunately, it is undesirable to directly apply the cyclical schedule in optimization to SG-MCMC for collecting samples at every step. SG-MCMC often requires a small stepsize in order to control the error induced by the noise from using a minibatch approximation. If the stepsize is too large, the stationary distribution of SG-MCMC might be far away from the true posterior distribution. To correct this error, it is possible to do stochastic Metropolis-Hastings (MH) (Korattikara et al., 2014; Bardenet et al., 2014; Chen et al., 2016b). However, stochastic MH correction is still computationally too expensive. Further, it is easy to get rejected with an aggressive large stepsize, and every rejection is a waste of gradient computations.
To alleviate this problem, we propose to introduce a system temperature T to control the sampler’s behaviour: p(θ|D) ∝ exp(−U(θ)/T ). Note that the setting T = 1 corresponds to sampling from the untempered posterior. When T → 0, the posterior distribution becomes a point mass. Sampling from limT→0 exp(−U(θ)/T ) is equivalent to minimizing U(θ); in this context, SG-MCMC methods become stochastic gradient optimization methods.
One may increase the temperature T from 0 to 1 when the step-size is decreasing. We simply consider T = 0 and perform optimization as the burn-in stage, when the completed proportion of a cycle r(k) = mod (k−1,dK/Me)dK/Me is smaller than a given threshold: r(k) < β. Note that β ∈ (0, 1) balances the proportion of the exploration and sampling stages in cSG-MCMC.
3.2 SAMPLING Algorithm 1 Cyclical SG-MCMC. Input: The initial stepsize α0, number of cycles M ,
number of training iterations K and the proportion of exploration stage β. for k = 1:K do α← αk according to Eq equation 1. if mod (k−1,dK/Me)dK/Me < β then
% Exploration stage θ ← θ − α∇Ũk(θ)
else % Sampling stage Collect samples using SG-MCMC methods
Output: Samples {θk}
The sampling stage corresponds to T = 1 of the exploration stage. When r(k) > β or step-sizes are sufficiently small, we initiate SG-MCMC updates and collect samples until this cycle ends.
SG-MCMC with Warm Restarts. One may consider the exploration stage as automatically providing warm restarts for the sampling stage. Exploration alleviates the inefficient mixing and inability to traverse the multimodal distributions of the traditional SG-MCMC methods. SG-MCMC with warm restarts explores different parts of the posterior distribution and captures multiple modes in a single training procedure.
In summary, the proposed cyclical SG-MCMC repeats the two stages, with three key advantages: (i) It restarts with a large stepsize at the beginning of a cycle which provides enough perturbation and encourages the model to escape from the current mode. (ii) The stepsize decreases more quickly inside one cycle than a traditional schedule, making the sampler better characterize the density of the local regions. (iii) This cyclical stepsize shares the advantage of the “super-convergence” property discussed in Smith & Topin (2017): cSG-MCMC can accelerate convergence for DNNs by up to an order of magnitude.
Connection to the Santa algorithm. It is interesting to note that our approach inverts steps of the Santa algorithm (Chen et al., 2016a) for optimization. Santa is a simulated-annealing-based optimization algorithm with an exploration stage when T = 1, then gradually anneals T → 0 in a refinement stage for global optimization. In contrast, our goal is to draw samples for multimodal distributions, thus we explore with T = 0 and sample with T = 1. Another fundamental difference is that Santa adopts the traditional stepsize decay, while we use the cyclical schedule.
We visually compare the difference between cyclical and traditional step size schedules (described in Section 2) in Figure 1. The cyclical SG-MCMC algorithm is presented in Algorithm 1.
Connection to Parallel MCMC. Running parallel Markov chains is a natural and effective way to draw samples from multimodal distributions (VanDerwerken & Schmidler, 2013; Ahn et al., 2014). However, the training cost increases linearly with the number of chains. Cyclical SG-MCMC can be seen as an efficient way to approximate parallel MCMC. Each cycle effectively estimates a different region of posterior. Note cyclical SG-MCMC runs along a single training pass. Therefore, its computational cost is the same as single chain SG-MCMC while significantly less than parallel MCMC.
Combining Samples. In cyclical SG-MCMC, we obtain samples from multiple modes of a posterior distribution by running the cyclical step size schedule for many periods. We provide a sampling combination scheme to effectively use the collected samples in Section F in the appendix.

heading: 4 THEORETICAL ANALYSIS
text: Our algorithm is based on the SDE characterizing the Langevin dynamics: dθt = −∇U(θt)dt +√ 2dWt , where Wt ∈ Rd is a d-dimensional Brownian motion. In this section, we prove nonasymptotic convergence rates for the proposed cSG-MCMC framework with a cyclical stepsize sequence {αk} defined in equation 1. For simplicity, we do not consider the exploration stage in the analysis as that corresponds to stochastic optimization. Generally, there are two different ways to describe the convergence behaviours of SG-MCMC. One characterizes the sample average over a particular test function (e.g., Chen et al. (2015); Vollmer et al. (2016)); the other is in terms of the Wasserstein distance (e.g., Raginsky et al. (2017); Xu et al. (2017)). We study both in the following.
Weak convergence Following Chen et al. (2015) and Vollmer et al. (2016), we define the posterior average of an ergodic SDE as: φ̄ , ∫ X φ(θ)ρ(θ)dθ for some test function φ(θ) of interest. For the corresponding algorithm with generated samples (θk)Kk=1, we use the sample average φ̂ defined as φ̂ = 1K ∑K k=1 φ(θk) to approximate φ̄. We prove weak convergence of cSGLD in terms of bias and MSE, as stated in Theorem 1.
Theorem 1. Under Assumptions 2 in the appendix, for a smooth test function φ, the bias and MSE of cSGLD are bounded as:
BIAS: ∣∣∣Eφ̃− φ̄∣∣∣ = O( 1
α0K + α0
) , MSE: E ( φ̃− φ̄ )2 = O ( 1
α0K + α20
) . (2)
Convergence under the Wasserstein distance Next, we consider the more general case of SGLD and characterize convergence rates in terms of a stronger metric of 2-Wasserstein distance, defined as:
W 22 (µ, ν) := inf γ {∫ Ω×Ω ‖θ − θ′‖22dγ(θ, θ′) : γ ∈ Γ(µ, ν) } where Γ(µ, ν) is the set of joint distributions over (θ, θ′) such that the two marginals equal µ and ν, respectively.
Denote the distribution of θt in the SDE as νt. According to Chiang & Hwang (1987), the stationary distribution ν∞ matches our target distribution. Let µK be the distribution of the sample from our proposed cSGLD algorithm at the K-th iteration. Our goal is to derive a convergence bound on W2(µK , ν∞). We adopt standard assumptions as in most existing work, which are detailed in Assumption 3 in the appendix. Theorem 2 summarizes our main theoretical result.
Theorem 2. Under Assumption 3 in the appendix, there exist constants (C0, C1, C2, C3) independent of the stepsizes such that the convergence rate of our proposed cSGLD with cyclical stepsize sequence equation 1 is bounded for all K satisfying (K mod M =0), as W2(µK , ν∞) ≤
C3 exp(− Kα0 2C4 ) +
( 6 +
C2Kα0 2
) 1 2
[(C1 3α20K
8 + σC0 Kα0 2 ) 1 2 + (C1
3α20K
16 + σC0 Kα0 4 ) 1 4 ] .
Particularly, if we further assume α0 = O(K−β) for ∀β > 1, W2(µK , ν∞) ≤ C3 +( 6 + C2
Kβ−1
) 1 2 [( 2C1
K2β−1 + 2C0 Kβ−1 ) 1 2 + ( C1 K2β−1 + C0 Kβ−1 ) 1 4 ].
Remark 1. i) The bound is decomposed into two parts: the first part measures convergence speed of exact solution to the stationary distribution, i.e., ν∑
k αk to ν∞; the second part measures the
numerical error, i.e., between µK and ν∑ k αk
. ii) The overall bound offers a same order of dependency on K as in standard SGLD (please see the bound for SGLD in Section E of the appendix. See also Raginsky et al. (2017)). iii) If one imposes stricter assumptions such as in the convex case, the bound can be further improved. Specific bounds are derived in the appendix. We did not consider this case due to the discrepancy from real applications.

heading: 5 EXPERIMENTS
text: We demonstrate cSG-MCMC on several tasks, including a synthetic multimodal distribution (Section 5.1), image classification on Bayesian neural networks (Section 5.2) and uncertainty estimation in Section 5.3. We also demonstrate cSG-MCMC can improve the estimate efficiency for uni-modal distributions using Bayesian logistic regression in Section A.2 in the appendix. We choose SLGD and SGHMC as the representative baseline algorithms. Their cyclical counterpart are called cSGLD and cSGHMC, respectively.
5.1 SYNTHETIC MULTIMODAL DATA
We first demonstrate the ability of cSG-MCMC for sampling from a multi-modal distribution on a 2D mixture of 25 Gaussians. Specifically, we compare cSGLD with SGLD in two setting: (1) parallel running with 4 chains and (2) running with a single chain, respectively. Each chain runs for 50k iterations. The stepsize schedule of SGLD is αk ∝ 0.05k−0.55. In cSGLD, we set M = 30 and the initial stepsize α0 = 0.09. The proportion of exploration stage β = 14 . Fig 2 shows the estimated density using sampling results for SGLD and cSGLD in the parallel setting. We observed that SGLD gets trapped in the local modes, depending on the initial position. In any practical time period,
SGLD could only characterize partial distribution. In contrast, cSGLD is able to find and characterize all modes, regardless of the initial position. cSGLD leverages large step sizes to discover a new mode, and small step sizes to explore local modes. This result suggests cSGLD can be a significantly favourable choice in the non-asymptotic setting, for example only 50k iterations in this case. The single chain results and the quantitative results on mode coverage are reported in Section A.1 of the appendix.

heading: 5.2 BAYESIAN NEURAL NETWORKS
text: We demonstrate the effectiveness of cSG-MCMC on Bayesian neural networks for classification on CIFAR-10 and CIFAR-100. We compare with (i) traditional SG-MCMC; (ii) traditional stochastic optimization methods, including stochastic gradient descent (SGD) and stochastic gradient descent
with momentum (SGDM); and (iii) Snapshot: a stochastic optimization ensemble method method with a the cyclical stepsize schedule (Huang et al., 2017). We use a ResNet-18 (He et al., 2016) and run all algorithms for 200 epochs. We report the test errors averaged over 3 runs, and the standard error (±) from the mean predictor. We set M = 4 and α0 = 0.5 for cSGLD, cSGHMC and Snapshot. The proportion hyper-parameter β =0.8 and 0.94 for CIFAR-10 and CIFAR-100, respectively. We collect 3 samples per cycle. In practice, we found that the collected samples share similarly high likelihood for DNNs, thus one may simply set the normalizing term wi in equation 33 to be the same for faster testing.
We found that tempering helps improve performance for Bayesian inference with neural networks. Tempering for SG-MCMC was first used by Li et al. (2016a) as a practical technique for neural network training for fast convergence in limited time1. We simply use the prescribed temperature of Li et al. (2016a) without tuning, but better results of the sampling methods can be achieved by tuning the temperature. More details are in Appendix J. We hypothesize that tempering helps due to the overparametrization of neural networks. Tempering enables one to leverage the inductive biases of the network, while representing the belief that the model capacity can be misspecified. In work on Safe Bayes, also known as generalized and fractional Bayesian inference, tempered posteriors are well-known to help under misspecification (e.g., Barron & Cover, 1991; de Heide et al., 2019; Grünwald et al., 2017).
For the traditional SG-MCMC methods, we found that noise injection early in training hurts convergence. To make these baselines as competitive as possible, we thus avoid noise injection for the first 150 epochs of training (corresponding to the zero temperature limit of SGLD and SGHMC), and resume SGMCMC as usual (with noise) for the last 50 epochs. This scheme is similar to the exploration and sampling stages within one cycle of cSG-MCMC. We collect 20 samples for the MCMC methods and average their predictions in testing.
Testing Performance for Image Classification We report the testing errors in Table 1 to compare with the non-parallel algorithms. Snapshot and traditional SG-MCMC reduce the testing errors on both datasets. Performance variance for these methods is also relatively small, due to the multiple networks in the Bayesian model average. Further, cSG-MCMC significantly outperforms Snapshot ensembles and the traditional SG-MCMC, demonstrating the importance of (1) capturing diverse modes compared to traditional SG-MCMC, and (2) capturing fine-scale characteristics of the distribution compared with Snapshot ensembles.
Diversity in Weight Space. To further demonstrate our hypothesis that with a limited budget cSGMCMC can find diverse modes, while traditional SG-MCMC cannot, we visualize the 12 samples we collect from cSG-MCMC and SG-MCMC on CIFAR-100 respectively using Multidimensional Scaling (MDS) in Figure 3 (a). MDS uses a Euclidean distance metric between the weight of samples. We see that the samples of cSG-MCMC form 4 clusters, which means they are from 4 different modes in weight space. However, all samples from SG-MCMC only form one cluster, which indicates traditional SG-MCMC gets trapped in one mode and only samples from that mode.
Diversity in Prediction. To further demonstrate the samples from different cycles of cSG-MCMC provide diverse predictions we choose one sample from each cycle and linearly interpolate between two of them (Goodfellow et al., 2014; Huang et al., 2017). Specifically, let J(θ) be the test error of a sample with parameter θ. We compute the test error of the convex combination of two samples J(λθ1 + (1− λ)θ2), where λ ∈ [0, 1]. We linearly interpolate between two samples from neighboring chains of cSG-MCMC since they are the most likely to be similar. We randomly select 4 samples from SG-MCMC. If the samples are from the same mode, the test error of the linear interpolation of parameters will be relatively
1https://github.com/ChunyuanLI/pSGLD/issues/2
smooth, while if the samples are from different modes, the test error of the parameter interpolation will have a spike when λ is between 0 and 1.
We show the results of interpolation for cSG-MCMC and SG-MCMC on CIFAR-100 in Figure 3 (b). We see a spike in the test error in each linear interpolation of parameters between two samples from neighboring chains in cSG-MCMC while the linear interpolation for samples of SG-MCMC is smooth. This result suggests that samples of cSG-MCMC from different chains are from different modes while samples of SG-MCMC are from the same mode.
Although the test error of a single sample of cSG-MCMC is worse than that of SG-MCMC shown in Figure 3 (c), the ensemble of these samples significantly improves the test error, indicating that samples from different modes provide different predictions and make mistakes on different data points. Thus these diverse samples can complement each other, resulting in a lower test error, and demonstrating the advantage of exploring diverse modes using cSG-MCMC.
Comparison to Parallel MCMC. cSG-MCMC can be viewed as an economical alternative to parallel MCMC. We verify how closely cSG-MCMC can approximate the performance of parallel MCMC, but with more convenience and less computational expense. We also note that we can improve parallel MCMC with the proposed cyclical stepsize schedule.
We report the testing errors in Table 2 to compare multiple-chain results. (1) Four chains used, each runs 200 epochs (800 epochs in total), the results are shown in the first 4 columns (Cyclical+Parallel vs Decreasing+Parallel). We see that cSG-MCMC variants provide lower errors than plain SG-MCMC. (2) We reduce the number of epochs (#epoch) of parallel MCMC to 100 epoch each for decreasing stepsize schedule. The total cost is 400 epochs. We compare its performance with cyclical single chain (200 epochs in total) in the last 4 columns (Decreasing+Parallel vs Cyclical+Single). We see that the cyclical schedule running on a single chain performs best even with half the computational cost! All the results indicate the importance of warm re-starts using the proposed cyclical schedule. For a given total cost budget, the proposed cSGMCMC is preferable to parallel sampling.
Comparison to Snapshot Optimization. We carefully compared with Snapshot, as our cSGMCMC can be viewed as the sampling counterpart of the Snapshot optimization method. We plot the test error wrt.various number of cycles M in Fig. 3. As M increases, cSG-MCMC and Snapshot both improve. However, given a fixed M , cSG-MCMC yields substantially lower test errors than Snapshot. This result is due to the ability of cSG-MCMC to better characterize the local distribution of modes: Snapshot provides a singe minimum per cycle, while cSG-MCMC fully exploits the mode with more samples, which could provide weight uncertainty estimate and avoid over-fitting.
Results on ImageNet. We further study different learning algorithms on a large-scale dataset, ImageNet. ResNet-50 is used as the architecture, and 120 epochs for each run. The results on the testing set are summarized in Table 3, including NLL, Top1 and Top5 accuracy (%), respectively. 3 cycles are considered for both cSGHMC and Snapshot, and we collect 3 samples per cycle. We see that cSGHMC yields the lowest testing NLL, indicating that the cy-
cle schedule is an effective technique to explore the parameter space, and diversified samples can help prevent over-fitting.
5.3 UNCERTAINTY EVALUATION
To demonstrate how predictive uncertainty benefits from exploring multiple modes in the posterior of neural network weights, we consider the task of uncertainty estimation for out-ofdistribution samples (Lakshminarayanan et al., 2017). We train a three-layer MLP model on the standard MNIST train dataset until convergence using different algorithms, and estimate the entropy of the predictive distribution on the notMNIST dataset (Bulatov, 2011). Since the samples from the notMNIST dataset belong to the unseen classes, ideally the predictive distribution of the trained model should be uniform over the notMNIST digits, which gives the maximum entropy.
In Figure 4, we plot the empirical CDF for the entropy of the predictive distributions on notMNIST. We see that the uncertainty estimates from cSGHMC and cSGLD are better than the other
methods, since the probability of a low entropy prediction is overall lower. cSG-MCMC algorithms explore more modes in the weight space, each mode characterizes a meaningfully different representation of MNIST data. When testing on the out-of-distribution dataset (notMNIST), each mode can provide different predictions over the label space, leading to more reasonable uncertainty estimates. Snapshot achieves less entropy than cSG-MCMC, since it represents each mode with a single point.
The traditional SG-MCMC methods also provide better uncertainty estimation compared to their optimization counterparts, because they characterize a local region of the parameter space, rather than a single point. cSG-MCMC can be regarded as a combination of these two worlds: a wide coverage of many modes in Snapshot, and fine-scale characterization of local regions in SG-MCMC.

heading: 6 DISCUSSION
text: We have proposed cyclical SG-MCMC methods to automatically explore complex multimodal distributions. Our approach is particularly compelling for Bayesian deep learning, which involves rich multimodal parameter posteriors corresponding to meaningfully different representations. We have also shown that our cyclical methods explore unimodal distributions more efficiently. These results are in accordance with theory we developed to show that cyclical SG-MCMC will converge faster to samples from a stationary distribution in general settings. Moreover, we show cyclical SG-MCMC methods provide more accurate uncertainty estimation, by capturing more diversity in the hypothesis space corresponding to settings of model parameters.
While MCMC was once the gold standard for inference with neural networks, it is now rarely used in modern deep learning. We hope that this paper will help renew interest in MCMC for posterior inference in deep learning. Indeed, MCMC is uniquely positioned to explore the rich multimodal posterior distributions of modern neural networks, which can lead to improved accuracy, reliability, and uncertainty representation.

heading: ACKNOWLEDGEMENTS
text: AGW was supported by an Amazon Research Award, Facebook Research, NSF I-DISRE 193471, NIH R01 DA048764-01A1, NSF IIS-1563887, and NSF IIS-1910266.

heading: A EXPERIMENTAL RESULTS

heading: A.1 SYNTHETIC MULTIMODAL DISTRIBUTION
text: The density of the distribution is
F (x) = 25∑ i=1 λN (x|µi,Σ),
where λ = 125 , µ = {−4,−2, 0, 2, 4} > × {−4,−2, 0, 2, 4}, Σ = [ 0.03 0 0 0.03 ] .
In Figure 5, we show the estimated density for SGLD and cSGLD in the non-parallel setting.
To quantitatively show the ability of different algorithms to explore multi-modal distributions, we define the mode-coverage metric: when the number of samples falling within the radius r of a mode center is larger than a threshold n̄, we consider this mode covered. On this dataset, we choose r = 0.25 and n̄ = 100. Table 4 shows the mode-coverage for several algorithms, based on 10 different runs.

heading: A.2 BAYESIAN LOGISTIC REGRESSION
text: We consider Bayesian logistic regression (BLR) on three real-world datasets from the UCI repository: Australian (15 covariates, 690 data points), German (25 covariates, 1000 data points) and Heart (14 covariates, 270 data points). For all experiments, we collect 5000 samples with 5000 burn-in iterations. Following the settings in Li et al. (2016a), we report median effective sample size (ESS) in Table 5.
Note that BLR is unimodal in parameter space. We use this experiment as an adversarial situation for cSG-MCMC, which we primarily designed to explore multiple modes. We note that even in the unimodal setting, cSG-MCMC more effectively explores the parameter space than popular alternatives. We can also use these experiments to understand how samplers respond to varying parameter dimensionality and training set sizes.
Overall, cSG-MCMC dramatically outperforms SG-MCMC, which demonstrates the fast mixing rate due to the warm restarts. On the small dataset Heart, SGHMC and cSGHMC achieve the
same results, because the posterior of BLR on this dataset is simple. However, in higher dimensional spaces (e.g., Australian and German), cSG-MCMC shows significantly higher ESS; this result means that each cycle in cSG-MCMC can characterize a different region of the posteriors, combining multiple cycles yields more accurate overall approximation.

heading: B ASSUMPTIONS

heading: B.1 ASSUMPTIONS IN WEAK CONVERGENCE ANALYSIS
text: In the analysis, we define a functional ψ that solves the following Poisson Equation:
Lψ(θk) = φ(θk)− φ̄, or equivalently, 1
K K∑ k=1 Lψ(θk) = φ̂− φ̄. (3)
The solution functional ψ(θk) characterizes the difference between φ(θk) and the posterior average φ̄ for every θk, thus would typically possess a unique solution, which is at least as smooth as φ under the elliptic or hypoelliptic settings (Mattingly et al., 2010). Following Chen et al. (2015); Vollmer et al. (2016), we make certain assumptions on the solution functional, ψ, of the Poisson equation equation 3. Assumption 2. ψ and its up to 3rd-order derivatives, Dkψ, are bounded by a function V , i.e., ‖Dkψ‖ ≤ HkVpk for k = (0, 1, 2, 3), Hk, pk > 0. Furthermore, the expectation of V on {θk} is bounded: supl EVp(θk) < ∞, and V is smooth such that sups∈(0,1) Vp (sθ + (1− s) θ′) ≤ C (Vp (θ) + Vp (θ′)), ∀θ, θ′, p ≤ max{2pk} for some C > 0.

heading: B.2 ASSUMPTIONS IN CONVERGENCE UNDER THE WASSERSTEIN DISTANCE
text: Following existing work in Raginsky et al. (2017), we adopt the following standard assumptions summarized in Assumption 3. Assumption 3. • There exists some constants A ≥ 0 and B ≥ 0, such that U(0) ≤ A and
∇U(0) ≤ B.
• The function U is LU -smooth : ‖∇U(w)−∇U(v)‖ ≤ LU‖w − v‖.
• The function U is (mu, b) − dissipative, which means for some mU > 0 and b > 0 〈w,∇U(w)〉 ≥mU‖w‖2 − b.
• There exists some constant δ ∈ [0, 1), such that E[‖∇Ũk(w) − ∇U(w)‖2] ≤ 2σ(M2U‖w‖2 +B2).
• We can choose µ0 which satisfies the requirement: κ0 := log ∫ e‖w‖ 2 µ0(w)dw <∞.

heading: C PROOF OF THEOREM 1
text: To prove the theorem, we borrow tools developed by Chen et al. (2015); Vollmer et al. (2015). We first rephrase the stepsize assumptions in general SG-MCMC in Assumption 4. Assumption 4. The algorithm adopts an N -th order integrator. The step sizes {hk} are such that 0 < hk+1 < hk, and satisfy 1) ∑∞ k=1 hk =∞; and 2) limK→∞ ∑K k=1 h N+1 k∑K
k=1 hk = 0.
Our prove can be derived by the following results from Chen et al. (2015). Lemma 1 (Chen et al. (2015)). Let SK , ∑K k=1 hk. Under Assumptions 2 and 4, for a smooth test function φ, the bias and MSE of a decreasing-step-size SG-MCMC with a N th-order integrator at time SL are bounded as:
BIAS: ∣∣∣Eφ̃− φ̄∣∣∣ = O( 1
SK +
∑K k=1 h N+1 k
SK
) (4)
MSE: E ( φ̃− φ̄ )2 ≤ C (∑ l h2k S2K E ‖∆Vl‖2 + 1 SK + ( ∑K k=1 h N+1 k ) 2 S2K ) . (5)
Note that Assumption 4 is only required if one wants to prove the asymptotically unbias of an algorithm. Lemma 1 still applies even if Assumption 4 is not satisfied. In this case one would obtain a biased algorithm, which is the case of cSGLD.
Proof of Theorem 1. Our results is actually a special case of Lemma 1. To see that, first note that our cSGLD adopts a first order integrator, thus N = 1. To proceed, note that SK = ∑K k=1 αk = O(α0K), and K−1∑ j=0 α2j+1 = α20 4 K−1∑ j=0 [cos( πmod(j − 1, [K/M ]) [K/M ] ) + 1]2
= α20 4 K−1∑ j=0 [cos2( πmod(j − 1,K/M) K/M ) + 1]2
= α20 4 K M ( M 2 +M) = 3α20K 8 . (6)
As a result, for the bias, we have∣∣∣Eφ̃− φ̄∣∣∣ = O( 1 SK + ∑K k=1 h N+1 k SK ) = O ( 1 α0K + 3α20K/8 α0K ) = O ( 1
α0K + α0
) .
For the MSE, note the first term ∑ l h2k S2K
E ‖∆Vl‖2 has a higher order than other terms, thus it is omitted in the big-O notation, i.e.,
E ( φ̃− φ̄ )2 = O ( 1
α0K + (
3α20K/8 α0K )2 )
= O
( 1
α0K + α20
) .
This completes the proof.

heading: D PROOF OF THEOREM 2
text: Proof of the bound for W2(µK , ν∞) in cSGLD. Firstly, we introduce the following SDE
dθt = −∇U(θt)dt+ √
2dWt , (7) Let νt denote the distribution of θt, and the stationary distribution of equation 34 be p(θ|D), which means ν∞ = p(θ|D).
θk+1 = θk −∇Ũk(θk)αk+1 + √
2αk+1ξk+1 (8) Further, let µk denote the distribution of θk.
Since W2(µK , ν∞) ≤W2(µK , ν∑K k=1 αk ) +W2(ν∑K k=1 αk , ν∞) (9)
, we need to give the bounds for these two parts respectively.
D.1 W2(µK , ν∑K k=1 αk )
For the first part, W2(µK , ν∑K k=1 αk ), our proof is based on the proof of Lemma 3.6 in Raginsky et al. (2017) with some modifications. We first assume E(∇Ũ(w)) = ∇U(w), ∀w ∈ Rd , which is a general assumption according to the way we choose the minibatch. And we define p(t) which will be used in the following proof:
p(t) = {k ∈ Z| k∑ i=1 αi ≤ t < k+1∑ i=1 αi} (10)
Then we focus on the following continuous-time interpolation of θk:
θ(t) =θ0 − ∫ t
0
∇Ũ θ(p(s)∑ k=1 αk) ds+√2 ∫ t 0 dW(d)s (11)
where ∇Ũ ≡ ∇Ũk for t ∈ [∑k i=1 αi, ∑k+1 i=1 αi ) . And for each k , θ( ∑k i=1 αi) and θk have the same probability law µk. Since θ(t) is not a Markov process, we define the following process which has the same one-time marginals as θ(t)
V (t) = θ0 − ∫ t
0
Gs (V (s)) ds+ √ 2 ∫ t 0 dW(d)s (12)
with
Gt(x) := E ∇Ũ θ(q(t)∑
i=1
αi)
 |θ(t) = x  (13)
Let PtV := L (V (s) : 0 ≤ s ≤ t) and P t θ := L (θ(s) : 0 ≤ s ≤ t) and according to the proof of Lemma 3.6 in Raginsky et al. (2017), we can derive a similar result for the relative entropy of PtV and Ptθ:
DKL(P t V ‖P t θ) =− ∫ dPtV log
dPtV dPtθ
= 1
4 ∫ t 0 E‖∇U(V (s))−Gs(V (s))‖2ds
= 1
4 ∫ t 0 E‖∇U(θ(s))−Gs(θ(s))‖2ds
The last line follows the fact that L(θ(s)) = L(V (s)), ∀s. Then we will let t = ∑K k=1 αk and we can use the martingale property of the integral to derive:
DKL(P ∑K k=1 αk V ‖P ∑K k=1 αk θ )
= 1
4 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖∇U(θ(s))−Gs(θ(s))‖2ds
≤ 1 2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖∇U(θ(s))−∇U(θ( q(s)∑ k=1 αi)‖2ds
+ 1
2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖∇U(θ( q(s)∑ k=1 αi)−Gs(θ( q(s)∑ k=1 αi)‖2ds
≤ L 2 U
2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖θ(s)− θ( q(s)∑ k=1 αi)‖2ds (14)
+ 1
2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖∇U(θ( q(s)∑ k=1 αi)−Gs(θ( q(s)∑ k=1 αi)‖2ds (15)
For the first part (14), we consider some s ∈ [ ∑j k=1 αk, ∑j+1 k=1 αk), for which the following holds:
θ(s)− θ( j∑
k=1
αk)
= −(s− j∑
k=1
αk)∇Ũk(θk) + √ 2(W(d)s −W
(d)∑j k=1 αk )
= −(s− j∑
k=1
αk)∇U(θk) + (s− j∑
k=1
αk)(∇U(θk)−∇Ũk(θk)) + √ 2(W(d)s −W
(d)∑j k=1 αk ) (16)
Thus, we can use Lemma 3.1 and 3.2 in Raginsky et al. (2017) for the following result:
E‖θ(s)− θ( j∑
k=1
αk)‖2 ≤ 3α2j+1E‖∇U(θj)‖2 + 3α2j+1E‖∇U(θj)−∇Ũj(θj)‖2 + 6αj+1d
≤ 12α2j+1(L2UE‖θj‖2 +B2) + 6αj+1d
Hence we can bound the first part, (choosing α0 ≤ 1),
L2U 2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖θ(s)− θ( q(s)∑ k=1 αi)‖2ds
≤ L 2 U
2 K−1∑ j=0 [ 12α3j+1(L 2 UE‖θj‖2 +B2) + 6α2j+1d ] ≤ L2U max
0≤j≤K−1
[ 6(L2UE‖θj‖2 +B2) + 3d ] ( K−1∑ j=0 α2j+1)
≤ L2U max 0≤j≤K−1
[ 6(L2UE‖θj‖2 +B2) + 3d ] 3α20K 8
(17)
The last line (17) follows from2 equation 6. The second part (15) can be bounded as follows:
1
2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖∇U(θ( q(s)∑ k=1 αi)−Gs(θ( q(s)∑ k=1 αi)‖2ds
= 1
2 K−1∑ j=0 αj+1E‖∇U(θj)−∇Ũ(θj)‖2
≤ σ max 0≤j≤K−1 (L2UE‖θj‖2 +B2) K−1∑ j=0 αj+1
≤ σ max 0≤j≤K−1 (L2UE‖θj‖2 +B2)( α0 2 K−1∑ j=0 (cos( πmod(j,K/M) K/M ) + 1))
≤ σ max 0≤j≤K−1
(L2UE‖θj‖2 +B2)( Kα0
2 )
Due to the data-processing inequality for the relative entropy, we have DKL(µK‖ν∑K
k=1 αk ) ≤ DKL(PtV ‖P t θ)
≤ L 2 U
2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖θ(s)− θ( q(s)∑ k=1 αi)‖2ds
+ 1
2 K−1∑ j=0 ∫ ∑j+1 k=1 αk∑j k=1 αk E‖∇U(θ( q(s)∑ k=1 αi)−Gs(θ( q(s)∑ k=1 αi)‖2ds
≤ L2U max 0≤j≤K−1
[ 6(L2UE‖θj‖2 +B2) + 3d ] 3α20K 8
+ σ max 0≤j≤K−1
(L2UE‖θj‖2 +B2)( Kα0
2 )
According to the proof of Lemma 3.2 in Raginsky et al. (2017), we can bound the term E‖θk‖2
E‖θk+1‖2 ≤ (1− 2αk+1mU + 4α2k+1M2U )E‖θk‖2 + 2αk+1b+ 4α2k+1B2 + 2αk+1d
β
Similar to the statement of Lemma 3.2 in Raginsky et al. (2017), we can fix α0 ∈ (0, 1 ∧ mU4M2U ). Then, we can know that
E‖θk+1‖2 ≤ (1− 2αminmU + 4α2minM2U )E‖θk‖2 + 2α0b+ 4α20B2 + 2α0d
β (18) , where αmin is defined as αmin , α02 [ cos ( π mod(dK/Me−1,dK/Me) dK/Me ) + 1 ] .
There are two cases to consider.
• If 1− 2αminmU + 4α2minM2U ≤ 0, then from equation 18 it follows that
E‖θk+1‖2 ≤ 2α0b+ 4α20B2 + 2α0d
β
≤ E‖θ0‖2 + 2(b+ 2B2 + d
β )
• If 0 ≤ 1− 2αminmU + 4α2minM2U ≤ 1, then iterating equation 18 gives
E‖θk‖2 ≤ (1− 2αminmU + 4α2minM2U )kE‖θ0‖2 + α0b+ 2α
2 0B 2 + α0dβ αminmU − 2α2minM2U
(19)
≤ E‖θ0‖2 + 2α0
mUαmin (b+ 2B2 +
d β ) (20)
2Note: we only focus on the case when K mod M = 0.
Now, we have
max 0≤j≤K−1
(L2UE‖θj‖2 +B2)
≤ (L2U (κ0 + 2(1 ∧ α0
mUαmin )(b+ 2B2 + d)) +B2) := C0
Due to the expression of α0αmin , C0 is independent of α0. Then we denote the 6L 2 U (C0 + d) as C1 and we can derive
DKL(µK‖ν∑K k=1 αk ) ≤ C1( 3α20K 8 ) + σC0( Kα0 2 )
Then according to Proposition 3.1 in Bolley & Villani (2005) and Lemma 3.3 in Raginsky et al. (2017), if we denote κ0 + 2b+ 2d as C2, we can derive the following result:
W2(µK , ν∑K k=1 αk ) ≤ (12 + C2( K∑ k=1 αk)) 1 2 · [DKL(µK‖ν∑K k=1 αk ) 1 2 +DKL(µK‖ν∑K k=1 αk ) 1 4 ]
≤ (12 + C2Kα0 2 ) 1 2 · [( 3C1α
2 0K
8 + KσC0α0 2 ) 1 2 + (
3C1α 2 0K
16 + KσC0α0 4 ) 1 4 ]
D.2 W2(ν∑K k=1 αk , ν∞)
We can directly get the following results from (3.17) in Raginsky et al. (2017) that there exist some positive constants (C3, C4),
W2(ν∑K k=1 αk , ν∞) ≤ C3 exp(− K∑ k=1 αk/C4)
Now combining the bounds for W2(µK , ν∑K k=1 αk ) and W2(ν∑K k=1 αk
, ν∞), substituting α0 = O(1/Kβ), and noting W2(ν∑K
k=1 αk , ν∞) decreases w.r.t. K, we arrive at the bound stated in the
theorem.

heading: E RELATION WITH SGLD
text: For the standard polynomially-decay-stepsize SGLD, the convergence rate is bounded as
W2(µ̃K , ν∞) ≤W2(µ̃K , ν∑K k=1 hk ) +W2(ν∑K k=1 hk , ν∞) (21)
where W2(µ̃K , ν∑K k=1 hk ) ≤ (6 + h0 ∑K k=1 1 k ) 1 2 ·
[(D1h 2 0 π2
6 + σD0h0 K∑ k=1 1 k ) 1 2 + (D1h 2 0 π2 16 + σD0 h0 2 K∑ k=1 1 k ) 1 4 ]
and W2(ν∑K k=1 hk , ν∞) ≤ C3 exp(− ∑K k=1 hk C4 ).
Proof of the bound of W2(µ̃K , ν∞) in the standard SGLD. Similar to the proof of W2(µK , ν∞) in cSGLD, we get the following update rule for SGLD with the stepsize following a polynomial decay i.e., hk = h0k ,
θk+1 = θk −∇Ũk(θk)hk+1 + √ 2hk+1ξk+1 (22)
Let µ̃k denote the distribution of θk.
Since
W2(µ̃K , ν∞) ≤W2(µ̃K , ν∑K k=1 hk ) +W2(ν∑K k=1 hk , ν∞) (23)
, we need to give the bounds for these two parts respectively.
E.1 W2(µ̃K , ν∑K k=1 hk )
We first assume E(∇Ũ(w)) = ∇U(w), ∀w ∈ Rd , which is a general assumption according to the way we choose the minibatch. Following the proof in Raginsky et al. (2017) and the analysis of the SPOS method in Zhang et al. (2018), we define the following p(t) which will be used in the following proof:
p(t) = {k ∈ Z| k∑ i=1 hi ≤ t < k+1∑ i=1 hi} (24)
Then we focus on the following continuous-time interpolation of θk:
θ(t) =θ0 − ∫ t
0
∇Ũ θ(p(s)∑ k=1 hk)  ds+√2∫ t 0 dW(d)s , (25)
(26)
where ∇Ũ ≡ ∇Ũk for t ∈ [∑k i=1 hi, ∑k+1 i=1 hi ) . And for each k , θ( ∑k i=1 hi) and θk have the same probability law µ̃k. Since θ(t) is not a Markov process, we define the following process which has the same one-time marginals as θ(t)
V (t) = θ0 − ∫ t
0
Gs (V (s)) ds+ √ 2 ∫ t 0 dW(d)s (27)
with
Gt(x) := E ∇Ũ θ(q(t)∑
i=1
hi)
 |θ(t) = x  (28)
Let PtV := L (V (s) : 0 ≤ s ≤ t) and P t θ := L (θ(s) : 0 ≤ s ≤ t) and according to the proof of Lemma 3.6 in Raginsky et al. (2017), we can derive the similar result for the relative entropy of PtV and Ptθ:
DKL(P t V ‖P t θ) =− ∫ dPtV log
dPtV dPtθ
= 1
4 ∫ t 0 E‖∇U(V (s))−Gs(V (s))‖2ds
= 1
4 ∫ t 0 E‖∇U(θ(s))−Gs(θ(s))‖2ds
The last line follows the fact that L(θ(s)) = L(V (s)), ∀s. Then we will let t = ∑K k=1 hk and we can use the martingale property of integral to derive:
DKL(P ∑K k=1 hk V ‖P ∑K k=1 hk θ )
= 1
4 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖∇U(θ(s))−Gs(θ(s))‖2ds
≤ 1 2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖∇U(θ(s))−∇U(θ( q(s)∑ k=1 hi)‖2ds
+ 1
2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖∇U(θ( q(s)∑ k=1 hi)−Gs(θ( q(s)∑ k=1 hi)‖2ds
≤ L 2 U
2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖θ(s)− θ( q(s)∑ k=1 hi)‖2ds (29)
+ 1
2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖∇U(θ( q(s)∑ k=1 hi)−Gs(θ( q(s)∑ k=1 hi)‖2ds (30)
For the first part (29), we consider some s ∈ [ ∑j k=1 hk, ∑j+1 k=1 hk), the following equation holds:
θ(s)− θ( j∑
k=1
hk)
= −(s− j∑
k=1
hk)∇Ũk(θk) + √ 2(W(d)s −W
(d)∑j k=1 hk )
= −(s− j∑
k=1
hk)∇U(θk) + (s− j∑
k=1
hk)(∇U(θk)−∇Ũk(θk)) + √ 2(W(d)s −W
(d)∑j k=1 hk ) (31)
Thus, we can use Lemma 3.1 and 3.2 in Raginsky et al. (2017) for the following result:
E‖θ(s)− θ( j∑
k=1
hk)‖2
≤ 3h2j+1E‖∇U(θj)‖2 + 3h2j+1E‖∇U(θj)−∇Ũj(θj)‖2 + 6hj+1d ≤ 12h2j+1(L2UE‖θj‖2 +B2) + 6hj+1d
Hence we can bound the first part, (choosing h0 ≤ 1),
L2U 2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖θ(s)− θ( q(s)∑ k=1 hk)‖2ds
≤ L 2 U
2 K−1∑ j=0 [ 12h3j+1(L 2 UE‖θj‖2 +B2) + 6h2j+1d ] ≤ L2U max
0≤j≤K−1
[ 6(L2UE‖θj‖2 +B2) + 3d ] ( K−1∑ j=0 h2j+1)
≤ L2U max 0≤j≤K−1
[ 6(L2UE‖θj‖2 +B2) + 3d ] π2 6 h20 (32)
where the last line follows from the fact that K−1∑ j=0 1 (j + 1)3 ≤ K−1∑ j=0 1 (j + 1)2 ≤ ∞∑ j=0
1 (j + 1)2 = π2 6 .
The second part (30) can be bounded as follows:
1
2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖∇U(θ( q(s)∑ k=1 hi)−Gs(θ( q(s)∑ k=1 hi)‖2ds
= 1
2 K−1∑ j=0 hj+1E‖∇U(θj)−∇Ũ(θj)‖2
≤ σ max 0≤j≤K−1 (L2UE‖θj‖2 +B2) K−1∑ j=0 hj+1
≤ σ max 0≤j≤K−1 (L2UE‖θj‖2 +B2)(h0 K∑ j=1 1 j )
Due to the data-processing inequality for the relative entropy, we have DKL(µ̃K‖ν∑K
k=1 hk ) ≤ DKL(PtV ‖P t θ)
≤ L 2 U
2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖θ(s)− θ( q(s)∑ k=1 hi)‖2ds
+ 1
2 K−1∑ j=0 ∫ ∑j+1 k=1 hk∑j k=1 hk E‖∇U(θ( q(s)∑ k=1 hi)−Gs(θ( q(s)∑ k=1 hi)‖2ds
≤ L2U max 0≤j≤K−1
[ 6(L2UE‖θj‖2 +B2) + 3d ] π2 6 h20
+ σ max 0≤j≤K−1 (L2UE‖θj‖2 +B2)(h0 K∑ j=1 1 j )
Similar to the proof of cSGLD , we have max
0≤j≤K−1 (L2UE‖θj‖2 +B2) ≤ D0
Then we denote the 6L2U (D0 + d) as D1 and we can derive
DKL(µ̃K‖ν∑K k=1 hk ) ≤ D1h20 π2 6 + σD0h0 K∑ j=1 1 j
Then according to Proposition 3.1 in Bolley & Villani (2005) and Lemma 3.3 in Raginsky et al. (2017), if we denote κ0 + 2b+ 2d as D2, we can derive the following result,
W2(µ̃K , ν∑K k=1 hk )
≤ [12 +D2( K∑ k=1 hk)] 1/2 · [(DKL(µ̃K‖ν∑K k=1 hk ))1/2 + (DKL(µ̃K‖ν∑K k=1 hk )/2)1/4]
= [12 +D2(h0 K∑ j=1 1 j )]1/2 · [(D1h20 π2 6 + σD0h0 K∑ j=1 1 j )1/2 + (D1h 2 0 π2 12 + σD0h0 K∑ j=1 1 2j )1/4]
Now we derive the bound for W2(µ̃K , ν∑K k=1 hk ).
E.2 W2(ν∑K k=1 hk , ν∞)
We can directly get the following results from (3.17) in Raginsky et al. (2017) that there exist some positive constants (C3, C4),
W2(ν∑K k=1 hk , ν∞) ≤ C3 exp(− K∑ k=1 hk/C4)
Based on the convergence error bounds, we discuss an informal comparison with standard SGLD. Consider the following two cases.We must emphasize that since the term W2(µK , ν∑K
k=1 αk ) in the
equation 9 increases w.r.t. K, our α0 must be set small enough in practice. Hence, in this informal comparison, we also set α0 small enough to make W2(µK , ν∑K
k=1 αk ) less important.
i) If the initial stepsizes satisfy α0 ≥ h0, our algorithm cSGLD runs much faster than the standard SGLD in terms of the amount of “diffusion time” i.e., the ”t” indexing θt in the continuous-time SDE mentioned above. This result follows from ∑K k=1 αk = Kα0 2 and ∑K k=1 hk = ∑K k=1 h0 k =
O(h0 logK) Kα02 . In standard SGLD, since the error described by W2(µ̃K , ν∑Kk=1 hk) increases w.r.t. K, h0 needs to be set small enough in practice to reduce the error. Following the general analysis of SGLD in Raginsky et al. (2017); Xu et al. (2017), the dominant term in the decomposition equation 21 will be W2(ν∑K
k=1 hk , ν∞) since it decreases exponentially fast with the increase of t
andW2(µ̃K , ν∑K k=1 hk ) is small due to the setting of small h0. Since ∑K k=1 αk increases much faster
in our algorithm than the term ∑K k=1 hk in standard SGLD, our algorithm thus endows less error
for K iterations, i.e., W2(ν∑K k=1 αk , ν∞) W2(ν∑K k=1 hk
, ν∞). Hence, our algorithm outperforms standard SGLD, as will be verified in our experiments.
ii) Instead of setting the h0 small enough, one may consider increasing h0 to make standard SGLD run as “fast” as our proposed algorithm, i.e., ∑K k=1 hk ≈ ∑K k=1 αk. Now the W2(ν∑Kk=1 hk , ν∞)
in equation 21 is almost the same as the W2(ν∑K k=1 hk
, ν∞) in equation 9. However, in this case, it is worth noting that h0 scales as O(α0K/ logK). We can notice that h0 is much larger than the α0 and thus the W2(µ̃K , ν∑K k=1 hk ) cannot be ignored. Now the h20 term in W2(µ̃K , ν∑K k=1 hk )
would scale as O(α20K2/ log 2K), which makes W2(µ̃K , ν∑K
k=1 αk ) in equation 21 much larger
than our W2(µK , ν∑K k=1 αk ) defined in equation 9 since O(α20K2/ log 2K) O(α20K). Again, our algorithm cSGLD achieves a faster convergence rate than standard SGLD.

heading: F COMBINING SAMPLES
text: In cyclical SG-MCMC, we obtain samples from multiple modes of a posterior distribution by running the cyclical step size schedule for many periods. We now show how to effectively utilize the collected samples. We consider each cycle exploring different part of the target distribution p(θ|D) on a metric space Θ. As we haveM cycles in total, themth cycle characterizes a local region Θm ⊂ Θ, defining the “sub-posterior” distribution: pm(θ|D) = p(θ|D)1Θmwm ,with wm = ∫ Θm
p(θ|D)dθ, where wm is a normalizing constant. For a testing function f(θ), we are often interested in its true posterior expectation f̄ = ∫ f(θ)p(θ|D)dθ. The sample-based estimation is
f̂ = M∑ m=1 wmf̂m with f̂m = 1 Km Km∑ j=1 f(θ (m) j ), (33)
where Km is the number of samples from the mth cycle, and θ(m) ∈ Θm. The weight for each cycle wi is estimated using the harmonic mean method (Green, 1995; Raftery et al., 2006): ŵm ≈ [ 1Km ∑Km j=1 1
p(D|θ(m)j ) ]−1. This approach provides a simple and consistent es-
timator, where the only additional cost is to traverse the training dataset to evaluate the likelihood p(D|θ(m)j ) for each sample θ (m) j . We evaluate the likelihood once off-line and store the result for testing.
If Θm are not disjoint, we can assume new sub-regions Θ̃m which are disjoint and compute the estimator as following
f̂m = 1
nm M∑ m=1 Km∑ j=1 f(θ(j)m )1Θ̃m(θ (m) j )
where
nm = M∑ m=1 Km∑ j=1 1Θ̃m(θ (m) j )
and 1Θ̃m(θ (m) j ) equals 1 only when θ (m) j ∈ Θ̃m. By doing so, our estimator still holds even if Θm are not disjoint.

heading: G THEORETICAL ANALYSIS UNDER CONVEX ASSUMPTION
text: Firstly, we introduce the following SDE
dθt = −∇U(θt)dt+ √ 2dWt , (34)
Let νt denote the distribution of θt, and the stationary distribution of equation 34 be p(θ|D), which means ν∞ = p(θ|D). However, the exact evaluation of the gradient∇U is computationally expensive. Hence, we need to adopt noisy evaluations of ∇U . For simplicity, we assume that at any point θk, we can observe the value
˜∇Uk = ∇U(θk) + ζk where ζk : k = 0, 1, 2, ... is a sequence of random (noise) vectors. Then the algorithm is defined as:
θk+1 = θk − αk+1 ˜∇Uk + √ 2αk+1ξk+1 (35)
Further, let µk denote the distribution of θk.
Following the existing work in Dalalyan & Karagulyan (2019), we adopt the following standard assumptions summarized in Assumption 5, Assumption 5.
• For some positive constants m and M, it holds
U(θ)− U(θ′)−∇U(θ′)T (θ − θ′) ≥ (m/2)‖θ − θ′‖22 ‖∇U(θ)−∇U(θ′)‖2 ≤M‖θ − θ′‖2
for any θ, θ′ ∈ Rd
• (bounded bias) E[‖E(ζk|θk)‖22] ≤ δ2d
• (bounded variance) E[‖ζk − E(ζk|θk)‖22 ≤ σ2d
• (independence of updates) ξk+1 in equation 35 is independent of (ζ1, ζ2, ..., ζk)

heading: G.1 THEOREM
text: Under Assumption 5 in the appendix and α0 ∈ (0, 1m ∧ 2 M ), if we define the αmin as α0 2 [ cos ( π mod(dK/Me−1,dK/Me) dK/Me ) + 1 ] , we can derive the the following bounds.
If mαmin +Mα0 ≤ 2, then W2(µk+1, ν∞) ≤
(1−mαmin)KW2(µ0, ν∞) + (1.65Mα
3/2 0 + α0δ)d 1/2
mαmin +
δ2α0d 1/2
1.65Mα 1/2 0 + δ + √ mαminδ
. (36)

heading: If mαmin +Mα0 > 2, then W2(µk+1, ν∞) ≤
text: (1− (2−Mα0))KW2(µ0, ν∞) + (1.65Mα
3/2 0 + α0δ)d 1/2
2−Mα0 +
δ2α0d 1/2
1.65Mα 1/2 0 + δ + √ 2−Mα0δ ,
(37)
where the M,m, δ, σ are some positive constants defined in Assumption 5

heading: G.2 PROOF
text: Proof. According to the equation 1, we can find that the stepsize αk varies from α0 to αmin, where αmin is defined as αmin , α02 [ cos ( π mod(dK/Me−1,dK/Me) dK/Me ) + 1 ] . When 0 < α0 <
min(2/M, 1/m), it is easy for us to know that 0 < αk < min(2/M, 1/m) for every k > 0. Then we can derive that all the ρk , max(1−mαk,Mαk − 1) will satisfy 0 < ρk < 1. Now according to the Proposition 2 in Dalalyan & Karagulyan (2019), we can derive the result that
W2(µk + 1, ν∞) 2 ≤ {ρk+1W2(µk, ν∞) + 1.65M(α3k+1d)1/2 + αk+1δ √ p}2 + δ2α2k+1d (38)
Then we will use another lemma derived from Dalalyan & Karagulyan (2019).
Lemma 2. If A,B,C are non-negative numbers such that A ∈ (0,1) and the sequence of non-negative numbers yk satisfies the following inequality
y2k+1 ≤ [(1−A)yk + C]2 +B2
for every integer k > 0. Then,
yk ≤ (1−A)ky0 + C
A +
B2
C + √ AB
Using Lemma 2, we can finish our proof now.
• If mαmin + Mα0 ≤ 2, the ρk will satisfy ρk ≤ 1 −mαmin for every k > 0. Then the equation 38 will turn into
W2(µk+1, ν∞) 2 ≤ {(1−mαmin)W2(µk, ν∞) + 1.65M(α30d)1/2 + α0δd1/2}2 + (δα0d1/2)2
for every k > 0. Then we can set A = mαmin, C = 1.65M(α30d) 1/2 + α0δd 1/2, B = δα0d 1/2 and we can get the result.
• If mαmin + Mα0 > 2, the ρk will satisfy ρk ≤ Mα0 − 1 for every k > 0. Then the equation 38 will turn into
W2(µk+1, ν∞) 2 ≤ {[1− (2−Mα0)]W2(µk, ν∞) + 1.65M(α30d)1/2 + α0δd1/2}2 + (δα0d1/2)2
for every k > 0. Then we can set A = 2 − Mα0, C = 1.65M(α30d)1/2 + α0δd1/2, B = δα0d 1/2 and we can get the result.

heading: H FUTURE DIRECTION FOR THE WASSERSTEIN GRADIENT FLOWS
text: We would like to point out that the convergence theorems developed in the above several sections can be potentially applied to study the convergence of the Wasserstein gradient flows (Santambrogio, 2017), which can be regarded as a continuous-time MCMC (Chen et al., 2018; Liu et al., 2019). The theorems may shed some lights on the stepsize choice of the Wasserstein gradient flows which is less studied in the literature. We leave it as an interesting future work.

heading: I HYPERPARAMETERS SETTING

heading: I.1 SENSITIVITY OF HYPERPARAMETERS
text: Compared to SG-MCMC, there are two additional hyperparameters in Algorithm 1: the number of cycles M and the proportion of exploration stage β. We now study how sensitive they are when comparing to the parallel MCMC. With the same setup as in Section 5.2, We compare our method with M cycles and L epochs per cycle with running M chains parallel MCMC for L epochs. The training budget is 200 epochs. In Table 2, M = 4 and β = 0.8 on CIFAR-10. We compare cSGLD and parallel SGLD with smaller and larger values ofM and β. In Table 6, we see that the conclusion that cSG-MCMC is better than parallel SG-MCMC holds with different values of M and β.

heading: I.2 HYPERPARAMETERS SETTING IN PRACTICE
text: Given the training budget, there is a trade-off between the number of cycles M and the cycle length. We find that it works well in practice by setting the cycle length such that the model with optimization methods will be close to a mode after running for that length. (e.g. the cycle length for CIFAR-10 is 50 epochs. The model optimized by SGD can achieve about 5% error after 50 epochs which means the model is close but not fully converge to a mode after 50 epochs.) Once the cycle length is fixed, M is fixed. β needs tuning for different tasks by cross-validation. Generally, β needs to be tuned so that the sampler has enough time to reach a good region before starting sampling.

heading: J TEMPERING IN BAYESIAN NEURAL NETWORKS
text: Tempering is common in modern Bayesian deep learning, for both variational inference and MCMC approaches (e.g., Li et al., 2016a; Nguyen et al., 2017; Fortunato et al., 2017). In general, tempering reflects the belief that the model capacity is misspecified. This combination of beliefs with data is what shapes the posterior we want to use to form a good predictive distribution.
Although we use the prescribed temperature in pSGLD (Li et al., 2016a) for all neural network experiments in the main text (T ≈ 0.0045), we here investigate the effect of temperature T on performance. We show negative log-likelihood (NLL) and classification error as a function of temperature on CIFAR-10 and CIFAR-100 using cSGLD with the same setup as in Section 5.2. We consider T ∈ [1, 0.5, 0.1, 0.05, 0.01, 0.005, 0]. Figure 6 and 7 show the results on CIFAR-10 and CIFAR-100, respectively. On CIFAR-10, the best performance is achieved at T = 0.1 with NLL 0.1331 and error 4.22%. On CIFAR-100, the best performance is achieved at T = 0.01 with NLL 0.7835 and error 20.53%. We find that the optimal temperature is often less than 1. We hypothesize that this result is due to the model misspecification common to neural networks.
Indeed, modern neural networks are particularly overparametrized. Tempering enables one to use a model with similar inductive biases to a modern neural network, but with a more well calibrated capacity (which is especially important when we are doing Bayesian integration instead of optimization). Indeed, we show that by sampling from the tempered posterior, we outperform optimization. Learning the amount of tempering by cross-validation is a principled way of aligning the tempering procedure with representing a reasonable posterior. We have shown that sampling with cSGMCMC with tempering helps in terms of both NLL and accuracy, which indicates that we are finding a better predictive distribution.

heading: K EXPERIMENTAL SETTING DETAILS

heading: K.1 BAYESIAN LOGISTIC REGRESSION
text: For both cSGLD and cSGHMC, M = 100, β = 0.01. For cSGLD, α0N = 1.2, 0.5, 1.5 for Austrilian, German and Hear respectively. For cSGHMC α0N = 0.5, 0.3, 1.0 for Austrilian, German and Hear respectively. For SG-MCMC, the stepsize is a for the first 5000 iterations and then switch to the decay schedule (2) with b = 0, γ = 0.55. aN = 1.2, 0.5, 1.5 for Austrilian, German and Hear respectively for SGLD and aN = 0.5, 0.3, 1.0 for Austrilian, German and Hear respectively for SGHMC. η = 0.5 in cSGHMC and SGHMC.
Assume that we collect {θb}Bb=1 samples. Effective sample size (ESS) is computed by
ESS = B 1 + 2 ∑B−1 s (1− s B )ρs
where ρs is estimated by
ρ̂s = 1
σ̂2(B − s) B∑ b=s+1 (θb − µ̂)(θb−s − µ̂)
Similar to Hoffman & Gelman (2014), σ̂2 and µ̂ are obtained by running an independent sampler. We use HMC in this paper.

heading: K.2 BAYESIAN NEURAL NETWORKS
text: For SG-MCMC, the stepsize decays from 0.1 to 0.001 for the first 150 epochs and then switch to the decay schedule (2) with a = 0.01, b = 0 and γ = 0.5005. η = 0.9 in cSGHMC, Snapshot-SGDM and SGHMC.

heading: K.3 UNCERTAINTY EVALUATION
text: For both cSG-MCMC and Snapshot, M = 4. β = 0.8 in cSG-MCMC. α0N = 0.01 and 0.008 for cSGLD and cSGHMC respectively. For SG-MCMC, the stepsize is a for the first 50 iterations and then switch to the decay schedule (2) with b = 0, γ = 0.5005. aN = 0.01 for SGLD and aN = 0.008 for SGHMC. η = 0.5 in cSGHMC, Snapshot-SGDM and SGHMC.

Review summaries:
This article presents cyclical stochastic gradient MCMC for Bayesian deep learning for inference in posterior distributions of network weights of Bayesian NNs. The posteriors of Bayesian NN weights are highly multi-modal and present difficulty for standard stochastic gradient MCMC methods. The proposed cyclical version periodically warm start the SG-MCMC process such that it can explore the multimodal space more efficiently.