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

arXiv: https://arxiv.org/abs/2605.08885 · HTML: https://mlsanity.com/paper/2605.08885/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 67 references that are in the map.
- First posted on arXiv 9 May 2026.

## Builds on (most cited first, 30 of 36)

- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Knowledge Distillation: A Survey](https://mlsanity.com/paper/2006.05525/) (2020)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [A Survey of Quantization Methods for Efficient Neural Network Inference](https://mlsanity.com/paper/2103.13630/) (2021)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [LLM-Pruner: On the Structural Pruning of Large Language Models](https://mlsanity.com/paper/2305.11627/) (2023)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [SliceGPT: Compress Large Language Models by Deleting Rows and Columns](https://mlsanity.com/paper/2401.15024/) (2024)
- [ShortGPT: Layers in Large Language Models are More Redundant Than You Expect](https://mlsanity.com/paper/2403.03853/) (2024)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [The Open Catalyst 2020 (OC20) Dataset and Community Challenges](https://mlsanity.com/paper/2010.09990/) (2020)
- [Structured Pruning of Large Language Models](https://mlsanity.com/paper/1910.04732/) (2019)
- [e3nn: Euclidean Neural Networks](https://mlsanity.com/paper/2207.09453/) (2022)
- [Compact Language Models via Pruning and Knowledge Distillation](https://mlsanity.com/paper/2407.14679/) (2024)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [UMA: A Family of Universal Models for Atoms](https://mlsanity.com/paper/2506.23971/) (2025)
- [The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials](https://mlsanity.com/paper/2205.06643/) (2022)
- [Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction](https://mlsanity.com/paper/2502.12147/) (2025)
- [SPICE, A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials](https://mlsanity.com/paper/2209.10702/) (2022)
- [DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics](https://mlsanity.com/paper/1712.03641/) (2017)
- [Compression of Deep Learning Models for Text: A Survey](https://mlsanity.com/paper/2008.05221/) (2020)
- [Matbench Discovery -- A framework to evaluate machine learning crystal stability predictions](https://mlsanity.com/paper/2308.14920/) (2023)
- [PET-MAD, a lightweight universal interatomic potential for advanced materials modeling](https://mlsanity.com/paper/2503.14118/) (2025)
- [High-performance training and inference for deep equivariant interatomic potentials](https://mlsanity.com/paper/2504.16068/) (2025)
- [Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians](https://mlsanity.com/paper/2501.09009/) (2025)
- [Accelerating Molecular Graph Neural Networks via Knowledge Distillation](https://mlsanity.com/paper/2306.14818/) (2023)

## Abstract (arXiv)

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.