STRATA

STRATA (Storm-resolving Tile-based autoRegressive Atmosphere Transformer Architecture) is a global, kilometer-scale AI atmosphere model from Nvidia, published as part of the Earth-2 checkpoints. It emulates the SCREAM (Simple Cloud-Resolving E3SM Atmosphere Model) storm-resolving simulation, autoregressively predicting the 3D atmospheric state on SCREAM's ne1024pg2 cubed-sphere grid (about 25.2 million columns, ~4.9 km spacing) with a 10-minute time step.

The prognostic state covers potential temperature, zonal and meridional wind, geopotential height, vertical pressure velocity and specific humidity on 24 model levels, plus 2-m air temperature. Surface pressure and surface precipitation fluxes are diagnosed at each step. The architecture is a two-stage "PixelDiT" transformer with neighborhood attention (331.7M parameters), a deterministic single-step predictor trained with a regression loss and fine-tuned with 4-step autoregressive rollouts; no diffusion sampling is used.

Inference is tile-based: a single GPU can advance any part of the globe, and multi-GPU domain decomposition advances the full globe. Intended uses are weather and climate research, such as short-range forecasting of the SCREAM state, emulation of storm-resolving simulations, and experiments in coupling or downscaling AI emulators. Weights are on Hugging Face under the OpenMDW 1.1 license; training recipes and inference code are in the NVlabs/STRATA repository, and the architecture is in NVIDIA PhysicsNeMo.