加州理工用AI模型将量子化学计算复杂度压至近线性

加州理工学院 Anima Anandkumar 团队于 8 月 24 日发表论文，提出一种基于傅里叶神经算子（FNO）的 AI 模型 Kohn-Sham FNO，将量子化学中最常用的密度泛函理论（DFT）计算复杂度从立方级 O(N³) 压缩至近线性级 O(N log N)。在单块 NVIDIA B300 GPU 上，团队完成了一个含 8250 个原子、82500 个价电子的镁位错模拟，而 2019 年一项获 ACM 戈登·贝尔奖提名的同类计算动用了 Summit 超算上约 7800 块 V100 GPU。

困扰量子化学 60 年的瓶颈

DFT 是药物筛选、催化剂设计、电池与芯片材料研发的底层工具，相关研究曾获诺贝尔化学奖。但其核心计算量随电子数呈立方增长：电子数翻 10 倍，计算量增加约 1000 倍。60 年来，物理学家一直在寻找加速办法，未取得实质突破。

此前 AI 介入 DFT 主要有两条路径，均存在明显短板：

• 直接预测路线：让模型一步输出最终结果。训练分布内精度尚可，但外推能力极差。论文实测显示，在从小分子外推到药物级大分子时，误差从不到 1% 飙升至约 10%，最高达 41%。
• 逆向映射路线：让模型学习反问题求解。数学上不稳定，论文对照实验中所有计算在迭代几步后均发散。

用「思维链式」迭代替代一步预测

新论文的核心思路是只替换 DFT 迭代流程中最昂贵的那一步，而非让模型端到端预测最终结果。具体做法是：

• 训练一个傅里叶神经算子（FNO），学会「给定输入条件，输出对应电子密度」的正向映射，将单步复杂度从 O(N³) 降至 O(N log N)。
• 将模型嵌入原有迭代循环，其他步骤保持不变，每轮迭代完成后自我校验，错误会在后续轮次中被修正。

这一设计类似大语言模型的「思维链」推理：分步推演而非一步猜答案。其额外优势是自带安全阀——若模型被推至能力范围之外，迭代会发散，研究者可立即判定结果不可信。论文中镁位错实验首次尝试即出现发散，发散本身成为警报。

训练数据仅用 8504 个结构，一个模型同时覆盖分子与固体材料，涵盖元素周期表前五行。这得益于团队设计的域不变 FNO 变体：不同尺寸的体系共享同一套学到的滤波器，小分子与大块晶体共用同一模型。

外推到训练集未见的药物级大分子时，Kohn-Sham FNO 的密度误差为 2.23%，直接预测模型为 9.97%；到 45 个重原子时，两者分别为 4% 与 41%，差距随体系增大而扩大。

单卡 GPU 完成 8 万电子模拟

大规模验证选用了镁的位错结构。该体系 2019 年在 Summit 超算上需约 7800 块 V100 完成全量 DFT 计算，Kohn-Sham FNO 在单块 B300 上即完成 8250 个原子、82500 个价电子的模拟并全部收敛。

• 实测缩放指数为 1.03（接近完美线性），传统方法为 3.37（立方级）。
• 两块硬件并非对等比较，但缩放指数与硬件无关：近线性与立方的差距随体系增大而不可逆。

研究者的下一步

论文发布同日，Anandkumar 宣布创办 AI 物理模拟公司 Accelerated Understanding，路透社随即做了报道。据其介绍，公司训练的模型在 4D 空间中模拟物理现象，推理上下文长度超过 5 万亿。

目前 Kohn-Sham FNO 仅替代了 DFT 流程中的一个环节，完整能量计算仍需一步后处理。团队计划补齐剩余环节，届时后处理也有望省去。AI 替代的不是物理本身，而是物理计算中最烧钱的那一步重复运算。

参考资料：https://tensorlab.cms.caltech.edu/users/anima/ks_fno.html