Building My Own Tiny AI Weather Model

This article explores how I built my own tiny AI weather model based on the Graphcast architecture.

I built and trained a tiny AI weather forecasting model inspired by the Graphcast architecture:

The research, figures and text are my own work, LLMs were used as spell check and for layout of the web page.

This article aims to provide an overview of how a basic AI weather model (based on GNNs) works. To build your own model, please refer to the Graphcast paper [2], or feel free to send me an email if you have any questions about the implementation aspects.

The Earth’s atmosphere: a chaotic system which impacts everyone on Earth, large parts of which we don’t yet fully understand. Yet, advances in meteorology, observations and computing have resulted in a spectacular increase in the accuracy of weather forecasts over the past decades [0]. The main drivers behind this progress have been numerical weather prediction (NWP) models. These models use physical relations between meteorological variables to make forecasts up to two weeks in the future. The codebases for these models have grown to millions of lines, and a single forecast requires up to an hour on a powerful supercomputer [1].

Now, a new era has arrived with data-driven weather forecasts, with many of the SOTA AI weather models now outperforming classical NWP models on medium-range weather forecasts, while being orders of magnitude more efficient [2, 3].

Some other advantages of data-driven models: their inner workings can be explained in a single article like this (and you don’t need a degree in meteorology to understand it), you can create your own in fewer than 1000 lines of code, you don’t need a supercomputer to run it, and as an added benefit, you don’t need to dabble in old Fortran code. Let’s see how I build my own tiny weather model based on the Graphcast architecture.

There are many different timescales when talking about weather/climate forecasting, from predicting if it will rain in 10 minutes (nowcasting) to predicting the change in the Earth’s temperature over the next 50 years (climate modelling). Here, I will be looking at medium-range weather forecasts, which concern the synoptic-scale evolution of the atmosphere over the coming 3–10 days.

To train a data-driven weather model, you obviously need a lot of historical weather data. Luckily, besides predicting future conditions, NWP methods have also been used to assimilate observations into a neatly structured dataset. This manifests itself in the ERA5 dataset [4]: the holy grail of weather data for data-driven medium-range weather forecasting.

The ERA5 dataset is, in essence, a 5-dimensional tensor, with the following dimensions:

In total, this comes down to about 5 petabytes of data. In my tiny weather model, I will, of course, not be using even a fraction of a percent of that. Even in the SOTA models, some of the dimensions are restricted to keep the training under control.

Many deep learning methods can be applied to weather data: from transformers and convolutional neural networks to diffusion models. Many of these methods have their strengths and weaknesses and shine in their own niche of data-driven weather modelling.

Regardless of the method, most data-driven medium-range weather models work with the encode-process-decode architecture:

For medium-range weather forecasting, the SOTA weather models (for example, ECMWF’s AIFS and Google’s WeatherNext) are based on graph neural networks (GNNs). The small model I will build is a GNN, largely inspired by Graphcast (2022) [2], which was one of the first data-driven weather models that went mainstream. Since then, many improvements and additions have been made, but the core idea remains the same.

As their name might suggest, GNNs operate on graphs. They use the topological structure of the graph to model spatial relationships and exchange information between nodes. Through the process of message passing, the data of neighboring edges and nodes is aggregated at all nodes on the graph to update their internal representations. This allows the model to attend to the information around it, while limiting the exchange of data to the neighbourhood to keep it efficient.

As discussed in the data section, the data comes in a latitude-longitude grid. Every cell of this grid has an input vector with the meteorological variables and some other useful metadata (position, time, …). For simplicity, the atmospheric levels are merged with the variable dimension. The input and output of the model are in this grid format, but for the internal representation in the processing step, we use a mesh.

The mesh is typically modelled as an icosahedral mesh. An icosahedral mesh positions the nodes to form (almost) equal-area triangles (unlike a typical latitude-longitude grid, which has a higher resolution near the poles). The smallest icosahedral mesh is an icosahedron (i.e. a D20 die), which can be recursively subdivided to increase the number of nodes and edges. Like in Graphcast, I will use a ratio of about 25 grid cells per mesh node; this ensures that the processing computation stays relatively efficient.

To use a GNN, we need a graph. More specifically, we need three separate graphs:

In the Graphcast model, the mesh graph is modelled as a multiscale graph by adding long-distance connections between nodes in the graph. This way, longer-distance connections (teleconnections) can be modelled, and the representation of bigger meteorological processes can be shared among nodes. To keep it simple, and because of the limited resolution of the data I will use, my model does not use a multiscale graph.

For a single message-passing iteration on a graph G=(V,E)G=(V,E)G=(V,E), consider a directed edge from node iii to node jjj. We first update its edge embedding with the embeddings of both adjacent nodes.

Then we update each node embedding by summing over its incoming edges, where N(j)\mathcal{N}(j)N(j) is the set of nodes with an edge to node jjj.

​[vj​,i∈N(j)∑​eij′​]​ This process updates the nodes with information from the neighbouring nodes. To use information from nodes which are further away, multiple processing layers are needed.

Now that we have established how to perform message passing, the encoder, processor and decoder can be implemented easily.

Now we can bring everything together in the full model architecture:

To make it slightly easier for the model, the output predicts the residuals of the meteorological variables we want to predict. These are then added to the initial values to get the actual forecast.

The full model consists of about 200k weights, about 115k of which are for the processor, and 30k for the decoder and encoder each. For comparison, the Graphcast model has about 37M weights (my tiny model is about 0.5% the size of Graphcast).

The main hyperparameters of the model are the following:

In my tiny model, I use 444 recursive subdivisions on the icosahedral graph (which results in 241 nodes on my domain), E=64E=64E=64, P=4P=4P=4, and a two-layer ReLU activated neural network with hidden size equal H=64H=64H=64.

To avoid waiting for multiple weeks for training to end, I chose a very small subset of the ERA5 dataset to use for training. My aim was to try to simulate the evolution of extratropical cyclones over the North Atlantic Ocean and Europe. I limited the domain to this area and chose 9 meteorological variables which should be relevant to this process (temperature, u and v wind components and geopotential at both the 500 hPa and 850 hPa levels, and mean sea level pressure), which I will call XtX_tXt​ for some timestep ttt. I also decided to use the 1° dataset (instead of the 0.25° full resolution) in timesteps of 6 h. I limited the time dimension to the years 2003–2016 for training and 2017–2019 for evaluation.

After limiting the size in all these dimensions, the total dataset came out to about 5 GB (about 1,000,000 times smaller than the full dataset). The data was downloaded through the WeatherBench 2 data repository [5].

Unlike the connections between nodes, the node positions are not explicitly encoded in the graph. The way we include them in the model is as extra input features. For every grid and mesh node, we add the sine and cosine of the longitude and latitude as input features, which are then embedded into the latent space. For the edges, we add the distance and the sine and cosine of the direction.

Besides the spatial features and the meteorological variables, there are other useful features to take into account for the input grid cells:

Unlike the meteorological variables, the spatial and static features are only used as input features and are not part of the output of the model. The input for the model becomes:

Since the meteorological variables used in the model have different scales and distributions, they have to be standardized (z-score normalization). The other input features are inherently normalized (cos/sin features) or are min-max normalized (distance and altitude).

To train the model, I will let the model predict the weather conditions for the next timestep Yt+1=f(Xt′)Y_{t+1}=f(X't)Yt+1​=f(Xt′​), where fff is my model and Xt′X'tXt′​ is the input data as discussed earlier. The predicted values are then compared with the target data from the next timestep Xt+1X{t+1}Xt+1​, and the loss is calculated as a function of the predicted and target values l(Xt+1,Yt+1)l(X{t+1},Y_{t+1})l(Xt+1​,Yt+1​).

The loss is then used to backpropagate the gradients and update the weights accordingly using stochastic gradient descent. I train the model in chronological order to avoid leaking future data as much as possible. This whole idea is very similar to the next-token prediction methodology used in pretraining LLMs.

The loss function calculates the root mean square error (RMSE) between the prediction and target for every grid cell. Graphcast uses some other tricks in its loss function, most of which I will omit here. The only addition to the standard RMSE loss is to weight the loss for every cell by its cell size a(ϕ)=cos⁡(ϕ)a(\phi)=\cos(\phi)a(ϕ)=cos(ϕ) to account for the grid cells being much smaller near the poles.

​ Training time

After some systematic tests with different hyperparameter settings, I ended up training the model for 20 epochs with a batch size of 8, an embedding size of E=64E=64E=64, and P=4P=4P=4 processing layers. The training took only 45 min on my 5-year-old gaming PC (the typical SOTA AI weather models take multiple weeks to train on dozens of powerful GPUs).

I evaluate the model on a held-out test dataset consisting of the years 2017–2019 (inclusive). The model is evaluated using the same metric as the loss function (latitude-weighted RMSE). To see if the model actually learned some patterns and acquired some skill, I test it against the persistence baseline, where Xt+1=XtX_{t+1}=X_tXt+1​=Xt​.

The goal of a weather forecast is to predict what the weather conditions will be in the future. We are typically not only interested in the next timestep (in this case, 6 h ahead), but also in the forecast for multiple days ahead. Very similarly to LLMs, we can roll out the prediction by using the prediction for timestep t+1t+1t+1 as input to predict timestep t+2t+2t+2. This process can be repeated to roll out the forecasts Yi=f(Yi−1,Yi−2)Y_i=f(Y_{i-1},Y_{i-2})Yi​=f(Yi−1​,Yi−2​). Of course, the errors will increase with every timestep, and at some point the forecast will no longer have any skill over climatological or persistence baselines (typically after about 10–14 days for current SOTA NWP and AI weather models).

When evaluating the model for a lead time of +72 h (12-timestep rollout), I find that the model outperforms the naive persistence baseline over the whole period for all meteorological variables. You can see that the error increases about linearly as a function of the lead time. The gain over the baseline decreases over time, which makes