Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction

• URL: http://arxiv.org/abs/2604.16590v1
• Date: Fri, 17 Apr 2026 17:20:33 GMT
• Title: Global Attention with Linear Complexity for Exascale Generative Data Assimilation in Earth System Prediction
• Abstract summary: We introduce a unified one-stage generative DA framework that reformulates assimilation as posterior sampling.<n>At the core is STORM, a novel transformer with a global attention linear-complexity scaling algorithm that breaks the quadratic attention barrier.<n>We further scale to 20 billiontemporal tokens, enabling scalable-scale global modeling regimes over 177k temporal frames, previously unreachable.
• Score: 9.39184913345145
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Accurate weather and climate prediction relies on data assimilation (DA), which estimates the Earth system state by integrating observations with models. While exascale computing has significantly advanced earth simulation, scalable and accurate inference of the Earth system state remains a fundamental bottleneck, limiting uncertainty quantification and prediction of extreme events. We introduce a unified one-stage generative DA framework that reformulates assimilation as Bayesian posterior sampling, replacing the conventional forecast-update cycle with compute-dense, GPU-efficient inference. At the core is STORM, a novel spatiotemporal transformer with a global attention linear-complexity scaling algorithm that breaks the quadratic attention barrier. On 32,768 GPUs of the Frontier supercomputer, our method achieves 63% strong scaling efficiency and 1.6 ExaFLOP sustained performance. We further scale to 20 billion spatiotemporal tokens, enabling km-scale global modeling over 177k temporal frames, regimes previously unreachable, establishing a new paradigm for Earth system prediction.

Related papers

• PnP-Corrector: A Universal Correction Framework for Coupled Spatiotemporal Forecasting [62.22445401483844]
  Existing methods are severely constrained by the bottleneck of compounding errors.<n>We propose a universal framework called -Cast-Corrector (Plug-and-Play Corrector)<n>Our method significantly enhances the long-term stability of coupled forecasting systems.
  arXiv  Detail & Related papers  (2026-05-09T13:12:33Z)
• Earth-o1: A Grid-free Observation-native Atmospheric World Model [81.55899748753434]
  We present Earth-o1, an observation-native atmospheric world model that overcomes structural limitations.<n>By integrating diverse sensor inputs into a unified, grid-free dynamical field, the model autonomously advances the atmospheric state in space and time.<n>We show that this fundamentally distinct paradigm enables direct, real-time forecasting and cross-sensor inference without the overhead of explicit numerical solvers.
  arXiv  Detail & Related papers  (2026-05-07T14:27:48Z)
• FuXiWeather2: Learning accurate atmospheric state estimation for operational global weather forecasting [15.10103712842682]
  We present FuXiWeather2, a unified end-to-end neural framework for assimilation and forecasting.<n>FuXiWeather2 generates high-resolution ($0.25circ$) global analysis fields and 10-day forecasts within minutes.<n>These high-quality analysis fields drive deterministic forecasts that exceed the skill of the HRES system in 91% of evaluated metrics.
  arXiv  Detail & Related papers  (2026-03-16T14:36:47Z)
• Searth Transformer: A Transformer Architecture Incorporating Earth's Geospheric Physical Priors for Global Mid-Range Weather Forecasting [23.666628869892737]
  YanTian is a global medium-range weather forecasting model.<n>It achieves higher accuracy than the high-resolution forecast of the European Centre for Medium-Range Weather Forecasts.<n>It requires roughly 200 times lower computational cost than standard autoregressive fine-tuning.
  arXiv  Detail & Related papers  (2026-01-14T13:20:17Z)
• DAWP: A framework for global observation forecasting via Data Assimilation and Weather Prediction in satellite observation space [60.729377189859]
  We propose our DAWP framework to enable AIWPs to operate in a complete observation space.<n>AIDA module applies a mask multi-modality autoencoder for assimilating irregular satellite observation tokens.<n>We show that AIDA significantly improves the roll out and efficiency of AIWP and holds promising potential to be applied in global precipitationresolution forecasting.
  arXiv  Detail & Related papers  (2025-10-13T03:13:35Z)
• FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale [91.84761739154366]
  FourCastNet 3 advances global weather modeling by implementing a scalable, geometric machine learning (ML) approach to probabilistic ensemble forecasting.<n>FourCastNet 3 delivers forecasting accuracy that surpasses leading conventional ensemble models and rivals the best diffusion-based methods.<n>Its computational efficiency, medium-range probabilistic skill, spectral fidelity, and rollout stability at subseasonal timescales make it a strong candidate for improving meteorological forecasting and early warning systems through large ensemble predictions.
  arXiv  Detail & Related papers  (2025-07-16T11:22:18Z)
• Advanced long-term earth system forecasting by learning the small-scale nature [74.19833913539053]
  We present Triton, an AI framework designed to address this fundamental challenge.<n>Inspired by increasing grids to explicitly resolve small scales in numerical models, Triton employs a hierarchical architecture processing information across multiple resolutions to mitigate spectral bias.<n>We demonstrate Triton's superior performance on challenging forecast tasks, achieving stable year-long global temperature forecasts, skillful Kuroshio eddy predictions till 120 days, and high-fidelity turbulence simulations.
  arXiv  Detail & Related papers  (2025-05-26T02:49:00Z)
• OneForecast: A Universal Framework for Global and Regional Weather Forecasting [67.61381313555091]
  We propose a global-regional nested weather forecasting framework (OneForecast) based on graph neural networks.<n>By combining a dynamic system perspective with multi-grid theory, we construct a multi-scale graph structure and densify the target region.<n>We introduce an adaptive messaging mechanism, using dynamic gating units, to deeply integrate node and edge features for more accurate extreme event forecasting.
  arXiv  Detail & Related papers  (2025-02-01T06:49:16Z)
• A Scalable Real-Time Data Assimilation Framework for Predicting Turbulent Atmosphere Dynamics [8.012940782999975]
  We introduce a generic real-time data assimilation framework and demonstrate its end-to-end performance on the Frontier supercomputer. This framework comprises two primary modules: an ensemble score filter (EnSF) and a vision transformer-based surrogate. We demonstrate both the strong and weak scaling of our framework up to 1024 GPUs on the Exascale supercomputer, Frontier.
  arXiv  Detail & Related papers  (2024-07-16T20:44:09Z)
• ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability [10.88886669820126]
  We introduce the Oak Ridge Base Foundation Model for Earth System Predictability (ORBIT) ORBIT is the largest model of its kind and surpasses the current climate AI foundation model size by a thousandfold. Performance scaling tests on the Frontier supercomputer have demonstrated that ORBIT achieves 684 petaFLOPS to 1.6 exaFLOPS sustained throughput.
  arXiv  Detail & Related papers  (2024-04-23T03:39:57Z)
• Towards an end-to-end artificial intelligence driven global weather forecasting system [57.5191940978886]
  We present an AI-based data assimilation model, i.e., Adas, for global weather variables. We demonstrate that Adas can assimilate global observations to produce high-quality analysis, enabling the system operate stably for long term. We are the first to apply the methods to real-world scenarios, which is more challenging and has considerable practical application potential.
  arXiv  Detail & Related papers  (2023-12-18T09:05:28Z)
• Residual Corrective Diffusion Modeling for Km-scale Atmospheric Downscaling [58.456404022536425]
  State of the art for physical hazard prediction from weather and climate requires expensive km-scale numerical simulations driven by coarser resolution global inputs. Here, a generative diffusion architecture is explored for downscaling such global inputs to km-scale, as a cost-effective machine learning alternative. The model is trained to predict 2km data from a regional weather model over Taiwan, conditioned on a 25km global reanalysis.
  arXiv  Detail & Related papers  (2023-09-24T19:57:22Z)
• Earthformer: Exploring Space-Time Transformers for Earth System Forecasting [27.60569643222878]
  We propose Earthformer, a space-time Transformer for Earth system forecasting. The Transformer is based on a generic, flexible and efficient space-time attention block, named Cuboid Attention. Experiments on two real-world benchmarks about precipitation nowcasting and El Nino/Southerntemporaltion show Earthformer achieves state-of-the-art performance.
  arXiv  Detail & Related papers  (2022-07-12T20:52:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.