Extending AI Architectures to Address Continuous Scientific Problems

Artificial intelligence is proving to be transformative in its ability to work with language and images. Now, with a growing push to apply AI to scientific discovery, Caltech's Anima Anandkumar says there is a crucial ingredient missing from most AI models: the ability to understand the physical world. Take as an example weather models, says Anandkumar, Caltech's Bren Professor of Computing and Mathematical Sciences. If you want an AI model to predict weather, it must understand chaotic physical systems, like how the atmosphere changes around the planet and over time.

The best way to gain that understanding is to learn the continuous functions, the underlying mathematical relationships, that fully describe those systems. But many AI models were created for language processing or computer vision (where AI learns to process visual data). In those categories, data can be treated as collections of isolated points, such as words in a sentence or a fixed set of pixels in an image.

To address this mismatch, Anandkumar and her colleagues from Caltech and tech giant NVIDIA recently described a framework that shows how to extend existing neural network architectures in a way that allows AI systems to learn continuous functions, which are needed to make predictions that are grounded in the physical world. A paper about the framework was published in the journal Nature Machine Intelligence.

Anandkumar first introduced these models, called neural operators, in 2020.

"If you think of weather, it happens everywhere around us. It's not just at the points where the satellites are or where measurements are taken at certain locations. It happens everywhere in between," Anandkumar says. "If you design neural architectures that also directly approximate everywhere in the continuous space—and that imagine what the continuous world looks like rather than only at a few points—that is much more faithful in capturing the physical world."

In the past, Anandkumar and colleagues successfully applied neural operators in many other scientific contexts, ranging from modeling the atomic scale of quantum chemistry to simulating a black hole. They have shown that, by using a neural operator, they can, for example, simulate when plasma in a fusion reactor will suddenly lose its heat or get disrupted—and do so a million times faster than with standard numerical simulations. They have also used a neural operator to model what happens to carbon dioxide when it is captured and stored underground.

In the new paper, Anandkumar and her colleagues offer a systematic way to modify existing neural network architectures—changing them from systems that learn relationships between a fixed grid of discrete data points to a neural operator that learns how an entire physical system evolves over space and time. This aligns the architectures more naturally with the continuous world. "We realized that there was this much more general abstraction where you don't have to do the hard work of coming up with a specific architecture from scratch," Anandkumar explains. "You can take what researchers have already done with computer vision and language when they weren't thinking of the physical world and adapt it to make it work in the physical world in a very principled way."

The team tested the framework on standard physics simulations that model how fluids move, including one that is commonly used to describe the motion of air or water, and another that models how fluids flow through porous materials. They found that neural operators designed according to the new framework maintained their accuracy even at resolutions they had never seen during training.

To understand this, consider the problem of creating a map of temperatures across California. Since data is not available at every single point, you could divide the state into a grid where each square in the grid records a single temperature value. Maybe you start with a 128 x128 grid for a total of 16,384 measurements. But you could divide the state into smaller squares, creating a 256 x 256 grid with 65,536 points. That would offer a higher-resolution method to represent the same overall temperature. Alternately, you could create a lower-resolution grid with only 4,096 points in a 64 x 64 grid.

In the paper, the team trained neural operators on data in 128 x 128 grids. But they found that their models remained accurate at higher and lower resolutions (64 x 64 grids and 256 x 256 grids) while conventional neural networks generally became less accurate when the resolution changed.

The result, Anandkumar says, is that neural operators produce more accurate outputs and predictions that are more trustworthy in a wide range of scientific fields where data are best represented as continuous systems. "This way, you don't have to make discrete approximations that could then be wrong on the points in between the grid," Anandkumar says. "This provides a blueprint for adapting advances in AI to the scientific domain."

The lead authors of the paper, "Principled approaches for extending neural architectures to function spaces for operator learning," are Julius Berner, a former Caltech postdoctoral scholar who is now at NVIDIA, and Miguel Liu-Schiaffini (BS '25), who completed the work as an undergraduate at Caltech and is now a graduate student at Stanford University. Other authors are Valentin Duruisseaux, a postdoctoral scholar at Caltech; Jean Kossaifi and Boris Bonev of NVIDIA; and Kamyar Azizzadenesheli of Project Prometheus. Anandkumar's work was supported by the Bren endowed chair, the Office of Naval Research, and the AI2050 Senior Fellow program at Schmidt Sciences. Liu-Schiaffini was supported in part by a Mellon Mays Undergraduate Fellowship.