[Paper Review] DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation

Duo Zhang, Hangrui Bi+4 more|arXiv (Cornell University)|Aug 17, 2022

Machine Learning in Materials ScienceMaterials Science27 citations

TL;DR

DPA-1 introduces an attention-based Deep Potential model pretrained on large-scale multi-element data and fine-tuned for downstream PES tasks, achieving superior transferability and sample efficiency. It also reveals interpretable element-type embeddings that align with periodic trends.

ABSTRACT

Machine learning assisted modeling of the inter-atomic potential energy surface (PES) is revolutionizing the field of molecular simulation. With the accumulation of high-quality electronic structure data, a model that can be pretrained on all available data and finetuned on downstream tasks with a small additional effort would bring the field to a new stage. Here we propose DPA-1, a Deep Potential model with a novel attention mechanism, which is highly effective for representing the conformation and chemical spaces of atomic systems and learning the PES. We tested DPA-1 on a number of systems and observed superior performance compared with existing benchmarks. When pretrained on large-scale datasets containing 56 elements, DPA-1 can be successfully applied to various downstream tasks with a great improvement of sample efficiency. Surprisingly, for different elements, the learned type embedding parameters form a $spiral$ in the latent space and have a natural correspondence with their positions on the periodic table, showing interesting interpretability of the pretrained DPA-1 model.

Motivation & Objective

• Motivate a scalable, accurate PES model that can leverage large-scale pretraining across many elements.
• Develop an attention-based Deep Potential (DPA-1) that preserves energy conservation, smoothness, and parallel scalability.
• Demonstrate pretrained-to-downstream transfer efficiency across AlMgCu, SSE, and HEA systems.
• Showcase interpretability of learned element-type embeddings and their relation to the periodic table.

Proposed method

• Introduce a local embedding with element-type information and a self-attention mechanism over neighbor atoms.
• Construct local descriptors via a transformed relative coordinate matrix and type embeddings, followed by an attention-based aggregation to build trainable descriptors.
• Train with a joint energy/force loss using Adam, with a scheduling scheme for balancing energy and force terms.
• Fine-tune pretrained models by adjusting the final energy bias and selectively freezing parameters, notably the type embeddings.
• Pretrain on OC2M (56 elements) and evaluate on downstream tasks including AlMgCu, SSE, and HEA datasets.

Experimental results

Research questions

• RQ1 Can a pretrained attention-based PES model generalize across varying compositions (single, binary, ternary) and different datasets?
• RQ2 Does pretraining reduce data and computational requirements for downstream tasks compared to scratch training?
• RQ3 Do learned element-type embeddings exhibit physically interpretable structure related to the periodic table?
• RQ4 Can the model maintain energy conservation, smoothness, and scalability for large-scale simulations?

Key findings

• DPA-1 outperforms DeepPot-SE in validation accuracy across AlMgCu, SSE, and HEA when transferring from limited composition data.
• Pretraining with OC2M enables strong sample efficiency gains, reducing the need for ternary data by large factors (e.g., in AlMgCu, 90% fewer ternary samples).
• OC2M-pretrained DPA-1 achieves competitive energy/force MAEs on OC2M validation (0.681 eV, 0.076 eV/Å) against other state-of-the-art models.
• Element-type embeddings in OC2M-pretrained DPA-1 arrange roughly in a spiral that correlates with periodic table groupings, indicating interpretable physics.
• DPA-1 scales to datasets with many elements without quadratic increases in training cost, unlike some baseline models.
• The pretraining approach yields better sample efficiency and maintains energy-conserving, smooth forces suitable for MD.