facebookresearch/schedule_free

What it solves

Training deep learning models typically requires a carefully tuned learning rate schedule (like cosine or linear decay) and a predefined stopping time to achieve optimal performance. This project provides optimizers that eliminate the need for these schedules, allowing for faster training without needing to know the total number of steps in advance.

How it works

Schedule-Free learning replaces standard optimizer momentum with a combination of interpolation and averaging. It maintains three distinct sequences: one for gradient evaluation ($y$), one for the primary iterates ($z$), and one for test/validation loss evaluations ($x$). By balancing these sequences, the method achieves fast early convergence while maintaining stability, matching or exceeding the performance of state-of-the-art schedules without requiring a decreasing learning rate.

Who it’s for

Machine learning practitioners and researchers using PyTorch (or Jax via Optax) who want to reduce the overhead of tuning learning rate schedules and avoid the constraint of fixed training durations.

Highlights

• Multiple Implementations: Includes schedule-free variants of SGD, AdamW, and RAdam.
• Flexible Integration: Offers a ScheduleFreeWrapper that can be applied to any base optimizer.
• Memory Efficient: Maintains the same memory requirements as the base optimizer (parameter buffer plus momentum).
• Coded for PyTorch: Provides easy-to-use PyTorch implementations, including versions compatible with optimizer step closures.