Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks

Introduction

Deep Neural Networks (DNNs) have recently generated significant interest, largely due to their state-of-the-art performance on a wide variety of tasks, such as image classification (?) and language modeling (?). Despite this significant empirical success, it remains a challenge to effectively train DNNs. This is due to two main problems: $(\textup{\it i})$ The function under consideration is often difficult to optimize and find a good local minima. It is believed that this is in large part due to the pathological curvature and highly non-convex nature of the function to be optimized (?). $(\textup{\it ii})$ Standard optimization techniques lead to overfitting, typically addressed through early stopping (?).

A Bayesian approach for learning neural networks incorporates uncertainty into model learning, and can reduce overfitting (?). In fact, it is possible to view the standard dropout technique (?) as a form of Bayesian approximation that incorporates uncertainty (?; ?). Many other recent works (?; ?; ?) advocate incorporation of uncertainty estimates during model training to help improve robustness and performance.

While a Bayesian approach can ameliorate the overfitting issue in these complicated models, exact Bayesian inference in DNNs is generally intractable. Recently, several approaches have been proposed to approximate a Bayesian posterior for DNNs, including a stochastic variational inference (SVI) method “Bayes by Backprop” (BBB) (?) and an online expectation propogation method (OEP) “probabilistic backpropagation” (PBP) (?). These methods assume the posterior is comprised of separable Gaussian distributions. While this is a good choice for computational reasons, it can lead to unreasonable approximation errors and underestimation of model uncertainty.

A popular alternative to SVI and OEP is to use Stochastic Gradient Markov Chain Monte Carlo (SG-MCMC) methods to generate posterior samples (?; ?; ?; ?). One of the most common SG-MCMC methods is the Stochastic Gradient Langevin Dynamics (SGLD) algorithm (?). One merit of this approach is that it is highly scalable; it requires only the gradient on a small mini-batch of data, as in the optimization method Stochastic Gradient Descent (SGD). It has been shown that these MCMC approaches converge to the true posterior by using a slowly-decreasing sequence of step sizes (?; ?). Costly Metropolis-Hasting steps are not required.

Unfortunately, DNNs often exhibit pathological curvature and saddle points (?), which render existing SG-MCMC methods inefficient. In the optimization literature, numerous approaches have been proposed to overcome this problem, including methods based on adapting a preconditioning matrix in SGD to the local geometry (?; ?; ?). These approaches estimate second-order information with trivial per-iteration overhead, have improved risk bounds in convex problems compared to SGD, and demonstrate improved empirical performance in DNNs. The idea of using geometry in SG-MCMC has been explored in many contexts (?; ?; ?) and includes second-order approximations. Often, these approaches use the expected Fisher information, adding significant computational overhead. These methods lack the scalability necessary for learning DNNs, as discussed further below.

We combine adaptive preconditioners from optimization with SGLD, to improve SGLD efficacy. To note the distinction from SGLD, we refer to this as the Preconditioned SGLD method (pSGLD). This procedure is simple and adds trivial per-iteration overhead. We first show theoretical properties of this method, including bounds on risk and asymptotic convergence properties. We demonstrate improved efficiency of pSGLD by demonstrating an enhanced bias-variance tradeoff of the estimator for small problems. We further empirically demonstrate its application to several models and large datasets, including deep neural networks. In the DNN experiments, pSGLD outperforms the results based on standard SGLD from (?), both in terms of convergence speed and the test-set performance. Futher, pSGLD generates state-of-the-art performance for the examples tested.

Related Work

Various regularization schemes have been developed to prevent overfitting in neural networks, such as early stopping, weight decay, dropout (?), and dropconnect (?). Bayesian methods are appealing due to their ability to avoid overfitting by capturing uncertainty during learning (?). MCMC methods work by producing Monte Carlo approximations to the posterior, with asymptotic consistency (?). Traditional MCMC methods use the full dataset, which does not scale to large data problems. A pioneering work in combining stochastic optimization with MCMC was presented in (?), based on Langevin dynamics (?). This method was referred to as Stochastic Gradient Langevin Dynamics (SGLD), and required only the gradient on mini-batches of data. The per-iteration cost of SGLD is nearly identical to SGD. Unlike SGD, SGLD can generate samples from the posterior by injecting noise into the dynamics. This encourages the algorithm to explore the full posterior, instead of simply converging to a maximum a posterior (MAP) solution. Later, SGLD was extended by (?), (?) and (?). Furthermore, higher-order versions of the SGLD with momentum have also been proposed, including stochastic gradient Hamiltionian Monte Carlo (SGHMC) (?) and stochastic gradient Nose-Hoover Thermostats (SGNHT) (?).

It has been shown that incoporating higher-order gradient information helps train neural networks when employing optimization methods (?). However, calculations of higher-order information is often cumbersome in most models of interest. Methods such as quasi-Newton, and those approximating second-order gradient information, have shown promising results (?). An alternative to full quasi-Newton methods is to rescale parameters so that the loss function has similar curvature along all directions. This strategy has shown improved performance in Adagrad (?), Adadelta (?), Adam (?) and RMSprop (?) algorithms. Recently, RMSprop has been explained as a diagonal preconditioner in (?). While relatively mature in optimization, these techniques have not been developed in sampling methods. In this paper, we show that rescaling the parameter updates according to geometry information can also improve SG-MCMC, in terms of both training speed and predictive accuracy.

Preliminaries

Given data $\mathcal{D}=\{\bm{d}_{i}\}^{N}_{i=1}$, the posterior of model parameters $\bm{\theta}$ with prior $p(\bm{\theta})$ and likelihood $\prod_{i=1}^{N}p(\bm{d}_{i}|\bm{\theta})$ is computed as $p(\bm{\theta}|\mathcal{D})\propto p(\bm{\theta})\prod_{i=1}^{N}p(\bm{d}_{i}|\bm{\theta})$. In the optimization literature, the prior plays the role of a penalty that regularizes parameters, while the likelihood constitutes the loss function to be optimized. The task in optimization is to find the MAP estimate, $\bm{\theta}_{\textsl{MAP}}=\arg\!\max\log p(\bm{\theta}|\mathcal{D})$. Let $\Delta\bm{\theta}_{t}$ denote the change in the parameters at time $t$. Stochastic optimization methods such as Stochastic Gradient Descent (SGD)For maximization, this is Stochastic Gradient Ascent. Here, we abuse notation because SGD is a more common term. update $\bm{\theta}$ using the following rule:

where $\{\epsilon_{t}\}$ is a sequence of step sizes, and $\mathcal{D}^{t}=\{\bm{d}_{t_{1}},\cdots,\bm{d}_{t_{n}}\}$ a subset of $n<N$ data items randomly chosen from $\mathcal{D}$ at iteration $t$. The convergence of SGD has been established (?).

Stochastic sampling methods such as SGLD incorporate uncertainty into predictive estimates. SGLD samples $\bm{\theta}$ from the posterior distributions via a Markov Chain with steps:

with ${\bf I}$ denoting the identity matrix. It also uses mini-batches to take gradient descend steps at each iteration. Rates of convergence are proven rigorously in (?). Given a set of samples from the update rule (2), posterior distributions can be approximated via Monte Carlo approximations as $p(y|x,\mathcal{D})\approx\frac{1}{T}\sum_{t=1}^{T}p(y|x,\bm{\theta}_{t})$, where $T$ is the number of samples.

Both stochastic optimization and stochastic sampling approaches have the requirement that the step sizes satisfy the the following assumption.The requirement for SGLD can be relaxed, see (?; ?) for more details.

The step sizes $\{\epsilon_{t}\}$ are decreasing, i.e., $0<\epsilon_{t+1}<\epsilon_{t}$, with 1) $\sum_{t=1}^{\infty}\epsilon_{t}=\infty$; and 2) $\sum_{t=1}^{\infty}\epsilon_{t}^{2}<\infty$.

If these step-sizes are not satisfied in stochastic optimization, there is no guarantee of convergence because the gradient estimation noise is not eliminated. Likewise, in stochastic sampling, decreasing step-sizes are necessary for asymptotic consistency with the true posterior, where the approximation error is dominated by the natural stochasticity of Langevin dynamics (?).

Preconditioned Stochastic Gradient Langevin Dynamics

As noted in the previous section, standard SGLD updates all parameters with the same step size. This could lead to slow mixing when the components of $\bm{\theta}$ have different curvature. Unfortunately, this is generally true in DNNs due to the composition of nonlinear functions at multiple layers. A potential solution is to employ a user-chosen preconditioning matrix $G(\bm{\theta})$ in SGLD (?). The intuition is to consider the family of probability distributions $p(\bm{d}|\bm{\theta})$ parameterised by $\bm{\theta}$ lying on a Riemannian manifold. One can use the non-Euclidean geometry implied by this manifold to guide the random walk of a sampler. For any probability distribution, the expected Fisher information matrix $\mathcal{I}_{\bm{\theta}}$ defines a natural Riemannian metric tensor. To further scale up the method to a general online framework stochastic gradient Riemannian Langevin dynamics (SGRLD) was suggested in (?). At position $\bm{\theta}_{t}$, it gives the stepThe update form in (?) is more complicated and seemingly different from (3); however, they can be shown to be equivalent.,

​​​where $\Gamma_{i}(\bm{\theta})=\sum_{j}\frac{\partial G_{i,j}(\bm{\theta})}{\partial\theta_{j}}$ describes how the preconditioner changes with respect to $\bm{\theta}_{t}$. . This term vanishes in SGLD because the preconditioner of SGLD is a constant ${\bf I}$. Both the direction and variance in Eq.(3) depends on the geometry of $G(\bm{\theta}_{t})$. The natural gradient in the SGRLD step takes the direction of steepest descent on a manifold. Convergence to the posterior is guaranteed (?; ?) as long as step sizes satisfy Assumption 1.

Unfortunately, for many models of interest, the expected Fisher information is intractable. However, we note that any positive definite matrix defines a valid Riemannian manifold metric. Hence, we are not restricted to using the exact expected Fisher information. Preconditioning aims to constitute a local transform such that the rate of curvature is equal in all directions. Following this, we propose to use the same preconditoner as in RMSprop. This preconditioner is updated sequentially using only the current gradient information, and only estimates a diagonal matrix. It is given sequentially as,

where for notational simplicity, $\bar{g}(\bm{\theta}_{t};\mathcal{D}^{t})=\frac{1}{n}\sum_{i=1}^{n}\nabla_{\bm{\theta}}\log p(\bm{d}_{t_{i}}|\bm{\theta}_{t})$, is the sample mean of the gradient using mini-batch $\mathcal{D}^{t}$, and $\alpha\in$. Operators $\odot$ and $\oslash$ represent element-wise matrix product and division, respectively.

RMSprop utilizes magnitudes of recent gradients to construct a preconditioner. Flatter landscapes have smaller gradients while curved landscapes have larger gradients. Gradient information is usually only locally consistent. Therefore, two equivalent interpretations for Eq. (3) can be reached intuitively: i) the preconditioner equalizes the gradient so that a constant stepsize is adequate for all dimensions. ii) the stepsizes are adaptive, in that flat directions have larger stepsizes while curved directions have smaller stepsizes.

In DNNs, saddle points are the most prevalent critical points, that can considerably slow down training (?), mostly because the parameter space tends to be flat in many directions and ill-conditioned in the neighborhood of these saddle points. Standard SGLD will slowly escape the saddle point due to the typical oscillations along the high positive curvature direction. By transforming the landscape to be more equally curved, it is possible for the sampler to move much faster.

In addition, there are two tuning parameters: $\lambda$ controls the extremes of the curvature in the preconditioner (default $\lambda\!\!=\!\!10^{-5}$), and $\alpha$ balances the weights of historical and current gradients. We use a default value of $\alpha\!\!=\!\!0.99$ to construct an exponentially decaying sequence. Our Preconditioned SGLD with RMSprop is outlined in Algorithm 1.

Preconditioned SGLD Algorithms in Practice

This section first analyzes the finite-time convergence properties of pSGLD, then proposes a more efficient variant for practical use. We note that prior work gave similar theoretical results (?), and we extend the theory to consider the use of preconditioners.

For a bounded function $\phi(\bm{\theta})$, we are often interested in its true posterior expectation $\bar{\phi}=\int_{\mathcal{X}}\phi(\bm{\theta})p(\bm{\theta}|\mathcal{D})d\bm{\theta}$. For example, the class distribution of a data point in DNNs. In our SG-MCMC based algorithm, this intractable integration is approximated by a weighted sample average $\hat{\phi}=\frac{1}{S_{T}}\sum_{t=1}^{T}\epsilon_{t}\phi(\bm{\theta}_{t})$ at time $S_{T}=\sum_{t=1}^{T}\epsilon_{t}$, with stepsizes $\{\epsilon_{t}\}$. These samples are generated from an MCMC algorithm with a numerical integrator (e.g., our pSGLD algorithm) that discretizes the continuous-time Langevin dynamics. The precision of the true posterior average and its MCMC approximation is characterized by the expected difference between $\bar{\phi}$ and $\hat{\phi}$. We analyze the pSGLD algorithm by extending the work of (?; ?) to include adaptive preconditioners. We first show the asymptotic convergence properties of our algorithm in Theorem 1 by the mean of the mean squared error (MSE)This is different from the optimization literature where the regret is studied, which is not straightforward in the MCMC framework.. To get the convergence result, some mild assumptions on the smoothness and boundness of $\psi$, the solution functional of $\mathcal{L}\psi=\phi(\bm{\theta}_{t})-\bar{\phi}$, is needed, where $\mathcal{L}$ is the generator of corresponding stochastic differential equation for pSGLD. We discuss these conditions and prove the Theorem in Appendix A.

Define the operator $\Delta V_{t}=(N\bar{g}(\bm{\theta}_{t};\mathcal{D}^{t})-g(\bm{\theta}_{t};\mathcal{D}^{t}))^{\top}G(\bm{\theta}_{t})\nabla_{\bm{\theta}}$. Under Assumption 1, for a test function $\phi$, the MSE of the pSGLD at finite time $S_{T}$ is bounded, for some $C>0$ independent of $\{\epsilon_{t}\}$, as:

Practical Techniques

In practice, there is always a tradeoff between bias and variance. In the case of infinite computation time, the traditional MCMC setting can reduce the bias and variance to zero. However, in practice, time is limited. Obtaining more effective samples can reduce the total risk (Eq. (6)), even if bias is introduced. In the following, we provide two model-independent practical techniques to further speed up the proposed pSGLD.

Though the evaluation of $\Gamma(\bm{\theta}_{t})$ in our case is manageable due to its diagonal nature, we propose to remove it during sampling to reduce the computation. It is interesting that in our case ignoring $\Gamma(\bm{\theta}_{t})$ produces a bias controlled by $\alpha$ on the MSE.

Omitting $\Gamma(\bm{\theta}_{t})$ introduces an extra term in the bound that is controlled by the parameter $\alpha$. The proof is in Appendix B. Since $\alpha$ is always set to a value that is very close to $1$, the term $(1-\alpha)^{2}/\alpha^{3}\approx 0$, the effect of $\Gamma(\bm{\theta}_{t})$ negligible. In addition, more samples per unit time are generated when $\Gamma(\bm{\theta}_{t})$ is ignored, resulting in a smaller variance on the prediction. Note that the term $\Gamma(\bm{\theta}_{t})$ is heuristically ignored in (?), but is only able to approximate the true posterior in the case of infinite data, which is not required in our algorithm.

Thinning samples

Making predictions using a whole ensemble of models is cumbersome and may be too computationally expensive to allow deployment for a large number of users, especially when the models are large neural nets. One practical technique is to average models using a thinned version of samples. By thinning the samples in pSGLD, the total number of samples is reduced. However, these thinned samples have a lower autocorrelation time and can have a similar ESS. We can also guarantee the MSE remains the same form under the thinning schema. The proof is in Appendix C.

By thinning samples from our pSGLD algorithm, the MSE remains the same form as in Theorem 1, and asymptotically approaches 0.

Experiments

Our experiments focus on the effectiveness of preconditioners in pSGLD, and present results in four parts: a simple simulation, Bayesian Logistic Regression (BLR), and two widely used DNN models, Feedforward Neural Networks (FNN) and Convolutional Neural Networks (CNN).

The proposed algorithm that uses the discussed practical techniques is denoted as pSGLD. The prior on the parameters is set to $p(\bm{\theta})=\mathcal{N}(0,\sigma^{2}{\bf I})$. If not specifically mentioned, the default setting for DNN experiments is shared as follows. $\sigma^{2}=1$, minibatch size is 100, thinning interval is 100, burn-in is 300. We employ a block decay strategy for stepsize; it decreases by half after every $L$ epochs.

We first demonstrate pSGLD on a simple 2D Gaussian example, \mathcal{N}\Big{(}\left[\begin{array}[]{c}0\\ 0\end{array}\right],\left[\begin{array}[]{cc}0.16&0\\ 0&1\end{array}\right]\Big{)}. Given posterior samples, the goal is to estimate the covariance matrix. A diagonal covariance matrix is used to show the algorithm can adjust the stepsize at different dimension.

We first compare SGLD and pSGLD with a large range of different stepsizes $\epsilon$. $2\times 10^{5}$ samples are collected for each algorithm. Reconstruction errors and autocorrelation time are shown in Fig. 1 (a). We see that pSGLD dominates the “vanilla” SGLD in that it consistently shows a lower error and autocorrelation time, particularly with larger stepsize. When the stepsize is small enough, the sampler does not move much, and the performances of the two algorithms become similar. The first $600$ samples of both methods for $\epsilon=0.3$ are shown in Fig. 1 (b). Because step sizes in pSGLD can be adaptive, it implies that even if the covariance matrix of a target distribution is mildly rescaled, a new stepsize is unnecessary for pSGLD. Meanwhile, the stepsize of the vanilla SGLD needs to be fine-tuned in order to obtain decent samples. See Appendix E for further details.

Bayesian Logstic Regression

To demonstrate that our pSGLD is applicable to general Bayesian posterior sampling, we demonstrate results on BLR. A small Australian dataset (?) is first used with $N=690$ and dimension $D=14$. We choose a minibatch size of $5$. The prior variance is $\sigma^{2}=100$. $5\times 10^{3}$ iterations are used. For both pSGLD and SGLD, we test stepsize $\epsilon$ ranging from $1\times 10^{-7}$ to $1\times 10^{-4}$, with 50 runs for each algorithm.

Following (?), we report the time per minimum Effective Sample ($\propto 1/\text{EES}$) in Fig. 2 (a), which is proportional to the variance. pSGLD generates much larger ESS compared to SGLD, especially when the stepsize is large. Meanwhile, Fig. 2 (b) shows that pSGLD provides smaller error in estimating weights, where the “groundtruth” is obtained by $10^{6}$ samples from HMC with Metroplis-Hastings. Therefore, the overall risk is reduced.

We then test BLR on a large-scale Adult dataset, $\mathtt{a9a}$ (?), with $N_{\text{train}}=32561,N_{\text{test}}=16281,$ and $D=123$. Minibatch size is set to 50, and the prior variance is $\sigma^{2}=10$. The thinning interval is 50, burn-in is 500, and $T=1.5\times 10^{4}$ iterations are used. Stepsize $\epsilon=5\times 10^{-2}$ for pSGLD and SGLD. The test errors are compared in Table 1Bayesian Logstic Regression, and learning curves are shown in Fig. 3Bayesian Logstic Regression. Both SG-MCMC methods outperform the recently proposed doubly stochastic variational Bayes (SDVI) (?), and higher-order variational autoencoder methods (L-BFGS-SGVI, HFSGVI) (?). Furthermore, pSGLD converges in less than $4\times 10^{3}$ iterations, while SGLD at least needs double the time to reach this accuracy.

Feedforward Neural Networks

The first DNN model we study is the Feedforward Neural Networks (FNN), or multilayer perceptron (MLP). The activation function is rectified linear unit (ReLU). A two-layer model, 784-X-X-10, is employed, where X is the number of hidden units for each layer. 100 epochs are used, with $L=20$. We compare our propose method, pSGLD, with representative stochastic optimization methods: SGD, RMSprop and RMSspectral (?). After tuning, we set the optimal stepsize for each algorithm as: for pSGLD and RSMprop as follows: $\epsilon=5\times 10^{-4}$, while for SGLD and SGD as $\epsilon=5\times 10^{-1}$.

We test the algorithms on the standard MNIST dataset, consisting of $28\times 28$ images (thus the 784-dimensional input vector) from $10$ different classes ( to $9$) with $60,000$ training and $10,000$ test samples. The test classification errors for network (X-X) size 400-400, 800-800 and 1200-1200 are shown in Table LABEL:tab:fnn. The results of stochastic sampling methods are better than their corresponding stochastic optimization counterparts. This indicates that incorporating weight uncertainty can improve performances. By increasing the variance $\sigma^{2}$ of pSGLD from $1$ to $100$, more uncertainty is introduced into the model from the prior, and higher performance is obtained. Figure 4 (a) displays the distribution histograms of weights in the last training iteration of the 1200-1200 model. We observe that smaller variance in the prior imposes lower uncertainty, by making the weights concentrate to ; while larger variance in the prior leads to a wider range of weight choices, thus higher uncertainty.

We also compare to other techniques developed to prevent overfitting (dropout) and weight uncertainty (BPB, Gaussian and scale mixtures). pSGLD provides state-of-the-art performance for FNN on test accuracy. We further note that pSGLD is able to give increasing performance with increasing network size, whereas BPB and SGD dropout do not. This is probably because overfitting is harder to be dealt with in large neural networks with pure optimization techniques.

Finally, learning curves of network configuration 1200-1200 are plot in Fig. 4 (b)RMSspectral is not shown because it uses larger batch sizes and so is difficult to compare on this scale.. We empirically find that pSGLD and SGLD take fewer iterations to converge, and the results are more stable than their optimization counterparts. Moreover, it can be seen that pSGLD consistently converges faster and to a better point than SGLD. Learning curves for other network sizes are provided in Appendix F. While the ensemble of samples requires more computation than a single FNN in testing, it shows significantly improved performance. As well, (?) showed that learning a single FNN that approximates the model average result gave nearly the same performance. We employ this idea, and suggest a fast version, distilled pSGLD. Its results for $\sigma^{2}=1$ show it can maintain good performances.

Our next DNN is the popular CNN model. We use a standard network configuration with 2 convolutional layers followed by 2 fully-connected layers (?). Both convolutional layers use $5\times 5$ filter size with 32 and 64 channels, respectively; $2\times 2$ max pooling is used after each convolutional layer. The fully-connected layers have 200-200 hidden nodes with ReLU nonlinearities, 20 epochs are used, and $L$ is set to 10. The stepsizes for pSGLD and RMSprop is set to $\epsilon=\{1,2\}\times 10^{-3}$ via grid search. For SGLD and SGD, this is $\epsilon=\{1,2\}\times 10^{-1}$. Additional results with CNNs are in Appendix G.

The same MNIST dataset is used. A comparison of test errors is shown in Table 3G, with the corresponding learning curves in Fig. 5G. We emphasize that the purpose of this experiment is to compare methods on the same model architecture, not to achieve overall state-of-the-art results. The CNN trained with traditional SGD gives an error of 0.82%. pSGLD shows significant improvement, with an error of 0.45%. This result is also comparable with some recent state-of-the-art CNN based systems, which have much more complex architectures. These include the stochastic pooling (?), Network in Network (NIN) (?) and Maxout Network(MN) (?).

Conclusion

A preconditioned SGLD is developed based on the RMSprop algorithm, with controllable finite-time approximation error. We apply the algorithm to DNNs to overcome their notorious problems of overfitting and pathological curvature. Extensive experiments show that our pSGLD can adaptive to the local geometry, allowing improved effective sampling rates and performance. It provides sample-based uncertainty in DNNs, and achieves state-of-the-arts performances on FNN and CNN models. Interesting future directions include exploring applications to latent variable models or recurrent neural networks (?).

This research was supported in part by ARO, DARPA, DOE, NGA, ONR and NSF.

References

Appendix A A. The proof for main theorem

In (?), the authors provide the convergence property for general SG-MCMC, here we follow their assumptions and proof techniques, with specific treatment on the 1st-order numerical integrator, and the case of preconditioner.

The solution functional $\psi(\bm{\theta}_{t})$ characterizes the difference between $\phi(\bm{\theta}_{t})$ and the posterior average $\bar{\phi}$ for every $\bm{\theta}_{t}$, thus would typically possess a unique solution, which is at least as smooth as $\phi$ under the elliptic or hypoelliptic settings (?). In the unbounded domain of $\bm{\theta}_{t}$, to make the presentation simple, we follow (?) and make certain assumptions on the solution functional, $\psi$, of the Poisson equation (10), which are used in the detailed proofs.

The mild assumptions of smoothness and boundedness made in the main paper are detailed as follows.

Proof of Theorem 1

Based on Assumption 2, we prove the main theorem.

where $\Delta V_{t}\triangleq(N\bar{g}(\bm{\theta}_{t};\mathcal{D}^{t})-g(\bm{\theta}_{t};\mathcal{D}^{t}))^{\top}G(\bm{\theta}_{t})\nabla_{\bm{\theta}}$, $g(\bm{\theta}_{t};\mathcal{D}^{t})$ is the full gradient, $\bar{g}(\bm{\theta}_{t};\mathcal{D}^{t})$ is the stochatic gradient calculated from the $t$-th minibatch.

In pSGLD, we use the Euler integrator, which is a first order integrator. As a result, according to (?), for a test function $\phi$, we can decompose it as:

According to the assumptions, there exists a functional $\psi$ that solves the following Poisson Equation:

where $\bar{\phi}$ is defined in the main text.

As a result, there exists some positive constant $C$, such that:

$A_{1}$ can be bounded by assumptions, and $A_{2}$ can be easily shown to be bounded by $O(\sqrt{\epsilon_{t}})$ due to the Gaussian noise. It turns out that the resulting terms have order higher than those from the other terms, thus can be ignored in the expression below. After some simplifications, (17) is bounded by:

for some $C>0$. It is easy to show under the assumptions, all the terms in the above bound approach zero. This completes the first part of the theorem. $\blacksquare$

Appendix B B. The proof for Corollary 2

To prove Corollary 2, we first show the following results.

Assume that the 1st-order and 2nd-order gradient are bounded, then there exists some constant $M$, for k-th component of $\Gamma(\bm{\theta}_{t})$, we have

Since $\Gamma(\bm{\theta})$ is a diagnal matrix, we focus on one of its elements thus omit the index $k$ in the following.

First, the iterative form of exponential moving average can be written as a function of the gradients at all the previous timesteps:

Based on this, for each component of $\Gamma(\bm{\theta}_{t})$, we have

With the assumption that the 1st-order and 2nd-order gradient are bounded, we have $\left|V^{-\frac{3}{2}}(\theta_{t-1})\frac{\partial\bar{g}(\theta_{t})}{\partial\theta_{t}}\right|\leq M$, where $M$ is a constant independent of $\{\epsilon_{t}\}$. Therefore, $\left|\sum_{t=1}^{T}\epsilon_{t}\Gamma(\theta_{t})\right|\ll MT(1-\alpha)/\alpha^{\frac{3}{2}}$. $\blacksquare$

Based on Lemma 4, we now proceed to the proof of Corollary 2.

By dropping the $\Gamma(\bm{\theta}_{t})$ terms, we get a modified version of the local generator corresponding to the SDE of the pSGLD, defined as

Following the proof of Theorem 1, we can derive the bound for $(\hat{\phi}-\bar{\phi})^{2}$, which is no more than (17) with an extra term as:

where the last inequality follows by using the bound from Lemma 4. Taking expectation on both sides, we arrive at the MSE:

By thinning samples from the pSGLD, we obtain a sequence of subsamples $\{\bm{\theta}_{t_{1}},\cdots,\bm{\theta}_{t_{m}}\}$ from the original samples $\{\bm{\theta}_{1},\cdots,\bm{\theta}_{n}\}$ where $m\leq n$ and $(t_{1},\cdots,t_{m})$ is a subsequence of $(1,2,\cdots,n)$. Since we use the 1st-order Euler integrator, based on the definition in (?), we have for the original samples:

where $A\circ B$ denotes the composition of the two operators $A$ and $B$, i.e., $A$ is evaluated on the output of B. Now substitute (28) into (29), and use the Baker-Campbell-Hausdorff formula (?) for commutators, we have

Variance term in risk of estimator

where the term $\sum_{|t|>2T}A(|t|)$ is omitted from (37) to (38), which is usually small according to the property of autocovariance function.

We repeat some defintions from the main paper (?).

is the autocovariance function, manifesting how strong two samples with a time lag $t$ are correlated. Its normalized version

is called the autocorrelation function (ACF).

is the integrated autocorrelation time (ACT), which measures the interval between independent samples.

Note that effective sample size (ESS) is defined as

Plugin the definition into the derivation for variance, we have

Appendix E E. More results on simulation

We demonstrate our pSGLD on a simple 2D gaussian example, \mathcal{N}\Big{(}\left[\begin{array}[]{c}0\\ 0\end{array}\right],\left[\begin{array}[]{cc}0.16&0\\ 0&a\end{array}\right]\Big{)}. The first 600 samples of both methods for different $a$ and $\epsilon$ are shown in Fig. 1.

Comparing the results for different stepsize $\epsilon$ at the same $a$, it can be seen that pSGLD can adapt stepsizes acorrding to the manifold geometry of different dimensions.

When $a$ is rescaled from $0.5$ to $2$, stepsize $\epsilon=0.1$ is appropriate for SGLD at $a=0.5$, but not a good choice at $a=2$, because the space is not fully explored. This also implies that even if the covariance matrix of a target distribution is mildly rescaled, we do not have to choose a new stepsize for pSGLD. Whilst, the stepsize of the standard SGLD needs to be fine-tuned in order to obtain decent samples.

Learning curves for network sizes of 400-400 and 800-800 on MNIST are provided in Fig. 7 (a) and (b), respectively. Similar with results of network size 1200-1200 in the main paper, stochastic sampling methods take less iterations to converge, and the results are more stable than their optimization counterparts. Moreover, it can be seen that pSGLD consistently converges faster and better than SGLD and others.

We use another fairly standard network configuration containing 2 convolutional layers on MNIST dataset. It is followed by a single fully-connected layer (?), containing 500 hidden nodes that uses ReLU. Both convolutional layers use $5\times 5$ filter size with 32 and 64 channels, respectively, $2\times 2$ max pooling are used after each convolutional layer. 100 epochs are used, and $L$ is set to 20. The stepsizes for pSGLD and RSMprop are set to $\epsilon=\{1,2\}\times 10^{-3}$ via grid search. For SGLD and SGD, this is $\epsilon=\{1,2\}\times 10^{-1}$.

A comparison of test errors is shown in Table 1G, with the corresponding learning curves in Fig. 3G. Again, under the same network architecture, CNN trained with traditional SGD gives an error of 0.81%, while pSGLD has a significant improvement, with an error of 0.56%.

We also tested a similar 3-layer CNN with 32-32-64 channels on Cifar-10 RGB image dataset (?), which consists of $50,000$ samples for training and $10,000$ samples for testing. No data augmentation is employed for the dataset. We keep the same setting for pSGLD and SGLD from MNIST, and show the comparison on Cifar-10 in Fig. 9. pSGLD converges faster and reach a lower error.