Anytime Pretraining: Horizon-Free Learning-Rate Schedules with Weight Averaging

TL;DR

It is suggested that weight averaging combined with simple, horizon-free step sizes offers a practical and effective anytime alternative to cosine learning rate schedules for large language model pretraining.

Abstract

Large language models are increasingly trained in continual or open-ended settings, where the total training horizon is not known in advance. Despite this, most existing pretraining recipes are not anytime: they rely on horizon-dependent learning rate schedules and extensive tuning under a fixed compute budget. In this work, we provide a theoretical analysis demonstrating the existence of anytime learning schedules for overparameterized linear regression, and we highlight the central role of weight averaging - also known as model merging - in achieving the minimax convergence rates of stochastic gradient descent. We show that these anytime schedules polynomially decay with time, with the decay rate determined by the source and capacity conditions of the problem. Empirically, we evaluate 150M and 300M parameter language models trained at 1-32x Chinchilla scale, comparing constant learning rates with weight averaging and $1/\sqrt{t}$ schedules with weight averaging against a well-tuned cosine schedule. Across the full training range, the anytime schedules achieve comparable final loss to cosine decay. Taken together, our results suggest that weight averaging combined with simple, horizon-free step sizes offers a practical and effective anytime alternative to cosine learning rate schedules for large language model pretraining.

Figures (10)

• Figure 1: 150M model: Cosine schedules do not transfer across horizons. The cosine envelope (red) is formed by independently tuning a horizon-aware cosine schedule for each terminal compute budget ($1\times$--$32\times$ Chinchilla) and taking the best validation loss at that horizon. Gray curves show the same cosine schedule evaluated at intermediate checkpoints when tuned for a single fixed terminal budget. The gap illustrates why cosine decay is not anytime: tuning for a long horizon can be far from optimal at shorter budgets. An analogous plot for the 300M model appears in Figure \ref{['fig:cos-transfer-300m']} (Appendix \ref{['app:additional_figures']}).
• Figure 2: Top row: 150M; bottom row: 300M. Left: Validation loss versus training compute for cosine decay, constant LR with averaging, WSD, and a $1/\sqrt{t}$-type schedule with averaging. Specifically, the $1/\sqrt{t}$ schedule uses a multiplicative factor $\sqrt{\alpha/(t+\alpha)}$, and we tune $\alpha$. Each point corresponds to training for $1\times, 2\times, 4\times, 8\times, 16\times,$ and $32\times$ the Chinchilla compute scale. Cosine baselines are trained as separate runs tuned for each duration; in contrast, the anytime schedules come from a single run trained to $32\times$. Red stars mark the per-duration optimal cosine envelope; the red curve shows the cosine schedule tuned for the full $32\times$ run (and $16\times$, respectively). For WSD, we apply a linear decay over the final $90\%$ of training, starting from the same run used for constant LR with averaging. Right: Loss difference relative to cosine at each compute multiple (negative = better than cosine). Hyperparameters are chosen to be near-optimal across intermediate checkpoints; per-checkpoint optimal losses are reported in Figure \ref{['fig:main_fig_optimal_hps']} (Appendix \ref{['app:additional_figures']}).
• Figure 3: For a 150M-parameter model trained with batch size $4096$, we compare cosine decay to constant learning rate with averaging, a $\sqrt{\alpha/(t+\alpha)}$ schedule with averaging, and WSD across end times ranging from $1\times$ to $16\times$ Chinchilla compute. Left: Validation loss versus training compute for each schedule. Right: Loss improvement over cosine at each compute multiple, where a negative value is an improvement over cosine. A per-duration optimal hyperparameter version of this plot is in Figure \ref{['fig:large_batch_optimal_hps']} (Appendix \ref{['app:additional_figures']}).
• Figure 4: Risk comparison between schedulers in SGD on linear regression. We plot the exact risk recursion from equation \ref{['eq:bias_and_variance_eta']}. The problem dimension is $d=500000$ and we train for a maximum of $N=50000$ samples at batch size $1$, with label noise $\sigma^2 = 0.01$. We plot source exponents $a = 1.1$, $1.5$ and $1.9$ on the columns, and the top row corresponds to the capacity exponent $b = a$, and the bottom row corresponds to $b=2a$. We sweep over learning rates $\eta \in \{ 0.0001,\ 0.0002,\ 0.0005,\ 0.0007,\ 0.001,\ 0.002,\ 0.005,\ 0.01,\ 0.02,\ 0.03,\ 0.05,\ 0.075,\ 0.1,\ 0.2,\ 0.3,\ 0.5,\\ 0.8,\ 1.0 \}$. For constant with averaging and $1/\sqrt{t}$ we average over the whole duration of the run and we only use the last iterate for WSD. For WSD, we fix intermediate points during the run at $1000, 2000, 3000, 5000, 8000, 10000, 20000, 50000$ samples and run until a fraction $p$ of each with constant learning rate, followed by a linear decay, where $p \in \{0.5, 0.6, 0.7, 0.8, 0.9\}$. For each run, hyperparameters are chosen such that they are close to anytime optimal by minimizing the average loss over intermediate points.
• Figure 5: Top plots correspond to 150M parameter, and bottom plots correspond to 300M models. At each intermediate point, we plot the best loss out of the whole hyperparameter sweep, then we linearly interpolate between the points. Note that for long training durations, constant with averaging and $1/\sqrt{t}$ offers a substantial improvement over cosine decay.

Theorems & Definitions (9)

• Theorem : Informal version of Theorem \ref{['thm:tgamma_rate']}
• Theorem 1
• Remark 1
• Corollary 1
• Theorem 2
• Lemma 1
• proof
• proof : Proof of Theorem \ref{['thm:tgamma_rate']}
• proof : Proof of Corollary \ref{['cor:optimal_hps']}