用于量子化学的可迁移神经网络量子态方法

《Computer Physics Communications》：Transferable neural network quantum state method for quantum chemistry

　　

# 摘要

我们提出了QiankunNet1.2，一种可迁移的神经网络量子态（NNQS）方法，用于求解量子化学中的多电子薛定谔方程。在QiankunNet1.0的基础上，QiankunNet1.2解决了其前身的关键局限性：在不同分子系统之间迁移学习效果差的问题。我们引入了两项主要创新，以提高准确性和可迁移性。

（1）**增强的波函数Ansatze与相位可迁移性**：QiankunNet1.2采用了一种改进的神经网络波函数Ansatze，具有统一相位表示。通过引入一个额外的固定输入维度多层感知机（MLP）来表示波函数相位，模型可以保留并复用先前学习到的系统知识。这种设计使QiankunNet1.2能够以比QiankunNet1.0显著更少的参数达到更低的基态能量。

（2）**基于继承的迁移学习**：我们实现了一种迁移学习策略，即通过多个参考系统对单个NNQS模型进行预训练，通过在一组相关分子上依次训练来获得丰富的"知识检查点"。使用该丰富检查点初始化新的计算，相比从头训练每个系统，在目标系统上实现了更快的收敛速度和更低的最终能量。

我们在多种分子基准上验证了QiankunNet1.2，实现了更快的收敛和更高的精度。我们的结果表明，QiankunNet1.2在可迁移性和效率上显著优于先前的NNQS方法，代表了神经网络量子化学计算的一项重要进步。

# 引言

预测分子系统的基态电子能量是量子化学和多体物理学的核心挑战。形式上，我们寻求求解含n个电子和m个原子核的系统的时间无关多体薛定谔方程 $H\psi = E\psi$，其中H是哈密顿量算符。在Born-Oppenheimer近似下，分子电子哈密顿量可以写为：

$$H = -\sum_{i=1}^{n} \frac{\nabla_i^2}{2} - \sum_{i=1}^{n} \sum_{I=1}^{m} \frac{Z_I}{|\mathbf{r}_i - \mathbf{R}_I|} + \sum_{i=1}^{n} \sum_{j>i} \frac{1}{|\mathbf{r}_i - \mathbf{r}_j|}$$

其中 $\mathbf{r}_i$ 和 $\mathbf{R}_I$ 分别是第i个电子和第I个原子核（具有核电荷 $Z_I$）的位置。原则上，求解此哈密顿量本征值问题将给出精确的基态波函数 $\psi$ 和能量 $E = \langle \psi | H | \psi \rangle$。然而，精确解的计算复杂度随系统规模（电子数和基函数数）呈指数增长。这激发了各种近似方法的发展，以在电子结构计算中平衡精度和可行性。

由于薛定谔方程的高维复杂性，各种方法被提出以平衡计算精度和效率。Hartree-Fock（HF）方法通过平均场近似简化电子-电子相互作用，将电子波函数处理为单电子波函数的乘积[1]，[2]。它构成了电子结构计算的基础。全组态相互作用（FCI）方法可以提供精确解[3]，但其计算成本随电子数呈指数增长，使其难以应用于大系统。在超级计算机上，FCI方法可以处理不超过24个轨道的系统[4]。组态相互作用（CI）[5]方法仅考虑与参考态不同且满足特定约束的组态，从而降低了计算复杂度，使该方法可以应用于更大的系统。耦合簇（CC）[6]，[7]方法对原始Slater行列式执行多重激发，为某些分子波函数产生更精确的结果，但对具有强静态关联的系统不太适用。密度矩阵重正化群（DMRG）[8]，[9]基于张量网络方法[10]，[11]，但当应用于高维系统时计算量较大且需要大量资源[12]。量子蒙特卡罗（QMC）[13]，[14]，[15]是一种使用随机采样技术求解量子多体问题的方法，包括变分蒙特卡罗（VMC）和扩散蒙特卡罗（DMC）[16]。VMC方法使用Ansatze来近似基态波函数，通过最小化能量的期望值来实现。

近年来，神经网络波函数Ansatze作为一种有前景的范式，被应用于VMC和多体模拟。Carleo和Troyer[17]率先将其用于自旋系统，NNQS利用深度神经网络来表示波函数，旨在结合高表达力和可扩展的优化。后续工作表明，神经网络Ansatze（例如基于受限玻尔兹曼机、卷积网络或自回归模型）在某些量子多体问题上可以达到甚至超过传统方法的精度[18]，[19]，[20]，[21]，[22]，[23]，[24]，[25]，[26]。在量子化学领域，神经网络量子态已被应用于晶格（二次量子化）表示和连续实空间电子坐标。值得注意的是，面向连续空间中分子的高精度费米子神经网络已被开发（例如Pfau等人[27]，[28]；Hermann等人[29]），通过在电子坐标上直接优化，为小分子展示了接近化学精度的结果。另一种方法，也是我们在本文中采用的方法，是在离散占据数基（自旋轨道基）中表示电子波函数[23]，[24]，[30]，[31]，[32]，[33]，[34]，[35]。在此表示中，波函数定义在 $N_{\text{orb}}$ 个自旋轨道被电子占据或未占据的组态上（类似于FCI基）。这种离散方法允许利用组合稀疏性和对称性，并且更有利于利用高性能计算（HPC）来处理更大的基组（尽管代价是要求使用固定的单粒子基）。

此前，我们提出了QiankunNet1.0[31]，[32]，使用深度神经网络进行二次量子化分子波函数。QiankunNet1.0采用基于Transformer的架构，使用多个解码器层来构建波函数振幅，并结合多层感知机（MLP）来输出复波函数的相位。它在多种分子上实现了高精度，在许多情况下接近精确的FCI能量。尽管取得了这些成就，QiankunNet1.0面临两个挑战。首先，对于每个新分子，需要大量的网络参数（尤其是在相位网络中）才能达到低能量误差。QiankunNet1.0中的相位MLP针对每个系统规模进行了定制：其输入维度等于自旋轨道数（$N_{\text{orb}}$ 或nqubits），因此网络架构（以及参数数量）因不同分子而异。这导致了效率低下，因为较小的系统仍然需要大量的参数，而较大的系统则需要更大的网络来保持精度。其次，也是更根本的是，变化的输入维度意味着在一个分子上学习到的知识无法直接转移到另一个分子上。在QiankunNet1.0中，如果两个分子具有不同的 $N_{\text{orb}}$，则相位MLP的结构不同，阻止了训练权重的直接重用。因此，在QiankunNet1.0中使用迁移学习——将预训练模型作为新系统的起点——是不切实际的。

为了克服这些挑战，我们认为一种可迁移的、低参数的神经网络波函数Ansatze是必不可少的。这样的模型可以保留先前系统的知识并快速适应新系统，从而通过迁移学习加速收敛到更低的能量。这与经典机器学习中的迁移学习类似，后者允许将大数据集上学习到的特征复用于相关任务，从而更快地训练并通常在更小的数据集上获得更好的性能。在量子化学中，初步研究表明，将优化后的较小系统的波函数迁移到较大或相关系统的优化初始化中，可以提高效率。我们使用QiankunNet1.2的目标就是将这种可迁移性原生地融入到面向量子化学的NNQS方法中。

我们引入QiankunNet1.2作为一个改进的NNQS框架，在保持高精度的同时增强了可迁移性和效率。本工作的核心贡献如下：

- **用于可迁移性的统一相位网络**：我们通过将波函数相位计算与系统规模解耦来重新设计了神经网络架构。QiankunNet1.2引入了一个额外的相位网络MLP，称为MLP2，具有固定输入大小，与QiankunNet1.0的依赖系统的相位设计相比，增强了可迁移性。在实践中，这是通过将每个Transformer解码器层的输出聚合为恒定大小的特征向量，然后将其送入相位MLP2来实现的。因此，相同的网络架构（权重和维度）可以跨不同分子使用，允许在系统之间直接转移学习到的参数。通过使用这种一致的相位表示，QiankunNet1.2可以从一个系统到另一个系统"携带"电子关联和波函数模式的知识，而不会丢失。重要的是，这一修改还使我们能够显著减少总参数数量。

- **基于继承的迁移学习**：我们开发了一种迁移学习策略，利用QiankunNet1.2的架构优势实现跨系统泛化。我们的方法依次在化学相关系统上训练单个网络，建立累积的知识储备以捕获基本的量子力学模式。为了证明这一点，我们在四个不同键长的N2分子上训练构建了一个丰富检查点，使模型能够学习电子结构如何随核间距变化。当应用于新的N2分子时，我们发现QiankunNet1.2配合丰富知识检查点即使在拉伸键的情况下也能可靠收敛，而从头训练无法达到化学精度，并且对于较短键也改善了最终能量精度。这些结果表明，QiankunNet1.2的统一相位架构实现了量子关联在不同分子组态之间的真正转移，标志着相对于此前需要在每个新系统上完全重新训练的NNQS方法的重大进步。

在介绍上述方法和实现后，我们对QiankunNet1.2的性能进行了全面评估。我们首先在小型分子上比较QiankunNet1.2和QiankunNet1.0，验证新模型以更少的参数实现了相同的基态精度。接下来，我们展示了QiankunNet1.2在从头训练16个分子时的稳健收敛性，在所有情况下都达到了化学精度。最后，我们分析了N2分子的迁移学习结果，在基于继承的迁移学习策略下比较了QiankunNet1.2和QiankunNet1.0。

本文其余部分组织如下：第2节详细描述了理论框架和QiankunNet1.2的Ansatze。第3节展示了计算结果和讨论，包括从头训练的性能和迁移学习测试。第4节总结了我们的主要发现及其对神经网络量子态方法在化学和物理领域未来发展的潜在影响。

## 方法节选

在本节中，我们回顾了使用变分蒙特卡罗算法求解多体问题的方法。

## 结果节选

我们通过三个全面的实验系统评估了QiankunNet1.2的性能，这些实验旨在测试其关键创新。首先，在第3.1节中，我们通过在多个超参数配置下比较QiankunNet1.2和QiankunNet1.0来分析双相位架构的影响，展示MLP2如何同时增强精度和参数效率。其次，在第3.2节中，我们通过在16个...上测试收敛性来建立QiankunNet1.2对从头计算的可靠性

## 结论节选

我们提出了一种改进的NNQS模型，称为QiankunNet1.2，其可迁移性显著优于其前身。通过在多个不同系统中引入一个具有固定输入维度的额外MLP2，我们确保其参数可以在迁移学习过程中完全转移，无需丢弃任何权重。此外，MLP2的引入使整个模型能够以比QiankunNet1.0更少的参数收敛到相同或更低的基态能量。

## CRediT作者贡献声明

Yangjun Wu：软件。Jiexuan Zhou：撰写-初稿、可视化、验证、调查、形式化分析、数据管理。Yumeng Zhou：调查、数据管理。Zhuozhao Xia：调查、数据管理。Honghui Shang：撰写-审阅与编辑、撰写-初稿、验证、监督、资源、项目管理、调查、资金获取、形式化分析、概念化。

## 利益冲突声明

作者声明他们没有已知的与本文所述工作相关的竞争性财务利益或个人关系。

Yangjun Wu | Jiexuan Zhou | Yumeng Zhou | Zhuozhao Xia | Honghui Shang