cs.LGOct 20, 2026
SaveGraphCast-Mamba: Learning Non-Markovian Corrections for Weather Forecasting
Organizations: Department of Physics, Princeton University, Princeton, NJ 08544, USA · Department of Electrical and Computer Engineering, Princeton University, Princeton, NJ 08544, USA · École Polytechnique Fédérale de Lausanne (EPFL), 1015 Lausanne, Switzerland · Princeton Quantum Initiative, Princeton University, Princeton, NJ 08544, USA ·
Abstract
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 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 is combined with the learned forecast correction to produce . The Mamba correction branch has separate parameters and a temporal memory state , comprising the SSM states and causal-convolution caches of its four Mamba layers.
Figure 3: Memory carry during training. Boxes show the four 24-step windows of one segment. The arrows marked 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: m temperature RMSE and RMSE reduction on the same forecast initializations. The Mamba correction branch uses , 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 grid. From left to right: 700 hPa specific humidity (Q700), m temperature (T2m), and 850 hPa zonal wind (U850). Blue shows the selected GraphCast–Mamba model M1 ( ), dashed gray its public GraphCast_small backbone ( ), and green the public GraphCast checkpoint ( ). 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 ; 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
| Mamba LR multiplier | Ten-day aggregate | Day 10 | |
|---|---|---|---|
| 16 | 14.70 | 21.98 | |
| 16 | 15.67 | 22.44 | |
| 16 | 16.34 | 23.68 | |
| 32 | 15.82 | 22.96 | |
| 32 | 16.29 | 23.32 | |
| 32 | 15.73 | 22.75 |
Table 2: Mamba capacity and learning-rate comparison at . 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 |
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 , 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 |
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 . 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 h (5 days), the training tail duration for . 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 . During the tail, uncorrected GraphCast forecasts supply subsequent inputs. GraphCast loss reduction exhibits a pronounced crossover, with substantially larger long-lead gains for . The vertical dashed line marks h (5 days), the training tail duration for ; 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 ( ), dashed gray its public GraphCast_small backbone ( ), and green the public GraphCast checkpoint ( ). 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 m temperature, U10 and V10 the 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 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 .
Figure 10: Ten-day forecasts from 00 UTC on 15 January 2022 at the grid points nearest Paris ( N, E; left) and New York City ( N, W; right). From top to bottom, rows show m temperature (T2m), hPa zonal wind (U850), and 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 feedback-conditioned steps, Mamba inner dimension , spatial width , and two graph message-passing steps; SWA averages checkpoints from updates 6k–11k. Both forecasts use the ERA5 states at and 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.
More in this topic
Weather Forecasting
Weather forecasting predicts future atmospheric conditions, from local storms to global patterns, using physical models, observations, and machine learning.
See all 176 papers