[论文解读] A Recipe for Charge Density Prediction

Xiang Fu, Andrew Rosen+6 more|arXiv (Cornell University)|May 29, 2024

Electron and X-Ray Spectroscopy TechniquesMaterials Science被引用 3

一句话总结

本文提出了一种用于密度泛函理论（DFT）中电荷密度预测的机器学习方法，其精度达到当前最先进水平，同时比现有方法快十倍以上。该方法结合了虚拟轨道、可学习基函数以及高容量等变神经网络，以高保真度和可扩展性建模电子密度。

ABSTRACT

In density functional theory, charge density is the core attribute of atomic systems from which all chemical properties can be derived. Machine learning methods are promising in significantly accelerating charge density prediction, yet existing approaches either lack accuracy or scalability. We propose a recipe that can achieve both. In particular, we identify three key ingredients: (1) representing the charge density with atomic and virtual orbitals (spherical fields centered at atom/virtual coordinates); (2) using expressive and learnable orbital basis sets (basis function for the spherical fields); and (3) using high-capacity equivariant neural network architecture. Our method achieves state-of-the-art accuracy while being more than an order of magnitude faster than existing methods. Furthermore, our method enables flexible efficiency-accuracy trade-offs by adjusting the model/basis sizes.

研究动机与目标

• 通过用机器学习替代迭代自洽场过程，加速密度泛函理论（DFT）中的电荷密度预测。
• 解决现有基于机器学习的电荷密度预测方法在精度与效率之间的权衡问题。
• 开发一种可扩展且精确的电荷密度表示方法，能够捕捉复杂电子结构，特别是化学键区域的特征。
• 通过调节模型和基集大小，实现灵活的效率-精度权衡。
• 在QM9数据集上验证该方法，并探索其向晶体材料的扩展应用。

提出的方法

• 使用原子和虚拟轨道坐标为中心的球形场表示电荷密度，其中虚拟节点被置于键中点，以改善对原子间电子密度的表征。
• 采用表达能力强、可学习的基函数（高斯型轨道）表示球形场，使模型能够自适应捕捉复杂的电子结构。
• 使用高容量、三维等变神经网络架构处理基于轨道的表示，并在三维网格上预测电荷密度。
• 为每个轨道引入可学习的缩放因子，以优化单个基函数的贡献，提升模型的表达能力。
• 通过可微分损失函数端到端训练模型，最小化预测电荷密度与DFT计算结果之间的L2差异。
• 通过利用轨道基函数的稀疏性和局域性，在推理过程中避免对整个体素进行处理，从而实现高效推理。

实验结果

研究问题

• RQ1 与标准原子轨道方法相比，引入虚拟节点的可学习轨道基表示是否能提升电荷密度预测的精度？
• RQ2 高容量等变神经网络在多大程度上增强了模型捕捉分子中复杂电子结构的能力？
• RQ3 虚拟轨道的引入在高电子密度梯度区域（如化学键区域）对预测精度有何影响？
• RQ4 该模型能否同时实现高精度与高推理速度，从而实现在大规模材料发现中的实际应用？
• RQ5 可学习缩放因子和基集大小对模型性能-效率权衡的影响如何？

主要发现

• 所提出的方法在QM9电荷密度基准测试中达到最先进精度，通过引入虚拟轨道，将平均相对误差从0.434%降低至0.237%（β=2.0）。
• 引入虚拟轨道使精度相对提升50%，尤其显著改善了共价键区域的预测表现。
• 该模型比现有方法快逾十倍，同时保持或超越其精度，可实现于从头算模拟中的实际部署。
• 更高容量的模型虽以效率为代价提升性能，但所有SCDP模型仍显著快于基线方法。
• 缩放因子在仅带来轻微效率损失的情况下提升了精度，位于Pareto前沿的模型展现出速度与精度之间的最优权衡。
• 该方法在原则上可推广至晶体材料，已提出基于CrystalNN或基于Voronoi的空隙填充策略，用于向周期性体系的扩展。