📊 statistics

ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models

This paper introduces ScheduleFree+, a learning-rate-free and schedule-free method that successfully scales Schedule-Free Learning to large language models, significantly outperforming state-of-the-art training schedules like Warmup-Stable-Decay, particularly in long-duration training scenarios.

Original authors: Aaron Defazio

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Aaron Defazio

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, super-intelligent robot (a Large Language Model) how to speak human language. To do this, you show it billions of sentences and let it guess the next word, correcting its mistakes over and over.

For a long time, the standard way to teach this robot was to use a Learning Rate Schedule. Think of this like a strict teacher who starts the class with a very loud, energetic voice (high learning rate) to get the robot's attention, but then slowly whispers and fades out (low learning rate) as the class ends. The problem? You have to guess exactly when to start whispering and how fast to fade out. If you guess wrong, the robot learns poorly.

This paper introduces a new method called ScheduleFree+. Instead of a teacher who changes their voice volume over time, this method uses a "smart averaging" technique that lets the robot learn at a steady, confident pace without needing a complex schedule.

Here is how the paper explains the magic behind ScheduleFree+, using simple analogies:

1. The "Smart Averaging" Trick (The Core Idea)

In traditional training, the robot updates its brain based on the very last thing it saw. This is like a student who only remembers the last question asked in class.

Schedule-Free changes this. It keeps two versions of the robot's brain:

The "Active" Brain ( $z$ ): This is the robot updating itself in real-time, learning from the latest mistakes.
The "Average" Brain ( $x$ ): This is a calm, steady version that remembers the average of everything the robot has learned so far.

The magic happens because the robot doesn't just look at the "Active" brain to decide what to do next. It looks at a mix of the "Active" brain and the "Average" brain.

Analogy: Imagine you are navigating a ship. The "Active" brain is the captain looking at the waves right now (very reactive). The "Average" brain is the navigator looking at the map of the whole journey so far (very stable). The new method tells the captain to steer based on a mix of both. This prevents the ship from swerving wildly when it hits a small wave, leading to a much smoother, faster journey.

2. Why It Failed Before (The "Big Batch" Problem)

The authors found that while this "Smart Averaging" worked great for small robots, it broke when they tried to train giant robots with huge batches of data.

The Problem: When you feed the robot too much data at once (a large "batch"), the "Active" brain gets confused and starts to wobble, causing the training to crash.
The Fix: They added Inner Momentum. Think of this as giving the robot a little bit of "inertia" or "momentum." Just like a heavy truck is harder to stop and turn than a bicycle, adding momentum helps the robot stay on course even when the data is huge and noisy. This allowed them to scale the method up to massive models.

3. The "Self-Adjusting" Speed (No More Guessing)

Usually, humans have to manually tune how fast the robot learns. The paper introduces a Polyak Step Size, which is like a self-driving speedometer.

How it works: Instead of a human setting the speed, the robot looks at how "confused" it is (the size of its mistakes) and automatically adjusts its speed. If the mistakes are big, it slows down to think carefully. If the mistakes are small, it speeds up.
The Result: You don't need to guess the learning rate anymore. The method is "learning-rate-free," meaning it figures out the perfect speed on its own.

4. Fixing the "Drifting" Weights

The paper discovered a weird side effect: when using this new method, the robot's internal "weights" (the strength of its connections) would sometimes grow too large or shrink too small, making the robot unstable.

The Analogy: Imagine the robot's brain cells are getting too big, stretching the wires until they snap.
The Fix: They applied a special "inverse-gradient" correction. Think of this as a smart regulator that tightens the screws whenever the robot's brain cells try to get too big, keeping everything stable and preventing the wires from snapping. This allows the robot to train for much longer without breaking.

5. The "Warm-Up" and "Cool-Down" Tweaks

To make this work perfectly for long training sessions, they added two final touches:

Warm-Start: At the very beginning, they let the robot learn normally without the "averaging" trick for a short while. This is like letting a runner jog a few laps before starting the race, so they don't stumble at the starting line.
Beta Annealing: As the training goes on, they slowly change the "mix" between the Active and Average brains. Early on, they listen more to the Active brain (to learn fast). Later, they listen more to the Average brain (to be precise). This is like a student who starts by taking notes frantically, but by the end of the semester, they rely more on their well-organized study guide.

The Big Result

When they tested ScheduleFree+ on massive language models:

It was faster: To reach the same level of intelligence, it took 31% less time than the best existing methods.
It was more stable: The learning curve was smooth and predictable, unlike the jagged, bumpy curves of older methods.
It works for "Anytime" training: You can stop the training at any second, and the robot will be at its best possible state for that moment. You don't have to wait for a specific "end of schedule" to get a good model.

In summary: The paper shows that by mixing the robot's current thoughts with its average thoughts, adding a little momentum for stability, and letting the robot adjust its own speed, we can train giant AI models faster, smoother, and without needing humans to constantly tweak the settings.

Technical Summary: ScheduleFree+ Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models

1. Problem Statement

Traditional machine learning optimization relies heavily on learning rate schedules (e.g., Linear Decay, Cosine, Warmup-Stable-Decay) to ensure convergence. While effective, these methods require significant tuning and often result in unstable or unpredictable training curves. Recent "Schedule-Free" learning methods, which replace schedules with aggressive iterate averaging, have shown promise on small-scale benchmarks but have historically failed to scale to Large Language Model (LLM) training. Specifically, prior Schedule-Free implementations suffered from:

Divergence at large batch sizes: Performance degraded significantly as batch sizes increased beyond small-scale limits.
Gradient norm drift: The method required unusually small weight decay values to function, leading to unstable weight norms and poor performance in long-duration training.
Suboptimal convergence in long runs: The method did not fully leverage the benefits of averaging over extended training horizons (e.g., 1000 tokens per parameter).

The paper addresses the challenge of scaling Schedule-Free Learning to match or exceed the performance of state-of-the-art (SOTA) schedule-based methods for LLM pretraining, aiming for a stable, "anytime" training procedure that requires no learning rate tuning.

2. Methodology: ScheduleFree+

The authors propose ScheduleFree+, a composite method that integrates several specific modifications to the base Schedule-Free algorithm to make it robust for LLMs. The core components are:

A. Reintroduction of Inner Momentum

While the original Schedule-Free work removed inner momentum (e.g., Adam's $\beta_1$ ) to save memory, the authors found this detrimental for large batch sizes.

Fix: Reintroduce inner momentum (e.g., $\beta_1 = 0.75$ or $0.9$) into the base optimizer.
Effect: This restores robustness to larger batch sizes and allows for higher learning rates without divergence, addressing the "scaling cliff" observed in standard Schedule-Free AdamW.

B. Inverse-Gradient Norm Weighting (Polyak Step Size)

Standard Schedule-Free learning often required small weight decay to compensate for increasing gradient norms, which caused instability. The authors derive a theoretical correction based on the Polyak step size.

Mechanism: The step size $\gamma_t$ is scaled inversely proportional to the L1 norm of the gradient ( $\gamma_t \propto 1/\|g_t\|_1$ ).
Implementation: An Exponential Moving Average (EMA) of the gradient L1 norm is used to estimate the denominator, approximating the optimal weighting derived from convex regret bounds.
Result: This correction counteracts the effective learning rate drift caused by normalization layers, allowing the use of standard, larger weight decay values (e.g., 0.05–0.5) which stabilize weight norms and prevent the "shrinking weight norm" phenomenon that destabilizes long runs.

C. Learning-Rate-Free Optimization via Polyak Step

Building on the inverse-gradient weighting, the authors implement a fully learning-rate-free variant using a modified Polyak step size:
$\gamma_t = \frac{f(y_t) - f^* + \beta \langle \nabla f(y_t), z_t - x_t \rangle}{\sqrt{\pi/2} \| \nabla f(y_t) \|_1}$

Stability: The numerator uses an estimate of the function value gap ( $f(z_t) - f^*$ ) rather than the noisy current iterate, and the denominator uses the smoothed L1 norm. This avoids the vanishing denominator issues common in stochastic Polyak steps.
Adaptivity: This approach automatically scales the learning rate with batch size (following $\sqrt{b}$ scaling for AdamW) and eliminates the need for manual learning rate grid searches.

D. Warm-Start and Annealing Strategies

C-Warmup: To prevent the averaging buffer from diluting early, rapid convergence phases, the authors delay the averaging update ( $c_t = 1$ ) for an initial period (e.g., $2\times$ the learning rate warmup duration).
Beta Annealing: The interpolation parameter $\beta$ (mixing the current iterate $z_t$ and average $x_t$ ) is annealed from a lower value (e.g., 0.9) to a higher value (e.g., 0.965) over the training run. This captures fast initial convergence while optimizing for long-term loss reduction.
Weighted Averaging ( $r=1$ ): For long runs, the averaging weights are set to $r=1$ (linearly increasing weights) rather than the default $r=0$ , which accelerates convergence for non-strongly convex problems.

3. Key Results

The authors conducted extensive scaling experiments on LLMs ranging from 120M to 1B parameters, trained on up to 1000 tokens per parameter.

Long-Duration Training (1000 TPP): At this scale, ScheduleFree+ significantly outperforms both Warmup-Stable-Decay (WSD) and Linear Decay schedules.
- Efficiency: It achieves the same final loss as a Linear Decay schedule in 31% less training time (a 45% reduction in tokens required to reach the target loss).
- Stability: The method produces smooth, predictable loss curves that follow a theoretical $1/\sqrt{t}$ decay rate, allowing for accurate prediction of final loss from early training data.
Short-Duration Training (20–100 TPP): ScheduleFree+ outperforms WSD schedules across all model sizes. Against Linear Decay, it matches or slightly underperforms at very short horizons (20 TPP) due to the time required to reach a steady state, but generally performs competitively.
Batch Size Scaling: With inner momentum, ScheduleFree+ scales robustly to batch sizes of 8M, matching the performance of standard AdamW with Linear Decay.
Weight Norm Stability: By using the Polyak step size and larger weight decay, the method maintains stable weight norms, avoiding the divergence and quantization issues associated with the small weight decay values required by unmodified Schedule-Free methods.

4. Significance and Claims

The paper claims that ScheduleFree+ provides a theoretical foundation for model averaging and checkpoint merging during pretraining, bridging the gap between online optimization and offline model souping.

Anytime Training: Unlike WSD schedules, which require a specific annealing phase at the end of training to realize gains, ScheduleFree+ offers a true "anytime" capability. The model is usable and optimal at any stopping point without a dedicated decay phase.
Elimination of Tuning: The method removes the need for learning rate grid searches, relying instead on the Polyak step size and inverse-gradient weighting to adapt automatically to batch size and training duration.
Predictability: The authors argue that the loss curves of Schedule-Free Learning are not just empirically stable but theoretically grounded, closely following the $1/\sqrt{t}$ convergence rate predicted by convex optimization theory, even in the non-convex LLM setting.

In summary, the paper demonstrates that with specific modifications—inner momentum, inverse-gradient weighting, and adaptive annealing—Schedule-Free Learning can be scaled to become a superior, stable, and tuning-free alternative to traditional learning rate schedules for training Large Language Models.