Equivariant Flow Matching for Electron Density Prediction

📄 arXiv: 2610.02651v1 📥 PDF

作者: Chenxing Liang, Chengdong Wang, Yuchao Lin, Xiaofeng Qian, Shuiwang Ji

分类: physics.chem-ph, cs.AI

发布日期: 2026-10-02

💡 一句话要点

提出OrbFlow以解决电子密度预测中的计算成本问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 电子密度预测 密度泛函理论 机器学习 生成模型 流匹配 量子化学 高斯型轨道 自洽场

📋 核心要点

  1. 现有的电子密度预测方法面临高计算成本和结构相关性捕捉不足的挑战。
  2. 本文提出OrbFlow，通过流匹配预测高斯型轨道系数，采用$ ext{SE}(3)$-等变生成模型。
  3. OrbFlow在QM9数据集上减少了13.6%的密度误差，并在MD基准上减少了51%至63%的误差，显著提高了预测效率。

📝 摘要（中文）

机器学习替代方法在密度泛函理论（DFT）中的应用日益增多，以降低第一性原理计算的成本。在这一领域，预测真实空间电子密度为自洽场（SCF）程序提供了可扩展和可转移的初始化。然而，现有方法面临明显的困境：基于网格的架构计算成本高，而基于基函数的方法无法捕捉系数空间中固有的结构相关性。为此，本文提出了OrbFlow，一种$ ext{SE}(3)$-等变生成模型，通过流匹配预测高斯型轨道（GTO）系数。OrbFlow保持了紧凑的原子中心基的高效性，同时用学习的概率路径替代了逐点回归。通过两阶段的轨迹课程进行训练，减轻了数值积分过程中的离散化漂移。OrbFlow在QM9数据集上实现了最先进的准确性，相较于之前最佳模型减少了13.6%的密度误差，并在MD基准的每个分子上相较于最强的先前方法减少了51%至63%的误差。预测的密度还将SCF迭代次数减少了多达68%，并在未见的交换-关联泛函上实现了零-shot迁移，恢复了偶极和四极矩，误差仅在DFT参考值的几个百分点内，无需任何SCF计算。

🔬 方法详解

问题定义：本文旨在解决电子密度预测中的高计算成本和现有方法无法有效捕捉结构相关性的问题。现有的网格基础方法计算开销大，而基函数方法则无法充分利用系数空间的结构信息。

核心思路：OrbFlow通过流匹配技术，利用$ ext{SE}(3)$-等变生成模型来预测高斯型轨道系数，旨在提高预测的准确性和效率。该模型通过学习概率路径而非逐点回归，能够更好地捕捉系数空间的结构特征。

技术框架：OrbFlow的整体架构包括两个主要阶段：首先是通过轨迹课程进行模型训练，以减轻数值积分中的离散化漂移；其次是利用训练好的模型进行电子密度的预测。

关键创新：OrbFlow的主要创新在于其$ ext{SE}(3)$-等变性和流匹配方法的结合，使得模型在保持高效性的同时，能够有效捕捉系数空间中的结构相关性。这一设计与传统方法有本质区别。

关键设计：在模型设计中，OrbFlow采用了紧凑的原子中心基，并通过特定的损失函数来优化流匹配过程。此外，模型的训练过程分为两个阶段，以确保在数值积分时的稳定性和准确性。通过这些设计，OrbFlow在准确性和效率上均有显著提升。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

OrbFlow在QM9数据集上实现了最先进的准确性，相较于之前最佳模型减少了13.6%的密度误差，并在MD基准的每个分子上减少了51%至63%的误差。此外，预测的密度将SCF迭代次数减少了多达68%，显示出其在实际应用中的巨大潜力。

🎯 应用场景

该研究的潜在应用领域包括量子化学计算、材料科学和药物设计等。通过提高电子密度预测的效率，OrbFlow能够显著降低第一性原理计算的成本，为相关领域的研究提供更为高效的工具，推动科学发现的进程。

📄 摘要（原文）

Machine learning surrogates for density functional theory (DFT) have been increasingly used to reduce the cost of first-principles calculations. In this arena, predicting real-space electron densities offers a scalable and transferable initialization for self-consistent field (SCF) procedures. However, current methods face a clear dilemma. That is, grid-based architectures incur a high computational cost, while basis-set methods fail to capture the structural correlations inherent in the coefficient space. Here, we develop OrbFlow, an $\mathrm{SE}(3)$-equivariant generative model that predicts Gaussian-type orbital (GTO) coefficients via flow matching. OrbFlow retains the efficiency of a compact atom-centered basis while replacing pointwise regression with a learned probability path over the full coefficient space. It is trained through a two-phase trajectory curriculum that mitigates discretization drift during numerical integration. OrbFlow achieves state-of-the-art accuracy on QM9, reducing density error by 13.6% relative to the previous best model, and reduces error by 51% to 63% on every molecule of the MD benchmark relative to the strongest prior method sharing its basis. The predicted density also cuts SCF iterations by up to 68% with zero-shot transfer to unseen exchange-correlation functionals and recovers dipole and quadrupole moments to within a few percent of DFT references without any SCF calculation.