Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory

六十多年来，科学家们一直依赖一套强大的规则来理解原子如何粘合在一起以形成我们周围的材料。这一框架被称为密度泛函理论，它并不将原子周围的电子云视为单个粒子的集合，而是将其视为一种平滑、连续的密度。这种方法使研究人员能够预测从新药到先进电池等各种事物的性质。然而，这其中有一个难点。为了获得准确的结果，该理论需要进行复杂的、重复的计算，涉及求解系统中每一个电子的行为。随着系统规模的增大，这些计算所需的时间会呈爆炸式增长，其增长速度之快，使得在即使是最强大的超级计算机上模拟中等规模的金属块或大型生物分子也变得不可能。这一瓶颈阻碍了许多重要的科学问题，迫使研究人员在准确性和规模之间做出选择。

加州理工学院的一个研究小组发现了一种打破这一障碍的方法，即通过教计算机跳过计算中最昂贵的部分而不损失精度。他们并没有尝试直接预测最终答案，也没有学习一种经常失败的困难数学捷径，而是教会了人工智能执行每一次计算循环中发生的特定重复步骤。通过专注于这个中间步骤，他们创造出一种可以在单个图形处理器（显卡）上模拟数万个原子的系统的方法，而此前这项任务需要数千个处理器。其结果是一种新的物质建模方式，它几乎像最简单的近似法一样快速，同时保留了最严谨方法的极高精度。

问题的核心在于该理论如何处理电子。为了找到材料的稳定状态，计算机必须反复调整电子密度图，直到其不再发生变化。在传统方法中，这种调整需要求解一个巨大的特征值问题，这本质上是在寻找系统中每个原子特定的电子运动模式。这种操作的计算量极大，其成本随电子数量呈立方级增长；这意味着系统规模翻倍，计算难度会增加到八倍。几十年来，科学家们一直试图通过不再追踪单个电子模式来绕过这一点，这种策略被称为无轨道密度泛函理论。然而，以往利用机器学习学习这种捷径的尝试都遇到了困难。一些人试图学习电子密度与能量之间的关系，但这种关系在数学上是不稳定的且容易出错。另一些人则试图直接从原子的排列预测最终的电子密度，但这些模型在面对比训练集更大或化学性质不同的系统时就会失效。

研究人员意识到，关键在于不要试图预测最终的目的地，而要学习旅程本身。在传统计算的每一步中，计算机都会获取一份作用于电子的力图，并生成一张新的电子可能出现的分布图。这个前向步骤在数学上是稳定且定义明确的。该团队训练了一个专门的神经网络——一种旨在理解空间模式的人工智能类型——来学习这种精确的转换。他们向网络输入了数千个这些力图及其对应的电子密度图的示例，这些示例取自包括有机化合物、绝缘体和金属在内的8,504种分子和固体材料。该网络学会了直接从力来预测新的电子密度，从而有效地用快速的学习预测取代了缓慢且沉重的计算。

这种方法的独特之处在于它如何处理不同系统的复杂性。研究人员设计了该网络使其具有领域不变性（domain-invariant），这意味着它可以将学到的规则应用于微小的分子和巨大的晶体，而无需为每个新尺寸重新训练。他们还确保了网络遵循空间的根本对称性，因此它理解旋转或移动分子不会改变其物理性质。在针对从未见过的系统进行测试时，该模型表现出了极强的鲁棒性。在一组包含训练数据中不存在的元素和尺寸的类药物分子上，新方法保持了高度的准确性，而之前的直接预测模型在分子变大时误差会剧增。这种自洽性的方法允许它在每一步中纠正自己的错误，就像一位导航员频繁检查位置，而不是试图一次性猜测整个路线。

当研究小组将其应用于晶体等周期性系统以及金属中的扩展缺陷时，该方法的威力显现了出来。在针对晶体的标准计算中，计算机必须在描述晶体重复模式的数学空间中的许多不同点上重复沉重的计算。新方法完全跳过了这种重复，直接在原胞上演化电子密度，而无需单独对这些点进行采样。这使得研究人员能够使用单一模型完成金属系统和半导体系统的计算收敛，并能重现与传统方法相同的电子光谱和结构性质。此外，该模型在不同理论近似之间表现出了惊人的迁移能力：一个基于某种类型计算训练的模型，可以通过简单地更换输入力，直接用于另一种更精确的力近似方法，而无需重新训练。

最终的测试来自于研究人员将该方法应用于镁位错——这是金属晶体中一种关键的缺陷，对于理解材料如何变形和断裂至关重要。这些缺陷会产生长程弹性场，需要通过模拟数千个原子才能准确建模。以往尝试以高精度模拟此类系统的尝试，需要拥有数千个图形处理器的超级计算机，并且需要运行数天。而该新方法在单个现代图形处理器上，成功完成了包含8,250个镁原子和82,500个价电子的系统计算。完成计算所需的时间随系统规模几乎呈线性增长，证实了计算瓶颈已被消除。所得出的电子密度与高精度参考值几乎无法区分，这使得研究人员能够解析决定金属强度的不同类型缺陷之间的细微能量差异。

这项工作证明了机器学习可以扩展电子结构计算的范畴，其方式不是取代整个物理理论，而是加速其最昂贵、最重复的操作。通过学习从力到密度的前向映射，研究人员创建了一个稳定、可跨不同物质类型迁移、并能处理此前无法触及的规模系统的工具。该方法为其自身的可靠性提供了直接诊断：如果计算无法收敛，它会发出信号，表明该系统超出了模型的信任范围，这是直接预测模型所缺乏的功能。虽然目前的模型预测的是电子密度，但总能量和其他性质仍可以通过最后一步单次计算获得，从而保留了完整理论的精度。这种方法为模拟复杂材料、生物系统和化学反应开辟了大门，使这些过程能够在曾经被视为理论上的不可能实现的细节水平和规模下进行。