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

Atomistic models are essential for simulating materials and molecules with high accuracy. Recent advances like NequIP and MACE incorporate SO(3) symmetry, greatly improving predictive performance. However, these models rely on high-order tensor operations, which are computationally expensive, especially for large systems. The high parameter count and complex calculations limit their practical deployment, prompting research into model compression. Existing techniques like pruning, quantization, and distillation have been explored in deep learning, but applying them to geometric models while preserving physical symmetries remains challenging. Ensuring equivariance during compression is critical to maintain physical correctness, making straightforward pruning approaches unsuitable.

Despite their accuracy, SO(3)-equivariant models suffer from high computational demands due to complex tensor operations over higher-order features. Current coarse-grained pruning methods reduce model size but often break the geometric symmetry, leading to degraded performance. The key challenge is to develop a fine-grained, structure-aware pruning strategy that preserves the irreducible representation structure and physical equivariance. Additionally, defining a robust importance metric that captures the contribution of features to both energy and forces is essential. Balancing model compression with accuracy retention, especially in large and complex systems, remains an open problem.

This paper introduces a block-level structural pruning method tailored for SO(3) equivariant GNNs, treating each (channel, l) irreducible representation as an atomic unit to preserve symmetry. It employs an energy-force sensitivity-based importance criterion, leveraging gradients of the potential energy surface, which is SO(3) invariant. The four-stage pipeline—calibration, target-driven pruning, retraining, and fine-tuning—enables effective compression without sacrificing physical fidelity. The approach supports multiple architectures and can be combined with quantization and distillation, offering a comprehensive framework for efficient geometric deep learning.

• �� Calibration: Use a subset of pre-training data to compute energy and force gradients, deriving importance scores for each feature block.
• �� Pruning: Define target architecture, rank feature blocks by importance, and physically slice weights and tensors to match the target.
• �� Reconfiguration: Adjust tensor products and weights to align with the pruned structure, ensuring computational consistency.
• �� Retraining: Fine-tune the pruned model on pre-training data to recover performance.
• �� Fine-tuning: Further adapt the model on downstream tasks for optimal accuracy.
• �� Importance metric: Based on the gradient of energy with respect to features, capturing the physical relevance of each block, invariant under rotations.

Experiments on MACE, SevenNet, and eSCN architectures validate the method. Models are pruned to various target sizes, with parameters reduced by up to 4×. Training costs decrease proportionally, and inference throughput improves significantly. The pruned models outperform from-scratch counterparts in accuracy, especially on out-of-distribution datasets like Matbench Discovery. Fine-tuning further enhances performance, and the method demonstrates transferability across architectures and datasets. Combining pruning with quantization and distillation yields additional efficiency gains, confirming the robustness and versatility of the approach.

Pruned models achieve up to 4× fewer parameters, 2.7× inference speedup, and 70.1% reduction in energy MAE. On Matbench Discovery, the pruned MACE-MP surpasses the official checkpoint on 7/9 metrics with only 25% of training data. Fine-tuned models on downstream datasets show significant error reductions, with energy errors dropping by over 70%. The approach also generalizes to other architectures, maintaining high accuracy with substantially less training data and computational resources.

The method enables efficient large-scale atomistic simulations in materials design, drug discovery, and catalysis. It provides a lightweight, high-accuracy initialization for downstream tasks, reducing training time and resource consumption. The approach is suitable for deployment on resource-constrained hardware, facilitating real-time simulations and high-throughput screening. Its compatibility with other compression techniques broadens its applicability in industry and academia, accelerating the development of predictive models for complex systems.

The importance metric relies on first-order approximations, which may misjudge feature significance in highly nonlinear or dynamic systems. The pruning process still requires retraining, adding complexity. In systems with extremely high-order interactions, the method may need further refinement to avoid losing critical information. Future work should focus on adaptive importance evaluation, automatic target architecture design, and hardware-aware optimization to address these limitations.