We formulate stochastic gradient descent (SGD) as a novel factorised Bayesian filtering problem, in which each parameter is inferred separately, conditioned on the corresopnding 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. Deep neural networks have recently shown huge success at a range of tasks including machine translation BID37 , dialogue systems BID29 , handwriting generation BID9 and image generation BID26 . These successes have been facilitated by the development of a broad range of adaptive SGD methods, including ADAGrad BID6 , RMSprop BID13 , Adam BID16 , and variants thereof, including Nesterov acceleratation (Nesterov, 1983; BID2 BID5 and AdamW BID18 . 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 BID23 BID30 BID24 BID25 BID7 BID22 .There have recently been attempts to connect adaptive SGD algorithms to natural gradient variational inference (VI) BID39 BID14 . These approaches give a momentum-free algorithm with a mean-square normalizer, in contrast to perhaps the most popular adaptive method, Adam BID16 , 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 BID39 BID15 , and the root-mean-square normalizer BID14 . 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 BID39 BID14 and Bayesian filtering BID22 . While BID22 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 BID39 BID14 . More problematically , BID22 introduces dynamics into the Kalman filter, but these dynamics correspond to the "addition of an artificial process noise Q t proportional to [the posterior covariance] P t−1 ". Thus, their generative model depends on inferences made under that model: a highly unnatural assumption that most likely does not correspond to any "real" generative process. DISPLAYFORM0 Figure 1: The heirarchy of generative models underlying our updates. A Full model for the gradients for a single parameter. The current estimate for all the other parameters, µ −i (t) vary slowly over time, and give rise to the current optimal value for the ith parameter, w * i . The gradient then arises from the current estimate of the ith parameter, µ i (t) (which is treated as an input here), and the optimal value, ith parameter, w * i . B The graphical model obtained by integrating over trajectories for the other parameter estimates, µ −i (t). In practice, we use a simplified model as reasoning about all possible trajectories of µ −i (t) is intractable. C To convert the model in B into a tractable hidden Markov model (HMM), we define a new variable, z i (t), which incorporates w * i along with other information about the dynamics.How might we obtain a principled Bayesian filtering approach that recovers the two key features of state-of-the-art adaptive SGD algorithms: momentum and the root-mean-square normalizer? Here, we note that past approaches including natural gradient VI take a complex generative model over all N parameters jointly, and use a very strong approximation: factorisation. Given that we know that the true posterior is a highly complicated, correlated distribution, it is legitimate to worry that these strong approximations might meaningfully disrupt the ability of Bayesian filtering to give closeto-optimal updates. Here we take an alternative approach, baking factorisation into our generative model, so that we can use Bayesian inference to reason about (Bayes) optimal updates under the constraints imposed by factorisation. In particular, we split up the single large inference problem over all N parameters, w, into N small inference problems over a single parameter. Remarkably, by incorporating factorisation into the problem setting, we convert intractable, high-dimensional correlations in the original posterior into tractable low-dimensional dynamics in the factorised model. This dynamical prior has a "natural" form, at least compared with BID22 , 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 BID13 and Adam BID16 . Finally, we discuss the general features of Bayesian adaptive SGD methods, including AdamW BID18 and Nesterov acceleration (Nesterov, 1983; BID5 , amongst others. 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 BID33 , or perhaps we wish to consider a biological setting, where the gradient is not present at all BID0 . 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 BID19 BID10 BID39 µ ← µ + Σg µ = µ post (t) 8:ĝ ← (1 − η g )ĝ + η g g Update average gradient 9: µ ← 1 − η 2 /(2σ 2 ) µ µ = µ prior (t + 1) 11: end while 12: return µ DISPLAYFORM0
We formulated SGD as a Bayesian filtering problem, and show that this gives rise to RMSprop, Adam, AdamW, NAG and other features of state-of-the-art adaptive methods The paper analyzes stochastic gradient descent through Bayesian filtering as a framework for analyzing adaptive methods. The authors attempt to unify existing adaptive gradient methods under the Bayesian filtering framework with the dynamical prior