How does FlashCart make equivariant interatomic potentials faster and more accurate?

FlashCart trades spherical coupling for Cartesian kernels

Equivariant interatomic potentials extend atomistic simulations beyond the length and timescales that electronic-structure methods can reach, but their computational cost limits the local correlations they can represent in practice, and therefore the accuracy they achieve. FlashCart attacks that cost directly: it expresses tensor products in independent Cartesian components, symbolically simplifies the products and their derivatives, and fuses the result into generated GPU kernels. The paper reports that these kernels often outperform optimized spherical counterparts.

The architecture recursively builds equivariant features and compresses them to a fixed width at each step, which is what makes higher-order correlations affordable. The paper's central claim is not that Cartesian math is cheaper in operation count; the authors state that operation counts show no arithmetic advantage over sparse spherical coupling. The runtime advantage comes from the complete kernel implementations, including memory access and aggregation of edge contributions.

Correlation order buys accuracy more cheaply than width

The paper compares increasing correlation order against increasing channel width, message-passing depth, or maximum tensor rank. Accuracy is measured by normalized error, which averages energy and force mean absolute errors after each is divided by a fixed reference value. On logarithmic axes, the fitted slope describes the relative change in error for a relative change in runtime or memory.

Across the tested configurations, runtime slopes range from -6.5 to -2.6 for correlation order and from -1.5 to -0.8 for the other architectural choices, so raising correlation order gives the largest error reduction per added runtime. With the rest of the architecture fixed, raising the correlation order from 2 to 5 reduces normalized error by 19 to 60 percent, while runtime rises by 4 to 37 percent and peak GPU memory by at most 1.7 percent. Most models on the observed accuracy-efficiency frontier use correlation orders of 4 and 5.

Higher tensor ranks amplify the gain from correlation order

The benefit of correlation order depends on the tensor ranks allowed. When both direction tensors and intermediate tensors in the correlation ladder are limited to rank 1, raising the order from 2 to 5 lowers energy error by 20 percent and increases runtime by 4 percent. At maximum ranks of 4, the same change lowers energy error by 58 percent and increases runtime by 37 percent. Higher tensor ranks therefore allow larger relative error reductions from increasing order, at greater computational cost.

The relative reductions in energy and force errors from increasing correlation order generally become smaller with more message-passing blocks or a higher maximum propagated rank, but vary little with channel width. Because increasing order also adds parameters, the authors also compare models at roughly fixed size near 0.93M parameters by reducing channel width as order rises. At maximum edge and product rank 3, an order-5 model with 64 channels reaches a force MAE of 8.29 meV per angstrom at 3.94 microseconds per atom, compared with 11.53 meV per angstrom at 5.49 microseconds per atom for order 2 with 108 channels.

Generated kernels beat spherical baselines at high order

In edge-convolution microbenchmarks with 1024 nodes, 128 feature channels, and 128 neighbors per node, FlashCart is on average 1.6x, 1.1x, and 1.9x faster than cuEquivariance for the forward pass, vector-Jacobian product, and vector-Jacobian product of the vector-Jacobian product, respectively, across a 30-configuration grid. Against OpenEquivariance it is 1.9x, 2.1x, and 3.1x faster on average, although OpenEquivariance is faster in some configurations at ranks 1 and 2. In the vector-Jacobian product of the vector-Jacobian product phase, FlashCart also uses about half the peak GPU memory of either baseline.

For node-local products with 4096 nodes, 128 feature channels, and maximum tensor ranks 3 and 4, the gap widens with order. At rank 3 and order 4, forward evaluation takes 0.40 milliseconds for path-reduced products and 2.07 milliseconds for the full ladder, compared with 280.61 milliseconds for cuEquivariance, making them roughly 695x and 136x faster. FlashCart reaches order 6 at every tested rank, while some cuEquivariance contractions could not be constructed or their kernels failed to launch. The authors note that these runtime differences also reflect how correlations are constructed: cuEquivariance assigns learned weights to complete symmetric contractions, while FlashCart applies weighted products successively and retains only ranks up to the chosen maximum after each step.

A 5.6M-parameter model beats a 189M-parameter transformer

On SPICE-MACE-OFF, the largest FlashCart model has 5.6M parameters and reaches energy and force MAEs of 0.072 plus or minus 0.002 meV per atom and 1.99 plus or minus 0.04 meV per angstrom. Evaluating energies and forces for 1000 diamond atoms takes 21.5 microseconds per atom and 5.2 MiB per atom of peak GPU memory. This improves on PET-L's errors of 0.106 meV per atom and 2.52 meV per angstrom, despite far fewer parameters and roughly ten times faster running. It also approaches DPA4-Plus's force MAE of 1.80 meV per angstrom while improving on its energy MAE of 0.091 meV per atom, with 7.7x faster evaluation and 3.1x less memory.

The timings use each model's own cutoff radius. In diamond, FlashCart processes 158 edges per atom compared with PET-L's 70, giving it an approximately 23x runtime advantage per edge. The authors state that this normalization accounts for the different numbers of neighbors, although the comparison still includes architectural differences. To isolate the contribution of the generated kernels, the paper compares them with the corresponding PyTorch implementations of tensor products and edge convolutions using the same 5.6M model checkpoint. With both implementations compiled, the generated kernels reduce inference time from 175.5 to 21.5 microseconds per atom and training-benchmark runtime from 442.6 to 82.2 microseconds per atom, making them 8.2x and 5.4x faster, respectively.

A 2.3M model stays competitive across 85 elements

The paper also evaluates a 2.3M-parameter FlashCart model trained on MAD, which contains fewer than 100,000 structures spanning 85 elements, against the 2.8M-parameter PET-MAD potential and other machine-learned interatomic potentials. FlashCart has lower force MAEs than PET-MAD on all six benchmarks reporting force errors and lower energy MAEs on six of seven benchmarks. On the full MAD test set, FlashCart reaches energy and force MAEs of 11.4 plus or minus 0.1 meV per atom and 63.9 plus or minus 0.5 meV per angstrom, compared with PET-MAD's 14.7 meV per atom and 72.9 meV per angstrom.

For energy and force evaluation on 1000 diamond atoms, FlashCart and PET-MAD take 8.4 and 45.1 microseconds per atom and use 2.0 and 3.3 MiB per atom of peak GPU memory, respectively. The authors describe these results as establishing FlashCart as a competitive architecture for general-purpose potentials, achieving accuracy comparable to state-of-the-art models across diverse materials and molecular environments.

Water simulations hold energy but overestimate density

To demonstrate applicability to atomistic simulations, the authors study bulk liquid water with the five models from the accuracy comparison, simulating 256 water molecules in LAMMPS with Kokkos acceleration. All models show energy drifts between -0.002 and 0.008 meV per atom per nanosecond, with energy fluctuations around the mean having standard deviations of 0.012 to 0.013 meV per atom and maximum absolute deviations below 0.056 meV per atom. In NPT simulations, the models predict water densities that decrease with temperature and remain above experimental values, consistent with the overestimation reported for MACE-OFF and some related variants trained at the same electronic-structure level.

The paper also locates a cost of compilation. Without model compilation, replacing PyTorch tensor products and edge convolutions with the generated Triton kernels makes the same model 26.8x faster for inference and 23.8x faster in the training benchmark. Compiling the model with Triton kernels further reduces runtime but increases peak GPU memory during training. Comparing memory allocation on 512 atoms with CUDA graphs disabled, GPU allocation after the energy and force errors are combined into the loss increased by 5.21 GiB with compilation and 5.72 GiB without it; during parameter-gradient computation it peaked a further 2.26 GiB above the preceding value with compilation, compared with 0.17 GiB without it. The additional peak memory therefore arises during parameter-gradient computation when the model is compiled, with the same generated Triton kernels used in both cases.

Choose correlation order before adding width or depth

The decision-relevant result is a priority ordering rather than a single accuracy number. When a fixed compute or memory budget is available, the paper's fitted slopes and controlled comparisons indicate that raising correlation order reduces normalized error more per unit of added runtime than increasing channel width, message-passing depth, or maximum tensor rank. The gain is largest when tensor ranks are also allowed to grow: at maximum ranks of 4, raising order from 2 to 5 cuts energy error by 58 percent for 37 percent more runtime, whereas at rank 1 the same order increase cuts energy error by only 20 percent for 4 percent more runtime.

Two boundaries qualify that ordering. First, the relative benefit shrinks as message-passing depth or maximum propagated rank increases, so the advantage of order is not uniform across architectures. Second, the runtime and memory comparisons use each model's own cutoff radius and include architectural differences, and the water simulations show densities above experimental values, so the accuracy-efficiency frontier should be read as measured on the paper's benchmarks rather than as a general guarantee. Within those limits, the paper's clearest guidance is that correlation order, not width or depth, is the first architectural knob to turn when the goal is lower energy and force errors at a given cost.

Sources used in this answer