The Road Less Scheduled — part 3/6
Cleaned-text code points [31637, 46508). Bibliography omitted here; source.md includes it.
Schedule-Free AdamW performs well across all considered tasks, out-performing the baseline on the WMT, VIT, FASTMRI and OGBG training, while tying on the Conformer and Criteo workloads, and marginally under-performing on the DeepSpeech workload. We attribute the performance on the Conformer and DeepSpeech tasks to their use of batch-norm - the AlgoPerf setup doesn’t easily allow us to update the BN running statistics on the x x sequence, which is necessary with our method to get the best performance (See Section 4.4 ).
### 4.3 Convex Problems
We validated the Schedule-Free learning approach on a set of standard logistic regression problems from the LibSVM repository. For each problem, and each method separately, we performed a full learning rate sweep on a power-of-two grid, and plotted mean and standard-error of the final train accuracy from 10 seeds using the best learning rate found.
Schedule-Free learning out-performs both averaging approaches and the state-of-the-art linear decay (LD) schedule baseline (Figure 8 ). It converges faster on all but 1 of 12 problems, has higher accuracy on 6 of the problems, and ties the baseline on the remaining problems. This demonstrates that the performance advantages of Schedule-Free methods are not limited to non-convex problems.
### 4.4 Implementation Concerns
The Schedule-Free variant of a method typically has the same memory requirements as the base method. For instance, Schedule-Free SGD requires no extra memory over standard SGD with momentum. Whereas SGDM tracks the current point x x and the momentum buffer m m , we can track x x and z z . The quantity y y can be computed directly from the latest values of x x and z z , and so doesn’t need to be explicitly stored. It’s also possible to instead store z z and y y , and then compute x x when needed. This low memory usage is the case for AdamW also, see Algorithm 1 .
Our efficient PyTorch implementation actually uses one buffer to always store z z and the primary parameter buffer to store either x x or y y , with the stored quantity flipping between the two for training and test/inference passes.
Our method requires extra code to handle models where batch norm is used. This is due to the fact that BatchNorm layers maintain a running_mean and running_var to track batch statistics which is calculated at y y . For model evaluation, these buffers need to be updated to match the statistics on the x x sequence. This can be done by evaluating a small number of training batches using x x right before each eval. More sophisticated approaches such as PreciseBN ( Wu and Johnson,, 2021 ) can also be used. This calculation is not needed for other normalization layers that do not use batch-statistics.
Learning rate warmup is still necessary for our method. We use a linear warmup for a fixed duration, and fuse the Adam bias-correction term into the learning rate for simplicity (this potentially impacts the effect of weight-decay during early iterations), giving a learning rate LR γ t = γ 1 − β 2 t min ( 1 , t / T warmup ) \gamma_{t}=\gamma\sqrt{1-\beta_{2}^{t}}\min(1,t/T_{\text{warmup}}) that approaches γ \gamma when the warmup and bias-correction period ends. We found that performance was greatly improved by using a weighted c t c_{t} sequence when warmup is used, weighted by the square of the γ t \gamma_{t} used during warmup: c t + 1 = γ t 2 ∑ i = 1 t γ i 2 . c_{t+1}=\frac{\gamma^{2}_{t}}{\sum^{t}_{i=1}\gamma^{2}_{i}}. (23) This sequence decreases at a 1 / t 1/t rate after the learning rate warmup. It is shifted by one from the indexing used in Theorem 2 , which is done to simplify the implementation. This sequence is motivated by Theorem 2 ’s weighting sequences, which suggest weights proportional to polynomials of the learning rate.
Weight decay for Schedule-Free methods can be computed at either the y y or z z sequences. We used decay at y y for our experiments, as this matches the interpretation of weight-decay as the use of an additional L2-regularizer term in the loss. We found that computing the regularization at y y gives significantly better performance on some problems including ImageNet and NanoGPT training.
## 5 Parameter Sensitivity
For Schedule-Free learning to be truly schedule-free , it’s important that the momentum hyper-parameter doesn’t implicitly have a dependence on the time-horizon. If tuning this parameter gave different values depending on the training duration, then the problem of setting the horizon has just been shifted to setting the momentum value. In Figure 9 we run ImageNet training with Schedule-Free SGD for a longer-then-standard 200 epochs with a variety of momentum values, with the LR fixed to 1.5. We find that the best choice of momentum ( β = 0.9 \beta=0.9 ) is the same for all durations of training.
Schedule-Free learning has a similar mild time-horizon dependency for the baseline learning rate value as schedule-based approaches. Figure 10 shows that the optimal learning rate stays the same for broad range of values, for both Schedule-Free and Schedule based training. For short duration training ( ≤ 25 \leq 25 epochs), larger LR values begin to show the best performance. Appendix I shows the sensitivity of the final test accuracy to the baseline learning rate for a selection of our test problems, in comparison to the baseline optimizer with a cosine schedule. We see that the overall sensitivity is similar to the baseline optimizer in each problem.
## 6 Conclusion
Two roads diverged in a wood, and I— I took the one less traveled by, And that has made all the difference. - Robert Frost We have presented Schedule-Free learning, an optimization approach that removes the need to specify a learning rate schedule while matching or outperforming schedule-based learning. The primary practical limitation is the need to sweep learning rate and weight decay, as the best values differ from the those used with a schedule. We provide a preliminary theoretical exploration of the method, but further theory is needed to fully understand the method.
## 7 Contributions
Aaron Defazio discovered the method, led research experimentation and proved initial versions of Theorems 1 and 3 , with experimental/theoretical contributions by Alice Yang. Alice Yang led the development of the research codebase. Ashok Cutkosky proved key results including Theorem 2 and led the theoretical investigation of the method. Ahmed Khaled developed preliminary theory for obtaining accelerated rates which was later supplanted by Theorem 2 , and investigated the utility of β \beta with large learning rates for quadratics. Additional derivations by Konstantin Mishchenko and Harsh Mehta are included in appendix sections. Discussions between Aaron Defazio, Ashok Cutkosky, Konstantin Mishchenko, Harsh Mehta, and Ahmed Khaled over the last year contributed to this scientific discovery.
## Appendix A Proof of Theorem 2
See 2
###### Proof.
Throughout this proof, we will use the notation w 1 : t = ∑ i = 1 t w i w_{1:t}=\sum_{i=1}^{t}w_{i} . The result is established by showing the following identity: w 1 : t F ( x t ) − w 1 : t − 1 F ( x t − 1 ) − w t F ( x ⋆ ) \displaystyle w_{1:t}F(x_{t})-w_{1:t-1}F(x_{t-1})-w_{t}F(x_{\star}) ≤ w t ⟨ ∇ F ( y t ) , z t − x ⋆ ⟩ . \displaystyle\leq w_{t}\langle\nabla F(y_{t}),z_{t}-x_{\star}\rangle. (24) Where here ∇ F ( y t ) \nabla F(y_{t}) indicates a subgradient of F F at y t y_{t} with 𝔼 [ g t | z t ] = ∇ F ( y t ) \mathbb{E}[g_{t}|z_{t}]=\nabla F(y_{t}) . Given the identity ( 24 ), we sum over all t t from 1 to T T . Then the LHS will telescope to obtain: w 1 : T ( F ( x T ) − F ( x ⋆ ) ) \displaystyle w_{1:T}(F(x_{T})-F(x_{\star})) ≤ ∑ t = 1 T w t ⟨ ∇ F ( y t ) , z t − x ⋆ ⟩ , \displaystyle\leq\sum_{t=1}^{T}w_{t}\langle\nabla F(y_{t}),z_{t}-x_{\star}\rangle, from which the conclusion immediately follows since 𝔼 [ g t | z t ] = ∇ F ( y t ) \mathbb{E}[g_{t}|z_{t}]=\nabla F(y_{t}) . So, let us establish ( 24 ). To do so, it will help to observe the following identities: w t z t \displaystyle w_{t}z_{t} = w 1 : t x t − w 1 : t − 1 x t − 1 \displaystyle=w_{1:t}x_{t}-w_{1:t-1}x_{t-1} w 1 : t − 1 ( x t − x t − 1 ) \displaystyle w_{1:t-1}(x_{t}-x_{t-1}) = w t ( z t − x t ) \displaystyle=w_{t}(z_{t}-x_{t}) (25) z t − y t \displaystyle z_{t}-y_{t} = β t 1 − β t ( y t − x t ) . \displaystyle=\frac{\beta_{t}}{1-\beta_{t}}(y_{t}-x_{t}). (26) Now, setting ∇ F ( x t ) \nabla F(x_{t}) to be an arbitrary subgradient of F F at x t x_{t} , we have: w 1 : t F ( x t ) − w 1 : t − 1 F ( x t − 1 ) − w t F ( x ⋆ ) \displaystyle w_{1:t}F(x_{t})-w_{1:t-1}F(x_{t-1})-w_{t}F(x_{\star}) = w 1 : t − 1 ( F ( x t ) − F ( x t − 1 ) ) + w t ( F ( x t ) − F ( x ⋆ ) ) \displaystyle=w_{1:t-1}(F(x_{t})-F(x_{t-1}))+w_{t}(F(x_{t})-F(x_{\star})) ≤ w 1 : t − 1 ⟨ ∇ F ( x t ) , x t − x t − 1 ⟩ + w t ( F ( x t ) − F ( x ⋆ ) ) \displaystyle\leq w_{1:t-1}\langle\nabla F(x_{t}),x_{t}-x_{t-1}\rangle+w_{t}(F(x_{t})-F(x_{\star})) using ( 25 ): = w t ⟨ ∇ F ( x t ) , z t − x t ⟩ + w t ( F ( x t ) − F ( x ⋆ ) ) \displaystyle=w_{t}\langle\nabla F(x_{t}),z_{t}-x_{t}\rangle+w_{t}(F(x_{t})-F(x_{\star})) = w t ⟨ ∇ F ( x t ) , z t − x t ⟩ + w t ( F ( x t ) − F ( y t ) ) + w t ( F ( y t ) − F ( x ⋆ ) ) \displaystyle=w_{t}\langle\nabla F(x_{t}),z_{t}-x_{t}\rangle+w_{t}(F(x_{t})-F(y_{t}))+w_{t}(F(y_{t})-F(x_{\star})) ≤ w t ⟨ ∇ F ( x t ) , z t − x t ⟩ + w t ⟨ ∇ F ( x t ) , x t − y t ⟩ + w t ⟨ ∇ F ( y t ) , y t − x ⋆ ⟩ \displaystyle\leq w_{t}\langle\nabla F(x_{t}),z_{t}-x_{t}\rangle+w_{t}\langle\nabla F(x_{t}),x_{t}-y_{t}\rangle+w_{t}\langle\nabla F(y_{t}),y_{t}-x_{\star}\rangle = w t ⟨ ∇ F ( x t ) − ∇ F ( y t ) , z t − y t ⟩ + w t ⟨ ∇ F ( y t ) , z t − x ⋆ ⟩ \displaystyle=w_{t}\langle\nabla F(x_{t})-\nabla F(y_{t}),z_{t}-y_{t}\rangle+w_{t}\langle\nabla F(y_{t}),z_{t}-x_{\star}\rangle using ( 26 ): = w t β t 1 − β t ⟨ ∇ F ( x t ) − ∇ F ( y t ) , y t − x t ⟩ + w t ⟨ ∇ F ( y t ) , z t − x ⋆ ⟩ \displaystyle=w_{t}\frac{\beta_{t}}{1-\beta_{t}}\langle\nabla F(x_{t})-\nabla F(y_{t}),y_{t}-x_{t}\rangle+w_{t}\langle\nabla F(y_{t}),z_{t}-x_{\star}\rangle Finally, recall that any convex function satisfies ⟨ ∇ F ( b ) − ∇ F ( a ) , a − b ⟩ ≤ 0 \langle\nabla F(b)-\nabla F(a),a-b\rangle\leq 0 for all a , b a,b . This classical fact can be established by adding the following two subgradient identities: F ( a ) ≥ F ( b ) + ⟨ ∇ F ( b ) , a − b ⟩ , \displaystyle F(a)\geq F(b)+\langle\nabla F(b),a-b\rangle, F ( b ) ≥ F ( a ) + ⟨ ∇ F ( a ) , b − a ⟩ . \displaystyle F(b)\geq F(a)+\langle\nabla F(a),b-a\rangle. Then, since β t ∈ [ 0 , 1 ] \beta_{t}\in[0,1] , we have w t β t 1 − β t ⟨ ∇ F ( x t ) − ∇ F ( y t ) , y t − x t ⟩ ≤ 0 w_{t}\frac{\beta_{t}}{1-\beta_{t}}\langle\nabla F(x_{t})-\nabla F(y_{t}),y_{t}-x_{t}\rangle\leq 0 , which establishes the desired identity ( 24 ). ∎
## Appendix B Recovering Prior Conversions, and Connections to Momentum
The following recursions provide an equivalent update to our main algorithm that casts the update in a more “momentum-like” form.
###### Theorem 4 .
Under the same assumptions and notation as Theorem 2 , set: Δ t \displaystyle\Delta_{t} = z t + 1 − z t , \displaystyle=z_{t+1}-z_{t}, m t \displaystyle m_{t} = x t + 1 − x t , \displaystyle=x_{t+1}-x_{t}, u t \displaystyle u_{t} = y t + 1 − y t . \displaystyle=y_{t+1}-y_{t}. Then: m t \displaystyle m_{t} = w t + 1 w 1 : t − 1 w t w 1 : t + 1 m t − 1 + w t + 1 w 1 : t + 1 Δ t \displaystyle=\frac{w_{t+1}w_{1:t-1}}{w_{t}w_{1:t+1}}m_{t-1}+\frac{w_{t+1}}{w_{1:t+1}}\Delta_{t} u t \displaystyle u_{t} = ( β t + ( β t − β t + 1 ) w 1 : t w t + 1 ) m t + ( 1 − β t ) Δ t . \displaystyle=\left(\beta_{t}+(\beta_{t}-\beta_{t+1})\frac{w_{1:t}}{w_{t+1}}\right)m_{t}+(1-\beta_{t})\Delta_{t}.
Here u t u_{t} is playing the role of the “update vector”, as the sequence of points y t y_{t} are where we will be evaluating gradients. The Δ t \Delta_{t} value can be interpreted as a “base update” value: for the case that the z t z_{t} sequence is specified by SGD (as in Theorem 1 ), Δ t = − η g t \Delta_{t}=-\eta g_{t} . Thus, the update can be interpreted as a momentum term m t m_{t} , plus an extra “push” in the direction of Δ t \Delta_{t} scaled by 1 − β t 1-\beta_{t} .
###### Proof.
Let’s solve for m t m_{t} in terms of previous values: m t \displaystyle m_{t} = x t + 1 − x t \displaystyle=x_{t+1}-x_{t} = w t + 1 w 1 : t + 1 ( z t + 1 − x t ) \displaystyle=\frac{w_{t+1}}{w_{1:t+1}}(z_{t+1}-x_{t}) = w t + 1 w 1 : t + 1 ( Δ t + z t − x t ) \displaystyle=\frac{w_{t+1}}{w_{1:t+1}}(\Delta_{t}+z_{t}-x_{t}) = w t + 1 w 1 : t + 1 ( Δ t + w 1 : t − 1 w t ( x t − x t − 1 ) ) \displaystyle=\frac{w_{t+1}}{w_{1:t+1}}(\Delta_{t}+\frac{w_{1:t-1}}{w_{t}}(x_{t}-x_{t-1})) = w t + 1 w 1 : t − 1 w t w 1 : t + 1 m t − 1 + w t + 1 w 1 : t + 1 Δ t . \displaystyle=\frac{w_{t+1}w_{1:t-1}}{w_{t}w_{1:t+1}}m_{t-1}+\frac{w_{t+1}}{w_{1:t+1}}\Delta_{t}.
Now let’s solve for u t u_{t} : u t \displaystyle u_{t} = β t + 1 x t + 1 + ( 1 − β t + 1 ) z t + 1 − β t x t − ( 1 − β t ) z t \displaystyle=\beta_{t+1}x_{t+1}+(1-\beta_{t+1})z_{t+1}-\beta_{t}x_{t}-(1-\beta_{t})z_{t} = β t m t + ( 1 − β t ) Δ t + ( β t − β t + 1 ) ( z t + 1 − x t + 1 ) \displaystyle=\beta_{t}m_{t}+(1-\beta_{t})\Delta_{t}+(\beta_{t}-\beta_{t+1})(z_{t+1}-x_{t+1}) = β t m t + ( 1 − β t ) Δ t + ( β t − β t + 1 ) w 1 : t w t + 1 ( x t + 1 − x t ) \displaystyle=\beta_{t}m_{t}+(1-\beta_{t})\Delta_{t}+(\beta_{t}-\beta_{t+1})\frac{w_{1:t}}{w_{t+1}}(x_{t+1}-x_{t}) = β t m t + ( 1 − β t ) Δ t + ( β t − β t + 1 ) w 1 : t w t + 1 m t \displaystyle=\beta_{t}m_{t}+(1-\beta_{t})\Delta_{t}+(\beta_{t}-\beta_{t+1})\frac{w_{1:t}}{w_{t+1}}m_{t} = ( β t + ( β t − β t + 1 ) w 1 : t w t + 1 ) m t + ( 1 − β t ) Δ t \displaystyle=\left(\beta_{t}+(\beta_{t}-\beta_{t+1})\frac{w_{1:t}}{w_{t+1}}\right)m_{t}+(1-\beta_{t})\Delta_{t} ∎
In the special case that w t = 1 w_{t}=1 for all t t , the updates simplify to: m t \displaystyle m_{t} = t − 1 t + 1 m t − 1 + 1 t + 1 Δ t \displaystyle=\frac{t-1}{t+1}m_{t-1}+\frac{1}{t+1}\Delta_{t} u t \displaystyle u_{t} = ( β t + t ( β t − β t + 1 ) ) m t + ( 1 − β t ) Δ t . \displaystyle=\left(\beta_{t}+t(\beta_{t}-\beta_{t+1})\right)m_{t}+(1-\beta_{t})\Delta_{t}. In the special case that β t = β \beta_{t}=\beta for all t t , the update for u t u_{t} simplifies to: u t = β m t + ( 1 − β ) Δ t . \displaystyle u_{t}=\beta m_{t}+(1-\beta)\Delta_{t}. From this, it is clear that if β = 1 \beta=1 and w t = 1 w_{t}=1 , then we recover the standard Polyak momentum with a time-varying momentum factor m t = t − 1 t + 1 m t − 1 + 1 t + 1 Δ t m_{t}=\frac{t-1}{t+1}m_{t-1}+\frac{1}{t+1}\Delta_{t} , while if β = 0 \beta=0 , then we have ordinary SGD without momentum.
### B.1 Recovering Linear Decay