Probabilistic Transformers for Joint Modeling of Global Weather Dynamics and Decision-Centric Variables

TL;DR

This work reframes global weather forecasting as a decision-centric problem by directly modeling diagnostic targets $Y_t$ that end users act upon, rather than solely predicting prognostic states $X_t$. The authors introduce GEM-2, a lightweight probabilistic transformer with periodic shifted-window attention that jointly emits $X_t$ and $Y_t$ on a global $0.25^\circ$ grid, trained with a strictly proper CRPS objective (including a spectral CRPS term) to ensure calibrated probabilistic forecasts. Empirically, GEM-2 delivers state-of-the-art decision-relevant skill, robust tail performance, subseasonal-to-seasonal stability, and strong relative economic value, while achieving large training-time speedups ($\sim$20–100x) and requiring modest compute compared with diffusion-based models. The approach reduces post-processing bias, supports end-user decision rules directly, and maintains stability across a broad range of architectural choices, making it a practical and impactful advance for operational forecasting and risk-informed decision-making.

Abstract

Weather forecasts sit upstream of high-stakes decisions in domains such as grid operations, aviation, agriculture, and emergency response. Yet forecast users often face a difficult trade-off. Many decision-relevant targets are functionals of the atmospheric state variables, such as extrema, accumulations, and threshold exceedances, rather than state variables themselves. As a result, users must estimate these targets via post-processing, which can be suboptimal and can introduce structural bias. The core issue is that decisions depend on distributions over these functionals that the model is not trained to learn directly. In this work, we introduce GEM-2, a probabilistic transformer that jointly learns global atmospheric dynamics alongside a suite of variables that users directly act upon. Using this training recipe, we show that a lightweight (~275M params) and computationally efficient (~20-100x training speedup relative to state-of-the-art) transformer trained on the CRPS objective can directly outperform operational numerical weather prediction (NWP) models and be competitive with ML models that rely on expensive multi-step diffusion processes or require bespoke multi-stage fine-tuning strategies. We further demonstrate state-of-the-art economic value metrics under decision-theoretic evaluation, stable convergence to climatology at S2S and seasonal timescales, and a surprising insensitivity to many commonly assumed architectural and training design choices.

Figures (31)

• Figure 1: Highlights of this work. (a) We show that modeling decision-centric variables natively eliminates structural bias which arises from post-hoc aggregation of snapshots; (b) We introduce a fine-tuning recipe that allows us to model variables of interest natively for different classes of diagnostic variables; (c) We obtain extreme efficiency, with $\sim$100x reduction in training cost and $\sim$10x reduction in inference costs relative to SOTA models; (d) Our predictions demonstrate state-of-the-art economic value for end-users of predictions of forecasts; (e) We obtain stability for free within our modeling paradigm with long-lead convergence to climatology. This stability is manifested as convergence toward climatological baseline in the long-lead domain Note: FGN-1 and FGN-4 refer to the 1-seed and 4-seed model of FGN, respectively.
• Figure 2: From weather dynamics modeling to decision-aligned forecasting. Left: Traditional forecasting frameworks propagate the full dynamical state forward in time by evolving prognostic variables that encode the system’s governing physics. Middle: Because these models provide only sparse state snapshots, downstream users often rely on post-processing to infer application-specific quantities. This workflow (i) ignores the structure of the desired targets during model training and (ii) introduces systematic bias when those targets cannot be reliably reconstructed from limited snapshots. Right: Decision-aligned forecasting integrates the evolution of weather states with the direct prediction of diagnostic targets—quantities of interest that do not influence the dynamics but are functionals of it. By coupling prognostic evolution with simultaneous diagnostic estimation, the approach yields the required outputs directly, eliminating dependence on post-hoc aggregation and reducing bias.
• Figure 3: Typology of diagnostic targets. Each class extends the forecasting system along a different structural axis—spatial, temporal, or aggregative—while leaving the underlying state transition unchanged. The primary mechanism of modeling weather states and optimizing the required user end-products at any spatial or temporal resolution jointly with the variables that drive the dynamics of the weather is the unifying component of all approaches.
• Figure 4: Backbone of GEM-2. Swin-style backbone with periodic shifted-window attention on a latitude--longitude grid.
• Figure 5: Backbone of GEM-2.1. Neighborhood Transformer with additional boundary conditions