PaperlayerView exact revision on arXiv

Source-linked AI summary

Demystifying Data-Driven Probabilistic Medium-Range Weather Forecasting

Jean Kossaifi, Nikola Kovachki, Morteza Mardani, Daniel Leibovici, Suman Ravuri, Ira Shokar, Edoardo Calvello, Mohammad Shoaib Abbas, Peter Harrington, Ashay Subramaniam, Noah Brenowitz, Boris Bonev, Wonmin Byeon, Karsten Kreis, Dale Durran, Arash Vahdat, Mike Pritchard, Jan Kautz

arXiv:2601.18111v1cs.LGcs.AI

TL;DR

Probabilistic weather forecasting is fragmented across specialized architectures and training methods, motivating a simpler general framework. The paper introduces ATLAS, combining directly downsampled latent modeling with history-conditioned high-resolution decoding and supporting multiple probabilistic estimators. ATLAS outperforms IFS and GenCast in the reported evaluations, with statistically significant aggregate improvements and reduced computational cost compared with GenCast.

  • Problem

    Probabilistic weather forecasting is fragmented across specialized architectures, objectives, and sampling strategies, creating a need to identify which design choices are genuinely necessary for state-of-the-art skill.

  • Method

    ATLAS combines a directly downsampled latent-space transformer with a history-conditioned local projector that decodes high-resolution residuals.

  • Results

    ATLAS outperforms both the operational IFS and GenCast, with statistically significant aggregate improvements and reduced computational cost compared with GenCast.

  • Takeaways & Limitations

    The results suggest that scaling general-purpose foundation models, rather than adding architectural complexity or tailored training recipes, can support high-fidelity medium-range probabilistic prediction.

  • Takeaways & Limitations

    Evaluation remains bounded by open ERA5 training data, a held-out year of 2020, selected phenomena, and a limited set of strong baselines.

Abstract

from arXiv · show

The recent revolution in data-driven methods for weather forecasting has lead to a fragmented landscape of complex, bespoke architectures and training strategies, obscuring the fundamental drivers of forecast accuracy. Here, we demonstrate that state-of-the-art probabilistic skill requires neither intricate architectural constraints nor specialized training heuristics. We introduce a scalable framework for learning multi-scale atmospheric dynamics by combining a directly downsampled latent space with a history-conditioned local projector that resolves high-resolution physics. We find that our framework design is robust to the choice of probabilistic estimator, seamlessly supporting stochastic interpolants, diffusion models, and CRPS-based ensemble training. Validated against the Integrated Forecasting System and the deep learning probabilistic model GenCast, our framework achieves statistically significant improvements on most of the variables. These results suggest scaling a general-purpose model is sufficient for state-of-the-art medium-range prediction, eliminating the need for tailored training recipes and proving effective across the full spectrum of probabilistic frameworks.

1. Introduction

Data-driven weather forecasting has progressed from matching numerical solvers to probabilistic prediction, but the field remains fragmented across specialized architectures and training methods. ATLAS addresses this fragmentation with standard scalable components and supports state-of-the-art probabilistic performance across multiple estimators.

  • Motivation: Data-driven models can match traditional solvers while generating forecasts several orders of magnitude faster, enabling massive ensembles and shifting attention toward probabilistic prediction.The probabilistic task targets uncertainty in partially observed, chaotic atmospheric systems.
  • Problem: The field uses fragmented architectural philosophies, including spherical-harmonic, convolutional, and graph-based inductive biases for atmospheric data.These approaches prioritize domain-specific constraints to represent spherical geometry and physical variables.
  • Problem: Probabilistic forecasting is also fragmented across CRPS objectives, spectral terms, diffusion, flow-matching, and consistency models, each with stated limitations in resolution, spectral fidelity, or inference speed.CRPS-based methods can underrepresent short-wavelength spectral power, while diffusion-based methods can be slow or resolution-limited.
  • Approach: ATLAS uses standard, scalable transformer components in compressed latent space to separate global-dynamics modeling from high-resolution synthesis.The framework intentionally removes domain-specific constraints while targeting state-of-the-art probabilistic performance.
  • Contributions: The proposed framework supports stochastic interpolants, diffusion models, and CRPS-based ensemble training as independent probabilistic estimation methods.This method-agnostic design is presented as a way to address the fragmented modeling landscape.

2. ATLAS: Medium-Range Latent Probabilistic Weather Forecasting

This section introduces ATLAS as a framework for medium-range probabilistic weather forecasting, combining latent-space modeling, transformer backbones, and three probabilistic generative methods.

  • ATLAS: Medium-Range Latent Probabilistic Weather Forecasting: ATLAS is introduced as a medium-range probabilistic weather-forecasting method built from latent-space modeling, backbone architectures, and three probabilistic generative methods.The section establishes the method components before describing their formulation and validation.

2.1. Problem Formulation

The paper models one-step conditional atmospheric-state distributions and repeatedly samples them to produce autoregressive forecasts. Under stationarity, the conditional law is treated as time-independent, although real weather may violate this assumption.

  • Problem Formulation: The forecasting task is to model the conditional distribution of the next atmospheric state given the current state from finite joint-distribution realizations.The atmospheric state is represented as a stationary, discrete-time stochastic process.
  • Evaluation: The supplied scorecard compares ATLAS-SI with IFS-ENS and GenCast using ensemble-mean RMSE and ensemble CRPS for fifteen-day, 56-member forecasts.The comparison averages results over ERA5 initial conditions from 2020.
  • Probabilistic Modeling: The three probabilistic methods define conditional transport maps that push known distributions to the next-state conditional distribution, enabling sampling by evaluating those maps.The transport maps may be defined through stochastic-process flow maps or other formulations described later.
  • Unrolling: Forecasts are generated by unrolling approximate one-step samples autoregressively, but approximation errors can compound through distribution shift at later steps.The authors report stable forecasts up to 15 days, corresponding to 60 autoregressive steps.

2.2. Latent Space Modeling

ATLAS reduces the dimensionality of atmospheric fields through direct bilinear downsampling, models latent residuals using history, and decodes them into high-resolution forecasts with a conditioned local projector.

  • Choice of latent space: The atmospheric dataset contains 75 variables on a quarter-degree 721 × 1440 grid, creating a state dimension of 77,868,000 that is costly to model directly.The method therefore adopts a latent-space approach with encoding, probabilistic inference, and decoding.
  • Residual prediction and history: The approach encodes current and historical atmospheric states into latent representations, predicts next-step latent residuals probabilistically, and decodes them into original space.The probabilistic model uses global attention, while the decoder uses local attention and conditions on the high-resolution initial state.
  • Choice of latent space: ATLAS uses a 16× compression scheme that bilinearly interpolates each variable to a one-degree 181 × 360 grid.At six-hour prediction intervals, scales below approximately seventy kilometers are treated as deterministically unpredictable.
  • Residual prediction and history: Predicting temporal residuals and including one additional historical state improves model robustness and stability at longer lead times.The decoder also uses the initial state to improve reconstruction of high-resolution details and physical dynamics.
  • Residual prediction and history: The decoder reconstructs high-resolution residuals from latent predictions, with conditioning on the initial state helping capture fine-scale details and physical dynamics.The decoded residual is added to the initial state to approximate the next atmospheric state.

2.3. Transformer Architectures as the backbone of ATLAS

ATLAS uses Diffusion Transformer architectures for both latent prediction and decoding. The decoder combines latent residuals with full-resolution initial conditions, while predictive models process noisy fields and historical states through patched token sequences.

  • Transformer backbone: ATLAS adapts the Diffusion Transformer for its decoder and predictive architecture.The architecture was originally popularized for image generation and is tailored here to parameterize both components.
  • Decoder architecture: The decoder maps a 181 × 360 latent residual and full-resolution 721×1440 initial condition to a full-resolution residual prediction.The decoder inputs contain 75 and 79 channels respectively, and its output contains 75 channels.
  • Decoder architecture: Decoder inputs are patched and concatenated into a 181 × 360 representation before positional encoding and DiT processing.The initial condition uses a learned strided convolution, while the residual uses a 1 × 1 convolution to expand channels.
  • Predictive architecture: Predictive models parameterizing b̂, ŝ, or f̂ ingest a noisy field together with two historical latent states.All three inputs are 75 × 181 × 360 spatial fields before patching and concatenation.
  • Predictive architecture: Using 2 × 3 patches provides a compromise between predictive accuracy, computational cost, memory use, and rollout stability.Smaller patches improved single-step performance but increased resource demands, while very small patches produced overfitting and unstable rollouts.
  • Predictive architecture: The patched predictive representation has sequence length 10,920 and is processed by DiT blocks with global attention.The representation is conditioned on the time input t and uses sine-cosine embeddings.

2.4. Three Approaches to Probabilistic Modeling in Latent Space

ATLAS performs probabilistic modeling in latent space, where three estimators approximate the conditional distribution through stochastic transport, diffusion, or direct generative mapping. A spectral CRPS regularizer addresses spatial-structure limitations while retaining expressive stochastic conditioning.

  • Latent-space formulation: The probabilistic methods sample a latent conditional distribution and decode samples into the conditional distribution of the full-resolution weather state.The history latent z−1 is supplied as an additional input to parameterized maps even when omitted from notation.
  • Latent-space formulation: ATLAS represents dataset states in a directly downsampled latent space before probabilistic inference and decoding.The dimensionality symbol d denotes the latent encoding dimension in this formulation.
  • Stochastic Interpolants: Stochastic Interpolants learn a stochastic process connecting the initial latent state to the conditional residual distribution, enabling approximate conditional sampling.The process starts at z0 and ends at r1, while a learned drift approximates the governing stochastic differential equation.
  • Diffusion models: Diffusion modeling estimates the score of a forward-noised process and uses a reverse-time stochastic differential equation to sample conditional latent residuals.The reverse process is discretized with a predictor-corrector method.
  • CRPS-based models: The CRPS-based model directly maps an initial latent state and Gaussian noise to a sample from the conditional residual distribution.Training minimizes a generalized CRPS with an ℓ1-norm chosen for robustness despite not being strictly proper for general distributions.
  • CRPS-based models: Aggregated CRPS does not explicitly enforce joint spatial structure, and CRPS-based models can under-represent high-frequency energy.The paper identifies possible pointwise accuracy with inconsistent spatial correlations as a consequence of this limitation.
  • Spectral regularization: ATLAS adds spectral regularization by computing CRPS on spherical-harmonic coefficient magnitudes, ensuring phase invariance.The spherical harmonic transform projects spatial fields onto orthonormal basis functions before regularization.
  • Spectral regularization: Spectral CRPS regularization supports high-dimensional noise and full expressive stochastic conditioning without biased–fair CRPS mixtures or low-dimensional noise constraints.This contrasts with prior strategies addressing under-constrained ensemble variability through estimator mixtures or restricted latent noise.

3. Experimental Setting

The experiments train and evaluate ATLAS on high-resolution ERA5 weather data using common optimization settings and compare it with IFS-ENS and GenCast. Evaluation covers probabilistic forecasting variables across the 2020 test year.

  • Protocol: The study evaluates models using standard metrics and comparisons against selected baselines.The evaluation protocol and metrics are detailed for quantitative results.
  • Dataset: Training uses ERA5 data at 721×1440 resolution sampled every 6 hours from 1980–2019, with 2020 reserved for testing.The dataset contains 75 input and output variables, including surface, atmospheric, sea-surface-temperature, and precipitation fields.
  • Training: Models are trained with elastic averaged stochastic gradient descent across 32 80GB A100 or H100 GPUs.The total batch size is 32, and the learning rate is 1.28 × 10−4.
  • Training: ATLAS-SI and ATLAS-EDM use 3 training cycles, whereas ATLAS-CRPS uses 6 cycles.Each cycle includes cosine-decayed optimization after linear warm-up, followed by a learning-rate reset for the next cycle.
  • Decoder training: All models share a decoder trained for 3 cycles with additional date- and surface-dependent fields to address potential non-stationarity.The decoder minimizes the ℓ1-norm against the residual field and receives four additional normalized fields.
  • Model configuration: ATLAS-SI and ATLAS-EDM use a 2.4B-parameter DiT architecture with shared predictive design.The architecture uses 12 DiT blocks, 13 attention heads, and a total embedding dimension of 3,328.
  • Evaluation: Evaluation compares ATLAS with the ECMWF IFS ENS physics-based ensemble and the data-driven GenCast baseline.The focus is on globally aggregated probabilistic metrics for the 2020 validation set.
  • Inference cost: A single A100 generates one timestep in 94s for ATLAS-SI, 88s for ATLAS-EDM, and 3.3s for ATLAS-CRPS.The corresponding GenCast time is 140s for a 12-hour timestep, compared with 6-hour ATLAS predictions.

3.2. Metrics

The evaluation reports probabilistic accuracy, ensemble-mean accuracy, calibration, and statistical significance across variables, lead times, locations, and initialization dates. Metrics use area-weighted global aggregation and paired comparisons with GenCast.

  • Evaluation metrics: Metrics are reported by variable, level, and lead time on a 721 × 1440 equiangular grid.Locations are indexed geographically and aggregated using normalized area weights.
  • Ensemble setup: The comparison uses 56 ensemble members for all compared baselines, making the non-standard ensemble size immaterial to intercomparison.The authors also report baseline results with 56 members.
  • Accuracy metrics: CRPS and ensemble mean RMSE assess pointwise probabilistic and mean-prediction accuracy.The fair CRPS is averaged across 28 initialization times and is unbiased with respect to ensemble size in the infinite-sample limit.
  • Calibration: The spread-skill ratio assesses calibration, with SSR = 1 indicating a properly calibrated model.Values below 1 indicate under-dispersion and values above 1 indicate over-dispersion.
  • Statistical significance: Paired t-tests compare ATLAS verification scores with GenCast across identical initialization dates for each lead time.The test uses paired score differences for ERMSE or CRPS and accounts for sampling error across dates.
  • Statistical significance: Results are statistically significant when the corresponding two-sided p-value is below 0.05.Under the null hypothesis, the test statistic follows a Student-t distribution with T−1 degrees of freedom.

4. Results

ATLAS is evaluated against the ECMWF Integrated Forecasting System and GenCast using probabilistic and ensemble-mean skill metrics across multiple model variants. It shows broad gains against IFS-ENS and variant-dependent advantages or trade-offs relative to GenCast.

  • Benchmark design: The evaluation benchmarks all ATLAS variants against the ECMWF Integrated Forecasting System and GenCast.The comparisons use RMSE and CRPS across the SI, EDM, and CRPS variants.
  • Comparison with IFS-ENS: All three ATLAS variants improve substantially over IFS-ENS at early lead times, especially for surface variables.The strongest relative gains occur for zonal and meridional surface winds.
  • Comparison with IFS-ENS: ATLAS retains an advantage over IFS-ENS through day 15 for most state variables, while z500 converges toward baseline skill after approximately 12 days.The persistent gains include q850, t850, and surface winds.
  • Comparison with GenCast: ATLAS-SI outperforms GenCast on CRPS and ensemble-mean RMSE to approximately seven days, while longer-range differences are statistically indistinguishable.ATLAS-EDM is nearly equivalent except for z500, whereas ATLAS-CRPS gains cluster around days 8–10 despite early underperformance.

15-day rollout comparison

ATLAS is compared with IFS and GenCast using probabilistic error growth, case studies, and tropical-cyclone tracking. It maintains strong skill across most variables and produces competitive or more accurate cyclone tracks in several examples.

  • 15-day rollout comparison: ATLAS-SI tends to achieve lower ensemble RMSE and CRPS than IFS and GenCast for most variables during the initial forecast week.ATLAS-SI is slightly more under-dispersive than GenCast, which is often over-dispersive at longer lead times.
  • 15-day rollout comparison: CRPS training leaves ATLAS-CRPS under-dispersive because forecast-error minimization appears to dominate the weaker signal encouraging ensemble spread.The resulting optimization balance prioritizes forecast accuracy over ensemble dispersion.
  • Storm Dennis case study: ATLAS-SI’s Storm Dennis ensemble spectra agree closely with ERA5 across lead times, with modest variance under-representation only at the finest scales.At horizontal wavenumbers greater than 200, spectral bias shows no lead-time-dependent trend.
  • Tropical cyclone evaluation: Across 46 tropical storms, ATLAS-SI improves ensemble-mean track error over GenCast beyond 80 hours, while average intensity error matches GenCast through 144 hours.Both models have similar tracking and intensity performance overall, with ATLAS-SI slightly lagging in intensity beyond 144 hours.
  • Tropical cyclone evaluation: ATLAS-SI predicts Hurricane Eta’s eastward swing more frequently and accurately than GenCast, whose ensemble assigns higher probability to a westward trajectory.This case illustrates qualitatively different probabilistic track distributions between the models.
  • Tropical cyclone evaluation: For Tropical Storm Krovanh, ATLAS-SI tracks span the cyclone’s full lifetime, whereas GenCast tracks terminate before the storm reaches southern Vietnam and Cambodia.The comparison covers the first seven forecast days.

5. Validating Latent Space Modeling

The latent-space validation tests whether straightforward compression and decoding preserve forecast quality. Learned autoencoders can achieve low reconstruction error yet degrade long autoregressive rollouts, whereas a history-conditioned DiT decoder preserves high-resolution temporal information.

  • Learned latent: The learned autoencoder has competitive 6- and 12-hour accuracy but accumulates autoregressive rollout error faster, despite relatively low reconstruction errors.Z500 reconstruction error is relatively high, and longer rollout performance lags IFS ENS.
  • Motivation: Standard VAE-style latent modeling is poorly suited to heterogeneous atmospheric variables, and even low reconstruction error can produce less accurate forecasts.The paper therefore treats decoder quality as important for weather’s pixel-level prediction task.
  • Baseline autoencoding models: The predictive latent is a 64 × 195 × 390 state used to model the next latent state or the next latent residual.A second autoencoder models latent residuals in the residual-to-residual approach.
  • Decoder validation: Direct bilinear interpolation preserves latent structure but loses small-scale information, while bilinear upsampling produces very high reconstruction errors.These results motivate a learned decoder rather than trivial decoding.
  • Decoder validation: The DiT decoder uses the low-resolution residual and previous high-resolution state to predict the high-resolution residual, retaining high-resolution temporal consistency during inference.The predictive backbone operates only on low-resolution states, while the decoder preserves information from the previous timestep.

6. Conclusion

ATLAS challenges the need for complex, domain-specific architectures and training methods in probabilistic weather forecasting, instead emphasizing scalable general-purpose models. Across probabilistic estimators, it improves forecasting skill, while diffusion-based approaches introduce a near-term versus medium-range trade-off that remains open to mitigation.

  • ATLAS combines a latent transformer and local projector with stochastic interpolants, diffusion models, and CRPS ensemble training, outperforming IFS and GenCast.The framework is designed to remain robust across probabilistic estimation methods rather than relying on one specialized recipe.
  • Scaling general-purpose foundation models, rather than increasing architectural complexity, is identified as the path toward higher-fidelity forecasting.
  • Statistically significant improvements in aggregate metrics, extreme-event emulation, and tropical-cyclone tracking are achieved at reduced computational cost compared with GenCast.
  • The evaluation is limited by comparison against one open-data baseline and by validation focused on selected phenomena and the held-out year 2020.Further comparisons with AIFS and FGN, broader phenomena, and additional years are recommended.
  • Diffusion-based forecasting can involve a trade-off between short-range and medium-range skill, with near-term skill deliberately sacrificed to improve stability.Improved noise-schedule design and tighter overfitting control are proposed as possible mitigations.
Loading 2601.18111v1…