Reducing the variance in online optimization by transporting past gradients

  • 2019-06-08 22:02:28
  • Sébastien M. R. Arnold, Pierre-Antoine Manzagol, Reza Babanezhad, Ioannis Mitliagkas, Nicolas Le Roux
  • 24

Abstract

Most stochastic optimization methods use gradients once before discardingthem. While variance reduction methods have shown that reusing past gradientscan be beneficial when there is a finite number of datapoints, they do noteasily extend to the online setting. One issue is the staleness due to usingpast gradients. We propose to correct this staleness using the idea of implicitgradient transport (IGT) which transforms gradients computed at previousiterates into gradients evaluated at the current iterate without using theHessian explicitly. In addition to reducing the variance and bias of ourupdates over time, IGT can be used as a drop-in replacement for the gradientestimate in a number of well-understood methods such as heavy ball or Adam. Weshow experimentally that it achieves state-of-the-art results on a wide rangeof architectures and benchmarks. Additionally, the IGT gradient estimatoryields the optimal asymptotic convergence rate for online stochasticoptimization in the restricted setting where the Hessians of all componentfunctions are equal.

 

Quick Read (beta)

Reducing the variance in online optimization by transporting past gradients

Sébastien M. R. Arnold
University of Southern California
Los Angeles, CA
[email protected]
&Pierre-Antoine Manzagol
Google Brain
Montréal, QC
[email protected]
&Reza Babanezhad
University of British Columbia
Vancouver, BC
[email protected]
&Ioannis Mitliagkas
Mila, Université de Montréal
Montréal, QC
[email protected]
&Nicolas Le Roux
Mila, Google Brain
Montréal, QC
[email protected]
Work done while at Mila.
Abstract

Most stochastic optimization methods use gradients once before discarding them. While variance reduction methods have shown that reusing past gradients can be beneficial when there is a finite number of datapoints, they do not easily extend to the online setting. One issue is the staleness due to using past gradients. We propose to correct this staleness using the idea of implicit gradient transport (IGT) which transforms gradients computed at previous iterates into gradients evaluated at the current iterate without using the Hessian explicitly. In addition to reducing the variance and bias of our updates over time, IGT can be used as a drop-in replacement for the gradient estimate in a number of well-understood methods such as heavy ball or Adam. We show experimentally that it achieves state-of-the-art results on a wide range of architectures and benchmarks. Additionally, the IGT gradient estimator yields the optimal asymptotic convergence rate for online stochastic optimization in the restricted setting where the Hessians of all component functions are equal.11 1 Open-source implementation available at: https://github.com/seba-1511/igt.pth

 

Reducing the variance in online optimization by transporting past gradients


  Sébastien M. R. Arnold thanks: Work done while at Mila. University of Southern California Los Angeles, CA [email protected] Pierre-Antoine Manzagol Google Brain Montréal, QC [email protected] Reza Babanezhad University of British Columbia Vancouver, BC [email protected] Ioannis Mitliagkas Mila, Université de Montréal Montréal, QC [email protected] Nicolas Le Roux Mila, Google Brain Montréal, QC [email protected]

\@float

noticebox[b]Preprint. Under review.\[email protected]

1 Introduction

We wish to solve the following minimization problem:

θ =argminθExp[f(θ,x)], (1)

where we only have access to samples x and to a first-order oracle that gives us, for a given θ and a given x, the derivative of f(θ,x) with respect to θ, i.e. f(θ,x)θ=g(θ,x). It is known robbins1951stochastic that, when f is smooth and strongly convex, there is a converging algorithm for Problem 1 that takes the form θt+1=θt-αtg(θt,xt), where xt is a sample from p. This algorithm, dubbed stochastic gradient (SG), has a convergence rate of O(1/t) (see for instance bubeck2015convex ), within a constant factor of the minimax rate for this problem. When one has access to the true gradient g(θ)=Exp[g(θ,x)] rather than just a sample, this rate dramatically improves to O(e-νt) for some ν>0.

In addition to hurting the convergence speed, noise in the gradient makes optimization algorithms harder to tune. Indeed, while full gradient descent is convergent for constant stepsize α, and also amenable to line searches to find a good value for that stepsize, the stochastic gradient method from robbins1951stochastic with a constant stepsize only converges to a ball around the optimum schmidt2014convergence .22 2 Under some conditions, it does converge linearly to the optimum (e.g., Vaswani19, ) Thus, to achieve convergence, one needs to use a decreasing stepsize. While this seems like a simple modification, the precise decrease schedule can have a dramatic impact on the convergence speed. While theory prescribes αt=O(t-α) with α(1/2,1] in the smooth case, practictioners often use larger stepsizes like αt=O(t-1/2) or even constant stepsizes.

When the distribution p has finite support, Eq. 1 becomes a finite sum and, in that setting, it is possible to achieve efficient variance reduction and drive the noise to zero, allowing stochastic methods to achieve linear convergence rates leroux2012stochastic ; johnson2013accelerating ; zhang2013linear ; mairal2013optimization ; Shalev-Shwartz2013sdca ; defazio2014saga . Unfortunately, the finite support assumption is critical to these algorithms which, while valid in many contexts, does not have the broad applicability of the standard SG algorithm. Several works have extended these approaches to the online setting by applying these algorithms while increasing N babanezhad2015wasting ; hofmann2015variance but they need to revisit past examples mutiple times and are not truly online.

Another line of work reduces variance by averaging iterates polyak1992acceleration ; lacoste2012simpler ; bach2013non ; flammarion2015averaging ; dieuleveut2017harder ; dieuleveut2017bridging ; JMLR:v18:16-595 . While these methods converge for a constant stepsize in the stochastic case33 3 Under some conditions on f., their practical speed is heavily dependent on the fraction of iterates kept in the averaging, a hyperparameter that is thus hard to tune, and they are rarely used in deep learning.

Our work combines two existing ideas and adds a third: a) At every step, it updates the parameters using a weighted average of past gradients, like in SAG (leroux2012stochastic, ; schmidt2017minimizing, ), albeit with a different weighting scheme; b) It reduces the bias and variance induced by the use of these old gradients by transporting them to “equivalent” gradients computed at the current point, similar to gower2017tracking ; c) It does so implicitly by computing the gradient at a parameter value different from the current one. The resulting gradient estimator can then be used as a plug-in replacement of the stochastic gradient within any optimization scheme. Experimentally, both SG using our estimator and its momentum variant outperform the most commonly used optimizers in deep learning.

2 Momentum and other approaches to dealing with variance

Stochastic variance reduction methods use an average of past gradients to reduce the variance of the gradient estimate. At first glance, it seems like their updates are similar to that of momentum polyak1964some , also known as the heavy ball method, which performs the following updates44 4 This is slightly different from the standard formulation but equivalent for constant γt.:

vt =γtvt-1+(1-γt)g(θt,xt),v0=g(θ0,x0)
θt+1 =θt-αtvt.

When γt=γ, this leads to θt+1=θt-αt(γtg(θ0,x0)+(1-γ)i=1tγt-ig(θi,xi)). Hence, the heavy ball method updates the parameters of the model using an average of past gradients, bearing similarity with SAG leroux2012stochastic , albeit with exponential instead of uniform weights.

Interestingly, while momentum is a popular method for training deep networks, its theoretical analysis in the stochastic setting is limited (sutskever2013importance, ), except in the particular setting when the noise converges to 0 at the optimum (loizou2017momentum, ). Also surprising is that, despite the apparent similarity with stochastic variance reduction methods, current convergence rates are slower when using γ>0 in the presence of noise Schmidt11_inexact , although this might be a limitation of the analysis.

2.1 Momentum and variance

We propose here an analysis of how, on quadratics, using past gradients as done in momentum does not lead to a decrease in variance. If gradients are stochastic, then Δt=θt-θ is a random variable. Denoting ϵi the noise at timestep i, i.e. g(θi,xi)=g(θi)+ϵi, and writing Δt-E[Δt]=αi=0tNi,tϵi, with Ni,t the impact of the noise of the i-th datapoint on the t-th iterate, we may now analyze the total impact of each ϵi on the iterates. Figure 1 shows the impact of ϵi on Δt-E[Δt] as measured by Ni,t2 for three datapoints (i=1, i=25 and i=50) as a function of t for stochastic gradient (γ=0, left) and momentum (γ=0.9, right). As we can see, when using momentum, the variance due to a given datapoint first increases as the noise influences both the next iterate (through the parameter update) and the subsequent updates (through the velocity). Due to the weight 1-γ when a point is first sampled, a larger value of γ leads to a lower immediate impact of the noise of a given point on the iterates. However, a larger γ also means that the noise of a given gradient is kept longer, leading to little or no decrease of the total variance (dashed blue curve). Even in the case of stochastic gradient, the noise at a given timestep carries over to subsequent timesteps, even if the old gradients are not used for the update, as the iterate itself depends on the noise.

(a) Stochastic gradient
(b) Momentum - γ=0.9
(c) Momentum - γt=1-1t
(d) Momentum - γt=1-1t with IGT.
Figure 1: Variance induced over time by the noise from three different datapoints (i=1, i=25 and i=50) as well as the total variance for SG (γ=0, top left), momentum with fixed γ=0.9 (top right), momentum with increasing γt=1-1t without (bottom left) and with (bottom right) transport. The impact of the noise of each gradient ϵi increases for a few iterations then decreases. Although a larger γ reduces the maximum impact of a given datapoint, the total variance does not decrease. With transport, noises are now equal and total variance decreases. The y-axis is on a log scale.

At every timestep, the contribution to the noise of the 1st, the 25th and the 50th points in Fig. 1 is unequal. If we assume that the ϵi are i.i.d., then the total variance would be minimal if the contribution from each point was equal. Further, one can notice that the impact of datapoint i is only a function of t-i and not of t. This guarantees that the total noise will not decrease over time.

To address these two points, one can increase the momentum parameter over time. In doing so, the noise of new datapoints will have a decreasing impact on the total variance as their gradient is multiplied by 1-γt. Figure 0(c) shows the impact Ni,t2 of each noise ϵi for an increasing momentum γt=1-1t. The peak of noise for i=25 is indeed lower than that of i=1. However, the variance still does not go to 0. This is because, as the momentum parameter increases, the update is an average of many gradients, including stale ones. Since these gradients were computed at iterates already influenced by the noise over previous datapoints, that past noise is amplified, as testified by the higher peak at i=1 for the increasing momentum. Ultimately, increasing momentum does not lead to a convergent algorithm in the presence of noise when using a constant stepsize.

2.2 SAG and Hessian modelling

The impact of the staleness of the gradients on the convergence is not limited to momentum. In SAG, for instance, the excess error after k updates is proportional to (1-min{116κ^,18N})k, compared to the excess error of the full gradient method which is (1-1κ)k where κ is the condition number of the problem. 55 5 The κ^ in the convergence rate of SAG is generally larger than the κ in the full gradient algorithm. The difference between the two rates is larger when the minimum in the SAG rate is the second term. This happens either when κ^ is small, i.e. the problem is well conditioned and a lot of progress is made at each step, or when N is large, i.e. there are many points to the training set. Both cases imply that a large distance has been travalled between two draws of the same datapoint.

Recent works showed that correcting for that staleness by modelling the Hessian wai2017curvature ; gower2017tracking leads to improved convergence. As momentum uses stale gradients, the velocity is an average of current and past gradients and thus can be seen as an estimate of the true gradient at a point which is not the current one but rather a convex combination of past iterates. As past iterates depend on the noise of previous gradients, this bias in the gradients amplifies the noise and leads to a non-converging algorithm. We shall thus “transport” the old stochastic gradients g(θi,xi) to make them closer to their corresponding value at the current iterate, g(θt,xi). Past works did so using the Hessian or an explicit approximation thereof, which can be expensive and difficult to compute and maintain. We will resort to using implicit transport, a new method that aims at compensating the staleness of past gradients without making explicit use of the Hessian.

3 Converging optimization through implicit gradient transport

Before showing how to combine the advantages of both increasing momentum and gradient transport, we demonstrate how to transport gradients implicitly. This transport is only exact under a strong assumption that will not hold in practice. However, this result will serve to convey the intuition behind implicit gradient transport. We will show in Section 4 how to mitigate the effect of the unsatisfied assumption.

3.1 Implicit gradient transport

Let us assume that we received samples x0,,xt in an online fashion. We wish to approach the full gradient gt(θt)=1t+1i=0tg(θt,xi) as accurately as possible. We also assume here that a) We have a noisy estimate g^t-1(θt-1) of gt-1(θt-1); b) We can compute the gradient g(θ,xt) at any location θ. We shall seek a θ such that

tt+1g^t-1(θt-1)+1t+1g(θ,xt)gt(θt).

To this end, we shall make the following assumption:

Assumption 3.1.

All individual functions f(,x) are quadratics with the same Hessian H.

This is the same assumption as (flammarion2015averaging, , Section 4.1). Although it is unlikely to hold in practice, we shall see that our method still performs well when that assumption is violated.

Under Assumption 3.1, we then have (see details in Appendix)

gt(θt) =tt+1gt-1(θt)+1t+1g(θt,xt)
tt+1g^t-1(θt-1)+1t+1g(θt+t(θt-θt-1),xt).

Thus, we can transport our current estimate of the gradient by computing the gradient on the new point at a shifted location θ=θt+t(θt-θt-1). This extrapolation step is reminiscent of Nesterov’s acceleration with the difference that the factor in front of θt-θt-1, t, is not bounded.

3.2 Combining increasing momentum and implicit gradient transport

We now describe our main algorithm, Implicit Gradient Transport (IGT). IGT uses an increasing momentum γt=tt+1. At each step, when updating the velocity, it computes the gradient of the new point at an extrapolated location so that the velocity vt is a good estimate of the true gradient g(θt).

We can rewrite the updates to eliminate the velocity vt, leading to the update:

θt+1 =2t+1t+1θt-tt+1θt-1-αt+1g(θt+t(θt-θt-1),xt). (IGT)

We see in Fig. 0(d) that IGT allows a reduction in the total variance, thus leading to convergence with a constant stepsize. This is captured by the following proposition:

Proposition 3.1.

If f is a quadratic function with positive definite Hessian H with largest eigenvalue L and condition number κ and if the stochastic gradients satisfy: g(θ,x)=g(θ)+ϵ with ϵ a random i.i.d. noise with covariance bounded by BI, then Eq. IGT with stepsize α=1/L leads to iterates θt satisfying

E[θt-θ2] (1-1κ)2tθ0-θ2+dα2Bν¯02t,

with ν=(2+2logκ)κ for every t>2κ.

The proof of Prop. 3.1 is provided in the appendix.

Despite this theoretical result, two limitations remain: First, Prop. 3.1 shows that IGT does not improve the dependency on the conditioning of the problem; Second, the assumption of equal Hessians is unlikely to be true in practice, leading to an underestimation of the bias. We address the conditioning issue in the next section and the assumption on the Hessians in Section 4.

3.3 IGT as a plug-in gradient estimator

We demonstrated that the IGT estimator has lower variance than the stochastic gradient estimator for quadratic objectives. IGT can also be used as a drop-in replacement for the stochastic gradient in an existing, popular first order method: the heavy ball (HB). This is captured by the following two propositions:

Proposition 3.2 (Non-stochastic).

In the non-stochastic case, where B=0, variance is equal to 0 and Heavyball-IGT achieves the accelerated linear rate O((κ-1κ+1)t) using the known, optimal heavy ball tuning, μ=(κ-1κ+1)2, α=(1+μ)2/L.

Proposition 3.3 (Online, stochastic).

When B>0, there exist constant hyperparameters α>0, μ>0 such that E[θt-θ]2 converges to zero linearly, and the variance is O~(1/t).

The pseudo-code can be found in Algorithm 3.3.

{algorithm}

Heavyball-IGT {algorithmic}[1] \ProcedureHeavyball-IGTStepsize α, Momentum μ, Initial parameters θ0 \Statev0g(θ0,x0),w0-αv0,θ1θ0+w0 \Fort=1,,T-1 \Stateγttt+1 \Statevtγtvt-1+(1-γt)g(θt+γt1-γt(θt-θt-1),xt) \Statewtμwt-1-αvt \Stateθt+1θt+wt \EndFor \Statereturn θT \EndProcedure

4 IGT and Anytime Tail Averaging

So far, IGT weighs all gradients equally. This is because, with equal Hessians, one can perfectly transport these gradients irrespective of the distance travelled since they were computed. In practice, the individual Hessians are not equal and might change over time. In that setting, the transport induces an error which grows with the distance travelled. We wish to average a linearly increasing number of gradients, to maintain the O(1/t) rate on the variance, while forgetting about the oldest gradients to decrease the bias. To this end, we shall use anytime tail averaging leroux2019anytime , named in reference to the tail averaging technique used in optimization JMLR:v18:16-595 .

Tail averaging is an online averaging technique where only the last points, usually a constant fraction c of the total number of points seen, is kept. Maintaining the exact average at every timestep is memory inefficient and anytime tail averaging performs an approximate averaging using γt=c(t-1)1+c(t-1)(1-1c1-ct(t-1)). We refer the reader to leroux2019anytime for additional details.

5 Impact of IGT on bias and variance in the ideal case

To understand the behaviour of IGT when Assumption 3.1 is verified, we minimize a strongly convex quadratic function with Hessian Q100×100 with condition number 1000, and we have access to the gradient corrupted by noise ϵt, where ϵtN(0,0.3I100). In that scenario where all Hessians are equal and implicit gradient transport is exact, Fig. 1(a) confirms the O(1/t) rate of IGT with constant stepsize while SGD and HB only converge to a ball around the optimum.

To further understand the impact of IGT, we study the quality of the gradient estimate. Standard stochastic methods control the variance of the parameter update by scaling it with a decreasing stepsize, which slows the optimization down. With IGT, we hope to have a low variance while maintaining a norm of the update comparable to that obtained with gradient descent. To validate the quality of our estimator, we optimized a quadratic function using IGT, collecting iterates θt. For each iterate, we computed the squared error between the true gradient and either the stochastic or the IGT gradient. In this case where both estimators are unbiased, this is the trace of the noise covariance of our estimators. The results in Figure 1(b) show that, as expected, this noise decreases linearly for IGT and is constant for SGD.

We also analyse the direction and magnitude of the gradient of IGT on the same quadratic setup. Figure 1(c) displays the cosine similarity between the true gradient and either the stochastic or the IGT gradient, as a function of the distance to the optimum. We see that, for the same distance, the IGT gradient is much more aligned with the true gradient than the stochastic gradient is, confirming that variance reduction happens without the need for scaling the estimate.

(a)
(b)
(c)
Figure 2: Analysis of IGT on quadratic loss functions. (\subreffig:quad) Comparison of convergence curves for multiple algorithms. As expected, the IGT family of algorithms converges to the solution while stochastic gradient algorithms can not. (\subreffig:noise) The blue and orange curves show the norm of the noise component in the SGD and IGT gradient estimates, respectively. The noise component of SGD remains constant, while it decreases at a rate 1/t for IGT. The green curve shows the norm of the IGT gradient estimate. (\subreffig:cosine) Cosine similarity between the full gradient and the SGD/IGT estimates.

6 Experiments

While Section 5 confirms the performance of IGT in the ideal case, the assumption of identical Hessians almost never holds in practice. In this section, we present results on more realistic and larger scale machine learning settings. All experiments are extensively described in the Appendix A and additional baselines compared in Appendix B.

6.1 Supervised learning

CIFAR10 image classification

We first consider the task of training a ResNet-56 model DBLP:journals/corr/HeZRS15 on the CIFAR-10 image classification dataset cifar10 . We use TF official models code and setup TfResnet , varying only the optimizer: SGD, HB, Adam and our algorithm with anytime tail averaging both on its own (ITA) and combined with Heavy Ball (HB-ITA). We tuned the step size for each algorithm by running experiments using a logarithmic grid. To factor in ease of tuning wilson2017marginal , we used Adam’s default parameter values and a value of 0.9 for HB’s parameter. We used a linearly decreasing stepsize as it was shown to be simple and perform well shallue2018measuring . For each optimizer we selected the hyperparameter combination that is fastest to reach a consistently attainable target train loss shallue2018measuring . Selecting the hyperparameter combination reaching the lowest training loss yields qualitatively identical curves. Figure 3 presents the results, showing that IGT with the exponential anytime tail average performs favourably, both on its own and combined with Heavy Ball: the learning curves show faster improvement and are much less noisy.

Figure 3: Resnet-56 on CIFAR10. Left: Train loss. Center: Train accuracy. Right: Test accuracy.
ImageNet image classification

We also consider the task of training a ResNet-50 modelDBLP:journals/corr/HeZRS15 on the larger ImageNet dataset ILSVRC15 . The setup is similar to the one used for CIFAR10 with the difference that we trained using larger minibatches (1024 instead of 128). In Figure 4, one can see that IGT is as fast as Adam for the train loss, faster for the train accuracy and reaches the same final performance, which Adam does not. We do not see the noise reduction we observed with CIFAR10, which could be explained by the larger batch size (see Appendix A.1).

Figure 4: ResNet-50 on ImageNet. Left: Train loss. Center: Train accuracy. Right: Test accuracy.
IMDb sentiment analysis

We train a bi-directional LSTM on the IMDb Large Movie Review Dataset for 200 epochs. maas2011learning We observe that while the training convergence is comparable to HB, HB-ITA performs better in terms of validation and test accuracy. In addition to the baseline and IGT methods, we also train a variant of Adam using the ITA gradients, dubbed Adam-ITA, which performs similarly to Adam.

Figure 5: Validation curves for different large-scale machine learning settings. Shading indicates one standard deviation computed over three random seeds. Left: Reinforcement learning via policy gradient on a LQR system. Right: Meta-learning using MAML on Mini-Imagenet.

6.2 Reinforcement learning

Linear-quadratic regulator

We cast the classical linear-quadratic regulator (LQR) kwakernaak1972linear as a policy learning problem to be optimized via gradient descent. This setting is extensively described in Appendix A. Note that despite their simple linear dynamics and a quadratic cost functional, LQR systems are notoriously difficult to optimize due to the non-convexity of the loss landscape. fazel2018global

The left chart in Figure 5 displays the evaluation cost computed along training and averaged over three random seeds. The first method (Optimal) indicates the cost attained when solving the algebraic Riccati equation of the LQR – this is the optimal solution of the problem. SGD minimizes the costs using the REINFORCE williams1992simple gradient estimator, averaged over 600 trajectories. ITA is similar to SGD but uses the ITA gradient computed from the REINFORCE estimates. Finally, GD uses the analytical gradient by taking the expectation over the policy.

We make two observations from the above chart. First, ITA initially suffers from the stochastic gradient estimate but rapidly matches the performance of GD. Notably, both of them converge to a solution significantly better than SGD, demonstrating the effectiveness of the variance reduction mechanism. Second, the convergence curve is smoother for ITA than for SGD, indicating that the ITA iterates are more likely to induce similar policies from one iteration to the next. This property is particularly desirable in reinforcement learning as demonstrated by the popularity of trust-region methods in large-scale applications. schulman2015trust ; openai2018learning

6.3 Meta-learning

Model-agnostic meta-learning

We now investigate the use of IGT in the model-agnostic meta-learning (MAML) setting. finn2017model We replicate the 5 ways classification setup with 5 adaptation steps on tasks from the Mini-Imagenet dataset ravi2016optimization . This setting is interesting because of the many sources contributing to noise in the gradient estimates: the stochastic meta-gradient depends on the product of 5 stochastic Hessians computed over only 10 data samples, and is averaged over only 4 tasks. We substitute the meta-optimizer with each method, select the stepsize that maximizes the validation accuracy after 10K iterations, and use it to train the model for 100K iterations.

The right graph of Figure 5 compares validation accuracies for three random seeds. We observe that methods from the IGT family significantly outperform their stochastic meta-gradient counter-part, both in terms of convergence rate and final accuracy. Those results are also reflected in the final test accuracies where Adam-ITA (65.16%) performs best, followed by HB-ITA (64.57%), then Adam (63.70%), and finally HB (63.08%).

7 Conclusion and open questions

We proposed a simple optimizer which, by reusing past gradients and transporting them, offers excellent performance on a variety of problems. While it adds an additional parameter, the ratio of examples to be kept in the tail averaging, it remains competitive across a wide range of such values. Further, by providing a higher quality gradient estimate that can be plugged in any existing optimizer, we expect it to be applicable to a wide range of problems. As the IGT is similar to momentum, this further raises the question on the links between variance reduction and curvature adaptation. Whether there is a way to combine the two without using momentum on top of IGT remains to be seen.

Acknowledgments

The authors would like to thank Liyu Chen for his help with the LQR experiments and Fabian Pedregosa for insightful discussions.

References

  • [1] The TensorFlow Authors. Tensorflow official resnet model. 2018.
  • [2] Reza Babanezhad, Mohamed Osama Ahmed, Alim Virani, Mark Schmidt, Jakub Konec̆ný, and Scott Sallinen. Stop wasting my gradients: Practical SVRG. In Advances in Neural Information Processing Systems, 2015.
  • [3] Francis Bach and Eric Moulines. Non-strongly-convex smooth stochastic approximation with convergence rate o (1/n). In Advances in Neural Information Processing Systems, pages 773–781, 2013.
  • [4] Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning, 8(3-4):231–357, 2015.
  • [5] Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. Saga: A fast incremental gradient method with support for non-strongly convex composite objectives. In Advances in Neural Information Processing Systems, pages 1646–1654, 2014.
  • [6] Aymeric Dieuleveut, Alain Durmus, and Francis Bach. Bridging the gap between constant step size stochastic gradient descent and markov chains. arXiv preprint arXiv:1707.06386, 2017.
  • [7] Aymeric Dieuleveut, Nicolas Flammarion, and Francis Bach. Harder, better, faster, stronger convergence rates for least-squares regression. The Journal of Machine Learning Research, 18(1):3520–3570, 2017.
  • [8] Maryam Fazel, Rong Ge, Sham M. Kakade, and Mehran Mesbahi. Global convergence of policy gradient methods for linearized control problems. CoRR, abs/1801.05039, 2018.
  • [9] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 1126–1135. JMLR. org, 2017.
  • [10] Nicolas Flammarion and Francis Bach. From averaging to acceleration, there is only a step-size. In Conference on Learning Theory, pages 658–695, 2015.
  • [11] Robert M Gower, Nicolas Le Roux, and Francis Bach. Tracking the gradients using the hessian: A new look at variance reducing stochastic methods. arXiv preprint arXiv:1710.07462, 2017.
  • [12] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. CoRR, abs/1512.03385, 2015.
  • [13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 630–645. Springer, 2016.
  • [14] Thomas Hofmann, Aurelien Lucchi, Simon Lacoste-Julien, and Brian McWilliams. Variance reduced stochastic gradient descent with neighbors. In Advances in Neural Information Processing Systems, pages 2305–2313, 2015.
  • [15] Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerating stochastic gradient descent for least squares regression. In Conference On Learning Theory, pages 545–604, 2018.
  • [16] Prateek Jain, Sham M. Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Parallelizing stochastic gradient descent for least squares regression: Mini-batching, averaging, and model misspecification. Journal of Machine Learning Research, 18(223):1–42, 2018.
  • [17] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Advances in Neural Information Processing Systems, pages 315–323, 2013.
  • [18] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
  • [19] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
  • [20] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
  • [21] Huibert Kwakernaak. Linear optimal control systems, volume 1. 1972.
  • [22] Simon Lacoste-Julien, Mark Schmidt, and Francis Bach. A simpler approach to obtaining an o (1/t) convergence rate for the projected stochastic subgradient method. arXiv preprint arXiv:1212.2002, 2012.
  • [23] Nicolas Le Roux. Anytime tail averaging. arXiv preprint arXiv:1902.05083, 2019.
  • [24] Nicolas Le Roux, Mark Schmidt, and Francis Bach. A stochastic gradient method with an exponential convergence rate for finite training sets. In Advances in Neural Information Processing Systems, pages 2663–2671, 2012.
  • [25] Yann LeCun and Corinna Cortes. Mnist handwritten digit database. AT&T Labs [Online]. Available: http://yann. lecun. com/exdb/mnist, 2010.
  • [26] Nicolas Loizou and Peter Richtárik. Momentum and stochastic momentum for stochastic gradient, newton, proximal point and subspace descent methods. arXiv preprint arXiv:1712.09677, 2017.
  • [27] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142–150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.
  • [28] Julien Mairal. Optimization with first-order surrogate functions. In Proceedings of The 30th International Conference on Machine Learning, pages 783–791, 2013.
  • [29] OpenAI, Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Józefowicz, Bob McGrew, Jakub W. Pachocki, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, Jonas Schneider, Szymon Sidor, Josh Tobin, Peter Welinder, Lilian Weng, and Wojciech Zaremba. Learning dexterous in-hand manipulation. CoRR, abs/1808.00177, 2018.
  • [30] Brendan O’Donoghue and Emmanuel Candes. Adaptive restart for accelerated gradient schemes. Foundations of computational mathematics, 15(3):715–732, 2015.
  • [31] Jeffrey Pennington, Richard Socher, and Christopher Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pages 1532–1543, 2014.
  • [32] Boris T Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1–17, 1964.
  • [33] Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4):838–855, 1992.
  • [34] Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. 2016.
  • [35] Herbert Robbins and Sutton Monro. A stochastic approximation method. Annals of Mathematical Statistics, 22(3):400–407, 1951.
  • [36] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211–252, 2015.
  • [37] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 4510–4520, 2018.
  • [38] Mark Schmidt. Convergence rate of stochastic gradient with constant step size. 2014.
  • [39] Mark Schmidt, Nicolas Le Roux, and Francis Bach. Convergence rates of inexact proximal-gradient methods for convex optimization. In Advances in Neural Information Processing Systems 24, 2011.
  • [40] Mark Schmidt, Nicolas Le Roux, and Francis Bach. Minimizing finite sums with the stochastic average gradient. Mathematical Programming, 162(1-2):83–112, 2017.
  • [41] John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization. CoRR, abs/1502.05477, 2015.
  • [42] Shai Shalev-Shwartz and Tong Zhang. Stochastic dual coordinate ascent methods for regularized loss. Journal of Machine Learning Research, 14(1):567–599, February 2013.
  • [43] Christopher J Shallue, Jaehoon Lee, Joe Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arXiv:1811.03600, 2018.
  • [44] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pages 1139–1147, 2013.
  • [45] Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of sgd for over-parameterized models and an accelerated perceptron. In Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics, 2019.
  • [46] Hoi-To Wai, Wei Shi, Angelia Nedic, and Anna Scaglione. Curvature-aided incremental aggregated gradient method. arXiv preprint arXiv:1710.08936, 2017.
  • [47] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229–256, 1992.
  • [48] Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nathan Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. arXiv preprint arXiv:1705.08292, 2017.
  • [49] Jian Zhang and Ioannis Mitliagkas. Yellowfin and the art of momentum tuning. In SysML, 2019.
  • [50] Lijun Zhang, Mehrdad Mahdavi, and Rong Jin. Linear convergence with condition number independent access of full gradients. In Advances in Neural Information Processing Systems, pages 980–988, 2013.

Appendix A Experimental Details

This section provides additional information regarding the experiments included in the main text.

For each experimental setting we strive to follow the reproducibility checklist, and provide:

  • a description and citation of the dataset,

  • a description of pre-processing steps,

  • training / validation / testing splits,

  • a description of the hyper-parameter search process and chosen values for each method,

  • the exact number of evaluation runs,

  • a clear definition of statistics reported, and

  • a description of the computing infrastructure.

A.1 CIFAR10 image classification

Dataset

The CIFAR10 dataset [20] consists 50k training and 10k testing images, partitioned over 10 classes. We download and pre-process the images using the TensorFlow models package, available at the following URL: https://github.com/tensorflow/models

Model

We use a residual convolutional network [12] with 56 layers as defined in the models package. Specifically, we use the second version whose blocks are built as a batch normalization, then a ReLU activation, and then a convolutional layer. [13]

Hyper-parameters

We use the exact setup from https://github.com/tensorflow/models/officials/resnet. As such, training is carried out with minibatches of 128 examples for 182 epochs and the training data is augmented with random crops and horizontal flips. Also note this setup multiplies the step size by the size of the minibatch. One deviation from the setup is our use of a linearly decaying learning rate instead of an explicit schedule. The linearly decaying learning rate schedule is simple and was shown to perform well [43]. This schedule is specified using two parameters: the decay rate, a multiplier specifying the final step size (0.1 or 0.01), and the decay step, specifying the step at which the fully decayed rate is reached (always set to 90% of the training steps). To factor in ease of tuning[48] we used Adam’s default parameter values and a value of 0.9 for HB’s parameter. We used IGT with the exponential Anytime Tail Averaging approach. For the tail fraction, we tried two values: the number of epochs and a tenth of that number (180 and 18). We ran using the following learning rate: (1e0, 3e-1, 1e-1, 3e-2, 1e-2) for SGD, HB and the IGT variants and (1e-2, 3e-3, 1e-3, 3e-4, 1e-4) for Adam. We ran a grid search over the base learning rate and its decay rate with a single run per combination. For each optimizer we selected the hyperparameter combination that is fastest to reach a consistently attainable target train loss of 0.2 [43]. Note that selecting the hyperparameter combination reaching the lowest training loss yields qualitatively identical curves.

The resulting hyper-parameters are:

  • SGD stepsize 0.3, decay 0.01

  • HB stepsize 0.03, decay 0.01

  • Adam stepsize 0.001, decay 0.01

  • ITA stepsize 0.3, decay 0.01, tail fraction 18

  • HB-ITA stepsize 0.03, decay 0.1, tail fraction 18

Infrastructure and Runs

The experiments were run using P100 GPUs (single GPU).

Additional Results

We provide all learning curves for the methods comparison presented in the main manuscript in figure 6.

Figure 6: Convergence and accuracy curves along training for the CIFAR10 experiments comparing baseline methods to ours. Left: Training. Right: Testing.
Ablation study: importance of IGT correction

We confirm the importance of the implicit gradient transport correction by running an experiment in which an increasing momentum is used without transport. The results appear in figure 7.

The resulting hyper-parameters are:

  • ATA stepsize 0.3, decay 0.01, tail fraction 18

  • HB-ATA stepsize 0.03, decay 0.01, tail fraction 18

Figure 7: Convergence and accuracy curves along training for the CIFAR10 experiments comparing the use of ATA combined with our proposed implicit transport mechanism. Left: Training. Right: Testing.
Effect of the batch size

We look into the effect of the batch size. To do so, we plot the number of steps required to reach a reliably attainable training loss of 0.4 as a function of the batch size. We ran using the following mini-batch sizes: 32, 128, 512 and 2048. Running with larger minibatches led to out of memory errors on our single GPU setup. The results presented in figure 8 show the benefit of IGT lowers as the batch size increases. Note that Adam’s ability to keep benefiting from larger batch sizes is consistent with previous observations.

Figure 8: Effect of mini-batch size on the number of steps to reach a target training loss.

A.2 ImageNet image classification

Dataset

We use the 2015 edition of the ImageNet Large-Scale Visual Recognition Challenge (ImageNet) [36] dataset. This dataset consists of 1.2M images partitioned into 1’000 classes. We use the pre-processing and loading utilities of the TensorFlow models package, available at the following URL: https://github.com/tensorflow/models

Model

Our model is again a large residual network, consisting of 50 layers. Similar to our CIFAR10 experiments above, we use the implementation described in [13].

Hyper-parameters

We used the same setup and approach as for the CIFAR-10 experiments. The setup trains for 90 epochs using mini-batches of 1024 examples. We used a larger grid for the learning rate schedule: decay (0.1, 0.01, 0.001) and decay step fraction (0.7, 0.8, 0.9).

The resulting hyper-parameters are:

  • SGD stepsize 0.3, decay 0.01, decay step 0.8

  • HB stepsize 0.03, decay 0.001, decay step 0.9

  • Adam stepsize 0.0001, decay 0.01, decay step 0.9

  • ITA stepsize 0.3, decay 0.01, tail fraction 90, decay step 0.9

  • HB-ITA stepsize 0.03, decay 0.01, tail fraction 90, decay step 0.9

Infrastructure and Runs

We ran these experiments using Google TPUv2.

Additional Results

We include error and accuracy curves for training and testing in Figure 9.

Figure 9: Convergence and accuracy curves along training for the ImageNet experiments. Left: Training. Right: Testing.

A.3 IMDb sentiment analysis

Dataset

The Internet Movie Database (IMDb) [27] consists of 25’000 training and 25’000 test movie reviews. The objective is binary sentiment classification based on the review’s text. We randomly split the training set in two folds of 17’536 and 7’552 reviews, the former being used for training and the latter for testing. The data is downloaded, splitted, and pre-processed with torchtext package, available at the following URL: https://github.com/pytorch/text More specifically, we tokenize the text at the word-level using the spaCy package, and embed the tokens using the 100-dimensional GloVe 6B [31] distributed representations.

Model

The model consists of an embedding lookup-table, followed by a bi-directional LSTM with dropout, and then by a fully-connected layer. The LSTM uses 256 hidden units and the dropout rate is set to 0.5. The whole model consists of 3.2M trainable parameters, with the embedding lookup-table initilized with the GloVe vectors. The model is trained to minimize the BCEWithLogitsLoss with a mini-batch size of 64.

Hyper-parameters

For each method, we used a grid-search to find the stepsize minimizing validation error after 15 epochs. The grid starts at 0.00025 and doubles until reaching 0.1024, so as to ensure that no chosen value lies on its boundaries. When applicable, the momentum factor is jointly optimized over values 0.1 to 0.95. The final hyper-parameters are displayed in the following table for each method.

Table 1: Hyperparameters for IMDb experiments.
HB Adam ASGD HB-IGT HB-ITA
α 0.032 0.0005 0.064 0.128 0.064
μ 0.95 0.95 n/a 0.9 0.9
ξ n/a n/a 100 n/a n/a
κ n/a n/a 105 n/a n/a
Infrastructure and Runs

All IMDb experiments use a single NVIDIA GTX 1080, with PyTorch v0.3.1.post2, CUDA 8.0, and cuDNN v7.0.5. We run each final configurations with 5 different random seeds and always report the mean tendency ± one standard deviation. Each run lasts approximately three hours and thirty minutes.

Additional Results

In addition to the results reported in the main text, we include training, validation, and testing curves for each method in Figure 10. Shading indicates the one standard deviation interval. Note that our focus is explicitly on optimization: in the specific case of IMDb, training for 200 epochs is completely unnecessary from a generalization standpoint as performance degrades rapidly after 15-20 epochs.

Figure 10: Convergence and accuracy curves along training for the IMDb experiments. Left: Convergence. Right: Accuracy.

A.4 Linear-quadratic regulator

Setup

Our linear-quadratic regulator [21] implements the following equations. The cost functional is evaluated at every timestep h and is given by

C(sh,ah)=shQsh+ahRah, (2)

for random symmetric positive definite matrices Q20×20 and R12×12 each with condition number 3. The initial state s0𝒩(0,3I20) is sampled around the origin, and the subsequent states evolve according to

sh+1=Ash+Bah, (3)

where entries of A20×20,B20×12 are independently sampled from a Normal distribution and then scaled such that the matrix has unit Frobenius norm. The actions are given by the linear stochastic policy ah=Ksh+ϵha, where ϵha𝒩(0,I) and K are the parameters to be optimized.

Gradient methods in this manuscript optimize the sum of costs using the REINFORCE estimate [47] given by

K𝔼h10C(sh,ah)=𝔼(h10KlogπK(ah|sh))(h10C(sh,ah)). (4)

In our experiments, the above expectation is approximated by the average of 600 trajectory rollouts. Due to the noisy dynamics of the system, it is possible for the gradient norm to explode leading to numerical unstabilities – especially when using larger stepsizes. To remedy this issue, we simply discard such problematic trajectories from the gradient estimator.

For each training iteration, we first gather 600 trajectories used for learning and then 600 more used to report evaluation metrics.

Hyper-parameters

Due to the simplicity of the considered methods, the only hyper-parameter is the stepsize. For each method, we choose the stepsize from a logarithmically-spaced grid so as to minimize the evaluation cost after 600 iterations on a single seed. Incidentally, the optimal stepsize for GD, SGD, and ITA is 0.0002.

Infrastructure and Runs

We use an Intel Core i7-5820K CPU to run the LQR experiments. All methods are implemented using numpy v1.15.4. We present results averaged over 3 random seeds, and also report the standard deviation. For stochastic gradient methods (SGD, ITA) training for 20K iterations takes about 3h, for full-gradient method (GD) about 10h, and computing the solution of the Riccati equation takes less than 5 seconds.

Additional Results

In addition to the evaluation cost reported in the main text, we also include the cost witnessed during training (and used for optimization) in Figure 11.

Figure 11: LQR costs along training iterations. Left: Costs used for learning. Right: Costs used for evaluation.

We notice that the training cost curve of ITA is not as smooth as the evaluation one. Similarly, the observed learning costs never reach as good a minima as the evaluation ones. This phenomena is easily clarified: during learning, ITA esimates the gradient using the shifted parameters Kt+γt1-γt(Kt-Kt-1) as opposed to the true parameters Kt. Those shifted parameters are not subject to a reduced variance, hence explaining the observed noise in the cost as well as deteriorated convergence.

A.5 Model-agnostic meta-learning

Dataset

We use the Mini-Imagenet dataset [34] in our model-agnostic meta-learning (MAML) [9] experiments. This dataset comprises 64 training, 12 validation, and 24 test classes. For each of train, validation, and test sets, tasks are constructed by sampling 5 classes from their respective split, and further sampling 5 images per class. Images are downsampled to 84x84x3 tensors of RGB values. For more details, please refer to the official code repository – which we carefully replicated – at the following URL: https://github.com/cbfinn/maml

Our implementation departs in two ways from the reference. First, we train our models for 100k iterations as opposed to 60k and only use 5 image samples to compute a meta-gradient whereas the reference implementation uses 15. Second, we only use 5 adaptation steps at evaluation time, when the reference uses 10.

Model

The model closely replicates the convolutional neural network of MAML [9]. It consists of 4 layers, each with 32 3x3 kernels, followed by batch normalization and ReLU activations. For specific implementation details, we refer the reader to the above reference implementation.

Hyper-parameters

We only tune the meta-stepsize for the MAML experiment. We set the momentum constant to 0.9, the adaptation-stepsize to 0.01, and average the meta-gradient of 4 tasks per iterations. Due to the reduced variance in the gradients, we found it necessary to increase the ϵ of Adam-ITA to 0.01.

For each method, we search over stepsize values on a logarithmically-spaced grid and select those values that maximize validation accuracy after 10k meta-iterations. These values are reported in Table 2.

HB Adam HB-ITA Adam-ITA
α 0.008 0.001 0.008 0.0005
Table 2: Stepsizes for MAML experiments.
Infrastructure and Runs

Each MAML experiment is run on a single NVIDIA GTX TITAN X, with PyTorch v1.1.0, CUDA 8.0, and cuDNN v7.0.5. We run each configuration with 3 different random seeds and report the mean tendency ± one standard deviation. Each run takes approximately 36 hours, and we evaluate the validation and testing accuracy every 100 iteration.

Additional Results

We complete the MAML validation curves from the main manuscript with training and testing accuracy curves in Figure 12. Moreover, we recall the final test accuracies for each method: Adam-ITA reaches 65.16%, HB-ITA 64.57%, Adam 63.70%, and HB 63.08%.

Figure 12: Training, validation, and testing accuracies for the MAML experiments along training. Shading indicates the 1 standard deviation interval. Left: Training. Center: Validation. Right: Testing.

Appendix B Additional Experiments

This section presents additional experiments to the ones included in the main text.

B.1 Baselines comparisons

While experiments in Section 5 highlighted properties of IGT and HB-IGT when the assumption of identical, constant Hessians was verified, we now turn to more realistic scenarios where individual functions are neither quadratic nor have the same Hessian to compare our proposed methods against popular baselines for the online stochastic optimization setting. We target optimization benchmarks and focus on training loss minimization.

Figure 13: Training loss curves for different optimization algorithms on several popular benchmarks. For each method, the hyper-parameters are tuned to minimize the training error after 15 epochs. Algorithms using the IGT gradient estimates tend to outperform their stochastic gradient counter-parts. Left: Logistic regression on MNIST. Center: LeNet5 on MNIST. Right: MobileNetv2 on CIFAR10.

We investigate three different scenarios: (a) linear-mnist: a logistic regression model on MNIST, (b) mnist: a modified version of LeNet5 [25] on MNIST and (c) cifar-small: the MobileNetv2 architecture [37] consisting of 19 convolutional layers on CIFAR10. All models are trained with a mini-batch size of 64, while the remaining hyper-parameters are available in Tables 45, and 6.

For each of the above tasks, models are trained for 200 epochs. We compare the following methods:

  • HB: the heavy ball method [33],

  • Adam [18],

  • ASGD [15],

  • SVRG [17],

  • SGD-dec: stochastic gradient method with an exponential learning rate schedule and exponential constant 0.999,

  • HB-IGT: the heavy ball using the IGT as a plug-in estimator, and

  • HB-ITA: same as HB-IGT but using the anytime tail averaging to forget the oldest gradients.

The hyperparameters of each method, and in particular the stepsize, are tuned independently according to a logarithmic grid so as to minimize the mean training error after epoch 15 on one seed. We then use those parameters on 5 random seeds and report the mean and standard deviation of the performance.

Figure 13 shows the training curves for the five algorithms in the three settings.

First, we observe that, for the logistic regression, HB-IGT performs on par with HB-ITA and far better than all the other methods, even though the assumption on the Hessians is violated. When using a ConvNet, however, we see that HB-IGT is not competitive with state-of-the-art methods such as Adam or ASGD. HB-ITA, on the other hand, with its smaller reliance on the assumption, once again performs much better than all other methods. In fact, HB-ITA not only converges to a lower final train error but also has a faster initial rate.

While our focus is on optimization, we also report generalization metrics in Table 3. For each algorithm, we computed the best mean accuracy after each epoch on the test set and report this value together with its standard deviation. The importance of the Anytime Tail-Averaging mechanism is again apparent: without it, Heavyball-IGT is unable to improve for more than a few epochs on the CIFAR10 validation set, regardless of the stepsize choice. On the other hand, it is evident from those results that the solutions found by Heavyball-ITA are competitive with the ones discovered by other optimization algorithms.

Linear-MNIST MNIST CIFAR10 IMDb
Heavyball 92.52 ± 0.04 99.08 ± 0.07 91.55 ± 0.25 86.90 ± 0.67
Adam 92.57 ± 0.10 98.99 ± 0.05 89.36 ± 0.75 85.62 ± 0.63
ASGD 92.47 ± 0.08 99.15 ± 0.07 91.45 ± 0.20 87.31 ± 0.21
SVRG 92.51 ± 0.04 99.06 ± 0.08 86.84 ± 0.17 n/a
SGD-dec 92.52 ± 0.06 99.11 ± 0.06 87.53 ± 0.23 86.73 ± 0.34
Heavyball-IGT 92.47 ± 0.10 99.00 ± 0.05 12.05 ± 0.21 86.61 ± 0.23
Heavyball-ITA 92.50 ± 0.10 99.19 ± 0.02 90.37 ± 0.31 87.26 ± 0.24
Table 3: Test accuracies from the best validation epoch.
HB Adam ASGD HB-IGT HB-ITA
α 0.0128 0.0002 0.0032 0.0032 0.0016
μ 0.1 0.95 n/a 0.9 0.1
ξ n/a n/a 10 n/a n/a
κ n/a n/a 104 n/a n/a
Table 4: Hyperparameters for linear-MNIST experiments.
HB Adam ASGD HB-IGT HB-ITA
α 0.0064 0.0016 0.0128 0.0032 0.0032
μ 0.9 0.95 n/a 0.95 0.95
ξ n/a n/a 10 n/a n/a
κ n/a n/a 104 n/a n/a
Table 5: Hyperparameters for MNIST experiments.
HB Adam ASGD HB-IGT HB-ITA
α 0.0512 0.0512 0.1024 0.0128 0.0512
μ 0.95 0.9 n/a 0.9 0.1
ξ n/a n/a 100 n/a n/a
κ n/a n/a 105 n/a n/a
Table 6: Hyperparameters for MobileNetV2 on CIFAR10 experiments.

Appendix C Proofs

C.1 Transport formula

gt(θt) =tt+1gt-1(θt)+1t+1g(θt,xt)
=tt+1(gt-1(θt-1)+H(θt-θt-1))+1t+1g(θt,xt) (Quadratic f)
=tt+1gt-1(θt-1)+1t+1(g(θt,xt)+tH(θt-θt-1))
=tt+1gt-1(θt-1)+1t+1g(θt+t(θt-θt-1),xt) (Identical Hessians)
tt+1g^t-1(θt-1)+1t+1g(θt+t(θt-θt-1),xt). (g^t-1 is an approximation)

Appendix D Proof of Prop. 3.1

In this proof, we assume that g is a strongly-convex quadratic function with hessian H.

At timestep t, we have access to a stochastic gradient g(θ,xt)=g(θt)+ϵt where the ϵt are i.i.d. with variance Cσ2H.

We first prove a simple lemma:

Lemma D.1.

If v0=g(θ0)+ϵ0 and, for t>0, we have

vt =tt+1vt-1+1t+1g(θt+t(θt-θt-1))+1t+1ϵt,

then

vt =g(θt)+1t+1i=0tϵi.
Proof.

Per our assumption, this is true for t=0. Now let us prove the result by induction. Assume this is true for t-1. Then we have:

vt =tt+1vt-1+1t+1g(θt+t(θt-θt-1))+1t+1ϵt
=tt+1g(θt-1)+1t+1i=0t-1ϵi
  +1t+1g(θt+t(θt-θt-1))+1t+1ϵt (recurrence assumption)
=tt+1g(θt-1)+1t+1i=0t-1ϵi
  +g(θt)-tt+1g(θt-1)+1t+1ϵt (g is quadratic)
=g(θt)+1t+1i=0tϵi.

This concludes the proof. ∎

Lemma D.2.

Let us assume we perform the following iterative updates:

vt =tt+1vt-1+1t+1g(θt+t(θt-θt-1))+1t+1ϵt
θt+1 =θt-αvt,

starting from v0=g(θ0)+ϵ0. Then, denoting Δt=θt-θ, we have

Δt =(I-αH)tΔ0-αi=0t-1Ni,tϵi

with

Ni,0 =0
Ni,t =(I-αH)Ni,t-1+1i<t1tI.
Proof.

The result is true for t=0. We now prove the result for all t by induction. Let us assume this is true for t-1. Using Lemma D.1, we have

vt-1 =g(θt-1)+1ti=0t-1ϵi

and thus, using g(θt-1)=HΔt-1,

Δt =Δt-1-αvt-1
=Δt-1-αHΔt-1-αti=0t-1ϵi
=(I-αH)Δt-1-αti=0t-1ϵi
=(I-αH)tΔ0-αi=0t-2(I-αH)Ni,t-1ϵi-αti=0t-1ϵi (recurrence assumption)
=(I-αH)tΔ0-αi=0t-1Ni,tϵi

with

Ni,t =(I-αH)Ni,t-1+1i<t1tI.

This concludes the proof. ∎

For the following lemma, we will assume that the Hessian is diagonal and will focus on one dimension with eigenvalue h. Indeed, we know that there are no interactions between the eigenspaces and that we can analyze each of them independently [30].

Lemma D.3.

Denote rh=1-αh. We assume α1L. Then, for any i and any t, we have

Ni,t 0 (Positivity)
Ni,t =0 if ti (Zero-start)
Ni,t log(2i(1-rh)) if i<t21-rh) (Constant bound)
Ni,t max{1+rh,2log(2i(1-rh))}t(1-rh)if 21-rht. (Decreasing bound)
Proof.

The Zero-start case it is immediate from the recursion of Lemma D.2. The Positivity property of Ni,t is also immediate from the recursion since the stepsize α is such that rh=1-αh is positive.

We now turn to the Constant bound property. We have, for t>i,

Ni,t =rhNi,t-1+1t
Ni,t-1+1t.

Thus, Ni,t-Ni,t-11t. Summing these inequalities, we get a telescopic sum and, finally:

Ni,t j=i+1t1j
x=itdxx
=log(ti).

This bound is trivial in the case i=0. In that case, we keep the first term in the sum separate and get

N0,t 1+logt.

In the remainder, we shall keep the log(ti) bound for simplicity. The upper bound on the right-hand size is increasing with t and its value for t=21-rh is thus an upper bound for all smaller values of t. Replacing t with 21-rh leads to

Ni,21-rh log(21-rhi)
=log(2i(1-rh)).

This proves the third inequality.

We shall now prove the Decreasing bound by induction. This bound states that, for t large enough, each Ni,t decreases as O(1/t). Using the second and third inequalities, we have

Ni,21-rh log(2i(1-rh))21-rh21-rh
=log(2i(1-rh))1-rh221-rh
max{1+rh,2log(2i(1-rh))}21-rh(1-rh).

The maximum will help us prove the last property. Thus, for t=21-rh, we have

Ni,t max{1+rh,2log(1i(1-rh))}t(1-rh)
νit,

with νi=max{1+rh,2log(1i(1-rh))}(1-rh). The Decreasing bound is verified for t=21-rh.

We now show that if, for any t>21-rh, we have Ni,t-1νit-1, then Ni,tνit. Assume that there is such at t. Then

Ni,t =rhNi,t-1+1t
rhνit-1+1t
=rhtνi+t-1t(t-1)
=(t-1)νi+(rh-1)tνi+νi+t-1t(t-1)
=νit+(rh-1)tνi+νi+t-1t(t-1).

We now shall prove that (rh-1)tνi+νi+t-1=[(rh-1)νi+1]t+νi-1 is negative. First, we have that

(rh-1)νi+1 =1-max{1+rh,2log(1i(1-rh))}
0.

Then,

[(rh-1)νi+1]t+νi-10 tνi-1(1-rh)νi-1

since (rh-1)νi+10. Thus, the property is true for every tνi-1(1-rh)νi-1. In addition, we have

νi 1+rh1-rh
νi(1-rh) 1+rh
2νi(1-rh)-2 νi(1-rh)-1+rh
21-rh νi-1νi(1-rh)-1,

and the property is also true for every t21-rh. This concludes the proof. ∎

Finally, we can prove the Proposition 3.1:

Proof.

The expectation of Δt is immediate using Lemma D.2 and the fact that the ϵi are independent, zero-mean noises. The variance is equal to V[Δt]=α2Bi=0tNi,t2. While our analysis was only along one eigenspace of the Hessian with associated eigenvalue h, we must now sum over all dimensions. We will thus define

ν¯i =max{2-αμ,2log(1iαμ)}αμ for i>0
ν¯0 =2+2log(1αμ)αμ,

which is, for every i, the maximum νi across all dimensions. We get

V[Δt] dα2Bi=0tν¯i2t2
dα2Bi=0tν¯02t2since νiνi+1i
dα2Bν¯02t.

Since we have

E[θt-θ] =(I-αH)t(θ0-θ),

we get

E[θt-θ2] =E[θt-θ]2+V[Δt]
(θ0-θ)(I-αH)2t(θ0-θ)+dα2Bν¯02t
(1-1κ)2tθ0-θ2+dα2Bν¯02t.

This concludes the proof. ∎

Appendix E Proof of Proposition 3.2 and Proposition 3.3

In this section we list and prove all lemmas used in the proofs of Proposition 3.2 and Proposition 3.3; all lemmas are stated in the same conditions as the proposition.

We start the following proposition:

Proposition E.1.

Let f be a quadratic function with positive definite Hessian H with largest eigenvalue L and condition number κ and if the stochastic gradients satisfy g(θ,x)=g(θ)+ϵ with ϵ a random uncorrelated noise with covariance bounded by BI.

Then, Algorithm 3.3 leads to iterates θt satisfying

E[θt-θ] =(I0)At(E[θ1-θ]E[θ0-θ]) (5)

where

A =(I-αH+μI-μII0) (6)

governs the dynamics of this bias. In particular, when its spectral radius, ρ(A) is less than 1, the iterates converge linearly to θ.

In a similar fashion, the variance dynamics of Heavyball-IGT are governed by the matrix

Di =((1-αhi+μ)2+2α2hi2μ2-2μ(1-αhi+μ)21001-αhi+μ0-μ)

If the spectral radius of Di, ρ(Di), is strictly less than 1 or all i, then there exist constants t0>0 and C>0 for which

Var(θt) 2α2dBClog(t)t,𝑓𝑜𝑟t>t0

where B is a bound on the variance of noise variables ϵi.

Lemma E.2 (IGT estimator as true gradient plus noise average).

If v0=g(θ0)+ϵ0 and for t>0 we have

vt=tt+1vt-1+1t+1g(θt+t(θt-θt-1))+1t+1ϵt,

then

vt=g(θt)+1t+1i=0tϵi.

This lemma is already proved in the previous section for the IGT estimator (Lemma D.1) and is just repeated here for completeness. We will use this result in the next few lemmas.

Lemma E.3 (The IGT gradient estimator is unbiased on quadratics).

For the IGT gradient estimator, vt, corresponding to parameters θt we have

𝔼[vt]=g(𝔼θt),

where the expectation is over all gradient noise vectors ϵ0,ϵ1,,ϵt.

Proof.

The proof proceeds by induction. The base case holds as we have

𝔼[v0]=𝔼[g0+ϵ0]=g(θ0).

For the inductive case, we can write

𝔼[vt] =𝔼[tt+1vt-1+1t+1g^(θt+t(θt-θt-1))]
=𝔼[tt+1vt-1+1t+1gt+tt+1gt-tt+1gt-1+1t+1ϵt]
=tt+1𝔼[vt-1-gt-1]+𝔼[gt]+tt+1𝔼[ϵt]
=𝔼[gt]=g(𝔼[θt]).

Where, in the third equality, 𝔼[vt-1-gt-1]=0 by the inductive assumption, and the last equality because the gradient of a quadratic function is linear. ∎

Lemma E.4 (Bounding the IGT gradient variance).

Let vt be the IGT gradient estimator. Then

Var[vt]2h2Var[θt-θ]+2Bt,

where B is the variance of the homoscedastic noise ϵt.

Proof.
Var[vt] =Var[gt+1t+1i=0tϵi]
=Var[hθt]+Var[1t+1i=0tϵi]
  +2Cov[hθt,1t+1i=0tϵi]
2Var[hθt]+2Var[1t+1i=0tϵi]
=2h2Var[θt-θ]+2Bt

Now that we have these basic results on the IGT estimator, we can analyze the evolution of the bias and variance of Heavyball-IGT. We use the quadratic assumption to decouple the vector dynamics of Heavyball-IGT into independent scalar dynamics. If the Hessian, H, has eigenvalues Lh1h2hn=L/κ, then we can assume without loss of generality that H is diagonal with Hii=hi.

Lemma E.5 (Evolution of bias for scalar quadratic).

Assume that the Hessian, second derivative, is h.

Starting with v0=g(θ0)+ϵ0 and w0=0, performing the following iterative updates (Heavyball-IGT, Algorithm 3.3):

vt=tt+1vt-1+1t+1g(θ+t(θt-θt-1))+1t+1ϵt,
wt+1=μwt+αvt,θt+1=θt-wt+1

results in

Δt=AtΔ0-αi=0t-1Ni,t[ϵi0]

where Nj,0=02×2,Ni,t=ANi,t-1+1i<t1tI,

Δt=[θt-θθt-1-θ] and A=(1-αh+μ-μ10).

Proof.

The proof proceeds by induction. First notice that for t=0 the equality naturally holds. We make the inductive assumption that it holds for t-1, and start by using Lemma E.2:

Δt =AΔt-1-αti=0t-1[ϵi0]
=A(At-1Δ0-αi=0t-2Ni,t[ϵi0])-αti=0t-1[ϵi0] (Inductive assumption)
=AtΔ0-α(i=0t-2ANi,t[ϵi0]+1ti=0t-1[ϵi0])
=AtΔ0-αi=0t-1Ni,t[ϵi0] (Def. of Ni,t)

Lemma E.6 (Evolution of variance).

Let Ut=Var[θt] and Vt=Cov[θt,θt-1], where θt is the t-th iterate of Heavyball-IGT on a 1-dimensional quadratic function with curvature h. The following matrix describes the variance dynamics of Heavyball-IGT.

D=((1-αh+μ)2+2α2h2μ2-2μ(1-αh+μ)21001-αh+μ0-μ) (7)

If the spectral radius of D, ρ(D), is strictly less than 1, then there exist constants t0>0 and C>0 for which

Var(θt) 2α2BClog(t)t

, where B is a bound on the variance of the noise.

Proof.

The proof (and lemma) is similar to the proof of Lemma 9 in [49]. We start by expanding Ut+1 as follows.

Ut+1 =𝔼[(θt+1-θ¯t+1)2]
=𝔼[(θt-αvt+μ(θt-θt-1)-θ¯t+αgt-μ(θ¯t-θ¯t-1))2]
=𝔼[(θt-αgt+μ(θt-θt-1)-θ¯t+αgt
  -μ(θ¯t-θ¯t-1)+α(gt-vt))2]
=𝔼[((1-αh+μ)(θt-θ¯t)-μ(θt-1-θ¯t-1))2]
  +α2𝔼[(gt-vt)2]
𝔼[((1-αh+μ)(θt-θ¯t)-μ(θt-1-θ¯t-1))2]
  +α2(2h2𝔼[(θt-θ¯t)2]+2Bt+1)
[(1-αh+μ)2+2α2μ2]𝔼[(θt-θ¯t)2]
  -2μ(1-αh+μ)𝔼[(θt-θ¯t)(θt-1-θ¯t-1)]
  +μ2𝔼[(θt-1-θ¯t-1)2]+α22Bt+1.

Where the fourth equality is obtained since we know that the IGT gradient estimator is unbiased, i.e. 𝔼[gt-vt]=0. The first inequality stems from Lemma E.4. We similarly expand Vt.

Vt =𝔼[(θt-θ¯t)(θt-1-θ¯t-1)]
=𝔼[((1-αh+μ)(θt-1-θ¯t-1)-μ(θt-2-θt-2¯)+α(gt-vt))
  (θt-1-θ¯t-1)]
=(1-αh+μ)𝔼[(θt-1-θ¯t-1)2]
  -μ𝔼[(θt-1-θ¯t-1)(θt-2-θ¯t-2)]

From the above expressions, we obtain

(Ut+1UtVt+1) D(UtUt-1Vt)+(α22Bt+100)
2α2Bi=0tDi(1t+1-i00)
2α2B(i=0s-1Di(1t+1-i00)+i=stDi(1t+1-i00))

where an inequality of vectors implies the corresponding elementwise inequalities.

If the spectral radius of D, ρ(D) is strictly less than 1, then there exists constant C>0 such that

(100)Ti=0s-1Di(1t+1-i00) Ci=0s-11t+1-i
Cst+2-s

If the spectral radius of D, ρ(D), is strictly less than 1, then there exists constant ζ>0 and constant C′′(ζ)>0 such that, ρ(D)+ζ<1 and

(100)Ti=stDi(1t+1-i00) (100)Ti=stDi(100)
C′′i=st(ρ(D)+ζ)s
= C′′(t-s+1)(ρ(D)+ζ)s

Let ρ=ρ(D)+ζ and s=2log1/ρt. Then (ρ(D)+ζ)s=1/t2, and putting the above two bounds together,

Ut+1 2α2B(2Clog1/ρtt+2-2log1/ρt+C′′t-2log1/ρt+1t2)
2α2BClog(t+1)t+1

where the last inequality holds for t>t0 for some t0 and some constant C>0.

We can now prove Proposition E.1.

Proof of Proposition E.1.

The bias statement of the proposition follows directly from taking an expectation on the bound of Lemma B.4, and the variance statement from summing up the d different variance terms given for each scalar component by Lemma B.5. ∎

E.1 Proof of Proposition 3.2

This Proposition follows from the observation that, in the noiseless case, ϵt=0 in our model. In that case, Lemma E.3 shows that Heavyball-IGT reduces to the heavy ball, and the rest follows from the optimal tuning of the heavy ball [49].

E.2 Proof of Proposition 3.3

Proof.

Like we did in previous proofs, we can assume without loss of generality that the Hessian, H, is diagonal with elements hi. For a diagonal H, matrix A can be permuted to be block diagonal with blocks

Ai=(1-αhi+μμ10).

To prove that ρ(A)<1 it suffices to prove that ρ(Ai)<1 for all i. For the rest of the proof we will focus on the dynamics along a single eigendirection with curvature hi. The rest of this proof used D to denote Di, A to denote Ai and h to denote hi.

To make explicit the dependence of matrices A and D on hyperparameters and curvature, we write A(α,μ,h) and D(α,μ,h). Let 0<α<2/(3h) and μ0=0. Using hyperparameters (α,μ0) one gets the results for gradient descent without momentum. In particular ρ(A(α,μ0,h))=|1-αh|<1, and the spectral radius of D is ρ(D(α,μ0,h))=|(1-αh)2+2α2h2|<1.

We will argue that there exists μ>0, such that ρ(A(α,μ,h))<1, and the spectral radius of D is ρ(D(α,μ,h))<1. Then the previous lemma implies that bias converges linearly, and variance is O(log(t)/t).

To argue the existence of μ>0, we will perform eigenvalue perturbation analysis using the Bauer-Fike theorem. Note that A(α,μ,h)=A(α,μ0,h)+μΔA where

ΔA=(1-100).

Similarly, D(α,μ,h)D(α,μ0,h)+μΔD where

ΔD=(2(1-αh)0-2(1-αh)00010-1).

This last approximate inequality is a first-order approximation, in the sense that we are working with arbitrarily small, positive values of μ, and we have kept terms linear in μ but ignored higher order terms, like μ2.

We will apply the Bauer-Fike theorem to bound the eigenvalues of D(α,μ,h). Consider the eigendecomposition D(α,μ0,h)=VΛV-1. We can compute

V=(001-2αh+3α2h21-αh0111-αh101)

and

V-1=(1-αh1-2αh+3α2h2-11-2αh+3α2h20010100).

Note that because we assume α<2/(3h) we get 1-αh>0. Also, 1-2αh+3α2h2>0 regardless of the choice of hyperparameters. This means that matrices V and V-1 are singular and of finite norm. The norm of ΔD is also finite. The Bauer-Fike theorem states that, if ν is an eigenvalue of D(α,μ0,h), then there exists an eigenvalue λ of D(α,μ,h) such that

|λ-ν|VpV-1pμΔDp,

for any p-norm. Since by construction |ν|ρ(D(α,μ0,h))<1, the above means that there exists a sufficiently small, but strictly positive value of μ, such that λ<1. By repeating this argument for all pairs of eigenvalues, we get the stated result. The same argument can be repeated to prove the existence of a strictly positive μ such that ρ(A(α,μ,h))<1.