Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU

Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU

N
News Editor
2026-08-26 12:42:10
A research team led by Anima Anandkumar at the California Institute of Technology published a paper on Aug. 24 describing an AI-assisted approach to density functional theory, or DFT, that shifts a key computational step from cubic scaling to near-linear behavior. The work targets one of the longest-running bottlenecks in quantum chemistry, where simulation costs rise sharply as molecular or material systems grow. Instead of asking a model to predict the final answer in one shot, the team trained a Fourier Neural Operator, or FNO, to learn the forward map from input conditions to electron density, then inserted that model back into the standard self-consistent DFT loop. The paper reports complexity dropping from O(N^3) to O(N log N) for that step. In tests on drug-scale molecules outside the training set, the Kohn-Sham FNO posted a density error of 2.23%, versus 9.97% for a direct-prediction baseline; at 45 heavy atoms, the gap widened to 4% versus 41%. For large-scale validation, the researchers used magnesium dislocation structures. The paper says the new method converged on a system with 8,250 atoms and 82,500 valence electrons using a single NVIDIA B300 GPU. It contrasts that with a 2019 full DFT calculation on 6,164 magnesium atoms that used about 7,800 NVIDIA V100 GPUs on the Summit supercomputer.

A team led by Anima Anandkumar at the California Institute of Technology said in a paper published on Aug. 24 that it used an AI model to rework a core step in quantum chemistry calculations, pushing the scaling of that step from cubic toward near-linear. The paper reports that a metal-defect simulation involving 82,500 valence electrons was completed on a single GPU, while a comparable 2019 calculation used about 7,800 GPUs.

Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU 2

Simulating quantum systems sits underneath work on drug discovery, batteries and chips. The article notes that a brute-force calculation for the quantum behavior of 100 electrons would take longer than the age of the universe, while a drug molecule can involve thousands of electrons and battery materials can involve even more.

DFT has remained limited by scaling for decades

One of the most widely used tools in quantum chemistry is density functional theory, or DFT. It is used in drug screening, catalyst design and battery-material research.

The problem is speed. As systems grow, the cost of DFT rises roughly with the cube of system size. In the article’s example, a 10x increase in electron count leads to a 1,000x increase in computation. Researchers have looked for ways to accelerate that process for 60 years, but the article says no substantial breakthrough had emerged.

Earlier AI routes ran into extrapolation and stability issues

The paper reviews two broad AI strategies that had already been tried.

Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU 3

One was direct prediction: ask the model to output the final result in one pass. That worked reasonably well inside the training distribution, but extrapolation was weak. Once molecules became larger than those seen during training, the error grew sharply. In the paper’s tests, the error for a direct-prediction model rose from under 1% to about 10% when moving from small molecules to drug-scale molecules, with the worst case reaching 41%.

The second route followed a more traditional physics path and had AI learn an inverse mapping. According to the paper, that approach was mathematically unstable. In the control experiments, every calculation diverged after a few iterations.

The new approach keeps the iterative loop and swaps out the expensive solve

The team took a different route. DFT itself is an iterative procedure: each round starts from a set of input conditions, computes the electron density, updates the input from that density, and repeats until the result stabilizes.

The expensive part in conventional methods is a large solve performed at every iteration, and that step carries cubic complexity. Anandkumar’s team replaced that single step with AI rather than asking the model for the final answer.

Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU 4

They trained a Fourier Neural Operator, or FNO, to learn the forward mapping from input conditions to electron density. The paper says the complexity of that part drops from O(N^3) to O(N log N). The trained model is then placed back into the original iterative loop, while the rest of the procedure stays the same.

The article compares the setup to chain-of-thought style reasoning. A direct-prediction model resembles a one-shot answer to a hard problem and can fail once the task gets larger. Kohn-Sham FNO works in stages, with each step checked indirectly by later iterations, which can pull the process back if it starts to drift.

Divergence becomes a warning signal

The paper also points to a practical safeguard. If the model is pushed beyond the range it can handle, the iterations diverge, giving researchers a visible sign that the result should not be trusted.

In the first attempt at the magnesium dislocation experiment, the team used a general pre-trained model directly and every calculation diverged immediately. In the article’s telling, that divergence was itself the alarm. A direct-prediction system could still return a number, but the output alone would not reveal whether it was right or wrong.

Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU 5

8,504 structures used for training, one model for molecules and solids

The paper says the training set contained 8,504 structures. One model handled both molecules and solid materials and covered the first five rows of the periodic table.

That came from a domain-invariant FNO variant designed by the team. Under that setup, systems of different sizes share the same learned filters, allowing the same model to work across small molecules and large crystals.

When extrapolating to drug-scale large molecules not seen during training, Kohn-Sham FNO recorded a density error of 2.23%, compared with 9.97% for the direct-prediction model. As molecules grew larger, the gap widened. At 45 heavy atoms, the direct-prediction error climbed to 41%, while Kohn-Sham FNO was at 4%.

One B300 GPU versus a 2019 run that used about 7,800 V100s

The team used magnesium dislocation structures for large-scale validation.

Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU 6

The article points back to a 2019 study that received an ACM Gordon Bell Prize nomination. That earlier work performed a full DFT calculation on 6,164 magnesium atoms using about 7,800 NVIDIA V100 GPUs on the Summit supercomputer.

In the new paper, Kohn-Sham FNO completed a calculation on 8,250 atoms and 82,500 valence electrons using a single NVIDIA B300 GPU, and all runs converged.

The measured scaling exponent was 1.03, close to linear, versus 3.37 for the conventional method. The article also notes that one B300 and 7,800 V100s are not a hardware-equivalent comparison, but the scaling exponent itself does not depend on hardware. One method is near linear; the other is close to cubic. As system size increases, that gap becomes harder to overcome.

Company launch announced the same day as the paper

On the same day the paper was released, Anandkumar said she had founded an AI physics-simulation company called Accelerated Understanding. Reuters later reported on the launch.

Caltech team uses AI to cut DFT scaling near linear, runs 82,500-electron simulation on one GPU 7

In a post on X, she said the company’s AI models simulate physical phenomena in 4D space and support inference context lengths above 5 trillion.

The model still covers only one part of the pipeline

According to the article, Kohn-Sham FNO currently learns only one stage in the full workflow, and complete energy calculation still requires an additional post-processing step.

Even so, the direction is clearer now. The AI is not replacing physics itself; it is replacing the most expensive repeated solve inside a physics computation. The team’s next step is to cover the remaining stages so that the post-processing step could also be removed.

The reference page listed in the article is https://tensorlab.cms.caltech.edu/users/anima/ks_fno.html. The original source attribution says the piece came from the WeChat account Xinzhiyuan, written by ASI Qishilu and edited by Ma Ke.

This article was originally published by Bit.Fan. For more cryptocurrency news and market insights, visit www.bit.fan.
140

Disclaimer:

The market information, project data, and third-party content displayed on this platform are for industry information sharing only and do not constitute any form of investment advice or return commitment.

Cryptocurrency trading carries high risks. Users should fully assess their risk tolerance and make independent decisions. All profits, losses, and legal responsibilities are borne by the users themselves.