Although traditional climate models have been improving over the decades, they often generate errors and have biases due to scientists’ incomplete understanding of how Earth’s climate works and how the models are built.
These models divide the globe into cubes — typically 50–100 km on each horizontal side — that extend from the surface up into the atmosphere, and then predict what happens to the weather in each cube over a stretch of time. To make predictions, they calculate how air and moisture move based on well-established laws of physics. But many important climate processes, including clouds and precipitation, vary over much smaller scales (millimeters to kilometers) than the cube dimensions used in current models and therefore cannot be calculated based on physics. Scientists also lack a complete physical understanding of some processes, such as cloud formation. So these traditional models don’t rely on first principles alone and instead use simplified models to generate approximations, called parameterizations, to simulate the small-scale and less understood processes. These simplified approximations inherently limit the accuracy of physics-based climate models.
Like a traditional model, NeuralGCM divides the Earth’s atmosphere into cubes and runs calculations on the physics of large-scale processes like air and moisture movement. But instead of depending on parameterizations formulated by scientists to simulate small-scale aspects like cloud formation, it uses a neural network to learn the physics of those events from existing weather data.
A key innovation of NeuralGCM is that we rewrote the numerical solver for large-scale processes from scratch in JAX. This allowed us to use gradient-based optimization to tune the behavior of the coupled system “online” over many time-steps. In contrast, prior attempts to enhance climate models with ML struggled greatly with numerical stability, because they used “offline” training, which ignores critical feedback between small- and large-scale processes that accumulates over time. Another bonus of writing the entire model in JAX is that it runs efficiently on TPUs and GPUs, in contrast to traditional climate models that mostly run on CPUs.