Ensemble Knowledge Distillation for Machine Learning Interatomic Potentials

2025-03-18

AI Summary

High-accuracy coupled-cluster (CC) energy data for small molecules are scarce, and such datasets typically lack corresponding force labels—leading to poor force prediction accuracy and unstable molecular dynamics simulations when training machine-learned interatomic potentials (MLIPs) on energy-only data. Method: We propose an Ensemble Knowledge Distillation (EKD) framework that leverages multiple pre-trained teacher models to generate high-quality force labels, enabling a student MLIP to jointly fit both reference energies and distilled forces. Contribution/Results: This is the first application of ensemble knowledge distillation to MLIP training; it eliminates the need for ground-truth force labels while substantially improving force prediction accuracy and MD stability. Evaluated on the COMP6 benchmark, EKD achieves new state-of-the-art performance, demonstrating strong generalization across chemical, biological, and materials systems.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Data Mining & Knowledge Management: Data Compression
• Humans and AI: Human-in-the-loop Machine Learning

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Humans versus LLMs for data annotation and labeling
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Machine learning interatomic potentials (MLIPs) are a promising tool to accelerate atomistic simulations and molecular property prediction. The quality of MLIPs strongly depends on the quantity of available training data as well as the quantum chemistry (QC) level of theory used to generate that data. Datasets generated with high-fidelity QC methods, such as coupled cluster, are typically restricted to small molecules and may be missing energy gradients. With this limited quantity of data, it is often difficult to train good MLIP models. We present an ensemble knowledge distillation (EKD) method to improve MLIP accuracy when trained to energy-only datasets. In our EKD approach, first, multiple teacher models are trained to QC energies and then used to generate atomic forces for all configurations in the dataset. Next, a student MLIP is trained to both QC energies and to ensemble-averaged forces generated by the teacher models. We apply this workflow on the ANI-1ccx dataset which consists of organic molecules with configuration energies computed at the coupled cluster level of theory. The resulting student MLIPs achieve new state-of-the-art accuracy on the out-of-sample COMP6 benchmark and improved stability for molecular dynamics simulations. The EKD approach for MLIP is broadly applicable for chemical, biomolecular and materials science simulations.

Problem

• Improves MLIP accuracy with limited energy-only datasets.
• Generates atomic forces using ensemble teacher models.
• Enhances molecular dynamics simulation stability and accuracy.

Innovation

• Ensemble knowledge distillation improves MLIP accuracy
• Teacher models generate atomic forces for training
• Student MLIP trained on QC energies and forces