Research map / Architectures and efficiency / RNNs and normalization
Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes why Schedule-Free AdamW tracks the river-valley loss landscape without LR decay or weight averaging (via edge-of-stability and central-flow analysis) and proposes a refined variant with a decoupling parameter C that improves momentum robustness and large-batch performance.
Problem
Fixed-budget LR schedules are inadequate for continuously scaling LM pretraining; WSD lacks visibility of model quality without decay, and weight averaging adds memory overhead. The paper asks whether Schedule-Free methods can offer flexibility, visibility, and minimal overhead, and why they are sen
Approach
Views Schedule-Free AdamW through the river-valley loss landscape, showing that its y iterates oscillate around the river (edge of stability) and that x is an implicit weighted average of y. Introduces a decoupling parameter C that sets c_t=(1-β)C/t, so averaging weights depend only on C and momentum β is set independently.
Key results
- Refined SF-AdamW with C=200 at beta1=0.95 lowers SlimPajama 0.5M-batch test perplexity from 25.12 (vanilla) to 23.60 (C=100).
- With 2M-token batches, vanilla SF-AdamW (beta1=0.98) reaches perplexity 26.09, while refined SF-AdamW with C=500 reaches 23.95 and C=1000 reaches 23.88, matching the cosine AdamW baseline.
- For suboptimal momentum beta1=0.1, refined SF-AdamW with C=10 reduces perplexity from 67.20 (vanilla) to 35.80.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add decoupling parameter C to SF-AdamW (ablation) | SF-AdamW · SlimPajama, 0.5M-token batch | test perplexity: Vanilla 25.12; C=100 23.60 | Helped |
| Add decoupling parameter C to SF-AdamW (ablation) | SF-AdamW · SlimPajama, 2M-token batch | test perplexity: Vanilla 26.09; C=1000 23.88 | Helped |
| Use refined SF-AdamW with C=200 (ablation) | SF-AdamW, β₁=0.95 · SlimPajama, 0.5M-token batch | validation loss | Helped |
| Use refined SF-AdamW with C=200 (ablation) | SF-AdamW, β₁=0.98 · SlimPajama, 2M-token batch | final performance | Helped |
| Apply LR decay or EWA to SF-AdamW (ablation) | SF-AdamW, tuned momentum · SlimPajama-6B | validation loss | No clear effect |
| Apply LR decay to SF-AdamW checkpoints with suboptimal β₁ (ablation) | SF-AdamW, β₁ ∈ {0.1, 0.5} · SlimPajama-6B | validation loss | Helped |
| Evaluate and use y_t instead of x_t for suboptimal momentum (ablation) | SF-AdamW, β₁ ∈ {0.1, 0.5} · SlimPajama-6B | validation loss | Helped |
| Apply EWA to y_t iterates (ablation) | SF-AdamW · SlimPajama-6B | validation loss | Helped |
| Increase β₁ momentum (ablation) | SF-AdamW central-flow dynamics · River-valley trajectory | oscillation variance | Helped |
Surprising findings
- Neither a short LR decay phase nor EWA improves SF-AdamW trajectories, unlike AdamW, which benefits substantially from both.
- Vanilla SF-AdamW with suboptimal momentum fails to follow the river, while the y iterates remain near the river on average.
Limitations
- Theoretical analysis relies on simplifying assumptions, and the central flow approximation is not validated in deep learning.
- Experiments are limited to small-scale (124M) language models due to compute constraints.
- Refined variant's gains are shown only with preliminary sweeps; the C parameter still requires tuning.
Experimental setup
Datasets: SlimPajama-6B OpenWebText2
Benchmarks: SlimPajama-6B held-out perplexity OpenWebText2 perplexity
Metrics: validation loss / perplexity test perplexity on 9M held-out tokens
Baselines: AdamW with cosine schedule AdamW with constant LR vanilla SF-AdamW WSD schedule exponential weight averaging
Builds on: Schedule-Free optimizer AdamW warmup-stable-decay schedule edge of stability central flow
Closest papers by meaning
- A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules (ICLR 2025 · similarity 0.90)
- PolyPythias: Stability and Outliers across Fifty Language Model Pre-Training Runs (ICLR 2025 · similarity 0.89)
- Resolving Discrepancies in Compute-Optimal Scaling of Language Models (NeurIPS 2024 · similarity 0.89)
- Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules (NeurIPS 2025 · similarity 0.89)
- Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (ICLR 2024 · similarity 0.89)
- Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.