Research map / Architectures and efficiency / Pruning and efficient tuning
Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes structural pruning of SO(3) equivariant atomistic foundation models along channel and irreducible-order blocks, using an energy-gradient importance score, to yield compact pre-trained MLIPs that beat same-size scratch-trained models at lower training cost.
Problem
SO(3) equivariant atomistic foundation models have high accuracy but expensive higher-order tensor operations, creating a trade-off between accuracy and inference efficiency; existing pruning for MLIPs is coarse (layer-level) and naive feature pruning breaks equivariance.
Approach
Prune features as complete (channel k, order l) irreducible blocks, which preserves SO(3) equivariance, ranked by a first-order Taylor energy-sensitivity score (|dE/dh * h| summed over m and atoms) that reuses gradients already computed for forces, followed by lightweight retraining and fine-tuning.
Key results
- Pruned MACE-MP (Layer2 L0 C128) beats official MACE-MP-0 small model on 7 of 9 Matbench Discovery metrics using 25% of MPtrj.
- Fine-tuning the pruned initialization reduces energy and force errors by 70.1% and 34.4% versus from-scratch models across eight inorganic downstream datasets.
- Pruned models give 2.7x inference throughput on A800 and 9.6x combined throughput with FP32 on RTX 3090; pipeline costs ~181 GPU-hours vs ~750 for scratch training.
Surprising findings
- Width (order) pruning outperforms depth (layer) pruning; depth pruning raised Energy MAE by 50% relative to scratch baseline.
- Retaining the linear readout after depth pruning gave much worse errors; a non-linear readout substitution reduced Energy MAE by 26.3%.
Limitations
- Target pruned architecture must be specified in advance; no automatic architecture search.
- Depth pruning performs poorly, so the method is mainly effective for width/order reduction.
- Largest gains reported are on MACE-MP; eSCN and SevenNet results are limited to a single comparison each.
Experimental setup
Datasets: MPtrj SPICE sAlex Matbench Discovery SSE-PBE H2O-PD AgAu AlMgCu Cu Ti V W
Benchmarks: Matbench Discovery MPtrj sAlex MACE-OFF test subsets 8 inorganic downstream datasets 3BPA AcAc rMD17 OC20
Metrics: MAE E MAE F MAE S F1 DAF RMSD K_srme RMSE throughput (atom/s) peak memory
Baselines: From-scratch MACE-MP-0 small models MACE-MP-0 official checkpoint Depth pruning (layer pruning) Knowledge distillation (Hessian KD)
Builds on: MACE structural pruning Hessian-based knowledge distillation
Closest papers by meaning
- ELoRA: Low-Rank Adaptation for Equivariant GNNs (ICML 2025 · similarity 0.87)
- POP: Online Structural Pruning Enables Efficient Inference of Large Foundation Models (NeurIPS 2026 · similarity 0.87)
- EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations (ICLR 2024 · similarity 0.86)
- Learning (Approximately) Equivariant Networks via Constrained Optimization (NeurIPS 2025 · similarity 0.86)
- Efficient and Scalable Density Functional Theory Hamiltonian Prediction through Adaptive Sparsity (ICML 2025 · similarity 0.86)
- Improving Equivariant Networks with Probabilistic Symmetry Breaking (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
