[论文解读] DeepErwin

Michael Scherbela, Leon Gerard+1 more|arXiv (Cornell University)|Mar 17, 2023

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

该论文提出DeepErwin，一种基础波函数模型，通过图神经网络将Hartree-Fock轨道映射为高精度神经网络轨道，实现在分子间的迁移学习。在多样化化合物上进行预训练后，仅需极少微调即可实现比标准训练快20倍的CCSD(T)级精度。

ABSTRACT

Deep neural networks have become a highly accurate and powerful wavefunction ansatz in combination with variational Monte Carlo methods for solving the electronic Schrödinger equation. However, despite their success and favorable scaling, these methods are still computationally too costly for wide adoption. A significant obstacle is the requirement to optimize the wavefunction from scratch for each new system, thus requiring long optimization. In this work, we propose a novel neural network ansatz, which effectively maps uncorrelated, computationally cheap Hartree-Fock orbitals, to correlated, high-accuracy neural network orbitals. This ansatz is inherently capable of learning a single wavefunction across multiple compounds and geometries, as we demonstrate by successfully transferring a wavefunction model pre-trained on smaller fragments to larger compounds. Furthermore, we provide ample experimental evidence to support the idea that extensive pre-training of a such a generalized wavefunction model across different compounds and geometries could lead to a foundation wavefunction model. Such a model could yield high-accuracy ab-initio energies using only minimal computational effort for fine-tuning and evaluation of observables.

研究动机与目标

• 为克服为每个新分子体系从头训练深度神经网络波函数所带来的高计算成本。
• 通过参数高效、尺寸不变的架构，实现波函数模型在不同分子和几何构型间的泛化能力。
• 建立一个在多样化化合物上预训练的基础波函数模型，减少对完整重新优化的需求。
• 证明广泛预训练可系统性提升精度与迁移能力。
• 通过微调实现显著降低计算成本的前提下，超越传统高精度方法（如CCSD(T)-ccpVTZ）的性能。

提出的方法

• 图神经网络（GNN）处理核坐标和初始轨道描述符，生成与体系尺寸无关的、学习得到的关联轨道。
• 模型采用消息传递架构，包含对称与反对称组件，以确保满足费米统计要求。
• 通过可学习函数结合电子嵌入与轨道描述符，构建电子轨道，确保自旋与空间对称性。
• 波函数通过这些学习到的轨道构建为斯莱特行列式，通过反对称激活函数强制实现对称性。
• 在小分子片段上的预训练阶段，使模型可通过微调迁移到更大、化学相似的化合物。
• 该方法将每批次的昂贵GNN和轨道网络计算解耦，降低计算开销。

实验结果

研究问题

• RQ1 是否可在一个神经网络波函数模型上对多个分子和几何构型进行预训练，以实现在新体系上的快速、高精度预测？
• RQ2 在小分子片段上进行预训练是否能提升在更大、化学相似化合物上的精度与泛化能力？
• RQ3 该模型是否仅通过极少微调即可达到CCSD(T)级精度，从而将计算成本降低数个数量级？
• RQ4 模型规模、数据规模和预训练步数如何影响波函数模型的最终精度与迁移能力？
• RQ5 与同期方法GLOBE相比，该方法在多样化化合物上的能量精度与泛化能力如何？

主要发现

• 在小分子片段上进行预训练，可成功将知识迁移到更大、化学相似的化合物上，在相对能量预测中保持高精度。
• 对预训练的DeepErwin模型进行微调，可实现约20倍于从头训练新模型速度的CCSD(T)-ccpVTZ级精度。
• 系统性地扩展模型规模、数据规模和预训练步数，可一致且显著提升波函数精度。
• 与同期方法GLOBE相比，该模型在绝对能量上更低，且在相对能量计算中精度更高，尤其在分布外分子上表现更优。
• 该架构的尺寸不变性使不同大小的分子间可共享权重，克服了先前波函数ansatz的主要局限。
• 该方法在分布内与分布外分子上均保持高精度，展现出强大的泛化能力。