Unresolved atmospheric processes can introduce history dependence beyond the short input windows used by many weather forecasting models. We introduce GraphCast--Mamba, which augments a frozen, pretrained 1∘ GraphCast backbone with a trainable Mamba correction branch combining graph message passing and Mamba layers. The branch retains a temporal memory state across successive six-hour forecast steps, learning corrections informed by temporal context beyond GraphCast's two-frame input. Training is confined to the correction branch, using ERA5 data from 2015--2021. After 2,500 optimizer updates, including warm-up, the selected model achieves substantial improvements relative to the uncorrected GraphCast backbone. It reduces the GraphCast loss (normalized weighted mean-squared error; MSE) by 14.42% when aggregated over the full ten-day forecast and by 21.66% at day 10 (240 h after forecast initialization). At this final forecast lead time, root-mean-square error (RMSE) decreases by 7.09--14.60% across all 11 evaluated variable groups, spanning temperature, wind, pressure, geopotential, specific humidity, and precipitation. These results demonstrate broad improvements in ten-day weather forecasts through a short correction-training schedule while retaining the pretrained backbone.
Figures & tables
Figure 1: GraphCast–Mamba: temporal-memory correction and forecast-skill overview. Panel (a) shows forecast feedback and memory carry. In (b), positive percentages indicate lower variable-group RMSE than the paired GraphCast baseline; (c) counts channels with lower RMSE, out of 83. Counts describe aggregate differences without channel-wise significance tests. Section 4.1 defines the evaluation protocol and RMSE aggregation.
Figure 2: Trainable Mamba correction branch. Two graph message-passing steps each feed a two-layer Mamba module, followed by decoding and an output head initialized to zero. The frozen GraphCast backbone forecast Gθ(Xt) is combined with the learned forecast correction ct to produce Y^t . The Mamba correction branch has separate parameters and a temporal memory state mt , comprising the SSM states and causal-convolution caches of its four Mamba layers.
minit=0⟶one segment: memory carried through 96 steps1:24sg25:48sg49:72sg73:96resetminit′=0
Figure 3: Memory carry during training. Boxes show the four 24-step windows of one segment. The arrows marked sg carry the SSM states and causal-convolution caches unchanged across window boundaries while stopping their gradients. Both are reset for the next segment.
Figure 4: GraphCast–Mamba forecast skill on the 122 evaluation initializations in 2023 (Section 4.1 ). Top: GraphCast loss and its per-lead reduction. Bottom: 2 m temperature RMSE and RMSE reduction on the same forecast initializations. The Mamba correction branch uses dinner=16 , a threefold Mamba learning rate (LR), and SWA over post-warm-up updates 500–2,000. Every reduction uses the paired GraphCast baseline shown in the absolute-error panels.
Figure 5: Forecast RMSE relative to the matched ECMWF ensemble mean (ENS) on a common 1.5∘ grid. From left to right: 700 hPa specific humidity (Q700), 2 m temperature (T2m), and 850 hPa zonal wind (U850). Blue shows the selected GraphCast–Mamba model M1 ( 1∘ ), dashed gray its public GraphCast_small backbone ( 1∘ ), and green the public GraphCast checkpoint ( 0.25∘ ). From 122 candidate 00 UTC initializations every third day in 2023, each diagnostic and lead uses the identical ENS-available dates for all four forecasts (92–105 starts per point in these panels). Squared errors are pooled before taking square roots. The horizontal line marks ENS at 100% ; lower values indicate lower RMSE. Appendix C specifies the protocol.
Policy
Temporal memory retained
Training/forecasting
Training
Forecasting
Carry/carry
✓
✓
Carry/reset
✓
×
Reset/reset
×
×
Table 1: Memory retained between successive steps: ✓ denotes retention; × denotes resetting both the SSM state and causal-convolution cache before every step. Carry/carry and carry/reset share the same trained weights. Reset/reset is trained separately with the same architecture and parameter count.
Figure 6: Comparison of carry/carry, carry/reset, and matched reset/reset memory policies defined in Table 1 . Curves show exact GraphCast loss reduction relative to a common GraphCast baseline on the 122 initializations in 2023. Reset/reset uses the matched training schedule described in the text.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
dinner
Mamba LR multiplier
Ten-day aggregate
Day 10
16
0.3×
14.70
21.98
16
1×
15.67
22.44
16
3×
16.34
23.68
32
0.3×
15.82
22.96
32
1×
16.29
23.32
32
3×
15.73
22.75
Appendix
Table 2: Mamba capacity and learning-rate comparison at 1∘ . Entries are GraphCast loss reductions (%) on the 32 forecast initializations from the 2022 development set; bold marks the largest value in each metric column. The multiplier applies only to Mamba parameters, with the spatial learning rate fixed. All six configurations use the same training schedule and SWA window.
Day-10 RMSE
RMSE reduction (%)
Variable [unit]
GraphCast
GraphCast– Mamba
Day 5
Day 7
Day 10
Upper-air variables (13 pressure levels)
Geopotential [m 2 s -2 ]
784.41
703.68
1.24
5.22
10.29
Temperature [K]
3.10
2.79
2.11
5.98
10.07
Specific humidity [g kg -1 ]
1.10
0.97
4.33
7.88
11.68
Zonal wind [m s -1 ]
8.67
7.63
2.76
7.19
11.94
Appendix
Table 3: Physical-unit RMSE and relative RMSE reduction for the principal GraphCast–Mamba model. Days 5, 7, and 10 denote leads of 120, 168, and 240 h. RMSE pools cosine-latitude-weighted squared errors across forecast initializations and, for upper-air groups, equally across 13 pressure levels before taking the square root. Each reduction is 100(1−RMSEmodel/RMSEGC) , computed before rounding. Positive values indicate lower RMSE; these percentages are distinct from GraphCast loss reductions.
Quantity
Selected model
Accelerator
1 NVIDIA A100 80GB PCIe
Allocated CPU cores
8
Elapsed allocation time
22 h 40 min 45 s
GPU allocation cost
22.68 GPU-hours
Peak host memory
97.1 GiB
Sampled peak GPU memory
23.15 GiB
Appendix
Table 4: Training resources and parameter counts for the selected correction model. GPU-hours equal elapsed allocation hours because the run uses one GPU. The GPU-memory value is a sampled peak.
Figure 7: Forecast skill after closed-loop training with different training tail lengths K . During the tail, corrected GraphCast–Mamba predictions supply subsequent inputs. Each curve shows GraphCast loss reduction relative to GraphCast over ten forecast days; positive values indicate improvement. The vertical dashed line marks 20×6=120 h (5 days), the training tail duration for K=20 . Gains for this model persist beyond its training tail to day 10, twice that duration.
Figure 8: Forecast skill after open-loop training with a teacher-forced prefix and different training tail lengths K . During the tail, uncorrected GraphCast forecasts supply subsequent inputs. GraphCast loss reduction exhibits a pronounced crossover, with substantially larger long-lead gains for K≥18 . The vertical dashed line marks 20×6=120 h (5 days), the training tail duration for K=20 ; this model retains positive gains beyond that duration, up to day 10. This figure and Figure 7 use the same vertical scale but compare separate sweeps with different experimental settings and loss reconstructions (Appendix A.2 ).
Figure 9: The comparison in Figure 5 extended to 12 diagnostics. Blue shows the selected GraphCast–Mamba model M1 ( 1∘ ), dashed gray its public GraphCast_small backbone ( 1∘ ), and green the public GraphCast checkpoint ( 0.25∘ ). Z denotes geopotential, T temperature, Q specific humidity, U zonal wind, and V meridional wind; a three-digit suffix gives pressure in hPa. T2m denotes 2 m temperature, U10 and V10 the 10 m wind components, and MSL mean sea-level pressure. All four forecasts use identical dates and ERA5 truth at each diagnostic and lead on the 1.5∘ grid; the available date set varies across points (91–108 starts). Curves show pooled RMSE as a percentage of the matched ENS-mean RMSE, with ENS at 100% .
Figure 10: Ten-day forecasts from 00 UTC on 15 January 2022 at the 1∘ grid points nearest Paris ( 49∘ N, 2∘ E; left) and New York City ( 41∘ N, 74∘ W; right). From top to bottom, rows show 2 m temperature (T2m), 850 hPa zonal wind (U850), and 6 h accumulated precipitation in physical units. ERA5 is shown in black, the frozen GraphCast forecast in dashed orange, and the GraphCast–Mamba forecast in blue. The correction was trained separately on 15 years of ERA5 data (2007–2021), with a training tail of K=20 feedback-conditioned steps, Mamba inner dimension 128 , spatial width 512 , and two graph message-passing steps; SWA averages checkpoints from updates 6k–11k. Both forecasts use the ERA5 states at −6 and 0 h (before and on the dotted line) as inputs, then feed back their own predicted fields at each step; the GraphCast–Mamba model begins with a zero temporal memory state. Precipitation is shown as predicted and may take small negative values.
Weather forecasting predicts future atmospheric conditions, from local storms to global patterns, using physical models, observations, and machine learning.
Weather forecasting offers an ideal testbed for artificial intelligence (AI) to learn complex, multi-scale physical systems. Traditional numerical weather prediction remains computationally costly for frequent regional updates, as high-resolution nests require intensive boundary coupling. We introduce Multi-Resolution Graph Neural Forecasting (MR-GNF), a lightweight, physics-aware model that performs short-term regional forecasts directly on an ellipsoidal, multi-scale graph of the Earth. The framework couples a 0.25° region of interest with a 0.5° context belt and 1.0° outer domain, enabling continuous cross-scale message passing without explicit nested boundaries. Its axial graph-attention network alternates vertical self-attention across pressure levels with horizontal graph attention across surface nodes, capturing implicit 3-D structure in just 1.6 M parameters. Trained on 40 years of ERA5 reanalysis (1980-2024), MR-GNF delivers stable +6 h to +24 h forecasts for near-surface temperature, wind, and precipitation over the UK-Ireland sector. Despite a total compute cost below 80 GPU-hours on a single RTX 6000 Ada, the model matches or exceeds heavier regional AI systems while preserving physical consistency across scales. These results demonstrate that graph-based neural operators can achieve trustworthy, high-resolution weather prediction at a fraction of NWP cost, opening a practical path toward AI-driven early-warning and renewable-energy forecasting systems. Project page and code: https://github.com/AndriiShchur/MR-GNF
Andrii Shchur, Inna Skarga-Bandurova
Kyiv School of Economics · Oxford Brookes University
Global medium-range weather forecasting requires modeling structured yet state-dependent interactions among heterogeneous atmospheric fields. Existing data-driven models largely learn these interactions implicitly, whereas equation-level physical constraints may inherit approximation and model-form biases. We present VeinCast, a physics-guided dynamic field graph and graph-conditioned fusion framework that jointly forecasts 69 surface and upper-air fields. Within each local window, its Physics-Guided Dynamic Field Graph combines predefined atmospheric relations with state-dependent Top-K residual edges and adapts Earth-window attention using the resulting graph context. Graph-Conditioned Latent Fusion further employs graph context and source-node centrality to guide field-to-latent aggregation, while bounded feedback preserves field-specific information. On the 1.5∘ ERA5 benchmark, VeinCast demonstrates competitive forecasting performance across all 69 meteorological fields at lead times of up to 14 days, compared with representative global weather forecasting models including FuXi, Pangu-Weather, GraphCast, FengWu, and ARROW. Ablations confirm that the two modules provide complementary gains, demonstrating the effectiveness of relational-level physical guidance for data-driven weather forecasting.
Zhisheng Chen, Jinhan Li, Yuxuan Li +6
1Nanyang Technological University · 2Tsinghua University · 3Peking University +1
Martian weather forecasting is important for future exploration, but atmospheric behaviour on Mars combines spatial, temporal, vertical, and dust-driven processes in ways that challenge current modelling and forecasting approaches. This paper introduces MaGMA (Martian Graph-based Multi-horizon Atmospheric Forecasting), a graph-based data engineering framework that transforms OpenMARS reanalysis fields into structured learning objects for Martian atmospheric forecasting. Local atmospheric patches are represented as graph nodes and linked through spatial neighbourhoods, temporal continuity, longer temporal dependencies, and dynamically similar atmospheric states. The model integrates recent atmospheric history, engineered physical descriptors, and vertical atmospheric information to support forecasting across multiple horizons. We evaluate MaGMA across five unseen Martian years, including regular years and a global dust storm year. In regular years, the model achieves overall R^2 values of approximately 0.73-0.85. For dust-column forecasting, it outperforms classical and deep temporal baselines in most year-horizon comparisons. During the global dust storm year, dust-column prediction remains strong at shorter horizons, with R^2 above 0.8 for the first two horizons, while broader multivariate performance declines. The results show that graph-based data engineering can create reusable and diagnostically useful representations for planetary atmospheric forecasting, while highlighting the need for better learning under rare extreme regimes and improved use of vertical atmospheric structure.