Year: 2018

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 RMSprop and Adam on MNIST, we find that Bayesian methods have the potential to considerably reduce test loss and classification error.

text: 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 RMSprop and Adam on MNIST, we find that Bayesian methods have the potential to considerably reduce test loss and classification error.

heading: 1 INTRODUCTION
text: 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 Qt proportional to [the posterior covariance] Pt−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.
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 allN parameters, w, intoN 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.

heading: 2 FACTORISATION IMPLIES A RICH DYNAMICAL PRIOR
text: The most obvious factorised approach is to compute the true posterior, P (wi|D), over a single weight, wi, 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 (wi|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,
w∗i = arg max wi E [log P (d|wi = wi,w−i = µ−i)] , (1)
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,
log P (d|wi = wi,w−i = µ−i) ≈ − 12Λlike,i (wi − µlike,i) 2 + const, (2)
where,
µlike,i = arg max wi
log P (d|wi = wi,w−i = µ−i) , (3)
Λlike,i = − ∂2 log P (d|wi = wi,w−i = µ−i)
∂2wi
∣∣∣∣ wi=µlike,i . (4)
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,
gi = ∂ log P (d|wi = wi,w−i = µ−i)
∂wi
∣∣∣∣ wi=µi ≈ Λlike,i (µlike,i − µi) . (5)
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 [gi|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,
E [µlike,i|Λlike,i = λlike,i,w∗i = w∗i ] ≈ E [µlike,i|w∗i = w∗i ] ≈ w∗i , (7)
with equality if the expected value of µlike,i is independent of Λi, Substituting this into Eq. (5), we obtain,
E [gi|Λlike,i = λlike,i,w∗i = w∗i ] ≈ λlike,i (w∗i − µi) . (8) 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,
P (gi = gi|Λlike,i = λlike,i,w∗i = w∗i ) ≈ N (gi;λlike,i (w∗i − µi) , λlike,i) . (9) 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,
log P (gi = gi|Λlike,i = λlike,i,w∗i = w∗i ) + const ≈
− (gi − λlike,i (w ∗ i − µi)) 2
2λlike,i = − 12λlike,i (w ∗ i − µlike,i) . (10)
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, wi. Also note that here we have specified the values for all the random variables: gi = gi 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, zi, incorporating w∗i , whose generative process is Markovian (Fig. 1C). For Adam, zi will have two elements representing a parameter, w∗i and the associated momentum, pi, whereas for RMSprop it will have only one element representing just the parameter,
zRMSprop,i(t) = (w ∗ i (t)) zAdam,i(t) = ( w∗i (t) pi(t) ) (11)
As zi represents the optimal setting for a parameter it will change over time, and we assume a simple Gaussian form for these changes,
P (zi(t)|zi(t− 1)) = N (zi(t); (I −A) zi(t− 1),Q) , (12) 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, zi (Eq. 11),
log P (gi|w∗i ) = log P (gi|zi) ≈ − 12 (zi − µlike,i) T Λlike,i (zi − µlike,i) , (13) 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,
gi = gi(t) w ∗ i = w ∗ i (t) zi = zi(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,
ΛRMSprop,like,i = (Λlike,i) ΛAdam,like,i =
( Λlike,i 0
0 0
) . (14)

heading: 3 BAYESIAN (KALMAN) FILTERING AS ADAPTIVE SGD
text: 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 (zi(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 (zi(t)|G(t− 1)) = ∫ dzi(t− 1) P (zi(t)|zi(t− 1)) P (zi(t− 1)|Gi(t− 1))
= N (zi(t);µprior,i(t),Σprior,i(t)) ,
where,
µprior,i = µprior,i(t) = (I −A)µpost,i(t− 1), (15a)
Σprior,i = Σprior,i(t) = (I −A) Σpost,i(t− 1) (I −A)T +Q, (15b)
and where Gi(t − 1) = {gi(1), gi(2), . . . , gi(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 (zi(t)|D(t)) = P (gi(t)|zi(t)) P (zi(t)|G(t− 1))
P (gi(t)|G(t− 1)) = N (zi(t);µpost,i(t),Σpost,i(t)) .
where,
Σpost,i = Σpost,i(t) = ( Σ−1prior,i + Λlike,i )−1 , (16a)
µpost,i = µpost,i(t) = µprior,i + Σpost,iΛlike,i (µlike,i − µprior,i) . (16b)
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 zi = µprior,i,
gi = ∂ log P (gi|z)
∂zi
∣∣∣∣ zi=µprior,i = Λlike (µlike,i − µprior,i) . (17)
Note that as the log-likelihood depends only on w∗i , we have,
gRMSprop,i = (gi) gAdam,i = ( gi 0 ) . (18)
Now, we identify this gradient (Eq. 17) in the mean updates (Eq. 16b),
µpost,i = µprior,i + Σpost,i gi. (19)
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,
Λlike,i = eie T i , where eRMSprop,i = (ei) eAdam,i = ( ei 0 ) (20)
As such, the updates for the posterior covariance (Eq. 16a) can be rewritten using the Sherman Morrison formula (Hager, 1989),
Σpost,i = Σprior,i − Σprior,ieie
T i Σprior,i
1 + eTi Σprior,iei . (21)
To estimate ei, 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.

heading: 4 BAYESIAN RMSPROP (BRMSPROP)
text: 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,
P (z(t+ 1)|z(t)) = N (( 1− η 2
2σ2
) z(t), η2 ) , (22)
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),
µprior = ( 1− η 2
2σ2
) µpost(t− 1) µpost = µprior + σ2postg, (23)
σ2prior = ( 1− η 2
2σ2
)2 σ2post(t− 1) + η2 σ2post = σ2prior ( 1− σ2priore 2
1 + σ2priore 2
) (24)
For an efficient implementation of the full algorithm, see SI Alg. 1.

heading: 4.1 RECOVERING RMSPROP FROM BRMSPROP
text: 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 σ2post to reach steady-state, at which point, σ2post = σ 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,
σ4post(t+ 1)〈e2〉 ≈ η2.
Solving for σ2post,
σ2post ≈ η√ 〈e2〉 .
Substituting this form into Eq. (23) recovers the root-mean-square normalization used in RMSprop.

heading: 5 BAYESIAN ADAM (BADAM)
text: 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 η2p ,
p(t+ 1) = (1− ηp) p(t) + ηp ξp(t), (25) 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− η 2+η2w 2σ2 ) w∗(t) + ηp(t) + ηw ξw(t), (26)
where ξw(t) is again standard Gaussian noise, η is the strength of the momentum coupling, η2+η2w 2σ2 is the strength of the weight decay, and η2w 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,
A =
( η2w+η 2
2σ2 −η 0 ηp
) Q = ( η2w 0 0 η2p ) , (27)
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.

heading: 5.1 RECOVERING ADAM FROM BADAM
text: 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,
µw(t+ 1) = µpost,w(t+ 1) = µw(t) + ηµp(t) + Σwwg(t), (28a) µp(t+ 1) = µpost,p(t+ 1) = (1− ηp)µp(t) + Σwpg(t). (28b)
These updates depend on two quantities, Σww and Σwp, which are related to 〈e2〉, 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 〈e2〉 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,
Q ≈ AΣ + ΣAT + Σ〈eeT 〉Σ. (29)
Substituting Eq. (27) and Eq. (20) into Eq. (29), and again using σ2 =∞ and η2w = 0, we obtain,( 0 0 0 η2p ) ≈ ηp ( 0 Σwp Σwp 2Σpp ) − η ( 2Σwp Σpp Σpp 0 ) + 〈e2〉 ( Σ2ww ΣwwΣwp ΣwwΣwp Σ 2 wp ) . (30)
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. η2p ≈ 2ηpΣpp + 〈e2〉Σ2wp),
Σwp ≈ ηp√ 〈e2〉 . (31)
To recover a very close approximation to Adam, we need a specific relationship between learning rates and evidence strength, such that,
Σww ≈ ηηp√ 〈e2〉 .
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)√ 〈e2〉 ≈ µw(t) + ηµp(t+ 1) (32a) µp(t+ 1) ≈ (1− ηp)µp(t) + ηp g(t)√ 〈e2〉 , (32b)
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,
µ̃p(t+ 1) = µp(t+ 1) √ 〈e2〉 (33)
giving,
µw(t+ 1) ≈ µw(t) + η µ̃p(t+ 1)√ 〈e2〉 , (34a)
µ̃p(t+ 1) ≈ (1− ηp) µ̃p(t) + ηpg(t), (34b) which recovers Adam. See SI Sec. 9.6 for a similar discussion for NAG/NAdam

heading: 6 EXPERIMENTS
text: 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/Ninputs. The network did not use dropout (or any other form of stochastic regularisation), which would increase the variance of the gradients, artificially inflating 〈e2〉. The key Bayesian-specific parameters are the variance of the stationary distribution and the initial uncertainty; both of which were set to 1/(2Ninputs). 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).

heading: 7 FEATURES OF BAYESIAN STOCHASTIC GRADIENT DESCENT
text: 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.

heading: 7.1 WEIGHT DECAY, L2 REGULARIZATION AND BAYESIAN PRIORS
text: 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.

heading: 7.2 NESTEROV ACCELERATED GRADIENTS/WEIGHT DECAY
text: 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).

heading: 7.3 CONVERGENCE
text: 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.

heading: 7.4 UPDATING THE PRECONDITIONER BEFORE APPLYING THE UPDATE
text: 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).

heading: 8 DISCUSSION
text: 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, BRMSprop 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 (HernándezLobato & 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).

heading: 9 SUPPLEMENTARY INFORMATION

heading: 9.1 STEADY-STATE COVARIANCE
text: 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 + eTΣprior(t)eT
) (I −A)T +Q. (35)
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,
(I −A) ( Σprioree TΣprior ) (I −A)T = ( AΣpriorA T −AΣprior −ΣpriorAT +Q ) ( 1 + eTΣpriore T )
(36)
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 inO and Θ. Note that for Adam we haveAAdam ∈ Θ(η) but for RMSprop we have ARMSprop ∈ Θ(η2), and we summarise both using the asymptotic upper bound above.
The right-hand-side of Eq. (36) is thus bounded above,( AΣpriorA T −AΣprior −ΣpriorAT +Q ) ( 1 + eTΣpriore 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η→0A = 0, (I −A) ( Σprioree TΣprior ) (I −A)T ∈ Θ ( Σ2(η) ) . (38)
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),
Σ2(η) ∈ O ( η2 + ηΣ(η) + ηΣ2(η) ) . (39)
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,
Σ2(η) ∈ O ( η2 )
=⇒ Σ(η) ∈ O (η) (40)
Alternatively, if it is bound by ηΣ(η), we have,
Σ2(η) ∈ O (ηΣ(η)) =⇒ Σ(η) ∈ O (η) (41)
And as it must be bound by one (or both) of these terms, we know that Σ(η) and hence Σprior are bounded by η,
Σ(η) ∈ O (η) =⇒ Σprior ∈ O (η) . (42)
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,
A ( Σprioree TΣprior ) ∈ O(η3), (43a)(
Σprioree TΣprior ) AT ∈ O(η3), (43b)
A ( Σprioree TΣprior ) AT ∈ O(η4), (43c)
and on the right-hand-side,
AΣpriorA T ( 1 + eTΣpriore ) ∈ O(η3), (43d)(
Q−AΣprior −ΣpriorAT ) ( eTΣpriore ) ∈ O(η3). (43e)
Neglecting these terms, and instead focusing only on the O(η2) terms we obtain,
Σprioree TΣprior ≈ Q−AΣprior −ΣpriorAT . (44)
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),
Σprioree TΣprior = ((I −A) Σpost (I −A) +Q) eeT ((I −A) Σpost (I −A) +Q)
≈ ΣposteeTΣpost. (45) And we do the same for the right-hand-side of Eq (44),
Q−AΣprior −ΣpriorAT = Q−A ((I −A) Σpost (I −A) +Q)− ((I −A) Σpost (I −A) +Q)AT
≈ Q−AΣpost −ΣpostAT (46) Combining these approximations to the left and right hand sides, we obtain the result used in the main text for Adam.
Σpostee TΣpost ≈ Q−AΣpost −ΣpriorAT . (47)
For RMSprop, remember that ARMSprop ∈ O(η2), so AΣpost ∈ O(η3), and we can obtain an even simpler result,
Σpostee TΣpost ≈ Q. (48)
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).

heading: 9.2 SETTING THE MOMENTUM DECAY
text: 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 inA,
Q ≈ AΨ + (AΨ)T , (49)
and as,
AΨ = ( η2+η2w 2σ2 −η 0 ηp )( Ψww Ψwp Ψwp Ψpp ) = ( η2+η2w 2σ2 Ψww − ηΨwp η2+η2w 2σ2 Ψwp − ηΨpp ηpΨwp ηpΨpp ) .
Eq. (49) becomes,( η2w 0 0 η2p ) ≈ ( 2 ( η2+η2w 2σ2 Ψww − ηΨwp ) η2+η2w 2σ2 Ψwp − ηΨpp + ηpΨwp
η2+η2w 2σ2 Ψwp − ηΨpp + ηpΨwp 2ηpΨpp
) (50)
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+η2w 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.

heading: 9.3 EFFICIENT IMPLEMENTATION OF BRMSPROP
text: 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 σ2prior and σ 2 post in memory. Second, we simplified the covariance updates in Eq. (24),
σ2post = σ 2 prior
( 1− σ2priore 2
1 + σ2priore 2
) =
σ2prior 1 + σ2priore 2 . (51)

heading: 9.4 EFFICIENT IMPLEMENTATION OF BADAM
text: 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,
T = I −A. (52)
Algorithm 1 BRMSprop Require: η . Learning rate Require: ηg . Learning rate for average gradient (0.01) Require: σ2 . Prior variance (1/(2Ninputs))
1: Σ← σ2 . Initialize uncertainties 2: µ ∼ N (0, 2/Ninputs) . Initialize parameters 3: ĝ ← 0 . Initialize average gradient 4: while not converged do . Σ = σ2prior(t), µ = µprior(t) 5: g ← ∇Lt(µ) . Gradients of minibatch t 6: Σ← Σ/(1 + (g − ĝ)2Σ) . Σ = σ2post(t) 7: µ← µ+ Σg . µ = µpost(t) 8: ĝ ← (1− ηg) ĝ + ηgg . Update average gradient 9: Σ← ( 1− η2/(2σ2) )2 Σ + η2 . Σ = σ2prior(t+ 1)
10: µ← ( 1− η2/(2σ2) ) µ . µ = µprior(t+ 1) 11: end while 12: return µ

heading: 9.5 TRICK TO ALLOW LEARNING RATES TO VARY ACROSS PARAMETERS
text: 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.

heading: 9.6 NESTEROV ACCELERATED GRADIENT
text: 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)
v(t) ≈ (1− ηp) v(t− 1) + ηηp g(t)√ 〈e2〉 , (53b)
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)),
gMom(t) = g (µw(t− 1)) = g (µpost, w(t− 1)) (54a) gNAG(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).

heading: 9.7 ESTIMATING THE NEGATIVE HESSIAN
text: 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-
Algorithm 2 BAdam Require: η . Learning rate Require: ηg . Learning rate for average gradient (0.01) Require: ηp . Learning rate for momentum (0.1) Require: ηw . Learning rate for weight alone (0) Require: σ2 . Prior variance (1/(2Ninputs))
1: Tww = 1− (η2 + η2w)/(2σ2) . Set dynamics matrix 2: Twp = η . Set dynamics matrix 3: Tpp = 1− ηp . Set dynamics matrix 4: 5: Σww ← σ2 . Initialize uncertainties 6: Σwp ← 0 . Initialize uncertainties 7: Σpp ← 0 . Initialize uncertainties 8: µw ∼ N (0, 2/Ninputs) . Initialize parameter 9: µp ← 0 . Initialize momentum 10: ĝ ← 0 . Initialize average gradient 11: while not converged do 12: g ← ∇Lt(µw) . Gradients of minibatch t 13: e2 ← ∇(g − ĝ)2

heading: 14: Σww ← Σww − Σ2wwe2/(1 + Σwwe2)

heading: 15: Σwp ← Σwp − ΣwwΣwpe2/(1 + Σwwe2) . Σ = Σpost(t)

heading: 16: Σpp ← Σpp − Σ2wpe2/(1 + Σwwe2)

heading: 17: µw ← µw + Σwwg . µ = µpost(t)
text: 18: µp ← µp + Σwpg 19: ĝ ← (1− ηg) ĝ + ηgg . Update average gradient 20: Σww ← T 2wwΣww + 2TwwTwpΣwp + T 2wpΣpp + η2w 21: Σwp ← TwwTppΣwp + TwpTppΣpp . Σ = Σprior(t+ 1) 22: Σpp ← T 2ppΣpp + η2p 23: µw ← Twwµw + Twpµp . µ = µprior(t+ 1) 24: µp ← Tppµp 25: end while 26: return µ
Information based estimate, of the negative Hessian, Λlike ≈ EP(dmodel|z) [ − ∂ 2
∂z2 log P (dmodel|z)
] = EP(dmodel|z) [ eeT ] , (55)
where
e = ∂ log P (dmodel|z)
∂z
∣∣∣∣ z=µprior . (56)
Importantly, this Fisher-Information based estimate does not depend on actual data, d. Instead, it uses expectations (and gradients) taken over data sampled from the model, dmodel, and this is fundamentally necessary for the the second equality in Eq. (55) to hold. As such, to use Eq. (55) we need to sample surrogate data from the model, and compute gradients with respect to that sampled data. Unfortunately, sampling surrogate data increases complexity and the required computational cost. As such, a simple low-cost alternative is to note that the key difference between e and g is that e has zero-mean, whereas g has non-zero mean, suggesting that we should correct the gradient by subtracting its mean,
e ≈ g − ĝ where ĝ is an empirical estimate of the mean gradient.

Review summaries:
In this work, the authors attempt to unify existing adaptive gradient methods under the Bayesian filtering framework with the dynamical prior.  In Ollivier, 2017, a framework is proposed to connect Bayesian filtering and natural gradient.  On the other hand,  in Khan et al., 2018. an approach is proposed to connect natural gradient and adaptive gradient methods.  The main contributions of this work are (1)  introducing a dynamical prior and (2) recovering RMSProp and Adam as special cases.