A Compact Tensor Regression Model for AcceleratingFirst-Principle Molecular Dynamics Simulation

ISSP, Ozaki Group
D1, Li Hengyu

Li Hengyu

On-the-fly AIMD

2023/11/8

1Outline• Background• Purpose of study• Method• Benchmark result• Conclusion• Future Prospects
Li Hengyu

On-the-fly AIMD

2023/11/8

2

BackgroundThe desired system size for firstprinciples

molecular(FPMD)

simulation

dynamics

islargelyincreased.

Machine learning techniques has been
proven to fit potential energy surface1.
For better performance, researchers
seek to develop more complex model,like deep neural network.
1 K. T. Butler et al., Nature 559, 547 (2018).

Li Hengyu

On-the-fly AIMD

2023/11/8

3

Overview of on-the-fly fitting (active learning)

Model is not fixed by pre-train, also updated along MD simulation.
Li Hengyu

On-the-fly AIMD

2023/11/8

4

Purpose of study
Main purpose
We try to fit the local PES for each atom with a compact atomic
decomposed model and on-the-fly update along MD simulation.
Highlight features• Fitting by solving linear equations problem (numerical robustness
and low retrain cost)• Independent model for each atom (O(N) time complexity and ability
for complex chemical environment)• Training data can be iteratively accumulated (Cold start and onlinetraining)
Li Hengyu

On-the-fly AIMD

2023/11/8

5

Tensor representation of atomic energy
Total energy is represented by summationof atomic energy:
model𝐸tot= ෍ 𝐸𝑖model𝑖

Them we consider the atomic energy foratom 𝑖 :𝑛 − 1 tiems

𝐸𝑖model = ൻ𝑉𝑖𝐴(𝑛) |𝑉𝑖 ۧ |𝑉𝑖 ۧ … |𝑉𝑖 ۧ ,𝐸 = ෍ 𝜔𝑖 ෍ 𝜔𝑗 𝑓𝑗 ෍ 𝑓𝑘 𝜔𝑘 𝑥𝑘𝑖(𝑛)𝐴(𝑛) |𝑉𝑖 ۧ = ෍ 𝐴𝑖,𝑗,𝑘…,𝑛 𝑣𝑗 = 𝐴(𝑛−1) ,𝑗

Where 𝐴(𝑛) is a n-rank tensor, |𝑉𝑖 ۧ isdescriptor vector. For 1st and 2nd rank case:
𝐸𝑖 𝐴 1
𝐸𝑖 𝐴 2𝑗,𝑘With Taylor expansion of 𝑓:
1𝑓𝑘 𝜔𝑘 𝑥𝑘 = 𝑓 𝑥0 + 𝑓 ′ 𝑥0 𝜔𝑘 𝑥𝑘 − 𝑥0 +2!
′2𝑓 ′ 𝑥0 𝜔𝑘 𝑥𝑘 − 𝑥0 + ⋯So, we have:= 𝐴1 𝑣1 + 𝐴2 𝑣2 + ⋯ + 𝐴𝑚 𝑣𝑚 ,= 𝐸𝑖 𝐴 1 + 𝐴22 𝑣1 𝑣1 + 𝐴23 𝑣1 𝑣2 + ⋯𝐸 = 𝑏0 + ෍ 𝑏𝑖 𝑥𝑖 + ෍ 𝑏𝑖𝑗 𝑥𝑖 𝑥𝑗 + ⋯𝑖𝑖,𝑗+ ෍ 𝑏𝑖𝑗…𝑛 𝑥𝑖 𝑥𝑗 … 𝑥𝑛𝑖,𝑗,…,𝑛Li Hengyu

On-the-fly AIMD

2023/11/86

Descriptor vector
We consider the atomic energy is contributed by 2body and 3-body interactions in a finite range.
𝑎𝑛𝑔|𝑉ۧ = (1, 𝑉1𝑟𝑎𝑑 , … , 𝑉𝑝𝑟𝑎𝑑 , 𝑉1
𝑉𝑝𝑟𝑎𝑑 =𝑎𝑛𝑔

, … , 𝑉𝑝)ℎ 𝑍𝑗 𝐺𝑝𝑟𝑎𝑑 (𝑟𝑖𝑗 )෍𝑗,𝑗∈𝑛𝑒𝑖𝑔ℎ𝑏𝑜𝑟

𝑎𝑛𝑔

𝑉𝑝=෍𝑎𝑛𝑔ℎ 𝑍𝑗 , 𝑍𝑘 𝐺𝑝(𝜃𝑖𝑗𝑘 )𝑗,𝑘;𝑗,𝑘∈𝑛𝑒𝑖𝑔ℎ𝑏𝑜𝑟

We use series of symmetry function 𝐺𝑝 to sample
all many body pairs inside cutoff range. 𝑍𝑗 is thenormalized atomic number for distinguish elementtype.
Li Hengyu

On-the-fly AIMD

2023/11/8

7

Orthogonal polynomial descriptor
Orthogonal polynomials with 𝒙 − 𝟏 𝒑 basis
The 𝑂𝑝 (𝑥) = Orth[ 𝑥 − 1 𝑝 ], 𝑝 ≥ 3 series polynomial is forced to be orthogonal.
𝑂𝑝 (𝑥)

Derivative of 𝑂𝑝 (𝑥)𝑂𝑝 , 𝑂′𝑝 , 𝑂′′𝑝 should converge to 0 for smooth energy change respect toneighboring atom change.
Li Hengyu

On-the-fly AIMD

2023/11/8

8

Loss function
The fitting process is to determine the tensor element in 𝐴𝑛𝑖 by minimizing theloss function:
𝐿 = ෍ 𝐸𝑁model − 𝐸𝑁DFT𝑁2modelDFT+ ෍ ෍ 𝐹𝑁,𝑎𝑥𝑖𝑠− 𝐹𝑁,𝑎𝑥𝑖𝑠
𝑁 𝑎𝑥𝑖𝑠2+ 𝜆2 ෍ 𝐴𝑛𝑖
𝑁Thus, to minimize the loss function we can let derivative to 0. (2nd rank tensorcase):model 𝑖model 𝑖𝜕𝐹𝑁,𝑎𝑥𝑖𝑠
𝜕𝐿modelDFTmodel
DFT 𝜕𝐸𝑁= ෍ 𝐸𝑁− 𝐸𝑁+ ෍ ෍ 𝐹𝑁,𝑎𝑥𝑖𝑠 − 𝐹𝑁,𝑎𝑥𝑖𝑠+ 2𝜆2 ෍ 𝐴𝑚′𝑛′
𝜕𝐴𝑚′𝑛′
𝜕𝑎𝑖,𝑗,𝑝𝜕𝐴𝑚′𝑛′
𝑁𝑁 𝑎𝑥𝑖𝑠

𝑁Then the fitting process could be done alone the MD simulation and solving thelinear equations problem:
𝜕𝐿𝐴𝑚 ′ 𝑛 ′ 𝑇 = 0
𝜕𝐴𝑚′ 𝑛′Li Hengyu

On-the-fly AIMD

2023/11/8

9Implementation
Decouple atomic force
modelmodelmodel
𝜕𝐸
𝜕𝐸
𝜕𝐸𝑗
𝑡𝑜𝑡𝑖
model
𝐹𝑖,𝑎𝑥𝑖𝑠==+෍= 𝐹self (𝐴𝑖,𝑁 ) + 𝐹neighbor (𝐴𝑗,𝑁 )𝜕𝑅𝑖,𝑎𝑥𝑖𝑠
𝜕𝑅𝑖,𝑎𝑥𝑖𝑠
𝜕𝑅𝑖,𝑎𝑥𝑖𝑠𝑗≠𝑖𝐹neighbor (𝐴𝑗,𝑁 ) contains parameters from neighboring atom which can not

solved independently. So, we use the parameters 𝐴𝑗,𝑁−1 to approximate
𝐹neighbor (𝐴𝑗,𝑁 ) and use autoregressive model to minimize residuals.

𝐹diff 𝐴𝑗,𝑁 , 𝐴𝑗,𝑁−1 = 𝛼1 𝐹diff 𝐴𝑗,𝑁−1 , 𝐴𝑗,𝑁−2 + 𝛼2 𝐹diff 𝐴𝑗,𝑁−2 , 𝐴𝑗,𝑁−3 + ⋯Li Hengyu

On-the-fly AIMD

2023/11/8

10Implementation
Online training
Cause the dimension of parameter tensor 𝐴𝑛𝑖 is fixed, thus we caniteratively accumulate the training data by:
𝜕𝐿
new= 𝐴𝑛𝑚′ 𝑛′ = 𝜔 ∗ 𝐴𝑛−1+𝐴′′𝑚 𝑛𝑚′ 𝑛 ′𝜕𝐴𝑚′ 𝑛′𝜔 is a forgetting ratio which is little smaller than 1 to discard history step
information that is too far away from current configuration.
Force residuals given by autoregressivemodelreflect

long-timescaleinteraction in MD trajectory. The idea
is same in recurrent neural networks(RNN) and graph neural networks(GNN).
Li Hengyu

On-the-fly AIMD

2023/11/8

11Workflow with OpenMX
Scheme of ML code and OpenMX1
Fully parallelized for each atom
Necessary data from OpenMX• Gxyzatom position• FNANneighboring atoms number• NATNneighboring atoms serial number• DecE
decomposed energy[1] OpenMX: https://www.openmx-square.org/
Li Hengyu

On-the-fly AIMD

2023/11/8

12

Benchmark calculation of various systems
Isolated System

Energy MAE

Force MAE

Methene (CH4)

0.08 meV/atom

7 meV/Å

Nitrobenzene (C6H5NO2)

0.73 meV/atom

68 meV/Å

Sialic Acid (C11H19NO9)

1.07 meV/atom

114 meV/Å

Adenosine Triphosphate(C10H16N5O13P3)

1.32 meV/atom

145 meV/Å

Bulk System

Energy MAE

Force MAE

2D Graphene

0.23 meV/atom

27 meV/Å

Aluminum0.47 meV/atom

76 meV/Å

Amorphous Silica (SiO2)

0.92 meV/atom

127 meV/Å

The typical accuracy range for ML potential is ~2 meV/atom and ~200 meV/Å.
Li Hengyu

On-the-fly AIMD

2023/11/813

Radial Distribution Function

Twopeaksareprecisely reproduced
byourtensorregression model.

RDF of amorphous silicon Si-Si distance under 300K.
(2000 fs in total, 400 fs for initial training, update every 10 steps).
Li Hengyu

On-the-fly AIMD

2023/11/8

14

Radial Distribution Function

RDF of amorphous silica under 300K.
(2000 fs in total, 500 fs for initial training, update every 10 steps).
Li Hengyu

On-the-fly AIMD

2023/11/8

15

Conclusion• We developed effective model of fitting local atomic energy
surface, the accuracy is comparable to other ML potentials. Also, it
able to reproduce the long-time scale statistical property.

• The accelerate ratio is around 10 by simple on-the-fly algorithm butlower than existing active learning method (～103).
• Our model shows enlightenment in the development of future machinelearning potentials. Compact models which based on the physicalsystem also possible for ML potentials compare with the more
complex and data-driven large models.

Li Hengyu

On-the-fly AIMD

2023/11/8

16

Future Prospects• Design on-the-fly algorithm based on our model, for example theanalysis of changes in atomic chemical environment.
• Application for electrode surface, because there are many possible

surface structures and a long AIMD time is required for screening.
• Extending the model prediction capabilities to more complex highentropy alloy surfaces that always face reconfiguration duringdynamics.
• Our final target system is electrochemical interface systems which
are still challenge for existing ML potentials and DFT packages .

Li Hengyu

On-the-fly AIMD

2023/11/8

17

Q&A

Thank you for your listening.

Li Hengyu

On-the-fly AIMD

2023/11/8

18

Density functional theory
Many-body Schrödinger equation
෡ 𝒓, 𝑹 = 𝐸Ψ 𝒓, 𝑹𝐻Ψ

With Born-Oppenheimer approximation, the electronic motion andthe nuclear motion in molecules can be separated.

Ψ 𝒓, 𝑹 = 𝜙𝑒𝑙𝑒 𝒓 𝜒𝑛𝑢𝑐 𝑹
Electron Schrödinger equation in Hartree unit

111𝑍𝐼
2− ෍ ∇𝑖 + ෍−෍𝜙𝑒𝑙𝑒 𝒓 = 𝐸𝜙𝑒𝑙𝑒 𝒓
22𝒓𝒊 − 𝑹𝑰𝒓𝒊 − 𝒓𝒋𝑖

Li Hengyu

𝑖≠𝑗𝑖,𝐼On-the-fly AIMD

2023/11/8

19

Density functional theory
Hohenberg-Kohn theorem I
For any system of interacting particles in an external potential 𝑉𝑒𝑥𝑡 (𝒓),the potential 𝑉𝑒𝑥𝑡 (𝒓) is determined uniquely by the ground-state particle
density 𝑛0 (𝒓).
Hohenberg-Kohn theorem IIA universal functional for the energy 𝐸[𝑛] can be defined in terms ofthe density. The exact ground state is the global minimum value of thisfunctional.
With Hohenberg-Kohn theorem, the solution of the system groundstate can be shifted from the solution of the wave function to theelectron density. The system degrees of freedom are reduced from
3N to 3.

Li Hengyu

On-the-fly AIMD

2023/11/820

Density functional theory
Kohn-Sham DFT
In Kohn-Sham scheme constructed an auxiliary system which the atoms are all
independent of each other, and in turn we can solve the single-electron
Schrödinger equation.
∇21𝑛 𝒓 𝑛 𝒓′′
𝐸 𝑛 = − + න 𝑑𝒓 𝑛 𝒓 𝑉𝑒𝑥𝑡 𝒓 + න 𝑑𝒓𝑑𝒓+ 𝐸𝑥𝑐 (𝑛)22𝒓 − 𝒓′With Hohenberg-Kohn theorem II, the electron density that minimizes 𝐸 𝑛 is
the real electron density of the system. By the Lagrange multiplier method wecan obtain the KS equation.
∇2
𝛿𝐸𝑥𝑐 𝑛− + 𝑉𝑒𝑥𝑡 𝒓 + 𝑉𝐻 𝒓 +2𝛿𝑛𝜓𝑖 𝒓 = 𝜀𝑖 𝜓𝑖 (𝒓)After determining the form of 𝐸𝑥𝑐 𝑛 , KS equation can be solved by selfconsistent field method.
Li Hengyu

On-the-fly AIMD

2023/11/821

Linear scaling DFTThe core idea of the O(N) method is based on the assumption that the electron
density 𝑛 is given by the summation over each site 𝑖.

The density matrix 𝜌 is calculated by solving following linear equations:where, 𝐻(𝑖) is the local Hamiltonian matrix, 𝑆 (𝑖) is the overlap matrix. By truncating
system the time complexity is reduced:
𝑂 𝑁 3 → 𝑁 × 𝑂(𝑁𝑐𝑜𝑛𝑠𝑡𝑎𝑛𝑡 )Further, the 𝑁𝑐𝑜𝑛𝑠𝑡𝑎𝑛𝑡 can be reduced by introducing Krylov subspace method, The
Krylov subspace method used in OpenMX is a combination of DC and recursive
methods to improve computational efficiency and robustness.
Li Hengyu

On-the-fly AIMD

2023/11/8

22

Energy decomposition in OpenMX
Because the pseudo-atomic orbitals 𝜒 for pseudopotential in OpenMX is in
finite range, the total energy can be decomposed into contributions from eachatomic site 𝑖 and localized orbital 𝛼:

Li Hengyu

On-the-fly AIMD

2023/11/8

23

Energy decomposition in OpenMX
The total energy in OpenMX is organized by:

Kinetic energy

withLi Hengyu

On-the-fly AIMD

2023/11/8

24

Energy decomposition in OpenMX
Electron-core Coulomb energy (Non-local part)

Neutral atom energy

Li Hengyu

On-the-fly AIMD

2023/11/825

Energy decomposition in OpenMX
Screened core correction energy

Electron-electron Coulomb energy

Li Hengyu

On-the-fly AIMD

2023/11/826

Energy decomposition in OpenMX
Exchange correlation energy

Li Hengyu

On-the-fly AIMD

2023/11/827