Controllable Probabilistic Forecasting with Stochastic Decomposition Layers

2025-12-21

AI Summary

Existing CRPS-oriented AI-based weather forecasting ensembles predominantly rely on global conditional normalization for noise injection, entailing high training costs and limited physical interpretability. To address this, we propose a Multi-Scale Stochastic Decomposition Layer (SDL), embedded at multiple decoder levels to inject physically interpretable, hierarchical noise—enabling latent-driven modulation, pixel-wise perturbation, and channel-wise scaling—thereby facilitating posterior uncertainty calibration and cross-scale uncertainty disentanglement. Inspired by StyleGAN, our approach integrates conditional modulation with transfer learning, achieving CRPS optimization within the WXFormer framework. Experiments demonstrate a 98% reduction in training cost relative to conventional methods, with a single-member model occupying only 5 MB. Ensemble dispersion–skill ratios approach unity, and rank histogram uniformity matches that of the operational IFS-ENS. Notably, our method is the first to empirically reveal the hierarchical structure of uncertainty across large- and mesoscales.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

AI weather prediction ensembles with latent noise injection and optimized with the continuous ranked probability score (CRPS) have produced both accurate and well-calibrated predictions with far less computational cost compared with diffusion-based methods. However, current CRPS ensemble approaches vary in their training strategies and noise injection mechanisms, with most injecting noise globally throughout the network via conditional normalization. This structure increases training expense and limits the physical interpretability of the stochastic perturbations. We introduce Stochastic Decomposition Layers (SDL) for converting deterministic machine learning weather models into probabilistic ensemble systems. Adapted from StyleGAN's hierarchical noise injection, SDL applies learned perturbations at three decoder scales through latent-driven modulation, per-pixel noise, and channel scaling. When applied to WXFormer via transfer learning, SDL requires less than 2% of the computational cost needed to train the baseline model. Each ensemble member is generated from a compact latent tensor (5 MB), enabling perfect reproducibility and post-inference spread adjustment through latent rescaling. Evaluation on 2022 ERA5 reanalysis shows ensembles with spread-skill ratios approaching unity and rank histograms that progressively flatten toward uniformity through medium-range forecasts, achieving calibration competitive with operational IFS-ENS. Multi-scale experiments reveal hierarchical uncertainty: coarse layers modulate synoptic patterns while fine layers control mesoscale variability. The explicit latent parameterization provides interpretable uncertainty quantification for operational forecasting and climate applications.

Problem

• Develops Stochastic Decomposition Layers for probabilistic weather forecasting
• Converts deterministic ML models into ensemble systems with interpretable uncertainty
• Reduces computational cost and improves calibration in medium-range forecasts

Innovation

• Hierarchical noise injection at decoder scales
• Latent-driven modulation with per-pixel noise
• Compact latent tensor for reproducible ensembles