A UNIFIED THEORY OF ADAPTIVE STOCHASTIC GRADIENT DESCENT AS B AYESIAN FILTERING
Anonymous authors
Paper under double-blind review
A BSTRACT
We formulate stochastic gradient descent (SGD) as a novel factorised Bayesian ﬁltering 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 ﬁnd that Bayesian methods have the potential to considerably reduce test loss and classiﬁcation error.
I NTRODUCTION
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 uniﬁed 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) ﬁltering 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 modiﬁed their natural gradient VI updates, without a principled justiﬁcation 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 ﬁltering (Ollivier, 2017).
While Ollivier (2017) did not examine the relationship between their ﬁltering updates and RMSprop/Adam, the equivalence of this particular ﬁltering 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 ﬁlter, but these dynamics correspond to the “addition of an artiﬁcial 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) optimal value, i th parameter, w ∗ i .
B The graphical model obtained by integrating over trajectories for the other parameter estimates, µ − i ( t ) .
In practice, we use a simpliﬁed 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 deﬁne a new variable, z i ( t ) , which incorporates w ∗ i along with other information about the dynamics.
How might we obtain a principled Bayesian ﬁltering 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 ﬁltering 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
F ACTORISATION 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 deﬁne a random variable, w ∗ i , representing the optimal value for the i th weight, conditioned on the other weights, w − i , being equal to our current estimate of those parameters, µ − i ,
[Equation]
where w ∗ i is a random variable because it depends on our current estimate of the other weights, µ − i , which is treated as a random variable, and w i is a “dummy” variable giving the location at which we evaluate the likelihood.
The random variable describing our current estimate of the other weights, µ − i , is generated from a complex distribution controlled by our overall optimization scheme.
However, we cannot now use the usual Bayesian approach of inferring w ∗ i based on the data, because the data depend on the full weight vector, w, not on 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 approximations (e.g. Zhang et al., 2017; Khan et al., 2017; 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 i th direction, keeping the other parameters, w − i , ﬁxed to their current estimates, µ − i ,
[Equation]
where,
[Equation]
Here, both the mode, µ like ,i , and the negative Hessian, Λ 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 random variables.
As such, the gradient of the log-likelihood, which depends on Λ i and µ − i , is also random variable,
[Equation]
Our goal is to understand the distribution of the gradients, conditioned on the random variable Λ i being equal to some speciﬁc value, λ i .
In practice, we will set Λ i , using the standard approximations in the natural gradient literature (Zhang et al., 2017; Khan et al., 2017; 2018).
To compute the expected gradient, we assume that the expected value of µ like ,i is independent of Λ i , and note that the expected gradient must be zero at µ i = w ∗ i , which gives,
[Equation]
substituting this into Eq .(5), we obtain,
[Equation]
However, obtaining an analytic form for the full distribution of the gradient is more difﬁcult.
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.
In particular, if we take,
[Equation]
then for any given approximation to the original log-likelihood (Eq. 2), the log-likelihood induced by conditioning on the gradient (Eq.
8) 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 speciﬁed the values for all the random variables: g i =
g i and Λ i
= λ i which ﬁx the value of µ like ,i to µ like =
g i λ 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 ﬁxed, and as such, Bayesian ﬁltering 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 difﬁcult to reason about, as it requires us to integrate over all possible trajectories for the other parameters, µ − i .
Instead, we write down simpliﬁed approximate dynamics over w ∗
i directly (Fig.
1B), and evaluate the quality of these simpliﬁed dynamics empirically.
B AYESIAN (K ALMAN ) FILTERING AS ADAPTIVE SGD
Note that in Fig.
1B, w ∗ i has arbitrary length temporal dependencies, which are difﬁcult to specify or to perform inference over.
Instead, we deﬁne 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. 10),
[Equation]
where have omitted time indices for clarity,
[Equation]
This likelihood is equivalent to our original likelihood (Eq. 9) because the gradients depend only on w ∗ i ( t ) .
As such, the negative Hessians must take the form,
[Equation] [Equation]
The Gaussian prior and approximate likelihood allows us to use standard two-step Kalman ﬁlter 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
[Equation]
where, µ prior ,i = µ prior ,
i
[Equation]
14a
[Equation]
14b
[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 deﬁned abbreviations for µ prior ,i ( t ) and Σ prior ,i ( t ) , omitting the temporal index, t .
Second we use Bayes theorem to incorporate new data,
[Equation]
where,
[Equation]
Thus far, we have simply repeated standard Kalman ﬁltering results.
To relate Kalman ﬁltering to gradient ascent, we compute the gradient of Eq.
(12) at z i = µ
[Equation]
,i ,
[Equation]
Note that as the log-likelihood depends only on w ∗ i , we have,
[Equation]
[Equation]
Now, we identify this gradient (Eq. 16) in the mean updates (Eq. 15b),
[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 conﬁdence in the current estimates.
As the precision is always rank-1 (Eq. 13), we can always write it as,
[Equation]
e Ti
[Equation]
where
[Equation]
As such, the updates for the posterior covariance (Eq. 15a) 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.
B AYESIAN RMS PROP (BRMS PROP )
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]
[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. 14) and the posterior (Eq. 18 and Eq. 20),
[Equation]
For an efﬁcient implementation of the full algorithm, see SI Alg.
4. R ECOVERING RMS PROP FROM BRMS PROP
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.
In particular we use the simpliﬁed covariance updates derived in SI Sec. 9.1,
[Equation]
As t → ∞ , we expect σ 2 post to reach steady-state, at which point, σ 2 post = σ 2 post ( t ) =
σ 2 post ( t + 1) .
Solving for σ 2 post ,
[Equation]
Substituting this form into Eq.
(22) recovers the root-mean-square normalization used in RMSprop.
B AYESIAN A DAM (BA DAM )
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 ) ,
[Equation]
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 η 2 p ,
[Equation]
where ξ p ( t ) is standard Gaussian noise.
This particular coupling of the injected noise variance and the decay rate ensures that the inﬂuence 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 ) ),
[Equation]
where ξ w ( t ) is again standard Gaussian noise, η is the strength of the momentum coupling, η 2 + η 2 w 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. 11), by using,
[Equation]
and these settings fully determine the updates, according to Eq.
(14) and Eq.
For an efﬁcient implementation of the full algorithm, see SI Alg.
5. R ECOVERING A DAM FROM BA DAM
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]
( t
[Equation] [Equation]
27b)
These updates depend on two quantities, Σ ww and Σ wp , which are related to (cid:104) e 2 (cid:105) , 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 (cid:104) e 2 (cid:105) in Adam have reached their steady-state values.
To ﬁnd the steadystates for Σ ww and Σ wp , we again use the simpliﬁed covariance updates derived in SI Sec. 9.1,
[Equation]
Substituting Eq.
(26) and Eq.
(19) into Eq.
(28), and again using σ 2 = ∞ and η 2 w = 0 , we obtain, (cid:18)
[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 (cid:28) η 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 ≈
[Equation]
p Σ pp + (cid:104) e 2 (cid:105) Σ
[Equation]
wp ),
[Equation]
To recover a very close approximation to Adam, we need a speciﬁc relationship between learning rates and evidence strength, such that,
[Equation]
While this may seem restrictive, it is only necessary to achieve the closest possible match between
Bayesian ﬁltering 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 ﬁltering and Adam are most similar, the ﬁltering updates become,
[Equation]
31b
[Equation]
where we have substituted for µ p ( t +1) into Eq. 31a, which assumes that η p (cid:28) 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]
32) giving,
[Equation]
which recovers Adam.
See SI Sec.
9.6 for a similar discussion for NAG/NAdam
E XPERIMENTS
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 ﬁrst 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
[Figure]
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, artiﬁcially inﬂating (cid:104)
e 2 (cid:105) .
The key Bayesian-speciﬁc parameters are the variance of the stationary distribution and the initial uncertainty; both of which were set to 1 / (2 N 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 speciﬁc 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 classiﬁcation 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 justiﬁed (correcting this is non-trivial in the dynamical setting, so we leave it for future work).
F EATURES OF B AYESIAN 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. W EIGHT DECAY , L2 REGULARIZATION AND B AYESIAN 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. 22), providing a potential explanation for its improved performance.
7. N ESTEROV 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. 16) is evaluated at µ prior (Eq. 14a), which is precisely a prediction formed by combining the current setting of the parameters, µ post ( t ) , with momentum and decay, embodied in A (Eq. 26) 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. C ONVERGENCE
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. U PDATING THE PRECONDITIONER BEFORE APPLYING THE UPDATE
ADAM (Kingma & Ba, 2015) ﬁrst 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 ﬁltering naturally recovers this choice as the gradient preconditioner in Eq.
(18) is the posterior, rather than the prior, covariance (i.e. updated with the current gradient).
D ISCUSSION
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 ﬁnd, BRMSprop and BAdam are superior to the standard methods in the tasks that we examined.
Finally, the novel interpretation of adaptive
SGD as ﬁltering 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´andezLobato & Adams, 2015) or assumed density ﬁltering (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 justiﬁcation for the introduction of rich dynamics, and they are unable to reason about optimal inference under these dynamics.
Bayesian ﬁltering presents a novel approach to neural network optimization, and as such, there are variety of directions for future work.
First, Bayesian ﬁltering 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 ﬁnd closed-form expressions for the dynamics in simpliﬁed 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 ﬁltering methods that represent the covariance of a full weight matrix by exploiting Kronecker factorisation (Martens & Grosse, 2015; Grosse &
Martens, 2016; Zhang et al., 2017).
S UPPLEMENTARY I NFORMATION 9. A PPROXIMATE COVARIANCE UPDATES
To recover previously known update rules (e.g. RMSprop and Adam), it is usually useful to collapse the two covariance updates (Eq. 14b and Eq. 20) into a single update, which we simplify by neglecting various third-order terms.
If we naively substitute Eq. 14b into Eq. 20, we obtain,
[Equation]
We have,
[Equation]
where η is the learning rate, which is assumed to be small.
Note that for Adam we have A ∼ η but for RMSprop we have, A ∼ η 2 .
As we are taking the limit of small learning rates, we use the pessimistic case of A ∼ η for the following derivations, and we can subsequently introduce a smaller scaling for A .
To obtain a self-consistent solution, we need,
[Equation]
As Σ post is O ( η ) , and updates to Σ post are O ( η 2 ) , we can neglect O ( η 3 ) terms,
[Equation]
which gives the required simpliﬁed form required for Adam.
Further, note that for RMSprop, A is actually O ( η 2 ) , and therefore we are able to obtain an even simpler form,
[Equation]
9. S ETTING THE MOMENTUM DECAY
In the main text, we brieﬂy note that we couple the injected noise and decay in Eq.
(24) 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.
(25) 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, Ψ ,
[Equation]
where we do not condition on data.
Neglecting the second-order term in A ,
[Equation]
34) and as,
[Equation]
Algorithm
Require: η (cid:46) Learning rate
Require: η g (cid:46) Learning rate for average gradient (0.01)
Require:
σ(cid:46)
Prior variance ( 1 / (2 N inputs ) ) 1: Σ ← σ(cid:46)
Initialize uncertainties 2: µ ∼ N (0 , 2 /N inputs ) (cid:46)
Initialize parameters 3: ˆ g ←(cid:46)
Initialize average gradient 4: while not converged
do (cid:46) Σ
= σ 2 prior ( t ) , µ = µ prior ( t ) 5: g ← ∇L t ( µ ) (cid:46)
Gradients of minibatch t 6: Σ ← Σ / (1 + ( g − ˆ g ) 2 Σ) (cid:46) Σ
= σ 2 post ( t ) 7: µ ← µ + Σ g (cid:46) µ = µ post ( t ) 8: ˆ g ← (1 − η g )
ˆ g + η g g (cid:46) Update average gradient 9: Σ ← (cid:0) 1 − η 2 / (2 σ 2 ) (cid:1)
2 Σ + η(cid:46) Σ
= σ 2 prior ( t + 1) 10: µ ← (cid:0) 1 − η 2 / (2 σ 2 ) (cid:1) µ (cid:46) µ = µ prior ( t + 1) 11: end while 12: return µ
Eq.
(34) becomes,
[Equation]
Now, we can solve for Ψ pp using the ppth element of Eq.
[Equation]
and we can solve for Ψ wp using the wpth element of Eq.
(35), in combination with our solution to Ψ pp ,
[Equation]
This approximation becomes exact in the limit we consider to recover Adam (i.e. σ 2 = ∞ ), but it also holds when σ 2 is ﬁnite, 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.
[Equation]
as required.
9. E FFICIENT IMPLEMENTATION OF BRMS PROP
See Alg.
1 for a complete speciﬁcation of our approach.
There are two important points to note.
First, the method is as memory efﬁcient 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 simpliﬁed the covariance updates in Eq.
[Equation]
9. E FFICIENT IMPLEMENTATION OF BA DAM
See Alg. 2 for a complete speciﬁcation 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 efﬁciency.
Second, we have simpliﬁed the derivation, by deﬁning,
[Equation]
[Equation]
Algorithm
Require: η (cid:46) Learning rate
Require: η g (cid:46) Learning rate for average gradient (0.01) Require: η p (cid:46) Learning rate for momentum (0.1) Require: η w (cid:46) Learning rate for weight alone (0)
Require:
σ(cid:46)
Prior variance ( 1 / (2 N inputs ) ) 1: {+T ww = 1 − ( η 2 + η 2 w ) / (2 σ 2 ) (cid:46) Set dynamics matrix 2: T wp = η (cid:46) Set dynamics matrix 3: T pp = 1 − η p (cid:46) Set dynamics matrix 4: 
 :+} Σ ww ← σ(cid:46) Initialize uncertainties 6: Σ wp ←(cid:46) Initialize uncertainties 7: Σ pp ←(cid:46)
Initialize uncertainties 8: µ w ∼ N (0 , 2 /N inputs ) (cid:46) Initialize parameter 9: µ p ←(cid:46)
Initialize momentum 10: ˆ g ←(cid:46)
Initialize average gradient 11: while not converged {+do 
+}
: g ← ∇L t ( µ w ) (cid:46) Gradients of minibatch t 13: e 2 ← ∇ ( g − ˆ g )14: Σ ww ← Σ ww − Σ 2 ww e 2 / (1 + Σ ww e 2 ) 15: Σ wp ← Σ wp − Σ ww Σ
wp e 2 / (1 + Σ ww e 2 ) (cid:46) Σ = Σ post ( t ) 16: Σ pp ← Σ pp − Σ 2 wp e 2 / (1 + Σ ww e 2 ) 17: µ w ← µ w + Σ ww g (cid:46) µ = µ post ( t ) 18: µ p ← µ p + Σ wp g 19: ˆ g ← (1 − η g ) ˆ g + η g g (cid:46) Update average gradient 20: Σ ww ← T 2 ww Σ ww + 2 T ww T wp Σ wp + T 2 wp Σ pp + η 2 w 21: Σ wp ← T ww T pp Σ wp + T wp T pp
Σ pp (cid:46) Σ
= Σ prior ( t + 1) {+
 : Σ pp+} ← T 2 pp Σ pp + η 2 p 23: µ w ← T ww µ w + T wp µ p (cid:46) µ = µ prior ( t + 1) 24: µ p ← T pp µ p 25: end while 26: return µ
9. T RICK 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. N ESTEROV A CCELERATED G RADIENT
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.
(31) in terms of v ( t + 1) =
ηµ p
[Equation]
[Equation]
[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
[Equation]
[Equation]
As noted in Eq.
(16), Bayesian ﬁltering requires us to evaluate the gradient at µ prior , implying that we use updates based on NAG (Eq. 39b), rather than updates based on standard momentum (Eq. 39a).
9. E STIMATING THE NEGATIVE H ESSIAN
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 FisherInformation based estimate, of the negative Hessian,
[-η (cid:46) Set dynamics matrix 3: T pp = 1 − η p (cid:46) Set dynamics matrix 4: 
-]
[-Initialize average gradient 11: while not converged do 
-]
[Equation]
where
[Equation]
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, d model , and this is fundamentally necessary for the the second equality in Eq.
(40) to hold.
As such, to use Eq.
(40) 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,
[Equation]
where ˆ g is an empirical estimate of the mean gradient.