Bridging the Gap between Stochastic Gradient MCMC and Stochastic Optimization

Introduction

Machine learning has made significant recent strides due to large-scale learning applied to “big data”. Large-scale learning is typically performed with stochastic optimization, and the most common method is stochastic gradient descent (SGD) (Bottou, 2010). Stochastic optimization methods are devoted to obtaining a (local) optima of an objective function. Alternatively, Bayesian methods aim to compute the expectation of a test function over the posterior distribution. At first glance, these methods appear to be distinct, independent approaches to learning. However, even the celebrated Gibbs sampler was first introduced to statistics as a simulated annealing method for maximum a posteriori estimation (i.e., finding an optima) (Geman and Geman, 1984).

Recent work on large-scale Bayesian learning has focused on incorporating the speed and low-memory costs from stochastic optimization. These approaches are referred to as stochastic gradient Markov chain Monte Carlo (SG-MCMC) methods. Well-known SG-MCMC methods include stochastic gradient Langevin dynamics (SGLD) (Welling and Teh, 2011), stochastic gradient Hamiltonian Monte Carlo (SGHMC) (Chen et al., 2014), and stochastic gradient thermostats (SGNHT) (Ding et al., 2014). SG-MCMC has become increasingly popular in the literature due to practical successes, ease of implementation, and theoretical convergence properties (Teh et al., 2014; Vollmer et al., 2015; Chen et al., 2015).

There are obvious structural similarities between SG-MCMC algorithms and stochastic optimization methods. For example, SGLD resembles SGD with additive Gaussian noise. SGHMC resembles SGD with momentum (Rumelhart et al., 1986), adding additive Gaussian noise when updating the momentum terms (Chen et al., 2014). These similarities are detailed in Section 2. Despite these structural similarities, the theory is unclear on how additive Gaussian noise differentiates a Bayesian algorithm from its optimization analog.

Just as classical sampling methods were originally used for optimization (Geman and Geman, 1984), we directly address using SG-MCMC algorithms for optimization. A major benefit of adapting these schemes is that Bayesian learning is (in theory) able to fully explore the parameter space. Thus it may find a better local optima, if not the global optima, for a non-convex objective function.

Specifically, in this work we first extend the recently proposed multivariate stochastic gradient thermostat algorithm (Gan et al., 2015) with Riemannian information geometry, which results in an adaptive preconditioning and momentum scheme with analogs to Adam (Kingma and Ba, 2015) and RMSprop (Tieleman and Hinton, 2012). We propose an annealing scheme on the system temperature to move from a Bayesian method to a stochastic optimization method. We call the proposed algorithm Stochastic AnNealing Thermostats with Adaptive momentum (Santa). We show that in the temperature limit, Santa recovers the SGD with momentum algorithm except that: i) adaptive preconditioners are used when updating both model and momentum parameters; ii) each parameter has an individual, learned momentum parameter. Adaptive preconditioners and momentums are desirable in practice because of their ability to deal with uneven, dynamic curvature (Dauphin et al., 2015). For completeness, we first review related algorithms in Section 2, and present our novel algorithm in Section 3.

We develop theory to analyze convergence properties of our algorithm, suggesting that Santa is able to find a solution for an (non-convex) objective function close to its global optima, shown in Section 4. The theory is based on the analysis from stochastic differential equations (Teh et al., 2014; Chen et al., 2015), and presents results on bias and variance of the annealed Markov chain. This is a fundamentally different approach from the traditional convergence explored in stochastic optimization, or the regret bounds used in online optimization. We note we can adapt the regret bound of Adam (Kingma and Ba, 2015) for our zero-temperature algorithm (with a few trivial modifications) for a convex problem, as shown in Supplementary Section F. However, this neither addresses non-convexity nor the annealing scheme that our analysis does.

In addition to theory, we demonstrate effective empirical performance on a variety of deep neural networks (DNNs), achieving the best performance compared to all competing algorithms for the same model size. This is shown in Section 5. The code is publicly available at https://github.com/cchangyou/Santa.

Preliminaries

Throughout this paper, we denote vectors as bold, lower-case letters, and matrices as bold, upper-case letters. We use $\odot$ for element-wise multiplication, and $\oslash$ as element-wise division; $\sqrt{\cdot}$ denotes the element-wise square root when applied to vectors or matricies. We reserve $(\cdot)^{1/2}$ for the standard matrix square root. $\textbf{I}_{p}$ is the $p\times p$ identity matrix, 1 is an all-ones vector.

The goal of an optimization algorithm is to minimize an objective function $U({\bm{\mathchar 28946\relax}})$ that corresponds to a (non-convex) model of interest. In a Bayesian model, this corresponds to the potential energy defined as the negative log-posterior, $U({\bm{\mathchar 28946\relax}})\triangleq-\log p({\bm{\mathchar 28946\relax}})-\sum_{n=1}^{N}\log p(\operatorname{\mathbf{x}}_{n}|{\bm{\mathchar 28946\relax}})$. Here ${\bm{\mathchar 28946\relax}}\in{\mathbf{R}}^{p}$ are the model parameters, and $\{\operatorname{\mathbf{x}}_{n}\}_{n=1,\dots,N}$ are the $d$-dimensional observed data; $p({\bm{\mathchar 28946\relax}})$ corresponds to the prior and $p(\operatorname{\mathbf{x}}_{n}|{\bm{\mathchar 28946\relax}})$ is a likelihood term for the $n^{th}$ observation. In optimization, $-\sum_{n=1}^{N}\log p(\operatorname{\mathbf{x}}_{n}|{\bm{\mathchar 28946\relax}})$ is typically referred to as the loss function, and $-\log p({\bm{\mathchar 28946\relax}})$ as a regularizer.

A recent idea in stochastic optimization is to use an adaptive preconditioner, also known as a variable metric, to improve convergence rates. Both ADAgrad (Duchi et al., 2011) and Adam (Kingma and Ba, 2015) adapt to the local geometry with a regret bound of $\mathcal{O}(\sqrt{N})$. Adam adds momentum as well through moment smoothing. RMSprop (Tieleman and Hinton, 2012), Adadelta (Zeiler, 2012), and RMSspectral (Carlson et al., 2015) are similar methods with preconditioners. Our method introduces adaptive momentum and preconditioners to the SG-MCMC. This differs from stochastic optimization in implementation and theory, and is novel in SG-MCMC.

Simulated annealing (Kirkpatrick et al., 1983; Černý, 1985) is well-established as a way of acquiring a local mode by moving from a high-temperature, flat surface to a low-temperature, peaky surface. It has been explored in the context of MCMC, including reversible jump MCMC (Andrieu et al., 2000), annealed important sampling (Neal, 2001) and parallel tempering (Li et al., 2009). Traditional algorithms are based on Metropolis–Hastings sampling, which require computationally expensive accept-reject steps. Recent work has applied simulated annealing to large-scale learning through mini-batch based annealing (van de Meent et al., 2014; Obermeyer et al., 2014). Our approach incorporates annealing into SG-MCMC with its inherent speed and mini-batch nature.

The Santa Algorithm

Santa extends the mSGNHT algorithm with preconditioners and a simulated annealing scheme. A simple pseudocode is shown in Algorithm 1, or a more complex, but higher accuracy version, is shown in Algorithm 2, and we detail the steps below.

The first extension we consider is the use of adaptive preconditioners. Preconditioning has been proven critical for fast convergence in both stochastic optimization (Dauphin et al., 2015) and SG-MCMC algorithms (Patterson and Teh, 2013). In the MCMC literature, preconditioning is alternatively referred to as Riemannian information geometry (Patterson and Teh, 2013). We denote the preconditioner as $\{\operatorname{\mathbf{G}}_{t}\in{\mathbf{R}}^{p\times p}\}$. A popular choice in SG-MCMC is the Fisher information matrix (Girolami and Calderhead, 2011). Unfortunately, this approach is computationally prohibitive for many models of interest. To avoid this problem, we adopt the preconditioner from RMSprop and Adam, which uses a vector $\{\operatorname{\mathbf{g}}_{t}\in{\mathbf{R}}^{p}\}$ to approximate the diagonal of the Fisher information matrixes (Li et al., 2016a). The construction sequentially updates the preconditioner based on current and historical gradients with a smoothing parameter $\sigma$, and is shown as part of Algorithm 1. While this approach will not capture the Riemannian geometry as effectively as the Fisher information matrix, it is computationally efficient.

Santa also introduces an annealing scheme on system temperatures. As discussed in Section 2, mSGNHT naturally accounts for a varying temperature by matching the particle momentum to the system temperature. We introduce $\beta=\{\beta_{1},\beta_{2},\cdots\}$, a sequence of inverse temperature variables with $\beta_{i}<\beta_{j}$ for $i<j$ and $\lim_{i\rightarrow\infty}\beta_{i}=\infty$. The infinite case corresponds to the zero-temperature limit, where SG-MCMCs become deterministic optimization methods.

The annealing scheme leads to two stages: the exploration and the refinement stages. The exploration stage updates all parameters based on an annealed sequence of stochastic dynamic systems (see Section 4 for more details). This stage is able to explore the parameter space efficiently, escape poor local modes, and finally converge close to the global mode. The refinement stage corresponds to the temperature limit, i.e., $\beta_{n}\rightarrow\infty$. In the temperature limit, the momentum weight updates vanish and it becomes a stochastic optimization algorithm.

We propose two update schemes to solve the corresponding stochastic differential equations: the Euler scheme and the symmetric splitting scheme (SSS). The Euler scheme has simpler updates, as detailed in Algorithm 1; while SSS endows increased accuracy (Chen et al., 2015) with a slight increase in overhead computation, as shown in Algorithm 2. Section 4.1 elaborates on the details of these two schemes. We recommend the use of SSS, but the Euler scheme is simpler to implement and compare to known algorithms.

According to Section 4, the exploration stage helps the algorithm traverse the parameter space following the posterior curve as accurate as possible. For optimization, slightly biased samples do not affect the final solution. As a result, the term consisting of $(1-\operatorname{\mathbf{g}}_{t-1}\oslash\operatorname{\mathbf{g}}_{t})$ in the algorithm (which is an approximation term, see Section 4.1) is ignored. We found no decreasing performance in our experiments. Furthermore, the term $\operatorname{\mathbf{g}}_{t-1}$ associated with the Gaussian noise could be replaced with a fixed constant without affecting the algorithm.

Theoretical Foundation

In this section we present the stochastic differential equations (SDEs) that correspond to the Santa algorithm. We first introduce the general SDE framework, then describe the exploration stage in Section 4.1 and the refinement stage in Section 4.2. We give the convergence properties of the numerical scheme in Section 4.3. This theory uses tools from the SDE literature and extends the mSGNHT theory (Ding et al., 2014; Gan et al., 2015).

The SDEs are presented with re-parameterized $\operatorname{\mathbf{p}}=\operatorname{\mathbf{u}}/\eta^{1/2}$, ${\bm{\Xi}}=\mbox{diag}({\bm{\alpha}})/\eta^{1/2}$, as in Ding et al. (2014). The SDEs describe the motion of a particle in a system where ${\bm{\mathchar 28946\relax}}$ is the location and $\operatorname{\mathbf{p}}$ is the momentum.

where $\operatorname{\mathbf{Q}}=\mbox{diag}(\operatorname{\mathbf{p}}\odot\operatorname{\mathbf{p}})$, $w$ is standard Brownian motion, $\operatorname{\mathbf{G}}_{1}({\bm{\mathchar 28946\relax}})$ encodes geometric information of the potential energy $U({\bm{\mathchar 28946\relax}})$, and $\operatorname{\mathbf{G}}_{2}({\bm{\mathchar 28946\relax}})$ characterizes the manifold geometry of the Brownian motion. Note $\operatorname{\mathbf{G}}_{2}({\bm{\mathchar 28946\relax}})$ may be the same as $\operatorname{\mathbf{G}}_{1}({\bm{\mathchar 28946\relax}})$ for the same Riemannian manifold. We call $\operatorname{\mathbf{G}}_{1}({\bm{\mathchar 28946\relax}})$ and $\operatorname{\mathbf{G}}_{2}({\bm{\mathchar 28946\relax}})$ Riemannian metrics, which are commonly defined by the Fisher information matrix (Girolami and Calderhead, 2011). We use the RMSprop preconditioner (with updates from Algorithm 1) for computational feasibility. Using the Fokker-Plank equation (Risken, 1989), we show that the marginal stationary distribution of (5) corresponds to the posterior distribution.

Denote $\operatorname{\mathbf{A}}:\operatorname{\mathbf{B}}\triangleq\mbox{tr}\left\{\operatorname{\mathbf{A}}^{T}\operatorname{\mathbf{B}}\right\}$.The stationary distribution of (5) is: $p_{\beta}({\bm{\mathchar 28946\relax}},\operatorname{\mathbf{p}},{\bm{\Xi}})\propto$

An inverse temperature $\beta=1$ corresponds to the standard Bayesian posterior.

We note that $\operatorname{\mathbf{p}}$ in (5) has additional dependencies on $\operatorname{\mathbf{G}}_{1}$ and $\operatorname{\mathbf{G}}_{2}$ compared to Gan et al. (2015) that must be accounted for. ${\bm{\Xi}}\operatorname{\mathbf{p}}$ introduces friction into the system so that the particle does not move too far away by the random force; the terms $\nabla_{{\bm{\mathchar 28946\relax}}}G_{1}({\bm{\mathchar 28946\relax}})$ and $\nabla_{{\bm{\mathchar 28946\relax}}}G_{2}({\bm{\mathchar 28946\relax}})$ penalize the influences of the Riemannian metrics so that the stationary distribution remains invariant.

The first stage of Santa, exploration, explores the parameter space to obtain parameters near the global mode of an objective functionThis requires an ergodic algorithm. While ergodicity is not straightforward to check, we follow most MCMC work and assume it holds in our algorithm.. This approach applies ideas from simulated annealing (Kirkpatrick et al., 1983). Specifically, the inverse temperature $\beta$ is slowly annealed to temperature zero to freeze the particles at the global mode.

Minimizing $U({\bm{\mathchar 28946\relax}})$ is equivalent to sampling from the zero-temperature limit $p_{\beta}({\bm{\mathchar 28946\relax}})\triangleq\frac{1}{Z_{\beta}}e^{-\beta U({\bm{\mathchar 28946\relax}})}$ (proportional to (6)), with $Z_{\beta}$ being the normalization constant such that $p_{\beta}({\bm{\mathchar 28946\relax}})$ is a valid distribution. We construct a Markov chain that sequentially transits from high temperatures to low temperatures. At the state equilibrium, the chain reaches the temperature limit with marginal stationary distribution $\rho_{0}({\bm{\mathchar 28946\relax}})\triangleq\lim_{\beta\rightarrow\infty}e^{-\beta U({\bm{\mathchar 28946\relax}})}$, a point massThe sampler samples a uniform distribution over global modes, or a point mass if the mode is unique. We assume uniqueness and say point mass for clarity henceforth. located at the global mode of $U({\bm{\mathchar 28946\relax}})$. Specifically, we first define a sequence of inverse temperatures, $(\beta_{1},\beta_{2},\cdots,\beta_{L})$, such that $\beta_{L}$ is large enoughDue to numerical issues, it is impossible to set $\beta_{L}$ to infinity; we thus assign a large enough value for it and handle the infinity case in the refinement stage.. For each time $t$, we generate a sample according to the SDE system (5) with temperature $\frac{1}{\beta_{t}}$, conditioned on the sample from the previous temperature, $\frac{1}{\beta_{t-1}}$. We call this procedure annealing thermostats to denote the analog to simulated annealing.

Generating exact samples from (5) is infeasible for general models. One well-known numerical approach is the Euler scheme in Algorithm 1. The Euler scheme is a 1st-order method with relatively high approximation error (Chen et al., 2015). We increase accuracy by implementing the symmetric splitting scheme (SSS) (Chen et al., 2015; Li et al., 2016b). The idea of SSS is to split an infeasible SDE into several sub-SDEs, where each sub-SDE is analytically solvable; approximate samples are generated by sequentially evolving parameters via these sub-SDEs. Specifically, in Santa, we split (5) into the following three sub-SDEs:

We then update the sub-SDEs in order $A$-$B$-$O$-$B$-$A$ to generative approximate samples (Chen et al., 2015). This uses half-steps $h/2$ on the $A$ and $B$ updatesAs in Ding et al. (2014), we define $h=\sqrt{\eta}$., and full steps $h$ in the $O$ update. This is analogous to the leapfrog steps in Hamiltonian Monte Carlo (Neal, 2011). Update equations are given in the Supplementary Section A. The resulting parameters then serve as an approximate sample from the posterior distribution with the inverse temperature of $\beta$. Replacing $\operatorname{\mathbf{G}}_{1}$ and $\operatorname{\mathbf{G}}_{2}$ with the RMSprop preconditioners gives Algorithm 2. These updates require approximations to $\nabla_{{\bm{\mathchar 28946\relax}}}G_{1}({\bm{\mathchar 28946\relax}})$ and $\nabla_{{\bm{\mathchar 28946\relax}}}G_{2}({\bm{\mathchar 28946\relax}})$, addressed below.

We propose a computationally efficient approximation for calculating the derivative vector $\nabla_{{\bm{\mathchar 28946\relax}}}\operatorname{\mathbf{G}}_{1}({\bm{\mathchar 28946\relax}})$ based on the definition. Specifically, for the $i$-th element of $\nabla_{{\bm{\mathchar 28946\relax}}}\operatorname{\mathbf{G}}_{1}({\bm{\mathchar 28946\relax}})$ at the $t$-th iteration, denoted as $(\nabla_{{\bm{\mathchar 28946\relax}}}\operatorname{\mathbf{G}}_{1}^{t}({\bm{\mathchar 28946\relax}}))_{i}$, it is approximated as:

where $\Delta\operatorname{\mathbf{G}}_{1}^{t}\triangleq\operatorname{\mathbf{G}}_{1}^{t}-\operatorname{\mathbf{G}}_{1}^{t-1}$. Step $A_{1}$ follows by the definition of a derivative, and $A_{2}$ by using the update equation for ${\bm{\mathchar 28946\relax}}_{t}$, i.e., ${\bm{\mathchar 28946\relax}}_{t}={\bm{\mathchar 28946\relax}}_{t-1}+\operatorname{\mathbf{G}}_{1}^{t}({\bm{\mathchar 28946\relax}})\operatorname{\mathbf{p}}_{t-1}h$. According to Taylor’s theory, the approximation error for the $\nabla_{{\bm{\mathchar 28946\relax}}}\operatorname{\mathbf{G}}_{1}({\bm{\mathchar 28946\relax}})$ is $O(h)$, e.g.,

for some positive constant $\mathcal{B}_{t}$. The approximation error is negligible in term of convergence behaviors because it can be absorbed into the stochastic gradients error. Formal theoretical analysis on convergence behaviors with this approximation is given in later sections. Using similar methods, $\nabla_{{\bm{\mathchar 28946\relax}}}G_{2}^{t}({\bm{\mathchar 28946\relax}})$ is also approximately calculated.

2 Refinement

The refinement stage corresponds to the zero-temperature limit of the exploration stage, where ${\bm{\Xi}}$ is learned. We show that in the limit Santa gives significantly simplified updates, leading to an stochastic optimization algorithm similar to Adam or SGD-M.

In the refinement stage Santa is a stochastic optimization algorithm. This relation is easier seen with the Euler scheme in Algorithm 1. Compared with SGD-M (Rumelhart et al., 1986), Santa has both adaptive gradient and adaptive momentum updates. Unlike Adagrad (Duchi et al., 2011) and RMSprop (Tieleman and Hinton, 2012), refinement Santa is a momentum based algorithm.

The recently proposed Adam algorithm (Kingma and Ba, 2015) incorporates momentum and preconditioning in what is denoted as “adaptive moments.” We show in Supplementary Section F that a constant step size combined with a change of variables nearly recovers the Adam algorithm with element-wise momentum weights. For these reasons, Santa serves as a more general stochastic optimization algorithm that extends all current algorithms. As well, for a convex problem and a few trivial algorithmic changes, the regret bound of Adam holds for refinement Santa, which is $\mathcal{O}(\sqrt{T})$, as detailed in Supplementary Section F. However, our analysis is focused on non-convex problems that do not fit in the regret bound formulation.

3 Convergence properties

Our convergence properties are based on the framework of Chen et al. (2015). The proofs for all theorems are given in the Supplementary Material. We focus on the exploration stage of the algorithm. Using the Monotone Convergence argument (Schechter, 1997), the refinement stage convergence is obtained by taking the temperature limit from the results of the exploration stage. We emphasize that our approach differs from conventional stochastic optimization or online optimization approaches. Our convergence rate is weaker than many stochastic optimization methods, including SGD; however, our analysis applies to non-convex problems, whereas traditionally convergence rates only apply to convex problems.

The goal of Santa is to obtain ${\bm{\mathchar 28946\relax}}^{*}$ such that ${\bm{\mathchar 28946\relax}}^{*}=\arg\!\min_{{\bm{\mathchar 28946\relax}}}U({\bm{\mathchar 28946\relax}})$. Let $\{{\bm{\mathchar 28946\relax}}_{1},\cdots,{\bm{\mathchar 28946\relax}}_{L}\}$ be a sequence of parameters collected from the algorithm. Define $\hat{U}\triangleq\frac{1}{L}\sum_{t=1}^{L}U({\bm{\mathchar 28946\relax}}_{t})$ as the sample average, $\bar{U}\triangleq U({\bm{\mathchar 28946\relax}}^{*})$ the global optima of $U({\bm{\mathchar 28946\relax}})$.

As in Chen et al. (2015), we require certain assumptions on the potential energy $U$. To show these assumptions, we first define a functional $\psi_{t}$ for each $t$ that solves the following Poisson equation:

Let $\left\|\cdot\right\|$ be the operator norm. Under Assumption 1, the bias and MSE of the exploration stage in Santa with respect to the global optima for $L$ steps with stepsize $h$ is bounded, for some constants $C>0$ and $D>0$, with:

To get convergence results right before the refinement stage, let a sequence of functions $\{g_{m}\}$ be defined as $g_{m}\triangleq-\frac{1}{L}\sum_{l=m}^{L+m-1}e^{-\beta_{l}\hat{U}({\bm{\mathchar 28946\relax}})}$; it is easy to see that $\{g_{m}\}$ satisfies $g_{m_{1}}<g_{m_{2}}$ for $m_{1}<m_{2}$, and $\lim_{m\rightarrow\infty}g_{m}=0$. According to the Monotone Convergence Theorem (Schechter, 1997), the bias and MSE in the limit exists, leading to Corollary 3.

Under Assumptions 1, the bias and MSE of the refinement stage in Santa with respect to the global optima for $L$ steps with stepsize $h$ are bounded, for some constants $D_{1}>0$, $D_{2}>0$, as

Corollary 3 implies that in the refinement stage, the discrepancy between annealing distributions and the global optima vanishes, leaving only errors from discretized simulations of the SDEs, similar to the result of general SG-MCMC (Chen et al., 2015). We note that after exploration, Santa becomes a pure stochastic optimization algorithm, thus convergence results in term of regret bounds can also be derived; refer to Supplementary Section F for more details.

Experiments

In order to demonstrate that Santa is able to achieve the global mode of an objective function, we consider the double-well potential (Ding et al., 2014),

As shown in Figure 1 (left), the double-well potential has two modes, located at $\mathchar 28946\relax=-3$ and $\mathchar 28946\relax=2$, with the global optima at $\mathchar 28946\relax=-3$. We use a decreasing learning rate $h_{t}=t^{-0.3}/10$, and the annealing sequence is set to $\beta_{t}=t^{2}$. To make the optimization more challenging, we initialize the parameter at $\mathchar 28946\relax_{0}=4$, close to the local mode. The evolution of $\mathchar 28946\relax$ with respect to iterations is shown in Figure 1(right). As can be seen, $\mathchar 28946\relax$ first moves to the local mode but quickly jumps out and moves to the global mode in the exploration stage (first half iterations); in the refinement stage, $\mathchar 28946\relax$ quickly converges to the global mode and sticks to it afterwards. In contrast, RMSprop is trapped on the local optima, and convergences slower than Santa at the beginning.

2 Feedforward neural networks

We first test Santa on the Feedforward Neural Network (FNN) with rectified linear units (ReLU). We test two-layer models with network sizes 784-X-X-10, where X is the number of hidden units for each layer; 100 epochs are used. For variants of Santa, we denote Santa-E as Santa with a Euler scheme illustrated in Algorithm 1, Santa-r as Santa running only on the refinement stage, but with updates on ${\bm{\alpha}}$ as in the exploration stage. We compare Santa with SGD, SGD-M, RMSprop, Adam, SGD with dropout, SGLD and Bayes by Backprop (Blundell et al., 2015). We use a grid search to obtain good learning rates for each algorithm, resulting in $4\times 10^{-6}$ for Santa, $5\times 10^{-4}$ for RMSprop, $10^{-3}$ for Adam, and $5\times 10^{-1}$ for SGD, SGD-M and SGLD. We choose an annealing schedule of $\beta_{t}=At^{\gamma}$ with $A=1$ and $\gamma$ selected from 0.1 to 1 with an interval of 0.1. For simplicity, the exploration is set to take half of total iterations.

We test the algorithms on the standard MNIST dataset, which contains $28\times 28$ handwritten digital images from $10$ classes with $60,000$ training samples and $10,000$ test samples. The network size (X-X) is set to 400-400 and 800-800, and test classification errors are shown in Table 1. Santa show improved state-of-the-art performance amongst all algorithms. The Euler scheme shows a slight decrease in performance, due to the integration error when solving the SDE. Santa without exploration (i.e., Santa-r) still performs relatively well. Learning curves are plotted in Figure 2, showing that Santa converges as fast as other algorithms but to a better local optimaLearning curves of FNN with size of 800 are provided in Supplementary Section G..

3 Convolution neural networks

We next test Santa on the Convolution Neural Network (CNN). Following Jarrett et al. (2009), a standard network configuration with 2 convolutional layers followed by 2 fully-connected layers is adopted. 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 activation. The same parameter setting and dataset as in the FNN are used. The test errors are shown in Table 1, and the corresponding learning curves are shown in Figure 2. Similar trends as in FNN are obtained. Santa significantly outperforms other algorithms with an error of 0.45%. This result is comparable or even better than some recent state-of-the-art CNN-based systems, which have much more complex architectures.

4 Recurrent neural networks

We test Santa on the Recurrent Neural Network (RNN) for sequence modeling, where a model is trained to minimize the negative log-likelihood of training sequences:

where ${\bm{\mathchar 28946\relax}}$ is a set of model parameters, $\{\operatorname{\mathbf{x}}_{t}^{n}\}$ is the observed data. The conditional distributions in (9) are modeled by the RNN. The hidden units are set to gated recurrent units (Cho et al., 2014).

We consider the task of sequence modeling on four different polyphonic music sequences of piano, i.e., Piano-midi.de (Piano), Nottingham (Nott), MuseData (Muse) and JSB chorales (JSB). Each of these datasets are represented as a collection of 88-dimensional binary sequences, that span the whole range of piano from A0 to C8.

The number of hidden units is set to 200. Each model is trained for at most 100 epochs. According to the experiments and their results on the validation set, we use a learning rate of 0.001 for all the algorithms. For Santa, we consider an additional experiment using a learning rate of 0.0002, denoted Santa-s. The annealing coefficient $\gamma$ is set to 0.5. Gradients are clipped if the norm of the parameter vector exceeds 5. We do not perform any dataset-specific tuning other than early stopping on validation sets. Each update is done using a minibatch of one sequence.

The best log-likelihood results on the test set are achieved by using Santa, shown in Table 2. Learning curves on the Piano dataset are plotted in Figure 3. We observe that Santa achieves fast convergence, but is overfitting. This is straightforwardly addressed through early stopping. The learning curves for all the other datasets are provided in Supplementary Section G.

Conclusions

We propose Santa, an annealed SG-MCMC method for stochastic optimization. Santa is able to explore the parameter space efficiently and locate close to the global optima by annealing. At the zero-temperature limit, Santa gives a novel stochastic optimization algorithm where both model parameters and momentum are updated element-wise and adaptively. We provide theory on the convergence of Santa to the global optima for an (non-convex) objective function. Experiments show best results on several deep models compared to related stochastic optimization algorithms.

Acknowledgements

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

References

Appendix A Solutions for the sub-SDEs

We provide analytic solutions for the split sub-SDEs in Section 4.1. For stepsize $h$, the solutions are given in (13).

Appendix B Proof of Lemma 1

For a general stochastic differential equation of the form

where $\operatorname{\mathbf{x}}\in{\mathbf{R}}^{N}$, $F:{\mathbf{R}}^{N}\rightarrow{\mathbf{R}}^{N}$, $D:{\mathbf{R}}^{M}\rightarrow{\mathbf{R}}^{N\times P}$ are measurable functions with $P$, and $\operatorname{\mathbf{w}}$ is standard $P$-dimensional Brownian motion. (5) is a special case of the general form (24) with

We write the joint distribution of $\operatorname{\mathbf{x}}$ as

A reformulation of the main theorem in Ding et al. gives the following lemma, which is used to prove Lemma 1 in the main text.

The stochastic process of $\vec{\mathchar 28946\relax}$ generated by the stochastic differential equation (24) has the target distribution $p_{{\bm{\mathchar 28946\relax}}}({\bm{\mathchar 28946\relax}})=\frac{1}{Z}\exp\{-U({\bm{\mathchar 28946\relax}})\}$ as its stationary distribution, if $\rho(\operatorname{\mathbf{x}})$ satisfies the following marginalization condition:

and if the following condition is also satisfied:

where $\nabla\triangleq\left(\partial/\partial{\bm{\mathchar 28946\relax}},\partial/\partial\operatorname{\mathbf{p}},\partial/{\bm{\Xi}}\right)$, “$\cdot$” represents the vector inner product operator, “$:$” represents a matrix double dot product, i.e., $\operatorname{\mathbf{X}}:\operatorname{\mathbf{Y}}\triangleq\mbox{tr}(\operatorname{\mathbf{X}}^{\top}\operatorname{\mathbf{Y}})$.

We first have reformulated (5) using the general SDE form of (24), resulting in (25). Lemma 1 states the joint distribution of $({\bm{\mathchar 28946\relax}},\operatorname{\mathbf{p}},{\bm{\Xi}})$ is

with $H(\operatorname{\mathbf{x}})=\frac{1}{2}\operatorname{\mathbf{p}}^{\top}\operatorname{\mathbf{p}}+U({\bm{\mathchar 28946\relax}})+\frac{1}{2}\mbox{tr}\left\{\left({\bm{\Xi}}-\operatorname{\mathbf{G}}_{2}({\bm{\mathchar 28946\relax}})\right)^{\top}\left({\bm{\Xi}}-\operatorname{\mathbf{G}}_{2}({\bm{\mathchar 28946\relax}})\right)\right\}$. The marginalization condition (35) is trivially satisfied, we are left to verify condition (36). Substituting $\rho(\operatorname{\mathbf{x}})$ and $F$ into (36), we have the left-hand side

$\displaystyle\mbox{LHS}=\sum_{i}\frac{\partial}{\partial\operatorname{\mathbf{x}}_{i}}\left(\rho F_{i}\right)$ $\displaystyle=$ $\displaystyle\sum_{i}\frac{\partial\rho}{\partial\operatorname{\mathbf{x}}_{i}}F_{i}+\frac{\partial F_{i}}{\partial\operatorname{\mathbf{x}}_{i}}\rho$ $\displaystyle=$ $\displaystyle\sum_{i}\left(\frac{\partial F_{i}}{\partial\operatorname{\mathbf{x}}_{i}}-\frac{\partial H}{\partial\operatorname{\mathbf{x}}_{i}}F_{i}\right)\rho$ $\displaystyle=$ $\displaystyle\left(\sum_{i}\nabla_{{\bm{\mathchar 28946\relax}}_{i}}(\operatorname{\mathbf{G}}_{1})_{i:}\operatorname{\mathbf{p}}-\sum_{i}\mbox{diag}({\bm{\Xi}})\right.$ $\displaystyle-\sum_{i}\beta\left(\nabla_{{\bm{\mathchar 28946\relax}}_{i}}U-\sum_{j}({\bm{\Xi}}_{ij}-(\operatorname{\mathbf{G}}_{2})_{ij})\nabla_{{\bm{\mathchar 28946\relax}}_{i}}(\operatorname{\mathbf{G}}_{2})_{ij}\right)(\operatorname{\mathbf{G}}_{1}\operatorname{\mathbf{p}})_{i}$ $\displaystyle-\beta p^{T}\left(-\operatorname{\mathbf{G}}_{1}\nabla_{{\bm{\mathchar 28946\relax}}}U-{\bm{\Xi}}\operatorname{\mathbf{p}}+\frac{1}{\beta}\nabla_{{\bm{\mathchar 28946\relax}}}\operatorname{\mathbf{G}}_{1}+\operatorname{\mathbf{G}}_{1}({\bm{\Xi}}-\operatorname{\mathbf{G}}_{2})\nabla_{{\bm{\mathchar 28946\relax}}}\operatorname{\mathbf{G}}_{2}\right)$ $\displaystyle\left.-\beta\sum_{i}\left({\bm{\Xi}}_{ii}-(\operatorname{\mathbf{G}}_{2})_{ii}\right)\left(\operatorname{\mathbf{Q}}_{ii}-\frac{1}{\beta}\right)\right)\rho$ $\displaystyle=\frac{1}{\beta}\mbox{tr}\left\{\operatorname{\mathbf{G}}_{2}(\operatorname{\mathbf{p}}\operatorname{\mathbf{p}}^{T}-\mathbf{I})\right\}\rho\,.$

It is easy to see for the right-hand side

According to Lemma 4, the joint distribution (B) is the equilibrium distribution of (5). ∎

Appendix C Proof of Theorem 2

We start by proving the bias result of Theorem 2.

For our $2$nd-order integrator, according to the definition, we have:

We divide both sides by $Lh$, use the Poisson equation (8), and reorganize terms. We have:

for some $C_{3}\geq 0$. According to the assumption, the term $C_{1}$ is bounded. As a result, collecting low order terms, the bias can be expressed as:

where the last equation follows from the finiteness assumption of $\psi$, $\|\cdot\|$ denotes the operator norm and is bounded in the space of $\psi$ due to the assumptions. This completes the proof. ∎

Similar to the proof of Theorem 2, for our 2nd–order integrator we have:

Sum over $t$ from 1 to $L+1$ and simplify, we have:

Substitute the Poisson equation (8) into the above equation, divide both sides by $Lh$ and rearrange related terms, we have

Taking the square of both sides, it is then easy to see there exists some positive constant $C$, such that

Appendix D Proof of Corollary 3

The refinement stage corresponds to $\beta\rightarrow\infty$. We can prove that in this case, the integration terms in the bias and MSE in Theorem 2 converge to 0.

To show this, define a sequence of functions $\{g_{m}\}$ as:

it is easy to see the sequence $\{g_{m}\}$ satisfies $g_{m_{1}}<g_{m_{2}}$ for $m_{1}<m_{2}$, and $\lim_{m\rightarrow\infty}g_{m}=0$. According to the monotone convergence theorem, we have

As a result, the integration terms in the bounds for the bias and MSE vanish, leaving only the terms stated in Corollary 3. This completes the proof. ∎

Appendix E Reformulation of the Santa Algorithm

In this section we give a version of the Santa algorithm that matches better than our actual implementation, shown in Algorithm 3–7.

Appendix F Relationship of refinement Santa to Adam

In the Adam algorithm (see Algorithm 1 of Kingma and Ba ), the key steps are:

Here, we maintain the square root form of $\operatorname{\mathbf{g}}_{t}$, so the square is equivalent to the preconditioner used in Adam. As well, in Adam, the vector $\operatorname{\mathbf{b}}_{1}$ is set to the same constant between 0 and 1 for all entries. An equivalent formulation of this is:

The only differences between these steps and the Euler integrator we present in our Algorithm 1 are that our $\operatorname{\mathbf{b}}_{1}$ has a separate constant for each entry, and the second term in $u$ does not include the $\operatorname{\mathbf{b}}_{1}$ in our formulation. If we modify our algorithm to multiply the gradient by $\operatorname{\mathbf{b}}_{1}$, then our algorithm, under the same assumptions as Adam, will have a similar regret bound of $O(\sqrt{T})$ for a convex problem.

Because the focus of this paper is not on the regret bound, we only briefly discuss the changes in the theory. We note that Lemma 10.4 from Kingma and Ba will hold with element-wise $\operatorname{\mathbf{b}}_{1}$.

which contains an element-dependent $\gamma_{i}$ compared to Adam.

Theorem 10.5 of Kingma and Ba will hold with the same modifications and assumptions for a $\operatorname{\mathbf{b}}$ with distinct entries; the proof in Kingma and Ba is already element-wise, so it suffices to replace their global parameter $\gamma$ with distinct $\gamma_{i}\triangleq\frac{b_{1,i}^{2}}{\sqrt{\sigma}}$. This will give a regret of $O(\sqrt{T})$, the same as Adam.

Appendix G Additional Results

Learning curves of different algorithms on MNIST using FNN with size of 800 are plotted in Figure 4. Learning curves of different algorithms on four polyphonic music datasets using RNN are shown in Figure 6.

We additionally test Santa on the ImageNet dataset. We use the GoogleNet architecture, which is a 22 layer deep model. We use the default setting defined in the Caffe package$\href https://github.com/cchangyou/Santa/tree/master/caffe/models/bvlc_googlenet$. We were not able to make other stochastic optimization algorithms except SGD with momentum and the proposed Santa work on this dataset. Figure 5 shows the comparison on this dataset. We did not tune the parameter setting, note the default setting is favourable by SGD with momentum. Nevertheless, Santa still significantly outperforms SGD with momentum in term of convergence speed.