On Neural Scaling Laws for Weather Emulation through Continual Training
2026-03-26 • Machine Learning
Machine Learning
AI summaryⓘ
The authors looked at how making weather forecasting models bigger and training them differently affects their accuracy. They used a simple yet flexible model called a Swin Transformer and trained it using a new method involving steady learning rates with breaks called cooldowns. Their results showed that this approach followed clear performance patterns and sometimes worked better than common training methods. They also tested various model sizes and training amounts to find the best ways to use computing power efficiently. This helps understand how to build better weather prediction models without wasting resources.
Neural scaling lawsSwin TransformerScientific Machine LearningWeather forecastingLearning rate scheduleCooldown phaseIsoFLOP curvesCompute-optimal trainingContinual training
Authors
Shashank Subramanian, Alexander Kiefer, Arnur Nigmetov, Amir Gholami, Dmitriy Morozov, Michael W. Mahoney
Abstract
Neural scaling laws, which in some domains can predict the performance of large neural networks as a function of model, data, and compute scale, are the cornerstone of building foundation models in Natural Language Processing and Computer Vision. We study neural scaling in Scientific Machine Learning, focusing on models for weather forecasting. To analyze scaling behavior in as simple a setting as possible, we adopt a minimal, scalable, general-purpose Swin Transformer architecture, and we use continual training with constant learning rates and periodic cooldowns as an efficient training strategy. We show that models trained in this minimalist way follow predictable scaling trends and even outperform standard cosine learning rate schedules. Cooldown phases can be re-purposed to improve downstream performance, e.g., enabling accurate multi-step rollouts over longer forecast horizons as well as sharper predictions through spectral loss adjustments. We also systematically explore a wide range of model and dataset sizes under various compute budgets to construct IsoFLOP curves, and we identify compute-optimal training regimes. Extrapolating these trends to larger scales highlights potential performance limits, demonstrating that neural scaling can serve as an important diagnostic for efficient resource allocation. We open-source our code for reproducibility.