A graphical transformer neural network force field for predicting atomic forces and energies in MD simulations

Abstract

A graph transformer neural network force field for predicting atomic forces and energies in MD simulations. The simulation involves converting a molecular dynamics snapshot of elements within a multi-element system into a graph with atoms as nodes of the graph; defining a matrix such that each column of the matrix represents a node in the graph; defining a distance matrix from a set of relative positions of each atom; iterating through the GTFF using an attention mechanism, operating on the matrix, and augmenting by merging the distance matrix to pass the hidden state from the current layer of the GTFF to the next layer of the GTFF; performing combinations over the columns of the matrix to produce scalar molecular energies; making a backward pass through the GTFF, iteratively computing derivatives at each of the layers of the GTFF to compute predictions of the forces acting on each atom; and returning predictions of the forces acting on each atom.

Description

Technical Field

[0001] The present invention generally relates to a Graph Transformer Neural Network Force Field (GTFF) computational algorithm for direct prediction of atomic forces in molecular dynamics computer simulations in material systems such as electrochemical and water filtration devices. Background Art

[0002] Molecular dynamics is a computational materials science method that is used to simulate the motion of atoms in a material system under realistic operating pressure and temperature conditions. There are methods for calculating the fundamental atomic forces used in simulating the motion of atoms. One method is the ab-initio quantum mechanics approach. This method is very accurate, but also very expensive because of the enormous amount of computing resources required to apply the method. Although there are other methods that consume less computing resources, these other methods do not provide the same high accuracy. Summary of the invention

[0003] According to one or more illustrative examples, a computational method for simulating the motion of elements within a multi-element system using a graph transformer neural network (GTFF) includes converting a molecular dynamics snapshot of elements within the multi-element system into a graph with atoms as graph nodes; defining a matrix so that each column of the matrix represents a node in the graph; defining a distance matrix based on a set of relative positions of each atom; iterating through the GTFF using an attention mechanism, operating on the matrix, and augmenting by merging the distance matrix to pass hidden states from a current layer of the GTFF to a next layer of the GTFF; performing combinations over the columns of the matrix to produce scalar molecular energies; making a backward pass through the GTFF, iteratively calculating derivatives at each of the layers of the GTFF to calculate predictions of forces acting on each atom; and returning predictions of forces acting on each atom.

[0004] According to one or more illustrative examples, a computing system for simulating the motion of elements within a multi-element system using a graph transformer neural network (GTFF) includes a memory storing instructions of a GTFF algorithm of a molecular dynamics (MD) software; and a processor programmed to execute the instructions to perform operations, the operations comprising converting a molecular dynamics snapshot of elements within the multi-element system into a graph with atoms as graph nodes; defining a matrix such that each column of the matrix represents a node in the graph; defining a distance matrix based on a set of relative positions of each atom; iterating through the GTFF using an attention mechanism, operating on the matrix, and augmenting by merging the distance matrix to pass hidden states from a current layer of the GTFF to a next layer of the GTFF; performing combinations over the columns of the matrix to produce scalar molecular energies; making a backward pass through the GTFF, iteratively calculating derivatives at each of the layers of the GTFF to calculate predictions of forces acting on each atom; and returning predictions of forces acting on each atom.

[0005] According to one or more illustrative examples, a non-transitory computer-readable medium includes instructions for simulating the motion of elements within a multi-element system using a graph transformer neural network (GTFF), which, when executed by a processor, causes the processor to: convert a molecular dynamics snapshot of elements within the multi-element system into a graph with atoms as graph nodes; define a matrix such that each column of the matrix represents a node in the graph; define a distance matrix based on a set of relative positions of each atom; iterate through the GTFF using an attention mechanism, operate on the matrix, and augment by merging the distance matrix to pass hidden states from a current layer of the GTFF to a next layer of the GTFF; perform combinations over the columns of the matrix to produce scalar molecular energies; make a backward pass through the GTFF, iteratively calculating derivatives at each of the layers of the GTFF to calculate predictions of forces acting on each atom; and return predictions of forces acting on each atom. BRIEF DESCRIPTION OF THE DRAWINGS

[0006] Figure 1 is an illustration of an example of using a neural network to perform a molecular dynamics simulation;

[0007] Figure 2 is an example illustration of the Graph Transformer Neural Force Field (GTFF) method for MD simulations; and

[0008] Figure 3 It can be used to implement the GTFF algorithm in one or more embodiments, such as Figure 2 Schematic illustration of the computational platform for the GTFF algorithm. DETAILED DESCRIPTION

[0009] As required, detailed embodiments of the present invention are disclosed herein; however, it will be understood that the disclosed embodiments are merely examples of the present invention, which may be embodied in various forms and alternative forms. The figures are not necessarily to scale; some features may be enlarged or minimized to show details of particular components. Therefore, the specific structural and functional details disclosed herein should not be interpreted as limiting, but merely as a representative basis for teaching those skilled in the art to employ the present invention in various ways.

[0010] The term "substantially" may be used herein to describe embodiments disclosed or claimed. The term "substantially" may modify a value or relative property disclosed or claimed in the present disclosure. In this case, "substantially" may mean that the value or relative property it modifies is within ±0%, 0.1%, 0.5%, 1%, 2%, 3%, 4%, 5%, or 10% of that value or relative property.

[0011] Molecular dynamics (MD) methods are useful for studying physical phenomena such as, but not limited to, ion transport, chemical reactions, and bulk and surface degradation of materials in material systems such as devices or functional materials. Non-limiting examples of such material systems include fuel cells, surface coatings, batteries, water desalination, and water filtration. Methods exist for calculating the fundamental atomic forces used in simulating atomic motion. Ab initio quantum mechanics methods are very accurate, but also very expensive, as a huge amount of computing resources are required to apply the method.

[0012] Neural networks have been used to fit and predict quantum mechanical energies. These methods are known as Neural Network Force Fields (NNFF). The derivatives of energy with respect to atomic positions and forces are predicted using quantum mechanical energies. However, these methods are also computationally extensive. In view of the above, what is needed is a computational method for calculating atomic forces that provides a sufficient level of accuracy while consuming a reasonable amount of computational resources.

[0013] Molecular dynamics uses atomic positions (and possible charge, bond or other structural information) to calculate the interatomic forces of each atom, which are used to modify the velocity of the atom in the simulation. The resulting atomic trajectories are used to describe physical phenomena such as, but not limited to, ion transport motions in batteries (e.g., Li-ion batteries) and fuel cells (e.g., fuel cell electrolytes), chemical reactions during bulk and surface material degradation, solid-state material phase transitions, molecular binding and protein folding for drug design, bioscience and biochemical design. Depending on the underlying method used to calculate atomic forces, there is a trade-off between the accuracy of the simulation and the size (measured by the number of atoms and the simulation dynamics time). As noted above, an accurate but expensive method uses ab initio quantum mechanics methods known as ab initio molecular dynamics (AIMD).

[0014] Existing work has been done to provide less expensive methods. In JP Mailoa et al. Nat. Mach. Intell. (2019) arXiv :1905.02791. A Fast Neural Network Approach for Direct Covariant Forces Prediction in Complex Multi-Element Extended Systems, a version of fast NNFF is proposed to directly predict atomic force vectors in extended atomic systems by using rotational covariant features. However, feature extraction relies on feature engineering, making it less accurate.

[0015] In T. Xie et al., Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties. Phys. Rev. Lett. 120, 145301 (2018), convolutional graph neural networks are used for automated feature extraction from atomic structures to enable prediction of total system energy. Using this approach, the message passing mechanism between atoms and bonds is controlled by the network designer.

[0016] In Path-Augmented Graph Transformer Network by B. Chen et al. arXiv :1905.12712 (2019), a graph transformer neural network architecture augmented with path (real chemical bond) features is used to predict the total system energy. The sequence and strength of message passing are controlled by an attention mechanism that exploits the path features.

[0017] In SchNet: A continuous-filter convolutional neuralnetwork for modeling quantum interactions by KT Schutt et al. arXiv:1706.08566 (2017), a convolutional graph network is used for automated feature extraction via message passing between atoms without an attention mechanism. The architecture enables prediction of molecular energies and atomic force vectors by taking the derivative of the entire molecular energy network with respect to position, an algorithm similar to the work of JPMailoa and can be used for actual molecular dynamics simulations.

[0018] The Graph Transformer neural network architecture was introduced in P. Velockovic et al. Graph Attention Networks. arxiv:1710.10903 (2018). So far, this architecture has not been adapted to solve molecular dynamics problems.

[0019] The methods described in the first four references above are directed to using deep learning to quickly regress atomic force vectors and / or energies. These bypass expensive quantum mechanical simulations such as density functional theory (DFT) or quantum chemistry. In the context of molecular dynamics (MD) simulations, this means that MD simulations can be performed with an accuracy close to that of ab initio MD (AIMD) at a significantly lower computational cost.

[0020] Figure 1An example diagram 100 of using a neural network to perform MD simulations is illustrated. Given an atomic input structure 102, quantum mechanics (such as expensive DFT) can be used to calculate system energy and atomic force vectors 104. Then, atoms can be moved based on force vectors calculated using Newton's equations of motion (this setup is called molecular dynamics simulation or MD simulation) to study phenomena such as ion transport, chemical reactions, etc. In MD simulations, iterations of the order of hundreds of millions of time steps are typical, making MD simulations based on DFT force calculations very expensive. Alternatively, given a sufficient amount of quantum mechanical atomic structure force vector training data 106, a neural network 108 can be trained to accurately perform the force calculation. This effectively bypasses the quantum mechanical force calculation and enables accurate force calculations at low computational cost, thereby allowing the performance of larger and longer MD simulations.

[0021] In contrast to previous approaches, the disclosed method uses a graph transformer network, rather than a graph neural network for use in a neural network force field approach, for force vector regression. As discussed in detail below, the graph transformer provides higher accuracy than a graph neural network for graph learning problems. In addition, in contrast to previous approaches, the described method uses a slightly different graph structure. Rather than constructing a graph with atoms as nodes and bonds as edges; in this architecture, atoms, bonds, or other molecular objects can be considered nodes in the graph. This provides greater flexibility and allows the incorporation of features that are incompatible with simpler graph structures. The described method can be used to make force vector predictions for use in molecular dynamics simulations, and can further be used to train machine learning systems that can be used for such applications.

[0022] In general, regarding the graph neural network architecture, in a standard feedforward neural network there are several layers where each layer modifies the hidden representation via multiplication with a weight matrix. In a graph neural network, operations are performed on a graph and each node in the graph has such a representation. Assume H is a matrix with each column acting as a node representation. H For input and In the neural network layer with the output graph, the two are related by:

[0023] .

[0024] here, is a learnable weight matrix, A is a fixed mixing matrix that combines information based on the structure of the graph, and is the activation function (e.g. tanh or ReLU ). More intuitively, The weights of the layer are applied to each node and then multiplied by A This information is combined based on the graph structure (for example, if A is the normalized adjacency matrix of the graph, then multiply by A The hidden states of the neighbors of each node in the graph are averaged). There exists a A There are multiple possible choices for , but once chosen it remains fixed throughout training and testing. Adding the layer input again after applying the activation function is a common technique used to improve training speed and convergence.

[0025] In the graph transformer, layers are operated on by:

[0026]

[0027] Among them, the current , , are all learnable weight matrices, and softmax The function is applied to each column independently. The essential difference here is that the expression to replace the fixed A This expression, known as the attention mechanism, involves H and some learnable parameters; thus, the network can adapt the mixing matrix to different types of inputs and learn it during training and to learn the best ways to combine information.

[0028] The attention mechanism in the graph transformer is further augmented by incorporating a distance term:

[0029]

[0030] in D is a distance matrix (that is, Is a node i and j The square of the distance between is a learnable scalar parameter. This has the effect of reducing the strength of the interaction between pairs of nodes as their distance increases, and Learn how much weight to assign to this reduction.

[0031] In the architecture described above, each layer is applied in sequence, resulting in the total network output ,in Each column of represents the hidden state of a specific node. If you are making node-level predictions (i.e., predicting the properties of each node), then this output can be used as the prediction. Conversely, if you are predicting properties of the graph, then The columns of can be combined (typically by summing the columns or by averaging the columns) to obtain a single graph-level output.

[0032] Back to MD simulations, given a snapshot of a molecule, the goal is to predict the i The force vector on It is possible to accomplish this goal using a simple graph neural network training scheme: the input to the network is some information about the molecule encoded as a graph with atoms as nodes, and the output is the predicted force vectors. However, in these predictions, nothing is enforced to conserve energy; therefore, they are not suitable for direct use in molecular dynamics simulations. One can use a neural network force field (NNFF) approach instead, which has better results.

[0033] Figure 2 is an example diagram 200 of a graphical transformer neural network force field (GTFF) method for MD simulations. In such a method, the input to the neural network is the position of each atom , and any other relevant information. The outputs of the final layer are summed to produce the molecular energy E The derivative can then be calculated , of which its Can be used to act on atoms i If the neural network is a set of positions If the derivative is a smooth function of Ensure that energy is conserved.

[0034] The NNFF method combined with the graph neural network architecture can be constructed by combining it with the graph transformer architecture. As shown, GTFF uses a L Once the graph transformer is trained, the algorithm can be executed as follows:

[0035] 1. Let is the position of the atom in the molecule. Organize the molecular information into a graph with atoms as nodes; let as a matrix with each column representing a node in the graph ( May or may not be included in middle). Other information may also be included, such as element type, partial charge, chemical environment, distance to neighbors, angle to neighbors, electronegativity, valence, any other quantum or classical information. It may be encoded in various ways using various bases such as one-hot-encoding, sine, cosine, exponential, etc. , for example as a scalar or as a vector. Let Also used to construct the distance matrix D .

[0036] 2.For i = 0 , …, L-1 :

[0037]

[0038] This is done by passing the graphics transformer layer Pass the hidden state;

[0039] 3. Sum or average over the columns to produce a scalar molecular energy E ;

[0040] 4. Make a backward pass through the network (as shown by the double-line arrows), iteratively computing the derivatives at each layer to i calculate ;and

[0041] 5. Return force prediction .

[0042] GTFF can be trained as a deep neural network. For example, the loss function can be formed as L , to predict the force This is compared to the ground truth on the labeled dataset, and the network weights can then be optimized using gradient descent with respect to this loss (see above). and ). Because gradient descent is performed on gradients, this requires the computation of higher-order derivatives; therefore, each training iteration will take approximately twice as long as for a feed-forward neural network.

[0043] Variations on the GTFF approach are possible. For example, while the graph described above assumes that all nodes are atoms, in some cases it may be desirable to include other molecular objects such as bonds, triplets of atoms, etc. as nodes in the graph. This is useful for incorporating molecular features that cannot be represented by atoms alone (e.g., bond order of bonds, bond angles of triplets, etc.), and can aid in calculating the total energy of the system. If this information is included as nodes, their distances to all other objects in the graph should also be defined, such that D remains well defined. However, the derivatives of the energy with respect to the positions of these objects may not need to be computed since we are only computing the force vectors for the atoms.

[0044] Regarding another possible variation, in step 3 above, the instruction to execute However, it should be noted that in other examples, it can be applied to Another function that is permutation invariant over the columns. As yet another possibility, this combination operation can be implemented as another small neural network.

[0045] Regarding another possible variation, in steps 3 and 4, predict E is assumed to be a single scalar value (and therefore, the output size of the final layer is assumed to be 1). In other examples, E Can be is a vector in , and the force can be taken as the sum of these derivatives: . Because each sum follows the law of conservation of energy, the entire sum will also follow the law of conservation of energy. Multidimensional Energy E Can be compared to scalar E Easier to learn.

[0046] The GTFF algorithm and / or method of one or more embodiments is implemented using a computing platform, such as Figure 3 . The computing platform 300 may include a memory 302, a processor 304, and a non-volatile storage device 306. The processor 304 may include one or more devices selected from a high-performance computing (HPC) system, including a high-performance core, a microprocessor, a microcontroller, a digital signal processor, a microcomputer, a central processing unit, a field programmable gate array, a programmable logic device, a state machine, a logic circuit, an analog circuit, a digital circuit, or any other device that manipulates a signal (analog or digital) based on computer executable instructions residing in the memory 302. The memory 302 may include a single memory device or multiple memory devices, including but not limited to random access memory (RAM), volatile memory, non-volatile memory, static random access memory (SRAM), dynamic random access memory (DRAM), flash memory, cache memory, or any other device capable of storing information. The non-volatile storage device 306 may include one or more persistent data storage devices, such as a hard drive, an optical drive, a tape drive, a non-volatile solid-state device, a cloud storage, or any other device capable of persistently storing information.

[0047] The processor 304 may be configured to read into the memory 302 and execute computer executable instructions, which reside in the GTFF software module 308 of the non-volatile storage device 306 and embody the GTFF algorithm and / or method of one or more embodiments. The processor 304 may further be configured to read into the memory 302 and execute computer executable instructions, which reside in the MD software module 310 (such as LAMMPS) of the non-volatile storage device 306 and embody the MD algorithm and / or method. The software modules 308 and 310 may include an operating system and applications. The software modules 308 and 310 may be compiled or interpreted from a computer program created using a variety of programming languages ​​and / or technologies, including but not limited to the following and the following alone or in combination: Java, C, C++, C#, Objective C, Fortran, Pascal, Java Script, Python, Perl, and PL / SQL. In one embodiment, PyTorch is a package for the Python programming language that can be used to implement the code of the GTFF of one or more embodiments. The code framework may be based on the Crystalline Graph Convolutional Neural Network (CGCNN) code available under license from the Massachusetts Institute of Technology, Cambridge, Massachusetts.

[0048] When executed by the processor 304, the computer executable instructions of the GTFF software module 308 and the MD software module 310 can cause the computing platform 300 to implement one or more of the GTFF algorithms and / or methods and the MD algorithms and / or methods disclosed herein, respectively. The non-volatile storage device 306 can also include GTFF data 312 and MD data 314 that support the functions, features, and processes of one or more embodiments described herein.

[0049] The program code embodying the algorithms and / or methods described herein can be distributed individually or collectively as a program product in a variety of different forms. The program code can be distributed using a computer-readable storage medium having computer-readable program instructions thereon for causing a processor to implement aspects of one or more embodiments. Inherently non-transitory computer-readable storage media may include volatile and non-volatile, removable and non-removable tangible media implemented in any method or technology for storing information, such as computer-readable instructions, data structures, program modules, or other data. The computer-readable storage medium may further include RAM, ROM, erasable programmable read-only memory (EPROM), electrically erasable programmable read-only memory (EEPROM), flash memory or other solid-state memory technology, portable compact disk read-only memory (CD-ROM) or other optical storage device, cassette, tape, disk storage device or other magnetic storage device, or any other medium that can be used to store the desired information and can be read by a computer. The computer-readable program instructions can be downloaded from the computer-readable storage medium to a computer, another type of programmable data processing device, or another device, or downloaded to an external computer or external storage device via a network.

[0050] Computer-readable program instructions stored in a computer-readable medium can be used to instruct a computer, other types of programmable data processing devices, or other devices to operate in a particular manner so that the instructions stored in the computer-readable medium produce an article of manufacture including instructions for implementing the functions, actions, and / or operations specified in the flowchart or diagram. In certain alternative embodiments, the functions, actions, and / or operations specified in the flowchart and diagram can be reordered, processed serially, and / or processed concurrently in accordance with one or more embodiments. In addition, any flowchart and / or diagram may include more or fewer nodes or boxes than those illustrated in accordance with one or more embodiments.

[0051] Although the entirety of the present invention has been illustrated by the description of various embodiments, and although these embodiments have been described in considerable detail, the applicant does not intend that the scope of the appended claims be limited or in any way restricted to such details. Additional advantages and modifications will be readily apparent to those skilled in the art. Therefore, the present invention in its broader aspects is not limited to the specific details, representative devices and methods, and illustrative examples shown and described. Therefore, deviations from such details may be made without departing from the spirit or scope of the overall inventive concept.

Claims

1. A computational method for simulating the motion of elements within a multi-element system using a graph transformer neural network, the method comprising: Convert molecular dynamics snapshots of elements within a multi-element system into graphs with atoms as graph nodes; Define a matrix so that each column of the matrix represents a node in the graph; Define a distance matrix based on the set of relative positions of each atom; Iterating through the graph transformer neural network using an attention mechanism, operating on the matrix, and augmenting by merging distance matrices to pass hidden states from a current layer of the graph transformer neural network to a next layer of the graph transformer neural network; performing a combination in the form of a sum or mean over the columns of the matrix to produce a scalar molecular energy; making a backward pass through the graph transformer neural network, iteratively computing derivatives at each of the layers of the graph transformer neural network to compute predictions of forces acting on each atom; as well as Returns a prediction of the forces acting on each atom.

2. The calculation method according to claim 1, wherein a neural network is used to implement the combination over the columns of the matrix. 3 . The computational method of claim 1 , further comprising including other molecular objects other than atoms as graph nodes in the matrix, the other molecular objects comprising one or more bonds or quantum states of the atoms. 4 . The computational method of claim 3 , further comprising defining distances from other molecular objects to other objects in the graph such that a distance matrix remains well defined for each element.

5. The computational method of claim 1 , further comprising training a graph transformer neural network as a deep neural network using a loss function, wherein the loss function is formed to compare force predictions with ground truth forces on a labeled dataset such that network weights are optimized using gradient descent with respect to the loss function. The computational method according to claim 1 , wherein the computational method is integrated into a molecular dynamics software.

7. The computational method of claim 1 further comprises utilizing softmax in an attention mechanism to propagate hidden states forward through a graph transformer neural network.

8. A computing system for simulating element motion within a multi-element system using a graph transformer neural network, the system comprising: a memory storing instructions for a graph transformer neural network force field algorithm of a molecular dynamics software; and A processor programmed to execute instructions to perform operations including Convert molecular dynamics snapshots of elements within a multi-element system into graphs with atoms as graph nodes; Define a matrix so that each column of the matrix represents a node in the graph; Define a distance matrix based on the set of relative positions of each atom; Iterating through the graph transformer neural network using an attention mechanism, operating on the matrix, and augmenting by merging distance matrices to pass hidden states from a current layer of the graph transformer neural network to a next layer of the graph transformer neural network; performing a combination in the form of a sum or mean over the columns of the matrix to produce a scalar molecular energy; making a backward pass through the graph transformer neural network, iteratively computing derivatives at each of the layers of the graph transformer neural network to compute predictions of forces acting on each atom; as well as Returns a prediction of the forces acting on each atom.

9. The computing system of claim 8, wherein the processor is further programmed to execute instructions to perform combinations over the columns of the matrix using a neural network.

10. The computing system of claim 8, wherein the processor is further programmed to execute instructions to include other molecular objects besides atoms as graph nodes in the matrix, the other molecular objects including one or more bonds or quantum states of the atoms.

11. The computing system of claim 10, wherein: The processor is further programmed to execute instructions to define distances of other molecular objects to other objects in the graph such that the distance matrix remains well defined for each element.

12. The computing system of claim 8, wherein the processor is further programmed to execute instructions to train the graph transformer neural network as a deep neural network using a loss function formed to compare force predictions to ground truth forces on a labeled dataset such that network weights are optimized with respect to the loss function using gradient descent.

13. The computing system of claim 8, wherein the processor is further programmed to execute instructions to utilize a softmax in an attention mechanism to propagate hidden states forward through the graph transformer neural network.

14. A non-transitory computer readable medium comprising instructions for simulating element motion within a multi-element system using a graph transformer neural network, the instructions, when executed by a processor, causing the processor to: Convert molecular dynamics snapshots of elements within a multi-element system into graphs with atoms as graph nodes; Define a matrix so that each column of the matrix represents a node in the graph; Define a distance matrix based on the set of relative positions of each atom; Iterating through the graph transformer neural network using an attention mechanism, operating on the matrix, and augmenting by merging distance matrices to pass hidden states from a current layer of the graph transformer neural network to a next layer of the graph transformer neural network; performing a combination in the form of a sum or mean over the columns of the matrix to produce a scalar molecular energy; making a backward pass through the graph transformer neural network, iteratively computing derivatives at each of the layers of the graph transformer neural network to compute predictions of forces acting on each atom; as well as Returns a prediction of the forces acting on each atom.

15. The non-transitory computer readable medium of claim 14, further comprising instructions that, when executed by a processor, cause the processor to: including other molecular objects other than atoms as graph nodes in the matrix, the other molecular objects including one or more bonds or quantum states of the atoms; and Define the distances of other molecular objects to other objects in the graph so that the distance matrix remains well defined for each element.

16. The non-transitory computer-readable medium of claim 14, further comprising instructions that, when executed by the processor, cause the processor to train a graph transformer neural network as a deep neural network using a loss function formed to compare force predictions to ground truth forces on a labeled dataset such that network weights are optimized using gradient descent with respect to the loss function.

Images