BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials
Laura Zichi ⋅ Gil Harari ⋅ Chuin W Tan ⋅ Albert Zhu ⋅ Marc Descoteaux ⋅ Menghang Wang ⋅ H. Kung ⋅ Boris Kozinsky
Abstract
Equivariant machine learning interatomic potentials (MLIPs) have revolutionized atomistic modeling, but accurate treatment of complex materials and molecules demands expensive models. This limits simulation length- and time-scales, with tensor products representing a key computational bottleneck. The recent emergence of foundation-scale MLIPs further exacerbates this challenge. We present BranchIP, a single-model architecture for learned adaptive tensor product computation, trained with a novel distillation loss. Across heterogeneous catalysis and proton-conducting solid acid electrolyte, BranchIP accelerates bespoke and foundation MLIPs by up to 5.7$\times$ at lower memory cost, with limited accuracy degradation while maintaining physical fidelity. Furthermore, the learned routing provides model interpretability by revealing which interactions demand deeper computation and showing how computational depth relates to chemical complexity, dynamics, and reactivity.
Chat is not available.
Successful Page Load