Idea Sharing: An Ingenious Solution of Schrodinger Equation: High-Accuracy Wave Function Solution Using FermiNet

Background

Schrodinger equation is a basic equation of quantum mechanics. By solving this equation, most physics and chemistry problems can be solved. However, there is a problem that the number of basis functions in the Schrodinger equation increases exponentially with the increase of the molecular system dimension. Take a methane molecule for example, if the atomic number of methane molecules is 5, the dimension is 9 = 3 x 5 – 6, and the number of basis functions is 109.

Since there is no accurate solution of Schrodinger equation, the chemical properties of molecules can be predicted by approximate solutions with high accuracy. The configuration interaction and coupled-cluster methods have high accuracy, but the calculation cost increases exponentially. Although the density functional theory (DFT) method has low calculation cost, the accuracy is limited. With the powerful nonlinear fitting capability of deep learning, DeepMind proposes FermiNet to implement approximate solution of wave functions.

Electrons not only interact with atomic nuclei and other electrons in a molecule, but also comply with Pauli exclusion principle. Two fermions cannot be in the same quantum state. The wave function after fermions exchange is antisymmetric, that is, two fermions exchange states, and the wave function needs to be reversed. Determinants naturally fit the antisymmetry of wave functions, so Slater determinants are widely used in quantum chemistry to represent wave functions.

1. Network Architecture

Figure 1 The Fermionic neural network (FermiNet) global architecture

Figure 2 Information flow transfer between network layers

Figure 1 shows the overall network architecture of FermiNet. Figure 2 shows the detailed view of a network layer. Each electron in the network has a separate information flow. Besides, when information is transferred at the network layer, each electron fuses information of other electrons and interaction relationships between electrons, and replaces the original one-electron orbital with a multi-electron wave function (see Formula 1) that meets permutation-equivariance. Then the final wave function (see Formula 2) is formed, which is more expressive than the traditional Slater determinant.

Before the network training, it is pretrained to improve the stability and reduce the training time. The solution of the Hartree-Fock equation of the STO-3G basis set is used as a reference for the loss, and the loss function is shown in Formula 3:

For the network training, based on the variational Monte Carlo, the energy expectation value is used as the loss function (see Formula 4). The energy can be represented by Formula 5, and the gradient of the energy is calculated in Formula 6. In addition, to efficiently optimize the network parameters, a modified version of Kronecker-factored approximate curvature (KFAC) is used.

2. Experimental Result

FermiNet accuracy surpasses traditional VMC method (see Table 1). The accuracy is better than that of CCSD (T) method in a finite basis set. Because FermiNet does not use basis sets, so there is no problem of basis set extrapolation.

Table 1 Ground state energy values (the values in bold are values from FermiNet, VMC and DMC that are closest to the exact values)

While CCSD(T) is very accurate for equilibrium geometries, it is limited and not as good as FermiNet for molecules with low-lying excited states or stretched, twisted or otherwise out-of-equilibrium geometries. The results are shown in Figure 3.

Figure 3 Energy curve of molecule H4

For nitrogen triple-bond dissociation of nitrogen molecules, FermiNet performs better than unrestricted CCSD(T) method, as shown in Figure 4.

Figure 4 The dissociation curve for the nitrogen triple-bond

3. Summary

By fusing information between a single electron and multiple electrons at the network layer and replacing a one-electron orbital with a multi-electron wave function that meets permutation-equivariance, FermiNet can solve wave functions with high accuracy. The accuracy of the solution is higher than that of the traditional VMC method. FermiNet is even better than the CCSD(T) method in dealing with challenging architectures such as non-equilibrium geometric structures and nitrogen triple-bond dissociation. The achievement of FermiNet will facilitate researchers to come up with new or better network architectures in quantum chemistry, for wave function solutions that are more efficient and accurate.

References

[1] Pfau D, Spencer J S, Matthews A G D G, et al. Ab initio solution of the many-electron Schrödinger equation with deep neural networks[J]. Physical Review Research, 2020, 2(3): 033429.

DOI: https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.2.033429