[Paper Review] PolyGET: Accelerating Polymer Simulations by Accurate and Generalizable Forcefield with Equivariant Transformer

Rui Feng, Tran Doan Huan+5 more|arXiv (Cornell University)|Sep 1, 2023

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

PolyGET introduces a generalizable, force-centric machine learning forcefield for polymer simulations using Equivariant Transformers. By training exclusively on forces rather than joint energy-force optimization, it achieves state-of-the-art accuracy and robustness in molecular dynamics simulations across 24 polymer families, generalizing effectively to unseen and larger polymers with ab initio-level fidelity.

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.

Motivation & Objective

• To develop a machine learning forcefield that generalizes across diverse polymer families while maintaining high accuracy and simulation robustness.
• To overcome the limitations of existing ML forcefields trained on single molecules or with competing energy-force optimization objectives.
• To enable reliable long-horizon molecular dynamics simulations of large, complex polymers using a unified, transferable model.
• To establish a training paradigm focused solely on force prediction to enhance generalization and stability in dynamic simulations.
• To demonstrate state-of-the-art performance on a large-scale polymer benchmark spanning 24 distinct polymer types.

Proposed method

• Employs an Equivariant Transformer as the backbone model to learn roto-translational equivariant atomic representations from 3D atomic coordinates.
• Introduces a novel force-centric training paradigm that optimizes only atomic forces, avoiding competing objectives with potential energy.
• Trains a single unified model on a diverse dataset of 24 polymer families, enabling generalization to unseen and larger polymers.
• Leverages the model's learned force predictions to estimate potential energy via numerical integration, maintaining consistency with ab initio reference data.
• Uses a large-scale benchmark dataset, Poly24, comprising 6.55 million conformations across four categories: cycloalkanes, lactones, ethers, and others.
• Enforces equivariance through attention mechanisms that preserve physical invariance, ensuring energy conservation and force accuracy.

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 generali

Experimental results

Research questions

• RQ1 Can a single, unified machine learning model trained on diverse polymer families generalize to unseen and larger polymers with high fidelity?
• RQ2 Does a force-only training objective lead to more robust and generalizable force predictions than joint energy-force optimization?
• RQ3 Can an Equivariant Transformer-based model achieve ab initio-level accuracy in force prediction while maintaining computational efficiency for large-scale polymer simulations?
• RQ4 How does the proposed training paradigm compare to existing methods in terms of simulation stability and out-of-distribution generalization?
• RQ5 To what extent can the model transfer knowledge across chemically distinct polymer families without fine-tuning?

Key findings

• PolyGET achieves state-of-the-art force accuracy on the Poly24 benchmark, outperforming existing ML forcefields in both in-distribution and out-of-distribution settings.
• The model generalizes effectively to previously unseen large polymers, maintaining high accuracy in molecular dynamics simulations without retraining.
• The force-centric training paradigm leads to more stable and robust simulations compared to joint energy-force optimization, reducing noise amplification in force predictions.
• Despite focusing only on forces during training, the model produces potential energy values that are linearly correlated with ab initio reference energies, enabling accurate energy estimation via integration.
• The unified multi-molecule training approach enables transferable knowledge of quantum mechanical interactions across diverse polymer families, including cycloalkanes, lactones, and ethers.
• PolyGET demonstrates high computational efficiency while preserving the accuracy of ab initio methods, enabling long-horizon simulations of complex polymer systems.

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.