Title: The Road Less Scheduled
Venue: NeurIPS 2024
Accepted: yes
Keywords: Stochastic Optimization

Optimization

Convex Optimization

Learning Rates

Learning Rate Schedules
Abstract: 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.

The Road Less Scheduled

Abstract

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 1 . 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 , z t +1 = z t -γg t where g t is the stochastic (sub-)gradient at time t , computed at the point z t (formally defined in Section 1.1) of a convex Lipschitz function f . Although standard practice for many classes of problems, classical convergence theory suggests that the expected loss of this z sequence is suboptimal , and that the Polyak-Ruppert (PR) average x of the sequence should be returned instead (Polyak, 1990; Ruppert, 1988):

where using c t +1 = 1 / ( t +1) results in x t = 1 T ∑ T t =1 z t . Despite their theoretical optimality, PR averages give much worse results in practice than using the last-iterate of SGD (Figures 2a,8) - 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.

0

20

40

60

80

100

Epoch

20

40

60

Test Accuracy (%)

ILSVRC 2012 ImageNet (ResNet-50)

0

200000

400000

600000

Step

2.8

2.9

3.0

3.1

3.2

Test Loss

OpenWebText (GPT-2 124M)

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 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.

Summary of Results

Our approach does not require the stopping time 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 largescale 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 2024 AlgoPerf Algorithmic Efficiency Challenge Self-Tuning (Adaptive Algorithm) Track, providing independent verification of its SOTA performance against other optimization algorithms in cases where hyperparametertuning is limited. We provide details of our entry and plots comparing it to the competition baseline.

1.1 Notation

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

0

10000

20000

30000

40000

50000

60000

Step

4.0

4.5

5.0

5.5

Test Loss

IWSLT14 (LSTM)

Polyak Averaging (4.35 SE 0.004)

Schedule-Free (4.18 SE 0.001)

2 Method

We propose the following method, which we call Schedule-Free SGD:

where c t +1 = 1 / ( t +1) and z 1 = x 1 is the initial point. Note that, with this weighting, the x sequence is just a running average of the z sequence. The y sequence is the gradient location sequence (on which gradients are evaluated at each step). The z sequence is the base sequence, which is where the base optimizer's update is performed (in this case SGD). The x sequence is the evaluation sequence, our best estimate of the weights so far.

This method has a momentum parameter β that interpolates between Polyak-Ruppert averaging ( β = 0) and Primal averaging ( β = 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 , instead of z :

this approach maintains the worst-case optimality of PR averaging but is generally considered to converge too slowly to be practical (Figures 2a,8). 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 sequence moves much quicker than the x sequence), while still keeping some coupling between the returned sequence x and the gradient-evaluation locations y , which increases stability. Values of β similar to standard momentum values β ≈ 0 . 9 appear to work well in practice. We will use the notation α = 1 -β when convenient.

In this formulation, β = 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 at step t is to introduce (1 -β ) g t = 0 . 1 g t into the iterate sequence y . This is similar to exponential-moving-average (EMA) momentum, where also (1 -β ) g t is added into the iterate sequence on step t . However, here the remainder of g t is very slowly added into y over time, via its place in the average x , whereas with an EMA with β = 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 β between 0 and 1 inclusive:

Theorem 1. Suppose F is a convex function, and ζ 1 , . . . , ζ T is an i.i.d. sequence of random variables such that F = E [ f ( x, ζ )] for some function f that is G -Lipschitz in x . For any minimizer x ⋆ , define

0.01

25.13

50.26

75.38

0.700

0.775

0.851

0.926

Minimal Loss for Different

and

values.

25

50

75

100

125

150

175

200

Minimal Loss

D = ∥ x 1 -x ⋆ ∥ and γ = D/ ( G √ T ) . Then for any β ∈ [0 , 1] , Schedule-Free SGD ensures:

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 . The proof is provided in Appendix A.

Theorem 2. Let F be a convex function. Let ζ 1 , . . . , ζ T be an iid sequence such that F ( x ) = E ζ [ f ( x, ζ )] . Let z 1 , . . . , z T be arbitrary vectors and let w 1 , . . . , w T and β 1 , . . . , β T be arbitrary numbers in [0 , 1] such that z t , w t and β t are independent of ζ t , . . . , ζ T . Set:

Then we have for all x ⋆ :

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 , β t a constant β and z t +1 = z t -γg t for all t . Next, observe that the sequence z 1 , . . . , z T is performing online gradient descent (Zinkevich, 2003), for which it is well-known that the regret ∑ T t =1 ⟨ g t , z t -x ⋆ ⟩ (appearing in the numerator of our result) is bounded by DG √ 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 t =1 ⟨ g t , z t -x ⋆ ⟩ 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 and β t = 0 for all 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 and β t = 1 for all 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 and β t = 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:

2.2 On Large Learning Rates

Under classical worst-case convergence theory, the optimal choice of γ for a fixed duration training time T is γ = D/ ( G √ 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:

However, learning rate sequences of this form have poor practical performance (Defazio et al., 2023). Instead, much larger steps of the form 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 O (1 / √ 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 . Let z t +1 = z t -γg t . Let D = ∥ z 1 -z ∗ ∥ for arbitrary reference point z ∗ and define G = max t ≤ T ∥ g t ∥ . Suppose that the chosen step-size is γ = D/G , then if it holds that:

This regret bound for SGD implies a convergence rate bound for Schedule-Free SGD by application of our online-to-batch conversion. Condition 31 can be checked during a training run (Using reference point z ∗ = x T , and so D = ∥ x 1 -x T ∥ ), and we find that it holds for every problem we consider in our experiments in Section 4. 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 even in quadratic minimization problems f ( x ) = 1 2 x ⊤ Ax -b ⊤ x . We generate 10 different such 20 -dimensional problems with eigenvalues drawn log-uniformly in [10 -6 , 1] . We plot the average minimal loss achieved as a function of the two parameters β and γ in Figure 2b. We can see that when the learning rate we use is small, what value of β we choose has little to no effect on the convergence of the algorithm. However, when γ is large, choosing β < 1 becomes crucial to achieving convergence.

then:

3 Related Work

The proposed method has a striking resemblance to Nesterov's accelerated method (Nesterov, 1983, 2013) for L -smooth functions, which can be written in the AC-SA form (Lan, 2012):

where c t +1 = 2 / ( t +2) . The averaging constant, and more generally

for any real r > -1 is equivalent to the weighted average (Shamir and Zhang, 2013; Defazio and Gower, 2021) x t ∝ ∑ T t =1 t ¯ r z t , where t ¯ r represents the r th factorial power of 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 and x t interpolations. We use a constant weight for y t and a decreasing weight for 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 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 inner steps:

The inner optimizer then starts at 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 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 O (log( T ) /T ) to O (1 /T ) (Rakhlin et al., 2012), followed by an outer step:

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

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.

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

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

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

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

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

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

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

NanoGPT 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 β over two values, 0 . 9 and 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 I.

Our approach shows very strong performance (Figure 3) out-performing existing state-of-the-art cosine schedules on CIFAR-10, CIFAR-100, SVHN, IWSLT-14 (Figure 2a) 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 β parameter works well at the default value of 0 . 9 for all problems except NanoGPT, where the loss started to increase rapidly when 0 . 9 was used (similar to the Polyak Averaging results in Appendix I). The larger β = 0 . 98 value in our sweep was stable.

4.1 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.

0

50

100

150

200

250

300

Epoch

75

80

85

90

95

Test Accuracy (%)

CIFAR-10 (WRN-16-8)

Schedule-Free (96.03% SE 0.04)

Step-Wise Schedule (95.59% SE 0.03)

0

50

100

150

200

250

300

Epoch

40

50

60

70

80

Test Accuracy (%)

CIFAR-100 (DenseNet)

Schedule-Free (78.71% SE 0.06)

Step-Wise Schedule (76.41% SE 0.14)

0

50

100

150

200

250

300

Epoch

94

95

96

97

98

99

Test Accuracy (%)

SVHN (ResNet-3-96)

Schedule-Free (98.40% SE 0.01)

Step-Wise Schedule (98.20% SE 0.01)

0

20

40

60

80

100

Epoch

40

50

60

70

Test Accuracy (%)

ILSVRC 2012 ImageNet (ResNet-50)

Schedule-Free (76.90% SE 0.03)

Step-Wise Schedule (76.49% SE 0.07)

0

10

20

30

40

50

Epoch

0.906

0.908

0.910

0.912

Test SSIM

fastMRI Knee (VarNet 2.0)

Schedule-Free (0.9112 SE 0.00012)

0

20

40

60

80

100

Epoch

0.780

0.785

0.790

Test Accuracy

Schedule-Free (0.7915 SE 0.00003)

0

20

40

60

80

100

Epoch

65

70

75

80

Test Accuracy (%)

MAE ImageNet Finetune (ViT)

Schedule-Free (83.54 SE 0.03)

0

200000

400000

600000

Step

2.8

2.9

3.0

3.1

3.2

Test Loss

OpenWebText (GPT-2 124M)

Schedule-Free (2.831 SE 0.008)

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.

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 sequence, which is necessary with our method to get the best performance (See Section 4.3).

4.2 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 7). 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.3 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 and the momentum buffer m , we can track x and z . The quantity y can be computed directly from the latest values of x and z , and so doesn't need

0.0

0.2

0.4

0.6

0.8

1.0

Normalized Time

0.4

0.6

0.8

1.0

Normalized Test Metric

WMT

Schedule-Free

0.0

0.2

0.4

0.6

0.8

Normalized Time

0.00

0.25

0.50

0.75

1.00

Normalized Test Metric

ViT

Schedule-Free

0.0

0.1

0.2

0.3

0.4

Normalized Time

0.990

0.995

1.000

Normalized Test Metric

fastMRI

Schedule-Free

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Normalized Time

0.97

0.98

0.99

1.00

Normalized Test Metric

Schedule-Free

0.0

0.1

0.2

0.3

0.4

Normalized Time

0.6

0.7

0.8

0.9

1.0

Normalized Test Metric

OGBG

Schedule-Free

0.0

0.2

0.4

0.6

0.8

Normalized Time

0.996

0.998

1.000

Normalized Test Metric

Criteo1TB

Schedule-Free

0.0

0.2

0.4

0.6

Normalized Time

0.98

0.99

1.00

Normalized Test Metric

Schedule-Free

to be explicitly stored. It's also possible to instead store z and y , and then compute 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 and the primary parameter buffer to store either x or 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 . For model evaluation, these buffers need to be updated to match the statistics on the x sequence. This can be done by evaluating a small number of training batches using 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 -β t 2 min(1 , t/T warmup ) that approaches γ when the warmup and bias-correction period ends. We found that performance was greatly improved by using a weighted c t sequence when

0

25

50

75

100

125

150

175

200

Epoch

60

65

70

75

Test Accuracy (%)

ILSVRC 2012 ImageNet (ResNet-50) Learning Rate Sweep

Schedule-Free LR 5.0 (76.27% SE 0.03)

Schedule-Free LR 3.0 (77.39% SE 0.05)

Schedule-Free LR 1.5 (77.83% SE 0.03)

Schedule-Free LR 1.0 (77.42% SE 0.02)

Schedule-Free LR 0.5 (76.17% SE 0.07)

0

25

50

75

100

125

150

175

200

Epoch

60

65

70

75

Test Accuracy (%)

ILSVRC 2012 ImageNet (ResNet-50) Learning Rate Sweep

SGD LR 0.5 (76.42% SE 0.05)

SGD LR 0.3 (77.05% SE 0.03)

SGD LR 0.15 (77.57% SE 0.02)

SGD LR 0.1 (77.68% SE 0.04)

SGD LR 0.05 (77.69% SE 0.04)

warmup is used, weighted by the square of the γ t used during warmup:

This sequence decreases at a 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. This sequence was used for both SGD and AdamW experiments.

Weight decay for Schedule-Free methods can be computed at either the y or z sequences. We used decay at 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 gives significantly better performance on some problems including ImageNet and NanoGPT training.

5 Parameter Sensitivity

For Schedule-Free learning to be truly schedulefree , it's important that the momentum hyperparameter 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 5 we run ImageNet training with Schedule-Free SGD for a longerthen-standard 200 epochs with a variety of momentum values, with the LR fixed to 1.5. We

0

25

50

75

100

125

150

175

200

Epoch

60

65

70

75

Test Accuracy (%)

ILSVRC 2012 ImageNet (ResNet-50) Momentum Sweep

Schedule-Free MOM 0.75 (76.60% SE 0.02)

Schedule-Free MOM 0.9 (77.78% SE 0.04)

Schedule-Free MOM 0.95 (76.99% SE 0.05)

Schedule-Free MOM 0.98 (75.37% SE 0.03)

find that the best choice of momentum ( β = 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 6 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 epochs), larger LR values begin to show the best performance. Appendix J 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, establishing its worst-case optimal performance for non-smooth Lipschitz convex optimization.

Funding Acknowledgments

AC is supported by NSF grant number CCF-2211718.

References

Bach, F. and Moulines, E. (2013). Non-strongly-convex smooth stochastic approximation with convergence rate O (1 /n ) . In Burges, C., Bottou, L., Welling, M., Ghahramani, Z., and Weinberger, K., editors, Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc.

Cesa-Bianchi, N., Conconi, A., and Gentile, C. (2004). On the generalization ability of on-line learning algorithms. IEEE Transactions on Information Theory , 50(9):2050-2057.

Cettolo, M., Niehues, J., Stüker, S., Bentivogli, L., and Federico, M. (2014). Report on the 11th IWSLT evaluation campaign. In IWSLT .

Chiang, C.-K., Yang, T., Lee, C.-J., Mahdavi, M., Lu, C.-J., Jin, R., and Zhu, S. (2012). Online optimization with gradual variations. In Conference on Learning Theory , pages 6-1. JMLR Workshop and Conference Proceedings.

Cutkosky, A. (2019). Anytime online-to-batch, optimism and acceleration. In International conference on machine learning , pages 1446-1454. PMLR.

Dahl, G. E., Schneider, F., Nado, Z., Agarwal, N., Sastry, C. S., Hennig, P., Medapati, S., Eschenhagen, R., Kasimbeg, P., Suo, D., Bae, J., Gilmer, J., Peirson, A. L., Khan, B., Anil, R., Rabbat, M., Krishnan, S., Snider, D., Amid, E., Chen, K., Maddison, C. J., Vasudev, R., Badura, M., Garg, A., and Mattson, P. (2023). Benchmarking Neural Network Training Algorithms.

Defazio, A. (2020). Momentum via primal averaging: Theoretical insights and learning rate schedules for non-convex optimization.

Defazio, A., Cutkosky, A., Mehta, H., and Mishchenko, K. (2023). When, why and how much? adaptive learning rate scheduling by refinement.

Defazio, A. and Gower, R. M. (2021). The power of factorial powers: New parameter settings for (stochastic) optimization. In Balasubramanian, V. N. and Tsang, I., editors, Proceedings of The 13th Asian Conference on Machine Learning , volume 157 of Proceedings of Machine Learning Research , pages 49-64. PMLR.

Defazio, A. and Jelassi, S. (2022). Adaptivity without compromise: A momentumized, adaptive, dual averaged gradient method for stochastic optimization. Journal of Machine Learning Research , 23:1-34.

Defazio, A. and Mishchenko, K. (2023). Learning-rate-free learning by D-adaptation. The 40th International Conference on Machine Learning (ICML 2023) .

Duchi, J., Hazan, E., and Singer, Y. (2011). Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research , 12(61).

Gokaslan, A. and Cohen, V. (2019). Openwebtext corpus. http://Skylion007.github.io/ OpenWebTextCorpus .

Hazan, E. (2022). Introduction to online convex optimization . MIT Press.

Hazan, E. and Kale, S. (2010). Extracting certainty from uncertainty: Regret bounded by variation in costs. Machine learning , 80:165-188.

He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R. (2021). Masked autoencoders are scalable vision learners. arXiv:2111.06377 .

He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition .

Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. (2017). Densely connected convolutional networks. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 2261-2269.

Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D., and Wilson, A. G. (2018). Averaging weights leads to wider optima and better generalization. In Conference on Uncertainty in Artificial Intelligence (UAI) .

Jean-Baptiste Tien, joycenv, O. C. (2014). Display advertising challenge.

Joulani, P., György, A., and Szepesvári, C. (2017). A modular analysis of adaptive (non-) convex optimization: Optimism, composite objectives, and variational bounds. In International Conference on Algorithmic Learning Theory , pages 681-720. PMLR.

Joulani, P., Raj, A., Gyorgy, A., and Szepesvári, C. (2020). A simpler approach to accelerated optimization: iterative averaging meets optimism. In International conference on machine learning , pages 4984-4993. PMLR.

Kaddour, J. (2022). Stop wasting my time! saving days of ImageNet and BERT training with latest weight averaging.

Kavis, A., Levy, K. Y., Bach, F., and Cevher, V. (2019). UniXGrad: A universal, adaptive algorithm with optimal guarantees for constrained optimization. Advances in neural information processing systems , 32.

Kingma, D. P. and Ba, J. (2014). Adam: a method for stochastic optimization. In International Conference on Learning Representations .

Lacoste-Julien, S., Schmidt, M., and Bach, F. (2012). A simpler approach to obtaining an o (1 /t ) convergence rate for the projected stochastic subgradient method.

Lan, G. (2012). An optimal method for stochastic composite optimization. Mathematical Programming , 133(1):365-397.

Loshchilov, I. and Hutter, F. (2019). Decoupled weight decay regularization. In International Conference on Learning Representations .

Naumov, M., Mudigere, D., Shi, H. M., Huang, J., Sundaraman, N., Park, J., Wang, X., Gupta, U., Wu, C., Azzolini, A. G., Dzhulgakov, D., Mallevich, A., Cherniavskii, I., Lu, Y ., Krishnamoorthi, R., Yu, A., Kondratenko, V., Pereira, S., Chen, X., Chen, W., Rao, V., Jia, B., Xiong, L., and Smelyanskiy, M. (2019). Deep learning recommendation model for personalization and recommendation systems. CoRR .

Nesterov, Y. (1983). A method for solving a convex programming problem with convergence rate O (1 /k 2 ) . Soviet Mathematics Doklady .

Nesterov, Y. (2013). Lectures on Convex Optimization . Springer Nature.

Nesterov, Y. and Shikhman, V. (2015). Quasi-monotone subgradient methods for nonsmooth convex minimization. Journal of Optimization Theory and Applications , 165(3):917-940.

Orabona, F. (2019). A modern introduction to online learning. arXiv preprint arXiv:1912.13213 .

Polyak, B. (1990). New stochastic approximation type procedures. Avtomatica i Telemekhanika , 7:98-107.

Portes, J., Blalock, D., Stephenson, C., and Frankle, J. (2022). Fast benchmarking of accuracy vs. training time with cyclic learning rates.

Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. (2019). Language models are unsupervised multitask learners. Technical report, OpenAI.

Rakhlin, A., Shamir, O., and Sridharan, K. (2012). Making gradient descent optimal for strongly convex stochastic optimization. In Proceedings of the 29th International Coference on International Conference on Machine Learning .

Rakhlin, A. and Sridharan, K. (2013). Online learning with predictable sequences. In Conference on Learning Theory , pages 993-1019. PMLR.

Reddi, S. J., Kale, S., and Kumar, S. (2018). On the convergence of Adam and beyond. In International Conference on Learning Representations .

Ruppert, D. (1988). Efficient estimations from a slowly convergent Robbins-Monro process. Technical Report, Cornell University .

Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L. (2015). ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV) , 115(3).

Sandler, M., Zhmoginov, A., Vladymyrov, M., and Miller, N. (2023). Training trajectories, mini-batch losses and the curious role of the learning rate.

Sanyal, S., Neerkaje, A., Kaddour, J., Kumar, A., and Sanghavi, S. (2023). Early weight averaging meets high learning rates for LLM pre-training.

Sebbouh, O., Gower, R. M., and Defazio, A. (2021). On the (asymptotic) convergence of stochastic gradient descent and stochastic heavy ball. In Conference on Learning Theory, COLT 2021 , Proceedings of Machine Learning Research. PMLR.

Shamir, O. and Zhang, T. (2013). Stochastic gradient descent for non-smooth optimization: Convergence results and optimal averaging schemes. In Proceedings of the 30th International Conference on Machine Learning .

Sriram, A., Zbontar, J., Murrell, T., Defazio, A., Zitnick, C. L., Yakubova, N., Knoll, F., and Johnson, P. (2020). End-to-end variational networks for accelerated MRI reconstruction. In International Conference on Medical Image Computing and Computer-Assisted Intervention . Springer.

Sutskever, I., Martens, J., Dahl, G., and Hinton, G. E. (2013). On the importance of initialization and momentum in deep learning. In Proceedings of the 30th International Conference on International Conference on Machine Learning - Volume 28 . JMLR.org.

Szegedy, C., Vanhoucke, V., Ioffe, S., Shlens, J., and Wojna, Z. (2016). Rethinking the inception architecture for computer vision. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 2818-2826.

Tao, W., Pan, Z., Wu, G., and Tao, Q. (2018). Primal averaging: A new gradient evaluation step to attain the optimal individual convergence. IEEE Transactions on Cybernetics , PP:1-11.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. (2017). Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc.

Wiseman, S. and Rush, A. M. (2016). Sequence-to-sequence learning as beam-search optimization. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing . Association for Computational Linguistics.

Wu, Y. and Johnson, J. (2021). Rethinking "batch" in batchnorm.

Zagoruyko, S. and Komodakis, N. (2016). Wide residual networks. In Proceedings of the British Machine Vision Conference (BMVC) .

Zamani, M. and Glineur, F. (2023). Exact convergence rate of the last iterate in subgradient methods.

Zbontar, J., Knoll, F., Sriram, A., Muckley, M. J., Bruno, M., Defazio, A., Parente, M., Geras, K. J., Katsnelson, J., Chandarana, H., et al. (2018). fastMRI: An open dataset and benchmarks for accelerated MRI. arXiv preprint arXiv:1811.08839 .

Zhang, M., Lucas, J., Ba, J., and Hinton, G. E. (2019). Lookahead optimizer: k steps forward, 1 step back. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc.

Zinkevich, M. (2003). Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the Twentieth International Conference on International Conference on Machine Learning , pages 928-935.

Contributions

Aaron Defazio discovered the method, led research experimentation and proved initial versions of Theorems 1 and 7, 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 β 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.

A Proof of Theorem 2

Theorem 2. Let F be a convex function. Let ζ 1 , . . . , ζ T be an iid sequence such that F ( x ) = E ζ [ f ( x, ζ )] . Let z 1 , . . . , z T be arbitrary vectors and let w 1 , . . . , w T and β 1 , . . . , β T be arbitrary numbers in [0 , 1] such that z t , w t and β t are independent of ζ t , . . . , ζ T . Set:

Then we have for all x ⋆ :

Proof. Throughout this proof, we will use the notation w 1: t = ∑ t i =1 w i . The result is established by showing the following identity:

Where here ∇ F ( y t ) indicates a subgradient of F at y t with E [ g t | z t ] = ∇ F ( y t ) . Given the identity (24), we sum over all t from 1 to T . Then the LHS will telescope to obtain:

from which the conclusion immediately follows since E [ g t | z t ] = ∇ F ( y t ) . So, let us establish (24). To do so, it will help to observe the following identities:

Now, setting ∇ F ( x t ) to be an arbitrary subgradient of F at x t , we have:

using (25):

using (26):

Finally, recall that any convex function satisfies ⟨∇ F ( b ) - ∇ F ( a ) , a -b ⟩ ≤ 0 for all a, b . This classical fact can be established by adding the following two subgradient identities:

Then, since β t ∈ [0 , 1] , we have w t β t 1 -β t ⟨∇ F ( x t ) -∇ F ( y t ) , y t -x t ⟩ ≤ 0 , which establishes the desired identity (24).

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:

Here u t is playing the role of the 'update vector', as the sequence of points y t are where we will be evaluating gradients. The ∆ t value can be interpreted as a 'base update' value: for the case that the z t sequence is specified by SGD (as in Theorem 1), ∆ t = -ηg t . Thus, the update can be interpreted as a momentum term m t , plus an extra 'push' in the direction of ∆ t scaled by 1 -β t .

Proof. Let's solve for m t in terms of previous values:

Then:

Now let's solve for u t :

In the special case that w t = 1 for all t , the updates simplify to:

In the special case that β t = β for all t , the update for u t simplifies to:

From this, it is clear that if β = 1 and 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 , while if β = 0 , then we have ordinary SGD without momentum.

B.1 Recovering Linear Decay

Let's take a look at the update for u t = y t +1 -y t in the special case that w t = 1 for all t :

Let us define α t = 1 -β t . Then we can re-write this update as:

It looks like we might be able to set α t such that the coefficient of m t vanishes. In this case, α t would play the role of a 'schedule' as the update would just be u t = α t ∆ t . Solving the recursion we get:

Amazingly, this recursion is satisfied by α t = T -t T , which is the linear decay schedule! Notably, this schedule has α T = 0 , which in turn implies that y T = x T , so that the last iterate of our algorithm is x T , for which Theorem 2 provides a convergence guarantee.

The recursion is also satisfied by α t = 1 for all t (which recovers standard Polyak-Ruppert averaging). Notably, this recursion shows that α 1 will determine all subsequent α values. The values will decease linearly to zero, and then they will try to go negative, which is not allowed. So the linear decay schedule is the value of α 1 that is 'just barely' allowed since it hits zero at α T .

In general with arbitrary w t , the recursion is:

If we insist that α T = 0 (so that y T = x T and we get a 'last iterate' guarantee), then solving the recursion yields:

which exactly recovers the main result of Defazio et al. (2023).

C Generalizing Theorem 2 via Bregman Divergences

Here, we provide a generalized version of Theorem 2 in the style of Joulani et al. (2020). This result employs Bregman divergences to tighten the inequality of Theorem 2 to an equality.

Theorem 5. Let F be a convex function. Let ζ 1 , . . . , ζ T be a sequence of i.i.d. random variables, and let g be a function such that E [ g ( x, ζ t )] ∈ ∂F ( x ) for all x and t . Let z 1 , . . . , z T be arbitrary vectors and let w 1 , . . . , w T and α 1 , . . . , α T be arbitrary non-negative real numbers with α t ≤ 1 such that z t , w t and α t are independent of ζ t , . . . , ζ T . Define the Bregman divergence of F as B F ( a, b ) = F ( a ) -F ( b ) -⟨∇ F ( b ) , a -b ⟩ 2 . Set:

Define the 'compressed sum' notation: w 1: t = ∑ t i =1 w i , with w 1:0 = 0 . Then we have for all x ⋆ :

Let's take a minute to unpack this result since it is depressingly complicated. Recall that the Bregman divergence for a convex function must be positive, and so all the subtracted Bregman divergence terms can be dropped to make the bound only looser. This recovers Theorem 2. However, in Section D, we show how to exploit the negative Bregman terms to achieve accelerated rates when F is smooth, and in Section E we show how to exploit the negative Bregman terms to achieve faster rates when F is strongly-convex.

Proof. The proof is nearly the same as that of Theorem 2. The only difference is that we keep track of all the error terms in the inequalities via Bregman divergences.

Throughout this proof, we use ∇ F ( x ) to indicate E ζ [ g ( x, ζ )] . When F is differentiable, this is simply the ordinary gradient at x . When F is non-differentiable, this reprents a specific choice of subgradient at x .

Recall that any convex function satisfies ⟨∇ F ( b ) -∇ F ( a ) , a -b ⟩ = -B F ( a, b ) -B F ( b, a ) for all a, b . This classical fact can be established by adding the following two subgradient identities:

The Theorem is established by showing the following identity:

Given the identity (28), we sum over all t from 1 to T . Then the LHS will telescope to obtain:

from which the conclusion immediately follows since E [ g t | g 1 , . . . , g t -1 ] = E [ ∇ F ( y t ) | g 1 , . . . , g t -1 ] . So, let us establish (28). To do so, it will help to observe the following identities:

So, we have:

using (29):

using (30):

using (27):

D Acceleration

In this section, we show that by instantiating our framework with an optimistic online learning algorithm (Rakhlin and Sridharan, 2013), we achieve accelerated convergence guarantees. Our results match those available in the prior literature (Kavis et al., 2019; Joulani et al., 2020). Our approach is inspired by Joulani et al. (2020),: their method is based upon a version of Theorem 5 for the special case that α t = 0 . Our result simply extends their analysis to α t = O (1 /t ) .

First, we establish an important technical Corollary that simplifies Theorem 5 in the case that F is smooth and α t is sufficiently small.

Corollary 1. Under the same conditions as Theorem 5, suppose additionally that F is L -smooth and suppose α t ≤ w t 10 w 1: t for all t . Then we have for all x ⋆ :

where above the value of y 0 is arbitrary (since the coefficient is w 1:0 = 0 ).

Proof. The key thing is to observe that smoothness implies B F ( a, b ) ≥ 2 L ∥∇ F ( a ) -∇ F ( b ) ∥ 2 . The rest of the argument is straightforward manipulation of the terms in Theorem 5:

Next, observe that for any vectors a, b, c , for any λ > 0 :

where we have used Young's inequality: |⟨ v, w ⟩| ≤ ∥ v ∥ 2 2 λ + λ ∥ w ∥ 2 2 . Therefore, setting λ t = 3 we obtain:

Now, since α t ≤ w t 10 w 1: t ≤ 1 , we obtain:

Now summing over t from 1 to T (and dropping one negative term), the sum telescopes to:

The result now follows from Theorem 5.

Now, we consider the case that z t is given by an optimistic mirror descent algorithm:

Corollary 2. Suppose F is L -smooth. Define g 0 = 0 and suppose also that for some D satisfying D ≥ ∥ y 1 -x ⋆ ∥ :

Finally, suppose E [ ∥ g t -g t -1 ∥ 2 ] ≤ ∥∇ F ( y t ) -∇ F ( y t -1 ) ∥ 2 + σ 2 t for some constants σ 1 , . . . , σ T (these are just variance bounds on the stochastic gradient oracle). Then with w t = t and α t ≤ 1 5( t -1) , we have:

where σ is uniform upper-bound on σ t . Note that the algorithm does not need to know L or σ .

Algorithms producing z sequences obtaining the guarantee stated here are called 'optimistic online learning algorithms'.

Proof. Applying Corollary 1, we obtain immediately:

Using the identity A √ C -BC ≤ A 2 4 B :

Divide by T ( T +1) 2 to conclude the result.

D.1 An Optimistic Regret Bound

In this section we provide an algorithm that achieves the optimistic regret bound required for our acceleration result Corollary 2. This algorithm is a mild variation on the established literature (Rakhlin and Sridharan, 2013; Chiang et al., 2012; Hazan and Kale, 2010; Joulani et al., 2017) to slightly improve a technical dependence on the maximum gradient value.

Lemma 1. For a sequence of vectors g 1 , . . . , g T , set η t = D √ 2 ∑ t i =1 ∥ g i -g i -1 ∥ 2 with g 0 = 0 , define m t = max i ≤ t ∥ g i -g i -1 ∥ and define the sequence of vectors z t , z ′ t and ˜ g t by the recursions:

Then:

Proof. For purposes of notation, define g 0 = 0 and z ′ 0 = 0 . Further, observe that:

Next, notice that z ′ t +1 = argmin ∥ z ∥≤ D ⟨ ˜ g t , z ⟩ + 1 2 η t ∥ z -z ′ t ∥ 2 . Therefore since ∥ x ⋆ ∥ ≤ D , by first order optimality conditions:

Similarly, we have z t = argmin ∥ z ∥≤ D ⟨ g t -1 , z ⟩ + 1 2 η t -1 ∥ z -z ′ t ∥ 2 . From this we have:

by Young's inequality:

So, combining these facts (and noticing that η t -1 ≥ η t :)

So, we have:

E Strongly Convex Losses

Suppose that the expected loss F is actually known to be µ -strongly convex. Then we'd like to have a convergence guarantee of O (1 /µT ) . This is achieved in Theorem 6 below.

Theorem 6. Under the same assumptions as Theorem 5, define ℓ t ( z ) = ⟨ g t , z ⟩ + µ 2 ∥ y t -z ∥ 2 . Define the 'regret' of the sequence z t as:

Then we have for x ⋆ = argmin F :

In particular, suppose ∥ x ⋆ ∥ ≤ D for some known bound D and ∥ g t ∥ ≤ G for all t for some G so long as ∥ y t ∥ ≤ D . Then if we define w t = t for all t and set z t by:

Proof. From Theorem 5, we have:

then we have:

Now, since F is µ -strongly convex, we have B F ( x ⋆ , y t ) ≥ µ 2 ∥ y t -x ⋆ ∥ 2 . Further, we have:

From this we obtain the desired result:

For the final statement, observe that with w t = t , w t ℓ t ( z ) = t ⟨ g t , z ⟩ + tµ 2 ∥ z -y t ∥ 2 is tµ -strongly convex. Therefore if we use learning rate η t = 1 µw 1: t = 2 µt ( t +1) , then standard analysis of projected OGD yields:

where in the last inequality we have observed that since ∥ z t ∥ ≤ D and y t is a linear combination of past z values, ∥ y t ∥ ≤ D as well. Finally, observing that w 1: T = T ( T +1) 2 , the result follows.

F Large Step size convergence

Theorem 7. Consider the online learning setting with bounded gradients g t . Let z t +1 = z t -γg t . Let D = ∥ z 1 -z ∗ ∥ for arbitrary reference point z ∗ and define G = max t ≤ T ∥ g t ∥ . Suppose that the chosen step-size is γ = D/G , then if it holds that:

Proof. Consider SGD with fixed step size γ :

then:

Let

Recall from D-Adaptation (Defazio and Mishchenko, 2023) theory that:

and:

Now suppose that the regret at time T is negative. Then trivially the theorem holds:

therefore, without loss of generality we may assume that ∑ T t =1 γ ⟨ g t , z t -z ∗ ⟩ ≥ 0 . Then from combining Equation 34 with Equation 33 we have:

This is a quadratic equation in ∥ s T +1 ∥ which we can solve explicitly via the quadratic formula, taking the largest root: √

Plugging in the values a = -1 2 , b = D , c = 1 2 ∑ T t =1 γ 2 ∥ g t ∥ 2 :

Substituting this into Equation 34:

Therefore, if ∑ T t =1 ⟨ g t , z t -z 1 ⟩ ≤ D √ ∑ T t =1 ∥ g t ∥ 2 then:

Therefore:

Plugging in γ = D/G :

and the theorem follows.

G Experimental Setup

G.1 Convex experiments

Each dataset is obtained from the LIBSVM repository and used without modifications.

G.2 CIFAR-10

Weused custom training code based on the PyTorch tutorial code for this problem. Following standard data-augmentation practises, we appliyed random horizontal flips and random offset cropping down to 32x32, using reflection padding of 4 pixels. Input pixel data was normalized by centering around 0.5.

G.3 CIFAR-100

We used the same codebase as for our CIFAR-10 experiments, with the same data augmentation.

We normalized each input image using fixed mean and standard error values derived from preprocessing the data.

G.4 SVHN

We used the same codebase as for our CIFAR experiments, and following the same data preprocessing.

G.5 ImageNet

We used the same code-base as for our CIFAR-10 experiments, and applied the same preprocessing procedure. The data-augmentations consisted of PyTorch's RandomResizedCrop, cropping to 224x224 followed by random horizontal flips. Test images used a fixed resize to 256x256 followed by a center crop to 224x224.

G.6 IWSLT14

We used the FairSeq framework 3 for our experiments. Rather than a vanilla LSTM we use the variant from Wiseman and Rush (2016) provided in the FairSeq codebase.

G.7 NanoGPT

We followed the NanoGPT codebase 4 as closely as possible, matching the default batch-size, training length and schedule. Our runs replicate the stated 2.85 loss in the documentation. Disabling gradient norm clipping is crucial for the Schedule-Free runs.

G.8 MAE

Our implementation uses the offical code 5 , with hyper-parameters following examples given in the repository.

G.9 DLRM

We used a custom implementation of the DLRM model based on the publicly available code. Our optimizer uses dense gradients for implementation simplicity, although sparse-gradients using AdaGrad is a more common baseline on this problem, we consider AdaGrad variants of our scheduling approach as future work.

G.10 MRI

We used the version of the the fastMRI code base at https://github.com/facebookresearch/ fastMRI/tree/main/banding_removal . Note that we found that training failed using PyTorch 2 or newer, and so we ran these experiments using PyTorch 1.9.

G.11 Algoperf

Our full algoperf entry is availiable at https://github.com/facebookresearch/schedule_ free/tree/main/schedulefree/algoperf . The hyper-parameters used for the self-tuning track submission are listed below.

H Convex Experiments

0

20

40

60

80

100

Epoch

40

60

80

Accuracy (%)

Sensorless

Polyak Averaging (85.5 SE 0.00)

Schedule-Free (89.7 SE 0.00)

0

20

40

60

80

100

Epoch

60

70

80

90

100

Accuracy (%)

Aloi

Polyak Averaging (95.7 SE 0.01)

Schedule-Free (97.2 SE 0.01)

0

20

40

60

80

100

Epoch

85

90

95

100

Accuracy (%)

DNA

Polyak Averaging (100.0 SE 0.01)

Schedule-Free (100.0 SE 0.00)

0

20

40

60

80

100

Epoch

50

55

60

65

70

Accuracy (%)

Polyak Averaging (68.8 SE 0.10)

Schedule-Free (72.1 SE 0.14)

0

20

40

60

80

100

Epoch

87.5

90.0

92.5

95.0

97.5

Accuracy (%)

Polyak Averaging (98.3 SE 0.08)

Schedule-Free (98.6 SE 0.00)

0

20

40

60

80

100

Epoch

40

50

60

70

80

Accuracy (%)

Letter

Polyak Averaging (77.8 SE 0.00)

Schedule-Free (77.8 SE 0.00)

0

20

40

60

80

100

Epoch

50

60

70

80

90

100

Accuracy (%)

Pendigits

Polyak Averaging (95.8 SE 0.01)

Schedule-Free (96.0 SE 0.01)

0

20

40

60

80

100

Epoch

60

70

80

90

100

Accuracy (%)

smallNORB

Polyak Averaging (90.8 SE 0.01)

Schedule-Free (98.1 SE 0.00)

0

20

40

60

80

100

Epoch

90

92

94

96

98

100

Accuracy (%)

USPS

Polyak Averaging (98.0 SE 0.01)

Schedule-Free (99.5 SE 0.00)

0

20

40

60

80

100

Epoch

60

70

80

Accuracy (%)

Vehicle

Polyak Averaging (80.8 SE 0.04)

Schedule-Free (83.4 SE 0.04)

0

20

40

60

80

100

Epoch

50

60

70

80

Accuracy (%)

Vowel

Polyak Averaging (76.8 SE 0.10)

Schedule-Free (77.5 SE 0.03)

0

20

40

60

80

100

Epoch

96

97

98

99

100

101

Accuracy (%)

Wine

Polyak Averaging (100.0 SE 0.00)

Schedule-Free (100.0 SE 0.00)

I Polyak and Primal Averaging Runs

These experiments follow the same tuning setup as Figure 3, where the learning rate and momentum is tuned separately for each method. In each case the c weighting sequence used for Schedule-Free training is also used to ensure a fair comparison. The Polyak averaging runs include momentum in the base optimizer as we found this gave the best results. We ran the NanoGPT experiment for a shorter 200,000 steps due to computational budget considerations. The NanoGPT Polyak averaging runs show a divergence in test loss for Polyak averaging.

0

50

100

150

200

250

300

Epoch

75

80

85

90

95

Test Accuracy (%)

CIFAR-10 (WRN-16-8)

Polyak Averaging (95.56% SE 0.05)

Schedule-Free Reference 96.03%

0

50

100

150

200

250

300

Epoch

40

50

60

70

80

Test Accuracy (%)

CIFAR-100 (DenseNet)

Polyak Averaging (76.99% SE 0.12)

Schedule-Free Reference 78.71%

0

50

100

150

200

250

300

Epoch

94

95

96

97

98

99

Test Accuracy (%)

SVHN (ResNet-3-96)

Polyak Averaging (98.32% SE 0.01)

Schedule-Free Reference 98.40%

0

20

40

60

80

100

Epoch

40

50

60

70

Test Accuracy (%)

ILSVRC 2012 ImageNet (ResNet-50)

Polyak Averaging (72.84% SE 0.02)

Schedule-Free Reference 76.90%

0

10

20

30

40

50

Epoch

0.906

0.908

0.910

0.912

Test SSIM

fastMRI Knee (VarNet 2.0)

Polyak Averaging (0.9088 SE 0.00141)

Schedule-Free Reference 0.9112

0

20

40

60

80

100

Epoch

0.780

0.785

0.790

Test Accuracy

Primal Averaging (0.7899 SE 0.00006)

Polyak Averaging (0.7901 SE 0.00004)

Schedule-Free Reference 0.7915

0

20

40

60

80

100

Epoch

65

70

75

80

Test Accuracy (%)

MAE ImageNet Finetune (ViT)

Polyak Averaging (82.72 SE 0.03)

Schedule-Free Reference 83.54%

0

50000

100000

150000

200000

Step

3

4

5

6

7

Test Loss

OpenWebText (GPT-2 124M)

Polyak Averaging (5.588 SE 0.134)

Schedule-Free Reference @ 200k 2.878

J Additional LR Sensitivity Plots

10

1

10 0

10 1

LR

94

95

96

Test Accuracy (%)

CIFAR-10 (WRN-16-8) LR Sensitivity

Schedule-Free

Cosine Schedule

10

2

10

1

10 0

10 1

LR

70

72

74

76

78

Test Accuracy (%)

CIFAR-100 (DenseNet) LR Sensitivity

Schedule-Free

Cosine Schedule

10

2

10

1

10 0

LR

98.1

98.2

98.3

98.4

Test Accuracy (%)

SVHN (ResNet-3-96) LR Sensitivity

Schedule-Free

Cosine Schedule

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our paper contains new theory as well as extensive experimental results for our method.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We detail the limitations of our method in the conclusion section.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All theorems stated in the paper are proven in the Appendix.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Detailed results concerning the experimental setup are provided in the ap- pendix.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have already open sourced our method. Experiments can be replicated using our open source PyTorch optimizer implementation together with existing open source code bases implementing each method.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: To the best of our knowledge, we provide sufficient information about hyperparameter settings for all experiments to be reproducible.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All experiments reported use multiple seeds, and error bars are included in all plots.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: As our method has no runtime overhead compared to existing approaches, we do not report computation times.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have abided by the code of ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA] .

Justification: There is no foreseeable societal impact from our work.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Not Applicable

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [No]

Justification: We do not distribute any models or datasets with this work. Datasets used for evaluation have license information available at the cited source.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA] .

Justification: No new assets are being released.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA] .

Justification: Not Applicable

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA] .

Justification: Not Applicable

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new optimization procedure for deep learning that yields good any-time performance instead of requiring learning rate schedules that decay to zero. The method can be performed on top of standard optimizers, making it easy to apply to existing networks. The paper theoretically describes the method, proving desirable properties. The method is evaluated across a broad range of tasks, including various deep learning tasks (supervised learning) and simpler convex optimization tasks, typically roughly matching or slightly exceeding the performance of standard baselines. The any-time performance of the method in a single run is comparable to the pareto frontier of standard methods when using different schedule lengths, a very large advantage of the proposed method.\n\nSTRENGTHS:\n- The problem studied is of significant relevance to the community and likely to have a large impact. The use of learning rate schedules is almost universal and this has the potential to change that. Strong anytime performance can significantly simplify experimental procedures.\n- Good theoretical contribution.\n- Strong empirical evaluation across a large number of tasks with sufficiently strong baselines with very good results.\n- Significant novelty as far as I know, I have not seen any works that can provide this type of any-time performance before.\n- Paper is relatively easy to follow with clear figures.\n- Experimental validation is done over multiple seeds.\n\nWEAKNESSES:\n- Deep learning evaluation could be improved slightly, especially with respect to hyperparameter tuning and transferability of the hyperparameter values. The effects of different hyperparameters and how they change between settings is not discussed. The full ranges of the hyperparameter tuning of the method and baseline is not provided.\n\nQUESTIONS:\nMinor Suggestions:\n- L20: Specify whether g_t is evaluated at x or z\n- L29: Maybe clarify that this “T in advance” applies to many popular schedules. There are many examples of schedules that don’t have this, e.g. trapezoidal schedules / warmup-stable-decay / reduce-on-plateau.\n- L41: No additional hyperparameters over SGD *with momentum* or Adam\n- Maybe include equation numbers, they make it much easier for people to refer to the equations e.g. when discussing the paper even if you do not use these references in the manuscript. Labeling equation 1 as such seems odd since nothing else is labeled.\n- L178: Why is Lookahead like EMA? I can see the resemblance to PA but not necessarily EMA.\n- L290: It would be nice to show some results for different types of weighing when using a warmup\n\nLIMITATIONS:\nYes, no concerns here.", "TITLE": "Official Review", "RECOMMENDATION": 10, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

comments: SUMMARY:
The paper proposes a scheduler free method for training, analyzes it theoretically to show that it matches the theoretical benefits of Polyak averaging while recovering the performance of standard cosine decay used in practice. Their approach can be interpreted as being an interpolation of Polyak averaging and Primal averaging (which is equivalent to momentum). Thus it gets the benefits of acceleration while maintaining the low variance of Polyak averaging.

STRENGTHS:
1. The empirical comparison to cosine decay scheduler is done in a thorough manner.
2. Another strength of this work is that alternative approaches such as Polyak averaging or EMA require maintaining one extra copy of the weights.

WEAKNESSES:
Since the focus of the paper is to match theoretical guarantees of Polyak averaging while matching empirical performance of schedulers they also need to compare to Polyak averaging and EMA variants[1]. The paper does not do this thoroughly. 
1. Line 22 states “Despite their theoretical optimality, PR averages give much worse results in practice than using the last-iterate of SGD with well-chosen learning rate sequences (Figure 2a)” but Figure 2a states the Polyak-averaging diverges. Clearly this could be fixed with a slightly smaller learning rate since cosine decay did not destabilize? 
2. The paper does not compare to EMA which is another method which can empirically reduces variance while not using schedules.

QUESTIONS:
Could the authors add comparison to (tuned) EMA? I would happy to increase my rating further if the authors do this.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method for training neural nets without needing to know the total training time T in advance. This contrasts with a standard training setup where one chooses a learning rate schedule in advance, and the schedule must include an a priori chosen stopping time T (e.g. cosine schedule or linear decay). If the method really works as advertised it therefore simplifies practical training setups and reduces the cost of training, by avoiding the need to run multiple training runs with different values of stopping time T.\n\nThe method works by tracking two sequences: one is the \"noisy\" sequence z that integrates noisy gradient updates, the second is the \"smoothed\" sequence x that tracks a uniform average over all past noisy iterates z. Gradients are evaluated at an interpolation between the current noisy z and smoothed x sequence, and added back to the noisy z sequence which is then averaged into the smoothed x sequence.\n\nAuthors present theoretical results about their technique for convex optimisation, and evaluate the method across a set of training benchmarks, some that involve tuning and some that don't. The results look promising.\n\nSTRENGTHS:\n- the idea is legitimately really cool and clever and seems very original to me. People know that weight averaging on top of standard training can boost performance, so it's really cool to ask \"is there a way to fold this back into the training loop\" as the authors do. \n- my broad understanding of the technique is that it's doing noise filtering in a clever way, and I think that this idea can inspire others to explore this direction, which is great\n- the paper is very clearly written and I like the mix of the intuitive description of the technique, the formal results, and then the experimental evaluation\n- it's great that the authors engaged with the MLCommons Algorithmic Efficiency benchmark, and from what I understand the results are very promising.\n\nWEAKNESSES:\nOkay as I've mentioned, I think the idea is cool and original and can inspire followup work, which is all we really want of a paper. So I'm going to focus on giving you feedback which is intended to be constructive and I hope can help you generally improve the quality of the work or followup work that you do. I'm not going to hold the paper hostage, but if I feel like you engage meaningfully with my critiques I'm willing to upgrade my score. All of my critiques can be addressed either by making minor amendments to text or by adding a limitations section at the end of the paper *or* by running more experiments and addressing them, but I won't insist on which and leave this up to the authors.\n\n### **Paper may be over claiming slightly on its results**\n\nThe paper makes claims to be **\"one of the largest and most comprehensive machine learning optimization algorithm evaluations”**. Can you quantify in what sense this is true? To me the evaluation doesn't feel that comprehensive---for instance, there is only one plot on hyperparameter sensitivity, while a broad swathe of the community has started to explore this question as standard in evaluations papers. The authors also state their method **\"has no notable memory, computation or performance limitations compared to scheduling approaches\"**. In my opinion, the truth value of this statement is essentially unknowable without more thorough evaluation. It would be fine to amend the statement by saying \"WE BELIEVE THAT our method has no notable memory, computation or performance limitations compared to scheduling approaches.\"\n\n### **Hyperparameter sensitivity is not thoroughly investigated**\n\nIn my opinion, the main potential limitation of the work is that tuning hyperparameters (learning rate, weight decay, interpolation constant $\\beta$) may implicitly be tuning an LR schedule. This is especially the case for this technique since the role of these hyperparameters is quite subtle given the unusual form of the update sequences. Now the authors may argue that they address this with their MLCommons Algorithmic Efficiency experiments that reportedly use the same set of hyperparameter across all tasks, but in my opinion this could be a fluke relating to all the MLCommons tasks involving similar training times for instance. What I would want to see to convince me otherwise is to see experiments that sweep across beta for different training times, and check if the optimal beta is invariant to training time, say.\n\n### **Paper applies a potentially ill-suited theoretical framework to deep learning, and as such there are gaps**\n\nThe paper comes from a part of the community that uses convex optimization frameworks to do algorithm design, and then extends the methods to deep network training. Sometimes the extension can feel a bit forced: e.g. let's switch to Adam instead of SGD as our base optimiser because it works better, let's use learning rate warmup because it works better, etc. The paper uses theorems within this convex framework to support its significance, but as far as I can tell these theorems don't answer basic questions a practitioner would have about the technique: e.g. Theorems 1 and 2 provided no guidance on setting the interpolation parameter $\\beta$ since they hold uniformly for all $\\beta \\in [0,1]$.\n\nI want to point out that another part of the community is trying to build up understanding and frameworks for deep learning fundamentals that involve less of a \"jump\" from convex optimisation to deep learning. To point out two examples:\n- https://arxiv.org/abs/2103.00065 \"Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability\"\n- https://arxiv.org/abs/2405.14813 \"Scalable Optimization in the Modular Norm\" (concurrent work)\n\nI think it would be great to at least read these works and potentially engage with them---not in this paper, but in the future\n\n### **Final note**\n\nI sometimes worry about giving frank feedback as I think I can sometimes have a blunt style. I want to re-emphasise that I really like this paper. I think the idea is clever and creative, and I really encourage you to pursue these directions further. The feedback is intended constructively. I am currently giving the paper a score of 5 because a score of 6 requires \"no major concerns with respect to evaluation\" and I am not there yet. I can get there if I become confident that you've engaged with my review.\n\nQUESTIONS:\n- do you know if beta is sensitive to e.g. training time? Fine if you don't know, but consider adding a limitation saying future work could investigate this closer\n- which sequence is used to do inference x, or y, or z? It might be worth clearly flagging this in this paper. Sorry if I just missed it.\n\nLIMITATIONS:\n\"Note that we excluded one benchmark problem, ResNet-50 training, as neither AdamW nor NAdamW can hit the target accuracy on that task.\" ---- this feels artificial to me. I still want to know how schedule free AdamW does!", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed an optimization style for stochastic optimization, termed schedule-free optimization, which is free of manually selected/tuned learning rate schedulers. The proposed method enjoys both the worst-case optimal last-iterate convergence and promising empirical performance. The authors also introduced a general online-to-batch analysis framework to analyze the proposed method. The empirical results show that the proposed method outperforms the state-of-the-art methods in various tasks, including training large language models (LLMs) and image classification tasks.\n\nSTRENGTHS:\n1. The mismatch between the theoretically optimal learning rate schedule and the practice is an important problem. This paper re-introduces and re-emphasizes this issue in a well-educated manner.\n2. The proposed online-to-batch analysis framework is general and insightful.\n3. The positive result in the empirical verification of equation (9) introduces an interesting new open problem that is worth further investigation: why could well-known convex/non-convex optimization problems as such an \"bounded-regret\" property under gradient descent?\n4. The proposed schedule-free method enjoys both the worst-case optimal last-iterate convergence and promising empirical performance.\n\nWEAKNESSES:\n1. In the update rule of $z_{t+1}$, using $y_t$ instead of $z_t$ could incur additional forward passes during the optimization process (of neural networks through backprop), which especially undermines the efficiency and applicability of the proposed method in the scenario of model parallelism.\n2. The adaptation from schedule-free SGD to schedule-free AdamW is still in a heuristic way.\n\nQUESTIONS:\n1. Why don't you report any empirical results that combining schedule-free methods and parameter-free methods? By the way, certain parameter-free literature is cited in the Appendix for the use of proof. The authors are encouraged to elaborate on the connection between this citation and the current work.\n2. Any conjecture/explanation on the empirical observation on Line 158?\n3. How do you deal with the \"decoupled weight decay\" issue in adapting the most popular optimizer for LLMs (i.e., AdamW) to the schedule-free style?\n\nLIMITATIONS:\nNo major societal limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}