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 unitswapping symmetries 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 new random variable, w ∗ i , sampled from the posterior distribution over w i , conditioned on the data, D , and conditioned on the other underlying, unknown random variables, w − i , being equal to our current best estimate of those parameters, µ − i ,
[Equation]
Note that w ∗ i is close to, but not quite the mode of the likelihood of the actual data.
To interpret w ∗ i as a mode, note that performing maximum likelihood with inﬁnite data recovers the underlying value of the parameter.
Therefore, we can consider we need to consider performing maximum likelihood with inﬁnite data, or equivalently with an expectation over data drawn from the true data distribution,
[Equation]
and w i is a “dummy” variable giving the location at which we evaluate the likelihood.
This picks out the true underlying value of the underlying parameter, w ∗ i , which is distributed according to Eq.
If we could directly compute P ( w ∗ i | µ − i , D ) , we would have solved the problem.
However,
P ( w ∗ i | µ − i , D ) is intractable because the relationship between w i and µ − i is determined by the intractable, high dimensional correlations in the underlying posterior, P ( w |D ) .
Instead, we could attempt to infer w ∗ i conditioned on just the data, i.e. P ( w ∗ i |D ) .
However, P ( w ∗ i |D ) is also vulnerable to unit-swapping symmetries and as such is also uselessly broad.
If the previous two approaches, computing P ( w i |D ) and P ( w ∗ i |D ) have failed, what other approaches remain viable?
Remember that the updates for almost all neural network optimization algorithms depend only on the backpropagated gradient for that parameter.
This gradient depends on and thus carries information about both the underlying data, D , and the current settings of the other parameters, µ − i .
Thus, a promising approach is to infer a posterior over w ∗ i from the gradient, P ( w ∗ i | g i ) .
To infer a parameter from the corresponding backpropagated gradient, g i , we need to understand the generative process for that gradient.
To do this, we consider the random variable, d, representing a minibatch drawn from the true underlying data distribution, and take the second-order Taylor series expansion for the corresponding log-likelihood,
[Equation]
where we perform the Taylor expansion for the i th parameter keeping the others ﬁxed.
The negative
Hessian in the i th direction, Λ i , and the mode, µ like ,i , are random variables that depend on the minibatch, and on the current estimate of all the other parameters, µ − i ,
[Equation]
Now, we can write down the gradient of the log-likelihood,
[Equation]
To understand the statistics of gradients (conditioned on the negative Hessian) we thus need to understand the statistics of µ like ,i .
In particular, we focus on its mean and variance
[Equation]
To obtain the mean, note that the expectation of the gradient in Eq.
(5) is zero at the mode (i.e. at µ i =
w ∗ i ), and assume that µ like ,i and Λ i are uncorrelated (but not independent).
The variance can be obtained directly by considering the Fisher Information, which tells us that the variance of the gradient should be Λ i .
Alternatively, note that for a well-calibrated model, w ∗ i must be a valid sample from the posterior distribution resulting from combining this likelihood with e.g. a uniform prior (Cook et al., 2006).
Using these moments, we can obtain expressions for the moments of the gradient,
[Equation]
Assuming that the gradient is Gaussian with these moments, we obtain a complete model for the backpropagated gradient,
[Equation]
The resulting generative process is shown in Fig.
1A (omitting Λ i ).
Importantly, if we choose to infer w ∗ i , then dynamics emerge automatically (indeed, they are already shown in Fig.
1A), because w ∗ i , depends on the current estimate of all the other parameters, and they vary over time as they are being optimized.
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.
To understand how this approach relates to previous approximations note that, up to second order, the likelihood for our new model (Eq.
8) is equivalent to that under the original model (Eq. 3),
[Equation]
where the last equality comes from substituting the gradient from Eq.
Thus, solving these two problems gives the same likelihood (up to our second order Taylor series expansion), reassuring us that we loose little when we move to our alternative problem: computing the posterior over w ∗ i conditioned on the gradient.
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 which are also being optimized.
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]
[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]
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 ,
[Equation]
where,
[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 = µ prior ,i ,
[Equation]
Note that as the log-likelihood depends only on w ∗ i , we have,
[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 i
[Equation]
Ti , where
[Equation]
As such, the updates for the posterior covariance (Eq. 15a) can be rewritten using the Sherman
Morrison formula (Hager, 1989),
[Equation] [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]
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]
24) 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]
[Equation]
where ξ w ( t ) is again standard Gaussian noise, η is the strength of the momentum coupling, η 2 +
[Equation] [Equation]
[Equation]
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]
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,
[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.
[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]
) giving,
[Equation]
33b
[Equation]
which recovers Adam.
See SI Sec.
9.6 for a similar discussion for NAG/NAdam
[Figure]
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 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) .
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).
7. C ENTERED GRADIENTS
BRMSprop and BAdam suggest that just as centred gradients (or gradients based on labels sampled from the current model) are used in work on natural gradients, they should be used in RMSprop and
Adam (and indeed, this is sometimes done, see Graves, 2013).
7. C OMBINE MOMENTUM WITH DIRECT GRADIENT
To recover Adam from BAdam we needed to set Σ ww ← 0 .
As such, BAdam differs from Adam in that it resembles a linear combination of RMSprop and Adam, suggesting that such linear combinations might provide fruitful directions for future investigation.
7. A UTOMATICALLY SETTING THE WEIGHT DECAY
A ﬁnal advantage of BRMSprop and BAdam is that they eliminate the need to set a separate weight decay parameter: the weight decay is speciﬁed by the stationary variance, σ 2 , which is very close to that used by standard initialization schemes, and the learning rate, η .
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, Σ post ( t + 1) = (cid:16)
[Equation]
I
[Equation]
A
[Equation]
post
[Equation]
I
[Equation]
A
[Equation]
−
[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
[Equation]
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]
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] [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]
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 µ
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 ( t + 1) ,
[Equation]
(38a)
[Equation]
(38b
[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]
= g
[Equation]
post
[Equation]
(39a)
[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 Fisher-
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
[Equation]
21: Σ
[Equation]
←
[Equation]
ww
[Equation]
22: Σ pp ← T 2 pp Σ
pp
+ η
[Equation]
p 23: µ w ← T ww µ w + T wp µ p
[Equation]
t
[Equation]
1) 24: µ p ← T pp µ p 25: end while 26: return µ
Information based estimate, of the negative Hessian,
[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, e
[Equation]
ˆ g where ˆ g is an empirical estimate of the mean gradient.