Preconditioned Stochastic Gradient Langevin Dynamics for
Deep Neural Networks

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.

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.

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)¹¹1For 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 (?).

For DNNs, the gradient is calculated by backpropagation (?). One data item $\bm{d}_{i}\triangleq(x_{i},y_{i})$ may consist of input $x_{i}\in\mathbb{R}^{D}$ and output $y_{i}\in\mathcal{Y}$, with $\mathcal{Y}$ being the output space (e.g., a discrete label space in classification). In the testing stage, the Bayesian predictive estimate for input $x$, is given by $p(y|x,\mathcal{D})=\mathbb{E}_{p(\bm{\theta}|\mathcal{D})}[p(y|x,\bm{\theta})]$. The MAP estimate simply approximates this expectation as $p(y|x,\mathcal{D})\approx p(y|x,\bm{\theta}_{\text{MAP}})$, ignoring parameter uncertainty.

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.²²2The requirement for SGLD can be relaxed, see (?; ?) for more details.

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 (?).

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 step³³3The 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[0,1]$. 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.

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.

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.