{"slug": "building-my-own-tiny-ai-weather-model", "title": "Building My Own Tiny AI Weather Model", "summary": "A developer built a working data-driven weather forecasting model in fewer than 1000 lines of code based on the Graphcast architecture, training it on a small fraction of the 5-petabyte ERA5 dataset rather than the full archive. The model follows the encode-process-decode pattern used by state-of-the-art systems such as ECMWF's AIFS and Google's WeatherNext, which are graph neural networks that now outperform classical numerical weather prediction on medium-range forecasts while running orders of magnitude more efficiently than NWP models that can take up to an hour on a supercomputer per forecast. The project shows that a single-GPU-scale AI weather model is reproducible without Fortran codebases or a meteorology degree.", "body_md": "# Building My Own Tiny AI Weather Model\n\nThis article explores how I built my own tiny AI weather model based on the Graphcast architecture.\n\n## Introduction and Background\n\nThe 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](#ref-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](#ref-1)].\n\nNow, 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](#ref-2), [3](#ref-3)].\n\nSome 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.\n\n## Data\n\n### The holy grail of weather data\n\nTo 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](#ref-4)]: the holy grail of weather data for data-driven medium-range weather forecasting.\n\n### The data structure\n\nThe ERA5 dataset is, in essence, a 5-dimensional tensor, with the following dimensions:\n\n- Latitude x longitude: the Earth is divided into latitude-longitude grid cells\n- Level: every grid cell has a number of vertical (atmospheric) levels\n- Variable: a number of variables are available for every cell (think temperature, wind, humidity, …)\n- Time: all this data is available in hourly timesteps since 1940\n\nIn 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.\n\n## Model\n\n### How deep learning tackles the problem\n\nMany 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.\n\nRegardless of the method, most data-driven medium-range weather models work with the encode-process-decode architecture:\n\n1. Encode: encode the grid data into an internal/latent representation (like in LLMs, this is typically called the embedding)\n2. Process: update the embeddings based on relevant (neighboring) data; this is where most of the difficult spatial connections are learned\n3. Decode: bring the latent representation back to the original grid to compute the actual forecast (often residuals instead of absolute values)\n\n### Graph neural networks\n\nFor 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](#ref-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.\n\n### Grids and meshes\n\nAs 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.\n\nThe 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.\n\n### The graphs\n\nTo use a GNN, we need a graph. More specifically, we need three separate graphs:\n\n1. Grid2Mesh graph: this bipartite graph connects the input grid cells (seen as nodes in the graph) with the mesh nodes and is used in the encoding step. Every grid node is connected to its closest mesh node.\n2. Mesh graph: as discussed previously, the mesh is the internal/latent representation of the model, used to do repeated message passing in the processor. As opposed to the other two graphs, the edges are bidirectional to allow message passing in both directions.\n3. Mesh2Grid graph: another bipartite graph connecting the mesh back to the grid nodes. Here, every grid node is connected to the mesh nodes of the triangle it is inside.\n\nIn 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.\n\n### A single layer of message-passing\n\nFor a single message-passing iteration on a graph , consider a directed edge from node to node . We first update its edge embedding with the embeddings of both adjacent nodes.\n\nThen we update each node embedding by summing over its incoming edges, where is the set of nodes with an edge to node .\n\nThis process updates the nodes with information from the neighbouring nodes. To use information from nodes which are further away, multiple processing layers are needed.\n\n### Encode-process-decode in GNNs\n\nNow that we have established how to perform message passing, the encoder, processor and decoder can be implemented easily.\n\n- Encoder: we perform a single message-passing step on the bipartite Grid2Mesh graph to update the mesh node embeddings\n- Processor: we perform P iterations of message passing on the mesh graph to update the mesh nodes with information from mesh nodes in the neighbourhood\n- Decoder: we perform a single message-passing step, this time on the Mesh2Grid graph, to update the grid nodes again\n\n### The full architecture\n\nNow we can bring everything together in the full model architecture:\n\n1. Embed: the input features of all nodes and edges are embedded using MLPs\n2. Encode: update grid and mesh nodes\n3. Process: multiple layers of message passing (with unshared weights) on the mesh graph\n4. Decode: update the grid nodes based on the processed mesh nodes\n5. Output: transform the embeddings on the grid nodes into output features\n\nTo 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.\n\nThe 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).\n\n### The hyperparameters\n\nThe main hyperparameters of the model are the following:\n\n- The number of subdivisions on the icosahedral mesh\n- Embedding size (the same for all nodes and edges for simplicity)\n- The number of processor layers\n- The architecture of the MLPs (typically a simple two-layer MLP)\n\nIn my tiny model, I use recursive subdivisions on the icosahedral graph (which results in 241 nodes on my domain), , , and a two-layer ReLU activated neural network with hidden size equal .\n\n## The Features\n\n### The ERA5 subset\n\nTo 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 for some timestep . 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.\n\nAfter 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](#ref-5)].\n\n### Spatial features\n\nUnlike 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.\n\n### Static features\n\nBesides the spatial features and the meteorological variables, there are other useful features to take into account for the input grid cells:\n\n- Time features: we add the local solar time and day of the year (both as cosine and sine values, because they are cyclical)\n- Geographic features: we add the altitude of every grid cell, a land/sea mask and the incident solar radiation at the top of the atmosphere for each grid cell\n\nUnlike 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:\n\n### Normalization\n\nSince 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).\n\n## Training\n\n### Next timestep prediction\n\nTo train the model, I will let the model predict the weather conditions for the next timestep , where is my model and is the input data as discussed earlier. The predicted values are then compared with the target data from the next timestep , and the loss is calculated as a function of the predicted and target values .\n\nThe 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.\n\n### Loss function\n\nThe 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 to account for the grid cells being much smaller near the poles.\n\n### Training time\n\nAfter 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 , and 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).\n\n## Evaluation and Results\n\n### Evaluation\n\nI 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 .\n\n### Rollout\n\nThe 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 as input to predict timestep . This process can be repeated to roll out the forecasts . 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).\n\n### Beating the baseline\n\nWhen 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 sense due to the accumulation of errors. Compared to Graphcast (I’m not comparing apples to apples, but it gives an idea), the RMSE of my model is about 5x worse than Graphcast at a lead time of 24 h, and about 10x worse at a lead time of 72 h. Of course, this makes sense given that I used a simpler and much smaller model, a couple of orders of magnitude less data and trained for only 45 min.\n\n### Some examples\n\nBelow are some examples of rollouts up to +72 h for some random dates in the test set. A single rollout is produced in a matter of seconds. You can see that the evolution and movement of the pressure systems look relatively realistic, even though at longer lead times the model starts diverging from the target.\n\n### Further improvements\n\nA variety of improvements can be made to improve the skill of this model. The obvious thing is to increase the amount of data (more variables, more years of data…), increase the model size and implement numerous other features (multiscale graphs, a more advanced loss function, …).\n\nAnother interesting observation is that the evaluation errors are highest near the western boundary of the domain. This is likely because the origin of the weather systems over the Atlantic lies over the United States (i.e. to the west of my bounding box). This is also visible in some of the example rollouts, in which whole pressure systems are missed in the western part of the bounding box. The forecasts are likely to improve significantly if the model is trained on a whole hemisphere or on the whole world, or boundary conditions are used.\n\nLastly, another interesting approach, which is the subject of my master's thesis, is to not use a static (icosahedral) graph, but a graph which changes its node positions. Since the weather patterns are very variable, I hypothesise that this will improve the skill of AI weather models by putting the compute where it is needed most (i.e. in highly dynamic areas).\n\n## References\n\n1. Ritchie, H. (2024, March 12). [*Weather forecasts have become much more accurate; we now need to make them available to everyone*](https://ourworldindata.org/weather-forecasts) . Our World in Data.\n2. Destination Earth. (2025). [*Artificial Intelligence in DestinE – The explainer*](https://destine.ecmwf.int/artificial-intelligence-in-destine-the-explainer/) . European Centre for Medium-Range Weather Forecasts.\n3. Lam, R., Sanchez-Gonzalez, A., Willson, M., et al. (2023). [Learning skillful medium-range global weather forecasting](https://doi.org/10.1126/science.adi2336) .*Science, 382* (6677), 1416–1421.\n4. Lang, S., Alexe, M., Chantry, M., et al. (2024). [AIFS—ECMWF’s data-driven forecasting system](https://doi.org/10.48550/arXiv.2406.01465) .*arXiv* .\n5. Hersbach, H., Bell, B., Berrisford, P., et al. (2020). [The ERA5 global reanalysis](https://doi.org/10.1002/qj.3803) .*Quarterly Journal of the Royal Meteorological Society, 146* (730), 1999–2049.\n6. Rasp, S., Hoyer, S., Merose, A., et al. (2024). [WeatherBench 2: A benchmark for the next generation of data-driven global weather models](https://doi.org/10.1029/2023MS004019) .*Journal of Advances in Modeling Earth Systems, 16* , e2023MS004019.\n\n[← Back to all notes](https://tcumulus.dev/blog)", "url": "https://wpnews.pro/news/building-my-own-tiny-ai-weather-model", "canonical_source": "https://tcumulus.dev/blog/building-my-own-tiny-ai-weather-model/", "published_at": "2026-10-28 13:59:21+00:00", "updated_at": "2026-10-28 14:12:02.597291+00:00", "lang": "en", "topics": ["machine-learning", "artificial-intelligence", "ai-research", "neural-networks"], "entities": ["Graphcast", "ERA5", "ECMWF", "AIFS", "WeatherNext", "Google"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/building-my-own-tiny-ai-weather-model", "markdown": "https://wpnews.pro/news/building-my-own-tiny-ai-weather-model.md", "text": "https://wpnews.pro/news/building-my-own-tiny-ai-weather-model.txt", "jsonld": "https://wpnews.pro/news/building-my-own-tiny-ai-weather-model.jsonld"}}