Source-linked AI summary
Diffusion Distillation for Efficient Weather Ensembles
Yiming Yang, Valentin Brekke, James Briant, Serge Guillas
TL;DR
Iterative diffusion sampling limits the efficiency of probabilistic weather forecasting, motivating a supervised distillation framework that combines teacher-distribution matching with ground-truth supervision. The resulting single-pass student outperforms existing distillation methods, matches GENCAST forecast skill, and uses 1 rather than 39 NFEs per autoregressive step.
Problem
Diffusion models offer strong probabilistic weather forecast skill, but iterative sampling requires many neural function evaluations, limiting large ensembles and real-time deployment.
Method
The method distills GENCAST with supervised distribution matching that combines teacher forecasts, ground-truth observations, and a sample-based discrepancy without auxiliary networks or GAN-style training.
Results
The single-pass student outperforms CM and DMD, matches GENCAST forecast skill, and uses 1 rather than 39 NFEs per autoregressive step.
Takeaways & Limitations
The method preserves ensemble spread and skillfully predicts Typhoon Nanmadol while substantially reducing inference cost.
Abstract
from arXiv · showhide
Diffusion models generate skillful weather ensembles but require costly iterative sampling. We introduce a supervised energy-distance distillation method that compresses a multi-step diffusion teacher into a single-step student by aligning student forecasts with teacher samples and ground-truth observations. Experiments on global forecasting and typhoon-track prediction show that our student outperforms existing distillation methods and preserves skill for extreme events. It matches or surpasses the teacher across key metrics using only one neural function evaluation per autoregressive step.
1 Introduction
Diffusion models provide strong probabilistic weather forecasts but iterative sampling requires many neural function evaluations. The paper proposes supervised distillation to align a single-pass student with teacher distributions and ground-truth observations.
- Diffusion models achieve strong probabilistic forecast skill, but iterative sampling requires many neural function evaluations, limiting large ensembles and real-time deployment.
- Existing trajectory- and distribution-matching distillation methods produce underdispersed ensembles with lower probabilistic forecast skill than the teacher.
- The proposed supervised framework aligns student and teacher distributions while incorporating ground-truth supervision.
- The framework uses a sample-based distribution-matching objective and applies it to GENCAST to create a single-pass ensemble generator.
2 Background
Diffusion-based weather forecasting models learn conditional atmospheric distributions by reversing noise, with GENCAST generating global ensembles autoregressively. Distribution matching frames distillation as minimizing a discrepancy between student and teacher predictive distributions.
- GENCAST applies conditional diffusion to gridded atmospheric states and autoregressively generates global weather ensembles at 12-hour intervals.
- EDM training perturbs future states with Gaussian noise across a decreasing noise schedule and trains a denoiser to recover the clean state.
- At inference, GENCAST starts from noisy samples and iteratively denoises them to obtain forecast realizations from its final predictive distribution.
- Distribution matching learns a student predictive distribution by minimizing a discrepancy from a target distribution, here the pretrained teacher distribution.
- KL divergence and Maximum Mean Discrepancy are common choices for the distributional discrepancy.
3 Methodology
The method distills an iterative diffusion teacher into a single-step student using supervised observations alongside teacher forecasts. It combines ensemble-mean supervision with sample-based energy-distance matching to retain distributional information while avoiding intractable density-based discrepancies.
- Student generator: The student uses the teacher denoiser’s architecture but generates each forecast in one step from historical states and Gaussian noise.The generator produces ˆx = Gϕ(h, σKϵ).
- Supervised distillation: Standard teacher-only matching can transfer both forecast skill and errors, so the method adds each paired ground-truth future state as a supervised target.
- Supervised distillation: The training objective balances a supervised loss with a student–teacher distribution discrepancy using coefficient β.
- Distribution matching: The distributional term transfers teacher knowledge and regularizes training, while sample-based matching avoids intractable KL divergence and unstable auxiliary density estimation.
- Training and inference: Algorithm 1 summarizes the training and inference procedures, with additional sample-count and loss-weight settings specified in Appendix C.1.2.
- Training and inference: Training draws student samples and estimates the distillation objective with an unbiased sample-based estimator before updating the student parameters.
- Supervised loss: The supervised loss uses ensemble-mean error because supervising every member against one observation may reduce diversity.This improves mean accuracy without directly penalizing ensemble spread.
- Energy-distance distillation: The energy score provides a strictly proper multivariate scoring rule whose expected value under the target distribution corresponds to squared energy distance.
4 Experiment
The method is evaluated on global probabilistic forecasting and typhoon-track prediction against GENCAST and distillation baselines. Results assess forecast skill, ensemble calibration, computational efficiency, and extreme-event performance.
- Global Weather Forecasting Skill: Forecasts use 25-member ensembles and compare the student with GENCAST, CM, DMD, GRAPHCAST, and ENS (WB2).Global forecasts cover 10 days from 12 initial conditions; ENS (WB2) is a 50-member ECMWF IFS ensemble.
- Global Weather Forecasting Skill: Across variables and lead times, the student achieves the best RMSE among distilled models and closely follows the teacher in CRPS.CM and DMD are competitive at short lead times, but their CRPS degrades more rapidly.
- Global Weather Forecasting Skill: The student better preserves teacher calibration and forecast diversity, while CM and DMD produce underdispersed ensembles.Spread–skill ratios closer to one indicate better ensemble calibration.
- Extreme-Event Forecasting: Typhoon Nanmadol forecasts preserve the storm trajectory and predict a stronger cyclone core than the teacher.Tracks are derived from each model’s fields and compared with the NOAA IBTrACS best track.
- Extreme-Event Forecasting: The student achieves the lowest position CRPS and deterministic track error among probabilistic models, particularly at longer lead times.Peak-intensity forecasts report mean sea-level pressure and 10-m winds alongside cyclone centers.
5 Conclusion
The paper positions GENCAST as a conditional diffusion teacher for autoregressive global weather ensembles and concludes that its supervised distillation produces a substantially cheaper single-pass student. The teacher uses spherical, weighted modeling and iterative PF-ODE sampling, while the student reduces each autoregressive step to one evaluation.
- Conclusion: The paper concludes that its single-pass student outperforms CM and DMD while matching GENCAST with 1 rather than 39 NFEs per autoregressive step.It also reports preserved ensemble spread and skillful Typhoon Nanmadol prediction.
- Diffusion Formulation: GENCAST generates 15-day global ensemble forecasts at 12-hour intervals using conditional diffusion and autoregressive atmospheric-state modeling.Its future-state distribution is factorized under a second-order Markov assumption.
- Diffusion Formulation: GENCAST predicts normalized residual states and reconstructs physical states using per-variable and per-level scaling, except for directly predicted accumulated precipitation.This stabilizes learning across atmospheric variables with different numerical scales.
- Diffusion Formulation: The diffusion process injects isotropic Gaussian noise on the sphere before projection to the grid and conditions denoising on the previous two atmospheric states.Noise-dependent EDM preconditioning supports denoising across σ ∈ [0.03, 80].
- Architecture: GENCAST uses a spherical graph-transformer architecture operating on a compressed icosahedral mesh to process high-dimensional global fields.The mesh contains 41,162 nodes and the processor has 16 transformer blocks with hidden dimension 512 and 4-head attention.
- Training Objective: Its training objective is a spatially and variable-weighted mean squared error over 40 years of ECMWF reanalysis data.Latitude-dependent area weights address equiangular-grid distortion and pole singularities, while variable-level weights balance gradient magnitudes.
C.1.1 Data
Experiments use ERA5 reanalysis for atmospheric states and verification targets on a 1° global grid. Training matches GENCAST’s 1979–2018 period, while evaluation uses selected 2022 initial dates.
- Data: ERA5 data are represented at 1° resolution on a 181 × 360 equiangular grid.The dataset includes six surface and six atmospheric variables at 13 pressure levels.
- Data: The student is trained on ERA5 from January 1, 1979, through December 31, 2018.This training period matches GENCAST’s training period.
- Data: Evaluation uses 12 initial dates: the 1st and 15th of each month from July to December 2022.
C.1.2 Training details
The 1° experiments use fixed optimization settings and a common tropical-cyclone tracking protocol. All models’ forecast fields are processed with the same tracker and compared against IBTRACS records.
- Training details: Training uses batch size 4 for 20k AdamW steps on 4 H100 GPUs with 96 GB each.
- Training details: The student and teacher each use M = N = 24 samples, with loss-balancing coefficient β = 2 after rescaling the loss terms.
- Training details: Optimizers use Optax, gradient clipping before AdamW updates, and settings shared by the student models and DMD critic network.The optimizer hyperparameters are summarized in Table 1.
- Tropical cyclone tracking: Tropical-cyclone forecasts use TEMPESTEXTREMES with identical detection criteria and temporal-linking rules across models.Forecast tracks are matched to IBTRACS best-track records for evaluation.
C.1.3 Metrics
The evaluation uses area-weighted RMSE, CRPS, and spread–skill ratio to assess forecast accuracy, probabilistic quality, and ensemble dispersion. Baseline distillation methods use trajectory or distribution matching, while the student generates forecasts in one evaluation.
- Metrics: Area-weighted RMSE, CRPS, and spread–skill ratio evaluate deterministic accuracy, probabilistic accuracy, and ensemble calibration, respectively.RMSE uses ensemble means for probabilistic models; CRPS assesses calibration and sharpness, while SSR summarizes dispersion relative to skill.
- Metrics: CRPS decomposes into member-wise skill and ensemble spread, with SSR = Spread/Skill.SSR near one indicates calibrated dispersion; values below or above one indicate underdispersion or overdispersion.
- Consistency-model Distillation: Consistency distillation trains the student on paired intermediate states from teacher reverse trajectories across a discrete noise schedule.The teacher is sampled independently for each forecast context, and the student is trained for consistency across adjacent noise levels.
- Consistency-model Distillation: At inference, the consistency-distilled student maps one noisy input directly to a forecast sample in one network evaluation.This reduces sampling to a single-step generator.
- Distribution Matching Distillation: DMD matches teacher and student score estimates using an auxiliary denoiser, then discards that auxiliary network at inference.The generator combines distribution-matching and regression losses to teacher samples, while forecast generation remains one evaluation.
C.3 Additional Experimental Results
Additional experiments compare the student with the teacher and alternative methods across variables, spatial fields, ensemble statistics, and typhoon-track forecasts. The student broadly preserves skill and achieves strong track-based probabilistic performance, while figures organize these comparisons by lead time and forecast quantity.
- Full-variable comparison: The student matches or improves upon the teacher for most variables at short-to-medium lead times, with degradation concentrated at longer leads.Figure 5 reports relative RMSE and CRPS across surface variables, pressure levels, and lead times.
- Typhoon-track errors: The student reduces the teacher’s Nanmadol track-position error over most of the forecast and remains comparable to CM.DMD diverges substantially after 16 September in the reported comparison.
- Typhoon-track errors: The student obtains the lowest position CRPS among probabilistic models for Nanmadol, combining track accuracy and ensemble dispersion.Figure 4 shows position error above and ensemble position CRPS below, computed as a multivariate energy score over track locations.
- Spatial forecast errors: Spatial error comparisons cover MSLP, Q700, T850, Z500, U10M, and V10M at 1-, 3-, 5-, and 10-day leads against ERA5.The figure compares GRAPHCAST, the GENCAST teacher, the student, CM, and DMD.
- Ensemble statistics: Spatial ensemble means and variances are compared across lead times for Z500, U10M, V10M, MSLP, Q700, and T850.Rows compare the multi-step teacher, distilled student, CM, and DMD.
- Efficiency: Table 2 reports wall-clock time per autoregressive step on one NVIDIA GH200 for the listed ensemble sizes.Values are means ± standard deviations over 10 runs, excluding JIT compilation.