Climate and weather modeling has long been a staple of high-performance computing, but as meteorologists look to improve the speed and resolution of their forecasts, they are increasingly embracing machine learning.
In a paper published in Nature this week, a team from Google and the European Centre for Medium-Range Weather Forecasts (ECMWF) detailed a new approach that uses machine learning to overcome the limitations of existing climate models and generate forecasts that are faster and more accurate than existing methods.
The model, called NeuralGCM, was developed using historical weather data collected by ECMWF and uses neural networks to augment more traditional HPC-style physics simulations.
As Stephan Hoyer, a member of the NeuralGCM staff, explained in a recent report, most climate models today make predictions by dividing the Earth into a cube measuring 50 to 100 kilometers on a side and simulating how air and moisture move within it based on known laws of physics.
NeuralGCM works in a similar way, but the added machine learning is used to track climate processes that are not always well understood or that occur on a smaller scale.
“Many important climate processes, including clouds and precipitation, vary on scales much smaller than the cube dimensions used in current models (millimeters to kilometers) and therefore cannot be calculated based on physics,” Heuer wrote.
Traditionally, these small-scale phenomena have been tracked using a series of simpler, second-order models, called parameterizations, Heuer explained. To compound the problem, he noted, “these simplistic approximations inherently limit the accuracy of physically-based climate models.”
In other words, these parameterizations are not necessarily the most reliable and may reduce the overall accuracy of the model.
NeuralGCM works by transposing these parameterizations into a neural network. The three models were trained on existing weather data collected by ECMWF between 1979 and 2019 at resolutions of 0.7, 1.4 and 2.8 degrees.
The results, according to the paper, are very promising: Using Google's WeatherBench2 framework, the team says NeuralGCM achieves accuracy on par with existing state-of-the-art forecast models up to five days at 0.7-degree resolution, and even more accuracy for 5- to 15-day forecasts at 1.4-degree resolution.
Meanwhile, the team concluded that their model could predict average temperatures between 1980 and 2020 within 2.8°C, with an average error rate that was one-third that of existing atmosphere-only models.
NeuralGCM also proved quite competitive against more targeted models like X-SHiELD, which Hoyer explains offers much higher resolution, albeit at a higher computational load.
Compared to X-SHiELD, the researchers found that the NeuralGCM 1.4 degree model was able to predict humidity and temperature with 15 to 20 percent less error in 2020. In the same tests, it was able to predict tropical cyclone patterns that matched the number and intensity of tropical cyclones observed that year.
Accelerated Prediction
The team didn't just translate these parameterizations into a neural network: the entire NeuralGCM was written in Google JAX, a machine learning framework that translates numerical functions into something usable in Python.
According to Hoyer, the move to JAX offered many benefits, including improved numerical stability during training and the ability to run models on TPUs or GPUs. In contrast, weather models have traditionally run on CPUs but are increasingly using GPUs, which we'll touch on a bit later.
Because NeuralGCM runs natively on the accelerator, Google claims its system can run orders of magnitude faster and cheaper.
“Our 1.4 degree model is over 3,500 times faster than X-SHiELD, meaning that if researchers simulate a year's worth of atmosphere with X-SHiELD, it would take 20 days compared to just eight minutes with the NeuralGCM,” Heuer wrote.
Additionally, Hoyer claims that he can run simulations on a single TPU rather than the 13,000 CPUs required to run X-SHiELD, and that he can even run NeuralGCM on a laptop if needed.
While promising, it's important to note that NeuralGCM is just a starting point, and Hoyer freely admits that it's not a complete climate model, but that seems to be the long-term goal.
“Ultimately, we hope to incorporate other aspects of the Earth's climate system, such as the oceans and carbon cycle, into the model, which would enable NeuralGCM to make predictions on longer timescales, allowing not only weather forecasts of days and weeks, but also on climate timescales,” Heuer wrote.
To support these efforts, the model source code and weights have been made publicly available on GitHub under a non-commercial license for amateur meteorologists to enjoy and use.
Machine learning gains momentum in climate modeling
This isn't the first time machine learning has been used for climate modeling — Nvidia's Earth-2 climate model is another example of how combining AI and HPC can not only improve the accuracy of predictions, but also accelerate them.
Announced at GTC this spring, Earth-2 is essentially a large-scale digital twin designed to use a combination of HPC and AI models to generate high-resolution simulations, down to two kilometers in resolution.
This is possible thanks to a diffusion model called CorrDiff, which Nvidia says can generate images of weather patterns at 12.5 times the resolution and up to 1,000 times faster than other numerical models. The result is a model that's fast and accurate enough to interest Taiwan, which is turning to the platform to improve its typhoon forecasts.
Meanwhile, more climate research centers are beginning to deploy GPU-accelerated systems, and climate research is one of several research areas targeted by the University of Bristol's 200 petaflops (FP64) Isambard-AI system.
Earlier this year, the European Mediterranean Climate Change Centre in Lecce, Italy, selected Lenovo to deploy a new Cassandra super, running on a small combination of Intel Xeon Max CPUs and Nvidia H100 GPUs, which the centre aims to use to run a range of AI-based climate simulations.®
