Scaling Storm-Resolving Atmospheric AI Simulation to the Entire Planet

STRATA is the first AI model to autoregressively simulate the entire planet's atmosphere at storm-resolving resolution — 4.9 km grid spacing, 10-minute timesteps, ~25 million horizontal grid columns — and it does so at roughly 50× the energy efficiency of the exascale physics model it emulates. The striking part is how little it needed: 17 days of training data. The trick that makes this possible is a bet on locality — that at 10-minute timescales the atmosphere is a local process, so you can train on small tiles and stitch a planet together at inference — and the paper is candid that the bet pays off at storm scales while the planetary scales slowly drift. That honesty, plus a genuinely interesting scaling result about the information density of convective data, makes this worth a careful read.

Why global storm-resolving simulation is a moonshot

Most global climate models run at 25–100 km resolution. At those scales a thunderstorm is smaller than a grid cell, so convection — the process shaping precipitation extremes, tropical cyclones, and the cloud feedbacks that dominate climate-sensitivity uncertainty — gets replaced by hand-tuned parameterizations. This is the largest known source of uncertainty in long-range climate projections.

The physics community's answer is global storm-resolving models (GSRMs) like SCREAM: run the fluid equations at ~3–5 km so storms are explicitly resolved. It works, and the fidelity gains are real. The cost is absurd. SCREAM achieves about 0.78 simulated days per megawatt-hour on the Frontier supercomputer — the paper's own back-of-envelope: roughly one MWh per simulated day, or about 2.5 simulated days per ton of CO₂ emitted. You cannot run century-scale ensembles that way.

AI emulators are the obvious escape hatch, and they've worked well at coarse resolution: models like GraphCast-era systems and ACE2 rival operational forecasts at 0.25° with 6-hour steps, and some roll out stably for decades. But nobody had achieved autoregressive rollout that is both global and kilometer-scale — the only regime that could actually replace a GSRM. The best prior km-scale result (Flora & Potvin) covered a (1000 km)² regional domain for 2 hours.

The obstacle isn't just model design. A day of GSRM output occupies thousands of times more memory than a day of ERA5. You can't fit global 25M-cell activations in GPU memory during training, and you don't have decades of training data — generating it costs a MWh per day. Constraining convection also demands 10-minute temporal sampling (36× denser than the 6-hour steps of coarse AI weather models) and 4.9 km spatial sampling (25× denser). Everything about the problem is bigger.

The key idea: trade time for space

Here is the aha. In 10 minutes, information in the atmosphere can only travel so far — even acoustic waves, the fastest signals, cover only ~200 km per step. So the 10-minute update at any point on Earth depends only on a modest neighborhood around it. The global update operator is really a local operator applied everywhere.

This reframes the data problem entirely. You have only 17 days of global snapshots — temporally scarce. But each snapshot contains 25 million grid columns, i.e., thousands of independent 576-km tiles, each a valid training example of "local atmospheric state → local state 10 minutes later." Scarce global-temporal samples become abundant local-spatial samples. Training happens on 128×128 tiles (small activations, fits in memory); only inference is global, done by decomposing each face of SCREAM's cubed-sphere grid into overlapping tiles, running the same model on each, and blending the overlaps with a smooth taper window (a Kaiser–Bessel-derived window, in signal-processing terms an overlap-add).

If you know patch-based diffusion super-resolution (CorrDiff-style), the tiling will feel familiar — but those models generate km-scale detail diagnostically from a coarse state, purely downscale. STRATA is autoregressive: the fine scales evolve themselves and can, in principle, feed back on larger scales. That's the regime that actually matters for climate.

The machinery that makes tiles work

Three pieces are needed to make "same model, every tile, any location" actually hold.

A local-attention 3D transformer. All fields (potential temperature, winds, vertical velocity, geopotential, humidity across 24 vertical levels) are stacked into one volumetric tensor and tokenized with 3D patch embedding — crucially with vertical patch size 1, treating altitude as a spatial axis rather than folding it into channels, because convection is genuinely 3D. The backbone is a DiT-style transformer using 3D neighborhood attention (NATTEN) over local (x, y, level) windows of about 160×160 km, alternating with column-wise vertical attention — a nice architectural echo of how the physics model itself computes radiation and microphysics column-by-column.

StereoRoPE. Standard RoPE indexes tokens by grid position, but the same 2D index offset means different physical displacements and orientations at different points on the sphere — a model trained near the equator would see garbage at a cube-face corner. STRATA instead projects each tile onto a local stereographic tangent plane centered on the tile and feeds those physical East–North coordinates into RoPE. Stereographic projection is conformal (angle-preserving), which matters for learning vector operators like advection \mathbf{u}\cdot\nabla q, where \mathbf{u} is wind and q a tracer. Wind vectors are additionally rotated into the tile-center frame so that the same physical wind has the same numerical representation everywhere. The payoff is genuine grid invariance: the trained model runs zero-shot on lat–lon and rotated stereographic grids it never saw (Section 4.4), with qualitatively consistent 3-hour rollouts. This positional encoding trick seems broadly reusable for any tile-based model on a sphere or irregular geometry.

A de-aliasing decoder, with a real explanation. Patch-tokenized models notoriously develop checkerboard artifacts during rollout. Appendix B.4 offers a satisfying linear-algebra account: the linearized update operator of a patch-embed/decode pipeline is block-diagonal with an identical block per patch, so its dominant eigenvectors are patch-periodic patterns — and at typical initialization the spectral radius exceeds 1 (≈2 under Kaiming init for state prediction; worse for tendency prediction, explaining a known empirical instability). Autoregressive rollout is literally power iteration, so it amplifies exactly the checkerboard mode. Training suppresses this slowly for smooth fields because they project weakly onto sub-patch variations. The fix is a lightweight PixelDiT-style pixel-space decoder with bilinear upsampling plus depthwise convolutions to enforce continuity across patch boundaries. If you build autoregressive vision transformers of any kind, this appendix is worth reading on its own.

Training is one-step Huber loss, then multi-step fine-tuning (2 then 4 rollout steps, loss on the final state), which the authors say is key to 24-hour stability.

Convective data is expensive per grid point — and they can say why

The iso-FLOP study (Section 4.2) sweeps three model tiers and three patch sizes and finds that stable 24-hour rollouts only emerge at ~5 TFLOPs per tile (DiT-M tier). Normalized per grid point, km-scale emulation needs ~10× more compute than you'd extrapolate from coarse AI weather models after accounting for the ~900× geometric factor from finer space-time resolution.

Their explanation is information-theoretic, and it's the most conceptually interesting result in the paper. The large-scale atmosphere behaves quasi-two-dimensionally (rotation plus large aspect ratio), making its increments highly compressible. Convective-scale increments are genuinely 3D and much closer to white noise. They quantify this with the differential entropy of 10-minute state increments under a Gaussian model:

For 8×8 patches, SCREAM increments carry ~6 more bits/dim than ERA5's — an enormous gap given that unit-variance white noise caps out at ~2 bits/dim. Convective data is nearly incompressible, so each grid point genuinely demands more model capacity. This gives the field a quantitative planning tool: entropy of the target increments predicts FLOP requirements. It also cuts both ways, as the authors note — high information density per sample is likely why 17 days suffices to learn convective physics at all.

The patch-size result is worth noting too: at equal FLOPs, patch size 4 both trains to lower loss than patch size 1 and runs 1.5–2.9× faster (better arithmetic intensity, larger effective attention window for free).

What the rollouts actually show

Since individual convective cells lose pointwise predictability within hours, the right tests are structural and statistical, not RMSE against truth. On those terms the storm-scale results are convincing: the 12-hour snapshots (Figure 1 and the supplementary videos) reproduce a tropical cyclone, marine stratocumulus decks, a midlatitude cold front, the ITCZ, diurnal island convection, and orographic precipitation with the right texture. Hovmöller diagrams show tropical rain systems propagating across the Indo-Pacific at the correct speeds through the full 24 hours. Rainfall Fractions Skill Score beats persistence, and precipitation power spectra match SCREAM closely except at synoptic scales. Global precipitation distributions broadly match, with a slight deficit in moderate-to-intense rates.

The large scales are where it breaks. RMSE for 850-hPa temperature and wind grows faster than the canonical error-doubling-per-day, and ACE2 — a coarse emulator trained on six decades of reanalysis — shows substantially slower synoptic error growth. Global-mean rainfall drifts up to 20% low. Most tellingly, the model needs a hand-coded spectral filter: because tiles are processed independently, nothing enforces that global-mean vertical velocity stays near zero (mass continuity), so the leading spherical-harmonic modes of \omega drift, and without zeroing them after every step, global precipitation collapses to half its correct value by 24 hours (Figure S6). Locality buys everything at storm scale and costs you the global constraints.

And the efficiency claim — the headline that makes all of this matter:

That's ~50× less energy per simulated day, or in the paper's units, 131 simulated days per ton of CO₂ versus 2. The distributed inference scales at 98.4% efficiency from 8 to 512 GPUs using peer-to-peer halo exchange — the tile decomposition maps naturally onto multi-GPU parallelism.

What to make of it, and what to be skeptical of

If this direction holds, the endgame is clear and the authors state it: couple a tile-based km-scale emulator (which owns the convective physics) with a coarse global emulator like ACE2 (which owns the circulation), and you get century-scale storm-resolving simulation for quantifying how extreme-weather frequency shifts with climate — currently impossible at any budget. The paper even floats the eventual prospect of 10–100 m global resolution to resolve low-cloud feedbacks directly.

Reasons for caution. The rollout horizon is 24 hours; nothing here demonstrates climate-length stability, and the error growth curves suggest the current model would degrade beyond that without the coarse-model coupling. The evaluation uses just 6 initializations from a held-out slice of one simulation — and that simulation contains an acknowledged dataset error (1994 atmosphere paired with 2020 SSTs). The \omega spectral filter and humidity/precipitation positivity fixers are pragmatic patches, not solutions; the authors themselves say a more principled large-scale constraint is needed. Whether the synoptic bias reflects insufficient data (17 days!) or a fundamental ceiling of tile-based training is explicitly left open — that's the live question for follow-up work. And the 50× energy figure compares a deterministic emulator of 24 tropospheric levels and 6 fields against a full physics model with 128 levels and complete microphysics; it's the right comparison for the intended use case, but it isn't apples-to-apples in state fidelity.

Still, the pieces most likely to outlive this specific system — tile-based training with overlap-blend inference for planetary-scale autoregression, StereoRoPE for grid-invariant spherical encoding, the spectral analysis of patch instability, and the entropy-based compute-scaling framework — are each independently useful. Code and the SCREAM dataset are public.

Where to spend your time: Section 3 (the method, especially 3.3 on StereoRoPE), Section 4.2 with Appendix J (the entropy/compute argument), and Appendix B.4 (the patch-instability analysis) are the intellectually dense parts. The supplementary rollout videos are the fastest way to calibrate your own judgment of the fidelity claims.