The Road Less Scheduled — part 2/6

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

Under classical worst-case convergence theory, the optimal choice of γ \gamma for a fixed duration training time T T is γ = D / ( G ​ T ) \gamma=D/(G\sqrt{T}) . This is the rate used in our bounds for Theorem 1 above. For any-time convergence (i.e. when stopping is allowed at any timestep), our proposed method can, in theory, be used with the standard learning rate sequence: γ t = D G ​ t . \displaystyle\gamma_{t}=\frac{D}{G\sqrt{t}}. (14) However, learning rate sequences of this form have poor practical performance ( Defazio et al.,, 2023 ) . Instead, much larger steps of the form D / G D/G give far better performance across virtually all problems in applications ( Defazio and Mishchenko,, 2023 ) — another theory-practice mismatch that is virtually undiscussed in the literature. Existing theory suggests that this step-size is too large to give 𝒪 ⁡ ( 1 / T ) \mathcal{O}(1/\sqrt{T}) convergence, however, as we show below, there is an important special case where such large step sizes also give optimal rates up to constant factors.

###### Theorem 3 .

Consider the online learning setting with bounded gradients g t g_{t} . Let z t + 1 = z t − γ ​ g t z_{t+1}=z_{t}-\gamma g_{t} . Let D = ‖ z 1 − z ∗ ‖ D=\left\|z_{1}-z_{*}\right\| for arbitrary reference point z ∗ z_{*} and define G = max t ≤ T ⁡ ‖ g t ‖ G=\max_{t\leq T}\left\|g_{t}\right\| . Suppose that the chosen step-size is γ = D / G \gamma=D/G , then if it holds that: ∑ t = 1 T ⟨ g t , z t − z 1 ⟩ ≤ D ​ ∑ t = 1 T ‖ g t ‖ 2 , \sum_{t=1}^{T}\left\langle g_{t},z_{t}-z_{1}\right\rangle\leq D\sqrt{\sum_{t=1}^{T}\left\|g_{t}\right\|^{2}}, (15) then: 1 T ​ ∑ t = 1 T ⟨ g t , z t − z ∗ ⟩ = 𝒪 ⁡ ( D T ​ ∑ t = 1 T ‖ g t ‖ 2 ) . \displaystyle\frac{1}{T}\sum_{t=1}^{T}\left\langle g_{t},z_{t}-z_{*}\right\rangle=\mathcal{O}\left(\frac{D}{T}\sqrt{\sum_{t=1}^{T}\left\|g_{t}\right\|^{2}}\right). (16)

This regret bound for SGD implies a convergence rate bound for Schedule-Free SGD by application of our online-to-batch conversion. Condition 15 can be checked during a training run (Using reference point z ∗ = x T z_{*}=x_{T} , and so D = ‖ x 1 − x T ‖ D=\left\|x_{1}-x_{T}\right\| ), and we find that it holds for every problem we consider in our experiments in Section 4.1 . More generally, the full conditions under which large learning rates can be used are not yet fully understood for stochastic problems. In the quadratic case, Bach and Moulines, (2013) established that large fixed step-sizes give optimal convergence rates, and we conjecture that the success of large learning rates may be attributed to asymptotic quadratic behavior of the learning process.

Empirically, we find that Schedule-Free momentum enables the use of larger learning rates γ > 0 \gamma>0 even in quadratic minimization problems f ⁡ ( x ) = 1 2 ​ x ⊤ ​ A ​ x − b ⊤ ​ x f(x)=\frac{1}{2}x^{\top}Ax-b^{\top}x . We generate 10 10 different such 20 20 -dimensional problems with eigenvalues drawn log-uniformly in [ 10 − 6 , 1 ] [10^{-6},1] . We plot the average minimal loss achieved as a function of the two parameters β \beta and γ \gamma in Figure 3 . We can see that when the learning rate we use is small, what value of β \beta we choose has little to no effect on the convergence of the algorithm. However, when γ \gamma is large, choosing β < 1 \beta<1 becomes crucial to achieving convergence.

### 2.3 How Schedule-Free Replaces a Schedule

Figure 4 illustrates the link between Polyak Averaging, Primal Averaging, Schedule-Free methods and the Linear decay schedule. Each of the averaging approaches have a crucial relationship to (worst-case optimal) Linear Decay schedules: their x x sequences simulate the effect of a Linear Decay schedule that ends at the current time-step t t , rather than the final time-step T T . This is why they can replace a schedule – they are implicitly applying one.

The four approaches differ primarily in how much each gradient contributes to the gradient location sequence y y where gradients are evaluated at each time-step, which greatly effects the stability and convergence of each method. Polyak averaging resembles the linear decay schedule’s behavior at the beginning, but near the end of training the y y sequence (blue) has much larger contributions from recent gradients. Primal averaging behaves in the opposite fashion – near the end of training it behaves similar to the Linear Decay schedules as gradients enter with a small weight, but at the early stages it significantly down-weights recent gradients, resulting in slow convergence. Schedule-Free (shown here with β = 0.6 \beta=0.6 ) remedies this down-side of Primal Averaging by modestly boosting the contribution of more recent gradients throughout all stages of training.

## 3 Related Work

The proposed method has a striking resemblance to Nesterov’s accelerated method ( Nesterov,, 1983 ; Nesterov,, 2013 ) for L L -smooth functions, which can be written in the AC-SA form ( Lan,, 2012 ) : y t \displaystyle y_{t} = ( 1 − c t + 1 ) ​ x t + c t + 1 ​ z t \displaystyle=(1-c_{t+1})x_{t}+c_{t+1}z_{t} (17) z t + 1 \displaystyle z_{t+1} = z t − k + 1 2 ​ L ∇ f ( y t ) \displaystyle=z_{t}-\frac{k+1}{2L}\nabla f(y_{t}) (18) 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}, (19) where c t + 1 = 2 / ( t + 2 ) c_{t+1}=2/(t+2) . The averaging constant, and more generally c t + 1 = r + 1 t + r + 1 , c_{t+1}=\frac{r+1}{t+r+1}, (20) for any real r > − 1 r>-1 is equivalent to the weighted average ( Shamir and Zhang,, 2013 ; Defazio and Gower,, 2021 ) x t ∝ ∑ t = 1 T t r ¯ ​ z t , x_{t}\propto\sum_{t=1}^{T}t^{\bar{r}}z_{t}, where t r ¯ t^{\bar{r}} represents the r r th factorial power of t t . Our framework is compatible with factorial power averages without sacrificing theoretical guarantees.

Our approach differs from conventional accelerated methods by using a different weight for the y t y_{t} and x t x_{t} interpolations. We use a constant weight for y t y_{t} and a decreasing weight for x t x_{t} . Accelerated methods for strongly-convex problems use a constant weight for both, and those for non-strongly convex use an decreasing weight for both, so our approach doesn’t directly correspond to either class of accelerated method. Accelerated methods also use a much larger step size for the z t z_{t} sequence than our approach.

The use of equal-weighted averages is less common than the use of exponential weighting in the practical deep learning optimization literature. Exponential moving averages (EMA) of the iterate sequence are used in the popular Lookahead optimizer ( Zhang et al.,, 2019 ) . In the case of SGD, it performs i = 1 ​ … ​ k i=1\dots k inner steps: z t , i = z t , i − 1 − γ ∇ f ( z t , i − 1 ) \displaystyle z_{t,i}=z_{t,i-1}-\gamma\nabla f(z_{t,i-1}) (21) followed by an outer step: x t = x t − 1 + α ⁡ ( z t , k − x t − 1 ) . \displaystyle x_{t}=x_{t-1}+\alpha\left(z_{t,k}-x_{t-1}\right). (22) The inner optimizer then starts at z t + 1 , 0 = x t − 1 z_{t+1,0}=x_{t-1} . The Lookahead method can be seen as the EMA version of primal averaging, just as exponential weight averaging is the EMA version of Polyak-Ruppert averaging.

Tail averaging, either using an exponential moving average or an equal-weighted average, is a common ‘folk-law’ technique that often yields a practical improvement. For instance, this kind of averaging is used without citation by the influential work of Szegedy et al., (2016) : “Model evaluations are performed using a running average of the parameters computed over time.”, and by Vaswani et al., (2017) : “…averaged the last 20 checkpoints”. Tail averages are typically “Polyak-Ruppert” style averaging as the average is not used for gradient evaluations during training.

More sophisticated tail averaging approaches such as Stochastic Weight Averaging ( Izmailov et al.,, 2018 ) and LAtest Weight Averaging ( Kaddour,, 2022 ; Sanyal et al.,, 2023 ) combine averaging with large or cyclic learning rates. They are not a replacement for scheduling, instead they aim to improve final test metrics. They generally introduce additional hyper-parameters to tune, and require additional memory. It is possible to use SWA and LAWA on top of our approach, potentially giving further gains.

Sandler et al., (2023) show via a stochastic quadratic analysis framework that averaging and learning rate decreases achieve the same effective learning rate. For instance, and average of two points along the training trajectory can give almost identical results to using a learning rate two times smaller. Stochastic quadratic problems are particularly special, Bach and Moulines, (2013) have shown that Polyak averaging gives optimal 𝒪 ⁡ ( 1 / T ) \mathcal{O}(1/T) rates without the use of decreasing time-dependent step size sequences in this setting.

Within optimization theory, tail averages can be used to improve the convergence rate for stochastic non-smooth SGD in the strongly convex setting from 𝒪 ⁡ ( log ⁡ ( T ) / T ) \mathcal{O}(\log(T)/T) to 𝒪 ⁡ ( 1 / T ) \mathcal{O}(1/T) ( Rakhlin et al.,, 2012 ) , although at the expense of worse constants compared to using weighted averages of the whole sequence ( Lacoste-Julien et al.,, 2012 ) .

Portes et al., (2022) use cyclic learning rate schedules with increasing cycle periods to give a method that explores multiple points along the Pareto frontier of training time vs eval performance. Each point at the end of a cycle is an approximation to the model from a tuned schedule ending at that time. Our method gives the entire frontier, rather than just a few points along the path. In addition, our method matches or improves upon best known schedules, whereas the “… cyclic trade-off curve underestimated the standard trade-off curve by a margin of 0.5% validation accuracy” ( Portes et al.,, 2022 ) .

## 4 Experiments

To validate the effectiveness of our method, we performed a large-scale comparison across multiple domains (computer vision, language, and categorical data) and covering a range of small scale to large-scale experiments (logistic regression to large language model training). Details of the implementation of our method for SGD and Adam used in the experiments are in Section 4.4 .

### 4.1 Deep Learning Problems

For our deep learning experiments, we evaluated Schedule-Free learning on a set benchmark tasks that are commonly used in the optimization research literature: CIFAR10 A Wide ResNet (16-8) architecture ( Zagoruyko and Komodakis,, 2016 ) on the CIFAR10 image classification dataset.

A DenseNet ( Huang et al.,, 2017 ) architecture on the CIFAR-100 (100-class) classification dataset.

A deep ResNet architecture (3-96) on the Street View House Numbers (SVHN) dataset.

A standard ResNet-50 architecture ( He et al.,, 2016 ) on the ILSVRC 2012 ImageNet ( Russakovsky et al.,, 2015 ) classification dataset.

A LSTM architecture ( Wiseman and Rush,, 2016 ) on the IWSLT14 German-English translation dataset ( Cettolo et al.,, 2014 ) .

The DLRM ( Naumov et al.,, 2019 ) architecture on the Criteo Kaggle Display Advertising dataset ( Jean-Baptiste Tien,, 2014 ) .

A stacked U-Net architecture ( Sriram et al.,, 2020 ) on the fastMRI dataset ( Zbontar et al.,, 2018 ) .

Fine-tuning a pretrained Masked Autoencoder ( He et al.,, 2021 ) ViT (patch16-512d-8b) on the ILSVRC 2012 ImageNet dataset.

A 124M parameter GPT-2 ( Radford et al.,, 2019 ) style decoder-only transformer on the OpenWebText dataset ( Gokaslan and Cohen,, 2019 ) .

For each problem, both the baseline and the Schedule-Free method were tuned by sweeping both the weight decay and learning rate on a grid. We also swept β \beta over two values, 0.9 0.9 and 0.98 0.98 . Final hyper-parameters are listed in the Appendix. Schedule-Free SGD was used for CIFAR10, CIFAR100, SVHN and ImageNet, and Schedule-Free AdamW ( Loshchilov and Hutter,, 2019 ) was used for the remaining tasks. We further include a step-wise schedule as a comparison on problems where step-wise schedules are customary. Further results for Polyak and Primal averaging are in Appendix H .

Our approach shows very strong performance (Figure 5 ) out-performing existing state-of-the-art cosine schedules on CIFAR-10, CIFAR-100, SVHN, IWSLT-14 (Figure 2 ) and OpenWebText GPT-2 problems, as well as the state-of-the-art Linear Decay schedules on the fastMRI and Criteo DLRM tasks. On the remaining two problems, MAE fine-tuning and ImageNet ResNet-50 training, it ties with the existing best schedules.

In general, the optimal learning rates for the Schedule-Free variants were larger than the optimal values for the base optimizers. The ability to use larger learning rates without diverging may be a contributing factor to the faster convergence of Schedule-Free methods. The β \beta parameter works well at the default value of 0.9 0.9 for all problems except NanoGPT, where the loss started to increase rapidly when 0.9 0.9 was used (similar to the Polyak Averaging results in Appendix H ). The larger β = 0.98 \beta=0.98 value in our sweep was stable.

Schedule-Free learning works particularly well on problems that are prone to gradient norm collapse during training. As an example, when training CIFAR-10 using SGD with traditional schedule-based approaches, at the latter stages of training the training-loss goes to near-zero and the gradient norms like-wise collapse to near-zero (Figure 6 ). From this point on learning slows significantly. This happens even when weight decay is optimally tuned to give the best final test accuracy. Schedule-Free SGD in contrast does not show a collapse of either the training loss or gradient norm sequence, and continues to reliably learn. This lack-of-collapse is likely a contributing factor to the particularly good performance of Schedule-Free learning on CIFAR-10, CIFAR-100 and SVHN in our experiments.

### 4.2 MLCommons Algorithmic Efficiency benchmark

The AlgoPerf challenge ( Dahl et al.,, 2023 ) is designed to be a large-scale and comprehensive benchmark for deep learning optimization algorithms, covering major data domains and architectures. It includes Transformers, ConvNets and U-Net models across image, language, graph and speech domains, and contains 8 problems total. We evaluated Schedule-Free AdamW following the competition guidelines, comparing against NAdamW, the competition reference Algorithm, running 10 seeds of each. As this is a time-to-target competition, traditional error bars are not appropriate so we instead plot all 10 seeds separately. Note that we excluded one benchmark problem, ResNet-50 training, as neither AdamW nor NAdamW can hit the target accuracy on that task. The remaining tasks are:

WMT A Encoder-Decoder Transformer Model on the WMT17 German-to-english translation task ( Bojar et al.,, 2017 ) .

A S/16 Vision Transformer ( Dehghani et al.,, 2023 ) model on the ILSVRC 2012 ImageNet classification task ( Russakovsky et al.,, 2015 ) .

The reference U-Net architecture from the fastMRI challenge Knee MRI dataset ( Zbontar et al.,, 2018 ) .

A Conformer ( Gulati et al.,, 2020 ) Speech Recognition model on the LibriSpeech ASR dataset ( Panayotov et al.,, 2015 ) .

A Graph-Neural Network in the style of Battaglia et al., (2018) on a Molecular property prediction task from the Open Graph Benchmark ( Hu et al.,, 2020 ) suite (PubChem BioAssay data).

Clickthrough-rate prediction on the criteo 1B dataset ( Criteo,, 2022 ) using the Deep Learning Recommendation Model (DLRM) architecture.

The Deep Speech model on the LibriSpeech ASR dataset.

The self-tuning track restricts participants to provide a single set of hyper-parameters to use for all 8 problems. Given the large number of problems, this gives performance representative of a good default configuration.