The Road Less Scheduled — part 1/6

Cleaned-text code points [0, 15819). Bibliography omitted here; source.md includes it.

# The Road Less Scheduled

###### Abstract

Existing learning rate schedules that do not require specification of the optimization stopping step T T are greatly out-performed by learning rate schedules that depend on T T . We propose an approach that avoids the need for this stopping time by eschewing the use of schedules entirely, while exhibiting state-of-the-art performance compared to schedules across a wide family of problems ranging from convex problems to large-scale deep learning problems. Our Schedule-Free approach introduces no additional hyper-parameters over standard optimizers with momentum. Our method is a direct consequence of a new theory we develop that unifies scheduling and iterate averaging. An open source implementation of our method is available 1 1 1 https://github.com/facebookresearch/schedule_free . Schedule-Free AdamW is the core algorithm behind our winning entry to the MLCommons 2024 AlgoPerf Algorithmic Efficiency Challenge Self-Tuning track.

## 1 Introduction

The theory of optimization, as applied in machine learning, has been successful at providing precise, prescriptive results for many problems. However, even in the simplest setting of stochastic gradient descent (SGD) applied to convex Lipschitz functions, there are glaring gaps between what our current theory prescribes and the methods used in practice.

Consider the stochastic gradient descent (SGD) step with step size γ > 0 \gamma>0 , z t + 1 = z t − γ ​ g t z_{t+1}=z_{t}-\gamma g_{t} where g t g_{t} is the stochastic (sub-)gradient at time t t , computed at the point z t z_{t} (formally defined in Section 1.2 ) of a convex Lipschitz function f f . Although standard practice for many classes of problems, classical convergence theory suggests that the expected loss of this z z sequence is suboptimal , and that the Polyak-Ruppert (PR) average x x of the sequence should be returned instead ( Polyak,, 1990 ; Ruppert,, 1988 ) : z t + 1 \displaystyle z_{t+1} = z t − γ ​ g t \displaystyle=z_{t}-\gamma g_{t} (1) x t + 1 \displaystyle x_{t+1} = ( 1 − c t + 1 ) ​ x t + c t + 1 ​ z t + 1 , \displaystyle=\left(1-c_{t+1}\right)x_{t}+c_{t+1}z_{t+1}, (2) where using c t + 1 = 1 / ( t + 1 ) c_{t+1}=1/(t+1) results in x t = 1 T ​ ∑ t = 1 T z t x_{t}=\frac{1}{T}\sum_{t=1}^{T}z_{t} . Despite their theoretical optimality, PR averages give much worse results in practice than using the last-iterate of SGD (Figures 2 , 11 ) — a folk-law result in the field of optimization, and a large theory-practice gap that is often attributed to the mismatch between this simplified problem class and the complexity of problems addressed in practice.

Recently, Zamani and Glineur, (2023) and Defazio et al., (2023) showed that the exact worst-case optimal rates can be achieved via carefully chosen learning rate sequences (also known as schedules ) alone, without the use of averaging. This result suggests that schedules have, in some sense, the same role to play as PR averaging in optimization. However, schedules have a critical disadvantage: they require setting the optimization stopping time T T in advance.

Motivated by the theory-practice gap for Polyak-Ruppert averaging, we ask the following question: Do there exist iterate averaging approaches that match the empirical performance of learning rate schedules, without sacrificing theoretical guarantees? By developing a new link between averaging and learning rate sequences, we introduce a new approach to averaging that maintains the worst-case convergence rate theory of PR averaging, while matching and often exceeding the performance of schedule-based approaches – firmly answering this question in the affirmative.

### 1.1 Summary of Results

• Our approach does not require the stopping time T T to be known or set in advance. It closely tracks the Pareto frontier of loss versus training time during a single training run (Figure 1 ), while requiring no additional hyper-parameters over the base SGD (with momentum) or Adam optimizer.

• Our approach uses an alternative form of momentum that replaces traditional momentum. This form has appealing theoretical properties: it is worst case optimal for any choice of the momentum parameter in the convex Lipschitz setting , a property that does not hold for traditional momentum.

• Our key theoretical result is a new online-to-batch conversion theorem, which establishes the optimality of our method while also unifying several existing online-to-batch theorems.

• We perform, to our knowledge, one of the largest machine learning optimization algorithm evaluations to date, consisting of 28 problems, ranging from logistic regression to large-scale deep learning problems. This evaluation contains more distinct and diverse large-scale machine-learning problems than any other optimizer evaluation we are aware of in the literature. Schedule-Free methods show strong performance, matching or out-performing heavily-tuned cosine schedules.

• Schedule-Free AdamW won the MLCommons 20204 AlgoPerf Algorithmic Efficiency Challenge Self-Tuning track, providing independent verification of it’s state-of-the-art performance against other optimization algorithms when hyperparameter-tuning is limited. We provide details of our entry and plots comparing it to the competition baseline.

### 1.2 Notation

Consider the stochastic convex minimization min x ∈ ℝ d ⁡ f ⁡ ( x ) = 𝔼 ζ ​ [ f ⁡ ( x , ζ ) ] , \min_{x\in\mathbb{R}^{d}}f(x)=\mathbb{E}_{\zeta}[f(x,\zeta)], where each f ⁡ ( x , ζ ) f(x,\zeta) is Lipschitz and convex in x x , and the expectation is taken over the random variable ζ \zeta . With a slight abuse of notation, we assume we are given, at time step t t and any point y y that we choose, an arbitrary sub-gradient ∇ f ​ ( y , ζ t ) \nabla f(y,\zeta_{t}) from the sub-differential of f f .

## 2 Method

We propose the following method, which we call Schedule-Free SGD: y t \displaystyle y_{t} = ( 1 − β ) ​ z t + β ​ x t , \displaystyle=(1-\beta)z_{t}+\beta x_{t}, (3) z t + 1 \displaystyle z_{t+1} = z t − γ ∇ f ( y t , ζ t ) , \displaystyle=z_{t}-\gamma\nabla f(y_{t},\zeta_{t}), (4) x t + 1 \displaystyle x_{t+1} = ( 1 − c t + 1 ) ​ x t + c t + 1 ​ z t + 1 , \displaystyle=\left(1-c_{t+1}\right)x_{t}+c_{t+1}z_{t+1}, (5) where c t + 1 = 1 / ( t + 1 ) c_{t+1}=1/(t+1) and z 1 = x 1 z_{1}=x_{1} is the initial point. Note that with this weighting, the x x sequence is just an online equal-weighted average of the z z sequence. The y y sequence is the gradient location sequence (on which gradients are evaluated at each step) and the x x sequence is the evaluation sequence, our current best estimate of the parameters. The z z sequence is the base sequence, which is where the base optimizer’s update is performed (in this case SGD).

This method has a momentum parameter β \beta that interpolates between Polyak-Ruppert averaging ( OPEN β = 0 ) \beta=0) and Primal averaging ( β = 1 \beta=1 ). Primal averaging ( Nesterov and Shikhman,, 2015 ; Tao et al.,, 2018 ; Cutkosky,, 2019 ; Kavis et al.,, 2019 ; Sebbouh et al.,, 2021 ; Defazio and Gower,, 2021 ; Defazio and Jelassi,, 2022 ) , is an approach where the gradient is evaluated at the averaged point x x , instead of z z : z t + 1 \displaystyle z_{t+1} = z t − γ ∇ f ( x t , ζ t ) \displaystyle=z_{t}-\gamma\nabla f(x_{t},\zeta_{t}) (6) x t + 1 \displaystyle x_{t+1} = ( 1 − c t + 1 ) ​ x t + c t + 1 ​ z t + 1 , \displaystyle=\left(1-c_{t+1}\right)x_{t}+c_{t+1}z_{t+1}, (7) this approach maintains the worst-case optimality of PR averaging but is generally considered to converge too slowly to be practical (Figures 2 , 11 ). The advantage of our interpolation is that we get the best of both worlds. We can achieve the fast convergence of Polyak-Ruppert averaging (since the z z sequence moves much quicker than the x x sequence), while still keeping some coupling between the returned sequence x x and the gradient-evaluation locations y y , which increases stability. Values of β \beta similar to standard momentum values β ≈ 0.9 \beta\approx 0.9 appear to work well in practice. We will use the notation α = 1 − β \alpha=1-\beta when convenient.

In this formulation, β = 0.9 \beta=0.9 gives the practical advantages of momentum, dampening the immediate impact of large gradients, resulting in more stable training. To see this, notice that the immediate effect of the gradient g t g_{t} at step t t is to introduce ( 1 − β ) ​ g t = 0.1 ​ g t (1-\beta)g_{t}=0.1g_{t} into the iterate sequence y y . This is similar to exponential-moving-average (EMA) momentum, where also ( 1 − β ) ​ g t (1-\beta)g_{t} is added into the iterate sequence on step t t . However, here the remainder of g t g_{t} is very slowly added into y y over time, via its place in the average x x , whereas with an EMA with β = 0.9 \beta=0.9 , the majority of the gradient is incorporated within the next 10 steps.

So from this viewpoint, the Schedule-Free updates can be seen as a version of momentum that has the same immediate effect, but with a greater delay for adding in the remainder of the gradient. This form of momentum (by interpolation) also has a striking advantage: it does not result in any theoretical slowdown; it gives the optimal worst case ( Nesterov,, 2013 ) convergence for the non-smooth convex setting (including constants), for any choice of momentum β \beta between 0 and 1 inclusive:

###### Theorem 1 .

Suppose F F is a convex function, and ζ 1 , … , ζ T \zeta_{1},\dots,\zeta_{T} is an i.i.d. sequence of random variables such that F = 𝔼 ⁡ [ f ⁡ ( x , ζ ) ] F=\mathbb{E}[f(x,\zeta)] for some function f f that is G G -Lipschitz in x x . For any minimizer x ⋆ x_{\star} , define D = ‖ x 1 − x ⋆ ‖ D=\left\|x_{1}-x_{\star}\right\| and γ = D / ( G ​ T ) \gamma=D/(G\sqrt{T}) . Then for any β ∈ [ 0 , 1 ] \beta\in[0,1] , Schedule-Free SGD ensures: 𝔼 ⁡ [ F ⁡ ( x T ) − F ⁡ ( x ⋆ ) ] ≤ D ​ G T \displaystyle\mathbb{E}[F(x_{T})-F(x_{\star})]\leq\frac{DG}{\sqrt{T}} (8)

In contrast, exponential-moving-average momentum in the non-smooth setting actually hurts the theoretical worst-case convergence rate. The Schedule-Free approach maintains the advantages of momentum ( Sutskever et al.,, 2013 ) without the potential worst-case slow-down.

### 2.1 General Theory

The method analyzed in Theorem 1 is actually a special-case of a more general result that incorporates arbitrary online optimization algorithms rather than only SGD, as well as arbitrary time-varying sequences of β t \beta_{t} . The proof is provided in Appendix A .

###### Theorem 2 .

Let F F be a convex function. Let ζ 1 , … , ζ T \zeta_{1},\dots,\zeta_{T} be an iid sequence such that F ⁡ ( x ) = 𝔼 ζ ​ [ f ⁡ ( x , ζ ) ] F(x)=\mathbb{E}_{\zeta}[f(x,\zeta)] . Let z 1 , … , z T z_{1},\dots,z_{T} be arbitrary vectors and let w 1 , … , w T w_{1},\dots,w_{T} and β 1 , … , β T \beta_{1},\dots,\beta_{T} be arbitrary numbers in [ 0 , 1 ] [0,1] such that z t z_{t} , w t w_{t} and β t \beta_{t} are independent of ζ t , … , ζ T \zeta_{t},\dots,\zeta_{T} . Set: x t \displaystyle x_{t} = ∑ i = 1 t w i ​ z i ∑ i = 1 t w i = x t − 1 ​ ( 1 − w t ∑ i = 1 t w i ) ⏟ ≜ 1 − c t + w t ∑ i = 1 t w i ⏟ ≜ c t ​ z t \displaystyle=\frac{\sum_{i=1}^{t}w_{i}z_{i}}{\sum_{i=1}^{t}w_{i}}=x_{t-1}\underbrace{\left(1-\frac{w_{t}}{\sum_{i=1}^{t}w_{i}}\right)}_{\triangleq 1-c_{t}}+\underbrace{\frac{w_{t}}{\sum_{i=1}^{t}w_{i}}}_{\triangleq c_{t}}z_{t} (9) y t \displaystyle y_{t} = β t ​ x t + ( 1 − β t ) ​ z t \displaystyle=\beta_{t}x_{t}+(1-\beta_{t})z_{t} (10) g t \displaystyle g_{t} = ∇ f ​ ( y t , ζ t ) . \displaystyle=\nabla f(y_{t},\zeta_{t}). (11) Then we have for all x ⋆ x_{\star} : 𝔼 ⁡ [ F ⁡ ( x T ) − F ⁡ ( x ⋆ ) ] ≤ 𝔼 ⁡ [ ∑ t = 1 T w t ​ ⟨ g t , z t − x ⋆ ⟩ ] ∑ i = 1 T w i . \displaystyle\mathbb{E}[F(x_{T})-F(x_{\star})]\leq\frac{\mathbb{E}[\sum_{t=1}^{T}w_{t}\langle g_{t},z_{t}-x_{\star}\rangle]}{\sum_{i=1}^{T}w_{i}}. (12)

To recover Theorem 1 from the above result, notice that the algorithm analyzed by Theorem 1 is captured by Theorem 2 with w t = 1 w_{t}=1 , β t \beta_{t} a constant β \beta and z t + 1 = z t − γ ​ g t z_{t+1}=z_{t}-\gamma g_{t} for all t t . Next, observe that the sequence z 1 , … , z T z_{1},\dots,z_{T} is performing online gradient descent ( Zinkevich,, 2003 ) , for which it is well-known that the regret ∑ t = 1 T ⟨ g t , z t − x ⋆ ⟩ \sum_{t=1}^{T}\langle g_{t},z_{t}-x_{\star}\rangle (appearing in the numerator of our result) is bounded by D ​ G ​ T DG\sqrt{T} and so the result of Theorem 1 immediately follows.

The regret is the principle object of study in online convex optimization ( Hazan,, 2022 ; Orabona,, 2019 ) . Viewed in this light, Theorem 2 provides a way to convert an online convex optimization algorithm into a stochastic optimization algorithm: it is a form of online-to-batch conversion ( Cesa-Bianchi et al.,, 2004 ) . Classical online-to-batch conversions are a standard technique for obtaining convergence bounds for many stochastic optimization algorithms, including stochastic gradient descent ( Zinkevich,, 2003 ) , AdaGrad ( Duchi et al.,, 2011 ) , AMSGrad ( Reddi et al.,, 2018 ) , and Adam ( Kingma and Ba,, 2014 ) . All of these algorithms can be analyzed as online convex optimization algorithms: they provide bounds on the regret ∑ t = 1 T ⟨ g t , z t − x ⋆ ⟩ \sum_{t=1}^{T}\langle g_{t},z_{t}-x_{\star}\rangle rather than direct convergence guarantees. It is then necessary (although sometimes left unstated) to convert these regret bounds into stochastic convergence guarantees via an online-to-batch conversion. Our result provides a more versatile method for effecting this conversion.

Theorem 2 actually provides a “grand unification” of a number of different online-to-batch conversions that have been proposed over the years. Most of these conversion methods were first developed specifically to provide convergence analysis for SGD (or some variant such as dual averaging or mirror descent), and then generalized into techniques that apply to any online convex optimization algorithm. For example, the classical Polyak averaging method can be generalized to form the “standard” online-to-batch conversion of Cesa-Bianchi et al., (2004) , and is immediately recovered from Theorem 2 by setting w t = 1 w_{t}=1 and β t = 0 \beta_{t}=0 for all t t . More recently Nesterov and Shikhman, (2015) ; Tao et al., (2018) derived an alternative to Polyak averaging that was later generalized to work with arbitrarily online convex optimization algorithms by Cutkosky, (2019) ; Kavis et al., (2019) , and then observed to actually be equivalent to the heavy-ball momentum by Defazio, (2020) ; Defazio and Gower, (2021) ; Defazio and Jelassi, (2022) . This method is recovered by our Theorem 2 by setting w t = 1 w_{t}=1 and β t = 1 \beta_{t}=1 for all t t . Finally, very recently Zamani and Glineur, (2023) discovered that gradient descent with a linear decay stepsize provides a last-iterate convergence guarantee, which was again generalized to an online-to-batch conversion by Defazio et al., (2023) . This final result is also recovered by Theorem 2 by setting w t = 1 w_{t}=1 and β t = t T \beta_{t}=\frac{t}{T} (see Appendix B ).

In Appendix C , we give a further tightening of Theorem 2 – it can be improved to an equality by precisely tracking additional terms that appear on the right-hand-side. This tightened version can be used to show convergence rate results for smooth losses, both with and without strong-convexity. As an example application, we show that schedule-free optimistic-gradient methods ( Rakhlin and Sridharan,, 2013 ) converge with accelerated rates: 𝔼 ⁡ [ F ⁡ ( x T ) − F ⁡ ( x ⋆ ) ] = O ⁡ ( D 2 ​ L T 2 + D ​ σ T ) . \displaystyle\mathbb{E}[F(x_{T})-F(x_{\star})]=O\left(\frac{D^{2}L}{T^{2}}+\frac{D\sigma}{\sqrt{T}}\right). (13)

### 2.2 On Large Learning Rates