Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning

About

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$\times$ to 4$\times$ fewer parameters and require 2.5$\times$ to 4$\times$ 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, Weile Jia• 2026

Related benchmarks

TaskDatasetResultRank
Energy and force predictionrMD17 Ethanol (test)
Energy3
11
Atomistic Potential PredictionMPtrj
MAE Energy (meV/atom)9
10
Materials Stability PredictionMatbench-Discovery
F1 Score68
10
Atomistic Potential PredictionsAlex
MAE E (meV/atom)67
8
Energy and force predictionrMD17 Aspirin
Energy MAE (meV)10.8
5
Energy and force predictionrMD17 Azobenzene
Energy MAE (meV)5.4
5
Energy and force predictionrMD17 Malonaldehyde
Energy MAE (meV)6.7
5
Energy and force predictionrMD17 Paracetamol
Energy MAE (meV)6.7
5
Energy and force predictionrMD17 Salicylic acid
Energy MAE (meV)4.5
5
Energy and force predictionrMD17 Toluene
Energy MAE (meV)2.1
5
Showing 10 of 17 rows

Other info

Follow for update