[论文解读] Generalizing Neural Wave Functions

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|Feb 8, 2023

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

本文提出Globe，一种基于图神经网络的轨道嵌入方法，通过GNN中的空间消息传递学习局域电子结构，使神经波函数能够跨不同分子实现泛化。此外，本文提出Moon，一种尺寸一致的神经波函数，其收敛速度比以往方法快达4.5倍，且在不同系统尺寸下保持准确的可加性能量缩放，从而实现仅用一个模型联合求解多个不同分子的薛定谔方程。

ABSTRACT

Recent neural network-based wave functions have achieved state-of-the-art accuracies in modeling ab-initio ground-state potential energy surface. However, these networks can only solve different spatial arrangements of the same set of atoms. To overcome this limitation, we present Graph-learned orbital embeddings (Globe), a neural network-based reparametrization method that can adapt neural wave functions to different molecules. Globe learns representations of local electronic structures that generalize across molecules via spatial message passing by connecting molecular orbitals to covalent bonds. Further, we propose a size-consistent wave function Ansatz, the Molecular orbital network (Moon), tailored to jointly solve Schrödinger equations of different molecules. In our experiments, we find Moon converging in 4.5 times fewer steps to similar accuracy as previous methods or to lower energies given the same time. Further, our analysis shows that Moon's energy estimate scales additively with increased system sizes, unlike previous work where we observe divergence. In both computational chemistry and machine learning, we are the first to demonstrate that a single wave function can solve the Schrödinger equation of molecules with different atoms jointly.

研究动机与目标

• 克服现有神经波函数的局限性，即仅能跨同一分子的不同几何构型泛化，而无法跨不同分子组成泛化。
• 实现单一神经波函数同时求解多个不同分子的薛定谔方程，包括原子数和轨道数不同的分子。
• 在神经波函数中实现尺寸一致性，确保能量随系统尺寸可加性缩放，这是量子化学物理准确性的一项关键要求。
• 通过设计利用空间消息传递和原子核锚点的波函数架构，加速变分量子蒙特卡洛训练的收敛速度。

提出的方法

• Globe通过图神经网络中的空间消息传递，将分子轨道嵌入三维空间，实现对不同分子中局域电子结构的学习。
• 轨道嵌入通过将分子轨道与共价键相连来学习，使模型能够泛化到轨道数不同的分子。
• Moon作为波函数架构，通过聚焦于局域电子-原子核相互作用并以原子核作为消息传递的锚点，实现尺寸一致性。
• 波函数被构建为多个行列式之和，其中轨道权重由Globe重新参数化，从而实现对多样化分子的联合训练。
• 该方法采用两级神经网络：一级用于重新参数化轨道权重（Globe），另一级用于计算最终波函数（Moon）。
• 训练通过变分蒙特卡洛进行，以能量最小化为目标，其中Globe的重新参数化在每一步中仅计算一次，与采样和能量评估解耦。

Figure 1: Globe steps. 1) Atom embeddings are obtained by message passing between atoms. 2) Orbital locations and embeddings are determined by core orbitals and molecular bonds. 3) Orbital embeddings are updated via a unidirectional message passing. For clarity, we omitted messages from each atom to

实验结果

研究问题

• RQ1 能否训练单一神经波函数，使其同时求解多个具有不同原子组成的分子的薛定谔方程？
• RQ2 如何在可微、可扩展的神经波函数架构中实现尺寸一致性？
• RQ3 GNN中的空间消息传递能否实现对轨道嵌入在轨道数不同的分子间泛化？
• RQ4 与单独训练相比，联合训练多个分子是否能加速收敛？
• RQ5 当神经波函数在不同分子间泛化时，其能量估计是否随系统尺寸可加性缩放？

主要发现

• Moon在训练步数上比以往方法减少4.5倍，同时达到相似或更低的能量估计，证明了其优化收敛速度更快。
• Moon的能量估计随系统尺寸可加性缩放，而以往方法存在发散缩放现象，证实其尺寸一致性。
• Globe使模型能够直接应用于更大分子体系，而无需额外的自洽场（SCF）计算，实现对更大链状体系的直接评估。
• 在氢链体系中，Globe与Moon保持了随系统尺寸增加而恒定的每原子能量，表明其具备稳健的广延性与物理一致性。
• Globe与Moon在无需预训练的情况下，成功实现了对乙烷和环丁二烯等较大分子的收敛，而其他方法在无预训练时无法收敛。
• Moon在仅使用60万参数的情况下，性能优于FermiNet并达到PsiFormer的精度，展现出极高的参数效率。

Figure 2: Illustration of Moon. 1) We initialize electron embeddings by aggregating their local neighborhood of electrons. 2) Nuclei aggregate electron embeddings via message passing. 3) Nuclei embeddings are iteratively updated. 4) Nuclei embeddings are structurally diffused towards the electrons v