Branchip speeds up atom simulation models while saving memory

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

Machine Learning

Summary

Simulating how atoms interact is important but can require a lot of computer power. The authors created BranchIP, a new way to speed up these simulations while using less memory. It decides smartly where to spend more computing effort based on the type of atomic interaction. This method keeps the results accurate and helps explain which atomic interactions are more complex. Their tests showed improvements on materials related to catalysts and special electrolytes.

What this means in practice

  • For materials scientists: Run more efficient atomistic simulations of complex materials by adaptively allocating computation where needed.
  • For chemical engineers: Accelerate simulations of catalytic and electrolyte materials to explore new designs with reduced computational cost.

Authors

Laura Zichi, Gil Harari, Chuin Wei Tan, Marc L. Descoteaux, Albert Zhu, Menghang Wang, Yoel Zimmermann, H. T. Kung, Boris Kozinsky

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.