Rotationally equivariant architectures improve the physical fidelity of machine-learned (ML) universal materials models based on graph neural networks (GNNs). However, their scalability is limited by the use of non-trivial, high-precision tensor operations during both training and inference. Adopting the reduced-precision formats of matrix multiplication units (MMUs), e.g., FP8 instead of FP32, provides massive increases in throughput, up to 30x (60x) on NVIDIA Hopper (Blackwell) GPUs, but may lead to drastic loss of accuracy. Here, we introduce tailored 'emulated' low-precision operations under the integer Ozaki scheme to accelerate the 'SO(2)-equivariant' processing layers of GNNs, which dominate the runtime of (1) many state-of-the-art ML interatomic potentials (MLIPs) and (2) most universal ML Hamiltonian (MLH) models. Our framework, FlashSO(2), outperforms TF32, the most accurate native low-precision MMU format, in throughput, accuracy, and memory footprint at once. At the SO(2) layer level, it reaches 40x (7x) the throughput of MALOQ's (UMA's) reference implementation, an MLH (MLIP) model, sustaining up to 236 (232) effective Tflop/s on an NVIDIA GH200, 3.5x the FP32 peak. At the model level, FlashSO(2) helps train MALOQ 3.9x faster than its pure-FP32 version while reducing peak memory by 58% at no loss in accuracy. In molecular dynamics simulations with UMA, the speedup factor is 3.3x, and the peak memory is reduced by 51%. FlashSO(2) thus pushes the speed-accuracy Pareto frontier of equivariant materials models beyond the native low-precision formats and substantially lowers the resources required to exploit their scaling laws.
Figures & tables
Figure 1: Historical trends in ML materials models and data scaling. (a) Training datasets by release year and size on disk. (b) MLIP (blue) and MLH (pink) models by release year and parameter count.
Figure 2: Integer emulation with block scales. (a) An FP32 matrix is split into p INT8 slices. (b) The slice products Ai′Bj′ with i+j≤p+1 are accumulated exactly in INT32 and summed in FP32 with weights 2−7(i+j) . (c) Block scaling refines the scales (yellow) from one per row of A and column of B in the Ozaki scheme to one per 128 entries of a row (1-D) or per 128×128 tile (2-D).
Figure 3: FlashSO(2): the SO(2) -convolution block and its fused implementation. (a) Forward pass; each dashed box S1 to S3 is one fused kernel: Only INT8 slices and their scales pass between kernels. (b) Backward pass (read from right to left); each kernel is placed under its forward counterpart. Dashed arrows mark the tensors saved in the forward pass.
Figure 4: Equivariance and roofline analysis of the SO(2) convolutions under low precision. Rotation-equivariance residual of the (a) MALOQ Hamiltonian output head and (b) UMA forces against the rotation angle. The dashed ring marks the 10−4 tolerance. Roofline of the forward pass of (c) MALOQ and (d) UMA-Large (UMA-L) models. Light blue symbols: native implementation; dark blue: addition of fused layer(-FL); purple: fused layer with INT8 emulation (Emu 1 –Emu 3 ). The arrows split the gain over the native layer into kernel fusion ( I1 in intensity, T1 in throughput) and INT8 emulation ( I2 , T2 ). I1 = 4.6, T1 = 7.0, I2 = 1.4 and T2 = 5.6 in (c), and 2.9, 1.3, 1.4 and 5.5 in (d).
Figure 5: Accelerating MALOQ’s training. (a) Training loss curve with different precisions on the nablaDFT-2K dataset. (b) Inference error against the 32 held-out conformers, weights are trained with Emu 2 in (a). Inset: MAE of the Hamiltonian elements of each molecule relative to FP32. (c) Training step time and (d) peak memory as a function of the input size (number of atoms/edges).
Equivariant neural networks encode transformations such as rotations, translations, and permutations so model outputs change consistently with their inputs.
SO(3) equivariant graph neural networks have become the dominant paradigm for atomistic foundation models, achieving high accuracy and data efficiency by building rotational symmetry directly into the architecture. Yet the computational cost of their higher-order tensor operations creates a tough trade-off between model accuracy and inference efficiency. In this paper, we propose a structural pruning method for SO(3) equivariant atomistic foundation models to bridge this accuracy-efficiency gap. The pruning is applied along the channel and order dimensions, with each irreducible representation kept or removed as a complete block, thereby retaining SO(3) equivariance. Starting from a large checkpoint, the pruned model substantially reduces the inference cost while retaining higher accuracy than an independently trained small model. The pruned MACE-MP model outperforms the official from-scratch trained small model on 7 of 9 metrics on the Matbench Discovery leaderboard. In terms of efficiency, compressed MACE-MP and MACE-OFF models contain 1.5× to 4× fewer parameters and require 2.5× to 4× less pre-training compute than training a small model from scratch. For downstream applications, fine-tuning the pruned model reduces energy and force errors by 70.1% and 34.4% compared to training task-specific models from scratch across eight representative downstream datasets. We demonstrate that the method generalizes to other SO(3) equivariant architectures (SevenNet, eSCN) and can be combined with quantization and knowledge distillation for further gains.
Chen Wang, Siyu Hu, Guangming Tan +1
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences
Equivariant Graph Neural Networks (GNNs) are essential for physically consistent molecular simulations but suffer from high computational costs and memory bottlenecks, especially with high-order representations. While low-bit quantization offers a solution, applying it naively to rotation-sensitive features destroys the SO(3)-equivariant structure, leading to significant errors and violations of conservation laws. To address this issue, in this work, we propose a Geometric-Aware Quantization (GAQ) framework that compresses and accelerates equivariant models while rigorously preserving continuous symmetry in discrete spaces. Our approach introduces three key contributions: (1) a Magnitude-Direction Decoupled Quantization (MDDQ) scheme that separates invariant lengths from equivariant orientations to maintain geometric fidelity; (2) a symmetry-aware training strategy that treats scalar and vector features with distinct quantization schedules; and (3) a robust attention normalization mechanism to stabilize gradients in low-bit regimes. Experiments on the rMD17 benchmark demonstrate that our W4A8 models match the accuracy of FP32 baselines (9.31 meV vs. 23.20 meV) while reducing Local Equivariance Error (LEE) by over 30x compared to naive quantization. On consumer hardware, GAQ achieves 2.39x inference speedup and 4x memory reduction, enabling stable, energy-conserving molecular dynamics simulations for nanosecond timescales.
Equivariant machine learning interatomic potentials (MLIPs) promise atomistic simulations with accuracy approaching that of quantum-mechanical calculations. However, the high computational cost of their core Clebsch--Gordan tensor product (CGTP) convolutions severely limits model scalability and efficiency. In this paper, we rethink the parallel execution of CGTP convolutions and introduce LunaTP. Built on node-centric execution, LunaTP exploits path-group parallelism and multi-level reuse to accelerate CGTP convolutions and higher-order derivative computation while preserving the underlying mathematical formulation. On an NVIDIA H100 GPU, LunaTP achieves kernel speedups of up to 250.0× over e3nn and 17.1× over OpenEquivariance. It also delivers end-to-end training throughput gains of up to 28.7× over e3nn and 3.11× over OpenEquivariance. Finally, LunaTP demonstrates scalability and robust performance across models, workload scales, and GPUs, delivering substantial acceleration for training and atomistic simulation with high-accuracy equivariant MLIPs.
Yuanchang Zhou, Mingzhen Li, Chen Wang +3
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences