Analysis review of replacing traditional DFT model and DeepE3 model based on deep learning model

The open source China community team made its first live broadcast, telling the story behind the open source China community in the name of sharing."

Author: Yu Fan

background

Deep learning has made significant progress in the field of computational quantum chemistry in recent years. Existing deep learning methods have proven their efficiency and expressiveness in solving various challenging quantum mechanical simulation tasks. With the continuous increase of today's computing power, Increasingly, the integration of deep learning and first principles is getting deeper and deeper. Starting from the basic laws of physics, equivariant neural networks are combined at the scale of atoms and electrons to predict values in the quantum field, especially in the field of DFT. Some deep learning models have been born, such as DeepE3 and QhNet. These networks are based on equivariant networks to predict DFT Hamiltonians. This article will introduce the origin of density functional theory (DFT), the principle of equivariant neural network and a more general equivariant neural network E3nn. Finally, it will outline the DeepHE3 model, a model proposed by the Tsinghua University team based on E3 and others. A model that changes the network to predict the DFT Hamiltonian [1].

**1.** Density functional theory

Density functional theory (DFT) is a quantum mechanical method that studies the electronic structure of multi-electron systems. Multi-electron systems are represented by the density functional of electrons. Density functional theory is widely used in physics and chemistry, especially to study the properties of molecules and condensed matter.

Before density functional theory, the system wave function was calculated by solving the Schrödinger equation. As the basic equation of quantum mechanics, the Schrödinger equation has the following form:

Equation 1. Time-dependent Schrödinger equation

Equation 2. Schrödinger equation without time

Where Ψ is the wave function describing the state of microscopic particles, E is kinetic energy, and H is the Hamiltonian. The Hamiltonian describes the evolution of the quantum system and can be expressed as the sum of the kinetic energy and potential energy of the particles.

For a multi-particle system composed of N electrons and M atoms, the stationary Schrödinger equation corresponding to the Hamiltonian can generally be written as:

Equation 3. Schrödinger equation

This wave function has 3*(M + N) variables, so it is very difficult to solve.

Density functional theory replaces the wave function with electron density, which is a method of studying the electronic structure of multi-electron systems through the density of electrons. Among them, density is a function of three-dimensional coordinates to describe the density of electrons, and functional function refers to a function that maps density into energy E. DFT takes electron density as a basic quantity and expresses energy as a functional function of electron density. Since the electron density is only a function of spatial coordinates, the dimension of the multi-electron system is directly reduced to 3, which simplifies the solution process of the Schrödinger equation.

1965年，美国加利福尼亚大学圣迭戈分校的沃尔特·科恩与沈吕九提出了Kohn–Sham equation。作为密度泛函理论里最普遍的通用的代表，KS方程把把相互作用的多粒子体系转换成一个无相互作用的单粒子体系，把电子之间的相互作用归结到未知的交换关联势中：

Equation 4. KS equation[4]

The exchange correlation potential term refers to the energy difference between the interacting multi-particle system and the non-interacting multi-particle system. And the exact functional form of this energy term is unknown, and it can only be expressed as an approximate functional of the electron density, such as the local density approximation (LDA). The electron density is determined by the solution of the above single-electron wave function equation, so the specific form of this equation depends on its solution itself, which needs to be solved by self-consistent iteration.

Figure 1. Approximate calculation process[4]

Its computational complexity is O(N^3), N is the number of electrons, and it is still difficult to solve large systems.

**2, ** Equivariant network

When using neural networks to calculate some quantum properties, it is usually necessary to take into account the transformation of these properties by the rotation of the particles. Some scalar values, such as energy values, distances between particles, etc., are not affected by the rotation of the particles. For some multi-dimensional vector characteristics, such as force, Hamiltonian, etc., the values need to be changed accordingly according to the rotation of the particles, and this change needs to be consistent from the beginning to the end of the network. Therefore, equivariant networks are used for most first-principles models.

2.1 What is equivariance?

Take a function as an example. If the transformation you apply to its input is also reflected on the output, then the function is equivariant. f(g(x)) = g(f(x)).

2.2 What is an equilateral network?

(1) The transformation of network input needs to be symmetrically mapped to the internal and output results.

(2) For example, if there is a three-dimensional atomic structure, we need to use a neural network to predict its various properties, such as potential energy, number of electrons, and direction of force. If we rotate the atomic structure, its potential energy and number of electrons should remain the same because they are scalars; and their force direction results should change accordingly because they are multidimensional vectors. This symmetrical mapping needs to be reflected in network intermediates and outcomes. Therefore, ensuring this mapping relationship requires an equivariant network.

2.3 Why should we achieve equivariance?

For a two-dimensional image, such as an animal picture, in order to make the model symmetrical, data enhancement is usually performed. The animal picture is rotated at 10 different angles and fed into the neural network for training to make the network train differently. A picture of an angle. But for a three-dimensional model, such as an atomic structure, this kind of enhancement is not realistic. Usually, if a simple three-dimensional model is to be augmented with data, it needs to have at least 500 rotations of data augmentation to adequately cover an atomic structure in different Characteristics of angles. If you use an equivariant network, you only need to pass in a structure.

Figure 2. Two-dimensional animal image

Figure 3. Three-dimensional model diagram[5]

**3, ** E3nn: Space transformation neural network based on three-dimensional Euclidean space

E3: The spatial transformation group of three-dimensional Euclidean space, which can be decomposed into translation, rotation (SO(3) special orthogonal group) and inversion. The equivariance of translation is already satisfied in convolution, so we focus on rotation And inversion -> SO(3)×Z2=O(3)

The main concepts in E3NN:

1. Group: Transformation types of space, such as rotation and inversion.

2. Representation: Defines the representation of which space transformation group (Group) the vector space belongs to.

3. Irreducible representation (irreps): Irreducible representation is equivalent to an irreducible Representation. Each irreps can be marked with (l,p), l=0,1,2,... is the order, p=e,o is the parity, and the dimension of the l-order irreducible representation is 2l+1. For example, a vector has an order of 1 (representing a dimension of 3) and odd evenness, so it can be abbreviated as 1o.

Figure 4. Introduction to irreps

For example, in the figure below, a1 – a9 represent 9 real numbers respectively. If a1 – a3 are regarded as 3 scalars respectively, a4 – a6 are regarded as a vector, and a7 – a9 are regarded as another vector, then the irreps of this matrix is "3 × 0e + 2 × 1o" to express. When we need to rotate this matrix, we need to perform different transformations according to the corresponding Group in irreps. For the three scalars a1-a3, the rotation does not affect their values, so they are multiplied by 1. And for a4 – The two vectors a6 and a7-a9 need to be multiplied by the corresponding rotation matrix to obtain the corresponding values.

Figure 5. Example of rotation matrix[5]

The following explains how to decompose two multiplied Irreps (how to decompose a tensor product)

Equation 5. Decomposition of Tensor product

For example: 2 ⊗ 1 = 1 ⊕ 2 ⊕ 3, 2 ⊗ 2 = 0 ⊕ 1 ⊕ 2 ⊕ 3. As can be seen from this example, the reason why e3nn can maintain equivariance is that it determines the irreps of network input, output and intermediate results in advance. To ensure that the transformation of Group is carried out according to the corresponding irreducible representation, thereby preventing confusion.

**4、** DeephE3

A general E{3} equivariant deep learning framework that predicts DFT Hamiltonians from atomic structures {R} with spin orbits through neural networks. DeephE3 can learn to predict electronic predictions of larger material systems by training the DFT results of small material systems. This method is applicable to various material systems, such as general magic-angle twisted bilayer graphene or twisted van der Waals materials, and is several orders of magnitude cheaper than direct DFT calculations.

The figure below shows the architecture of the entire network. Among them, {Zi} represents the atomic number, | rij | represents the distance between atoms, and is used to construct a vector with order equal to 0. ^rij represents the relative position between atoms, and the vector is used to construct a vector with order equal to 1, 2. {Zi} is passed in the elemental embedding module (Elemental embedding) as the initial vertex; | rij | is passed in the Gaussian extension (Gaussian Bias) as the edge feature; ^rij is passed in the spherical harmonic function as the relative position between atoms Perform mapping to generate Y(^rij). The spherical harmonic function Y^l maps a 3-dimensional vector into a 2l+1-dimensional vector, which represents the coefficient when the input vector is decomposed into 2l+1 basic spherical harmonics.

Figure 6. Overall structure of DeephE3[1]

The generated vertex and edge features are updated L times through vertex update and edge update update blocks. The update block encodes inter-atom distance and relative unknown information through equivariant convolution. The "⋅" symbol represents channel multiplication, and || represents vector concatenation.

Then use the message-passing method to update the vectors of edges and vertices by obtaining information about adjacent edges.

The final edge vectors are passed to the Wigner-Eckart layer to display the DFT Hamiltonian. If spin-orbit coupling (SOC) is ignored, the output vector of the neural network is converted into a Hamiltonian through the Wigner Eckart layer using the rule 1 ⊕ 2 ⊕ 3 = 1 ⊗ 2. If SOC is included, the output consists of two sets of real vectors that combine to form a complex-valued vector. These vectors are converted into spin-orbit DFT Hamiltonians using another rule: (1 ⊕ 2 ⊕ 3) ⊕ (0 ⊕ 1 ⊕ 2) ⊕ (1 ⊕ 2 ⊕ 3) ⊕ (2 ⊕ 3 ⊕ 4) = (1 ⊕ 1/2) ⊕ (2 ⊕ 1/2 ). ⊕ refers to tensor add, ⊗ refers to tensor product.

Figure 7. Wigner-Eckart layer[1]

**5, ** Summary

This article introduces the application of deep learning from first principles, as well as the related physical background. With the deeper combination of deep learning and equivariant networks, more and more quantum properties that are difficult to calculate with traditional methods can be predicted through neural networks, thus better helping scientific research institutions to research new materials, build material databases, etc., to achieve More application innovations.

references

[1]https://www.nature.com/articles/s41467-023-38468-8

[2]https://www.nature.com/articles/s43588-022-00265-6

[3]https://arxiv.org/abs/2207.09453

[4]https://www.bilibili.com/video/BV1vU4y1f7gQ/?spm_id_from=333.337.search-card.all.click

[5]https://www.youtube.com/watch?v=9rS8gtey_Ic

Analysis review of replacing traditional DFT model and DeepE3 model based on deep learning model

references

Guess you like