HEADLINE: Beyond Fixed Durations: A New Approach to Compute-Optimal Training and Scaling Laws

In a recent paper titled "Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations," researchers from EPFL (Ecole Polytechnique Fédérale de Lausanne) challenge the conventional reliance on cosine learning rate schedules in machine learning. The authors, Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi, argue that this approach unnecessarily complicates scaling experiments and increases compute costs. Instead, they propose a simpler method using constant learning rates with cooldowns, which they show to be equally effective and more efficient.

What Changed Technically?

1. Constant Learning Rate with Cooldowns:

   • The authors introduce a training setup that uses a constant learning rate for the majority of training, followed by a cooldown phase where the learning rate is gradually reduced.
   • This approach allows for consistent scaling across different training durations without the need to adjust hyperparameters for each run.

2. Stochastic Weight Averaging (SWA):

   • They incorporate SWA, which averages model weights over multiple points in the training trajectory.
   • SWA improves performance without additional training costs and is particularly beneficial when used with constant learning rates and cooldowns.

Why It Matters to Practitioners

• Reduced Compute Costs:

  • By using fewer but more reusable training runs, practitioners can significantly reduce the number of GPU hours required for scaling experiments.
  • This makes it more feasible to conduct extensive hyperparameter tuning and model exploration, especially for resource-constrained teams.

• Predictable Scaling:

  • The constant learning rate with cooldowns provides a predictable and reliable way to scale models across different training durations.
  • This predictability is crucial for designing efficient training setups and optimizing compute resources.

Key Findings

• Performance Parity:

  • The authors demonstrate that the constant learning rate with cooldowns approach achieves performance comparable to the cosine schedule, both in terms of final model accuracy and scaling behavior.
  • They show this across various model sizes and datasets, ensuring the robustness of their findings.

• Improved Efficiency:

  • By leveraging SWA, they achieve better performance along the training trajectory without increasing computational overhead.
  • This is particularly useful for models that benefit from longer training periods but are constrained by compute budgets.

Implementation Details

• Training Setup:

  • The constant learning rate phase typically covers about 90% of the total training duration.
  • The cooldown phase gradually reduces the learning rate over the remaining 10%, ensuring a smooth convergence.

• Benchmarks:

  • The authors provide detailed benchmarks comparing their approach to the cosine schedule on popular datasets like ImageNet and CIFAR-10.
  • They show that the constant learning rate with cooldowns achieves similar or better performance with significantly fewer training runs.

Conclusion

The paper "Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations" offers a compelling alternative to the traditional cosine schedule for machine learning practitioners. By simplifying the training process and reducing compute costs, this approach makes it easier to conduct scaling experiments and optimize models. The authors' findings are supported by robust benchmarks and practical implementation details, making it a valuable resource for anyone looking to improve their training efficiency.