Schedule-Free Learning in PyTorch

When using Schedule-Free learning, consider the following guidance:

• Learning Rate Warmup: Highly recommended. Use the warmup_steps parameter.
• Momentum ($\beta$): Training is sensitive to $\beta$. The default is $0.9$, but for very long training runs, you may need to increase this to $0.95$ or $0.98$.
• SGD Learning Rates: A good starting point is $10x$ to $50x$ larger than classical rates.
• AdamW Learning Rates: A good starting point is $1x$ to $10x$ larger than schedule-based approaches.
• Regularization: This method requires tuning; it may not outperform scheduled approaches without proper tuning of regularization and learning rate parameters.

To run the MNIST example provided in this repository, install the required dependencies and execute the main.py script. You can specify a specific GPU using the CUDA_VISIBLE_DEVICES environment variable.

pip install -r requirements.txt
python main.py
# To specify a GPU ID (e.g., 2):
# CUDA_VISIBLE_DEVICES=2 python main.py

Because Schedule-Free optimizers use two different points for gradient calls and test/val loss calculations, you must switch the parameter buffer between the two during training.

To do this, call optimizer.train() at the same time you call model.train(), and optimizer.eval() at the same time you call model.eval(). Additionally, place the optimizer in eval() mode when storing checkpoints.

If your training loop supports PyTorch Optimizer step closures, you can use the ScheduleFreeClosure versions of the optimizers, which do not require these manual .train() and .eval() calls.

Install the schedulefree package using pip to access Schedule-Free optimizer implementations for PyTorch.

pip install schedulefree

The optimizer manages internal weight sequences (x, y, and z) through its train() and eval() methods:

• .train(): Switches the optimizer to training mode. It restores the model parameters p from the y sequence (the sequence used during training steps).
• .eval(): Switches the optimizer to evaluation mode. It updates the model parameters p by copying the x sequence (the averaged weights) into the parameters. This is the state used for inference/evaluation.

Warning: If you call .step() without having called .train() first, the optimizer will raise an Exception.

RAdamScheduleFree is a Schedule-Free RAdam optimizer that eliminates the need for manual warmup hyperparameters or learning rate schedulers.

Critical Usage Requirement: This optimizer relies on switching between training and evaluation modes to manage its internal state. You must call .train() and .eval() on the optimizer instance at the appropriate times:

1. Call optimizer.train() before starting the training phase.
2. Call optimizer.eval() before starting the evaluation/validation phase.
3. When saving checkpoints, the optimizer should be in .eval() mode.

If step() is called while the optimizer is not in .train() mode, it will raise an exception.

To use ScheduleFreeWrapper correctly, you must switch between .train() and .eval() modes. This manages the interpolation between the weight points ($x$, $y$, and $z$).

• Call optimizer.train() before the training loop/step to prepare the weights.
• Call optimizer.eval() during evaluation to set the weights to the averaged version.

Note: If you call .step() while the optimizer is not in .train() mode, it will raise an Exception.

If your model uses BatchNorm, standard test/val evaluations may not work correctly because the training_mean/training_var cache is updated during the training phase at the $y$ sequence. To ensure evaluations use values calculated at the $x$ sequence, you must perform a small number of forward passes in training mode before switching to eval mode.

Alternatively, using PreciseBN will avoid this issue.

model.train()
optimizer.eval()
with torch.no_grad():
  for batch in itertools.islice(train_loader, 50):
    model(batch)
model.eval()

The experimental ScheduleFreeWrapper allows you to apply schedule-free logic to any base optimizer. When using this wrapper, you can disable the base optimizer's momentum as the wrapper manages its own.

If you set weight decay on the base optimizer, it is computed at $z$. You can use the weight_decay_at_y parameter to compute weight decay at $y$ instead, which may yield better results.

There is also a ScheduleFreeWrapperReference version for increased numerical stability during research.

base_optimizer = torch.optim.RMSprop(model.parameters(), lr=0.0025)
optimizer = ScheduleFreeWrapper(
  base_optimizer, momentum=0.9, weight_decay_at_y=0.1
)

When initializing ScheduleFreeWrapperReference, you can tune the following parameters: