BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

1 INTRODUCTION

• Branchlegro is demonstrated on CO/Pt catalysis and proton transport in CDP, which has no obvious chemically motivated spatial partition.Both systems require large length- and time-scale molecular dynamics to sample rare events.Source: §1 · p. 2
• BranchIP is a single-model framework for feature-level adaptive tensor-product computation in MLIPs.Features that are sufficiently refined avoid further tensor-product computation.Source: §1 · p. 2
• Learned dynamic routers use a novel distillation objective to route features through tensor products while obeying exact energy conservation.Routers smoothly determine which features continue through tensor-product computation.Source: §1 · p. 2
• The Allegro implementation retains a twice continuously differentiable potential energy surface and equivariance.Source: §1 · p. 2
• BranchIP accelerates inference and reduces memory with limited accuracy degradation.The introduction demonstrates the method on bespoke and foundation MLIPs for CO/Pt and CDP.Source: §1 · p. 2
• Learned routing reveals which interactions require deeper equivariant computation and how this varies with chemistry and dynamics.Source: §1 · p. 2

2 BACKGROUND AND RELATED WORK

• MLIPs predict total energy from atomic positions and species, while Allegro uses local equivariant pair features and neighbor-pair tensor products to incorporate many-body geometry.Automatic differentiation yields forces and potentially stresses.Source: §2 · p. 2
• Tensor-product acceleration has focused on cheaper formulations and optimized kernels, approaches described as complementary to BranchIP.Source: §2 · p. 3
• Existing adaptive potentials use predefined regions, hand-crafted descriptors, global or element-wise expert routing, or fixed dynamic cutoffs; BranchIP adapts computation within one model.These architectures are complementary to BranchIP.Source: §2 · p. 3
• Early exits and finer-grained adaptive methods allocate computation dynamically, but learned dynamic allocation of equivariant computation across model depth had not been explored.Prior examples include skipping blocks, selecting feature channels, and varying graph-neural-network message passing.Source: §2 · p. 3

3 LEARNING ADAPTIVE EQUIVARIANT COMPUTATION

• BranchIP uses lightweight routers to gate features before tensor-product layers, omitting zero-gated features at inference while retaining them during training.This allocates depth adaptively across features rather than applying the full depth uniformly.Source: §3 · p. 3
• Continuous gates scale features before the next tensor product; a flat exact-zero region lets removed features’ values and derivatives vanish.This permits zero-gated features to be omitted without changing energy or its derivatives.Source: §3 · pp. 3–4
• With invariant, same-per-feature C2 gates and removal only at exact zero, sparse evaluation preserves C2 regularity, exact derivatives, E(3)-equivariance, and permutation symmetry.The sparse path reproduces the dense evaluation exactly.Source: §3 · pp. 3–4
• Lightweight router MLPs are trained by distilling a frozen full-compute teacher, with the student BranchIP initialized from that teacher.The loss uses prediction discrepancy to supervise stopping or continuing, with a tolerance interval left unsupervised.Source: §3 · pp. 4–5
• In Branchlegro, edge-feature routers sparsify only tensor products; terminated edges retain scalar updates, and environment weights are evaluated over all edges at every layer.The schematic shows zero-gated tensor features routed out and routed scalar features returned to the energy readout.Source: Figure 2 · p. 4, §3 · p. 5

4 RESULTS

• Relaxing routing tolerances increased throughput with moderate error increases up to a model-dependent knee, and the selected models were near validation-set knees.Some profiles near routing-regime transitions had greater seed-to-seed throughput variance; two of seven CDP FM seeds had worse validation and test errors across all tolerance profiles.Source: §4 · pp. 5–6
• 1.6–2.4× speedups and 1.2–2.6× memory reductions were achieved across model sizes and systems relative to full-sized Allegro.Figure 4 shows the throughput and memory scaling comparison; the reported values reflect selected models and systems, not an inherent performance limit.Source: Figure 4 · p. 6, §4 · p. 6
• Branchlegro maintained comparable energy, force, and stress MAEs and physical observables, with CDP proton-diffusion activation energy of 0.40 eV versus the experimental range of 0.39–0.43 eV.The CDP L2 teacher was close to Branchlegro, while L1 Allegro was outside the experimental range at 0.46 eV; Branchlegro models obeyed energy conservation in MD.Source: §4 · pp. 6–7, Table 1 · p. 6

4.1 INTERPRETABILITY

• In CO/Pt, later routers preferentially retain adsorbed-CO and surface Pt interactions while routing out most gas-phase CO and distinguishing Pt environments.The second router separates bulk, bare-surface, and adsorbate-bound Pt; the third further distinguishes top and multiple adsorption configurations, consistent with greater need for higher body-order information.Source: §4.1 · p. 7, Figure 5 · p. 8
• In CDP, framework edges are mostly retained early and terminated by the final router, while a subset of H- and O-involving edges remains for deeper computation.The reported routing analysis covers H–P and O–O interactions across the four routers.Source: §4.1 · p. 7, Figure 5 · p. 8
• For H–P edges, nearly all unrelated edges terminate at Router 1, covalent edges at Router 3, and hydrogen-bonded edges at Router 4.The covalent P–O–H geometry is near-rigid, whereas hydrogen-bonded geometry depends on the arrangement of two reorienting PO4 groups.Source: §4.1 · pp. 7, 9
• For O–O edges, same-PO4 pairs terminate at Router 2, but O-pair hydrogen-bond edges remain at full amplitude through all four routers.The classes nearly coincide in distance; O-pair edges depend on proton position and require a larger neighborhood than intramolecular edges.Source: §4.1 · p. 9

5 CONCLUSION

• BranchIP is a learned dynamic adaptive inference framework for equivariant MLIPs that obeys energy conservation and uses routing to avoid further computation when features have sufficient tensor-product refinement.The conclusion reports Branchlegro speed and memory gains on CO/Pt and CDP without loss of physical fidelity or reliance on hand-crafted routing.Source: §5 · p. 9
• Branchlegro only reduces tensor-product computations, retaining other layer computations, and its learned routing requires a pretrained teacher that increases total training cost.The speedup constraint is especially relevant when other operations, such as large edge-wise MLPs, are expensive; experiments cover two systems.Source: §5 · p. 9

A.1 FORCE DISTILLATION LOSS FORMULATION

• For edge-based Allegro features, force distillation converts atom-resolved force discrepancies into edge discrepancies because routers predict edge-wise probabilities.The stated force discrepancy is the L2 norm between teacher and student force predictions; Equation 1 is then adapted for Branchlegro’s edge features.Source: §A.1 · p. 14

A.2 CONTINUATION-GATING FUNCTION

• Branchlegro’s continuation gate uses p=3 for a C2 potential energy surface, with T=0.3 and w=0.1 controlling the transition.The black curve marks the parameter setting used in this work.Source: §A.2 · pp. 14–15, Figure 6 · p. 15

B THEORETICAL GUARANTEES

• Branchlegro is algebraically equivalent to dense gated evaluation and preserves energy gradients, Hessians, equivariance, and permutation symmetries.The guarantees concern the sparse routing implementation of Allegro.Source: B THEORETICAL GUARANTEES · pp. 15–16

B.1 ENERGY EQUIVALENCE

• Under A1–A2, dense gated and sparse active-set evaluation produce identical scalar latent features at every layer and therefore identical energies.Only tensor-product computation is sparsified; scalar features and environment weights remain evaluated for all edges.Source: §B.1 · p. 16

B.2 DERIVATIVE EQUIVALENCE

• Under A1–A3, sparse differentiation matches dense gated energy gradients and Hessians; the equivalence also extends to stresses and to higher orders when the components have matching smoothness.At active-set boundaries, an omitted branch has zero value and zero first- and second-derivative contribution.Source: §B.2 · p. 18

B.3 SYMMETRY PRESERVATION

• An E(3)-invariant continuation gate from a shared edge-wise router preserves Allegro’s E(3) equivariance, energy invariance, and permutation symmetry.In Branchlegro, the router acts on invariant scalar edge features.Source: §B.3 · p. 19

C TRAINING SETUP

• Training used NVIDIA A100-SXM4-40GB, A100-SXM4-80GB, and H200 GPUs; across Branchlegro runs, seeds varied router initialization and batch ordering while teacher and data splits stayed fixed.Hyperparameter tables report Allegro and Branchlegro training settings for both systems.Source: C TRAINING SETUP · pp. 19–20
• Tables 2 and 4 cover bespoke model architecture and training hyperparameters for CO/Pt and CDP, respectively.Tables 3 and 5 report foundation-model fine-tuning hyperparameters for CO/Pt and CDP.Source: Table 2 · p. 20, Table 3 · p. 20, Table 4 · p. 21
• A configuration excerpt lists stop tolerances 0.12 and 0.155, router MLP depth 2, width 8, and batch size 8.The supplied excerpt does not identify which model or system this configuration belongs to.Source: C TRAINING SETUP · p. 20

D MD THROUGHPUT AND MEMORY MEASUREMENTS

• Throughput benchmarks used LAMMPS NVT simulations at 525 K for CDP and 300 K for CO/Pt, with respective timesteps of 0.5 fs and 1 fs.Runs included a 104-step warm-up and up to 9.99 × 10^6 production steps, with a 1200 s benchmark limit.Source: D MD THROUGHPUT AND MEMORY MEASUREMENTS · p. 22
• Peak allocated GPU memory was recorded after each LAMMPS force evaluation, with the reported value taken as the maximum across evaluations during the simulation.CUDA peak-memory counters were reset before each force evaluation.Source: D MD THROUGHPUT AND MEMORY MEASUREMENTS · p. 22

E.1 ENERGY PROFILE FOR CO/PT

• Branchlegro predicts a 212.4 GPa bulk modulus (+0.90% versus DFT’s 210.5 GPa), compared with bespoke one- and two-layer Allegro deviations of −4.97% and +2.85%.The CO/Pt equation-of-state figure reports results as mean ± standard deviation over seven seeds.Source: Figure 7 · p. 22, §E.1 · p. 23
• All models show close agreement with DFT across physical-fidelity evaluations, reproducing the relevant energy landscapes and derived material properties.Source: §E.1 · p. 22
• For foundation models, Allegro and Branchlegro remain within 1% of the DFT bulk modulus, with deviations of −0.75% and −0.10%, respectively.Source: §E.1 · p. 23

E.2 MD SIMULATIONS FOR CDP ACTIVATION ENERGY

• CDP radial distribution functions for O–H, O–O, and P–O pairs are compared with ab initio MD references at 525 K.The bespoke-model RDFs average three 1000-atom trajectories; foundation-model RDFs average three 144-atom trajectories.Source: Figure 8 · p. 23, §E.2 · p. 23
• Proton-diffusion activation energies are fitted from temperature-dependent diffusion coefficients for bespoke models only, because long-timescale foundation-model simulations are extremely expensive.Diffusion coefficients are evaluated at 525 K, 550 K, 575 K, and 625 K.Source: §E.2 · p. 23

• Branchlegro NVE simulations for CO/Pt and CDP were initialized at 300 K and 525 K and propagated for 100 ps, with energy drift monitored after 10 ps equilibration.Source: F ENERGY CONSERVATION IN MD · pp. 23–24
• Energy drift was assessed over time for different Branchlegro and Allegro models.Source: F ENERGY CONSERVATION IN MD · p. 24
• All cases show stable, conservative dynamics; Figure 10 reports total-energy drift per atom relative to the value at 10 ps.Source: Figure 10 · p. 24

• CO/Pt site classification distinguishes adsorbate-bound, bare-surface, and bulk Pt, then separates adsorbate-bound Pt into top or shared-carbon sites.The rules use 2.6 Å Pt–C and 3.2 Å Pt–Pt cutoffs, with the latter compared against bulk coordination of 12.Source: G INTERPRETABILITY DETAILS · p. 24, Figure 11 · p. 25
• CDP P–H edges are labeled covalent, hydrogen-bonded, or unrelated; O–O edges distinguish shared-PO4 pairs from proton-bridged pairs.The hydrogen-bond and proton-bridge rules use a 2.20 Å threshold.Source: G INTERPRETABILITY DETAILS · p. 24, Figure 12 · p. 25
• The 2.6 Å CO/Pt distance cutoff is determined by JmolNN using radii from Cordero et al. (2008).Source: G INTERPRETABILITY DETAILS · p. 24