🤖 machine learning

Timestep-Conditioned Transformers for Global Weather Forecasting

The paper introduces GEM-3, a lightweight, timestep-conditioned transformer model that resolves the trade-off between short-range resolution and long-range error accumulation in weather forecasting by enabling flexible, multi-timestep inference with a single set of trained weights.

Original authors: Sam Levang, Fran Bartolic, Ty Dickinson, Chase Dwelle, Paulius Rauba, Viktor Cikojevic

Published 2026-08-07

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Sam Levang, Fran Bartolic, Ty Dickinson, Chase Dwelle, Paulius Rauba, Viktor Cikojevic

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the future of a giant, swirling ocean of air that covers the entire planet. This is the world of weather forecasting, a field where scientists use massive computers to simulate how the atmosphere moves. For decades, the best tools for this job have been complex physics equations, but recently, a new kind of "smart" computer program—powered by machine learning—has started to beat them. These programs learn from history to guess what the weather will do next. However, there is a tricky puzzle they face: how fast should the computer take a "step" into the future? If it takes tiny, quick steps (like every hour), it can see the details of a sudden storm forming, but it might get tired and make mistakes as it tries to walk a long distance. If it takes giant, slow steps (like every day), it can walk far without tripping, but it misses all the exciting details of the journey.

This paper introduces a clever new weather model called GEM-3 that solves this puzzle. Instead of forcing the computer to choose between being a fast, detailed sprinter or a slow, steady marathon runner, GEM-3 is like a magical chameleon that can change its stride length on the fly. The researchers found that by teaching the model to understand how much time it is supposed to skip in a single step, they could use the exact same brain (the same set of computer weights) to predict the weather for the next hour or the next month. They discovered that mixing different step sizes during training makes the model more stable, and they proved that this flexible approach works better than older models that are stuck with just one speed.

The Great Step-Size Dilemma

To understand why this is a big deal, picture a video game character trying to cross a vast, stormy ocean. If the character takes tiny, one-inch steps, they can carefully navigate every single wave and rock. But if they have to cross a thousand miles, they will take millions of steps. With every single step, there is a tiny chance of a tiny mistake. If you take a million steps, those tiny mistakes pile up, and the character might end up in the wrong country entirely. This is the problem with "fine" weather models: they are great for the next few days but get messy over long periods.

On the other hand, imagine the character takes giant, 100-mile leaps. They cross the ocean in just ten jumps. Because they take so few steps, they don't have many chances to make a mistake. But here's the catch: they can't see the small waves or the specific rocks in between. They might miss a sudden storm or a cool breeze because they are too busy jumping over the whole day at once. This is the problem with "coarse" models: they are stable for long trips but miss the short-term details.

For years, scientists had to pick one or the other. They built one model for short-term, high-detail forecasts and a different model for long-term, stable forecasts. It was like having a sports car for city driving and a tank for off-roading, but you had to buy two separate vehicles.

The Magic Chameleon: GEM-3

The authors of this paper, working at a company called Salient, built a new model named GEM-3. Think of GEM-3 not as a single vehicle, but as a magical shapeshifter. The secret sauce is a feature called timestep conditioning.

In the old models, the computer was trained to only know how to take one specific step size. If you wanted it to predict the weather for 24 hours later, it had to take 24 one-hour steps, or one 24-hour step, but it couldn't do both. GEM-3 is different. The researchers taught the model to listen to a "whisper" (a number) that tells it, "Hey, for this specific prediction, take a 6-hour step," or "No, take a 24-hour step."

The model has a single brain, but it can instantly reconfigure itself. If you ask for a forecast for tomorrow morning, it can switch to a fast, detailed mode to catch the morning breeze. If you ask for a forecast for next month, it can switch to a slow, steady mode to ensure it doesn't drift off course. This means you don't need two different models; you just need one flexible one that can adapt to the job.

The Secret Sauce: Training Like a Gymnast

How did they teach this model to be so flexible? They used a trick called mixed-timestep training. Imagine a gymnast training for the Olympics. Instead of only practicing on the balance beam (one specific skill), they practice on the beam, the floor, and the uneven bars all in the same session.

The researchers trained GEM-3 by showing it weather data and asking it to predict the future using random step sizes—sometimes 1 hour, sometimes 6, sometimes 24. By forcing the model to handle all these different speeds at once, it learned a more robust way to understand the atmosphere. The paper suggests that this "mixed" training acts like a stabilizer. It turns out that a model trained on many different speeds is actually better at taking short steps than a model trained only on short steps. It's as if the gymnast, by learning to balance on many different surfaces, becomes steadier on the beam than someone who only ever practiced on the beam.

The Anomaly Trick: Removing the Noise

There was another hurdle. Weather has a rhythm: it's warmer in summer and colder in winter, and warmer during the day and cooler at night. If a model tries to predict the exact temperature, it has to constantly remember these huge, predictable cycles. This can be distracting and cause the model to get "drunk" on the long-term trends, losing its way over time.

To fix this, the authors used a technique called anomaly-space modeling. Instead of asking the model, "What will the temperature be?", they asked, "How much different will the temperature be from the usual average for this time of year?"

Imagine you are trying to guess how much a friend will eat at dinner. Instead of guessing the total weight of the food (which changes wildly between a salad and a feast), you guess how much more or less they will eat than their usual portion. By focusing on the "surprise" (the anomaly) rather than the total, the model stays focused on the actual weather changes and doesn't get confused by the predictable seasons. The paper shows that this helps the model stay on track for much longer, especially when making long-range predictions.

What They Found

The results are quite impressive. When they tested GEM-3 against the current world champions (like the European weather service's models and other top AI weather tools), GEM-3 held its own or even won.

Short-term: It predicts the next few days with high detail, catching the nuances of daily weather.
Long-term: It stays stable for up to 46 days, beating many other models in keeping the forecast accurate without drifting off.
Flexibility: The most exciting finding is that the model can switch between a 6-hour step and a 24-hour step at the very last second, depending on what the user needs.

However, the paper is careful to note that this isn't a magic wand for everything. They found that if you try to make the steps too huge (like 48 hours), the model gets confused because the weather changes too much in that time for a single step to handle. Similarly, if you try to make the steps too tiny (like 1 hour) for a very long forecast, the tiny mistakes still pile up, even with the new tricks. But within a "sweet spot" of 1 to 24 hours, the model works beautifully.

The Bottom Line

GEM-3 is a practical breakthrough because it stops forcing meteorologists to choose between detail and stability. It suggests that the future of weather forecasting isn't about building bigger, specialized computers for every job, but about building smarter, more adaptable ones that can change their stride to fit the terrain. By teaching a single model to speak many "time languages," the researchers have created a tool that is not only more accurate but also more useful for real-world decisions, from planning a picnic to preparing for a storm weeks in advance.

Technical Summary: Timestep-Conditioned Transformers for Global Weather Forecasting (GEM-3)

1. Problem Statement

Existing machine learning (ML) weather forecasting models face a fundamental architectural trade-off regarding autoregressive timesteps ( $\Delta t$ ).

Fine Timesteps (e.g., 1–6 hours): These resolve sub-daily atmospheric dynamics and diurnal cycles effectively. However, in an autoregressive setting, they require deep rollouts to reach medium or extended ranges, leading to significant error accumulation and drift.
Coarse Timesteps (e.g., 24 hours): These reduce the number of autoregressive steps required for long horizons, mitigating error accumulation. However, they fail to capture short-term variability and sub-daily trajectories, limiting their utility for short-range forecasting.

Current approaches typically force practitioners to choose between models optimized for short leads, models optimized for long leads, or post-processing methods that introduce structural errors. There is a lack of a single model capable of dynamically adapting its temporal resolution at inference time to balance predictability and usability across a broad forecast horizon.

2. Methodology

2.1 Multi-Timestep Inference Architecture

The core contribution of GEM-3 is an architectural mechanism that treats the forecasting timestep $\Delta t$ as an explicit conditioning variable rather than a fixed hyperparameter.

Timestep Conditioning: The model accepts $\Delta t$ as an input alongside the atmospheric state $X_t$ and conditioning fields $C$ . This allows a single set of trained weights to dynamically switch temporal resolutions.
Joint Temporal and Aleatoric Modulation: The model utilizes an Adaptive Layer Normalization (AdaLN-Zero) pathway. Both the timestep $\Delta t$ $Δ t$ and the stochastic noise vector $z$ $z$ (used for ensemble generation) are routed through this pathway.
- $\Delta t$ is Fourier-embedded and projected to produce shift, scale, and gate parameters.
- $z$ is processed similarly to modulate the network for ensemble diversity.
- Sequential Modulation: The timestep modulation is applied first to define the primary transition operator, followed by the noise modulation as a second-order perturbation. This allows the network to learn a transition kernel that adapts to the required scale of evolution while maintaining stochastic diversity.

2.2 Anomaly-Space Modeling

To address the issue of mass, heat, and moisture budget imbalances that cause drift over long horizons, GEM-3 employs anomaly-space modeling for a subset of variables.

Transformation: Selected variables (temperature, geopotential, and pressure) are transformed into climatological anomalies by subtracting a pre-computed, low-pass filtered climatology (indexed by hour-of-year) before entering the transformer backbone.
Reconstruction: The climatology is added back after the model output.
Rationale: This shifts the learning target to a zero-centered residual state space, which is more stable for autoregressive rollouts. Variables with hard physical bounds or non-linear connections to solar forcing (e.g., precipitation, wind) are kept in state space.

2.3 Model Architecture and Training

Backbone: GEM-3 is a lightweight Neighborhood Attention Transformer (NAT) with $\sim$ 134M parameters, operating on an equirectangular grid. It utilizes SwiGLU MLPs, QK-normalization, and 2D rotary positional embeddings.
Mixed-Timestep Training: The model is trained from scratch on a mixed distribution of timesteps (e.g., 1, 3, 6, 12, 24 hours). To prevent gradient imbalance caused by the increasing magnitude of uncertainty with larger $\Delta t$ , a timestep-dependent loss scaling factor is applied.
Loss Function: The objective combines a per-gridpoint fair CRPS (Continuous Ranked Probability Score) and a spectral CRPS computed in the spherical harmonic domain to ensure fidelity across spatial scales.
Optimization: A hybrid strategy using Muon for most parameters and AdamW for a subset found to be unstable under Muon.

3. Key Contributions

Configurable Multi-Timestep Inference: Introduction of a mechanism allowing a single model to adaptively change its autoregressive timestep at inference time, eliminating the need for separate models for different lead times.
Mixed-Timestep Training Stability: Demonstration that training on a curriculum of mixed timesteps acts as a regularizer, improving the rollout stability of short-step forecasts compared to timestep-specialist models.
Anomaly-Space Integration: Successful application of anomaly-space modeling for key thermodynamic variables to reduce long-horizon drift without sacrificing short-range fidelity.
Operational System (GEM-3): A practical forecasting system that couples near-state-of-the-art medium-range probabilistic skill with stable extended-range rollouts and native sub-timestep diagnostic outputs (e.g., min/max temperature).

4. Experimental Results

4.1 Medium-Range Performance

GEM-3 achieves near-state-of-the-art CRPS scores, outperforming the ECMWF ENS across all analyzed variables and lead times.
It is broadly comparable to top ML models like AIFS-ENS on core dynamical benchmarks. However, the Google FGN model outperforms GEM-3 on specific variables including 500 hPa geopotential ( $z_{500}$ ), mean sea-level pressure (mslp), and precipitation. Conversely, GEM-3 outperforms FGN on 2m temperature.
GEM-3 shows significant advantages in native sub-timestep diagnostics (minimum and maximum temperature) compared to models that must estimate these from snapshots.

4.2 Extended-Range and Rollout Stability

Hybrid Schedule: The operational model utilizes a hybrid schedule: 6-hour timesteps for the first 14 days (to resolve diurnal cycles) and 24-hour timesteps for the extended range (to minimize error accumulation).
Performance: GEM-3 generally outperforms the ECMWF extended-range ensemble out to 46 days and maintains positive skill relative to climatology up to 126 days.
Ablation Findings:
- Generalist vs. Specialist: A generalist model trained on mixed timesteps (1–24h) exhibits lower error accumulation over multi-day rollouts compared to timestep-specialists, particularly at 1h and 6h steps, despite slightly higher initial step errors.
- Timestep Limits: The useful range for timestep conditioning appears bounded. While 1–24h works well, 48h steps degrade performance significantly, suggesting the single learned operator cannot adequately represent the highly non-linear evolution of a 48h transition.
- Drift vs. Spectral Fidelity: There is a trade-off where the shared timestep-conditioned operator effectively suppresses large-scale mean-state drift (benefiting fine-step rollouts) but injects a small excess of high-wavenumber power (spectral noise) compared to specialists, particularly at coarser timesteps.

5. Significance and Claims

The paper claims that GEM-3 represents a practical shift in ML weather forecasting design. By decoupling the inference schedule from the model architecture, it resolves the traditional trade-off between short-range detail and long-range stability.

Operational Viability: The system demonstrates that a single parameter set can serve both high-resolution short-range needs and stable extended-range needs, reducing the operational complexity of maintaining multiple specialized models.
Regularization Effect: The study highlights that mixed-timestep training itself acts as a stabilizer, improving the robustness of short-step forecasts beyond what is achieved by timestep-specialist training.
Efficiency: The model achieves these results with 134M parameters, significantly smaller than some competing models (e.g., Atlas at 4.3B), while maintaining competitive or superior skill.

The authors conclude that within an appropriate timestep range (roughly 3–24 hours), multi-timestep inference is a viable and superior design for probabilistic weather forecasting, allowing the choice of timestep to be an operational decision based on cost, usability, and required horizon rather than a fixed architectural constraint.