The Road Less Scheduled

Aaron Defazio; Xingyu Alice Yang; Harsh Mehta; Konstantin Mishchenko; Ahmed Khaled; Ashok Cutkosky

Source-linked AI summary

The Road Less Scheduled

Aaron Defazio, Xingyu Alice Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky

arXiv:2405.15682v4May 24, 2024cs.LGcs.AImath.OCstat.ML

TL;DR

Learning-rate schedules can require advance knowledge of the stopping time T, motivating an alternative that removes schedules while retaining theoretical guarantees. Schedule-Free learning links iterate averaging with learning-rate sequences and achieves strong performance across convex and large-scale deep-learning problems, including a winning MLCommons AlgoPerf submission.

Problem
Learning-rate schedules can require setting the optimization stopping time T in advance, while existing optimization theory leaves gaps between prescribed methods and practice.
Method
Schedule-Free learning replaces schedules with an averaging approach derived from a new online-to-batch conversion theory, without adding hyperparameters over momentum-based optimizers.
Results
Schedule-Free methods match or outperform heavily tuned cosine schedules across 28 problems, and Schedule-Free AdamW won the MLCommons 2024 AlgoPerf Self-Tuning track.
Takeaways & Limitations
Schedule-Free optimization avoids specifying T while tracking the loss–training-time Pareto frontier during a single run.
Takeaways & Limitations
The method requires sweeping learning rate and weight decay, and extra code is needed for models using BatchNorm.

Abstract

from arXiv · show

Existing learning rate schedules that do not require specification of the optimization stopping step T are greatly out-performed by learning rate schedules that depend on 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 at 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 paper addresses gaps between optimization theory and practice by replacing stopping-time-dependent schedules with Schedule-Free averaging methods that retain theoretical guarantees and perform competitively across diverse problems.

Theory–practice gap: Classical theory recommends returning the Polyak-Ruppert average for stochastic gradient descent, although practitioners commonly use the last iterate.This mismatch contributes to a prominent theory–practice gap in optimization.
Schedules and stopping time: Learning-rate schedules can match worst-case optimal rates without averaging, but they require the optimization stopping time T to be set in advance.The stopping-time requirement is the central practical disadvantage identified for schedules.
Schedule-Free approach: Schedule-Free methods use a new link between iterate averaging and learning-rate sequences to preserve PR averaging’s worst-case convergence rate while matching or exceeding schedule-based performance.The approach is intended to answer whether averaging can achieve schedule-like empirical performance without sacrificing theory.
Practical behavior: Schedule-Free optimization does not require knowing T in advance, tracks the loss–training-time Pareto frontier in one run, and adds no hyperparameters over momentum-based SGD or Adam.Figure 1 reports this behavior for Schedule-Free SGD and AdamW against cosine schedules of varying lengths.
Theoretical properties: The method replaces traditional momentum with an alternative form that is worst-case optimal for every momentum parameter in the convex Lipschitz setting.Traditional momentum does not share this property.
Empirical evaluation: The evaluation covers 28 problems from logistic regression to large-scale deep learning, with Schedule-Free methods matching or outperforming heavily tuned cosine schedules.The paper presents this as one of the largest and most diverse optimizer evaluations reported in the literature.
Empirical evaluation: Schedule-Free AdamW won the MLCommons 2024 AlgoPerf Algorithmic Efficiency Challenge Self-Tuning track under limited hyperparameter tuning.The result is presented as independent verification of state-of-the-art performance against other optimization algorithms.

2 Method

Schedule-Free learning replaces predefined schedules with an averaging approach that retains theoretical guarantees while matching or exceeding schedule-based performance. Its theory connects online optimization, iterate averaging, and learning-rate choices, including larger practical step sizes.

Schedule-Free SGD: Schedule-Free SGD introduces an averaging method whose x sequence averages the base z sequence while y determines gradient-evaluation locations.The z sequence receives the base optimizer update; x is the current parameter estimate, and y is the gradient location.
Schedule-Free SGD: The momentum parameter β interpolates between Polyak-Ruppert averaging at β = 0 and Primal averaging at β = 1.Values near β ≈ 0.9 are reported to work well in practice.
Schedule-Free SGD: The approach preserves worst-case optimality while retaining momentum’s practical advantages without the potential worst-case slowdown of exponential-moving-average momentum.The interpolation couples the returned sequence x with gradient locations y, increasing stability while allowing the z sequence to move faster.
Empirical Behavior: Schedule-Free learning converges faster than classical averaging and often outperforms tuned schedules, with momentum enabling convergence at larger learning rates on quadratic problems.Figure 3 reports that choosing β < 1 becomes crucial for convergence when γ is large.
General Theory: Theorem 2 converts arbitrary online optimization algorithms into stochastic optimization guarantees and unifies several existing online-to-batch conversions.The result generalizes beyond SGD and includes time-varying momentum parameters β_t.
Large Learning Rates: Schedule-Free learning avoids specifying the stopping time T, whereas practical learning-rate choices often use larger steps than classical fixed-duration theory recommends.The proposed method can use large learning rates such as γ = D/G in a special case, with conditions checkable during training.

3 Related Work

The paper relates Schedule-Free learning to averaging, accelerated methods, and learning-rate schedules. Its framework uses averaging to reproduce schedule-like behavior without fixing the final horizon.

Relation to acceleration: Unlike conventional accelerated methods, Schedule-Free uses constant weighting for y but decreasing weighting for x.The paper distinguishes this design from accelerated methods for strongly convex and non-strongly convex problems.
Related averaging methods: Tail averaging methods such as SWA and LAWA supplement schedules rather than replace them, while generally requiring additional hyperparameters and memory.The paper notes they can potentially be combined with Schedule-Free learning for further gains.
Theoretical context: Prior work connects averaging and learning-rate decreases through equivalent effective learning rates, while tail averages can improve stochastic SGD rates in strongly convex settings.These results provide theoretical context for the paper’s averaging-based approach.
Training-time trade-offs: Cyclic schedules approximate several points on the training-time versus evaluation-performance frontier, whereas the proposed method provides the entire frontier.The paper also reports matching or improving upon known schedules in contrast to a reported 0.5% validation-accuracy underestimation.

4 Experiments

Experiments evaluate Schedule-Free methods across logistic regression, vision, language, recommendation, graph, speech, and imaging tasks. The approach generally matches or outperforms tuned schedules and baselines, with implementation caveats for batch normalization.

Experimental scope: The evaluation spans multiple domains and scales, from logistic regression to large language model training.The deep-learning suite includes computer vision, language, categorical data, and several standard architectures and datasets.
Deep learning problems: Schedule-Free methods outperform state-of-the-art cosine schedules on six listed deep-learning problems and linear-decay schedules on fastMRI and Criteo DLRM.They tie the best schedules on MAE fine-tuning and ImageNet ResNet-50 training.
Deep learning problems: The optimal Schedule-Free learning rates were generally larger than those of the base optimizers.The paper suggests that using larger non-divergent learning rates may contribute to faster convergence.
MLCommons AlgoPerf: Schedule-Free AdamW outperforms the AlgoPerf baseline on WMT, VIT, FASTMRI, and OGBG, ties on Conformer and Criteo, and marginally underperforms on DeepSpeech.The evaluation follows the self-tuning setup across eight problems and uses 10 seeds per method.
Convex problems: On 12 stochastic logistic-regression problems, Schedule-Free learning converges faster on 11, has higher accuracy on 6, and ties on the remainder against the linear-decay baseline.The comparison uses separate learning-rate sweeps and final train accuracy averaged over 10 seeds.
Implementation: Schedule-Free variants typically retain the base optimizer’s memory requirements, including no extra memory for SGD with momentum.The implementation stores x and z or z and y, computing the third sequence when needed.
Implementation: Models using batch normalization require extra code because evaluation buffers must match statistics from the x sequence.The paper describes updating running statistics with x before evaluation or using PreciseBN.

5 Parameter Sensitivity

The sensitivity experiments test whether Schedule-Free hyperparameters depend on the training horizon. Momentum remains broadly horizon-independent, while the baseline learning rate shows only mild dependence.

Momentum sensitivity: Schedule-Free SGD was evaluated for 200 epochs with fixed learning rate 1.5 across multiple momentum values.This tests whether momentum must be retuned when training lasts longer than the standard duration.
Learning-rate sensitivity: The baseline learning rate has a similar mild time-horizon dependency for Schedule-Free and schedule-based training.For short runs of ≤25 epochs, larger learning rates begin to perform best.

6 Conclusion

Schedule-Free learning removes the need to specify a learning rate schedule while matching or outperforming schedule-based learning. Its practical limitation is the need to sweep learning rate and weight decay, and its theory remains preliminary.

Schedule-Free learning removes the need to specify a learning rate schedule while matching or outperforming schedule-based learning.
The best learning rate and weight decay values differ from those used with a schedule, requiring hyperparameter sweeps.
The paper provides a preliminary theoretical exploration, with further theory needed to fully understand the method.

7 Contributions

The contributions include research experimentation, theoretical analysis, codebase development, and proofs of several theorem results.

The authors contributed to research experimentation, theoretical investigation, research code, and proofs of key theorems.

A Proof of Theorem 2

The proof develops an online-to-batch identity for weighted iterates and momentum-like updates. Special choices recover standard Polyak momentum, ordinary SGD, Polyak-Ruppert averaging, and a linear decay schedule.

Theorem 2 is established by deriving an identity, summing it over t, and using telescoping together with E[gt|zt] = ∇F(yt).
The weighted-iterate relations connect zt, xt, and yt and express the update through differences between these sequences.
The equivalent update has a momentum term plus an extra push in the direction of the base update, scaled by 1 −βt.
β = 1 with wt = 1 recovers standard Polyak momentum, whereas β = 0 gives ordinary SGD without momentum.
αt = 1 recovers Polyak-Ruppert averaging, while the barely allowed choice produces values that decrease linearly to zero.
When αT = 0, the final averaged and auxiliary iterates coincide, yielding a last-iterate convergence guarantee and recovering a prior schedule result.

C Generalizing Theorem 2 via Bregman Divergences

Bregman divergences generalize the main theorem from an inequality to an equality and expose negative terms that can support accelerated and strongly-convex rates. The section then connects these guarantees to optimistic online learning and regret bounds.

The generalized theorem uses Bregman divergences to tighten Theorem 2’s inequality into an equality.
Dropping the subtracted Bregman divergence terms recovers Theorem 2, while retaining them supports accelerated and faster strongly-convex rates.
The analysis tracks error terms through Bregman divergences and uses convexity identities involving gradients and divergences.
An optimistic online learning algorithm with αt = O(1/t) achieves accelerated convergence guarantees matching prior literature.
For smooth F, the method’s guarantee applies with wt = t and sufficiently small αt, and the algorithm need not know L or σ.
The proposed algorithm supplies the optimistic regret bound required for the acceleration result and is a mild variation on established methods.

G.1 Convex experiments

The convex experiments use datasets and model-specific implementations drawn from established repositories or codebases, with preprocessing and training details varying by task. Several experiments follow publicly documented configurations, while DLRM and fastMRI include implementation-specific caveats.

Datasets were drawn from the LIBSVM repository without modifications.
Image experiments used task-specific augmentation and normalization, including flips, crops, reflection padding, resizing, and fixed statistics.CIFAR-style runs used 32x32 crops with reflection padding, while ImageNet-style runs used 224x224 crops and center-cropped test images.
Experiments reused established codebases, including FairSeq, NanoGPT, fastMRI, and official repository implementations.The NanoGPT setup matched documented defaults and reproduced the stated 2.85 loss.
The DLRM implementation used dense gradients for simplicity, whereas sparse-gradient AdaGrad is a more common baseline and related AdaGrad variants were left for future work.
The fastMRI experiments used PyTorch 1.9 because training failed with PyTorch 2 or newer.

H Polyak and Primal Averaging Runs

The section presents Polyak and primal averaging experiments alongside an SVHN learning-rate sensitivity analysis. The reported setup includes method-specific tuning and a shortened NanoGPT run, with some Polyak NanoGPT runs diverging.

The experiments tune learning rate and momentum separately for each method while using the same c weighting sequence for Schedule-Free training.
Polyak averaging runs include momentum in the base optimizer because this produced the best results.
NanoGPT experiments were shortened to 200,000 steps because of computational budget considerations.
The NanoGPT Polyak averaging runs show divergence.
Figure 11 reports Polyak and primal averaging experiments, while a separate panel concerns SVHN ResNet-3-96 learning-rate sensitivity.