A unified theory of adaptive stochastic gradient descent as Bayesian filtering

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 . 1 INTRODUCTION . Deep neural networks have recently shown huge success at a range of tasks including machine translation ( Wu et al. , 2016 ) , dialogue systems ( Serban et al. , 2016 ) , handwriting generation ( Graves , 2013 ) and image generation ( Radford et al. , 2015 ) . These successes have been facilitated by the development of a broad range of adaptive SGD methods , including ADAGrad ( Duchi et al. , 2011 ) , RMSprop ( Hinton et al. , 2012 ) , Adam ( Kingma & Ba , 2015 ) , and variants thereof , including Nesterov acceleratation ( Nesterov , 1983 ; Bengio et al. , 2013 ; Dozat , 2016 ) and AdamW ( Loshchilov & Hutter , 2017 ) . However , such a broad range of approaches raises the question of whether it is possible to obtain a unified theoretical understanding of adaptive SGD methods . Here we provide such a theory by reconciling state-of-the-art adaptive SGD algorithms with very early work that used Bayesian ( Kalman ) filtering to optimize the parameters of neural networks ( Puskorius & Feldkamp , 1991 ; Sha et al. , 1992 ; Puskorius & Feldkamp , 1994 ; 2001 ; Feldkamp et al. , 2003 ; Ollivier , 2017 ) . There have recently been attempts to connect adaptive SGD algorithms to natural gradient variational inference ( VI ) ( Zhang et al. , 2017 ; Khan et al. , 2017 ; 2018 ) . These approaches give a momentum-free algorithm with a mean-square normalizer , in contrast to perhaps the most popular adaptive method , Adam ( Kingma & Ba , 2015 ) , which combines momentum with a root-meansquare normalizer . To achieve a closer match to Adam , they modified their natural gradient VI updates , without a principled justification based on approximate inference , to incorporate momentum ( Zhang et al. , 2017 ; Khan et al. , 2018 ) , and the root-mean-square normalizer ( Khan et al. , 2017 ; 2018 ) . As such , there appears to be only a loose connection between successful adaptive SGD algorithms such as Adam , and natural gradient VI . There is a formal correspondence between natural gradient VI ( Zhang et al. , 2017 ; Khan et al. , 2017 ; 2018 ) and Bayesian filtering ( Ollivier , 2017 ) . While Ollivier ( 2017 ) did not examine the relationship between their filtering updates and RMSprop/Adam , the equivalence of this particular filtering approach and natural gradient VI indicates that they would encounter the issues described above , and thus be unable to obtain momentum or the root-mean-square normalizer ( Zhang et al. , 2017 ; Khan et al. , 2017 ; 2018 ) . More problematically , Ollivier ( 2017 ) introduces dynamics into the Kalman filter , but these dynamics correspond to the “ addition of an artificial process noise 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 . 2 FACTORISATION IMPLIES A RICH DYNAMICAL PRIOR . 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 can not 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 )

Summary: 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.