Title: A UNIFIED THEORY OF ADAPTIVE STOCHASTIC GRADIENT DESCENT AS BAYESIAN FILTERING

Abstract: We formulate stochastic gradient descent (SGD) as a novel factorised Bayesian filtering problem, in which each parameter is inferred separately, conditioned on the corresponding backpropagated gradient. Inference in this setting naturally gives rise to BRMSprop and BAdam: Bayesian variants of RMSprop and Adam. Remarkably, the Bayesian approach recovers many features of state-of-the-art adaptive SGD methods, including amongst others root-mean-square normalization, Nesterov acceleration and AdamW. As such, the Bayesian approach provides one explanation for the empirical effectiveness of state-of-the-art adaptive SGD algorithms. Empirically comparing BRMSprop and BAdam with naive RM-Sprop and Adam on MNIST, we find that Bayesian methods have the potential to considerably reduce test loss and classification error.

1. INTRODUCTION:

Deep neural networks have recently shown huge success at a range of tasks including machine translation (Wu et al., 2016) , dialogue systems (Serban et al., 2016) , handwriting generation (Graves, 2013) and image generation (Radford et al., 2015) . These successes have been facilitated by the development of a broad range of adaptive SGD methods, including ADAGrad (Duchi et al., 2011) , RMSprop (Hinton et al., 2012) , Adam (Kingma & Ba, 2015) , and variants thereof, including Nesterov acceleratation (Nesterov, 1983; Bengio et al., 2013; Dozat, 2016) and AdamW (Loshchilov & Hutter, 2017) . However, such a broad range of approaches raises the question of whether it is possible to obtain a unified theoretical understanding of adaptive SGD methods. Here we provide such a theory by reconciling state-of-the-art adaptive SGD algorithms with very early work that used Bayesian (Kalman) filtering to optimize the parameters of neural networks (Puskorius & Feldkamp, 1991; Sha et al., 1992; Puskorius & Feldkamp, 1994; 2001; Feldkamp et al., 2003; Ollivier, 2017) .

There have recently been attempts to connect adaptive SGD algorithms to natural gradient variational inference (VI) (Zhang et al., 2017; Khan et al., 2017; 2018) . These approaches give a momentum-free algorithm with a mean-square normalizer, in contrast to perhaps the most popular adaptive method, Adam (Kingma & Ba, 2015) , which combines momentum with a root-meansquare normalizer. To achieve a closer match to Adam, they modified their natural gradient VI updates, without a principled justification based on approximate inference, to incorporate momentum (Zhang et al., 2017; Khan et al., 2018) , and the root-mean-square normalizer (Khan et al., 2017; 2018) . As such, there appears to be only a loose connection between successful adaptive SGD algorithms such as Adam, and natural gradient VI.

There is a formal correspondence between natural gradient VI (Zhang et al., 2017; Khan et al., 2017; 2018) and Bayesian filtering (Ollivier, 2017) . While Ollivier (2017) did not examine the relationship between their filtering updates and RMSprop/Adam, the equivalence of this particular filtering approach and natural gradient VI indicates that they would encounter the issues described above, and thus be unable to obtain momentum or the root-mean-square normalizer (Zhang et al., 2017; Khan et al., 2017; 2018) . More problematically, Ollivier (2017) introduces dynamics into the Kalman filter, but these dynamics correspond to the "addition of an artificial process noise Q t proportional to [the posterior covariance] P t-1 ". Thus, their generative model depends on inferences made under that model: a highly unnatural assumption that most likely does not correspond to any "real" generative process.

A:

µ -i (t -2) µ -i (t -1) µ -i (t) µ -i (t + 1) µ -i (t + 2)

w * i (t -2) w * i (t -1) w * i (t) w * i (t + 1) w * i (t + 2) g i (t -2) g i (t -1) g i (t) g i (t + 1) g i (t + 2) µ i (t -2) µ i (t -1) µ i (t) µ i (t + 1) µ i (t + 2) B w * i (t -2) w * i (t -1) w * i (t) w * i (t + 1) w * i (t + 2) g i (t -2) g i (t -1) g i (t) g i (t + 1) g i (t + 2) µ i (t -2) µ i (t -1) µ i (t) µ i (t + 1) µ i (t + 2) C z i (t -2) z i (t -1) z i (t) z i (t + 1) z i (t + 2) g i (t -2) g i (t -1) g i (t) g i (t + 1) g i (t + 2) µ i (t -2) µ i (t -1) µ i (t) µ i (t + 1) µ i (t + 2)

Figure 1 : The heirarchy of generative models underlying our updates. A Full model for the gradients for a single parameter. The current estimate for all the other parameters, µ -i (t) vary slowly over time, and give rise to the current optimal value for the ith parameter, w * i . The gradient then arises from the current estimate of the ith parameter, µ i (t) (which is treated as an input here), and the optimal value, ith parameter, w * i . B The graphical model obtained by integrating over trajectories for the other parameter estimates, µ -i (t). In practice, we use a simplified model as reasoning about all possible trajectories of µ -i (t) is intractable. C To convert the model in B into a tractable hidden Markov model (HMM), we define a new variable, z i (t), which incorporates w * i along with other information about the dynamics.

How might we obtain a principled Bayesian filtering approach that recovers the two key features of state-of-the-art adaptive SGD algorithms: momentum and the root-mean-square normalizer? Here, we note that past approaches including natural gradient VI take a complex generative model over all N parameters jointly, and use a very strong approximation: factorisation. Given that we know that the true posterior is a highly complicated, correlated distribution, it is legitimate to worry that these strong approximations might meaningfully disrupt the ability of Bayesian filtering to give closeto-optimal updates. Here we take an alternative approach, baking factorisation into our generative model, so that we can use Bayesian inference to reason about (Bayes) optimal updates under the constraints imposed by factorisation. In particular, we split up the single large inference problem over all N parameters, w, into N small inference problems over a single parameter. Remarkably, by incorporating factorisation into the problem setting, we convert intractable, high-dimensional correlations in the original posterior into tractable low-dimensional dynamics in the factorised model. This dynamical prior has a "natural" form, at least compared with Ollivier (2017) , in that it does not depend on the posterior. Next, we give a generic derivation showing that Bayesian SGD is an adaptive SGD method, where the uncertainty is used to precondition the gradient. We then adapt the generic derivation to the two cases of interest: RMSprop (Hinton et al., 2012) and Adam (Kingma & Ba, 2015) . Finally, we discuss the general features of Bayesian adaptive SGD methods, including AdamW (Loshchilov & Hutter, 2017) and Nesterov acceleration (Nesterov, 1983; Dozat, 2016) , amongst others.

2. FACTORISATION IMPLIES A RICH DYNAMICAL PRIOR:

The most obvious factorised approach is to compute the true posterior, P (w i |D), over a single weight, w i , conditioned on all the data, D. However, this approach immediately fails, because symmetries involving permuting the hidden units (Sussmann, 1992) imply that marginal posteriors such as P (w i |D) are always so broad as to be useless.

To obtain an informative, narrow posterior, one approach is to eliminate these symmetries by conditioning on our current estimate of the other parameters. In particular, we define a random variable, w * i , representing the optimal value for the ith weight, conditioned on the other weights, w -i , being equal to our current estimate of those parameters, µ -i ,

EQUATION

where d is a random minibatch drawn from the true, unknown underlying data distribution (and not just the data we have), so this is analogous to maximum-likelihood with infinite data. As such, even if µ -i is fixed, w * i will be unknown if only have finite data, and therefore we do not know the true underlying data distribution. The optimal weight, w * i , is a random variable because it depends on our current estimate of the other weights, µ -i , which is a random variable because it depends on a random initialization, and potentially on a randomised optimization algorithm. However, we cannot now use the usual Bayesian approach of inferring w * i based on the data, because, in the Bayesian setting, the data are assumed to be generated from a model with parameters, w, not w * i . Instead, note that the updates for almost all neural network optimization algorithms depend only on the current value of that parameter, and the (history of) backpropagated gradients. This suggests a potential approach: writing down a generative model for the backpropagated gradients that depends on w * i , then inverting that model to infer w * i from those gradients. Following standard Laplace/Extended Kalman Filter-like approximations (e.g. Zhang et al., 2017; Khan et al., 2017; Ollivier, 2017; Khan et al., 2018) as closely as possible, we approximate the likelihood as a second-order Taylor series expansion. In particular, we consider the log-likelihood for a single minibatch, d, where the minibatch is treated as a random variable drawn from the underlying true data distribution, and we expand in the ith direction, keeping the other parameters, w -i , fixed to their current estimates, µ -i ,

EQUATION

where,

EQUATION

EQUATION

Here, both the mode, µ like,i , and the negative Hessian, Λ like,i , are random variables as they are deterministic functions of the minibatch, d, and the current estimate of all the other parameters, µ -i , which are both treated as random variables. As such, the gradient of the log-likelihood, which depends on Λ like,i and µ -i , is also a random variable,

EQUATION

)

where the approximation comes from the second-order Taylor expansion in Eq. (2). Our goal is to understand the distribution of the gradients, conditioned on the random variable Λ like,i being equal to some specific value, λ like,i . In practice, we will set Λ like,i , using the standard approximations in the natural gradient literature (Zhang et al., 2017; Khan et al., 2017; 2018) . The expected gradient is zero at

µ i = w * i , 0 = E [g i |w * i = w * i ] µi=w * i ≈ E [Λ like,i µ like,i |w * i = w * i ] -E [Λ like,i |w * i = w * i ] w * i (6)

where the approximation again comes only from the second-order Taylor expansion. Thus,

EQUATION

with equality if the expected value of µ like,i is independent of Λ i , Substituting this into Eq. ( 5), we obtain,

EQUATION

However, obtaining an analytic form for the full distribution of the gradient is more difficult. Given that all work in this domain makes fairly strong assumptions (e.g. Zhang et al., 2017; Khan et al., 2017; 2018) , an alternative approach opens up: choosing a model for the gradients to match as closely as possible any given approximation of the original likelihood in Eq. ( 2). In particular, we take,

EQUATION

Note that the variance is the negative Hessian, λ like,i , which is reminiscent of Fisher-Information like results (Sec. 9.7). However, we chose this form such that, for any given approximation to the original log-likelihood (Eq. 2), the log-likelihood induced by conditioning on the gradient (Eq. 9) has the same form,

EQUATION

)

The only material difference between the right-hand-side of this expression and the original likelihood (Eq. 2) is that here the parameter of interest is the optimal weight, w * i , as opposed to the weight in the underlying generative model, w i . Also note that here we have specified the values for all the random variables: g i = g i and Λ like,i = λ like,i which fix the value of µ like,i to µ like = gi λlike,i + µ i . The key difference between the methods is that in the original problem, where we infer a single distribution over all weights jointly, the underlying weights, w, were fixed, and as such, Bayesian filtering reduces to recursively applying Bayes theorem, which does not give rise to interesting dynamics. In contrast, when we consider N separate inference problems, in which w * i is inferred from the gradient, then we are forced to introduce dynamics, because w * i varies over time, as it depends on all the other parameters in the network, µ -i , which are also being optimized (Fig. 1A ). However, the full generative process in Fig. 1A is very difficult to reason about, as it requires us to integrate over all possible trajectories for the other parameters, µ -i . Instead, we write down simplified approximate dynamics over w * i directly (Fig. 1B ), and evaluate the quality of these simplified dynamics empirically. We choose the model in Fig. 1 such that it has finite-length temporal dependencies, and hence can be written in terms of a Markov model with an expanded state-space. In particular, we define a new random variable, z i , incorporating w * i , whose generative process is Markovian (Fig. 1C ). For Adam, z i will have two elements representing a parameter, w * i and the associated momentum, p i , whereas for RMSprop it will have only one element representing just the parameter,

EQUATION

As z i represents the optimal setting for a parameter it will change over time, and we assume a simple Gaussian form for these changes,

EQUATION

)

where Q is the covariance of Gaussian perturbations and A is the dynamics matrix incorporating weight decay and momentum. We can write the second-order Taylor expansion of the likelihood as a function of the expanded latent, z i (Eq. 11),

EQUATION

where we have omitted conditioning on Λ like,i , as in practice Λ like,i is estimated from the gradients, and we have omitted time indices for clarity,

g i = g i (t) w * i = w * i (t) z i = z i (t) µ like,i = µ like,i (t) Λ like,i = Λ like,i (t)

. This likelihood is equivalent to our original likelihood (Eq. 10) because the gradients depend only on w * i (t). As such, the negative Hessians must take the form,

EQUATION

)

3. BAYESIAN (KALMAN) FILTERING AS ADAPTIVE SGD:

The Gaussian prior and approximate likelihood allows us to use standard two-step Kalman filter updates. First, we propagate the previous time-step's posterior, P (z i (t -1)|D(t -1)), with mean µ post,i (t -1) and covariance Σ post,i (t -1), forward in time to obtain a prior at time t,

P (z i (t)|G(t -1)) = dz i (t -1) P (z i (t)|z i (t -1)) P (z i (t -1)|G i (t -1)) = N (z i (t); µ prior,i (t), Σ prior,i (t)) ,

where,

EQUATION

EQUATION

and where G i (t -1) = {g i (1), g i (2), . . . , g i (t -1)} is all gradients up to time t -1. Note that we have also defined abbreviations for µ prior,i (t) and Σ prior,i (t), omitting the temporal index, t. Second we use Bayes theorem to incorporate new data,

P (z i (t)|D(t)) = P (g i (t)|z i (t)) P (z i (t)|G(t -1)) P (g i (t)|G(t -1)) = N (z i (t); µ post,i (t), Σ post,i (t)) .

where,

EQUATION

EQUATION

Thus far, we have simply repeated standard Kalman filtering results. To relate Kalman filtering to gradient ascent, we compute the gradient of Eq. ( 13) at z i = µ prior,i ,

EQUATION

)

Note that as the log-likelihood depends only on w * i , we have,

EQUATION

Now, we identify this gradient (Eq. 17) in the mean updates (Eq. 16b),

EQUATION

)

This form is extremely intuitive: it states that the uncertainty should be used to precondition gradient updates, such that the updates are larger when there is more uncertainty, and smaller when past data gives high confidence in the current estimates.

As the precision is always rank-1 (Eq. 14), we can always write it as,

EQUATION

)

As such, the updates for the posterior covariance (Eq. 16a) can be rewritten using the Sherman Morrison formula (Hager, 1989) ,

EQUATION

)

To estimate e i , we use the Fisher information (see SI Sec. 9.7). In particular, we could use the gradient itself, but could also (under weaker approximations) use the centred gradient (Graves, 2013) , or the gradient for data sampled under the prior.

4. BAYESIAN RMSPROP (BRMSPROP):

Here, we develop a Bayesian variant of RMSprop, which we call BRMSprop. We consider each parameter to be inferred by considering a separate Bayesian inference problem, so the latent variable, z = w * , is a single scalar, representing a single parameter (we omit the index i on z for brevity). We use A = η 2 /(2σ 2 ) and Q = η 2 giving a dynamical prior,

EQUATION

)

such that the stationary distribution over z has standard-deviation σ. To obtain a complete description of our updates, we substitute these choices into the updates for the prior (Eq. 15) and the posterior (Eq. 19 and Eq. 21),

EQUATION

EQUATION

For an efficient implementation of the full algorithm, see SI Alg. 1.

4.1. RECOVERING RMSPROP FROM BRMSPROP:

Now we show that with in steady state, BRMSprop closely approximates RMSprop. Making this comparison is non-trivial because the "additional" variables in RMSprop and BRMSprop (i.e. the average squared gradient and the uncertainty respectively) are not directly comparable. However, the implied learning rate is directly comparable. We therefore look at the learning rate when the average squared gradient and the uncertainty have reached steady-state. As t → ∞, we expect σ 2 post to reach steady-state, at which point, σ 2 post = σ 2 post (t) = σ 2 post (t + 1). In SI Sec. 9.1, we consider steady-state in the general case. Applying those results to RMSprop, we obtain,

σ 4 post (t + 1) e 2 ≈ η 2 . Solving for σ 2 post , σ 2 post ≈ η e 2 .

Substituting this form into Eq. ( 23) recovers the root-mean-square normalization used in RMSprop.

5. BAYESIAN ADAM (BADAM):

We now develop a Bayesian variant of Adam (Graves, 2013; Kingma & Ba, 2015) , which we call BAdam. To introduce momentum into our Bayesian updates, we introduce an auxiliary momentum variable, p(t), corresponding to each parameter, w * (t),

z(t) = w * (t) p(t) ,

then we infer both the parameter and momentum jointly. Under the prior, the momentum, p(t), evolves through time independently of w * (t), obeying an AR(1) (or equivalently a discretised Ornstein-Uhlenbeck) process, with decay η p and injected noise variance

EQUATION

where ξ p (t) is standard Gaussian noise. This particular coupling of the injected noise variance and the decay rate ensures that the influence of the momentum on the weight is analogous to unitvariance Gaussian noise (see SI Sec. 9.2). The weight obeys similar dynamics, with the addition of a momentum-dependent term which in practice causes changes in w * i (t) to be correlated across time (i.e. multiple consecutive increases or decreases in w * i (t)), w * (t + 1) = 1 -

EQUATION

where ξ w (t) is again standard Gaussian noise, η is the strength of the momentum coupling,

η 2 +η 2 w 2σ 2

is the strength of the weight decay, and η 2 w is the variance of the noise injected to the weight. It is possible to write these dynamics in the generic form given above (Eq. 12), by using,

EQUATION

and these settings fully determine the updates, according to Eq. ( 15) and Eq. ( 16). For an efficient implementation of the full algorithm, see SI Alg. 2.

5.1. RECOVERING ADAM FROM BADAM:

Now we show that with suitable choices for the parameters, BAdam closely approximates Adam. In particular, we compare the updates for the (posterior) parameter estimate, µ w , and the momentum, µ p , when we eliminate weight decay by setting σ 2 = ∞, and eliminate noise injected directly into the weight by setting η w = 0,

EQUATION

EQUATION

These updates depend on two quantities, Σ ww and Σ wp , which are related to e 2 , but have no direct analogue in standard Adam. As such, to make a direct comparison, we use the same approach as we used previously for RMSprop: we compare the updates for the parameter and momentum, when Σ ww , Σ wp in BAdam and e 2 in Adam have reached their steady-state values. To find the steadystates for Σ ww and Σ wp , we again use the simplified covariance updates derived in SI Sec. 9.1,

EQUATION

Substituting Eq. ( 27) and Eq. ( 20) into Eq. ( 29), and again using σ 2 = ∞ and η 2 w = 0, we obtain,

EQUATION

)

We now assume that the data is informative, in the sense that it is strong enough to give a narrow posterior relative to the prior (without which any neural network training algorithm is unlikely to be able to obtain good performance). This implies Σ pp η p /2 (see SI Sec. 9.2), allowing us to solve for Σ wp , using the lower-right element of the above expression (i.e. η 2 p ≈ 2η p Σ pp + e 2 Σ 2 wp ),

EQUATION

)

To recover a very close approximation to Adam, we need a specific relationship between learning rates and evidence strength, such that,

Σ ww ≈ ηη p e 2 .

While this may seem restrictive, it is only necessary to achieve the closest possible match between Bayesian filtering and plain Adam. However, our goal is not to match Adam exactly, given that Adam does not even converge (Reddi et al., 2018) . Instead, our goal is to capture the essential dynamical insights of Adam in a Bayesian method. Indeed, we hope that the resulting Bayesian method constitutes an improvement over Adam, which implies that it must exhibit some differences. Nonetheless, focusing on the regime where filtering and Adam are most similar, the filtering updates become,

µ w (t + 1) ≈ µ w (t) + ηµ p (t) + ηη p g(t)

e 2 ≈ µ w (t) + ηµ p (t + 1) (32a)

EQUATION

)

where we have substituted for µ p (t + 1) into Eq. 32a, which assumes that η p 1. This is very close to the Adam updates, except that the root-mean-square normalization is in the momentum updates, rather than the parameter updates. To move the location of the root-mean-square normalization, we rewrite the updates in terms of a rescaled momentum,

EQUATION

giving,

EQUATION

)

EQUATION

which recovers Adam. See SI Sec. 9.6 for a similar discussion for NAG/NAdam

6. EXPERIMENTS:

We compared Bayesian and standard methods on MNIST. In particular, we trained a CNN with relu nonlinearities and maxpooling for 50 epochs. The network had two convolutional layers with 10 channels in the first layer and 20 in the second, with 5 × 5 convolutional kernels, followed by a single fully connected layer with 50 units, and was initialized with draws from a Gaussian with variance 2/N inputs . The network did not use dropout (or any other form of stochastic regularisation), which would increase the variance of the gradients, artificially inflating e 2 .

The key Bayesian-specific parameters are the variance of the stationary distribution and the initial uncertainty; both of which were set to 1/(2N inputs ). In principle, they should be similar to the initialization variance, but in practice we found that using a somewhat lower value gave better performance, though this needs further investigation. For the RMSprop and Adam specific parameters, we used the PyTorch defaults.

Comparing the Bayesian and non-Bayesian methods, we do not expect to see very large discrepancies, because they are approximately equivalent in steady-state. Nonetheless, the Bayesian methods show considerably lower test loss, somewhat lower classification error, and similar training loss and training error (Fig. 2 ). The local maximum in test-loss may arise because we condition on each datapoint multiple times, which is not theoretically justified (correcting this is non-trivial in the dynamical setting, so we leave it for future work).

7. FEATURES OF BAYESIAN STOCHASTIC GRADIENT DESCENT:

Here we summarise the features of Bayesian stochastic gradient descent schemes, pointing out where we recover previously known best practice and where our approach suggests new ideas.

7.1 WEIGHT DECAY, L2 REGULARIZATION AND BAYESIAN PRIORS Loshchilov & Hutter (2017) recently examined different forms of weight decay in adaptive SGD methods such as Adam. In particular, they asked whether or not the root-mean-square normalizer should be applied to the weight decay term. In common practice, we take weight decay as arising from the gradient of an L2 regularizer, in which case it is natural to normalize the gradient of the objective and the gradient of the regularizer in the same way. However, there is an alternative: to normalize only the gradient of the objective, but to keep the weight decay constant (in which case, weight decay cannot be interpreted as arising from an L2 normalizer). Loshchilov & Hutter (2017) show that this second method, which they call AdamW, gives better test accuracy than the standard approach. Remarkably, AdamW arises naturally in BRMSprop and BAdam (e.g. Eq. 23), providing a potential explanation for its improved performance.

7.2. NESTEROV ACCELERATED GRADIENTS/WEIGHT DECAY:

In Nesterov accelerated gradients (NAG), we compute the gradient at a "predicted" location formed by applying a momentum update to the current parameters (Nesterov, 1983) . These updates arise naturally in our Bayesian scheme, as the required gradient term (Eq. 17) is evaluated at µ prior (Eq. 15a), which is precisely a prediction formed by combining the current setting of the parameters, µ post (t), with momentum and decay, embodied in A (Eq. 27) to form a prediction, µ prior (t + 1). It should be noted that the original Nesterov acceleration (Nesterov, 1983) had no gradient preconditioning. Instead, our approach corresponds to NAdam (Dozat, 2016) . Interestingly, as we also implement weight decay through the dynamics matrix, A, we should also apply the updates from weight decay before computing the gradient, giving, to our knowledge, a novel method that we denote "Nesterov accelerated weight decay" (NAWD).

7.3. CONVERGENCE:

A series of recent papers have discussed the convergence properties of RMSprop and Adam, noting that they may fail to converge (Wilson et al., 2017; Reddi et al., 2018) if the exponentially decaying average over the squared gradients is computed with a too-small timescale (Reddi et al., 2018) .

Our method circumvents these issues by coupling the learning rate to the timescale over which the implicit exponential moving average for the mean-square gradient is performed. As such, in the limit as the learning rates go to zero (i.e. η → 0 for BRMSprop and η w → 0 and η p → 0 for BAdam), our method becomes SGD with adaptive learning rates that scale as 1/t and is therefore likely to be convergent for convex functions (Robbins & Monro, 1951; Bottou, 1998) , though we leave a rigorous proof to future work.

7.4. UPDATING THE PRECONDITIONER BEFORE APPLYING THE UPDATE:

ADAM (Kingma & Ba, 2015) first updates the root-mean-square gradient normalizer before computing the parameter update. This is important, because it ensures that updates are bounded in the pathological case that gradients are initially all very small, such that the root-mean-square normalizer is also small, and then there is a single large gradient. Bayesian filtering naturally recovers this choice as the gradient preconditioner in Eq. ( 19) is the posterior, rather than the prior, covariance (i.e. updated with the current gradient).

8. DISCUSSION:

Here, we developed BRMSprop and BAdam, Bayesian versions of RMSprop and Adam respectively. These methods naturally recover many features of state-of-the-art adaptive SGD methods, including the root-mean-square normalizer, Neseterov acceleration and AdamW. As such, BRM-Sprop and BAdam provide a possible explanation for the empirical success of RMSprop, Adam, Nesterov acceleration and AdamW. Experimentally, we find, BRMSprop and BAdam are superior to the standard methods in the tasks that we examined. Finally, the novel interpretation of adaptive SGD as filtering in a dynamical model opens several avenues for future research, including new stochastic regularisation methods and combinations with Kronecker factorisation.

We believe that it should be possible to combine the temporal changes induced by the factorised approximations with other types of inference, including probabilistic backpropagation (Hernández-Lobato & Adams, 2015) or assumed density filtering (Ghosh et al., 2016) . At present, our method bears closest relation to natural gradient variational inference methods (Zhang et al., 2017; Khan et al., 2018) , as they also use the Fisher Information to approximate the likelihood. Indeed, our method becomes equivalent to these approaches in the limit as we send the learning rate, η to zero.

The key difference is that because they do not formulate a factorised generative model, they are unable to provide a strong justification for the introduction of rich dynamics, and they are unable to reason about optimal inference under these dynamics.

Bayesian filtering presents a novel approach to neural network optimization, and as such, there are variety of directions for future work. First, Bayesian filtering converts the problem of neural network optimization into the statistical problem of understanding the dynamics of changes in the optimal weight induced by optimization in the other parameters. In particular, we can perform an empirical investigation in large scale systems, or attempt to find closed-form expressions for the dynamics in simplified domains such as linear regression. Second, here we wrote down a statistical model for the gradient. However, there are many circumstances where the gradient is not available. Perhaps a low precision or noisy gradient is available due to noise in the parameters (e.g. due to dropout Srivastava et al., 2014) , or perhaps we wish to consider a biological setting, where the gradient is not present at all (Aitchison et al., 2014) . The Bayesian approach presented here gives a straightforward recipe for developing (Bayes) optimal algorithms for such problems. Third, stochastic regularization has been shown to be extremely effective at reducing generalization error in neural networks. This Bayesian interpretation of adaptive SGD methods presents opportunities for new stochastic regularization schemes. Fourth, it should be possible to develop filtering methods that represent the covariance of a full weight matrix by exploiting Kronecker factorisation (Martens & Grosse, 2015; Grosse & Martens, 2016; Zhang et al., 2017) .

9. SUPPLEMENTARY INFORMATION:



9.1. STEADY-STATE COVARIANCE:

To understand the steady-state behaviour of the covariance, we begin by substituting Eq. 21 into Eq. 15b to form a combined update, Σ prior (t + 1) = (I -A) Σ prior (t) -Σ prior (t)ee T Σ prior (t) 1 + e T Σ prior (t)e T (I -A)

EQUATION

)

Note that this gives updates for Σ prior (t), rather than Σ post (t), as this simplifies the algebra, but we will move back to Σ post (t) at the end. As t → ∞, we reach steady state for which, Σ prior = Σ prior (t) = Σ prior (t + 1). Substituting this asympototic covariance into Eq. ( 35) and rearranging, we obtain,

EQUATION

Now our goal is to find the asymptotic scaling of Σ prior with η so as to obtain a valid approximation to this steady-state expression. To begin, we know that,

A ∈ O (η) Q ∈ Θ η 2 ,

where O denotes an asymptotic upper bound, and Θ denotes an asymptotic bound on both sides (Knuth, 1976) . Here, η is the learning rate, which is assumed to be small, and hence we take the limit of η → 0 for the purposes of the limits in O and Θ. Note that for Adam we have A Adam ∈ Θ(η) but for RMSprop we have A RMSprop ∈ Θ(η 2 ), and we summarise both using the asymptotic upper bound above.

The right-hand-side of Eq. ( 36) is thus bounded above,

AΣ prior A T -AΣ prior -Σ prior A T + Q 1 + e T Σ prior e T ∈ O η 2 + ηΣ(η) + ηΣ 2 (η) (37)

where Σ(η) describes the scaling of Σ prior ,

Σ prior ∈ Θ(Σ(η)).

The left-hand-side can be asymptotically bounded on both sides as lim η→0 A = 0,

EQUATION

The asymptotic upper bound (Eq. 37) for the right-hand-side of Eq. ( 36) must bound the asymptotic lower bound for the left-hand-side (Eq. 38),

EQUATION

Now, we consider the case in which each of the terms η 2 , ηΣ(η) or ηΣ 2 (η) is asymptotically dominant (potentially along with other terms). The ηΣ 2 (η) term cannot dominate, because it cannot form an asymptotic upper bound on Σ 2 (η). As such, Σ 2 (η) could be asymptotically bound either by η 2 , or by ηΣ(η) (or both simultaneously). If it is bound by η 2 , we have,

EQUATION

Alternatively, if it is bound by ηΣ(η), we have,

EQUATION

And as it must be bound by one (or both) of these terms, we know that Σ(η) and hence Σ prior are bounded by η,

EQUATION

)

Now that we have the scaling of Σ prior , deriving a simplified approximation to Eq. ( 36) is straightforward. In particular, we consider only terms that scale with η 2 , neglecting those that scale with η 3 . We now consider terms to be eliminated in Eq. ( 36). In particular, on the left-hand-side,

EQUATION

EQUATION

EQUATION

and on the right-hand-side,

EQUATION

Q -AΣ prior -Σ prior A T e T Σ prior e ∈ O(η 3 ).

Neglecting these terms, and instead focusing only on the O(η 2 ) terms we obtain,

EQUATION

However, for the purposes of our updates, we are interested in the asymptotic behaviour of Σ post , not Σ prior . We therefore substitute for Σ prior from Eq. (15b) in Eq. ( 44), and neglect third-order terms. First, we consider the the left-hand-side of Eq. ( 44),

EQUATION

And we do the same for the right-hand-side of Eq (44),

EQUATION

Combining these approximations to the left and right hand sides, we obtain the result used in the main text for Adam.

EQUATION

For RMSprop, remember that A RMSprop ∈ O(η 2 ), so AΣ post ∈ O(η 3 ), and we can obtain an even simpler result,

EQUATION

)

These equations bear a remarkable resemblance to the continuous time Kalman filter that emerges when you take the limit of small timesteps (Xie et al., 2007) .

Under review as a conference paper at ICLR 2019

9.2. SETTING THE MOMENTUM DECAY:

In the main text, we briefly note that we couple the injected noise and decay in Eq. ( 25) such that the effect of p(t) on w(t) is analogous to that of unit-variance Gaussian noise. In particular, we require that the stationary variance of w(t) is again σ 2 , despite the momentum term in Eq. ( 26) being treated as if it was unit-variance Gaussian noise. To show that the stationary variance of w(t) is indeed σ 2 , we solve the Lyapunov equation for the stationary covariance, Ψ,

Ψ = (I -A) Ψ (I -A) T + Q,

where we do not condition on data. Neglecting the second-order term in A,

EQUATION

and as,

AΨ = η 2 +η 2 w 2σ 2 -η 0 η p Ψ ww Ψ wp Ψ wp Ψ pp = η 2 +η 2 w 2σ 2 Ψ ww -ηΨ wp η 2 +η 2 w 2σ 2 Ψ wp -ηΨ pp η p Ψ wp η p Ψ pp .

Eq. ( 49) becomes,

EQUATION

Now, we can solve for Ψ pp using the ppth element of Eq. ( 50),

Ψ pp = η p 2 ,

and we can solve for Ψ wp using the wpth element of Eq. ( 50), in combination with our solution to Ψ pp ,

Ψ wp = ηη p /2 η p + η 2 +η 2 w 2σ 2 ≈ η/2.

This approximation becomes exact in the limit we consider to recover Adam (i.e. σ 2 = ∞), but it also holds when σ 2 is finite, because η, η w and η p are all smaller than 1, and η p -the "learning rate" for the momentum -is usually larger than η and η w (in fact, η p can be as high as 0.1). Finally, we can solve for Ψ ww using the wwth element of Eq. ( 50),

Ψ ww ≈ σ 2 ,

as required.

9.3. EFFICIENT IMPLEMENTATION OF BRMSPROP:

See Alg. 1 for a complete specification of our approach. There are two important points to note. First, the method is as memory efficient as RMSprop, as the updates for Σ and µ can be done inplace, so there is no need to maintain separate µ prior and µ post and σ 2 prior and σ 2 post in memory. Second, we simplified the covariance updates in Eq. ( 24), (51)

9.4 EFFICIENT IMPLEMENTATION OF BADAM See Alg. 2 for a complete specification of our approach. There are two important points to note. First, we have carefully ordered the updates so that they can all be done in-place, improving memory efficiency. Second, we have simplified the derivation, by defining, ĝ ← (1 -η g ) ĝ + η g g Update average gradient 9:

EQUATION

Σ ← 1 -η 2 /(2σ 2 ) 2 Σ + η 2 Σ = σ 2 prior (t + 1) 10:

µ ← 1 -η 2 /(2σ 2 ) µ µ = µ prior (t + 1) 11: end while 12: return µ

9.5. TRICK TO ALLOW LEARNING RATES TO VARY ACROSS PARAMETERS:

We require that the initialization for Σ varies across parameters, according to the number of inputs (as in typical neural network initialization schemes). While this is possible in automatic differentiation libraries, including PyTorch, it is extremely awkward. As such, we reparameterise the network, such that all parameters are initialized with a draw from a standard Gaussian, and to ensure that the outputs have the same scale, we explicitly scale the output.

9.6. NESTEROV ACCELERATED GRADIENT:

To begin, we note that accelerated stochastic gradient descent remains an open research problem, in which the optimality of plain NAG is unclear (Cohen et al., 2018) . As such, our goal is again, not to match NAG exactly, but to capture its essential insights within a Bayesian framework, so that we can suggest improved methods. To highlight the link between our approach and Nesterov accelerated gradient, we rewrite Eq. ( 32) in terms of v(t + 1) = ηµ p (t + 1), µ w (t) ≈ µ w (t -1) + v(t) (53a)

EQUATION

)

which results in updates with a very similar form to standard momentum and Nesterov accelerated gradient (e.g. Sutskever et al., 2013) , with the addition of root-mean-square normalization for the gradient. The key difference between momentum and Nesterov accelerated gradient is where we evaluate the gradient: for momentum, we evaluate the gradient at µ w (t -1) (i.e. at µ post (t -1)), and for Nesterov accelerated gradient, we evaluate the gradient at a "predicted" location (i.e. at µ prior (t)),

g Mom (t) = g (µ w (t -1)) = g (µ post, w (t -1)) (54a)

g NAG (t) = g (µ w (t -1) + (1 -η p )v(t -1)) ≈ g (µ w (t -1) + v(t -1)) = g (µ prior, w (t)) (54b)

As noted in Eq. ( 17), Bayesian filtering requires us to evaluate the gradient at µ prior , implying that we use updates based on NAG (Eq. 54b), rather than updates based on standard momentum (Eq. 54a).

9.7. ESTIMATING THE NEGATIVE HESSIAN:

One approach that works well in practice in natural-gradient methods (Amari, 1998; Martens & Grosse, 2015; Grosse & Martens, 2016; Song & Ermon, 2018; Smith et al., 2018) is to use a Fisher-

:



Figures and Tables:

Figure 2: Comparing RMSprop and BRMSprop on MNIST. The vertical bars represent one standard error and are based on 30 independent runs. Adam and BAdam both use η w = 0.1.

)

Table content: <table><tr><td colspan="2">Algorithm 1 BRMSprop</td><td/></tr><tr><td colspan="2">Require: η</td><td>Learning rate</td></tr><tr><td colspan="2">Require: η g</td><td>Learning rate for average gradient (0.01)</td></tr><tr><td colspan="2">Require: σ 2</td><td>Prior variance (1/(2N inputs ))</td></tr><tr><td colspan="2">1: Σ ← σ 2</td><td>Initialize uncertainties</td></tr><tr><td colspan="2">2: µ ∼ N (0, 2/N inputs )</td><td>Initialize parameters</td></tr><tr><td colspan="2">3: ĝ ← 0</td><td>Initialize average gradient</td></tr><tr><td colspan="2">4: while not converged do 5: g ← ∇L t (µ)</td><td>Σ = σ 2 prior (t), µ = µ prior (t) Gradients of minibatch t</td></tr><tr><td>6:</td><td>Σ ← Σ/(1 + (g -ĝ) 2 Σ)</td><td>Σ = σ 2 post (t)</td></tr><tr><td>7:</td><td>µ ← µ + Σg</td><td>µ = µ post (t)</td></tr><tr><td>8:</td><td/><td/></tr></table>

Peer review (iclr2019):
 * Description The paper considers the following random process on the parameters z (modeled as Gaussians): - shrink z towards zero and add Gaussian i.i.d. noise to it. - update the parameters to the posterior w.r.t. a batch, where the likelihood is approximated as a diagonal multivariate normal distribution. This results in a Kalman filter like updates. There have been related methods proposed performing Bayesian learning in the form of assumed density filtering, considered as separate learning algorithms. At the same time methods such as RMSprop and Adam were previously derived from completely different considerations. The work can derive these methods in the Bayesian framework with certain additional assumptions / simplifications. It allows to naturally explain tracking the gradient statistics as uncertainties and the normalization of the gradient in the existing methods as the update of the mean parameters in the Kalman filter taking into account these uncertainties. The experiments on MNIST show that derived more Bayesian variants of RMSprop and Adam can improve generalization in terms of test likelihood and test error. * Assessment The provided derivation of Bayes like learning algorithms is relatively simple and could be very useful in practice and in further improvement of the learning methods. The approximations used are not completely clear. The clarification of the idea of a separate optimization problem per variable is necessary. The provided experiments, if there is nothing subtle, are clearly done and would be sufficient. There are some open questions such as: does the method in fact learn useful variances of the parameters, i.e. really performs an approximate Bayesian learning? Overall if find it a promising novel research direction of high practical relevance. * Clarity Intro: Why is the unnumbered equation on page 1 is called a “Bayesian optimization problem”? There is so many sings called Bayesian that one cannot be sure what it means. In the context of the paper it should be a Bayesian learning problem, but I do not see a posterior distribution over the parameters. Overall, I did not get the point of the discussion in the introduction and Figure 1 altogether. Everything it says to me is that global minimize coordinates are dependent through the objective. I do not see what the unnumbered equation on page 1 has to do with Bayesian inference and how the correlation of parameters in the posterior distribution is related to the dependencies in the minimizer. Could authors please seriously consider clarifying this section? In what follows the paper keeps a factorize approximation to the posterior of parameters of a NN in the form of a Gaussian distribution per coordinate. It thus does not in any way avoid making this restrictive assumption. Results: Sorry, I am not familiar with the background behind (6). Which value of z is assumed in the conditional expectation, is it conditioning on “z = .mu_{prior}”? How come the approximation to the variance of the data likelihood does not depend on the data? If we make this approximation, how much it is still relevant to the Bayesian learning? What are the overheads of the proposed methods? I expect they scale as easily to large problems as SGD? * Experiments From Figure 2 it seems that BRMSprop and BAdam can achieve relatively good results for large range of eta in 10^-5 to 10^-2 and it seems from the trend that even smaller eta would work. Does it mean they do not need in fact tuning of the learning rate? The experiment uses 50 epochs, do the compared methods reach the convergence? Could the authors consider an experiment running best setting of parameters per method with twice as many epochs? Some artificial toy experiments could be of interest. For example, consider a classification problem with a 1D Gaussian data distribution in each class and the logistic regression model with 2 parameters. Does the method approximate the posterior distribution? * Related work The approach to Bayesian learning taken in the paper needs to be better discussed. I think it is from the family of methods known as “assumed density filtering”, occurring in: Ghosh et al. “Assumed Density Filtering Methods for Scalable Learning of Bayesian Neural Networks” with earlier works well described in Minka T. “Expectation propagation for approximate Bayesian inference”. In particular equation (5) of the submission is well known. The work Khan et al. 2018 “Fast and scalable Bayesian deep learning by weight-perturbation in Adam” also derives Bayesian learning algorithms in the forms closely similar to RMSprop and Adam and interprets the running statistics as uncertainties. However it takes the variational Bayesian learning approach, which means the reverse KL divergence is used somewhere. Could the authors discuss conceptual similarities and differences to this work?