ArXiv

Timestep-Conditioned Transformers for Global Weather Forecasting

Authors
Sam Levang, Fran Bartolic, Ty Dickinson...
Categories
cs.LG, cs.OS
arXiv
https://arxiv.org/abs/2608.06241v1
PDF
https://arxiv.org/pdf/2608.06241v1

Brief

GEM-3 is a probabilistic global weather model that enables timestep-conditioned multi-timestep inference with a single set of trained weights, letting users trade off short-range (1–6 h) resolution versus reduced error accumulation of longer steps (e.g., 24 h). It uses a lightweight neighborhood-attention transformer (~134M parameters) on an equirectangular grid; mixed‑timestep training improves rollout stability vs timestep-specialists and yields near‑SOTA medium-range probabilistic skill and stable extended-range rollouts.

Why it matters

GEM-3 enables timestep-conditioned multi-timestep inference with a single set of trained weights, addressing the trade-off between short timesteps (examples given: 1–6 hours) and longer timesteps (example: 24 hours) at inference time.

Key details

  • Architecture: a lightweight neighborhood-attention transformer with ~134M parameters on an equirectangular grid, representing architectural advances beyond predecessor GEM-2.
  • Training and performance: mixed-timestep training improves rollout stability relative to timestep-specialist models; GEM-3 achieves near‑SOTA medium-range probabilistic skill, stable extended-range rollouts, efficient training/inference, and provides decision-relevant diagnostics.
Source evidence

Abstract

Existing machine-learning weather forecasting models rely on predetermined and fixed autoregressive timesteps. The choice of model timestep involves a fundamental trade-off: shorter timesteps (e.g. 1 to 6 hours) finely resolve atmospheric dynamics within the diurnal cycle but increase error accumulation for a given forecast horizon, while longer timesteps (e.g. 24 hours) reduce error accumulation but limit the usability of short-range forecasts where sub-daily predictability is high. In this work, we present GEM-3, a probabilistic global weather model that addresses this trade-off through explicit multi-timestep inference. With a single set of trained weights, the model timestep can be configured at inference time to balance predictability and usability across a broad forecast horizon. Additionally, we find that mixed-timestep training consistently improves rollout stability relative to timestep-specialist models. Under the hood, GEM-3 is a lightweight neighborhood-attention transformer with ~134M parameters on an equirectangular grid with a number of architectural advancements beyond its predecessor GEM-2. The result is a practical forecasting system that couples near-SOTA medium-range probabilistic skill, stable extended-range rollouts, efficient training and inference, and decision-relevant diagnostics.