AERIS: Modelo de Sistemas Terrestres de Argonne para Previsões Confiáveis e Precisas
AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions
September 16, 2025
Autores: Väinö Hatanpää, Eugene Ku, Jason Stock, Murali Emani, Sam Foreman, Chunyong Jung, Sandeep Madireddy, Tung Nguyen, Varuni Sastry, Ray A. O. Sinurat, Sam Wheeler, Huihuo Zheng, Troy Arcomano, Venkatram Vishwanath, Rao Kotamarthi
cs.AI
Resumo
O aprendizado de máquina generativo oferece novas oportunidades para compreender melhor a dinâmica complexa do sistema terrestre. Métodos recentes baseados em difusão abordam vieses espectrais e melhoram a calibração de conjuntos na previsão do tempo em comparação com métodos determinísticos, mas até agora têm se mostrado difíceis de escalar de forma estável em altas resoluções. Apresentamos o AERIS, um transformador de difusão Swin em nível de pixel com 1,3 a 80 bilhões de parâmetros, para preencher essa lacuna, e o SWiPe, uma técnica generalizável que combina paralelismo de janelas com paralelismo de sequência e pipeline para fragmentar transformadores baseados em janelas sem custo adicional de comunicação ou aumento do tamanho global do lote. No Aurora (10.080 nós), o AERIS mantém 10,21 ExaFLOPS (precisão mista) e um desempenho de pico de 11,21 ExaFLOPS com tamanho de patch 1x1 no conjunto de dados ERA5 de 0,25{\deg}, alcançando 95,5% de eficiência de escalonamento fraco e 81,6% de eficiência de escalonamento forte. O AERIS supera o IFS ENS e permanece estável em escalas sazonais até 90 dias, destacando o potencial de modelos de difusão com bilhões de parâmetros para previsão do tempo e clima.
English
Generative machine learning offers new opportunities to better understand
complex Earth system dynamics. Recent diffusion-based methods address spectral
biases and improve ensemble calibration in weather forecasting compared to
deterministic methods, yet have so far proven difficult to scale stably at high
resolutions. We introduce AERIS, a 1.3 to 80B parameter pixel-level Swin
diffusion transformer to address this gap, and SWiPe, a generalizable technique
that composes window parallelism with sequence and pipeline parallelism to
shard window-based transformers without added communication cost or increased
global batch size. On Aurora (10,080 nodes), AERIS sustains 10.21 ExaFLOPS
(mixed precision) and a peak performance of 11.21 ExaFLOPS with 1 times 1
patch size on the 0.25{\deg} ERA5 dataset, achieving 95.5% weak scaling
efficiency, and 81.6% strong scaling efficiency. AERIS outperforms the IFS ENS
and remains stable on seasonal scales to 90 days, highlighting the potential of
billion-parameter diffusion models for weather and climate prediction.