BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

2026-10-01

AI Summary

This study addresses the prohibitive computational overhead of tensor product operations in equivariant machine learning interatomic potentials, which constrains the simulation scale of complex materials. We propose BranchIP, a single-model framework that introduces an adaptive tensor product computation mechanism driven by a novel distillation loss. By revealing the intrinsic relationship between chemical complexity and computational depth, BranchIP enables dynamic compute allocation while preserving physical fidelity. Experimental results demonstrate that BranchIP accelerates inference by up to 2.4× and reduces memory consumption by 2.6×, all while maintaining high-precision physical simulations. This work effectively overcomes the computational bottlenecks in large-scale materials modeling and enhances model interpretability.

Abstract

Equivariant machine learning interatomic potentials (MLIPs) have revolutionized atomistic modeling, but accurate treatment of complex materials and molecular systems demands expensive models. This limits simulation length- and time-scales, with tensor products a key computational bottleneck. The recent emergence of foundation-scale MLIPs further exacerbates this challenge. We present Branch Interatomic Potential (BranchIP), a single-model framework for learned adaptive tensor product computation, trained with a novel distillation loss. In our experiments on two systems of physical interest, a heterogeneous catalysis system and a proton-conducting solid acid electrolyte, BranchIP accelerates MLIPs across model sizes by up to $2.4\times$ while reducing memory usage by up to $2.6\times$. This is achieved while maintaining physical fidelity. Furthermore, the learned adaptive computation provides model interpretability by revealing which interactions demand deeper computation and showing how computational depth relates to chemical complexity and dynamics.

Problem

• Machine Learning Interatomic Potentials
• Equivariant Models
• Tensor Products
• Computational Efficiency
• Molecular Dynamics

Innovation

• Equivariant Machine Learning Interatomic Potentials
• Adaptive Tensor Product Computation
• Distillation Loss
• Model Interpretability
• BranchIP