Reproduce “Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Minhak Song, Beomhan Baek, Kwangjun Ahn, Chulhee Yun

Researchers from KAIST, SNU, and Microsoft Research investigate Schedule-Free (SF) methods for language model training, revealing SF-AdamW's ability to navigate the loss landscape's "river" without explicit learning rate decay or additional memory for weight averaging. They propose a refined SF variant that enhances robustness to momentum choices and improves performance under large batch sizes, achieving results comparable to AdamW with cosine schedules.

Key insights

  • SF-AdamW implicitly performs weight averaging, providing generalization benefits without requiring additional memory overhead.
  • SF methods operate at the 'Edge of Stability,' where the (preconditioned) sharpness at `y_t` iterates stabilizes near a theoretical threshold.
  • The `y_t` iterates in SF are more robust to momentum parameter choices and remain better aligned with the loss landscape's 'river' compared to `x_t` iterates.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.