cs.LGOct 1, 2026
SaveBranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials
Organizations: Harvard University · Robert Bosch LLC Research and Technology Center
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 while reducing memory usage by up to . 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.
Figures & tables
Figure 1: Overview of BranchIP. Routers placed before model layers terminate sufficiently refined features to reduce computation. Features represent local atomic interactions.
Figure 2: (a) Schematic of the Branchlegro network. Blue and red arrows represent scalar and tensor features. is the atomic species of atom , is the relative displacement vector from atom to . For simplicity, the residual update is not presented. Continuation gates modulate the tensor features and they are routed out in inference when they have a zero-gate (dashed red arrows). Routed scalar features are concatenated back for the energy readout (dashed blue arrows). (b) Pseudocode of the Branchlegro forward pass. For edge and layer , are the scalar features, are the environment weights, are the tensor features, are the spherical harmonic embeddings, is the number of layers, and is a gating function. The initial embedding is the standard Allegro one.
Figure 3: Accuracy–throughput trade-offs for Branchlegro on CO/Pt and CDP using bespoke models and FMs. Each point reports the mean test weighted-sum MAE and MD throughput over seven seeds; error bars denote standard deviation. The inverted -axis indicates higher accuracy (higher is better). Weights for CO/Pt are 1:1:1 for per-atom energy, forces, and stress; for CDP are 10:1 for per-atom energy and forces. Throughput was measured as described in Appendix D using system sizes of 4,160 and 512 atoms for CO/Pt bespoke models and FMs, and 14,336 and 768 atoms for CDP bespoke models and FMs. Black markers denote Allegro baselines. Only Pareto frontier points are shown; tolerance scaling stops once no new frontier points are found. The markers indicate the configurations used for the main results.
Figure 4: MD throughput and memory usage scaling of Allegro and Branchlegro on CO/Pt and CDP. Branchlegro values are averaged over seven seeds, with error bars indicating the standard deviation. All evaluations followed the setup in Appendix D . Speedups and memory reductions are relative to the corresponding full-sized Allegro model. Crosses denote out-of-memory (OOM).
| CO/Pt | CDP | ||||
| Architecture | [meV/atom] | [meV/Å] | [meV/Å 3 ] | [meV/atom] | [meV/Å] |
| Allegro {0} | |||||
| Allegro {0,1} | |||||
| Branchlegro {}_{\bm{\{\prescript{\scalebox{0.5}{\hookleftarrow}}{}{0},\prescript{\scalebox{0.5}{\hookleftarrow}}{}{1}\}}} | {\scriptstyle\pm0.1042370563} | {\scriptstyle\pm0.2161352951} | {\scriptstyle\pm0.005536132909} | {\scriptstyle\pm0.02604456412} | {\scriptstyle\pm0.6819899932} |
| Allegro {0,1,2,3} | |||||
| Branchlegro {}_{\bm{\{\prescript{\scalebox{0.5}{\hookleftarrow}}{}{0},\prescript{\scalebox{0.5}{\hookleftarrow}}{}{1},\prescript{\scalebox{0.5}{\hookleftarrow}}{}{2},\prescript{\scalebox{0.5}{\hookleftarrow}}{}{3}\}}} | {\scriptstyle\pm0.2981581678} | {\scriptstyle\pm1.21551682} | {\scriptstyle\pm0.007910900502} | {\scriptstyle\pm0.05709050628} | {\scriptstyle\pm3.172767509} |
Table 1: Accuracy (MAE) of Allegro and Branchlegro on test datasets. Branchlegro errors are shown as mean standard deviation over seven seeds. Superscript denotes a router before a layer.
Figure 5: Adaptive routing in CO/Pt and CDP Branchlegro FMs. Atom opacity represents the cumulative gating of its connected edges retained in the TP operations after each router. (a) and (b) CO/Pt and CDP structures before routing and after Routers 1–4. For clarity, the CO/Pt visualization retains only the two closest planes of Pt atoms. (c) Top views of CO/Pt at Router 3 for configurations with varying CO adsorption. Only the upper two Pt layers and CO molecules within 2.2 Å of the uppermost Pt layer are shown. (d) Unweighted normalized radial distribution functions (RDFs), , and retention-weighted RDFs, , after Routers 1–4 for Pt–C/O and Pt–Pt interactions in CO/Pt and H–P and O–O in CDP. RDFs are normalized once using the initial peak value.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: The continuation gate function used in Branchlegro. The black curve denotes the parameter setting used in this work.
Table 2: Bespoke model architecture and training hyperparameters used for CO/Pt.
Table 3: FM fine-tuning hyperparameters used for CO/Pt (architecture from Kavanagh et al. (2026) ).
Table 4: Bespoke model architecture and training hyperparameters used for CDP.
Table 5: FM fine-tuning hyperparameters used for CDP (architecture from Kavanagh et al. (2026) ).
Figure 7: (a) Equation of state of fcc Pt, energy per atom as a function of atomic volume. Equilibrium volumes ( ) and bulk moduli ( ) were obtained using ASE’s stabilized-jellium equation-of-state fit. (b) Absolute energy of an isolated CO dimer as a function of bond length, with denoting the lowest-energy sampled distance. Crosses mark the corresponding energy minima. Branchlegro results are reported as mean standard deviation over seven seeds.
Figure 8: Radial distribution functions for O-H, O-O, and P-O pairs in molecular dynamics (MD) simulations of CDP at 525 K with ab initio MD reference.
Figure 9: Arrhenius plot of diffusion coefficients for L1 Allegro, L2 Allegro, and Branchlegro bespoke models. The activation energy, reported in the legend, is determined from the slope of the fitted linear regression.
Figure 10: Energy conservation during molecular dynamics in the microcanonical ensemble. Each panel shows the drift in total energy, reported as the change in total energy in meV per atom relative to the value at 10 ps. All cases show stable and conservative dynamics.
Figure 11: Schematic of CO adsorption sites on Pt.
Figure 12: Edge classifications in CDP.