Machine learning approach for vibronically renormalized electronic band structures

The fundamental assumption of most ab initio electron-phonon calculations, especially the frozen-phonon methods, is the Born-Oppenheimer (BO) or adiabatic approximation. The fact that masses of nuclei are much larger than that of electron indicates that the electron velocities are much larger than the nuclear ones, suggesting that electrons can follow the motion of the slow nuclei almost instantaneously. The well-separated time scales of the two sets of degrees of freedom allows one to write the full nuclei-electron wave function as a product form $|\Psi\rangle=|\chi\rangle\otimes|\psi\rangle$, where $|\chi\rangle$ and $|\psi\rangle$ are the phonon and electron wave functions, respectively. Under the adiabatic approximation, one can integrate out the electron degrees of freedom to obtain a potential energy surface $\mathcal{E}(\mathbf{R}_{1},\mathbf{R}_{2},\cdots,\mathbf{R}_{N})$ as a function of the coordinates $\mathbf{R}_{i}$ of the nuclei. Practically, this involves solving an electronic structure problem with each atomic configuration frozen in $\mathbf{R}_{i}$ based on DFT calculations, or other electronic structure methods.

The Hamiltonian of nuclei in BO approximation is given by $\mathcal{H}_{\rm nucl}=\sum_{i}\mathbf{P}_{i}^{2}/2M_{i}+\mathcal{E}(\{\mathbf{R}_{i}\})$, where $\mathbf{P}_{i}$ is the momentum operator of the $i$-th nuclei and $M_{i}$ is the corresponding mass. However, since the computation of the potential energy for each atomic configuration $\{\mathbf{R}_{i}\}$ requires solving an electronic structure problem, direct ab initio determination of the phonon modes is computationally intractable, if not impossible. As discussed in Sec. I, ML force-field approaches offer an efficient and accurate approximation to the potential energy surface. However, due to the lack of explicit analytical expressions for $\mathcal{E}(\{\mathbf{R}_{i}\})$, the phonon modes cannot be directly derived from the ML model. Vibrational thermal average has to be done in conjunction with dynamical simulations based on the ML force-field models.

Direct phonon or lattice models can be obtained based on the harmonic approximation to the nuclear Hamiltonian $\mathcal{H}_{\rm nucl}$. For crystalline systems, nuclei only undergo small amplitude oscillations about their equilibrium positions. We can then write the nuclear positions as $\mathbf{R}_{i}=\mathbf{R}^{(0)}_{i}+\mathbf{u}_{i}$, where $\mathbf{R}^{(0)}_{i}$ is the equilibrium position and $\mathbf{u}_{i}$ is the displacement vector. Assuming small displacements, one can then express the potential energy $\mathcal{E}(\{\mathbf{R}_{i}\})$ in a power series of the displacement vectors. Since the equilibrium positions minimize the potential energy, the linear term in the expansion vanishes. To the leading second-order, the potential energy is $\mathcal{E}=\mathcal{E}_{0}+\sum_{ij}\sum_{\alpha,\beta}D_{i\alpha,j\beta}u_{i}^{\alpha}u_{j}^{\beta}$, where $u_{i}^{\alpha}$ denotes the $\alpha$-component of the displacement vector of the $i$-th nucleus ($\alpha=x,y,$ and $z$), and $D_{i\alpha,j\beta}=\partial^{2}\mathcal{E}/\partial u_{i,\alpha}\partial u_{j,\beta}$, evaluated at the equilibrium positions $\{\mathbf{R}^{(0)}\}$, is the dynamical matrix which encapsulates the normal modes of the lattice.

The normal modes $\mathbf{V}^{(\nu)}_{s,\mathbf{k}}$ of the crystalline systems are obtained by diagonalizing the harmonic Hamiltonian, here $\nu$ is the branch-index, $s$ is the sublattice index, and $\mathbf{k}$ is the reciprocal lattice wave vector. The displacement vectors can then be expressed as

where $\xi_{\nu,\mathbf{k}}$ are the effective or normal-mode coordinates. In terms of the effective coordinates, the vibrational dynamics of the lattice is described by the following phonon Hamiltonian

Here $\Omega_{\nu{\bf k}}$ is the eigen-frequency of the normal modes. Ab initio calculation of the phonon modes can now be routinely computed using DFT methods, both based on the DFPT or frozen-phonon formalisms. Computation of the dynamical matrices using methods beyond DFT have also been demonstrated [70, 71].

Next we consider the computation of observables such as in the frozen-phonon framework. Quantum mechanically, an observable corresponds to a Hermitian operator $\mathcal{O}(\{\mathbf{r}_{l}\},\{\mathbf{R}_{i}\})$ which depends on both electron and nuclei coordinates. By integrating out the electronic degrees of freedom within the Born-Oppenheimer approximation, the observable

becomes a function of nuclear coordinates only. The electron quantum state $|\psi(\{\mathbf{r}_{l}\})\rangle$ in general also implicitly depends on the nuclei configuration $\{\mathbf{R}_{i}\}$. The vibrational expectation value of the observable quantity at finite temperature $T$ is then given by [49]

where the summation is carried over nuclear quantum states $|\chi_{M}(\{\mathbf{R}_{i}\})\rangle$ whose energy is $\mathcal{E}_{M}$, $\mathcal{Z}=\sum_{M}\exp[-\mathcal{E}_{M}/k_{B}T]$ is the partition function of phonons and $k_{B}$ is the Boltzmann constant. In this paper, the observable of interest is the electronic band gap $E_{g}$ but the formalism applies to any other observables. Similar formulas can also be derived for matrix elements $\mathcal{A}_{jk}(\{\mathbf{R}_{i}\})$ between different electronic states; the resultant expectation values will determine the transition rates such as optical absorption and dielectric coefficients [72, 73].

In the harmonic approximation, the quantum number of the nuclei is given by an array of integers $M=\{n_{\nu,\mathbf{k}}\}$, and the corresponding wave function becomes a product state: $\chi_{M}(\{\mathbf{R}_{i}\})\propto\prod_{\nu,\mathbf{k}}H_{n_{\nu,\mathbf{k}}}(\xi_{\nu,\mathbf{k}})\exp[-\xi^{2}_{\nu,\mathbf{k}}/2]$, where $H_{n}(x)$ is the $n$-th order Hermite polynomial, $\mathbf{R}_{i}=\mathbf{R}^{(0)}_{i}+\mathbf{u}_{i}$, and the displacement vectors $\mathbf{u}_{i}$ is expressed in terms of the normal mode coordinates $\xi_{\nu,\mathbf{k}}$ through Eq. (1). As the atomic displacements will play a central role in this work, we introduce the notation $\mathbf{U}=\{\mathbf{u}_{1},\mathbf{u}_{2},\cdots,\mathbf{u}_{N}\}$ to denote the collection of displacement vectors. Within the quadratic approximation, the summation over the integer quantum numbers $n_{\nu,\mathbf{k}}$ in Eq. (3) can be recast into integrals over the normal mode coordinates:

where the width $\sigma_{\nu,\mathbf{k}}$ of the Gaussian integral is given by:

and $n_{\nu{\bf k},T}=[e^{\hbar\Omega_{\nu{\bf k}}/k_{B}T}-1]^{-1}$ is the Bose-Einstein occupation of the $(\nu,\mathbf{k})$ phonon. The displacement configuration $\mathbf{U}$ in $\mathcal{O}(\mathbf{U})$ of the Gaussian integrals is again expressed as superposition of the phonon modes as in Eq. (1).

The computation of the multidimensional Gaussian integral, however, is computationally challenging for observables with a complex dependence on the nuclei configuration. The calculation can be simplified by expanding the observables as a power series expansion of displacement vectors, often up to second order, the Gaussian integral can then be analytically evaluated. Within the DFT framework, the expansion coefficients can be similarly computed using DFPT [30, 74, 75, 76] or frozen-phonon methods [43, 46].

An alternative approach, which also allows one to go beyond perturbation theory and DFT, is the stochastic Monte Carlo method [52, 51, 53]. To this end, we rewrite the integral in Eq. (4) as

Here $\mathcal{D}\mathbf{U}$ is a simplified notation for the multiple integrals over the effective coordinates $\xi_{\nu,\mathbf{k}}$, and $\pi_{G}(\mathbf{U})$ is an effective probability density function of $\mathbf{U}$, corresponding to the product of Gaussian density functions. The Metropolis-Hastings algorithm is then used to implement a transition probability $P(\mathbf{U}\to\mathbf{U}^{\prime})={\rm min}\left[1,\pi_{G}(\mathbf{U}^{\prime})/\pi_{G}(\mathbf{U})\right]$ for sampling the displacement configurations. A series of displacement configurations $\mathbf{U}^{(k)}$ sampled from the Markov-Chain Monte Carlo is used to compute the vibrational average as indicated by the approximation in Eq. (6) above.

Having obtained an ab initio phonon model, sampling of the nuclei configuration, or more specifically the displacement field $\mathbf{u}$, is relatively straightforward. The bottleneck of the stochastic approach to vibrational thermal averaging is the computation of the observable $\mathcal{O}(\mathbf{U})$ for the sampled configurations. As discussed in Sec. I, to overcome this computational complexity, a fully deterministic method based on the idea of effective vibrational configurations is proposed to replace the stochastic MC sampling. It has been shown that the vibrational average can be approximated by either a single nuclear configuration, called special displacement configuration [64], or a reduced MC sampling along so-called thermal lines [62]. While these methods greatly simplify the problem and have been shown to work reasonably well for a wide range of materials [77], they are essentially based on quadratic approximation of electron-phonon theory and therefore may fail to capture higher order effects in certain class of materials [78]. In addition, it has been shown very recently that the special displacement method fails to correctly account for the renormalization of energy levels of nitrogen vacancy centers in diamond [79]. Hence, it is desirable to develop an efficient vibrational averaging method without sacrificing too much accuracy.

Modern ML methods, especially supervised learning, offer a promising solution to the issue of computational efficiency by providing an accurate and efficient mapping from the phonon configuration $\mathbf{U}$ to the electronic observable $\mathcal{O}$ of interest. Fundamentally, this approach relies on the universal approximation theorem which shows that deep multilayer neural networks (NN) can be trained to accurately approximate any given high-dimensional functions [80, 81]. A schematic diagram of our proposed ML model is shown in FIG. 1; the input of the model is a frozen phonon configuration represented by a displacement field $\mathbf{U}$ and the temperature $T$, while the output is the predicted observable corresponding to the expectation value $\mathcal{O}(\mathbf{U})=\langle\psi|\hat{\mathcal{O}}|\psi\rangle$ where $|\psi\rangle$ is calculated from some electronic structure methods. Using an ab initio phonon model discussed in Sec. II.1 to generate a series of displacement fields $\mathbf{U}^{(r)}_{T}$ at a few pre-determined temperatures $T$, the electronic states $|\psi^{(r)}\rangle$ are then solved consistently based on the same first-principle methods. A dataset of the form $\left\{\mathbf{U}^{(r)}_{T},T;\,\mathcal{O}^{(r)}\right\}$ is used to train the above ML model.

Broadly speaking, our proposed ML model for $\mathcal{O}(\mathbf{U})$ can be viewed as a special case of ML-based modeling of structure-property relationships in materials science [65, 66, 67, 82, 69, 83, 84]. Yet, while the inputs to most conventional ML structure-property model are static disordered or quasi-random structures of the system, a frozen phonon configuration as represented by a displacement field $\mathbf{U}$ describes an instantaneous structure of nuclei. In this sense, our ML model is also similar to the ML-based interatomic potential models which maps a dynamical atomic configuration in a finite range to a local atomic energy [7, 8, 9, 10, 11, 12, 13, 14, 15].

As shown in Fig. 1, there are two central components of the ML model: (i) a learning model based on deep multi-layer neural networks (NN) and (ii) a descriptor which provides generalized or effective coordinates of frozen phonons. The ML model works as follows. First, the set of displacement vectors $\mathbf{U}=\{\mathbf{u}_{i}\}$ that characterizes the atomic configuration within a cubic box is mapped to a collection of symmetry-invariant feature variables $\bm{G}=\{G_{1},G_{2},\cdots\}$, also known as a descriptor. These feature variables are then fed into the NN which produces the predicted observable $\mathcal{O}(\mathbf{U})$ at the output node, thus bypassing the time-consuming electronic structure calculations. Integrating this ML phonon-to-property model into the stochastic framework offers an efficient yet accurate vibrational thermal average.

First principles density-functional-theory (DFT) calculations were done using Quantum Espresso (QE) [99] package. Perdew-Zunger (PZ) exchange-correlation functional [100] within the local density approximation (LDA) was used in all the calculations. Full lattice relaxation of primitive cell of Si crystal was performed which gave lattice parameter of 5.398 Å consistent with previous works [49]. The primitive BZ was sampled by using k-mesh of 24 $\times$ 24 $\times$ 24 and energy cutoff of 30 Rydberg was used after careful convergence tests. The relaxed lattice was used to perform ab initio phonon calculation using DFPT [30, 31, 32] within the harmonic approximation as implemented in QE. For phonon calculations, we used a q-mesh of 8 $\times$ 8 $\times$ 8. From the information of phonon eigen-frequencies $\Omega_{\nu,\mathbf{k}}$ and eigenvectors $\mathbf{V}^{(\nu)}_{s,\mathbf{k}}$, we generated atomic supercell configurations of size 6 $\times$ 6 $\times$ 6 primitive unit cells containing 432 Si atoms using importance sampling Monte Carlo according to Eq. (4). Subsequently, DFT self-consistent calculation were performed on about 10% of the distorted atomic supercell configurations using k-mesh of 3 $\times$ 3 $\times$ 3 to generate a database of electronic band gap as a function of atomic configurations.

In FIG. 3, we show the electron and phonon dispersion relations in the ground state of silicon crystal obtained from DFT calculations. These results are consistent with previous literature [101]. DFT calculations show an indirect band gap of 500 meV which is severely underestimated compared to the experimental band gap of 1.12 eV. Such band gap underestimation of insulators and semiconductors is a well known problem in DFT and more sophisticated methods like GW calculations are known to give better estimation of the band gap. In practice, GW correction of band gap is similar to the “scissor shift” of bands. Since we are interested in the band gap reduction induced by thermal phonons, this underestimation of band gap should not be an issue as the same GW-correction (to leading order) appears in both the perfect and perturbed crystal structures. Hence, in the following, we will use phonon induced band gap correction $\Delta E_{g}=E_{g}^{\rm eq}-E_{g}^{\rm distorted}$ as our observable of interest, where $E_{g}^{\rm eq}$ is the band gap of the perfect crystalline structure and $E_{g}^{\rm distorted}$ is the band gap of the distorted structure.

The phonon dispersion shown in FIG. 3 is also consistent with previous studies. By sampling the normal mode coordinates $\xi_{\nu\mathbf{k}}$ according to Eq. (4), a displacement configuration $\mathbf{U}=\{\mathbf{u}_{i}\}$, or a frozen phonon configuration, within the supercell is constructed according to Eq. (1). Fixing the nuclei positions at $\mathbf{R}_{i}=\mathbf{R}_{i}^{(0)}+\mathbf{u}_{i}$ in the supercell, DFT calculations were performed to obtain the optimized electron density distribution. An energy gap $E_{g}(\mathbf{U})$ corresponding to a particular phonon configuration is computed as the difference between the lowest unoccupied and the highest occupied states. In FIG. 4, we show histogram plot of the band gap correction $\Delta E_{g}$ for a sample size of 100 at two different temperature values. As seen from the plots, the MC method has not fully converged and needs large number of samples to get a better estimation of the band gap. In the following, we show that it is possible to get a better estimation (i.e. smaller standard deviation) of the band gap without incurring additional computational cost by first using the DFT calculated dataset to train a ML phonon-to-band gap mapping, and then using the ML model to predict band gap for a large number of configurations to obtain thermal average.

As shown in FIG. 1, our learning model is based on a fully-connected neural network (NN), also known as a multilayer perceptron model. A fully-connected NN comprises of a series of fully connected linear layers with non-linear activation functions that are applied at each layer. The parameters of the $j$-th layer are the matrix of network weights $\mathbf{W}^{(j)}$ and bias vector $\mathbf{b}^{(j)}$. The weights introduce coupling between neurons of adjacent layers. The mapping of neurons $\bm{x}$ from one layer to the next is achieved by a linear transformation followed by a nonlinear activation function as

A commonly used activation function is the rectified linear unit (ReLU) [102] defined as $f_{\rm av}(x):=\max(x,0)$. However, sometimes a large number of ReLU neurons remains inactivated in the model. Hence, often some variants of leaky ReLU [103] are used instead. In this work, we used the Gaussian Error Linear Units (GeLU) [104] activation function, which shows better performance than the standard ReLU.

We build a fully-connected NN with four hidden layers. The dimension of the hidden layers are (2048, 1024, 512, 256); see Table I for details. Moreover, a dropout layer with 0.3 dropout rate is introduced. As for the initialization of each layers, we utilize the Xavier uniform for the weights and the Normal distribution for the bias. Adam optimizer with initial 0.0001 cosine learning rate and $l_{2}$ regularization coefficient $5\times 10^{-9}$ is used.

As noted above, the phonon configurations are sampled from a supercell consisting of $6\times 6\times 6$ primitive unit cells. Here a primitive unit cell corresponds to a parallelepiped formed from primitive vectors $\mathbf{a}_{1}=(0,a/2,a/2)$, $\mathbf{a}_{2}=(a/2,0,a/2)$, and $\mathbf{a}_{3}=(a/2,a/2,0)$, where $a$ is the length of a conventional cubic unit cell. The skewed primitive unit cell, however, does not preserve the cubic symmetry of the silicon crystal. To facilitate the incorporation of the point-group symmetry, here we instead use an input block comprised of $3\times 3\times 3$ cubic cells with a total of 279 atoms.

To further improve the statistics of the ML predictions, we apply the same ML model to multiple overlapping cubic blocks of the supercell. Specifically, we choose sites with coordinate $(i,j,k)a$ and $(i+1/2,j+1/2,k+1/2)a$ as the center, where $i,j,k$ are integers ranging from 0 to 2. For these 54 sites, we consider their surrounding sites within a block of 3 $\times$ 3 $\times$ 3 cubic cells which includes 279 sites. Since our descriptors keep the same dimension as the configuration, the number of the input features is $838=279\times 3+1$, where first 837 features are the descriptors of 279 sites configuration and the rest feature is the temperature. The band gap energy corresponding to the phonon configuration $\mathbf{U}$ of the supercell is obtained by averaging ML predictions from these 54 blocks. Accordingly, the loss function used for training the NN is given by the following mean square error (MSE)

where $\hat{E}_{g}^{(k)}$ denotes the ML predicted band gap energy for the $k$-th cubic block, $M=54$ is the number of blocks, and $\langle\cdots\rangle$ indicates averaging over the training dataset. We have generated 103, 108, 113, and 204 distinct phonon configurations for four different temperatures at $T=0,100,200$, and 300K, respectively, and performed fully self-consistent DFT calculations on these configurations to estimate the band gap $E^{\rm DFT}_{g}$. For each temperature, 80 randomly chosen configurations and the corresponding band gap energies are used to train the ML model, while the rest of the configurations are kept as the validation and test dataset. After 600 epochs of training process with each batch including only one configuration, the loss function value is converged to around 0.0024297 for the training dataset and 0.00410986 for the validation dataset. The whole process takes 540 seconds on NVIDIA GPU A100 workstation.

FIG. 5 shows the parity plot of ML predicted band gap corrections $\Delta_{g}$ versus those from DFT calculations for four temperatures used in NN training. As expected, the phonon-induced correction to the band gap energy is enhanced at higher temperatures due to the increased number of thermal phonons. While an overall agreement was obtained between the ML predicted values and DFT results, the prediction accuracy varies with the temperatures. For a more detailed comparison, histogram pots of the prediction error $\delta=\Delta E_{g}^{\rm ML}-\Delta E_{g}^{\rm DFT}$ are shown in FIG. 6(a)–(d) for four different temperatures. The error is larger for higher temperature because the width of the Gaussian distribution used in MC sampling increases with temperature. Yet, it is remarkable that the error from NN trained on $\sim$ 100 data points is merely of the order of 10 meV, thus demonstrating the robustness of our ML model.

After NN training and optimization, we used the trained NN to predict band gap for a large number of MC generated supercell configurations at the four temperature values at which DFT calculations were performed to train the ML model. FIG. 7 shows the temperature dependence of the phonon-induced band gap correction of Si using DFT and NN prediction. For comparison, we also show the experimentally measured band gap correction. It should be noted that, although there are noticeable discrepancies between experiment and computational results for higher temperature region ($T\gtrsim 150$ K), the ML predictions agree quite well with the DFT calculations. Importantly, by using ML-predicted data set, we were able to reduce the errorbars at $T=0,100,200$ and 300 K significantly compared to the DFT calculation.

It is remarkable that both DFT calculation and ML prediction correctly capture the temperature dependence of band gap, especially in the low temperature regime. The discrepancy at higher temperature is likely due to the anharmonic and lattice thermal expansion effects which are not incorporated in our calculations. Inclusion of these effects in ab initio calculations and application of more accurate exchange correlation functionals could yield better comparison with experiments[105, 106]. As discussed above, by training a ML model using more accurate ab initio data, our proposed framework can also incorporate such many-body effects into the electron-phonon couplings and the phonon-induced temperature dependence.

Moreover, since temperature serving as a conditional control is another input to our NN (see FIG. 1), we also apply the trained ML model to predict the band gaps at intermediate temperature values (see FIG. 7). The inclusion of temperature as a control parameter to the ML model, which was trained by datasets from multiple temperatures, could impose further constraints by enforcing a consistency between predictions at different temperatures. We have also trained ML models only trained by dataset from a particular temperature and obtained better prediction accuracies at that temperature. However, such models not only lack the generality for other temperatures, but also likely suffer from over-fitting. On the other hand, the ML models with additional conditioning from the temperature is shown to produce consistent trend of temperature dependent band gap energy. The prediction accuracies at the intermediate temperatures are of the same order as those used for the training. We note that the overall prediction accuracy can be improved by increasing the number of training data points or the scale of the neural network. Our results show that, by leveraging the power of transfer learning, band gap corrections at intermediate temperatures (where training data are not available) can also be accurately predicted by our trained ML model without incurring any significant computational cost. The only cost is in generating supercell configurations which is nominal.