# PolyGET: Accelerating Polymer Simulations by Accurate and Generalizable
Forcefield with Equivariant Transformer

Rui Feng<EMAIL_ADDRESS>Huan Tran<EMAIL_ADDRESS>Aubrey Toland
<EMAIL_ADDRESS>Binghong Chen<EMAIL_ADDRESS>Qi Zhu<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>*  [

###### Abstract

Polymer simulation with both accuracy and efficiency is a challenging task.
Machine learning (ML) forcefields have been developed to achieve both the
accuracy of ab initio methods and the efficiency of empirical force fields.
However, existing ML force fields are usually limited to single-molecule
settings, and their simulations are not robust enough. In this paper, we
present PolyGET, a new framework for Polymer Forcefields with Generalizable
Equivariant Transformers. PolyGET is designed to capture complex quantum
interactions between atoms and generalize across various polymer families,
using a deep learning model called Equivariant Transformers. We propose a new
training paradigm that focuses exclusively on optimizing forces, which is
different from existing methods that jointly optimize forces and energy. This
simple force-centric objective function avoids competing objectives between
energy and forces, thereby allowing for learning a unified forcefield ML model
over different polymer families. We evaluated PolyGET on a large-scale dataset
of 24 distinct polymer types and demonstrated state-of-the-art performance in
force accuracy and robust MD simulations. Furthermore, PolyGET can simulate
large polymers with high fidelity to the reference ab initio DFT method while
being able to generalize to unseen polymers.

[1]Rampi Ramprasad [1]Chao Zhang

[1]Georgia Institute of Technology, Georgia Institute of Technology, 225 North
Avenue NW, Atlanta, 30332, GA, USA

2]University of Illinois Urbana-Champaign, University of Illinois Urbana-
Champaign, 601 East John Street, Champaign, 61820, IL, USA

##  1 Introduction

Efficient and accurate forcefield computation is essential for polymer
simulation, but it remains a challenging task. Although ab initio methods
offer excellent accuracy, they are computationally slow. For example, Density
Functional Theory (DFT)[1, 2] has a time complexity of $N_{e}^{3}$, where
$N_{e}$ is the number of electrons, rendering it too expensive for large
polymers or long-horizon simulations. In contrast, empirical forcefields are
more computationally efficient, as they simplify chemical interactions between
atoms to a sum over bonded and unbonded pairwise terms[3, 4]. This approach
reduces the time complexity to $N_{a}^{2}$, where $N_{a}$ is the number of
atoms, allowing for simulating larger molecular systems. However, their
accuracy is limited due to the simplification of chemical interactions, which
results in suboptimal molecular dynamics (MD) simulation performance and
chemical insights [5].

Machine learning (ML) forcefields combine ab initio accuracy with empirical
forcefield computational efficiency. Early methods, such as kernel regression
[6, 7, 8] and feedforward neural networks [9], rely on hand-crafted features,
limiting flexibility and expressivity for complex molecular structures and
interactions [5]. Recent Equivariant Graph Neural Networks (EGNNs)[10, 11, 12,
13] enable automatic feature learning, achieving state-of-the-art accuracy in
predicting energy and forces. However, existing EGNNs face two main
limitations: (1) They are trained in a single-molecule setting, necessitating
separate models for different molecules and lacking generalization to unseen
molecules. In Section5.2, we will demonstrate that adapting EGNNs to the
multi-molecule paradigm is nontrivial due to conflicting objectives in
optimizing energy and force. (2) They are not robust enough for simulations
[14], despite achieving high accuracy in static forcefield prediction. The
single-molecule training paradigm can lead to overfitting to individual
molecules, compromising the model’s ability to perform robustly in dynamic
simulations.

Figure 1: PolyGET learns a single forcefield model across various polymer
families by capturing quantum mechanical interactions. The unified model
captures generalizable knowledge from ab-initio reference calculations.
PolyGET enables accurate and reliable MD simulations while being able to
generalize to large, unseen polymers.

We introduce a new ML forcefield approach for polymer simulation, named _Poly
mer Forcefields with Generalizable Equivariant Transformers_ (PolyGET). As
depicted in Figure 1, PolyGET has two key advantages compared to existing EGNN
models: (1) _Multi-molecule training and generalization_ : The model trains
over a diverse set of polymer families and can generalize to previously unseen
and larger polymers from smaller ones. (2) _Accurate and robust simulations_ :
The model can produce reliable and precise forces during molecular dynamics
(MD) simulations. PolyGET achieves these advantages through two novel designs:

First, we train a unified multi-molecule model using Equivariant Transformer
[10] as the backbone model. As a variant of the Transformer model [15] based
on the attention mechanism, the Equivariant Transformer facilitates learning
interactions between atoms. By incorporating multiple layers of attention
blocks, the model effectively captures complex quantum interactions.
Furthermore, the Equivariant Transformer learns latent atomic features that
are roto-translational equivariant to atomic positions, ensuring the
invariance of the learned energy, which is necessary for energy conservation.
Its expressivity and capacity make it a suitable backbone for PolyGET.

Second, we propose a novel training paradigm that facilitates multi-molecule
training with a unified ML model. While existing deep learning approaches [10,
16] jointly optimize potential energy and forces, our training paradigm
focuses solely on the optimization of forces, which we have found to be
generalizable across multiple polymer families. With our training paradigm,
the underlying ML model acquires robust and transferable knowledge about
quantum mechanical interactions between atoms, enabling both accurate and
robust MD simulations and the capacity to generalize to unseen and larger
polymers with similar quantum interactions. Although our model concentrates on
force optimization, it can also output energy that is linearly correlated with
ground truth energy, allowing the model to estimate ground truth energy with
numerical integration.

We have trained and evaluated the PolyGET model on Poly24, a novel benchmark
comprising 24 polymer types and 6,552,624 conformations from four categories:
cycloalkanes, lactones, ethers, and others. PolyGET demonstrates state-of-the-
art force accuracy and generalizes to previously unseen large polymers.
Furthermore, we have assessed the capacity of PolyGET to perform accurate MD
simulations. We found that even for large 15-loop polymers not present in the
training data, PolyGET can simulate polymers with high fidelity to the
reference DFT data.

In summary, we introduce a novel PolyGET for polymer forcefields, achieving:
(1) a unified Equivariant Transformer model that can transfer across multiple
polymer families for forcefield and energy prediction, (2) accurate and robust
MD simulation for large unseen polymers, and (3) state-of-the-art performance
in force prediction and MD simulation on a large benchmark Poly24. This
paradigm can be applied to various deep learning models, enabling the
development of general ML forcefields and robust MD simulation for a broader
range of chemical families.

##  2 The PolyGET Model

###  2.1 Force-centric Multi-molecule Optimization

Let $x$ denote a molecule consisting of $n_{x}$ atoms. The conformations of
these atoms are represented by $\mathbf{r}_{x}\in\mathbb{R}^{n_{x}\times 3}$
and their atom types are denoted by $\mathbf{z}_{x}\in\mathbb{R}^{n_{x}}$. The
potential energy, $E$, is a function of the conformations and atom types. The
forces exerted on the atoms are derived by
$F=-\frac{\partial{E}}{\partial\mathbf{r}}$. In molecular dynamics (MD)
simulations, molecular conformations are influenced by the forces applied and
the simulation thermostat. Define
$\mathbb{T}^{\mathrm{DFT}}(x,\mathbf{r}_{0},\mathcal{T})$ as the distribution
of conformations in a density functional theory (DFT)-based MD simulation for
a specific thermostat $\mathcal{T}$ and initial conformation $\mathbf{x}_{0}$.
Machine learning (ML) force fields aim to learn a model, $\Phi_{\theta}$,
parameterized by $\theta$, which approximates DFT-generated reference data.
Existing ML force field models optimize model parameters $\theta$ by
minimizing the following objective:

$\theta=\arg\min\mathbb{E}^{\mathrm{DFT}}\left[\|E^{\mathrm{ML}}-E^{\mathrm{DFT}}\|^{2}_{2}+\|F^{\mathrm{ML}}-F^{\mathrm{DFT}}\|^{2}_{2}\right].$
(1)

This optimization process is typically conducted in single-molecule settings
by existing EGNN models, which learn a separate model for each molecule. As
discussed in Section 1 and will be demonstrated in Section 5, models trained
in a single-molecule setting not only fail to generalize to unseen molecules
but also perform poorly in simulations due to their tendency to overfit the
single molecules. Although it is natural to question whether a unified model
can be trained to leverage DFT reference data for multiple molecules, this
task is nontrivial. The challenge of training a unified model for multiple
molecules can be understood by analyzing the per-atom potential energy and
force distributions. As illustrated in Figure 2, although forces follow
similar distributions for different types of polymers, the per-atom potential
energy does not. We hypothesize that the mismatch between potential energy and
force distributions is the primary cause of the optimization bottleneck and
the reason for the preference for single-molecule settings in existing
methods. In Section 5.2, we demonstrate that when extending existing EGNN
methods for multi-molecule training— _i.e._ , training a unified ML force
field for various types of molecules—their performance can deteriorate instead
of improving.

Figure 2: The distribution of per-atom potential energies and forces across
different polymers. Forces exhibit similar distributions for various types of
polymers, whereas the per-atom potential energy does not.

In our approach, we address this challenge and learn generalizable force
fields from multiple molecules by proposing a new learning objective:

$\mathbb{E}_{x,\mathbf{r}0}\left[\mathbb{E}_{\mathbf{r}\sim\mathbb{T}^{\mathrm{DFT}}(x,\mathbf{r}0,\mathcal{T})}\left[\left\|-\nabla\Phi\theta(\mathbf{r},\mathbf{z})-F(\mathbf{r},\mathbf{z})\right\|^{2}_{2}\right]\right].$
(2)

In this objective, we minimize the differences between the predicted forces
$-\nabla\Phi_{\theta}(\mathbf{r},\mathbf{z})$ and the ground-truth forces
$F(\mathbf{r},\mathbf{z})$. The loss is optimized across multiple DFT
trajectories for different types of polymers.

The primary distinction between our approach and existing methods lies in the
exclusive focus on forces, rather than jointly optimizing the potential energy
and forces. This design choice is grounded in two key rationales: 1) forces,
being a local atomic attribute, exhibit greater transferability across diverse
molecules [17, 18], and 2) forces demonstrate similar distributions among
different molecules, as evidenced in Figure 2. Conversely, the potential
energy, a molecular-level attribute, displays considerable variation in its
distributions—even after normalization—due to the intricate quantum mechanical
processes involved. By focusing on forces during training, our model can
efficiently transfer information across DFT-generated data without being
hindered by competing objectives between energy and force. In order to perform
well across a wide array of molecules and conformation distributions, the
unified model must effectively capture generic atomic interactions, rather
than overfitting to a single molecule. As a result, this comprehensive
training on various molecules bolsters our model’s capacity to generalize to
unseen molecules and conduct more robust simulations.

Figure 3: The model output is not directly optimized to be aligned with the
ground truth potential energy but can approximate the latter with a linear
transform.

###  2.2 Energy Prediction

The objective function (2) provides a supervision signal to the model only on
the forces. In order to predict the potential energy for molecules during
simulation, we observe that by training exclusively on forces with (2), the
model output $\Phi_{\theta}$ is linearly correlated with the ground-truth
potential energy $E$ (Figure 3). We thus call the output energy from
$\Phi_{\theta}$ as _uncalibrated potential energy_. To predict the true
potential energy $E$ from $\Phi_{\theta}$, we can estimate $E_{t}$ with
${E}_{t}=E_{0}+\int_{0}^{t}\frac{\partial\Phi_{\theta}}{\partial\mathbf{r}}(\mathbf{r}_{s},\mathbf{z})ds$
for a fixed time window $t\in[0,T]$, where $E_{0}$ is the initial energy of at
the start of the simulation. The integral can be approximated with numeric
Taylor expansion:

$\overline{E}_{t}=\overline{E}_{t-1}-\left\langle\frac{\partial{\Phi}}{\partial{\mathbf{r}}}(\mathbf{r}_{t-1}),\mathbf{r}_{t}-\mathbf{r}_{t-1}\right\rangle.$

Then given uniformly sampled $0=t_{0}<\ldots<t_{m}=T$, we can estimate the
potential energy function with

$\displaystyle E^{*}(\mathbf{r},\mathbf{z})$
$\displaystyle=\mathbf{W}_{E}\Phi_{\theta}(\mathbf{r},\mathbf{z})+\mathbf{b}_{E},$
$\displaystyle\mathrm{where}\ $ $\displaystyle(\mathbf{W}_{E},\mathbf{b}_{E})$
$\displaystyle=\arg\min\sum_{i=0}^{m}\|(\mathbf{W}_{E}\Phi_{\theta}+\mathbf{b}_{E})-\overline{E}_{i}\|.$
(3)

Figure 4:  Equivariant Transformer learns to predict a conservative vector
field optimized to be atomic forces by taking the negative gradient of an
invariant scalar output. The atomic forces can be used for MD simulations with
an integrator.

##  3 Model architecture: Equivariant Transformer

In conjunction with our new learning objective, we use Equivariant Transformer
[10] as our backbone model for PolyGET. Equivariant Transformer uses the
powerful _self-attention mechanism_ [15] and the Transformer architecture for
molecular data. The self-attention mechanism allows the model to capture
interatomic interactions and the attraction-repulsion dynamics. Furthermore,
Equivariant Transformer builds roto-translational equivariant to atomic
positions — this not only allows it to learn in a more data-efficient way, but
also ensures the invariance of the learned energy for the energy conservation
law.

As shown in Figure 4, the Equivariant Transformer architecture operates on
molecular data in the form of graphs and consists of three main blocks: an
embedding layer, an updated layer, and an output network.

#### Embedding Layer

The first component is an atom embedding layer, which transforms atoms to
vector representations that encode the quantum mechanical information about
the atoms and their surroundings. The atom embedding is further split to
_scalar_ embeddings $\mathbf{h}_{0}\in\mathbb{R}^{D}$ and _vector_ embeddings
$\mathbf{v}_{0}\in\mathbb{R}^{3\times D}$, where $D$ is the embedding size.
The scalar embedding is derived by integrating both the intrinsic vector,
which captures information specific to the atoms themselves, and the
neighborhood vector, which encapsulates the interactions within the atomic
neighborhoods. The intrinsic vector, denoted as $\mathbf{x}^{\mathrm{int}}$,
is acquired by associating each atom type with a corresponding dense feature
vector in $\mathbb{R}^{D}$. To compute the neighborhood embedding
$\mathbf{x}^{\mathrm{nbr}}$, we initially define the distance filter using a
radial basis function. Subsequently, a continuous-filter convolution technique
is employed to aggregate information from the surrounding atoms within a given
neighborhood. The atom embedding $\mathbf{x}$ is obtained by the concatenation
of intrinsic and neighborhood embeddings. The vector embedding $\mathbf{v}$ is
initialized as zero vectors in $\mathbb{R}^{3\times D}$.

#### Update Layers

The update layers defines the incremental change of the _scalar embeddings_
$\mathbf{x}$ and _vector embeddings_ $\mathbf{v}$ in each network layer. For
each atom $i$, the update layers are defined as follows:

$\displaystyle\nabla\mathbf{x}_{i}$
$\displaystyle=f^{x}(\mathbf{x}_{i})+\sum_{j}\phi^{x}(\mathbf{x}_{i},\mathbf{x}_{j})+\psi^{xv}(\mathbf{v}_{i})$
(4) $\displaystyle\nabla\mathbf{v}_{i}$
$\displaystyle=f^{v}(\mathbf{v}_{i})+\sum_{j}\frac{\mathbf{r}_{j}-\mathbf{r}_{i}}{\|\mathbf{r}_{j}-\mathbf{r}_{i}\|}\phi^{v}(\mathbf{x}_{i},\mathbf{x}_{j})+\psi^{vx}(\mathbf{x}_{i}).$

For each atom $i$, the attention mechanism $\phi^{x}$ and $\phi^{v}$ will
model its interatomic potentials to all other atoms. Based on the embeddings
of two atoms and their relative distance vectors, $\phi^{x}$ and $\phi^{v}$
define the interaction between the two atoms. $f^{x}$ and $f^{v}$ are the
residual connections [19] that prevent the network to suffer from performance
degradation. $\psi^{xv}$ and $\psi^{vx}$ define the information exchange
between $\mathbf{x}$ and $\mathbf{v}$ in each layer. $\psi^{xv}$ is used to
update the vector embedding $\mathbf{v}$ from the scalar embedding
$\mathbf{x}$, and $\psi^{vx}$ is used to update the scalar embedding
$\mathbf{x}$ from the vector embedding $\mathbf{v}$.

#### Output Network

The output network computes scalar atom-wise predictions using gated
equivariant blocks, which is a modified version of the gated residual network
[10]. The gated equivariant blocks aggregate the atomic embeddings into a
single molecular prediction. This prediction can be matched with a scalar
target variable or differentiated against atomic coordinates to provide force
predictions.

Index | Type | # Atoms | Monomer | Polymer
---|---|---|---|---
CK1 | cyclopropane | 9 |  |
CK2 | cyclobutane | 12 |  |
CK3 | cyclopentane | 15 |  |
CK4 | cyclohexane | 18 |  |
CK5 | cycloheptane | 21 |  |
CK6 | cycloctane | 24 |  |
ETH | ethylene oxide | 7 |  |
LAC1 |  $\gamma$-butyrolactone | 9 |  |
LAC2 |  $\gamma$-butyrolactone | 12 |  |
LAC3 | 1,4-dioxan-2-one | 13 |  |
LAC4 |  $\delta$-valerolactone | 15 |  |
LAC5 | 3-methyl-1,4-dioxan-2-one | 16 |  |
LAC6 |  $\beta$-methyl-$\delta$-valerolactone | 18 |  |
LAC7 |  $\delta$-caprolactone | 18 |  |
LAC8 |  $\delta$-decalactone | 30 |  |
LAC9 | (-)-Menthide | 30 |  |
OTH1 | n-alkane sub $\delta$-valerolactone | 18 |  |
OTH2 |  $\alpha$-Methylene-$\gamma$-butyrolactone | 13 |  |
OTH3 | n-alkane sub $\delta$-valerolactone | 21 |  |
OTH4 | n-alkane sub $\delta$-valerolactone | 24 |  |
OTH5 | n-butyl $\delta$-valerolactone | 27 |  |
OTH6 | n-alkane sub $\delta$-valerolactone | 30 |  |
OTH7 | n-alkane sub $\delta$-valerolactone | 33 |  |
OTH8 | n-alkane sub $\delta$-valerolactone | 42 |  |
Table 1: Polymers in our used Poly24 dataset for training PolyGET. We follow
[20] and classify polymers into 4 broad categories: cycloalkanes, lactones,
ethers, and others.

##  4 Dataset

Our PolyGET model is trained on a DFT-generated dataset [21] for computing
ring-opening enthalpy ($\Delta H^{\rm ROP}$), which is a crucial thermodynamic
concept governing ring-opening polymerization. This process involves breaking
a ring of cyclic monomers and adding the “opened” monomer to a long chain,
which forms a polymer chain. The data were generated by utilizing MD
simulations of multiple monomer and polymer models at a uniform level of DFT
[1, 2] computations. More details about our data generation process can be
found in Refs. [21, 22].

Given a cyclic monomer, several polymers models were created using Polymer
Structure Predictor (PSP) package [23]. Each of them was obtained by
multiplying the monomer with a small integer $L$, e.g., $L=3,4,5$, and $6$,
forming a loop of size $L$ (larger loops are better models of polymers). For
each monomer or polymer models, about 10 or more maximally diversified
configurations were selected as the initial configurations of the DFT-based MD
simulations.

As all of our models are non-periodic, we used the $\Gamma$-point version of
Vienna Ab initio Simulation Package (vasp),[24, 25] employing a basis set of
plane waves with kinetic energy up to 400 eV to represent the Kohn-Sham
orbitals. The ion-electron interactions were computed using the projector
augmented wave (PAW) method [26] while the exchange-correlation (XC) energies
were computed using the generalized gradient approximation Perdew-Burke-
Ernzerhof (PBE) functional.[27]

We show the polymers used in this study in Table 1, which displays all 24
polymer types which have been broadly classified into cycloalkanes, ethers,
lactones, and others, following the classification used in [21]. Averagely,
the dataset includes 10 DFT trajectories for both monomers and polymers. The
polymers are formed by polymerizing the monomers, resulting in a comprehensive
collection of data for each polymer type. The average DFT trajectory lengths
for polymers up to 6-loop are listed in Table 2. In total, we have 1311 DFT
trajectories and 6,552,624 molecular conformations.

| Monomer | 2-loop | 3-loop | 4-loop | 5-loop | 6-loop
---|---|---|---|---|---|---
# Trajectory. | 233 | 20 | 280 | 270 | 282 | 220
Avg. Length | 9526.31 | 8464.15 | 6182.98 | 4293.98 | 3180.02 | 1614.87
Table 2: Average trajectory length w.r.t. polymer size.

##  5 Results

(a) Forces Correlation for 15 Loop CK6.

(b) Forces Correlation for 15 Loop OTH4.

(c) Forces Correlation for 15 Loop LAC2.

(d) Forces Correlation for 15 Loop LAC5.

(e) RMSD and Potential Energy for 15 Loop CK6.

(f) RMSD and Potential Energy for 15 Loop OTH4.

(g) RMSD and Potential Energy for 15 Loop LAC2.

(h) RMSD and Potential Energy for 15 Loop LAC5. (i) Dynamics Illustration for
CK6. (j) Dynamics Illustration for OTH4. (k) Dynamics Illustration for LAC2.
(l) Dynamics Illustration for LAC5.

Figure 5: Accuracy and Robustness of PolyGET in MD simulations for large
polymers. We select 15 loop CK6 (cycloctane), OTH4 (n-alkane substituted
$\delta$-valerolactone), LAC2 ($\gamma$-butyrolactone) and LAC5
(3-methyl-1,4-dioxan-2-one) as examples. In (a)-(d) it is observed that
PolyGET forces are highly correlated with DFT forces during its own MD
simulations. (e)-(h) shows the RMSD and potential energy curves of PolyGET
during MD simulations. The curves suggest that PolyGET simulations experience
equilibration. (i)-(j) presents a visualization of sample conformations
captured at various time points during the simulation process, which shows
PolyGETcould effectively explore and reveal the dynamic changes in polymer
structures over time.

In this section, we present the empirical performance of PolyGET on Poly24.
First, we report on the simulation performance of PolyGET on large 15-loop
unseen polymers in Section 5.1, demonstrating the accuracy, robustness, and
speed of PolyGET for large polymer simulation. Then, we explore the multi-
molecule training performance of PolyGET in Section 5.2. We show that our
training paradigm is crucial in enhancing force prediction accuracy and the
generalization capabilities of the model.

###  5.1 MD Simulation Accuracy

#### Model Training and Simulation Setup.

We have developed a 6-layer Equivariant Transformer with an embedding
dimension of 128 and 8 attention heads. During training, unless specified
otherwise, we trained the model with 5-loop polymers and smaller ones.
Polymers of size 6 and above were unseen during training and were used as test
data. We used 5% of the training data as the validation set and recorded the
validation loss after every 500 training steps. The model has trained with the
AdamW [28] optimizer for 3 epochs with an initial learning rate of $0.0001$,
and the learning rate was reduced by a factor of $0.8$ whenever the validation
loss increased for 30 consecutive validation steps.

For testing the model’s MD simulation performance, we take the 15-loop
polymerization of 4 example polymers from Table 1, CK6, OTH4, LAC2, and LAC5,
and run MD simulations on them with PolyGET-generated forcefields for a
maximum of 600K steps. The number of atoms for these 15-loop polymers are,
respectively: 360, 360, 180, and 240. The maximum loop size is seen in the
training data of PolyGET is 5-loop. Therefore, PolyGET is required to perform
accurate MD simulations on unseen polymers, testing both its simulation and
generalization ability.

For simulation, we have integrated our method into the ASE [29] simulation
environment. A diagram of the simulation is shown in Figure 4. We use
Nosé–Hoover thermostat for all simulations. The velocities for molecules are
initialized with Maxwell-Boltzmann distribution to simulate the desired
temperature environment. Unless otherwise specified, we use a temperature of
300 Kelvin and a timestep of 0.5fs.

#### Results.

Figure 5 (a)-(d) visualize the comparison results between PolyGET predicted
forces and DFT reference data. For all 15 loop polymers, PolyGET obtains
highly accurate forces to DFT calculations, with 0.01 eV/$\AA$ MAE in energy
and close to zero cosine distance. Hence, PolyGET can perform MD simulations
on multiple types of polymers with uniformly high correlation with DFT
references. The test 15-loop polymers contain up to 360 atoms, unseen in the
training data. This shows that PolyGET can extrapolate well to large unseen
polymers with known monomers. Furthermore, in practice, PolyGET only needs
training data from small polymers, which are cheap to generate _ab initio_
data and thus greatly reduce the cost of DFT data generation.

In Figure 5 (e) to (h), the Root Mean Square Deviation (RMSD) and potential
energy graphs are shown. RMSD curve measures how much the polymer structure
deviates from the initial conformation, and the potential energy graph
indicates the energy state of the simulated polymers. Both RMSD and potential
energy curves indicate that the polymer in simulation converges to a stable
distribution around the equilibrium. Furthermore, PolyGET can explore
different possible conformation states of the polymers, as illustrated in
Figure 5 (i) to (l).

#### Simulation Efficiency.

The PolyGET model was trained on all polymers with loop size 5 and below using
two A5000 GPUs, which took approximately 30 hours. However, this is a one-time
cost, as the trained model can serve as a unified model for all future
simulations of different polymers.

For 15-loop polymers, the model predicts atomic forces for 20-30 steps per
second on average. This means that the simulations for 600K steps, as
mentioned in the previous section, can be completed in 6-8 hours. In contrast,
the DFT model takes hours on CPU parallelization and still 5-10 minutes on GPU
parallelization for one single-step calculation — we speed up DFT simulations
by at least 6000 times. Depending on the available GPU memory, the PolyGET
model can be parallelized to perform multiple simulations simultaneously. We
leave the implementation of this acceleration for future work.

(a) SchNet (b) NequIP (c) TorchMDNet (d) PolyGET

Figure 6: Performance of deep learning forcefield, including SchNet, NequIP,
TorchMDNet, and PolyGET on multi-molecule training. Two metrics, Mean Absolute
Error (MAE) and Cosine Distance (CosD) are reported. On the left, (a) shows
the MAE comparison between models across 4 types of monomers: cycloalkanes,
lactones, ethers, and others, as in Table 1. On the right, (b) - (e) shows the
scatter plots of DFT forces compared with predicted forces for the 4 methods.

###  5.2 Effectiveness of Multi-molecule Training

In this subsection, we show that the multi-molecule training paradigm can
crucially improve the prediction performance of PolyGET. To this end, we
compare PolyGET with 3 state-of-the-art machine learning forcefields models:
SchNet [16], NequIP [12], and TorchMDNet [10]. These methods feature different
implementations of EGNN and have achieved accurate forcefield prediction on
static DFT trajectories.

(a) Multi-molecule vs. Single-molecule.

(b) Extrapolation to Unseen Polymers.

(c) Validation Forces MAE.

Figure 7: Figure 7(c): validation forces MAE of the Equivariant Transformers
during training, where “Original” refers to the original paradigm of joint
energy and forces optimization, and “PolyGET” refers to our force-centered
multi-molecule training paradigm. Figure 7(a): comparison of forces accuracy
of models trained on single molecules and multiple molecules. Figure 7(b): we
analyze PolyGET’s extrapolation performance to polymers not seen during
training.

#### Our model achieves state-of-the-art multi-molecule force prediction
accuracy.

To demonstrate this, we train PolyGET and the baseline models on all polymers
with loop sizes below 6 (Table 1). We then evaluate the force prediction
performance of different models on unseen 6-loop polymers of the same
monomers. Table 6 reports the results. As shown, our model PolyGET attains
0.073 eV/Å MAE and a near-zero cosine distance compared to DFT forces, which
is 6 times better than the previous state-of-the-art model TorchMDNet.

#### Our model benefits the most from diverse training data.

In Figure 7(a), we compared the performance of baseline models when trained
separately or jointly on different polymers. For NequIP and TorchMDNet, joint
training on all polymers brings no significant improvement and could even harm
force prediction accuracy. In contrast, our model benefits from having
multiple types of polymers in the training data, significantly reducing force
prediction errors. It is worth noting that our model shares a similar
Equivariant Transformer architecture with TorchMDNet [10]. The performance
gain is mainly because of our force-centered training paradigm. As depicted in
Figure 7(c), This paradigm allows the validation MAE on forces to decrease
more consistently in multi-molecule training, leading to a significantly
improved model performance. Such results demonstrate the effectiveness of our
force-centric multi-molecule training paradigm.

#### The multi-molecule paradigm improves model extrapolation.

In Figure 7(b), we report the baseline models’ performance on “unseen”
polymers compared to “seen” polymers. “Seen” polymers refer to polymers with
loop sizes below 6 present in the training data, while “unseen” polymers
pertain to 6-loop polymers. In line with previous multi-molecule performance
observations, our model experiences minimal performance deterioration (0.012
eV/Å to 0.013 eV/Å) when applied to unseen polymers. In contrast, NequIP’s
force MAE drops from 0.080 eV/Å to 0.103 eV/Å, while TorchMDNet drops from
0.043 to 0.056. This result shows that the multi-molecule paradigm plays a
vital role in enhancing the model’s extrapolation capabilities.

#### Our model accurately predicts energy with linear transformation.

Although our model does not directly optimize potential energy, it still
captures the potential energy distributions for different polymers, as shown
in Figure 8(a). Averaged across different polymers, PolyGET achieves an
average of 0.0067 cosine distance between the model output and the DFT
potential energy. After learning a linear regression on eight ground-truth
conformations for each polymer, our model attains an average of 0.0006
eV/Å/atom MAE between the transformed model energy prediction and the DFT
potential energy and near 0 cosine distance. We can use this adjusted model
output to predict accurate potential energies for polymers, as shown in 8(b).

(a) Model Output vs DFT Potential Energy.

(b) Adjusted Model Output vs DFT Potential Energy.

Figure 8: PolyGET’s Energy Prediction Performance. Figure 8(a): PolyGET model
output plotted against DFT potential energy in eV/Å/atom. Figure 8(b): PolyGET
model predicts accurate potential energy after applying linear transforms for
each polymer.

##  6 Related Work

Machine learning forcefields have been attracting increasing attention, in
order to achieve both high accuracy of _ab initio_ methods and the computation
efficiency of empirical forcefields [17, 5, 18]. Traditional machine learning
approaches rely on “fingerprints”, pre-engineered interatomic features
describing atomic environments. Similar to classical forcefields, these
features describe pairwise radial or angular relations between atoms, and
relatively simple machine learning algorithms such as kernel regression [6, 7,
8] and neural networks (NN) [5, 17, 18] can transform these features to find
the optimal statistical relations between the features and the potential
energy and forces. Particularly, neural networks have the universal
approximation property [30] that in theory can approximate arbitrarily complex
functions, implying that with good enough features, the NN models can
approximate _ab initio_ forces with high accuracy. Of course, in practice,
optimality is difficult to achieve due to problems such as overfitting, model
capacity, and the quality of feature engineering. Still, neural networks have
been used in calculating forces for many different types of materials [31, 32,
33, 34, 35, 36]. Despite their success, traditional machine learning
approaches rely on hand-crafted features that are not flexible enough to
represent the complex quantum mechanical interactions between atoms. Imagine
if the features are all based on pairwise distances between atoms. In such a
case, no matter how complex the subsequent model is, the information it uses
to make predictions would not extend beyond the pairwise relations defined in
the features. This is not desirable if more complex quantum mechanics are
involved.

The state-of-the-art approaches in machine learning forcefield take advantage
of the recent advances in Graph Neural Networks (GNN). Instead of pre-defined
fingerprint features, these algorithms model directly the interaction between
atoms in the latent space. Properly designed with equivariant constraints,
they have theoretically universal approximation ability to equivariant and
invariant physical systems [37]. By treating molecules as graphs with 3D
coordinates, GNN enables the learning of quantum mechanical interactions
between atoms directly with equivariant operators, such as neural message
passing [38], EGNN [13], SchNet [16], DIME [39], GemNet [40], and NequIP [12],
and TorchMDNet [10]. With stronger model capacity, they could achieve
comparable accuracy with _ab initio_ forcefield. However, these models
typically are only trained and evaluated on given conformations from single
molecules [10, 12, 16, 40] (i.e., in-distribution data), and it is nontrivial
to directly extend them to train in multi-molecule settings. Furthermore,
existing machine learning methods are prone to instability on out-of-
distribution data, limiting their performance in MD simulations where noises
are accumulated for atomic positions. The instability of existing methods is
empirically analyzed by [14].

Among existing machine learning forcefields, AGNI [17, 18] and sGDML [6, 7]
agree with us in focusing on forces optimization. Several different reasons
are stated. [6] observes the amplification effect of the differentiation
operator on noises. Combined with the inherently noisy nature of machine
learning algorithms, [6] claims that models trained on potential energy would
cause considerable noise in the forces predictions. [17, 18] present a
learning scheme where the _local_ forces of atoms are learned and used to
simulate geometry optimizations and MD simulations, instead of the _global_
potential energies. Adding to their observations, we find empirical evidence
that the potential energy and forces are competing in the optimization of
machine learning algorithms and that the forces of atoms, being localized
features that can be derived from local atomic environments, have the benefit
of better generalization across different molecules.

##  7 Conclusion

We introduce PolyGET, a novel polymer forcefields learning framework with
Generalizable Equivariant Transformers. Focusing on force optimization
exclusively, our training paradigm enables accurate molecular dynamics
simulations and generalization across polymer families and unseen polymers.
Our multi-molecule approach fosters robust, transferable knowledge of
interatomic interactions, contrasting with existing forcefields that jointly
optimize potential energy and forces. Our results on large polymer simulation
datasets demonstrate PolyGET’s accuracy, robustness, and speed, showcasing its
potential for material science and industry applications.

In the future, we aim to extend PolyGET’s training to polymers beyond carbon,
hydrogen, and oxygen. Additionally, we will investigate the extent to which
PolyGET performs on vastly different chemical families and explore efficient
fine-tuning strategies to enhance the model’s adaptability with minimal
effort.

## References

  *   * [1] Hohenberg, P. & Kohn, W.  Inhomogeneous electron gas.  _Phys. Rev._ 136, B864–B871 (1964).
  * [2] Kohn, W. & Sham, L.  Self-consistent equations including exchange and correlation effects.  _Phys. Rev._ 140, A1133–A1138 (1965).
  * [3] Unke, O. T., Koner, D., Patra, S., Käser, S. & Meuwly, M.  High-dimensional potential energy surfaces for molecular simulations: from empiricism to machine learning.  _Machine Learning: Science and Technology_ 1, 013001 (2020).
  * [4] González, M.  Force fields and molecular dynamics simulations.  _École thématique de la Société Française de la Neutronique_ 12, 169–200 (2011).
  * [5] Unke, O. T. _et al._ Machine learning force fields.  _Chemical Reviews_ 121, 10142–10186 (2021).
  * [6] Chmiela, S. _et al._ Machine learning of accurate energy-conserving molecular force fields.  _Science advances_ 3, e1603015 (2017).
  * [7] Chmiela, S., Sauceda, H. E., Müller, K.-R. & Tkatchenko, A.  Towards exact molecular dynamics simulations with machine-learned force fields.  _Nature communications_ 9, 3887 (2018).
  * [8] Chmiela, S. _et al._ Accurate global machine learning force fields for molecules with hundreds of atoms.  _arXiv preprint arXiv:2209.14865_ (2022).
  * [9] Svozil, D., Kvasnicka, V. & Pospichal, J.  Introduction to multi-layer feed-forward neural networks.  _Chemometrics and intelligent laboratory systems_ 39, 43–62 (1997).
  * [10] Thölke, P. & De Fabritiis, G.  Torchmd-net: Equivariant transformers for neural network based molecular potentials.  _arXiv preprint arXiv:2202.02541_ (2022).
  * [11] Fuchs, F., Worrall, D., Fischer, V. & Welling, M.  Se(3)-transformers: 3d roto-translation equivariant attention networks.  _Advances in Neural Information Processing Systems_ 33, 1970–1981 (2020).
  * [12] Batzner, S. _et al._ E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials.  _Nature communications_ 13, 1–11 (2022).
  * [13] Satorras, V. G., Hoogeboom, E. & Welling, M.  E (n) equivariant graph neural networks 9323–9332 (2021).
  * [14] Fu, X. _et al._ Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations.  _arXiv preprint arXiv:2210.07237_ (2022).
  * [15] Vaswani, A. _et al._ Attention is all you need.  _Advances in neural information processing systems_ 30 (2017).
  * [16] Schütt, K. T., Sauceda, H. E., Kindermans, P.-J., Tkatchenko, A. & Müller, K.-R.  Schnet-a deep learning architecture for molecules and materials.  _The Journal of Chemical Physics_ 148, 241722 (2018).
  * [17] Ramprasad, R., Batra, R., Pilania, G., Mannodi-Kanakkithodi, A. & Kim, C.  Machine learning in materials informatics: recent applications and prospects.  _npj Computational Materials_ 3, 54 (2017).
  * [18] Botu, V., Batra, R., Chapman, J. & Ramprasad, R.  Machine learning force fields: construction, validation, and outlook.  _The Journal of Physical Chemistry C_ 121, 511–522 (2017).
  * [19] He, K., Zhang, X., Ren, S. & Sun, J.  Identity mappings in deep residual networks 630–645 (2016).
  * [20] Huan, T. D. _et al._ A universal strategy for the creation of machine learning-based atomistic force fields.  _NPJ Computational Materials_ 3, 37 (2017).
  * [21] Tran, H. _et al._ Toward recyclable polymers: Ring-opening polymerization enthalpy from first-principles.  _J. Phys. Chem. Lett._ 13, 4778–4785 (2022).
  * [22] Huan, T. D. _et al._ A polymer dataset for accelerated property prediction and design.  _Sci. data_ 3, 1–10 (2016).
  * [23] Sahu, H., Shen, K.-H., Montoya, J. H., Tran, H. & Ramprasad, R.  Polymer structure predictor (psp): a python toolkit for predicting atomic-level structural models for a range of polymer geometries.  _J. Chem. Theory Comput._ 18, 2737–2748 (2022).
  * [24] Kresse, G. & Furthmüller, J.  Efficiency of ab-initio total energy calculations for metals and semiconductors using a plane-wave basis set.  _Comput. Mater. Sci._ 6, 15–50 (1996).
  * [25] Kresse, G. & Furthmüller, J.  Efficient iterative schemes for ab initio total-energy calculations using a plane-wave basis set.  _Phys. Rev. B_ 54, 11169–11186 (1996).
  * [26] Blöchl, P. E.  Projector augmented-wave method.  _Phys. Rev. B_ 50, 17953–17979 (1994).
  * [27] Perdew, J. P., Burke, K. & Ernzerhof, M.  Generalized gradient approximation made simple.  _Phys. Rev. Lett._ 77, 3865–3868 (1996).
  * [28] You, Y. _et al._ Large batch optimization for deep learning: Training bert in 76 minutes.  _arXiv preprint arXiv:1904.00962_ (2019).
  * [29] Larsen, A. H. _et al._ The atomic simulation environment—a python library for working with atoms.  _Journal of Physics: Condensed Matter_ 29, 273002 (2017).
  * [30] Baker, M. R. & Patil, R. B.  Universal approximation theorem for interval neural networks.  _Reliable Computing_ 4, 235–239 (1998).
  * [31] Kruglov, I., Sergeev, O., Yanilkin, A. & Oganov, A. R.  Energy-free machine learning force field for aluminum.  _Scientific reports_ 7, 1–7 (2017).
  * [32] Artrith, N. & Behler, J.  High-dimensional neural network potentials for metal surfaces: A prototype study for copper.  _Physical Review B_ 85, 045439 (2012).
  * [33] Behler, J.  Atom-centered symmetry functions for constructing high-dimensional neural network potentials.  _The Journal of chemical physics_ 134, 074106 (2011).
  * [34] Chen, C. _et al._ Accurate force field for molybdenum by machine learning large materials data.  _Physical Review Materials_ 1, 043603 (2017).
  * [35] Deringer, V. L., Caro, M. A. & Csányi, G.  A general-purpose machine-learning force field for bulk and nanostructured phosphorus.  _Nature communications_ 11, 1–11 (2020).
  * [36] Hajibabaei, A., Ha, M., Pourasad, S., Kim, J. & Kim, K. S.  Machine learning of first-principles force-fields for alkane and polyene hydrocarbons.  _The Journal of Physical Chemistry A_ 125, 9414–9420 (2021).
  * [37] Sannai, A., Takai, Y. & Cordonnier, M.  Universal approximations of permutation invariant/equivariant functions by deep neural networks.  _arXiv preprint arXiv:1903.01939_ (2019).
  * [38] Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O. & Dahl, G. E.  Neural message passing for quantum chemistry 1263–1272 (2017).
  * [39] Gasteiger, J., Groß, J. & Günnemann, S.  Directional message passing for molecular graphs.  _arXiv preprint arXiv:2003.03123_ (2020).
  * [40] Gasteiger, J., Becker, F. & Günnemann, S.  Gemnet: Universal directional graph neural networks for molecules.  _Advances in Neural Information Processing Systems_ 34, 6790–6802 (2021).