Cloudy with a Chance of Convergence — Optimizer Benchmarking for the PARADIS Weather Model.

Project Member · Machine Learning & Scientific Computing · M2PI 2026 (PIMS × Environment and Climate Change Canada)

Note: this page is based on a working draft of the project report; the final version will replace it once published.

Figure 1. Global spatial forecast-error maps for the Adam, SGD and MUD optimizers. MUD (right) produces visibly the smallest error, best preserving the underlying weather field.

Overview

PARADIS is a physics-inspired machine-learning model developed by Environment and Climate Change Canada (ECCC) for global weather forecasting. Instead of a single monolithic neural network, it uses a structured advection–diffusion–reaction operator decomposition to embed physically motivated inductive biases. Because training such models is computationally expensive and the choice of optimizer strongly affects both convergence speed and forecast quality, this project studies which optimizers are most effective when training a PARADIS-based model — using data generated from the shallow water equations (SWE) as a lower-cost proxy for full ERA5-scale forecasting.

Approach

We benchmarked six optimizers under a shared model architecture and training setup, tuning each one independently and evaluating them across several complementary metrics:

• Optimizers compared: Adam, AdamW, Stochastic Gradient Descent (SGD), Muon, MUD, and Gauss–Newton.
• Test problem: data generated from the shallow water equations on the sphere (integrated with a spectral method using spherical harmonics), retaining key atmospheric dynamics at a fraction of the cost.
• Hyperparameter tuning: Optuna with a tree-structured Parzen estimator sampler and early-stopping pruning to select each optimizer's hyperparameters fairly.
• Evaluation metrics: training and validation loss, wall-clock runtime, autoregressive rollout stability, and spectral fidelity diagnostics.
• Compute: models were trained on ECCC's GPU cluster; the implementation was written in Python with PyTorch and PyTorch Lightning.

Figure 2. Spectral fidelity — percent difference from ground truth across wavenumbers for the rotational, divergent, potential and total components. MUD and Muon stay closest to zero, while SGD deviates strongly.

Results

Consistent with the "no free lunch" principle, the most accurate optimizers (MUD and Muon) were also the slowest to train, but found superior minima that translated into better spectral fidelity and spatial accuracy even under long autoregressive rollouts. We recommend MUD when training time is moderately constrained (near-top accuracy at a modest cost over Adam-type methods) and Muon when accuracy matters most and compute is less constrained. Adam and AdamW remain attractive for their ease of use and speed, SGD showed interesting stability but lagged overall, and Gauss–Newton proved too costly at this dimensionality.

Technologies Used

Team & Acknowledgments

This project was completed during the 2026 Math to Power Industry (M2PI) program, organized by the Pacific Institute for the Mathematical Sciences (PIMS), with Environment and Climate Change Canada as the industry partner. It was developed by a team of seven — Carlos Ballesteros, Claire Curry, Yuxuan Mao, Angela R. Morrison, Lingjun Qian, Mohammadreza Sadeghian, and Leonidas Theocharous — with guidance from ECCC mentors Valentin Dallerit, Stéphane Gaudreault, Shoyon Panday, Carlos Pereira, and Siqi Wei.