HiRO-ACE: Fast and skillful AI emulation and downscaling trained on a 3 km global storm-resolving model

Kilometer-scale simulations of the atmosphere are an important tool for assessing local weather extremes and climate impacts, but computational expense limits their use to small regions, short periods, and limited ensembles. Machine learning offers a pathway to efficiently emulate these high-resolution simulations. Here we introduce HiRO-ACE, a two-stage AI modeling framework combining a stochastic version of the Ai2 Climate Emulator (ACE2S) with diffusion-based downscaling (HiRO) to generate 3 km precipitation fields over arbitrary regions of the globe. Both components are trained on data derived from a decade of atmospheric simulation by X-SHiELD, a 3 km global storm-resolving model. HiRO performs a 32x downscaling–generating 3 km 6-hourly precipitation from coarse 100 km inputs by training on paired high-resolution and coarsened X-SHiELD outputs. ACE2S is a $1^\circ \times 1^\circ$ ($\sim$100 km) stochastic autoregressive global atmosphere emulator that maintains grid-scale precipitation variability consistent with coarsened X-SHiELD, enabling its outputs to be ingested by HiRO without additional tuning. HiRO-ACE reproduces the distribution of extreme precipitation rates through the 99.99th percentile, with time-mean precipitation biases below 10% almost everywhere. The framework generates plausible tropical cyclones, fronts, and convective events from poorly resolved coarse inputs. Its computational efficiency allows generation of 6-hourly high-resolution regional precipitation for decades of simulated climate within a single day using one H100 GPU, while the probabilistic design enables ensemble generation for quantifying uncertainty. This establishes an AI-enabled pathway for affordably leveraging the realism of expensive km-scale simulations to support local climate adaptation planning and extreme event risk assessment.

💡 Research Summary

The paper introduces HiRO‑ACE, a two‑stage artificial‑intelligence framework designed to emulate and downscale a 3 km global storm‑resolving climate model (X‑SHiELD) with dramatically reduced computational cost. The authors first train a stochastic autoregressive emulator, ACE2S, on 6‑hourly outputs from a ten‑year X‑SHiELD simulation that have been coarsened to a 1° × 1° (~100 km) grid. ACE2S builds on the deterministic ACE2 model but adds conditional layer normalization and injected noise, allowing it to generate probabilistic forecasts that preserve the grid‑scale variability of precipitation and other atmospheric fields. Pre‑training on the ERA5 reanalysis dataset provides a physical backbone, after which fine‑tuning on the coarsened X‑SHiELD data aligns the model with the high‑resolution climate dynamics.

The second stage, HiRO, is a diffusion‑based generative model that performs a 32‑fold spatial downscaling from the 100 km ACE2S (or directly from coarsened X‑SHiELD) inputs to 3 km precipitation fields. HiRO follows a CorrDiff approach: a simple bicubic interpolation of the coarse input serves as a “perfect‑prediction” mean field, and the diffusion network learns to correct this mean by iteratively removing noise, thereby reconstructing realistic fine‑scale precipitation structures. Training uses paired high‑resolution and coarsened fields, optimizing a combination of mean‑squared error and a KL‑divergence term to match the full probability distribution of 3 km precipitation.

Evaluation focuses on a hold‑out year (2023) and compares histograms, zonal power spectra, and spatial patterns among X‑SHiELD, a “perfect‑prediction” HiRO run (downscaling directly from coarsened X‑SHiELD), and the full HiRO‑ACE pipeline (ACE2S + HiRO). The results show that HiRO‑ACE reproduces the precipitation probability density function (PDF) of the reference model up to the 99.99 th percentile (≈500 mm day⁻¹) with only a few percent error, and maintains mean‑bias below 10 % globally. Extreme events (>1000 mm day⁻¹) are somewhat under‑represented (≈78 % of the reference’s 99.9999 th percentile), reflecting the loss of information when extreme convective storms are averaged to 100 km resolution. Nonetheless, the system generates plausible tropical cyclones, fronts, and mesoscale convective systems, preserving realistic spatial organization.

Computationally, ACE2S runs at roughly 1,500 simulated years per day on a single NVIDIA H100 GPU, while HiRO downscales a 16° × 16° patch of one year’s data in about 45 minutes on the same hardware. Training the two models required 7 days (ACE2S) and 4 days (HiRO) on eight H100 GPUs, consuming only ~1 % of the energy (≈6–8 GJ) used by the original X‑SHiELD simulation (≈800 GJ). Inference for a global 10‑year 3 km precipitation dataset costs ~0.5 % of the original simulation’s energy and can be completed in a single day using one H100 GPU, representing a >100‑fold speedup.

The authors discuss limitations: (1) the most extreme precipitation tails are not perfectly captured due to coarse‑grid averaging of convective spikes; (2) the framework currently focuses solely on precipitation, leaving temperature, humidity, and wind fields for future work; (3) polar regions (>65° S/N) are excluded because of data sparsity; and (4) further model compression (e.g., denoiser distillation) could reduce inference time by an order of magnitude.

In conclusion, HiRO‑ACE demonstrates that a probabilistic coarse‑grid emulator combined with a diffusion‑based super‑resolution model can faithfully reproduce the statistics and dynamics of a 3 km global storm‑resolving simulation at a fraction of the computational cost. This opens a practical pathway for generating large ensembles of high‑resolution precipitation data for climate‑impact assessments, adaptation planning, and extreme‑event risk analysis, and provides a template for extending AI‑driven downscaling to additional climate variables and regions.