Graph transformer neural network force field for prediction of atomic forces and energies in molecular dynamic simulations

Abstract

A simulation includes 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 according to a set of relative positions of each of the atoms; iterating through the GTFF using an attention mechanism, operating on the matrix and augmented by incorporating the distance matrix, to pass hidden state from a current layer of the GTFF to a next layer of the GTFF; performing a combination over the columns of the matrix to produce a scalar molecular energy; making a backward pass through the GTFF, iteratively calculating derivatives at each of the layers of the GTFF to compute a prediction of force acting on each atom; and returning the prediction of the force acting on each atom.