Source-linked AI summary
DPA-2: a large atomic model as a multi-task learner
Duo Zhang, Xinzijian Liu, Xiangyu Zhang, Chengqian Zhang, Chun Cai, Hangrui Bi, Yiming Du, Xuejian Qin, Anyang Peng, Jiameng Huang, Bowen Li, Yifan Shan, Jinzhe Zeng, Yuzhi Zhang, Siyuan Liu, Yifan Li, Junhan Chang, Xinyan Wang, Shuo Zhou, Jianchuan Liu, Xiaoshan Luo, Zhenyu Wang, Wanrun Jiang, Jing Wu, Yudi Yang, Jiyuan Yang, Manyi Yang, Fu-Qiang Gong, Linshuang Zhang, Mengchao Shi, Fu-Zhi Dai, Darrin M. York, Shi Liu, Tong Zhu, Zhicheng Zhong, Jian Lv, Jun Cheng, Weile Jia, Mohan Chen, Guolin Ke, Weinan E, Linfeng Zhang, Han Wang
TL;DR
Constructing DFT-labeled datasets is costly, and conventional single-task models have limited generalization across applications and labeling settings. DPA-2 addresses this gap with multi-task pre-training as a large atomic model, showing broad downstream generalization while reducing data requirements and motivating an open, application-oriented ecosystem.
Problem
DFT-labeled dataset construction requires substantial effort, while single-task training limits simultaneous use of differently labeled datasets and downstream generalizability.
Method
DPA-2 is a large atomic model trained through multi-task pre-training across diverse chemical and materials datasets, with a pipeline supporting fine-tuning, knowledge distillation, and deployment.
Results
DPA-2 demonstrates exceptional generalization across downstream tasks, with multi-task pre-training reducing data requirements by approximately 1–2 orders of magnitude.
Takeaways & Limitations
DPA-2 supports a model-centric approach to molecular and materials simulation and underscores the value of open, application-oriented evaluation and collaboration.
Takeaways & Limitations
Current pre-training datasets are deficient in 2-D materials, limiting DPA-2’s generalizability to such systems and making LAM development a long-term endeavor.
Abstract
from arXiv · showhide
The rapid advancements in artificial intelligence (AI) are catalyzing transformative changes in atomic modeling, simulation, and design. AI-driven potential energy models have demonstrated the capability to conduct large-scale, long-duration simulations with the accuracy of ab initio electronic structure methods. However, the model generation process remains a bottleneck for large-scale applications. We propose a shift towards a model-centric ecosystem, wherein a large atomic model (LAM), pre-trained across multiple disciplines, can be efficiently fine-tuned and distilled for various downstream tasks, thereby establishing a new framework for molecular modeling. In this study, we introduce the DPA-2 architecture as a prototype for LAMs. Pre-trained on a diverse array of chemical and materials systems using a multi-task approach, DPA-2 demonstrates superior generalization capabilities across multiple downstream tasks compared to the traditional single-task pre-training and fine-tuning methodologies. Our approach sets the stage for the development and broad application of LAMs in molecular and materials simulation research.
1 Introduction
DPA-2 addresses the costly, poorly generalizing model-generation process by combining multi-task pre-training with fine-tuning and distillation. It generalizes across chemical and materials applications while reducing downstream data and computational requirements.
- Motivation: DFT labeling is expensive, and conventional machine-learning potentials often generalize poorly beyond their training applications.Labeling 141K alloy configurations required more than 10 million CPU hours.
- Motivation: Single-task training cannot readily combine datasets with different DFT labeling settings, limiting generalization across molecular and materials applications.Examples include incompatibilities between MPtrj and ANI-1x datasets.
- Approach: DPA-2 uses multi-task pre-training to learn a unified descriptor across diverse datasets, then applies it through downstream fine-tuning and distillation.The descriptor connects to task-specific fitting networks and can be compressed into a more efficient student model.
- Results: DPA-2 shows strong single-task benchmark accuracy, matching or exceeding several state-of-the-art architectures.It outperforms ANI-1x on ANI-1x and is comparable to GemNet-OC while exceeding Equiformer V2, NequIP, Allegro, and MACE on pre-training datasets.
- Results: Multi-task pre-training substantially improves zero-shot generalization compared with single-task DPA-2 and MPtrj-trained models.Comparable MACE-MP-0 and DPA-2 performance suggests the gain primarily reflects the multi-task scheme rather than architecture differences.
- Results: Fine-tuning saves two orders of magnitude of H2O-PBE0TS-MD training data for equivalent energy accuracy, especially when downstream data are scarce.The advantage diminishes as downstream sample size increases.
- Results: Knowledge distillation improves time-to-solution and the maximum single-GPU system size by nearly two orders of magnitude while retaining fine-tuned-model accuracy.The distilled models preserve accuracy while reducing the computational inefficiency associated with DPA-2’s extensive parameter set.
- Representation: DPA-2 representations organize chemical and configurational environments consistently with periodic-table chemistry and local environments.They also overlap across differently labeled sulfur datasets, indicating insensitivity to the DFT labeling functional in that example.
3 Discussion
DPA-2 is presented as a multi-task large atomic model and workflow for broad molecular and materials simulation. The authors report strong generalization but identify insufficient pre-training data as an ongoing limitation requiring continued ecosystem development.
- DPA-2 combines multi-task pre-training, fine-tuning, knowledge distillation, and practical deployment in a large atomic model workflow.
- 18 datasets spanning 73 chemical elements support DPA-2’s generalization across datasets with differing labeling methodologies.
- Downstream applications require approximately 1–2 orders of magnitude less data without sacrificing accuracy.
- Existing pre-training datasets are insufficient, particularly for 2-D materials, limiting generalizability to such systems.
- The authors frame LAM development as a long-term effort requiring diverse data, application-specific tests, automated workflows, and an open collaborative ecosystem.
4 Methods
DPA-2 represents atomic environments with symmetry-preserving local descriptors, updates these representations through repinit and repformer layers, and predicts atomic energies for PES modeling. Its multi-task protocol shares a unified descriptor across datasets with different DFT settings while retaining dataset-specific fitting networks.
- 4.1 Formulation: The PES is expressed as a function of atomic coordinates and numbers, with atomic energies summed into the total energy and forces obtained as negative energy gradients.For periodic systems, the virial tensor is also computed from the energy and cell vectors.
- 4.2.1 The overall architecture of the DPA-2 model: DPA-2 predicts atomic energy contributions from atomic numbers and coordinates using a descriptor followed by an energy-fitting network.The atomic-number-dependent energy-bias layer adds a learned constant to each atomic contribution.
- 4.2.1 The overall architecture of the DPA-2 model: The descriptor maps local atomic environments to representations invariant or equivariant under translation, rotation, and same-element permutation operations.Its environment matrix uses relative Cartesian coordinates and a switched inverse-distance function within a smooth cutoff region.
- 4.2.2–4.2.3 Representation layers: Repinit updates single-atom and pair-atom representations from local configuration information, while repformer layers repeatedly transform these representations with shared shape dimensions.The model uses a cutoff-neighbor set, and the equivariant pair representation is kept unchanged during repinit.
- 4.2.3 The repformer layer: The repformer attention uses a smooth modified softmax, and its gate term is reported as critical to the model’s generalization ability.The attention projections use trainable Q, K, and V matrices; updating the equivariant representation was not retained because it did not improve accuracy and often destabilized training.
- 4.4–4.5 Multi-task training and fine-tuning: Multi-task training concurrently optimizes one unified descriptor across datasets while each dataset-specific fitting network is trained only on its associated dataset.The shared descriptor is intended to encode common configurational and elemental knowledge, reducing the data needed for downstream fine-tuning.
5 Data and Code Availability
The study’s datasets and models are publicly available through AIS Square, with code, datasets, and input scripts also archived on Zenodo.
- AIS Square hosts the datasets and models used in the study.The paper also provides a Bohrium Notebook and DP Combo web server for model exploration.
S1 Datasets
The study assembles diverse datasets spanning molecular, metallic, semiconductor, electrolyte, perovskite, battery, cluster, and drug-related systems for pre-training and downstream evaluation.
- Twenty-five datasets, including training and test splits for pre-training and downstream tasks, are accessible through AIS Square.
- Materials datasets cover alloys, metallic elements, semiconductors, cathodes, perovskite oxides, electrolytes, water, silver-gold systems, and metal nanoclusters.These datasets span varied compositions, structures, temperatures, pressures, and DFT labeling settings.
- Several datasets are explicitly partitioned between pre-training and downstream use, including Cluster-P/Cluster-D and SSE-PBE-P/SSE-PBE-D.Other datasets are used entirely for pre-training or downstream tasks, depending on their stated role.
- The datasets use varied DFT software, exchange-correlation functionals, and simulation procedures, including VASP, ABACUS, PBE, SCAN, and DP-GEN.
- The collection includes molecular datasets such as ANI-1x and Drug, covering over 5 million conformations and over 1.4 million structures, respectively.ANI-1x contains organic molecules with up to 13 heavy atoms, while Drug includes up to 70 heavy atoms across eight elements.
S2 Single-task benchmark of the generalizability
Single-task DPA-2 benchmarks evaluate accuracy on molecular test sets and compare several models across pre-training datasets under a constrained training budget.
- DPA-2 shows superior accuracy to ANI-1x on six molecular test datasets, including ANI-MD, DrugBank, GDB07-09, GDB10-13, S66x8, and Tripeptide.Energy and force RMSEs are reported for the comparison.
- Six models are compared using energy and force test RMSEs after 1 million training steps with batch size 1.The models are GemNet-OC, Equiformer V2, NequIP, Allegro, MACE, and DPA-2.
- The comparison uses a limited time and computational-resources budget, with one million steps described as insufficient for a large dataset such as OC2M.
- DPA-2 achieves the smallest energy WARMSE and the second smallest force WARMSE among the compared single-task models.WARMSE is a weighted average across the pre-training datasets.
S3 Multi-task training of the DPA-2 on the pre-training dataset
The multi-task DPA-2 model is trained jointly across pre-training datasets with weighted sampling and shared optimization, using a protocol matched to single-task comparisons.
- The multi-task DPA-2 model is trained for 1,000K steps with batch size 8, while single-task comparisons use effective steps adjusted by dataset weights.The adjustment uses weight/13.2×8×1,000K, where 13.2 is the sum of dataset weights.
- Different pre-training datasets receive different sampling weights, and gradients across heads on each GPU are aggregated and averaged.The learning rate decays exponentially from 2e-4 to 3e-8.
- The multi-task model is harder to train because it must fit all datasets with the same model capacity and number of training steps as the single-task model.
S4 Downstream learning curves
Sample efficiency was compared across all 15 downstream tasks, with the results illustrated in Fig. S1.
- Sample efficiency was analyzed across all 15 downstream tasks.The comparative results are illustrated in Fig. S1.
S5 Choices of fine-tuning heads
Downstream fine-tuning can use a task-relevant pre-trained head, a randomly selected head, or a newly initialized head. On ANI-1x, fitting-head choice had minimal influence on learning because most parameters were pre-trained in the descriptor.
- Fine-tuning may select a task-relevant pre-trained head, choose a head randomly, or initialize a new head from scratch.
- On ANI-1x, fitting-head choice had minimal influence on learning.Most parameters were pre-trained within the descriptor.
- The main-text experiments used fitting heads aligned with each downstream task.
S6 Distillation and validation by applications
Model distillation was evaluated on three representative downstream tasks by comparing fine-tuned DPA-2 teachers with distilled DPA-1 students and full-data DPA-1 benchmarks.
- Distillation was conducted on H2O-PBE0TS-MD, SSE-PBE-D, and FerroEle-D.
- Test-data RMSEs were reported for fine-tuned DPA-2 teacher models and distilled DPA-1 student models.
- DPA-1 models with the same architecture were trained on full downstream datasets for comparison.
S7 Ablation study
The ablation study examined structures in the repformer layer by removing components from single-atom and pair-atom representations and tracking average RMSE changes across datasets.
- Sequential component removal was evaluated using average RMSE across all datasets during multi-task pre-training.The study independently examined single-atom and pair-atom channels.
- For the single-atom representation, every structure except sym_f was substantially important.
- For the pair-atom representation, removing components produced varying increases in RMSE.The passage identifies the gate component as part of the pair-atom analysis, but its complete result is truncated.
S8 Energy conservation
NVE simulations on three randomly chosen OC2M structures show that DPA-2 preserves total energy, whereas the compared non-conservative models exhibit significant energy drift despite higher reported accuracy.
- DPA-2 simulations showed no energy drift in NVE simulations of three randomly chosen OC2M structures.The model used for this illustration had energy MAE 0.912 eV and force MAE 0.116 eV/Å after 1 million training steps.
- Despite higher energy and force accuracy, GemNet-OC and Equiformer-V2 produced significant total-energy drift because they lack energy conservation.Their reported energy MAEs were 0.286 eV and 0.236 eV, with force MAEs of 0.026 eV/Å and 0.016 eV/Å, respectively.
S9 Hyper-parameters of the DPA-2 model
The DPA-2 configuration uses a shared descriptor and repformer architecture for single-task and multi-task training, with task-specific fitting networks and a staged optimization schedule. Dataset preprocessing and task-specific distillation settings address differences across datasets and downstream applications.
- Descriptor and representation: Single-task and multi-task DPA-2 training used consistent descriptor configurations, including a 9.0 Å repinit cutoff and maximum neighbor count of 120.The single-atom input representation has dimension 8.
- Network architecture: The descriptor uses three-layer embedding and fitting components, 12 repformer layers, and feature dimensions of 128 for single-atom and 32 for pair-atom representations.The local and gated attention layers use four heads, with hidden dimensions 128 and 32, respectively.
- Parameterization: The multi-task model has approximately 7.68 million parameters, including a 5-million-parameter descriptor and 2.68 million parameters across 18 fitting networks.The single-task model has approximately 5.15 million parameters, including nearly 5 million descriptor parameters and around 0.15 million fitting-network parameters.
- Optimization: Training starts at a learning rate of 2e-4 and decays exponentially to 3e-8 over 1 million steps while energy and force prefactors are adjusted.The energy prefactor changes from 0.02 to 1, and the force prefactor changes from 1,000 to 1.
- Dataset handling and distillation: Dataset-specific procedures include subtracting fitted per-atom energy biases, using task-specific distillation parameters, and splitting OC2M for sequential Allegro or NequIP training when memory is insufficient.For H2O-PBE0TS-MD distillation, the DPA-1 architecture excludes attention layers and uses a 6.0 Å cutoff with at most 120 neighbors.