微软基于AI的新电子结构计算框架在Nature子刊发表

编辑 | ScienceAI

编者按：为了使电子结构方法突破当前广泛应用的密度泛函理论（KSDFT）所能求解的分子体系规模，微软研究院科学智能中心的研究员们基于人工智能技术和无轨道密度泛函理论（OFDFT）开发了一种新的电子结构计算框架M-OFDFT。这一框架不仅保持了与KSDFT相当的计算精度，而且在计算效率上实现了显著提升，并展现了优异的外推性能，为分子科学研究中诸多计算方法的基础——电子结构方法开辟了新的思路。相关研究成果已在国际知名学术期刊《自然-计算科学》（Nature Computational Science）上发表。

然而，各种电子结构方法的精度和计算效率通常很难达到平衡。目前，Kohn-Sham形式的密度泛函理论（Kohn-Sham density functional theory, KSDFT）被广泛应用，因为它在精度和效率之间取得了相对合理的平衡。不过，KSDFT仍然存在较高的计算复杂度，这使其无法很好地应对不断增长的求解大规模分子体系的需求。

微软研究院科学智能中心的研究员们提出了一种基于深度学习和无轨道密度泛函理论（OFDFT）的电子结构计算框架M-OFDFT。这一框架不仅在计算效率方面取得显著突破，同时也保持了高精度的计算解决能力。该成果突显了人工智能在提升电子结构计算中精度和效率的平衡方面的独特潜力。这一新框架有望加速相关行业问题的研究与开发进程。

M-OFDFT的相关研究成果以「Ovecoming the Barrier of Orbital-Free Density Functional Theory for Molecular Systems Using Deep Learning」为题，于 2024 年 3 月 11 日发表在国际知名学术期刊《自然-计算科学》（Nature Computational Science）上。

人工智能给电子结构方法带来新机会

电子结构方法是研究分子体系物理化学性质的基本工具。由于多电子体系的复杂性，高精度的电子结构方法往往计算代价较高，难以在工业界应用。为了处理较大分子系统，通常引入了一些近似，从而牺牲了一定的精度。目前，KSDFT方法在精度和效率方面取得了平衡，因此被广泛应用。

不过，近期人工智能技术的喜人进展也为其他电子结构计算框架带来了新的机会。为了使电子结构方法突破KSDFT所能求解的分子体系规模，微软研究院的研究员们利用人工智能技术，开发了M-OFDFT，该方法比KSDFT效率更高，同时又能保有其精度。基于OFDFT的开发，让M-OFDFT成为了一种比KSDFT理论复杂度更低的电子结构计算框架，因为它只需优化电子密度函数这一个函数来求解电子状态即可，KSDFT则需要优化与电子数相同的多个函数。

然而，OFDFT仍然面临着一个巨大的挑战——需要关于密度函数的电子动能泛函，但其确切形式尚不明确，并且构建适用于分子体系的高精度近似方法也相当困难。

针对这一难题，M-OFDFT使用一个深度学习模型来近似动能泛函。借助深度学习模型的强大拟合能力，M-OFDFT可实现比基于近似物理模型设计的经典动能泛函更高的准确度。对于一个待求解的分子体系结构，M-OFDFT会使用动能泛函模型以及其他可直接计算的能量项构造出一个电子密度的优化目标，然后通过优化过程求解最优（基态）电子密度（图1），进而可计算能量、力、电荷分布等分子属性。

图1：对于待求解的分子体系结构，M-OFDFT通过最小化电子能量来求解电子密度（以其向量化系数表示），其中难以近似的动能部分由深度学习模型来近似。

M-OFDFT实现兼具精度与效率的电子结构方法

研究员们对M-OFDFT进行了一系列的实验验证。首先考察的是M-OFDFT在常见小分子体系上的求解精度。

结果显示，M-OFDFT在乙醇分子构象以及QM9数据集的分子上可以达到与KSDFT相当的精度（能量达到化学精度1 kcal/mol）。相较于经典OFDFT方法，精度提高了两个数量级（图2-a）。M-OFDFT解得的电子密度也与KSDFT的结果重合（图2-b），特别是得到了电子壳层结构，而经典OFDFT的结果则有明显偏差。由M-OFDFT解得的乙醇构象空间上的势能面（每个点都是通过密度优化得到的，并不是直接预测）也与KSDFT的结果一致（图2-c）。

之后，研究员们又验证了M-OFDFT不仅保有KSDFT级别的精度，其更低的理论计算复杂度还使其在效率上也超越了KSDFT。在实际计算中M-OFDFT取得了的复杂度（图3），比KSDFT 的实际复杂度低了一阶，且其所需绝对时间也明显少于KSDFT。在两个更大的蛋白质体系上（包含2676和2750个电子），M-OFDFT实现了25.6倍和27.4倍的加速。

M-OFDFT具有更强的泛化能力

深度学习模型在科学任务中的应用面临一大挑战是，在具有与训练数据不同特点的数据上的泛化问题。但采用了OFDFT框架后，动能泛函模型遇到的泛化问题就会减轻，从而使 M-OFDFT可以在比训练集分子规模更大的体系上展现出良好的外推能力。

实验结果表明，M-OFDFT的能量预测误差显著低于基于深度学习的端到端能量预测模型（图4-a）。此外，研究员们还利用在多肽片段上训练的M-OFDFT模型求解完整蛋白结构，并取得了超越端到端模型和经典OFDFT的泛化性能（图4-c）。不仅如此，相较端到端模型，M-OFDFT还可以用更少的大分子体系训练数据取得更好的泛化表现（图4-b与图4-d）。

M-OFDFT的工作原理

「神龙见首又见尾」：高效捕获非局域效应的动能泛函模型

动能密度泛函具有明显的非局域效应，而用经典的基于格点（grid）的方式表征电子密度则会带来高昂的非局域计算代价。为此，M-OFDFT将电子密度在一组原子基组函数上展开，并使用展开系数作为电子密度表征。由于基函数叠加的形状与电子分布接近，所以其数量可远小于格点数，使得非局域计算代价大大降低，并有助于刻画电子密度中的壳层结构。

M-OFDFT将每个原子上的电子密度系数和类型与坐标作为节点特征,并基于Graphormer模型 [1]预测电子动能（图5），其自注意力机制显式刻画了荷载在每两个原子上的电子密度特征之间的相互作用，从而可捕捉非局域性质。此外，为了保证动能的旋转不变性，M-OFDFT使用了以各个原子为中心、基于其相邻原子的局部坐标系，将电子密度系数转换为旋转不变的特征。

「横看成岭侧成峰，远近高低各不同」：高效学习电子能量曲面的训练策略

与传统机器学习任务不同，动能泛函模型是被当作其输入变量的优化目标使用的，而非用于在一些单点上做预测，这对模型的学习提出了更高的要求：模型必须捕捉到每个分子结构上电子能量曲面的轮廓。

为此，研究员们深入分析了用来生成数据的电子结构方法，发现它其实可以为每个分子结构生成多个数据点，而且还能提供梯度标注，从而让模型可以拥有更丰富的曲面轮廓特征。然而梯度的巨大范围也使神经网络难以优化。对此，研究员们还提出了一系列增强模块，让模型能够更容易地表达巨大的梯度。

开启未来电子结构方法的新篇章

M-OFDFT成功突破了无轨道密度泛函框架在分子体系中的瓶颈，将其求解精度提升到了常用的KSDFT的水平，同时保有了其更低的计算代价，推进了电子结构方法在「精度-效率」方面的权衡，为分子科学研究提供了一种更有潜力的研究工具。

尽管M-OFDFT已经在某些分子体系上展现了出色的泛化性能，但在更大的分子体系上实现长时间且稳定的高精度模拟仍是一个巨大的挑战。微软研究院期待M-OFDFT可以沿着这一方向激发更多研究与创新，并在未来和其他方法一起为电子结构计算带来更多突破性的成果和影响。

相关文章：

[1] Do Transformers really perform badly for graph representation? Advances in Neural Information Processing Systems 34 (NeurIPS 2021).（https://proceedings.neurips.cc/paper/2021/hash/f1c1592588411002af340cbaedd6fc33-Abstract.html） 

注：封面来自网络。