GraphCast: Learning skillful medium-range global weather forecasting

In plain words

Weather forecasts today come mostly from giant physics simulations run on supercomputers. They are very accurate, but they do not learn from the decades of recorded weather that already exist. GraphCast takes the other route: it is a neural network trained on past weather records to predict how the atmosphere changes from one short time step to the next.

The model looks at the current state of the air around the planet and predicts what it will look like six hours later. It then feeds its own prediction back in to get the next step, and repeats this until it has a forecast about ten days ahead. To spread information across the whole globe, it works on a mesh of connected points that is dense everywhere and has a few long-range links that let distant regions talk to each other quickly. The authors report that it beats the best operational deterministic forecast on most of the variables they checked, and that it is fast enough to run in about a minute on one accelerator.

TL;DR

• GraphCast is a learned one-step weather model, a graph neural network with an encode, process and decode structure, rolled out autoregressively in six-hour steps § GraphCast.
• Its processor runs 16 unshared message-passing layers on a multi-mesh built from a refined icosahedron, which gives long-range links within few steps app. §3.3 app. §3.5.
• It is trained with an MSE loss over 12-step rollouts on ERA5 (development data 1979–2017), after a curriculum that ramps the rollout length up to 12 app. §4.2 app. §4.5.
• Author-reported: it beats HRES-fc0 on 90.3% of 1380 verification targets over 10 days § Forecast verification results.
• The released code is Apache-2.0 and covers the model and loss, but not the training loop (see Paper vs code).

Problem

Numerical weather prediction (NWP) gets more accurate with more compute, but it cannot use historical archives to improve its own physics model. Improving it takes expert effort § Introduction. Machine-learning forecasts had been improving but, at medium range, had not matched the top operational system at high resolution. The paper's stated target is a deterministic 10-day global forecast at 0.25 degrees that beats HRES on a broad verification set and is cheap to run § Introduction § GraphCast.

The approach assumes that ERA5 reanalysis is a usable ground truth for training and evaluation. HRES is evaluated against a separate HRES-based ground truth, because ERA5 is built from the same inputs as the model and would penalise HRES at short lead times app. §1.2 app. §5.2.

Method

Data and task. The model works on a 0.25-degree latitude-longitude grid of 721 by 1440 points. Each point holds 5 surface variables and 6 atmospheric variables on 37 pressure levels, which gives 227 target values per point app. §2.3. Time steps are 6 hours, and the horizon is up to 40 steps (10 days) app. §2.2. The model reads two consecutive states, the current one and the one six hours earlier, and predicts the next state app. §3.1. Forcing fields that can be computed analytically (incident solar radiation, and sine and cosine of time of day and year progress) and static fields (land-sea mask, surface geopotential) are added as inputs app. §3.3 tab. 1.

Forecast loop. Each one-step prediction is fed back as an input for the next step, so a forecast of any length is a chain of identical calls app. §3.1 fig. 1c. The pseudocode below follows this structure; it is our own rendering of the paper, not the released code.

# Roll the one-step model forward  [§ GraphCast] [app. §3.1]
# Input:  prev_state (t-6h), state (t), forcings_for(time)
# Output: list of forecast states, one per 6-hour step

def forecast(prev_state, state, num_steps):
    forecasts = []
    for step in range(num_steps):
        # Forcings are analytic functions of time  [app. §3.3]
        next_state = one_step(prev_state, state, forcings_for(step))
        forecasts.append(next_state)
        # Feed the prediction back in as the newest input  [app. §3.1]
        prev_state, state = state, next_state
    return forecasts

# One learned step: encode, process, decode, residual  [app. §3.4-§3.7]
# Input:  prev_state, state, forcings (three times), static fields
# Output: state at the next time step

def one_step(prev_state, state, forcings):
    # Normalise inputs per variable and level  [app. §3.7]
    grid_features = concat(
        normalise(prev_state),
        normalise(state),
        forcings,
        static_fields,
    )
    # Encoder: one message-passing step, grid to multi-mesh  [app. §3.4]
    mesh_latent, grid_latent = grid_to_mesh(grid_features)
    # Processor: 16 layers with unshared weights on the multi-mesh  [app. §3.5]
    for layer in processor_layers:
        mesh_latent = layer(mesh_latent)
    # Decoder: one message-passing step, mesh back to grid  [app. §3.6]
    grid_latent = mesh_to_grid(mesh_latent, grid_latent)
    # Output MLP gives 227 normalised values per grid point  [app. §3.6]
    normalised_change = output_mlp(grid_latent)
    # Scale the change back to physical units, then add it to the state  [app. §3.7]
    change = normalised_change * std_of_six_hour_change_by_level
    return state + change

Multi-mesh. Start from a regular icosahedron with 12 nodes and 20 faces. Refine it six times, splitting each triangle into four and projecting new nodes onto the sphere. The finest mesh has 40,962 nodes. Edges from every refinement level are kept, so coarse edges give long-range links and fine edges give local ones app. §3.3 fig. 1g.

Encoder app. §3.4 fig. 1d. Node, edge and mesh features are embedded by MLPs. One interaction-network step moves information from grid nodes to mesh nodes along grid-to-mesh edges. Grid-to-mesh edges connect every grid point to mesh nodes within a radius of 0.6 times the longest finest-mesh edge. Residual connections are used.

Processor app. §3.5 fig. 1e. The same message-passing layer structure is applied 16 times, with unshared weights in each layer. Each layer updates the mesh edges from their endpoints, then the mesh nodes from their incoming edges, with residual connections.

Decoder app. §3.6 fig. 1f. One interaction-network step moves information from mesh nodes back to grid nodes along mesh-to-grid edges. Each grid point is linked to the three vertices of the finest-mesh triangle that contains it. An output MLP gives 227 values per grid point.

Residual output. The output is a normalised change, which is scaled by a per-level standard deviation of six-hour differences and added to the current state app. §3.7 eq. output in app. §3.6.

Normalisation. Inputs are normalised per variable and level with mean and standard deviation over 1979–2015 app. §3.7. Mesh edge features are normalised by the longest edge app. §3.3.

MLPs. Every MLP has one hidden layer of width 512 and uses the swish activation. LayerNorm follows every MLP except the output MLP of the decoder app. §3.7.

Training objective. The loss is a mean squared error weighted by latitude cell area, by pressure level, by per-variable loss weights and by a per-variable inverse variance of six-hour differences. It is averaged over lead times and the batch, and it is computed on rollouts of up to 12 steps app. §4.2 app. §4.3. The paper gives the loss as a formula; the sketch below is our own rendering.

# Autoregressive weighted-MSE training loss  [app. §4.2] [app. §4.3]
# Input:  batch of trajectories, rollout length (up to 12)
# Output: scalar loss; gradients flow through every rollout step

def training_loss(batch, rollout_steps):
    prev_state, state = batch.initial_states
    total = 0.0
    for step in range(rollout_steps):
        pred = one_step(prev_state, state, batch.forcings[step])
        target = batch.targets[step]
        # Squared error, weighted per variable, per level and per latitude
        # cell  [app. §4.2]
        total = total + weighted_squared_error(pred, target)
        prev_state, state = state, pred
    # Average over lead times and the batch  [app. §4.2]
    return total / rollout_steps

Training schedule. There are three phases app. §4.5:

1. 1,000 updates with one autoregressive step. The learning rate rises linearly from 0 to 1e-3.
2. 299,000 updates with one autoregressive step. The learning rate follows a half-cosine decay to 0.
3. 11,000 updates in which the rollout length rises from 2 to 12, by one step every 1,000 updates, at a fixed learning rate of 3e-7.

Evaluation. Skill is measured by latitude-weighted RMSE and anomaly correlation coefficient (ACC) on targets every 12 hours, for forecasts initialised at 06z and 18z app. §5.2 app. §5.3. The headline scorecard covers 69 variable-level pairs at 20 lead times, which gives the 1380 targets § Verification methods fig. 2.

Hyperparameters

Results

Author-reported. Headline setup: forecasts initialised at 06z and 18z in 2018, ERA5 as ground truth for GraphCast and HRES-fc0 for HRES, RMSE latitude-weighted, 1380 targets § Verification methods app. §5.2.

Setup for the severe-event results: cyclones are tracked with an ECMWF-style tracker and scored against IBTrACS § Severe event forecasting results app. §8.1. Atmospheric-river results cover coastal North America and the Eastern Pacific in cold months (October to April) app. §8.2. Extreme-heat results cover summer months over land app. §8.3. Exact figure values are not stated in the text.

Paper vs code

Code reviewed: https://github.com/deepmind/graphcast at f2f2c5117d2f864e2d5e7c2f9f220db5e1049dd3 (HEAD on the clone date). The repository is now a multi-model monorepo, and the GraphCast code lives in weathernext/weathernext1_graph/. It contains inference-side model code and loss code only.

• Training code is absent. The repository has no optimiser, learning-rate schedule, rollout training loop or gradient clipping. A search for AdamW, Optax and clipping returned nothing. Everything in the training section is paper-only and cannot be checked against code.
• Model size is not in the repository. ModelConfig takes mesh_size, latent_size, gnn_msg_steps and hidden_layers as arguments graphcast.py. The released checkpoint values are not in the repo. The demo notebook's slider defaults (mesh 4, latent 32, four steps) are not the paper's values graphcast_demo.ipynb. The paper's 6, 512 and 16 can therefore not be confirmed or refuted from code.
• Released weights include a model not in the paper. The demo lists GraphCast_operational, described as "fine-tuned on HRES data from 2016 to 2021" graphcast_demo.ipynb. The paper describes only ERA5-trained models, so the two are not the same model.
• Loss weights match. 2 m temperature has weight 1.0 and the 10 m winds, MSL and precipitation have 0.1 graphcast.py app. §4.2.
• Level weighting matches. Weight is proportional to pressure, normalised to mean 1 losses.py app. §4.2.
• Residual output and output normalisation match. Predictions are added to the last input, and the output is scaled by a per-level standard deviation of differences normalization.py app. §3.7.
• Grid-to-mesh radius matches. The demo sets the fraction to 0.6 graphcast_demo.ipynb app. §3.3.
• Mesh-to-grid links match. Each grid point connects to the vertices of the triangle that contains it graphcast.py app. §3.3.
• Processor and mesh layers match in structure. The mesh GNN is applied gnn_msg_steps times graphcast.py. The value 16 is paper-only (see above).
• Pressure levels. The 37-level task and the 13-level WeatherBench task are both defined graphcast.py variables.py. The 13 levels match the scored levels § Verification methods.
• Precipitation. Precipitation is a target in the code and is weighted 0.1 graphcast.py. The paper also outputs it and excludes it from scoring app. §3.3 § Verification methods. A variant without precipitation input exists in the code graphcast.py.
• Edge normalisation. Mesh-edge normalisation uses the longest finest-mesh edge graphcast.py graphcast.py app. §3.3.

Pitfalls

• Training is not reproducible from the code. The inverse-variance factor per variable and level is described in words only, and its values are not stated app. §4.2. The reported per-variable weights are stated, but there is no training loop to check them against. An implementer must write the rollout trainer from the paper alone.
• Training years are ambiguous. The text gives 1979–2015 for training and 2016–2017 for validation, and calls 1979–2017 the development data. It does not say whether the final weights were trained on the full 1979–2017 range. Pick one and state it in your implementation.
• Learning-rate phases are easy to get wrong. Phase 1 is a linear rise, phase 2 is a half-cosine decay to zero, and phase 3 is a fixed 3e-7 with a rollout that grows by one step every 1,000 updates app. §4.5. The paper does not state an optimiser warm-up beyond phase 1.
• Memory limits the rollout. A 12-step rollout with backpropagation through time does not fit in 32 GB without bfloat16 activations and gradient checkpointing app. §4.3 app. §4.6.
• Protocol changes the answer. The headline comparison uses HRES-fc0 and 06z/18z initialisations, with targets every 12 hours. The Pangu comparison uses 00z/12z initialisations and 6-hourly targets, so its 99.2% figure is not comparable to the 90.3% headline app. §5.2 app. §6.
• Discontinuity at 3.5 days. HRES 06z/18z forecasts stop at 3.75 days, so the comparison switches to 00z/12z initialisations after about 3.5 days. The skill-score curve jumps there § Verification methods fig. 2b.
• Blur at long lead times. The MSE objective pushes the model toward smooth predictions at long leads. Optimal blurring of HRES narrows the gap but does not close it § Forecast verification results app. §7.4. Sharp-field users should check this.
• Stratosphere is weaker. HRES wins most often at 50 hPa. The loss weight for levels at or above 50 hPa is only 0.66% of the total app. §4.2 § Forecast verification results.
• Deterministic only. The authors note that the MSE objective makes GraphCast express uncertainty by blurring, and that no ensemble is provided § Conclusions.
• Precipitation is not scored. ERA5 precipitation has known biases, so precipitation is excluded from verification even though it is an output § Verification methods app. §3.3.
• Mesh edge count depends on refinement. The multi-mesh has 327,660 edges. The finest-mesh-only ablation has 245,760 app. §3.3 app. §7.3.1. Implement all refinement levels, not just the finest one.
• Grid-to-mesh radius is a trade-off. The paper says 0.6 reduces the number of grid points feeding several mesh nodes, which cuts memory, while larger values give better predictions. The code exposes this as a config value graphcast.py.
• Poles repeat. At the poles the 1440 longitudes are identical, so the number of distinct grid points is smaller than 1,038,240 app. §2.3.
• Forcing details. The incident solar radiation forcing is accumulated over one hour, and time-of-day and year-progress features are sine and cosine normalised to [0, 1) app. §3.3.
• Evaluation is 12-hourly, rollout is 6-hourly. Targets are scored only every 12 hours, so the lead-time axis has 20 points, not 40 app. §5.2 fig. 2.
• Figure-only numbers. Cyclone track errors, atmospheric-river RMSE and extreme-heat precision-recall values are shown in figures; the text gives only the summary percentages fig. 3.
• Appendix numbering. Section and subsection numbers above come from the arXiv LaTeX source (appendix sections 1 to 8). The journal typesetting may number them differently. Main-text sections are unnumbered in the source, so they are cited by heading.

Lineage

• Builds on: GNN-based learned simulators for fluids and physics, which the authors adapt to weather app. §3.2. These are cited, not carded: Sanchez-Gonzalez et al. (2020) and Pfaff et al. (2021).
• Builds on: the encode-process-decode GNN framework, from Battaglia et al. [arXiv:1806.01261], cited for the architecture app. §3.2.
• Builds on: a GNN weather forecaster by Keisler [arXiv:2202.07575], cited as a similar approach with promising results app. §3.2.
• Compared against: ECMWF's high-resolution deterministic forecast (HRES) § Forecast verification results, scored against HRES-fc0 app. §1.2.
• Compared against: Pangu-Weather [arXiv:2211.02556], on 252 targets under the Pangu protocol app. §6.