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

cs.LGOct 5, 2026

FlashCart: Fast Cartesian Tensor Products for Equivariant Interatomic Potentials

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

Organizations: INM – Leibniz Institute for New Materials · Saarland University · German Research Center for Artificial Intelligence (DFKI) · University of Stuttgart · NEC Laboratories Europe

Abstract

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.

Figures & tables

Figure 1: Accuracy–efficiency tradeoff evaluated on SPICE-MACE-OFF. (a,c) Normalized error as a function of evaluation time and peak GPU memory, with FlashCart models on the measured accuracy–efficiency frontier. Local evaluation costs in (a,c) use a

1000

-atom diamond cell. (b,d) Evaluation time and peak GPU memory per atom as a function of system size, reaching a plateau at smaller system sizes for larger FlashCart models.

\dagger

All timings use an NVIDIA RTX PRO 6000 Blackwell (96 GB), except for eSEN literature values measured on an NVIDIA A100 (80 GB) [ 23 ] .

Figure 3: Runtime and memory of edge convolutions and node-local tensor products. (a) Results for edge convolutions with

1024

nodes and

128

edges per node as a function of the maximum tensor rank. (b) Results for node-local tensor products with

4096

nodes as a function of the correlation order. Solid FlashCart curves show path-reduced products that sum contributions within the kernel, with a scalar weight per path and feature channel. Dashed curves show the full ladder, including the linear maps that mix paths and channels between steps. cuEquivariance evaluates symmetric contractions with a learned weight for each contraction term and feature channel. Both implementations use the indicated maximum rank for their input and output tensors. Grey crosses denote cases where contraction construction was not completed, or a kernel failed to launch (Appendix G ). Dashed extensions indicate trends beyond the available measurements. All results are averaged over three independent runs and use

128

feature channels and single precision.

Figure 4: Accuracy–efficiency tradeoffs across architectural choices. Normalized error on SPICE-MACE-OFF is shown against runtime (a) and peak GPU memory (b). Filled symbols denote

\nu=5

, while unfilled ones denote lower correlation orders. The insets show the effect of varying one architectural choice at a time, with the maximum edge and product ranks varied together. Larger negative fitted slopes indicate larger relative error reductions per relative increase in runtime or memory. We average errors over three independent runs and measure runtime and memory for energy and force evaluation on

4096

diamond atoms using one trained model per configuration.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure C1: Independent Cartesian components and the cost of generated tensor products. (a) Stored components for ranks

2

and

3

are blue, with light blue entries obtained by index symmetry and grey entries from tracelessness. (b) Full, symmetric, and independent component counts as a function of tensor rank. (c) Additions and multiplications for distinct couplings up to rank

5

, ordered by

J=l_{1}+l_{2}+l_{3}

, before compiler transformations. Independent components reduce Cartesian operation counts to values closer to sparse Clebsch–Gordan coupling. (d) Lines of generated source code for the forward convolution as a function of maximum tensor rank.

		Cartesian		Spherical
Maximum rank	Couplings	Full	Independent	Sparse CG
3	15	1,299	615	609
4	27	13,802	2,774	2,141
5	42	79,128	7,505	4,951

Appendix

Table C1: Operation counts for Cartesian and spherical tensor products.

Figure D1: Memory layout of Cartesian tensor features. Each row

X_{r,:}

stores tensor ranks in contiguous blocks, with

\operatorname{off}(l)

giving the offset of the block for rank

l

within the row. The expanded rank-

2

block shows how each stored Cartesian component contains

P_{2}

coupling paths and

d_{X}

feature channels. The darker cells highlight the contiguous channels of one component–path pair.

Figure 8

Figure E1: Model architecture and fixed-width correlation ladder. (a) Each message-passing block combines a convolution with the product ladder. Blue bars show tensor features grouped by rank

l

, starting at

l=0

. Light grey lines show shared geometric inputs and blockwise scalar features combined in the readout to give

E

. The output of each block is also passed to the next block. (b) The same atomic basis

\bm{A}_{u}

is reused at every tensor product. Grey bars show the intermediate tensor product outputs with the retained coupling paths, compressed by

W^{(q)}

d

channels per retained rank, shown in blue. The output of each step enters the next product and also contributes to the final sum. Bar heights schematically indicate the number of tensor components. We skip residual updates, normalization, and nonlinear gates for simplicity.

Configuration	$T$	$\nu$	$d$	$l_{\max}^{\mathrm{prop}}$	$l_{\max}^{\mathrm{edge}}=l_{\max}^{\mathrm{prod}}$	Product
Models shown in Figures 4 and G1
Default values	2	2	64	1	3	Ladder
Explored values	$1$ – $4$	$2$ – $5$	$32,64,96,128$	$0$ – $3$	$1$ – $4$	Ladder
Models shown in Figure 1 and Table G1
FlashCart-69k	2	3	20	2	3	Reduced
FlashCart-0.5M	3	5	32	2	3	Ladder

Appendix

Table F1: Model configurations.

T

\nu

, and

d

denote the number of message-passing blocks, correlation order, and channel width, respectively, while

l_{\max}^{\mathrm{edge}}=l_{\max}^{\mathrm{prod}}

gives the shared maximum rank of the edge and product tensors. For the models in Figures 4 and G1 , default values apply to parameters that are held fixed. The labels “Ladder” and “Reduced” denote the correlation ladder and path-reduced products described in Appendix E .

Hyperparameter	Figures 4 and G1	Figure 1 and Table G1	Tables 1 and G2
Cutoff radius	6.0	6.0	5.0
Maximum epochs	150	400	400
Warmup epochs	20	20	20
Cosine-decay epochs	130	380	380
Target batch size	128	128	192
Number of random seeds	3	3	3

Appendix

Table F2: Training parameters. We use the target batch size and training-set averages to determine how many atoms and edges each batch can contain, so the number of structures per batch can vary. For distributed training, we divide the target batch size across GPUs. EMA denotes the exponential moving average of model parameters.

Figure G1: Effect of correlation order across model architectures. Results for the models trained on SPICE-MACE-OFF and shown in Figure 4 are presented as energy MAE (top), force MAE (middle), and runtime (bottom) against correlation order

\nu

. From left to right, the columns show results for different maximum edge and product ranks (varied together), maximum propagated ranks, channel widths

d

, and numbers of message-passing blocks

T

. Unless varied in a column, models use maximum edge and product ranks of

l_{\mathrm{max}}^{\mathrm{prod}}=l_{\mathrm{max}}^{\mathrm{edge}}=3

, maximum propagated rank

l_{\mathrm{max}}^{\mathrm{prop}}=1

d=64

, and

T=2

. Percentages show the relative changes in error or runtime as

\nu

increases from

2

5

. Prediction errors are averaged over three independent training runs. Shaded bands show one standard deviation but are not visible at the plotted line width. Runtimes are measured for energy and force evaluation on

4096

diamond atoms using one trained model per configuration.

Model		Params.	PubChem	DES Mon.	DES Dim.	Dipept.	Solv. AA	Water	QMugs	Mean	NE	$\mu$ s/atom	MiB/atom
MACE-OFF23-S [ 31 ]	E	694k	1.41	1.04	0.98	0.84	1.60	1.67	1.03	1.224	1.000	3.6	0.83
	F		35.68	17.63	16.31	25.07	38.56	28.53	41.45	29.03	1.000	3.6	0.83
MACE-OFF23-M [ 31 ]	E	1.4M	0.91	0.63	0.58	0.52	1.21	0.76	0.69	0.757	0.593	6.7	1.48
	F		20.57	9.36	9.02	14.27	23.26	15.27	23.58	16.48	0.593	6.7	1.48
MACE-OFF23-L [ 31 ]	E	4.7M	0.88	0.59	0.54	0.42	0.98	0.83	0.45	0.670	0.490	14.9	3.79
	F		14.75	6.58	6.62	10.19	19.43	13.57	16.93	12.58	0.490	14.9	3.79

Appendix

Table G1: Energy and force errors on SPICE-MACE-OFF. Results are shown separately for the seven test subsets for the models in Figure 1 . Rows give energy (E) MAEs in meV/atom and force (F) MAEs in meV/Å. We also report the arithmetic mean of the seven subset MAEs, separately for energy and forces. Normalized error (NE) combines these means as defined in Appendix F . FlashCart results show the mean and standard deviation across three independent training runs. We measure runtime and peak GPU memory for energy and force evaluation on a periodic diamond cell containing

1000

atoms using an NVIDIA RTX PRO 6000 Blackwell GPU (96 GB). The lowest energy and force MAEs and normalized errors are shown in bold, taking the reported standard deviations into account.

Figure G2: Runtime and memory with generated kernels and model compilation. Runtime (a) and peak GPU memory (b) per atom are compared for the same

5.6

M-parameter FlashCart model during inference (left) and training (right). Baseline denotes the PyTorch implementation of tensor products and edge convolutions, while Triton denotes the generated kernels. We evaluate both implementations with and without model compilation, using the fastest measured compilation mode. The factors above the bars show the runtime and memory improvements relative to the uncompiled PyTorch baseline. Inference evaluates energies and forces on

1000

diamond atoms. Training evaluates the energy and force loss and its parameter gradients on

512

atoms, excluding optimizer updates. Runtime error bars show standard deviations over repeated evaluations.

Subset	FlashCart		PET-MAD
MC2D	12.56 $\bm{\pm}$ 0.30 $/$	57.2 $\bm{\pm}$ 1.2	15.54 $/$	72.7
MC3D	7.02 $\bm{\pm}$ 0.07 $/$	35.1 $\bm{\pm}$ 0.4	9.74 $/$	46.3
MC3D Clusters	28.41 $\bm{\pm}$ 0.54 $/$	127.6 $\bm{\pm}$ 1.0	43.06 $/$	148.3
MC3D Random	68.09 $\pm$ 0.71 $/$	316.2 $\pm$ 2.0	63.87 $/$	249.6
MC3D Rattled	10.04 $\bm{\pm}$ 0.20 $/$	142.2 $\pm$ 1.1	11.73 $/$	140.3
MC3D Surfaces	10.58 $\bm{\pm}$ 0.07 $/$	88.7 $\bm{\pm}$ 1.0	13.61 $/$	102.6

Appendix

Table G2: Energy and force errors on the full MAD test set. Each energy/force pair gives the energy MAE in meV/atom followed by the force MAE in meV/Å. FlashCart results show the mean and standard deviation across three models trained with different random seeds on a fixed data split. The overall MAEs are computed over the full test set. The lowest energy and force MAEs for each dataset are shown in bold.

Figure G3: Energy conservation and density of bulk liquid water. (a) Total energy per atom during NVE dynamics, shown relative to its mean. We add a constant vertical offset to keep the curves distinguishable. (b) Density from NPT simulations at

1

bar, with shaded bands showing the standard error estimated by block averaging over

50

ps intervals. The bands are narrower than the plotted lines and are therefore not visible. The dashed curve shows experimental densities at

1

bar [ 36 ] . All results use one trained model per configuration and the final

1

ns of each trajectory.

Explore similar work

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials
Oct 1, 2026Laura Zichi, Gil Harari, Chuin Wei Tan +6Interatomic Potentials Molecular Dynamics
Harvard University · Robert Bosch LLC Research and Technology Center
Geometric Algebra Meets Cartesian Tensors: Higher-Order Equivariance for Interatomic Potentials
Jun 28, 2026Can Polat, Erchin Serpedin, Mustafa Kurban +1Interatomic Potentials Cosine Similarity
Dept. of Electrical & Computer Eng. Texas A&M Univ. College Station, TX, USA · Dept. of Prosthetics & Orthotics College of Science & Eng. Ankara Univ., Ankara, Turkey · Hamad Bin Khalifa Univ. and Texas A&M Univ. at Qatar Doha, Qatar
Scalar-pathway fidelity improves physical accuracy in short-range equivariant interatomic potentials
Jun 14, 2026Jia Bi, Alin Marin Elena, Samuel PinillaInteratomic Potentials Molecular Dynamics
Science and Technology Facilities Council, Harwell Campus, Didcot, OX11 0QX, Oxfordshire, United Kingdom. · Science and Technology Facilities Council, Keckwick Lane, Daresbury, WA4 4AD, Cheshire, United Kingdom. · Diamond Light Source, Harwell Science and Innovation Campus,2026 Didcot, OX11 0QX, Oxfordshire, United Kingdom.

Oct 1, 2026, cs.LGJ/K move · Enter open · S save

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

Laura Zichi, Gil Harari, Chuin Wei Tan, Marc L. Descoteaux+5

Harvard University · Robert Bosch LLC Research and Technology Center

Interatomic Potentials Molecular Dynamics High-Performance Computing Rotation Equivariance

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

Geometric Algebra Meets Cartesian Tensors: Higher-Order Equivariance for Interatomic Potentials

Scalar-pathway fidelity improves physical accuracy in short-range equivariant interatomic potentials