Online test-time adaptation for better generalization of interatomic potentials to out-of-distribution data

Area of Science:

• Computational chemistry
• Materials science
• Machine learning

Background:

• Machine learning interatomic potentials (MLIPs) offer efficient simulations with high accuracy.
• Distribution shift between training and test data degrades MLIP performance and can cause simulation collapse.

Purpose of the Study:

• To develop an online framework, Test-time Adaptation Interatomic Potential (TAIP), to improve MLIP generalization on unseen test data.
• To address the performance deterioration of MLIPs due to distribution shifts.

Main Methods:

• Proposed a dual-level self-supervised learning approach within the TAIP framework.
• Leveraged global structure and atomic local environment information for model adaptation.
• Implemented online adaptation to align the model with test data distribution.

Main Results:

• TAIP effectively bridges the domain gap between training and test datasets without requiring additional data.
• Demonstrated enhanced test performance across diverse benchmarks, including small molecules and complex periodic systems.
• Enabled stable molecular dynamics (MD) simulations where baseline models previously failed.

Conclusions:

• TAIP significantly improves the robustness and generalization of MLIPs in molecular dynamics simulations.
• The proposed online adaptation method is effective for diverse chemical and material systems.
• TAIP offers a viable solution to the distribution shift problem in MLIP applications.