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

Problem Statement

SO(3) equivariant graph neural networks (GNNs) have established themselves as the dominant paradigm for atomistic foundation models, achieving high accuracy and data efficiency by explicitly encoding rotational symmetry. However, these models rely on higher-order tensor product operations to couple geometric features, resulting in computational costs that scale up to $O(L_{max}^6)$, where $L_{max}$ is the maximum order of the irreducible representations. This creates a significant trade-off between model accuracy and inference efficiency.

Existing model compression strategies—specifically Knowledge Distillation (KD) and Quantization—have limitations in this domain. KD often requires training a student from scratch and does not directly reuse teacher weights, while quantization reduces bit-precision but leaves the parameter count and tensor-product paths unchanged. Furthermore, standard structural pruning methods developed for Large Language Models (LLMs) cannot be directly applied to SO(3) equivariant models. Naively removing feature dimensions breaks the rotational symmetry of the tensor product path, violating the fundamental equivariance property required for physical consistency. Previous attempts at pruning atomistic models have been limited to coarse-grained block removal (e.g., entire interaction layers), which reduces the model's receptive field and is particularly detrimental to shallow architectures.

Methodology

The authors propose a structural pruning framework specifically designed for SO(3) equivariant atomistic foundation models. The method operates on pre-trained checkpoints to produce smaller, dense models that retain SO(3) equivariance and high accuracy.

1. Equivariant Pruning Granularity

The core innovation lies in the definition of the pruning unit. In SO(3) equivariant models, features are decomposed into irreducible representations of order $l \in [0, L_{max}]$. A feature vector for a specific channel $k$ and order $l$ consists of $2l+1$ components indexed by $m \in [-l, l]$.

• Constraint: The rotation operation in SO(3) couples all $m$ components within a specific order $l$. Removing individual components (element-wise pruning) breaks equivariance.
• Solution: The method treats the subspace spanned by the $2l+1$ components of a specific order $l$ and channel $k$ as an atomic unit. Pruning is performed on the tuple $(k, l)$, represented by a binary mask $z^{(t)}_{k,l}$. This ensures that if a feature block is retained, all its rotational components are preserved, strictly maintaining SO(3) equivariance.

2. Physics-Informed Importance Criterion

To determine which $(k, l)$ blocks to retain, the authors propose an importance score based on the sensitivity of the potential energy surface.

• Dual Objective: MLIPs must predict both potential energy ($E$) and forces ($F_i = -\nabla_{r_i} E$).
• Metric: The importance score $I^{(t)}_{k,l}$ is derived from a first-order Taylor expansion of the energy change upon removing a feature block. It aggregates the product of the feature magnitude and the gradient of the energy with respect to that feature over all atoms and representation dimensions:
  $$I^{(t)}_{k,l} = \mathbb{E}_{R,Z \sim \mathcal{D}} \left| \sum_{j,m} \frac{\partial E}{\partial h^{(t)}_{j,klm}} \cdot h^{(t)}_{j,klm} \right|$$
• Advantage: This criterion is SO(3) invariant and incurs negligible overhead because the gradients $\frac{\partial E}{\partial h}$ are already computed during standard force inference.

3. Structural Alignment and Pipeline

The framework follows a four-stage pipeline:

1. Calibration: A small subset of pre-training data is passed through the large model to compute importance scores.
2. Target-Driven Pruning: A target architecture is defined. The top-ranked $(k, l)$ blocks are retained based on importance scores.
3. Structural Alignment: The model structure is physically altered to match the reduced dimensions. This involves:
   • Tensor Product Reconfiguration: Slicing radial weights and Clebsch-Gordan paths to align with retained channels.
   • Weight Slicing: Physically removing rows and columns from linear transformation weights to create a dense, smaller weight tensor.
4. Retraining & Fine-tuning: The pruned model undergoes lightweight retraining on a subset of pre-training data to recover from approximation errors, followed by fine-tuning on specific downstream tasks.

Key Contributions

1. Equivariant Pruning Method: Introduction of a structural pruning method for higher-order SO(3) equivariant tensors across channel and order dimensions that strictly preserves SO(3) equivariance by treating $(k, l)$ blocks as atomic units.
2. Four-Stage Framework: A complete pipeline (Calibration, Pruning, Retraining, Fine-tuning) for compressing foundation models.
3. Generalizability: Demonstration that the method applies to multiple architectures, including MACE, SevenNet, and eSCN.
4. Combinability: Validation that structural pruning can be combined with quantization and knowledge distillation for compounding efficiency gains.

Experimental Results

The method was evaluated primarily on MACE-MP (inorganic materials) and MACE-OFF (organic molecules), with generalizability tests on SevenNet and eSCN.

Accuracy vs. From-Scratch Training

• Superiority: Pruned models consistently outperform models of the same size trained from scratch.
  • On MACE-MP, pruned models achieved lower Force MAE and better out-of-distribution performance on the sAlex dataset compared to from-scratch baselines.
  • On Matbench Discovery, a pruned MACE-MP model (using only 25% of the pre-training data) outperformed the official from-scratch trained small model on 7 of 9 metrics.
  • On MACE-OFF, pruned models matched or exceeded from-scratch baselines across diverse molecular subsets, even when using only 50% of the pre-training data.
• Downstream Performance: Fine-tuning pruned models on eight inorganic downstream datasets reduced energy and force errors by 70.1% and 34.4%, respectively, compared to training task-specific models from scratch.

Efficiency Gains

• Parameter Reduction: Compressed models contain 1.5× to 4× fewer parameters.
• Training Cost: The full pipeline requires 2.5× to 4× less pre-training compute than training a small model from scratch (e.g., 181 GPU-hours vs. 750 GPU-hours for a specific configuration).
• Inference Speed: Pruned models achieve a 2.7× to 3.0× inference speedup and up to 5.7× reduction in peak memory usage.

Combination with Other Methods

• Quantization: Combining pruning with FP32 quantization yielded a 9.6× throughput improvement (relative to the base FP64 model).
• Knowledge Distillation: Combining pruning with Hessian-based KD further refined model performance, demonstrating that the two methods are complementary.

Significance and Claims

The paper claims that structural pruning offers a natural and effective pathway to bridge the accuracy-efficiency gap in atomistic foundation models. By directly inheriting weights from large, pre-trained checkpoints and compressing them while preserving physical symmetries, the method produces compact models that are:

1. More accurate than models of the same size trained from scratch.
2. Significantly cheaper to train and deploy.
3. Physically consistent, maintaining the strict SO(3) equivariance required for reliable molecular dynamics simulations (validated via MD simulations of liquid water showing stable energy drift and accurate radial distribution functions).

The authors emphasize that the core principle—evaluating the importance of coupled irreducible representations—is broadly applicable to any equivariant architecture representing features as sums of SO(3) irreducible representations. The work suggests that pruning allows practitioners to deploy high-fidelity foundation models in resource-constrained environments without sacrificing the transferability learned during large-scale pre-training.