PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashCart: Fast Cartesian Tensor Products for Equivariant Interatomic Potentials

Viktor Zaverkin, Payman Goodarzi, Sergey V. Sukhomlinov, Davit Hovhannisyan, Roland Aydin, Martin H. Müser, Mathias Niepert

arXiv:2610.06409v1cs.LG

TL;DR

Equivariant architectures can represent higher-order local correlations, but their computational cost limits practical correlation order and achievable accuracy. FlashCart combines generated Cartesian tensor-product kernels with a recursive fixed-width equivariant architecture. On SPICE-MACE-OFF, its 5.6M-parameter model improves on PET-L’s energy and force errors while running approximately ten times faster.

  • Problem

    The computational cost of equivariant architectures limits the local correlation orders practical models can represent and their achievable accuracy.

    Source: §1 · p. 1

  • Method

    FlashCart symbolically generates Cartesian tensor-product kernels and recursively builds higher-order equivariant features while compressing them to fixed width at each step.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    The 5.6M-parameter FlashCart model reaches energy MAE 0.072±0.002 meV/atom and force MAE 1.99±0.04 meV/Å, improving on PET-L’s errors while running approximately ten times faster.

    Source: §6 · p. 9

  • Takeaways & Limitations

    On SPICE-MACE-OFF, increasing correlation order improves accuracy more efficiently than increasing width, depth, or tensor rank.

    Source: §1 · p. 2

    Fixed channel widths and truncated tensor ranks limit representable correlations, while retaining higher ranks increases generated-kernel code size and compilation cost.

    Source: §7 · p. 10

Abstract

from arXiv · show

Machine-learned interatomic potentials extend atomistic simulations beyond the length- and timescales accessible to electronic-structure methods. However, the computational cost of equivariant architectures limits the local correlations they can represent in practice and therefore their achievable accuracy. Here we introduce FlashCart, which makes higher-order correlations affordable by combining generated GPU kernels with an architecture that recursively builds equivariant features and compresses them to a fixed width at each step. We express tensor products in independent Cartesian components and symbolically simplify them and their derivatives, producing fused kernels that often outperform optimized spherical counterparts. We then show that increasing correlation order improves accuracy more efficiently than increasing width, depth, or tensor rank. On SPICE-MACE-OFF, FlashCart models advance the measured accuracy-efficiency frontier: a model with $5.6$ million parameters achieves lower energy and force errors and $10\times$ faster inference than a transformer with $189$ million parameters.

1 INTRODUCTION

Equivariant interatomic potentials can represent increasingly complex local correlations through repeated tensor products, but higher order raises evaluation cost. FlashCart combines Cartesian tensor products, generated GPU kernels, and fixed-width recursive features to make higher-order correlations more affordable.

1 INTRODUCTION

  • Equivariant models encode rotationally transforming atomic features and can build higher-order correlations through repeated products, but those correlations increase evaluation cost.Making higher-order correlations affordable offers a complementary route to improved accuracy–efficiency tradeoffs while retaining explicit equivariance.Source: §1 · p. 1
  • FlashCart models lie on the measured SPICE-MACE-OFF accuracy–efficiency frontier, with the figure comparing normalized error against evaluation time and peak GPU memory.The local evaluation costs use a 1000-atom diamond cell.Source: Figure 1 · p. 2, §1 · p. 2
  • FlashCart expresses tensor products and derivatives in independent Cartesian components and generates fused GPU kernels for node-local products and neighbor aggregation.The kernels often outperform optimized spherical counterparts in the reported benchmarks.Source: §1 · pp. 1–2
  • FlashCart recursively builds higher-order correlations by mixing paths and channels back to fixed width, with each added order requiring one product and one linear map.At fixed width and maximum tensor rank, correlation cost grows linearly with order.Source: §1 · p. 2

2 BACKGROUND

The background describes how interatomic potentials build invariant energies from local neighborhoods while preserving the required transformation behavior. Equivariant convolutions aggregate neighbor information, and repeated node-local products build higher-order correlations.

2 BACKGROUND

  • In the absence of external fields, the potential energy is invariant to translations, rotations, reflections, and permutations of atoms of the same species, while forces transform as vectors.Deriving forces from scalar energy yields a conservative force field; training on energies and forces also requires differentiating force evaluation with respect to model parameters.Source: §2 · p. 2
  • Equivariant architectures combine neighbor-aggregating convolutions with repeated node-local products; when neighbor contributions depend only on the central atom and that neighbor, order-q products have body order at most q + 1.The recurrence acts independently at each atom, unlike convolution, which sums contributions from neighbors.Source: §2 · pp. 2–3
  • Rank-l irreducible Cartesian tensors are symmetric and trace-free, have 2l + 1 independent components, and represent the same O(3) irreducible representation as degree-l real spherical tensors.Their natural parity is (-1)^l.Source: §2 · p. 3
  • For admissible paths, spherical and Cartesian products express the same equivariant coupling in different coordinates, so representation choice changes evaluation arithmetic and memory access rather than expressive power.The Cartesian coupling contracts index pairs, symmetrizes remaining indices, and removes traces.Source: §2 · p. 3

3 GENERATING CARTESIAN TENSOR PRODUCTS

The section generates Cartesian tensor-product kernels from independent components and symbolically derives their gradients. The independent-component expressions reduce operation counts, while symbolic differentiation supports force training.

3 GENERATING CARTESIAN TENSOR PRODUCTS

  • The generator selects 2l + 1 independent components per rank and reconstructs all 3l tensor components symbolically, while model evaluation reads and writes only independent components.Trace relations determine the remaining symmetric-tensor components.Source: §3 · p. 3
  • For each admissible path, the generator forms the Cartesian coupling, contracts index pairs, symmetrizes and removes traces, then emits expressions for the independent output components.Paths related by exchanging input ranks share an expression with exchanged operands.Source: §3 · pp. 3–4
  • Symbolic differentiation produces input VJPs and VJP-of-VJP formulas for weighted bilinear products, which are embedded in Triton kernels before compilation.A PyTorch forward implementation with automatic differentiation is generated to verify the derivatives.Source: §3 · p. 4
  • For ranks up to 3, independent-component expressions reduce the summed operation count across distinct couplings from 1,299 to 615, near the 609 operations required by sparse Clebsch–Gordan coupling.Through ranks 4 and 5, the expressions use 1.3–1.5 times as many operations as spherical counterparts.Source: §3 · p. 4

4 EXECUTING CARTESIAN TENSOR PRODUCTS

The execution strategy uses generated Triton kernels for node-local products and edge convolutions, with receiver-sorted accumulation as the default forward design. Derivative kernels reuse the generated expressions, while atomic gradient accumulation does not guarantee bitwise reproducibility.

4 EXECUTING CARTESIAN TENSOR PRODUCTS

  • The kernels support node-local products and edge convolutions, and the feature layout groups entries by tensor rank, independent component, coupling path, and channel.Increasing rank blocks are ordered so retaining lower ranks needs no feature rearrangement or copy.Source: §4 · p. 4
  • The receiver-sorted forward kernel assigns each program a receiver, channel tile, and output-rank group, accumulates incoming-edge contributions locally, and stores each result once.This avoids forward atomic additions and intermediate edge-message tensors.Source: §4 · pp. 4–5
  • The VJP reuses receiver-sorted edges and generated input VJPs, while the VJP-of-VJP accumulates each receiver’s incoming-edge contributions before storing the result.Shared input gradients use atomic additions, while weight gradients are stored directly.Source: §4 · p. 5
  • FlashCart keeps edges receiver-sorted, but its atomic gradient accumulation does not guarantee bitwise reproducibility.The cited comparison describes Sobek as ordering edges by sender for deterministic sender-feature gradient accumulation.Source: §4 · pp. 5–6

5 MODEL ARCHITECTURE

FlashCart’s node-local ladder builds higher-order correlations while keeping a fixed channel width. It uses channel-diagonal Cartesian products followed by learned mixing, with tensor-rank cutoffs controlling what is retained at each step.

5 MODEL ARCHITECTURE

  • The ladder builds correlations recursively: each additional order adds one pairwise product and one learned linear map, without a separate higher-order operator.The recurrence reuses the same atomic basis at every step and can be viewed as a tensor-train factorization of polynomial coefficients for retained coupling-path sequences.Source: §5 · p. 6
  • At each step, the cutoff lq sets the highest output tensor rank retained; the final block returns only scalars.The resulting degree-q polynomial is in A, but because normalization and nonlinear gating make A neighborhood-dependent, q does not bound the number of interacting neighbors.Source: §5 · p. 6
  • Each ladder step couples same-index channels, then a learned map mixes paths and channels to return d channels per retained tensor rank.The map uses Pq d^2 parameters per step, compared with Pq d^3 for a dense map combining products across all pairs of input channels.Source: §5 · p. 6
  • With maximum tensor ranks fixed, parameter count and per-atom computational cost grow quadratically with channel width and linearly with correlation degree.The full ladder retains separate path contributions for learned mixing, while path-reduced products sum weighted paths within each step and mix channels only after the recurrence.Source: §5 · p. 6

6 RESULTS AND DISCUSSION

FlashCart’s higher-order correlations improve accuracy efficiently, and its models compare favorably with baselines in kernel speed and benchmark accuracy. The reported results span kernel microbenchmarks, SPICE-MACE-OFF, and MAD evaluations.

Kernel microbenchmarks

  • Across 30 edge-convolution configurations, FlashCart is 1.6×, 1.1×, and 1.9× faster than cuEquivariance for forward, VJP, and VJP-of-VJP, respectively.Against OpenEquivariance, it is 1.9×, 2.1×, and 3.1× faster on average, though OpenEquivariance is faster in some rank-1 and rank-2 configurations; FlashCart uses about half the peak memory in VJP-of-VJP.Source: §6 · p. 7, Figure 3 · p. 7
  • At rank 3 and order 4, node-local forward evaluation takes 0.40 ms for path-reduced products and 2.07 ms for the full ladder, versus 280.61 ms for cuEquivariance.FlashCart reaches order 6 at every tested rank, while some cuEquivariance contractions could not be constructed or their kernels failed to launch.Source: §6 · p. 7, Figure 3 · p. 7

Efficiency of higher-order correlations

  • Raising correlation order from ν = 2 to 5 reduces normalized error by 19–60%, with runtime increasing by 4–37% and peak GPU memory by at most 1.7%.Runtime slopes range from −6.5 to −2.6 for correlation order, compared with −1.5 to −0.8 for other architectural choices; most observed frontier models have ν = 4 or 5.Source: §6 · p. 8, Figure 4 · p. 8
  • At approximately 0.93M parameters and rank 3, order 5 with 64 channels reaches a force MAE of 8.29 meV/Å at 3.94 µs/atom, versus 11.53 meV/Å at 5.49 µs/atom for order 2 with 108 channels.Across maximum edge and product ranks 1–4, increasing order while reducing width lowers force error by 9–29% and runtime by 10–34% at approximately the same parameter count.Source: §6 · p. 8

Model accuracy and computational cost

  • 5.6M-parameter FlashCart reports energy and force MAEs of 0.072±0.002 meV/atom and 1.99±0.04 meV/Å, improving on PET-L’s 0.106 meV/atom and 2.52 meV/Å while running approximately 10× faster.It also approaches DPA4-Plus’ force MAE of 1.80 meV/Å while improving on its energy MAE of 0.091 meV/atom, with 7.7× faster evaluation and 3.1× less memory.Source: §6 · p. 9

Kernel microbenchmarks

  • Generated kernels reduce inference time from 175.5 to 21.5 µs/atom and training-benchmark runtime from 442.6 to 82.2 µs/atom, making them 8.2× and 5.4× faster, respectively.The comparison uses the same 5.6M-parameter checkpoint and compiled implementations; compilation increases peak GPU memory during training by 24%.Source: §6 · p. 9

Accuracy across materials and molecules

  • 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’s energy and force MAEs are 11.4 ± 0.1 meV/atom and 63.9 ± 0.5 meV/Å, versus PET-MAD’s 14.7 meV/atom and 72.9 meV/Å.Source: §6 · p. 9, Table 1 · p. 9

7 CONCLUSIONS

FlashCart combines symbolic Cartesian tensor products and generated GPU kernels with a fixed-width correlation ladder. The conclusion reports accuracy–efficiency gains and identifies representation, width, rank, and compilation constraints.

7 CONCLUSIONS

  • FlashCart combines symbolic Cartesian tensor products, generated GPU kernels, and a correlation ladder whose cost grows linearly with correlation order at fixed width and maximum tensor rank.The generated kernels often outperform optimized spherical counterparts in the reported benchmarks and reduce evaluation and training-loss differentiation time versus a compiled PyTorch implementation of the same model.Source: §7 · pp. 9–10
  • FlashCart models lie on the SPICE-MACE-OFF accuracy–efficiency frontier, and higher correlation order improves the accuracy–runtime tradeoff more than width, depth, or tensor rank in tested configurations.Trading width for higher order also lowers force error and runtime at approximately fixed parameter counts; MAD results show competitive accuracy across diverse materials and molecular environments.Source: §7 · p. 10
  • Natural-parity products, fixed channel widths, and truncated ranks constrain the ladder, while higher ranks increase generated-code size and compilation cost.The authors note that opposite-parity products can be added using Levi–Civita contractions and parity labels.Source: §7 · p. 10

ETHICS STATEMENT

  • The work develops computational methods for atomistic simulations using existing datasets and involves neither human participants nor personal data.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

APPENDICES

The related work situates FlashCart among recursive product-basis methods, channel-mixing and coefficient-factorization approaches, Cartesian tensor representations, and GPU kernel libraries for spherical tensors.

A RELATED WORK

  • ACE and MACE construct many-body features from an atomic basis, while recursive evaluation reuses lower-order intermediates; NICE and TRACE limit or mix features between correlation orders.Source: A RELATED WORK · p. 15
  • Learned channel mixing and channel-diagonal products appear in Clebsch–Gordan networks, Cormorant, Allegro, and factorizations including tensor-reduced ACE and GRACE [2, 3, 5, 29, 40–43].Source: A RELATED WORK · p. 15
  • The paper follows ICTP's higher-rank irreducible Cartesian framework, alongside Cartesian potentials including TensorNet, HotPP, Cartesian-3j, and CarNet [8, 9, 19, 24–26, 44–49].Source: A RELATED WORK · p. 15
  • cuEquivariance, OpenEquivariance, Sobek, and an Allegro Triton implementation accelerate spherical tensor operations through sparse couplings, fusion, or reordered accumulation [15–18].Source: A RELATED WORK · p. 15

B BACKGROUND

The background describes equivariant convolutions and recursive node-local products as ways to build many-body correlations from neighboring atoms. It defines the tensor couplings and the rank, body-order, and parity scope relevant to the method.

B BACKGROUND

  • Equivariant architectures vary in body order, retained angular ranks, channel count, and depth; a radial envelope makes edge contributions vanish smoothly at the cutoff, preserving twice-differentiability of energy with respect to positions.Source: B BACKGROUND · p. 15
  • The convolution combines sender features with direction tensors along admissible rank-coupling paths, using distance-dependent maps for radial weighting and channel mixing.The output-rank selection and equivariant linear maps determine which ranks are retained at each recursive product step.Source: B BACKGROUND · p. 16
  • A degree-q recursive product has maximum body order q+1 when each neighbor contribution depends only on the central atom and that neighbor; stacking T message-passing layers extends interaction range to Trmax.Source: B BACKGROUND · p. 16
  • For each tensor rank, normalized direction tensors are formed from symmetrized repeated outer products of the unit direction, with identity terms removing traces while preserving rank.The prefactor normalizes the full contraction with the unit direction.Source: B BACKGROUND · p. 16
  • Admissible tensor-product paths satisfy the triangle inequalities, and their equivariant bilinear coupling is unique up to normalization; spherical and Cartesian products therefore encode the same coupling in different coordinates.Source: B BACKGROUND · p. 16
  • The work restricts products to natural-parity outputs; odd-J couplings require Levi–Civita contractions and produce opposite-parity outputs, though the symbolic construction can be extended to include them.Source: B BACKGROUND · p. 16

C GENERATING CARTESIAN TENSOR PRODUCTS

The appendices describe symbolic Cartesian product generation and differentiation, then detail the data layout and GPU execution strategies for node-local products and receiver-sorted convolutions.

C GENERATING CARTESIAN TENSOR PRODUCTS

  • The generator constructs each Cartesian coupling from contracted inputs, symmetrizes the remaining indices, removes traces, and applies a path-dependent factor to match the product normalization.The factor accounts for the difference between averaging over permutations and summing distinct index placements.Source: C GENERATING CARTESIAN TENSOR PRODUCTS · pp. 16–17
  • Figure C1 compares component counts, per-coupling operation counts before compiler transformations, and generated forward-convolution source lines as tensor rank increases.Its caption states that independent components bring Cartesian operation counts closer to sparse Clebsch–Gordan coupling.Source: Figure C1 · p. 17
  • For a weighted bilinear product, the explicit VJP computes gradients for both inputs and the scalar weight; bilinearity lets the VJP-of-VJP combine contributions into linear-map evaluations.The weight gradient is the inner product of the incoming output gradient with the unweighted output.Source: C GENERATING CARTESIAN TENSOR PRODUCTS · pp. 17–18
  • The operation-count totals exclude memory accesses and compiler transformations, so they do not directly predict evaluation time; symbolic generation takes seconds, while first-use Triton compilation took a few minutes at rank 4 on the reported device.Source: C GENERATING CARTESIAN TENSOR PRODUCTS · p. 18

D EXECUTING CARTESIAN TENSOR PRODUCTS

  • The kernels support channel-diagonal products with separate dense mixing over channels and paths, as well as path retention or summation and compile-time derivative options.Source: D EXECUTING CARTESIAN TENSOR PRODUCTS · p. 18
  • Feature rows store rank blocks contiguously, with independent components, coupling paths, and channels ordered within each rank; retaining ranks through L selects a row prefix without rearrangement.Figure D1 depicts the rank-block offsets and contiguous channels for a component–path pair.Source: D EXECUTING CARTESIAN TENSOR PRODUCTS · pp. 18–19, Figure D1 · p. 19
  • The receiver-sorted forward kernel accumulates path and incoming-edge contributions locally for each receiver, channel tile, and rank group, avoiding intermediate edge-message tensors and atomic additions for its outputs.The implementation groups output ranks to limit simultaneous sums and uses channel tiles.Source: D EXECUTING CARTESIAN TENSOR PRODUCTS · p. 19
  • The VJP uses receiver-sorted edges and atomic additions for shared input gradients, while the VJP-of-VJP sums receiver contributions to the incoming-gradient derivative; VJP edge chunking parallelizes receivers with many neighbors.The VJP-of-VJP keeps all incoming edges for a receiver in one program because its output requires their sum.Source: D EXECUTING CARTESIAN TENSOR PRODUCTS · pp. 20–21

E MODEL ARCHITECTURE

FlashCart combines equivariant neighbor aggregation with node-local products that recursively build higher-order features. Its product ladder mixes paths and channels back to fixed width, while residual updates and equivariant normalization support each message-passing block.

E MODEL ARCHITECTURE

  • Each block applies residual updates and equivariant RMS normalization after convolution and the product ladder; scalar outputs from all blocks feed an MLP that predicts atomic energies.Normalization combines scalar invariants across ranks and channels without batch statistics; the total energy is the sum of atomic energies.Source: E MODEL ARCHITECTURE · pp. 21–22
  • Neighbor aggregation combines each sender’s rank-specific features with direction tensors and radial weights before mixing coupling-path contributions into receiver features.The radial weights vanish with their first two derivatives at the cutoff, avoiding discontinuities in energy and its first two derivatives as edges enter or leave.Source: E MODEL ARCHITECTURE · p. 22
  • At each atom, the product ladder repeatedly combines basis features and sums contributions through degree ν, building higher-order correlations.The recurrence evaluates products along retained coupling paths and produces contributions of degrees 1 through ν.Source: E MODEL ARCHITECTURE · pp. 22–23
  • The fixed-width ladder accumulates path products without channel mixing, then applies a rank-wise linear map to mix paths and channels into d output channels.For ranks 0–3, the general product intermediates contain 99d independent Cartesian components, reduced to 16d after mixing.Source: E MODEL ARCHITECTURE · pp. 22–23

F EXPERIMENTAL DETAILS

The experiments evaluate FlashCart on SPICE-MACE-OFF and MAD, with architectural comparisons varying correlation order and other model choices. Runtime and memory benchmarks use specified GPU configurations, while training selects models by validation energy and force errors.

F EXPERIMENTAL DETAILS

  • SPICE-MACE-OFF contains 951,005 training and 50,195 test structures, with three runs using random validation sets and a seven-subset test set.The dataset covers ten elements, and each run uses 47,550 validation structures and the remaining 903,455 for training.Source: F EXPERIMENTAL DETAILS · p. 23
  • MAD experiments use its released 76,483/9,566/9,546 train/validation/test split and evaluate the trained model on the 1,562-structure PET-MAD benchmark.The MAD experiments vary model initialization across three runs.Source: F EXPERIMENTAL DETAILS · p. 23
  • Architectural comparisons vary correlation order, width, message-passing depth, and tensor ranks, including a near-fixed-size order comparison with 0.913M–0.945M parameters.The near-fixed-size comparison uses two blocks, propagated rank 1, and edge/product ranks from 1 to 4.Source: F EXPERIMENTAL DETAILS · p. 23
  • Models are selected at the epoch minimizing validation energy MAE plus force MAE, after fixed energy shifts, output scale, and convolution normalization are determined.Training uses EMA parameters and the hyperparameters summarized in Table F2.Source: F EXPERIMENTAL DETAILS · pp. 23–24, Table F2 · p. 24
  • Runtime and memory benchmarks use an NVIDIA RTX PRO 6000 Blackwell GPU, single precision, and precomputed neighbor lists; energy and force tests use 1,000 or 4,096 diamond atoms.The generated-kernel training comparison uses 512 atoms because uncompiled PyTorch runs out of GPU memory at 1,000 atoms; timings exclude optimizer updates and data loading.Source: F EXPERIMENTAL DETAILS · p. 25

G RESULTS AND DISCUSSION

On SPICE-MACE-OFF, the largest FlashCart model compares favorably with DPA4-Plus in accuracy and efficiency, while correlation-order increases yield accuracy gains across architectural settings. Generated kernels substantially reduce runtime, and water simulations show small energy drift alongside density overestimation relative to experiment.

G RESULTS AND DISCUSSION

  • Some cuEquivariance kernel comparisons could not be constructed or launched at higher ranks and orders, including a rank-4, order-4 construction that reached 767 GiB of CPU RAM before failing.The construction failed under a 1,000 GiB memory limit; a rank-3, order-4 VJP-of-VJP kernel also failed with a CUDA out-of-memory error.Source: G RESULTS AND DISCUSSION · p. 25
  • For models with two message-passing blocks, 64 channels, and maximum propagated rank 1, fitted normalized-error-versus-runtime slopes across correlation orders ν=2–5 were −5.0, −6.5, −3.7, and −2.6 at maximum edge/product ranks 1–4.Source: G RESULTS AND DISCUSSION · p. 26
  • 7.7× faster energy and force evaluation and 3.1× less peak GPU memory: the 5.6M-parameter FlashCart model has lower mean energy MAE and combined normalized error than 8.8M-parameter DPA4-Plus.Against the similarly sized 5.1M-parameter DPA4-Air, FlashCart has lower energy and force MAEs on all seven subsets.Source: G RESULTS AND DISCUSSION · p. 26, Table G1 · p. 27
  • 26.8× faster inference and 23.8× faster training: generated Triton kernels replace PyTorch tensor products and edge convolutions without model compilation.With compilation, runtime falls further, but peak training memory increases; the additional peak arises during parameter-gradient computation.Source: G RESULTS AND DISCUSSION · pp. 26, 28, Figure G2 · p. 27
  • Energy drifts by −0.002 to 0.008 meV/atom/ns in bulk-water NVE simulations, with fluctuations of 0.012–0.013 meV/atom standard deviation and maximum deviations below 0.056 meV/atom.In NPT simulations, predicted densities decrease with temperature but remain above experimental values, consistent with reported overestimation for related potentials.Source: G RESULTS AND DISCUSSION · pp. 28–29, Figure G3 · p. 28
Loading 2610.06409v1…