Electric Polarization from Many-Body Neural Network Ansatz
Ab initio calculation of dielectric response with high-accuracy electronic
structure methods is a long-standing problem, for which mean-field approaches
are widely used and electron correlations are mostly treated via approximated
functionals. Here we employ a neural network wavefunction ansatz combined with
quantum Monte Carlo to incorporate correlations into polarization calculations.
On a variety of systems, including isolated atoms, one-dimensional chains,
two-dimensional slabs, and three-dimensional cubes, the calculated results
outperform conventional density functional theory and are consistent with the
most accurate calculations and experimental data. Furthermore, we have studied
the out-of-plane dielectric constant of bilayer graphene using our method and
re-established its thickness dependence. Overall, this approach provides a
powerful tool to consider electron correlation in the modern theory of
polarization.
Xiang Li, Yubing Qian, Ji Chen
# Electric Polarization from Many-Body Neural Network Ansatz

###### Abstract

_Ab initio_ calculation of dielectric response with high-accuracy electronic structure methods is a long-standing problem, for which mean-field approaches are widely used and electron correlations are mostly treated via approximated functionals. Here we employ a neural network wavefunction ansatz combined with quantum Monte Carlo to incorporate correlations into polarization calculations. On a variety of systems, including isolated atoms, one-dimensional chains, two-dimensional slabs, and three-dimensional cubes, the calculated results outperform conventional density functional theory and are consistent with the most accurate calculations and experimental data. Furthermore, we have studied the out-of-plane dielectric constant of bilayer graphene using our method and re-established its thickness dependence. Overall, this approach provides a powerful tool to consider electron correlation in the modern theory of polarization.

Electric polarization plays a crucial role in electromagnetic phenomena such as ferroelectricity and piezoelectricity. Despite its significance, a proper microscopic definition of polarization was only formulated in the 1990s [1; 2], which revealed the hidden relation between physical polarization and the Berry phase of solid systems. This theoretical advance leads to successful calculations of the dielectric response of solid materials from first principles [3; 4; 5], which is critical in several fields of condensed matter physics, such as the ferroelectric and topological materials [6]. However, the underlying electronic structure methods are mostly mean-field approaches such as density functional theory (DFT) [7], which has its limitation because the result depends heavily on the so-called exchange-correlation functional. Exchange-correlation functionals can not fully account for the exact correlation effects of electrons. In particular, widely used semi-local functionals often produce an excessive overestimate of electric susceptibility [5; 8]. Although correlated wavefunction methods, such as coupled-cluster theory can also be employed to calculate polarization [9], their high computational complexity hinders their application in solid systems. Furthermore, most of these correlated electronic structure methods are limited in open boundary conditions (OBC) for polarization calculations, which leads to slow convergence and heavy computational costs towards the thermodynamic limit (TDL), see Fig. 1 for a summary of the state-of-the-art methods in polarization calculations.

In addition to the conventional deterministic electronic structure methods mentioned above, quantum Monte Carlo (QMC) methods are also widely adopted for electronic structure calculations, showing favorable computational scaling and high accuracy [10; 11; 12]. Pioneering works to study electric susceptibility using QMC have been reported [13; 14], in which traditional Slater Jastrow type wavefunction is combined with diffusion Monte Carlo (DMC) to study polarization of hydrogen chains in periodic boundary conditions (PBC). The main difficulty for DMC is to write down the local self-consistent Hamiltonian under a finite electric field and run calculations iteratively. Despite the promising results on hydrogen chain [13], there are still grand challenges: multiple loops of DMC are needed for the self-consistent procedure, a complex forward walking strategy is required for evaluat

Figure 1: A brief illustration of the computational cost and accuracy of different electronic structure methods in polarization calculations. \(N\) denotes the number of electrons in the system. High-level correlated wavefunction methods are not shown in the PBC panel because they have not been applied to PBC polarization calculations so far. QMC: quantum Monte Carlo; DFT: density functional theory; HF: Hartree-Fock; MP2: second-order Moller-Plesset perturbation theory; CCSD: coupled cluster with single and double excitations; FCI: full configuration interaction.
ing the polarization, and the quality of trial wavefunction affects the accuracy of DMC. Therefore, it is desirable to develop more accurate and efficient approaches to calculate the electric polarization of solid systems.

In recent years, there has been significant progress in the application of neural network in the electronic structure community. Neural network wavefunction ansatz combined with QMC simulations has demonstrated higher accuracy with lower computational complexity than conventional high-order wavefunction methods [15; 16; 17; 18; 19; 20; 21; 22; 23; 24; 25; 26; 27; 28; 29; 30; 31; 32; 33]. The expressiveness of neural network overcomes the main bottleneck of traditional wavefunction ansatz in QMC, making the approach a competitive option for state-of-the-art electronic structure calculation. So far, the neural network QMC calculations have shown great power in treating spin systems [15; 16; 17], molecules [18; 19; 20; 21], periodic models [25; 26; 27; 28], and real solids [29; 28].

In this work, we extend the neural network QMC calculation to the electric polarization of solid systems. Specifically, we employ a recently developed solid neural network, dubbed DeepSolid [28], in conjunction with variational Monte Carlo (VMC). Anthetic sampling [34] is developed for efficient computation of the Berry phase and thus the electric polarization. Our approach has been tested on a diverse range of systems, including isolated atoms, one-dimensional chains, two-dimensional slabs, three-dimensional cubes, and bilayer graphene. The results demonstrate clear advantages of our approach over traditional methods.

To introduce our methodology, let us consider a crystal system under a finite electric field \(\mathbf{E}\), the enthalpy of this system is formulated below [3; 4; 13; 37]

\[F[\psi]=\frac{\langle\psi|\hat{H}_{S}|\psi\rangle}{\langle\psi|\psi\rangle}- \Omega_{S}\mathbf{E}\cdot\mathbf{P}[\psi]\, \tag{1}\]

where \(\hat{H}_{S}\) denotes the supercell Hamiltonian in the absence of electric field \(\mathbf{E}\), and \(\Omega_{S}\) is the supercell volume. The term \(-\Omega_{S}\mathbf{E}\cdot\mathbf{P}\) represents the interaction between electric polarization density \(\mathbf{P}\) and electric field. However, a proper microscopic definition of \(\mathbf{P}[\psi]\) remains absent for decades since the ordinary position operator \(\hat{\mathbf{r}}\) violates the periodic boundary condition. This problem was finally solved after recognizing the polarization as the Berry phase in the Brillouin zone, according to which the polarization can be extracted from a general wavefunction \(\psi\) as follows [38]

\[\begin{split}\mathbf{P}[\psi]&=-\frac{1}{\Omega_{S}} \sum_{i}\frac{\mathbf{a}_{i}}{2\pi}\text{Im}\ln\frac{\langle\psi|\hat{U}_{i}| \psi\rangle}{\langle\psi|\psi\rangle}\,\\ \hat{U}_{i}&=\exp\left[\text{i}\mathbf{b}_{i}\cdot \left(\sum_{e}\hat{\mathbf{r}}_{e}-\sum_{I}Z_{I}\mathbf{R}_{I}\right)\right]\,\end{split} \tag{2}\]

where \(\mathbf{a}_{i},\mathbf{b}_{i}\) denote lattice and reciprocal lattice vectors of the supercell. \(\hat{U}_{i}\) serves as a periodic generalization of the position operator \(\hat{\mathbf{r}}\) in solid systems and \(\text{Im}\ln(x)\) is used to extract the Berry phase within \(x\). Note that \(\hat{U}_{i}\) is an intrinsic many-body operator which includes all the electron coordinates in the exponent. A charge-weighted sum of ion coordinates \(Z_{I}\mathbf{R}_{I}\) is also included to achieve translation invariance of polarization.

With the enthalpy functional formulated above, traditional methods usually start with a Hartree-Fock (HF) ansatz, which is typically expressed as follows

\[\psi_{\text{HF}}(\mathbf{r})=\text{Det}\left[e^{\text{i}\mathbf{k}_{i}\cdot \mathbf{r}_{j}}u_{\mathbf{k}_{i}}(\mathbf{r}_{j})\right]. \tag{3}\]

Electrons are treated independently of each other with a mean-field interaction in Eq. (3), simplifying quantum many-body problems but also deviating from the ground truth. To fully treat the electron correlation effects, we employ a correlated neural network wavefunction \(\psi_{\text{net}}\) from DeepSolid [28], whose general form reads

\[\psi_{\text{net}}(\mathbf{r})=\text{Det}\left[e^{\text{i}\mathbf{k}_{i}\cdot \mathbf{r}_{j}}u_{\mathbf{k}_{i}}(\mathbf{r}_{j};\mathbf{r}_{\neq j})\right], \tag{4}\]

where \(\mathbf{r}_{\neq j}\) denotes all the electron coordinates except \(\mathbf{r}_{j}\). Eq. (4) resembles the form of the traditional Bloch function, while cell-periodic functions \(u_{\mathbf{k}}\) are now represented using deep neural networks that rely on all electrons to accommodate electron correlations [19]. Electron features \(\mathbf{r}_{i}\) are converted to be periodic and permutation equivariant before being fed into neural networks, and complex-valued orbitals \(u_{\mathbf{k}}\) are constructed with a pair of neural networks outputting the real and imaginary part respectively. As a result, Fermionic anti-symmetry, periodicity, and complex-valued nature are all encoded in our network, promoting it to be a legitimate and expressive ansatz for solid. See Ref. [28] for more details of the architecture.

Using the neural network we have constructed, the enthalpy functional outlined in Eq. (1) can be efficiently minimized through variational Monte Carlo, allowing for gradual convergence to the ground truth. However, there have been significant fluctuations observed in \(\hat{U}_{i}\) evaluation, which seriously impedes optimization. As a solution, antithetic sampling is employed in the Monte Carlo evaluation, which reads

\[\begin{split}\langle\hat{U}_{i}\rangle&=\frac{ \int d\mathbf{r}\ |\psi(\mathbf{r})|^{2}\ U_{i}(\mathbf{r})}{\int d\mathbf{r}\ |\psi(\mathbf{r})|^{2}}=\frac{\int d\mathbf{r}\ |\psi(\mathbf{r})|^{2}\ \widetilde{U}_{i}(\mathbf{r})}{ \int d\mathbf{r}\ |\psi(\mathbf{r})|^{2}},\\ &\widetilde{U}_{i}(\mathbf{r})=\frac{1}{2}\left[U_{i}(\mathbf{r} )+\frac{|\psi(-\mathbf{r})|^{2}}{|\psi(\mathbf{r})|^{2}}U_{i}(-\mathbf{r}) \right].\end{split} \tag{5}\]

And thus the fluctuations are significantly reduced through the cancellation between \(U_{i}(\mathbf{r})\) and its inverted image \(U_{i}(-\mathbf{r})\). It's worth noting that centrosymmetric cells are assumed in Eq. (5), and one can choose other images for cancellation if central symmetry is not satisfied. To further improve efficiency, we have employed a Kronecker-factored curvature estimator (KFAC) optimizer [39], which effectively integrates second-order information into the optimization process, surpassing traditional optimizers. See the Supplementary Material for more computational details, and the code of this work is
developed at the open-source repository of DeepSolid 1.

Footnote 1: [https://github.com/bytedance/DeepSolid](https://github.com/bytedance/DeepSolid)

Isolated atoms are the first systems selected for direct comparison with the most accurate methods and experimental data. In our calculations, we place a single atom in a large enough box to eliminate periodic image interactions. The calculated polarizability \(\alpha\) is shown in Tab. 1, which measures the linear response of the dipole moment to the applied field and has some subtle relation with the bulk susceptibility \(\chi\) (see Supplementary Material). Results from DFT with the B3LYP functional, HF, and CCSD(T) under OBC are also listed for comparison. P-state atoms (B, C, O, F) are skipped because their anisotropy requires special treatments [40]. As can be seen from Tab. 1, although B3LYP is widely-trusted functional belonging to the fourth rung of the so-called Jacob's ladder of DFT, it consistently deviates from the ground truth and has a relatively large mean absolute error (MAE). The behavior of DFT is due to the inaccuracy in treating the exchange-correlation effects, which can be very different for energy and polarization calculations. In HF calculations, because of the explicit treatment of non-local exchange, deviations in polarization are significantly reduced. CCSD(T) is the coupled cluster theory with single, double, and perturbative triple excitations, and is considered a very accurate method in the literature. It further incorporates correlation effects on top of HF wavefunction and achieves smaller MAE than HF results. Overall, DeepSolid results are comparable with CCSD(T), showing that the exchange-correlation treatments in our neural network are accurate and reliable for polarization calculations.

Having demonstrated our technique with single atoms, we proceed to simulate periodic systems by arranging bonded molecules into a one-dimensional chain and a two-dimensional slab. These systems are widely known as challenging cases for conventional DFT methods, which would have a serious overestimation of their longitudinal susceptibility. This problem stems from the fact that surface charges are insensitive to the bulk charge within the system when non-local interactions are absent, and this can be solved using more accurate _ab initio_ methods [5]. For the one-dimensional case, hydrogen chain (n H\({}_{2}\)) and polyyne (n C\({}_{2}\)) are studied, and the simulation size is pushed to 22 H\({}_{2}\) and 7 C\({}_{2}\) respectively for TDL convergence. Correlation-consistent effective core potential (ccECP) is employed for polyyne to accelerate neural network optimization and reduce fluctuation [31; 44]. The final results are plotted in Fig. 2, which show that susceptibility calculated by DeepSolid agrees well with correlated wavefunction methods CCSD(T) and random phase approximation (RPA). Local-density approximation (LDA) functional deviates severely from the ground truth for one-dimensional chains [9], but the use of hybrid functions such as B3LYP leads to partial recovery of non-local exchange effects and, consequently, a reduction in the overshot. HF is much better than DFT calculations, which further proves the importance of the

\begin{table}
\begin{tabular}{c c c c c c|c|c}  & H & He & Li & Be & N & Ne & MAE \\ \hline B3LYP & 5.187 & 1.485 & 142.727 & 43.090 & 7.711 & 2.838 & 4.669 \\ HF & 4.484 & 1.318 & 169.231 & 45.441 & 7.138 & 2.365 & 2.243 \\ CCSD(T) & 4.484 & 1.372 & 165.803 & **37.707** & **7.212** & 2.642 & 0.326 \\ DS & **4.51** & **1.39** & **165.0** & 36.95 & 7.16 & **2.67** & 0.32 \\ Recommended & 4.5(exact) & 1.38375(2) & 164.1125(5) & 37.74(3) & 7.4(2) & 2.66110(3) & 0 \\ \end{tabular}
\end{table}
Table 1: Calculated atom polarizability in atomic unit (Bohr\({}^{3}\)). DeepSolid results are labeled as DS. B3LYP, HF, and CCSD(T) results are calculated with PySCF [35] in the def2-qxyppd basis set and non-relativistic limit. Recommended data is taken from Ref. [36], which is deduced from experiment data and the most accurate calculations.

Figure 2: Calculations of chains and slabs. DeepSolid results are labeled as DS. (**a**) illustrations of hydrogen chain, polyyne, hydrogen slab, and the applied electric field. (**b**) hydrogen chain, (**c**) polyyne, (**d**) hydrogen slab susceptibilities \(\chi\). \(\Omega_{p}\) denotes the volume of the primitive cell. For the hydrogen system, LDA and HF results were calculated with the 3-21G basis set under PBC [5]. CCSD(T) calculations were performed with the 6-311G** basis set under OBC [9]. Intrapair and interpair distances are set to 2 and 3 Bohr respectively for the hydrogen chain. The interchain distance is set to 4.724 Bohr for the slab. For polyyne, the alternating distance between carbon atoms is set to 1.18 and 1.4 Å. LDA, HF, and RPA results were calculated with 6-31G, 3-21G, and 4-31G basis set respectively under OBC [41; 42; 43].
non-local exchange effect for electric polarization calculation in this system. As we arrange hydrogen chains periodically to form hydrogen slabs, the computational cost of high-level deterministic wavefunction methods, such as CCSD(T), grows rapidly and is soon beyond reach. However, our approach has a lower scaling and we can obtain the first accurate polarization calculation for such a hydrogen slab (Fig. 2d). And for the slab, the performances of DFT and HF compared with our accurate neural network results are similar to those observed for the chains.

To further test our method, we applied it to alkali metal hydrides and calculated their dielectric constants, allowing direct comparison with experimental results. These systems have a simple structure, consisting of alternating cations and anions, but they are of considerable research significance due to their relevance in hydrogen storage applications [49]. The high-frequency dielectric constant \(\epsilon_{\infty}\) can be extracted through optical experiments from the following relations,

\[\begin{split}&\mathbf{D}=\epsilon_{\infty}\mathbf{E}=\mathbf{E}+4 \pi\mathbf{P}\,\\ &\epsilon_{\infty}=1+4\pi\chi=n_{\mathrm{D}}^{2}\,\end{split} \tag{6}\]

where \(n_{\mathrm{D}}\) denotes the corresponding refractive index. In the visible light regime, ions are almost frozen relative to the incident light frequency and this leads to the dominance of electric polarization in \(\epsilon_{\infty}\). These three-dimensional systems are also qualitatively different from the chains and slabs, because the fluctuation of \(\hat{U}_{i}\) increases as we tile the cells in all three direction to form the simulation cell. To balance the influence from finite-size error and \(\hat{U}_{i}\) fluctuations, we tile the conventional cell in the direction of the applied field \(\mathbf{E}\) and the transverse directions remain unchanged. Moreover, Burkatzki-Filippi-Dolg (BFD) pseudopotential [50] is used to remove inertial core electrons [31]. Our calculations have been pushed to the \(4\times 1\times 1\) supercell and the results are plotted in Fig. 3. LDA and Perdew-Burke-Ernzerhof (PBE) results are also plotted for comparison, while more accurate conventional wavefunction methods are not applicable due to computational costs. As we can see, numerical simulations and experiments agree that \(\epsilon_{\infty}\) decreases as the alkali metal atom becomes heavy, since \(\epsilon_{\infty}\) is inversely proportional to the cell volume in Eq. (6). However, LDA and PBE functionals [51] tend to overestimate \(\epsilon_{\infty}\), and the error is largest in CsH. In contrast, our DeepSolid results agree well with the experiment for all systems, which manifests the capability of neural network wavefunction to capture non-local exchange and correlation effects.

After demonstrating the accuracy of our methods in previous sections, we now proceed to apply our method to bilayer graphene (BLG), an extensively studied 2D material system known for its rich electronic properties. Despite its fundamental importance, the precise value of the dielectric constant of BLG remains elusive and has been an important subject of both experimental and theoretical works [52; 53]. Specifically, theoretical calculations reported were either restricted to DFT level [52] or based on values calculated with monolayer graphene [53]. Here we use DeepSolid to directly calculate the out-of-plane dielectric constant \(\epsilon_{\infty}^{\perp}\) of bilayer graphene. \(2\times 2\) supercells containing monolayer and equilibrium AA-stacked bilayer graphene were used. The calculated monolayer polarizability equals 5.7(1) Bohr\({}^{3}\) and bilayer polarizability equals 11.6(1) Bohr\({}^{3}\), which agrees with the linear dependence of polarizability on the number of layers as shown in Ref. [52]. Based on this linear dependence and following Ref. [54], one can derive the expression of the out-of-plane dielectric constant as a function of the layer separation \(d\) (see Supplementary Material):

\[\epsilon_{\infty}^{\perp}(d)=\left(1-\frac{2\pi\alpha_{\mathrm{equil}}^{ \mathrm{BLG}}}{S\cdot d}\right)^{-1}\, \tag{7}\]

where \(S=5.25\) A\({}^{2}\) denotes the area of the primitive cell. Using the computed polarizability \(\alpha_{\mathrm{equil}}^{\mathrm{BLG}}\), we can re-establish the relation of \(\epsilon_{\infty}^{\perp}\), which is plotted in Fig. 4. To further check Eq. (7), we also calculate the polarizability of bilayer graphene at slightly larger (4A) and smaller

Figure 4: Calculated effective 2D dielectric constant \(\epsilon_{\infty}^{\perp}\) of bilayer graphene. (**a**) plot of bilayer graphene under electric field. (**b**) calculated \(\epsilon_{\infty}^{\perp}\) as a function of graphene layer distance \(d\). The equilibrium separation of BLG is set to 3.347 Å.

Figure 3: Calculated high-frequency dielectric constant \(\epsilon_{\infty}\) of alkali metal hydrides XH. DeepSolid results are labeled as DS. LDA and PBE results in PBC are taken from Ref. [45]. Experimental data are taken from Refs. [46; 47; 48], which are derived from refractive indexes \(n_{\mathrm{D}}\) for sodium doublet (589.29 nm) via Eq. (6). RbH experiments are absent.
(3A) layer distances and plot the corresponding dielectric constant in Fig. 4, and the results agrees well with each other. Moreover, there are two notable limits when varying the layer separation \(d\): as \(d\) decreases, two graphene layers coincide with each other and the system becomes metallic, which explains the diverging of \(\epsilon_{\infty}^{\perp}\); as \(d\) becomes large, BLG polarization becomes negligible and vacuum contribution dominants in \(\epsilon_{\infty}^{\perp}\) which approaches unity. The thickness-dependent dielectric constant will be valuable for further understanding and tuning the stacked multilayer graphene systems.

In conclusion, this work proposes an efficient and accurate method for investigating solid polarization based on the recently developed solid neural network wavefunction combined with quantum Monte Carlo. Our approach demonstrates superiority over the conventional state-of-the-art electronic structure methods. In the future, with the proposed framework, it is promising to investigate a wide range of phenomena, including ferroelectricity, topological electronic transport, quantum Hall effect, and orbital magnetization, among others, on a higher level of accuracy and with electron correlations accounted for properly. Furthermore, this work provides more possibilities for utilizing neural network applications in condensed matter physics.

## Acknowledgements

We want to thank ByteDance Research Group for inspiration and encouragement. This work is directed and supported by Hang Li and ByteDance Research. J.C. is supported by the National Natural Science Foundation of China under Grant No. 92165101.
