🔬 physics

Scaling Storm-Resolving Atmospheric AI Simulation to the Entire Planet

The paper introduces STRATA, the first autoregressive AI emulator capable of generating stable, energy-efficient, global storm-resolving atmospheric simulations at 4.9-km resolution by leveraging a tile-based training strategy and novel architectural components to overcome the computational limitations of traditional physics models.

Original authors: Zeyuan Hu, Akshay Subramaniam, Noel Keen, Tao Ge, Jaideep Pathak, Mohammad Shoaib Abbas, Suman Ravuri, Karthik Kashinath, Naser Mahfouz, Peter Caldwell, Mike Pritchard, Noah Brenowitz

Published 2026-07-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zeyuan Hu, Akshay Subramaniam, Noel Keen, Tao Ge, Jaideep Pathak, Mohammad Shoaib Abbas, Suman Ravuri, Karthik Kashinath, Naser Mahfouz, Peter Caldwell, Mike Pritchard, Noah Brenowitz

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 trying to predict the weather for the entire planet. For decades, scientists have used massive supercomputers to run physics simulations. These simulations are incredibly accurate, but they are also incredibly slow and expensive. Running them for just one day of simulated weather consumes as much electricity as a small town uses in a day. It's like trying to paint a masterpiece by hand, one tiny brushstroke at a time, when you need to finish the whole canvas in an hour.

Enter STRATA, a new artificial intelligence model developed by researchers at NVIDIA and other institutions. Think of STRATA not as a painter, but as a master "weather chef" who has tasted the recipe so many times that they can instantly recreate the dish without measuring every single ingredient.

Here is how STRATA works, explained through simple analogies:

1. The Problem: Too Much Detail, Too Little Time

Traditional weather models try to simulate every single cloud and raindrop across the whole Earth. To do this, they need to break the planet into a grid of tiny squares (about 5 kilometers wide). This creates a massive puzzle with 25 million pieces.

The Old Way: The physics supercomputer tries to calculate the movement of every single piece of the puzzle simultaneously. It's like trying to solve a 25-million-piece jigsaw puzzle by looking at every piece individually. It takes forever and uses a ton of energy.
The New Way (STRATA): STRATA learns the patterns of the weather. Instead of calculating the physics from scratch every time, it predicts what happens next based on what it has seen before.

2. The Strategy: The "Tiled" Approach

The biggest challenge is that the Earth is round, and the weather data is huge. You can't fit the whole planet into the AI's memory at once.

The Analogy: Imagine you are trying to learn the layout of a giant city, but you can only look at one neighborhood at a time through a small window.
How STRATA Does It: STRATA breaks the Earth into small, overlapping "tiles" (like tiles on a floor). It learns how the weather changes in one small neighborhood (a 128x128 tile) over 10 minutes.
The Magic: Because weather on a 10-minute scale mostly depends on what's happening right next to it (local dynamics), the AI doesn't need to see the whole planet to make a good guess about the next 10 minutes. It learns the "local rules" of the weather.
The Assembly: When it needs to predict the whole planet, it runs the simulation on all these small tiles at the same time and then blends the edges together, like a patchwork quilt, to create a seamless global forecast.

3. The Special Tools: Seeing the World Correctly

The Earth is a sphere, but computers usually think in flat squares. This causes distortion, especially near the poles (like stretching a map of the world).

The "Stereographic" Lens: STRATA uses a special mathematical trick called Stereographic Rotary Position Embedding. Imagine looking at the Earth through a fisheye lens that centers on your specific tile. This allows the AI to understand that "North" and "East" change direction as you move around the globe, keeping its predictions accurate no matter where it is looking.
The "De-aliasing" Filter: When you stitch small tiles together, you sometimes get a "checkerboard" pattern where the edges don't match perfectly. STRATA has a special decoder (a cleanup crew) that smooths out these jagged edges, ensuring the temperature and rain look natural and continuous, not blocky.

4. The Results: Fast, Efficient, and Realistic

The paper claims STRATA is a massive leap forward in efficiency and realism:

Speed: While the old physics model takes about 1 day of supercomputer time to simulate 1 day of weather, STRATA can simulate 48 days of weather using the same amount of energy. That's a 50x improvement in energy efficiency.
Detail: It captures "storm-resolving" details. This means it can see individual thunderstorms, tropical cyclones, and mountain rain patterns that older, coarser models miss. It's like switching from a blurry, low-resolution photo to a sharp 4K image.
Stability: It can run for 24 hours straight without falling apart. While it isn't perfect (it sometimes gets the big, slow-moving weather patterns slightly wrong over time), it successfully recreates the chaotic, messy behavior of real storms.

5. The Catch (What the Paper Says)

The paper is honest about the limitations:

Data Hunger: Because storms are chaotic and full of detail (high "information density"), STRATA needs about 10 times more computing power per grid point than simpler, low-resolution AI weather models.
Training Data: It was trained on only 17 days of high-resolution physics data. Despite this small amount of training data, it managed to produce stable, realistic forecasts, which is a significant achievement.
The "Drift": Over long periods (beyond 24 hours), the model starts to drift away from reality, particularly with large-scale wind patterns. The authors suggest that in the future, they might need to pair this high-speed AI with a slower, low-resolution model to keep the big picture accurate.

Summary

STRATA is like taking a high-definition, slow-motion video of the Earth's weather and teaching an AI to predict the next frame instantly. By breaking the planet into manageable neighborhoods and using smart math to handle the globe's curvature, it achieves a level of detail and speed that was previously impossible, making high-resolution weather simulation affordable and fast.

Technical Summary: STRATA – Scaling Storm-Resolving Atmospheric AI Simulation to the Entire Planet

1. Problem Statement

Global atmospheric models capable of explicitly resolving convective storms (Global Storm-Resolving Models or GSRMs) operate at kilometer-scale resolutions (e.g., ~4.9 km) and 10-minute timesteps. While these models offer unprecedented physical fidelity for simulating precipitation extremes and cloud feedbacks, they are computationally prohibitive. Leading GSRMs, such as SCREAM, require exascale supercomputers and consume approximately 1 MWh to simulate just one day of atmospheric time, limiting their utility for long-duration ensembles and climate projections.

Existing AI weather emulators have achieved high efficiency at coarse resolutions (~0.25°) or limited regional domains, but they fail to replicate the explicit convective dynamics required for storm-resolving scales. The challenge lies in the high information density of 3D convective dynamics, which demands significantly more computational resources per grid point than large-scale synoptic forecasting. Furthermore, training autoregressive models on global, high-resolution data is hindered by memory bottlenecks (activation memory) and the scarcity of long-duration global simulation data.

2. Methodology: STRATA Architecture

The authors introduce STRATA (Storm-resolving Tile-based autoRegressive Atmosphere Transformer Architecture), the first AI emulator designed for autoregressive global rollout at kilometer-scale resolution. The methodology is built on the premise that atmospheric dynamics on 10-minute timescales are predominantly local, allowing the conversion of scarce global temporal samples into abundant local spatial samples.

2.1 Tile-Based Training and Inference

Training Strategy: To overcome activation memory bottlenecks, STRATA is trained on independent $128 \times 128$ tiles (~576 km across) cropped from global snapshots. This trades scarce temporal coverage for abundant spatial samples.
Inference Strategy: During inference, the global cubed-sphere grid is decomposed into overlapping tiles. The model processes these tiles independently, and the results are blended into a continuous global state.
Boundary Handling: To address geometric discontinuities at cubed-sphere face boundaries, the authors employ Duo-grid halo padding. This extends the equiangular grid of a face past its boundary and interpolates values from adjacent faces, providing a locally coherent neighborhood for boundary tiles.

2.2 Core Architecture

The model is based on a Diffusion Transformer (DiT) backbone adapted for deterministic atmospheric emulation, featuring three key innovations:

3D Local-Attention Backbone:
- 3D Patch Embedding: Unlike coarse-resolution models that treat the vertical dimension as a channel, STRATA treats it as a spatial axis ( $p_v=1$ ). This allows the modeling of 3D convective dynamics at the token level.
- Neighborhood Attention (NATTEN): The model uses 3D local neighborhood attention within windows of $36 \times 36$ native grid cells (~160 km).
- Column-wise Attention: Alternating blocks of column-wise attention capture vertical processes (e.g., radiation, microphysics) that lack horizontal coupling.
Stereographic Rotary Position Embedding (StereoRoPE):
- Standard position embeddings fail on spherical grids due to orientation changes. STRATA projects each tile to a local stereographic tangent plane centered on the tile.
- Wind vectors are rotated into the tile-centered East-North frame to ensure grid-invariant encoding. This enables the model to generalize across different grid topologies without retraining.
Pixel-Space De-aliasing Decoder:
- Patch tokenization introduces checkerboard artifacts (spectral bias) in smooth fields.
- STRATA employs a lightweight decoder inspired by PixelDiT, replacing linear upsampling with bilinear upsampling followed by depthwise convolution. This enforces continuity across patch boundaries and suppresses rollout artifacts.

2.3 Training Objective

STRATA is first trained for one-step prediction using a channel-normalized Smoothed $\ell_1$ (Huber) loss. To ensure stability in autoregressive rollouts, the model undergoes multi-step fine-tuning (up to 4 steps), where the loss is computed only on the final prediction of the sequence. This suppresses the accumulation of errors and instabilities.

3. Key Results

3.1 Physical Realism and Predictability

Rollout Stability: STRATA produces stable 24-hour global rollouts with realistic km-scale dynamics across diverse regimes, including tropical cyclones, stratocumulus, midlatitude fronts, and orographic precipitation.
Statistical Fidelity: The model maintains organized rainfall structures better than a persistence baseline (measured by Fractions Skill Score). Tropical precipitation propagation speeds and phase tilts match the ground truth (SCREAM).
Limitations: Large-scale biases develop over time. Global mean rainfall is biased low by ~20%, and error growth rates for large-scale circulations (e.g., 850 hPa temperature and wind) are faster than expected for a 24-hour forecast, though slower than the model's own instability. The model struggles with synoptic-scale dynamics compared to coarse-resolution emulators trained on decades of data.

3.2 Computational Efficiency and Scaling

Energy Efficiency: STRATA achieves 48 simulated days per MWh, representing a ~50x improvement in energy efficiency over the SCREAM physics model (0.78 SD/MWh).
Throughput: On 512 H100 GPUs, STRATA achieves a throughput of 741 simulated days per wall-clock day.
Compute Requirements: An iso-FLOP scaling study reveals that km-scale emulation requires ~10x more FLOPs per grid point than coarse-resolution AI weather models. This is attributed to the higher information entropy of 3D convective data compared to the quasi-2D large-scale atmosphere.
Scaling Laws: Within a fixed FLOP budget, larger patch sizes ( $p_h=4$ ) yield lower loss and faster inference (1.5–2.9x faster) than smaller patches, as they improve Model FLOP Utilization (MFU) and effective attention windows.

3.3 Grid Invariance

The model demonstrates zero-shot generalization to unseen grid topologies. When deployed on latitude-longitude and rotated oblique stereographic grids without retraining, STRATA produces qualitatively consistent precipitation and humidity fields, validating the effectiveness of the StereoRoPE and wind vector rotation mechanisms.

4. Significance and Claims

The paper claims STRATA is the first AI model to successfully demonstrate autoregressive rollout at global kilometer-scale resolution (4.9 km, 10-minute timesteps), expanding the state of the art from limited regional domains (900 km, 2-hour rollouts) to the entire planet.

The significance of this work lies in:

Feasibility of Planetary-Scale Convection-Resolving Simulation: It proves that AI can generate convection-resolving simulations at a cost (energy and time) that is orders of magnitude lower than physics-based GSRMs, making large-ensemble simulations for climate uncertainty quantification potentially tractable.
Quantitative Framework for High-Resolution Emulation: The study establishes that storm-resolving data has significantly higher entropy, requiring ~10x more compute per grid point than synoptic models. This provides a design framework for future high-resolution emulators.
Architectural Innovations: The introduction of StereoRoPE for grid-invariant spherical encoding and the tile-based training strategy offers a scalable solution for global high-resolution modeling that avoids the memory bottlenecks of traditional global solvers.

The authors remain modest regarding the current limitations, noting that while convective-scale physics are well-learned, large-scale circulation biases persist. They posit that coupling STRATA with a complementary low-resolution global emulator is a necessary next step to achieve century-scale simulations and eventually resolve low-cloud feedbacks at ~10–100 m scales.