Self-Tuning Stochastic Optimization with Curvature-Aware Gradient Filtering
Abstract
Standard ﬁrst-order stochastic optimization algorithms base their updates solely on the average mini-batch gradient, and it has been shown that tracking additional quantities such as the curvature can help de-sensitize common hyperparameters.
Based on this intuition, we explore the use of exact per-sample Hessian-vector products and gradients to construct optimizers that are self-tuning and hyperparameter-free.
Based on a dynamics model of the gradient, we derive a process which leads to a curvature-corrected, noise-adaptive online gradient estimate.
The smoothness of our updates makes it more amenable to simple step size selection schemes, which we also base off of our estimates quantities.
We prove that our model-based procedure converges in the noisy quadratic setting.
Though we do not see similar gains in deep learning tasks, we can match the performance of well-tuned optimizers and ultimately, this is an interesting step for constructing self-tuning optimizers.
Introduction
Stochastic gradient-based optimization is plagued by the presence of numerous hyperparameters.
While these can often be set to rule-of-thumb constants or manually-designed schedules, it is also common belief that a more information regarding the optimization landscape can help present alternative strategies such that manual tuning has less of an impact on the end result.
For instance, the use of curvature information in the form of Hessian matrices or Fisher information can be used to de-sensitize or completely remove step size parameter (Ypma, 1995; Amari, 1998; Martens, 2014), and the momentum coefﬁcient can be set to reduce the local gradient variance (Arnold et al., 2019b).
[Figure]
Based on these intuitions we investigate the use of efﬁcient curvature and variance estimates during training to construct a self-tuning optimization framework.
Under a Bayesian paradigm, we treat the true gradient as the unobserved state of a dynamical system and seek to automatically infer the true gradient conditioned on the history of parameter updates and stochastic gradient observations.
Our method is enabled by evaluations of exact persample gradients and Hessian-vector products.
With recent improvements in automatic differentiation tooling ( e.g. , Bradbury et al., 2018; Agarwal and
Ganichev, 2019; Dangel et al., 2020), this matches the asymptotic time cost of minibatch gradient and
Hessian-vector product evaluations.
∇ f t −∇ f t ∇ f t + g t −g t g t + Σ t −Σ t Σ t +
B t
B t +
Q t
Q t +
While our framework contains the good properties of both curvature-based updates and variance reduction—which are attested in toy and synthetic scenarios—we do not observe signiﬁcant improvements empirically in optimizing deep neural networks.
Notably, our approach can be viewed as an explicit form of the implicit gradient transport of Arnold et al.
(2019b), yet
it does not achieve the same acceleration empirically observed in practice.
While we do not fully understand this behavior, we analyze the estimated quantites along the training trajectory and hypothesize that our method has a higher tendency of going down high-variance high-curvature regions whereas standard stochastic gradient descent is repelled from such regions due to gradient variance.
This potentially serves as a downside of our method in the deep learning setting.
Regardless, the use of efﬁcient variance estimation and interpretation of gradient estimation as Bayesian ﬁltering are useful constructs in the development of self-tuning stochastic optimization.
Bayesian Filtering for Stochastic Gradients
We consider stochastic optimization problems of the general form
[Equation]
(1) where we only have access to samples ξ .
Stochastic gradient descent—the prototypical algorithm for this setting—iteratively updates θ t +
=
θ t − α t g t , where
[Equation]
ξ
[Equation] iid [Equation]
p [Equation] , (2) and α t is a scalar step size.
We may use notational shorthands like f t = f
( θ t ) , ∇ f t = ∇ f
( θ t ) .
SGD is hampered by the effects of gradient noise.
It famously needs a decreasing step size schedule to converge; used with a constant step size, it goes into diffusion in a region around the optimum (see, e.g. , Bottou et al., 2018).
Gradient noise also makes stochastic optimization algorithms difﬁcult to tune.
In particular, unreliable directions are not amenable to step size adaptation.
To stabilize update directions, we build a framework for estimating the true gradient ∇ f based on
Kalman ﬁltering.
This can also be viewed as a variance reduction method, but does not require the typical ﬁnite-sum structure assumption of e.g. Schmidt et al.
; Johnson and Zhang (2013).
2.1 Dynamical System Model
We treat the true gradient ∇ f t as the latent state of a dynamical system.
This dynamical system is comprised of an observation model p ( g t | f t ) and a dynamics model p ( ∇ f t | ∇ f t − 1 , δ t − 1 ) where δ t − 1 = θ t − θ t − 1 is the update direction.
We will later choose δ t to be depend on our variance-reduced gradient estimates, but the gradient inference framework itself is agnostic to the choice of δ t .
The observation model p ( g t | ∇ f t ) describes how the gradient observations relate to the state of the dynamical system.
In our case, it is relatively straight-forward, since g t is simply an unbiased stochastic estimate of ∇ f t , but the exact distribution remains to be speciﬁed.
We make the assumption that g t follows a Gaussian distribution,
[Equation]
with covariance Σ t .
Since g t is the mean of iid terms (Eq. 2), this assumption is supported by the central limit theorem when sufﬁciently large batch sizes are used.
The dynamics model p ( ∇ f t | ∇ f t − 1 ) describes how the gradient evolves between iterations.
We base our dynamics model on a ﬁrst order Taylor expansion of the gradient function centered at θ t , ∇ f
( θ t − 1 ) ≈ ∇ f ( θ t ) − ∇ 2 f ( θ t ) δ t − 1 .
We propose to approximate the gradient dynamics by computing a stochastic estimate of the Hessian-vector product, B t δ t − 1 , where E [ B t ] = ∇ 2 f ( θ t ) .
Again, we make a Gaussian noise assumption.
This implies the dynamics model
[Equation]
where Q t is the covariance of B t δ t − 1 , taking into account the stochasticity in B t .
A key insight is that the parameters B t δ t − 1 , Q t , Σ t of the model can all be “observed” directly using automatic differentiation of the loss on each minibatch of samples.
We use the Hessian at θ t so that the
Hessian-vector product can be simultaneously computed with g t with just one extra call to automatic differentiation (or “backward pass”) in each iteration (note this does not require constructing the full matrix B t ).
The variances Q t and Σ t can also be empirically estimated with some memory overhead by using auto-vectorized automatic differentiation routines.
We discuss implementation details later in Section 4.
2.2 Filtering Framework for Gradient Inference
As Equations (3) and (4) deﬁne a linear-Gaussian dynamical system, exact inference on the true gradient conditioned on the history of gradient observations p ( ∇ f t | g 1: t , δ 1: t − 1 ) takes the form of the well-known Kalman ﬁltering equations (Kalman, 1960)
(review in Särkkä, 2013):
We deﬁne parameters m − t , m t , P − t and P t such that
[Equation]
Starting from a prior belief ∇ f 0 ∼ N ( m 0 , P 0 ) , these parameters are updated iteratively:
[Equation]
Equation (6) is referred to as the prediction step as it computes mean and covariance of the predictive distribution p ( ∇ f t | g 1: t − 1 ) .
In our setting, it predicts the gradient ∇ f t based on our estimate of the previous gradient ( m t − 1 ) and the Hessian-vector product approximating the change in gradient from the step θ t = θ t − 1 + δ t − 1 .
Equation (8) is the correction step.
Here, the local stochastic gradient evaluation g t is used to correct the prediction.
Importantly, the Kalman gain (7) determines the blend between the prediction and the observations according to the uncertainty in each.
The resulting algorithm gives an online estimation of the true gradients as the parameters θ t are updated.
We refer to this framework as M EKA , loosely based on model-based Kalman-adjusted gradient estimation .
During optimization, we may use the posterior mean m t as a variance-reduced gradient estimator and take steps in the direction of δ t = − α t m t .
We note two key insights enabling M EKA : First, all parameters of the ﬁlter are not set ad hoc , but are directly evaluated or estimated using automatic differentiation.
Secondly, the dynamics model makes explicit use of the Hessian to predict gradients.
This is a ﬁrst-order update.
In contrast to second-order methods, like quasi-Newton methods, M EKA does not try to estimate the Hessian from gradients, but instead leverages a (noisy) projection with the actual Hessian to improve gradient estimates.
This is both cheaper and more robust than second-order methods, because it does not involve solving a linear system.
A DAM -style Update Directions
While M EKA produces variance-reduced gradient estimates, it does not help with ill-conditioned optimization problems, a case where full batch gradient descent can perform poorly.
To alleviate this, we may instead take update directions motivated by the A DA G RAD (Duchi et al., 2011) line of optimizers.
We follow A DAM (Kingma and Ba, 2014) which proposes dividing the ﬁrst moment of the gradient element-wise by the square root of the second moment, to arrive at
[Equation]
where ε is taken for numerical stability and simply set to 10 − 8 .
Whereas A DAM makes use of two exponential moving averages to estimate the ﬁrst and second moments of g t , we have estimates automatically inferred through the ﬁltering framework.
We refer to this variant as A DA M EKA .
Uncertainty-informed Step Size Selection
We can adopt a similar Bayesian ﬁltering framework for probabilistic step size adaptation.
Our step size adaptation will be a simple enhancement to the quadratic rule, but takes into account uncertainty in the stochastic regime and is much more robust to stochastic observations.
The standard quadratic rule if the objective f can be computed exactly is α quadratic :=
[Equation]
, which is based on minimizing a local quadratic approximation f [Equation]
[Equation]
However, since we only have access to stochastic estimates of ∇ f and ∇ 2 f , naïvely taking this step size with high variance samples results in unpredictable behavior and can cause divergence during optimization.
To compensate for the stochasticity and inaccuracy of a quadratic approximation, adaptive step size approaches often include a “damping” term ( e.g. Martens (2010))—where a constant is added to the denominator—and an additional scaling factor on α t , both of which aim to avoid large steps but introduces more hyperparameters.
As an alternative, we propose a scheme that uses the variance of the estimates to adapt the step size, only taking steps into regions where we are conﬁdent about minimizing the objective function.
Once again leveraging the availability of Q t and Σ t , our approach allows automatic trade-off between minimizing a local quadratic approximation and the uncertainty over large step sizes, foregoing manual tuning methods such as damping.
We adopt a similar linear-Gaussian dynamics model for tracking the true objective f t , with the same assumptions as in Section 2.
Due to its similarity with Section 2, we delegate the derivations to
Appendix B. We again deﬁne the posterior distribution,
[Equation]
where u t and s t are inferred using the Kalman update equations.
Finally, setting f t +1 = f
( θ t + α t δ t ) for some direction δ t , we have a predictive model of the change in function value as
[Equation]
Contrasting this with the simple quadratic approximation, the main difference is now we take into account the uncertainty in f t , ∇ f t , and ∇ 2 f t .
Each term makes different contributions to the variance as α t increases, corresponding to different trade-offs between staying near where we are more certain about the function value and exploring regions we believe have a lower function value.
Explicitly specifying this trade-off gives an acquisition function .
These decision rules are typically used in the context of Bayesian optimization (Shahriari et al., 2016), but we adopt their use for step size selection.
3.1 Acquisition Functions for Step Size Selection
Computing the optimal step size in the context of a long but ﬁnite sequence of optimization steps is intractable in general, but many reasonable heuristics have been developed.
These heuristics usually balance immediate progress against information gathering likely to be useful for later steps.
α
E s t i m a t ed f
( θ + α δ )
Quadratic Mean/Std Expected Improvement
Prob.
of Improvement (a) Positive curvature
α
E s t i m a t ed f
( θ + α δ )
Quadratic Mean/Std Expected Improvement
Prob.
of Improvement (b) Negative curvature
One natural and hyperparameter-free heuristic is maximizing the probability of improvement (PI) (Kushner, 1964),
[Equation]
(12) which is simply the cumulative distribution function of (11) evaluated at zero.
Figure 3 visualizes the different step sizes chosen by maximizing different acquisition functions.
The heuristic of choosing the minimum of the quadratic approximation can be a poor decision when the uncertainty rises quickly.
The optimum for PI interpolates between zero and the quadratic minimum in such a way that avoids regions of high uncertainty.
Expected improvement (Jones et al., 1998) is another popular acquisition function; however, in tests we found it to not be as robust as PI and often results in step sizes that require additional scaling.
Maximizing probability of improvement is equivalent to the following optimization problem
[Equation]
We numerically solve for α PI using Newton’s method, which itself is only a small overhead since we only optimize in one variable with ﬁxed constants: no further evaluations of f are required.
We also note that there is exactly one optimum for α ∈ R + .
For optimization problems where negative curvature is a signiﬁcant concern, we include a third-order correction term that ensures ﬁnite and positive step sizes (details in Appendix B.2).
A Practical Implementation
While the above derivations have principled motivations and are free of hyperparameters, a practical implementation of M EKA is not entirely straightforward.
Below we discuss some technical aspects, simpliﬁcations and design choices that increase stability in practice, as well as recent software advances that simplify the computation of quantities of interest.
Computing Per-Example Quantities for Estimating Variance Recent extensions for automatic differentiation in the machine learning software stack (Bradbury et al., 2018; Agarwal and Ganichev, 2019; Dangel et al., 2020) implement an automatic vectorization map function.
Vectorizing over minibatch elements allows efﬁcient computation of gradients and Hessian-vector products of neural network parameters with respect to each data sample independently.
These advances allow efﬁcient computation of the empirical variances of gradients and Hessian-vector products, and enable our ﬁltering-based approach to gradient estimation.
Stabilizing Filter Estimates Instead of working with the full covariance matrices Σ t and Q t , we approximate them as scalar objects σ t I and q t I , with σ t , q t ∈ R + by averaging over all dimensions.
We have experimented with diagonal matrices, but found that the scalar form increases stability, generally performing better on our benchmarks.
Furthermore, we use an exponential moving average for smoothing the estimated gradient variance σ t as well as the adaptive step sizes α t .
The coefﬁcients of these exponential moving average are kept at 0 .
999 in our experiments and seem to be quite insensitive, with values in { 0 . 9 , 0 . 99 , 0 .
999 } all performing near identically (see Appendix F).
Related Work
Designing algorithms that can self-tune its own parameters is a central theme in optimization (Eiben and Smit, 2011; Yang et al., 2013); we focus on the stochastic setting, building on and merging ideas from several research directions.
The Bayesian ﬁltering framework itself has previously been applied to stochastic optimization.
To the best of our knowledge, the idea goes back to Bittner and
Pronzato (2004) who used a ﬁltering approach to devise an automatic stopping criterion for stochastic gradient methods.
Patel (2016) proposed ﬁltering-based optimization methods for large-scale linear regression problems.
Vuckovic (2018) and Mahsereci (2018) used Kalman ﬁlters on general stochastic optimization problems with the goal of reducing the variance of gradient estimates.
In contrast to our work, none of these existing approaches leverage evaluations of Hessian-vector products to give curvature-informed dynamics for the gradient.
In terms of online variance reduction, Gower et al.
(2017) have discussed the use of Hessian-vector products to correct the gradient estimate; however, they propose methods that approximate the
Hessian whereas we compute exact Hessian-vector products by automatic differentiation.
Arnold et al. (2019a) recently proposed an implicit gradient transport formula analogous to our dynamics model, but they require a rather strong assumption that the Hessian is the same for all samples and parameter values.
In contrast, we focus on explicitly transporting via the full Hessian.
This allows us to stay within the ﬁltering framework and automatically infer the gain parameter, whereas the implicit formulation of Arnold et al. (2019a) requires the use of a manually-tuned averaging schedule.
Step size selection under noisy observations is a difﬁcult problem and has been tackled from multiple viewpoints.
Methods include meta-learning approaches (Almeida et al., 1999; Schraudolph, 1999;
Plagianakos et al., 2001; Yu et al., 2006; Baydin et al., 2017) or by assuming the interpolation regime (Vaswani et al., 2019; Berrada et al., 2019).
Rolinek and Martius (2018) proposed extending a linear approximation to adapt step sizes but introduces multiple hyperparameters to adjust for the presence of noise, whereas we extend a quadratic approximation and automatically infer parameters based on noise estimates.
Taking into account observation noise, Mahsereci and Hennig (2017) proposed a probabilistic line search that is done by ﬁtting a Gaussian process to the optimization landscape.
However, inference in Gaussian processes is more costly than our ﬁltering approach.
Convergence in the Noisy Quadratic Setting
[Figure]
As a motivating example, consider a simple toy problem, where
[Equation]
[Equation] i.e., a mixture of quadratic functions with identical Hessian but varying location determined by the “data” ξ .
The full gradient is ∇ f ( θ ) =
H ( θ − E [ ξ ]) and perexample gradients evaluate to ∇ f ( θ, ξ ) =
H ( θ − ξ ) = ∇ f ( θ ) − H ( ξ − E [ ξ ] ).
Hence, we have additive gradient noise with covariance Σ
= H Cov [ ξ ] H T independent of θ .
Moreover, since the Hessian ∇ 2 f ( θ, ξ ) =
H is independent of ξ , we have that B t δ t − 1 ≡ ∇ f t − ∇ f t − 1 .
The covariance Q t is zero and the ﬁlter equations simplify to
[Equation]
initialized with m 0 =
g 0 , P 0 = Σ .
The ﬁlter covariance P t contracts in every step and in fact, shrinks at a rate of O (1 /t ) , meaning that the ﬁlter will narrow in on the exact gradient.
We show that this enables O (1 /t ) convergence with a constant step size.
[Figure]
Proposition 1.
Assume a problem of the form (14) with µI (cid:22) H (cid:22) LI .
If we update θ t +1 = θ t − αm t with α ≤ 1 /L and m t obtained via Eq.
(15) , then E [ f ( θ t ) − f ∗ ]
∈ O (1 /t ) .
Figure 4 shows experimental results for such a noisy quadratic problem of dimension d = 20 with a randomly-generated Hessian (with condition number > 1000 ) and ξ ∼ N (0 , I ) .
Using SGD with a high learning rate simply results in diffusion, and setting the learning rate smaller results in slow convergence.
Gradient descent (GD) converges nicely with the high learning rate, and using adaptive steps sizes leads to a better convergence rate.
The ﬁltered gradients from M EKA converge almost as well as gradient descent, and adaptive step sizes provide an improvement.
On the other hand, SGD produces unreliable gradient directions and does not work well with adaptive step sizes.
We note that the stochastic gradient has a full covariance matrix and does not match our modeling assumptions, as our model uses a diagonal covariance for efﬁciency.
Even so, the training loss of M EKA follows that of gradient descent very closely after just a few iterations.
Classiﬁcation Experiments
Next we test and diagnose our approach on classiﬁcation benchmarks, MNIST and CIFAR-10.
We use JAX’s (Bradbury et al., 2018) vectorized map functionality for efﬁcient per-example gradients and
Hessian-vector products.
For MNIST, we test using a multi-layer perceptron (MLP); for CIFAR-10, a convolutional neural network (CNN) and a residual network (ResNet-32) (He et al., 2016a,b).
One key distinction is we replace the batch normalization layers with group normalization (Wu and He, 2018) as batch-dependent transformations conﬂict with our assumption that the gradient samples are independent.
We note that the empirical per-iteration cost of M EKA is 1 . 0 – 1 . 6 × that of SGD due to the computation of Hessian-vector products.
Full experiment details are provided in Appendix E. A detailed comparison to tuned baseline optimizers is presented in Appendix D.1.
[Figure]
Online Variance Reduction We test whether the ﬁltering procedure is correctly aligning the gradient estimate with the true gradient.
For this, we use CIFAR-10 with a
CNN and no data augmentation, so that the true full-batch gradient over the entire dataset can be computed.
Figureshows the L 2 norm difference between the gradient estimators and the full-batch gradient ∇ f t .
M EKA ’s estimated gradients are closer to the true around by around a factor of 5 compared to the minibatch gradient sample.
Adaptive Step Sizes are Appropriately Scaled Without uncertainty quantiﬁcation, the quadratic minimum step size scheme tends to result in step sizes too large.
As such, one may include a scaling factor such that the update is modiﬁed as θ t[-+
-]
+1 =
θ t − cα t δ t .
In contrast, we ﬁnd that the adaptive step sizes based on probability of improve-
[Figure]
ment (PI) are already correctly scaled in the sense that a c different from 1 . 0 will generally result in worse performance.
Figure 6 shows a comparison of different values for c for the quadratic and
PI (12) adaptive schemes.
We plot expected improvement in Appendix D, which performs poorly and requires non-unit scaling factors.
7.1 Adaptive Step Sizes Dives into High-curvature, High-variance Regions
A core aspect of our ﬁltering approach is the ability to estimate quantities of interest during optimization.
We now use these to help understand the loss landscape of ResNet-32 on CIFAR-10.
We ﬁnd that a cause for slow convergence of M EKA with adaptive step sizes is due to an abundance of minima that are usually too high variance for standard SGD.
Figure 7 shows estimates of the normalized curvature along the descent direction δ T B t δ δ T δ as well as the per-sample gradient variance, averaged over parameters.
To understand the loss landscape along the trajectory of optimization, we use multiple runs of M EKA with the same initialization.
Each run takes a different ﬁxed number of constant-size steps before switching to the adaptive step size scheme.
It is clear that immediately after switching to adaptive step sizes, M EKA falls into an increasingly high curvature region and remains there.
The gradient variance also remains high.
As our optimization procedure can handle relatively high variance and curvature, it proceeds to optimize within this sharp but potentally non-local minimum.
On the other hand, it may be an advantage of ﬁxed-step-size SGD that it skips over both high-variance and high-curvature minima.
This failing of adaptive step sizes during the initial phase of training may be related to the “short horizon bias” (Wu et al., 2018) of our one-step-ahead acquisition function.
If so, compute budget can be used to approximate multi-step-ahead gains to help reduce this bias.
Additionally, the ability to optimize within high-curvature high-variance regions could potentially be an advantage on problems with fewer local minima, yet this may not be the case for deep learning.
Conclusion
We introduced an online gradient estimation framework for stochastic gradient-based optimization, which leverages Hessian-vector products and variance estimates to perform automatic online gradient estimation and step size selection.
The result is a stochastic optimization algorithm that can self-tune many important parameters such as momentum and learning rate schedules, in an online fashion without checkpointing or expensive outer-loop optimization.
While the required additional observables can be computed efﬁciently with recent advances in automatic differentiation tooling, they are of course not free, increasing computational cost and memory usage compared to SGD.
What one gains in return is automation, so that it sufﬁces to run the algorithm just once, without tedious tuning.
Given the amount of human effort and computational resources currently invested into hyperparameter tuning, we believe our contributions are valuable steps towards fully-automated gradient-based optimization.
A The Full M EKA Algorithm with Adaptive Step Sizes
{+Algorithm 1 M EKA with adaptive step sizes based on maximizing the probability of improvement. Hyperparameters: decay rates β r =0.999, β Σ =0.999, β α =0. m 0 , P 0 ← (cid:126) 0 , 10(cid:46) large variance initialization ensures ﬁrst Kalman gain is one u 0 , s 0 ← 0 , 10δ 
 = (cid:126)t ←repeat t ← t +f ( i ) t , ∇ f ( i ) t , ∇ 2 f ( i ) t δ t − 1 ← VectorizedMap ( f, { x i } Mi =1 ; θ t − 1 ) (cid:46) compute per-example quantities+}
[Table]
[Table]
←
[Table]
[Table] [Table] (cid:46) exponential moving average (EMA) on the variances g [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
, y [Table] [Table] [Table]
[Table]
(cid:46) ﬁlter update equations m [Table] , P [Table]
←
[Table]
[Table]
[Table]
[Table]
t ) (cid:46) ﬁlter update equations
[Table]
min
[Table]
13) (cid:46) with an EMA (decay rate β
[Table]
) on the coefﬁcients
[Table]
m
[Table]
←
[Table]
until convergence
B The Function Value Dynamics Model
We discuss inferring the function value f t , taking into account uncertainty due to changes in function value and observation noise during optimization.
The gradient dynamics (5) imply the following dynamics model for the function value itself:
[Equation]
where we again use a quadratic approximation using Taylor expansion.
Instead of the intractable ∇ f and ∇ 2 f , we use the estimates from Section 2.
The observations y t and r t are the empirical mean and variance of f t from a minibatch.
The variance terms in the dynamics model are due to the uncertainty associated with m t − 1 and B t δ t − 1 .
Here we have included a scalar term λ t , which acts as a correction to the local quadratic model.
This acts similar to a damping component, except we can automatically infer an optimal λ t by maximizing the likelihood of p ( y t | y 1: t − 1 ) , with a closed form solution (see Appendix B.1).
While damping terms are usually difﬁcult to set empirically (Choi et al., 2019), we note that including our λ t term is essentially free, and it automatically decays after optimization stabilizes.
B.1 Adaptively Correcting the Dynamics Model
We construct a dynamics model of the function value as follows (repeated for convenience):
[Equation]
We include a scalar parameter λ t in case the local quadratic approximation is inaccurate, ie.
when y t is signiﬁcantly different from the predicted value.
If this occurs, a high value of λ t causes the
Kalman gain to become large, throwing away the stale estimate and putting more weight on the new observed function value.
We pick a value for λ t by maximizing the likelihood of p ( y t | y 1: t − 1 ) .
Marginalizing over f t , we get p ( y t | y 1:
[Equation]
−
[Equation]
) =
N
[Equation]
Taking the log and writing it out, we get
[Equation]
and its derivative is
[Equation]
Setting this to zero, we get
[Equation]
Since the role of λ t is to ensure we are not overconﬁdent in our predictions, and we don’t want to deal with negative variance values, we set
[Equation]
As can be seen from Figure 8, this λ t term goes to zero when it is not needed, ie.
when the dynamics model is correct, which occurs on MNIST after convergence.
It is also a quantity that shows us just how incorrect our dynamics model is, and for the problems we tested, we ﬁnd that it is signiﬁcantly smaller than the posterior variance s t .
This suggests that it has minimal impact if removed, but we keep it in the algorithm for cases when a quadratic approximation is not sufﬁcient.
[Figure]
B.2 Dealing with negative curvature in step size adaptation
When the gradient estimating is pointing in a direction of negative curvature, there is a chance that the optimal step size is inﬁnity (Figure 9a).
[Figure]
This occurs when the variance of f t +1 − f t rises slower than the expectation.
To handle this situation, we can add a third order correction term, which appears in the variance as a term that scales with || δ || 6 .
Using the same procedure as inferring a constant λ t , we can instead add the term λ t ||
δ || 6 to the variance.
We then choose λ t as
[Equation]
This extra term (if λ t > 0 ) in the variance ensures that variance increases faster than the expectation.
Adaptive step sizes based on the probability of improvement will then have an optimal step size that is ﬁnite in value (Figure 9b).
An additional damping effect may be added by lower bounding λ t .
We did not fully test this approach as the exponential moving averaged curvature used in practice was always positive for our test problems.
Incidentally, Figures 9a and 9b show that the expected improvement is not a good heuristic as it is extremely large even with this extra term.
Moreover, the quadratic approximation will result in negative step sizes.
C Using the Current Hessian vs the Previous Hessian
For the dynamics model, we make the choice to use the Hessian at the updated location B t , which is an unbiased estimate of ∇ 2 f ( θ t ) , instead of B t − 1 , the Hessian at θ t − 1 .
Firstly, we made this choice for computational reasons: it is easier to compute the Hessian at θ t since we are already computing the gradient g t evaluated at θ t .
Secondly, we found that using the current Hessian results in better performance and more stability.
Figure 10 shows this on MNIST.
For CIFAR-10, we found that using the previous Hessian B t − 1 resulted in immediate divergence and NaNs, so we do not show those plots.
[Figure]
D Comparison of Adaptive Step Size Schemes
[Figure]
T r a i n i n g L o ss
MLP on MNIST
[Figure]
CNN on CIFAR-
[Figure]
c
[Figure]
EI c
[Figure]
c
[Figure]
EI c
[Figure] PI [Figure]
[Figure]
D.1 Comparison with Tuned Optimizers
We measure the performance of our method against a variety of other approaches.
We compare with ﬁxed step size versions of SGD, SGD with momentum (Polyak, 1964)
[Figure]
and A DAM (Kingma and Ba, 2014).
For these, we tune the step size using grid search.
We also compare against A DADELTA (Zeiler, 2012), which is a competing learning rate-free algorithm, and SVRG (Johnson and Zhang, 2013), an optimization algorithm focused around estimating the full batch gradient, using the same step size as the tuned SGD.
For comparison, we implemented M EKA using the same constant learning rate as the tuned SGD, and A DA M EKA with the same learning rate as the tuned A DAM .
We applied fully adaptive step sizes with M EKA update directions; on ResNet-32, we used A DA M EKA update directions to mitigate poor conditioning.
Figure 12 shows the resulting loss and accuracy curves.
As this includes optimizers across a wide range of motivations, we highlight some speciﬁc comparisons.
M EKA generally performs better than SGD in terms of test accuracy, showing that too much stochasticity can hurt generalization.
Though both designed with gradient estimation in mind, M EKA seems to compare favorable against
SVRG in terms of performance.
We note that the per-iteration costs of M EKA were also cheaper as
SVRG requires two gradient evaluations.
With the default learning rate of 1 . 0 , A DADELTA performs decently on MNIST but ends up diverging on CIFAR-10.
In comparison, our adaptive step sizes converge well and generally outperform ﬁxed step sized M EKA (or A DA M EKA ) in both trainng loss and test accuracy.
E Experiment Details E.1 Dataset Description
We used the ofﬁcial train and test split for MNIST and CIFAR-10.
We did not do any data augmentation for MNIST.
For CIFAR-10, we normalized the images by subtracting every pixel with the global mean and standard deviation across the training set.
In addition to this, unless speciﬁed otherwise, we pre-processed the images following He et al.
(2016a)
by padding the images by 4 pixels on each side and applying random cropping and horizontal ﬂips.
E.2 Architecture Description
All models use the the ReLU (Nair and Hinton, 2010) activation function.
MLP We used an MLP with 1 hidden layer of 100 hidden units.
CNN We used the same architecture as the “3c3d” architecture in Schneider et al.
(2019), which consists of 3 convolutional layers with max pooling, followed by 3 fully connected layers.
The ﬁrst convolutional layer has a kernel size of 5 × 5 with stride 1, “valid” padding, and 64 ﬁlters.
The second convolutional layer has a kernel size of 3 × 3 with stride 1, “valid” padding, and 96 ﬁlters.
The third convolutional layer has a kernel size of 3 × 3 with stride 1, “same” padding, and 128 ﬁlters.
The max pooling layers have a window size of 3 × 3 with stride 2.
The 2 fully connected layers have 512 and 256 units respectively.
ResNet-32 Our ResNet-32 (He et al., 2016a) model uses residual blocks based on He et al.
(2016b).
We replaced the batch normalization layers with group normalization (Wu and He, 2018) as batchdependent transformations conﬂict with our assumption that the gradient samples are independent, and hinder our method to estimate gradient variance.
E.3 Optimizer Comparisons Description
We tuned the step size of SGD in a grid of { 0 . 001 , 0 . 01 , 0 . 1 , 1 . 0 } .
We tuned the step size of SGD with momentum, and A DAM in a grid of { 0 . 0001 , 0 . 001 , 0 . 01 , 0 . 1 } .
We chose the best step size of
SGD for the variant of M EKA with a constant learning rate, and for SVRG.
The chosen step size for SGD was 0.1 for MNIST, and 0.1 for CIFAR-10.
For SGD with momentum, the chosen step size was 0.01 for MNIST, 0.1 for ResNet-32 on CIFAR-10, and 0.001 for CNN on
CIFAR-10.
The best step size for A DAM was 0.001 for MNIST and ResNet-32 on CIFAR-10, and 0.0001 for CNN on CIFAR-10.
For SGD with momentum, the momentum coefﬁcient γ was ﬁxed to 0.9.
For A DAM , β 1 , β 2 , ε were ﬁxed to 0.9, 0.999, and 10 − 8 respectively.
For A DADELTA , ρ and ε were ﬁxed to 0.95 and 10 −respectively.
F Sensitivity of M EKA
’s Hyperparameters
[Figure]
cc
[Figure]
[Figure]
A cc u r a c
[Figure]
G Additional Cost of M EKA
[Table] [Table]
H Proofs
Proof of Proposition 1.
A standard Lipschitz bound yields
[Equation]
Using strong convexity ( (cid:107)∇ f t (cid:107) [Equation] f t − f ∗ ) ) and subtracting f ∗ from both sides results in
[Equation]
So in each step, we get a multiplicative decrease in the expected function value (left term) but we add a term that depends on the variance of our ﬁltered gradient estimate m t .
So, in essence, to establish convergence, we have to show that E [ (cid:107) m t − ∇ f t (cid:107) 2 ] decreases to zero sufﬁciently fast.
Since all assumptions of the Kalman ﬁlter are satisﬁed, we know that E [ m t ] =
E [ ∇ f t ] and E [( m t − ∇ f t )( m t − ∇ f t ) T ] =
P t .
Hence, E [ (cid:107) m t − ∇ f t (cid:107) 2 ] = tr( P t ) .
We now show inductively that
P t = 1 t +1 Σ .
This holds for t = 0 by construction.
Assume it holds for arbitrary but ﬁxed t − 1 .
Then
[Equation]
and, thus,
P t =
[Equation]
I [Equation] Plugging E [ (cid:107) m t − ∇ f t (cid:107) 2 ] = tr( P t ) =
1 t +1 tr(Σ) back into Eq.
and introducing the shorthands e t =
E [ f t − f ∗ ] and σ 2 := tr(Σ) reads
[Equation]
[Equation]
Iterating backwards results in
[Equation]
Lemma 1 shows that the sum term is O (1 /t ) .
The ﬁrst (exponential) term is trivially O (1 /t ) , which concludes the proof.
Lemma 1.
Let 0 < c < 1 and deﬁne the sequence (for t ≥ 1 )
[Equation]
Then a t ∈ O ( 1 t ) .
Proof.
Let T be the smallest index such that c T +1 T
[Equation]
This ensures that a T ≤ MT and
[Equation]
) for all t ≥ T .
We now show inductively that a t ≤ Mt for all t ≥ T .
It holds for t =
T by construction of M .
Assume it holds for some t ≥ T .
Then
[Equation]