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

• URL: http://arxiv.org/abs/2605.08885v1
• Date: Sat, 09 May 2026 11:07:03 GMT
• Title: Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning
• Abstract summary: We propose a structural pruning method for SO(3) equivariant atomistic foundation models.<n>The pruning is applied along the channel and order dimensions, with each irreducible representation kept or removed as a complete block.<n>The pruned MACE-MP model outperforms the official from-scratch trained small model on 7 of 9 metrics on the Matbench Discovery leaderboard.
• Score: 4.793059213046564
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Event-Driven Language Models with Sparse Neural Activity for Neuromorphic Hardware [5.066906328188419]
  Inference with transformer-based large language models (LLMs) is often limited by the memory-bound KV cache and quadratic attention cost.<n>We introduce a method that induces sparse neural activity in heavily quantized linear-attention models with minimal performance loss.
  arXiv  Detail & Related papers  (2026-08-31T08:28:01Z)
• Solve the Loop: Attractor Models for Language and Reasoning [4.8720589853137435]
  Looped Transformers offer a promising alternative to purely feed-forward computation.<n>We introduce Attractor Models, in which a backbone module first proposes output embeddings, then an attractor module refines them by solving for the fixed point.<n>We show that Attractor Models outperform existing models across two regimes, large-scale language-model pretraining and reasoning with tiny models.
  arXiv  Detail & Related papers  (2026-05-12T17:51:26Z)
• Three Roles, One Model: Role Orchestration at Inference Time to Close the Performance Gap Between Small and Large Agents [0.4666493857924357]
  We study whether inference-time scaffolding alone, without any additional training compute, can improve the performance of a small model in complex multi-step environments.<n>We evaluate Qwen3-8B on the AppWorld benchmark under both full-precision and 4-bit quantized configurations.<n>On full-precision inference, our scaffolded 8B model surpasses DeepSeek-Coder 33B Instruct (7.1%) from the original AppWorld evaluation.
  arXiv  Detail & Related papers  (2026-04-13T13:40:33Z)
• Equi-GSPR: Equivariant SE(3) Graph Network Model for Sparse Point Cloud Registration [2.814748676983944]
  We propose a graph neural network model embedded with a local Spherical Euclidean 3D equivariance property through SE(3) message passing based propagation. Our model is composed mainly of a descriptor module, equivariant graph layers, match similarity, and the final regression layers. Experiments conducted on the 3DMatch and KITTI datasets exhibit the compelling and robust performance of our model compared to state-of-the-art approaches.
  arXiv  Detail & Related papers  (2024-10-08T06:48:01Z)
• FINE: Factorizing Knowledge for Initialization of Variable-sized Diffusion Models [35.40065954148091]
  FINE is a method based on the Learngene framework to initializing downstream networks leveraging pre-trained models. It decomposes pre-trained knowledge into the product of matrices (i.e., $U$, $Sigma$, and $V$), where $U$ and $V$ are shared across network blocks as learngenes'' It consistently outperforms direct pre-training, particularly for smaller models, achieving state-of-the-art results across variable model sizes.
  arXiv  Detail & Related papers  (2024-09-28T08:57:17Z)
• Tangent Model Composition for Ensembling and Continual Fine-tuning [69.92177580782929]
  Tangent Model Composition (TMC) is a method to combine component models independently fine-tuned around a pre-trained point. TMC improves accuracy by 4.2% compared to ensembling non-linearly fine-tuned models.
  arXiv  Detail & Related papers  (2023-07-16T17:45:33Z)
• SqueezeLLM: Dense-and-Sparse Quantization [80.32162537942138]
  Main bottleneck for generative inference with LLMs is memory bandwidth, rather than compute, for single batch inference. We introduce SqueezeLLM, a post-training quantization framework that enables lossless compression to ultra-low precisions of up to 3-bit. Our framework incorporates two novel ideas: (i) sensitivity-based non-uniform quantization, which searches for the optimal bit precision assignment based on second-order information; and (ii) the Dense-and-Sparse decomposition that stores outliers and sensitive weight values in an efficient sparse format.
  arXiv  Detail & Related papers  (2023-06-13T08:57:54Z)
• Gradient-Free Structured Pruning with Unlabeled Data [57.999191898036706]
  We propose a gradient-free structured pruning framework that uses only unlabeled data. Up to 40% of the original FLOP count can be reduced with less than a 4% accuracy loss across all tasks considered.
  arXiv  Detail & Related papers  (2023-03-07T19:12:31Z)
• Hyperspherical Quantization: Toward Smaller and More Accurate Models [17.154801913113566]
  Vector quantization aims at reducing the model size by indexing model weights with full-precision embeddings. Binary and other low-precision quantization methods can reduce the model size up to 32$times$, however, at the cost of a considerable accuracy drop. We propose an efficient framework for ternary quantization to produce smaller and more accurate compressed models.
  arXiv  Detail & Related papers  (2022-12-24T04:42:15Z)
• Load-balanced Gather-scatter Patterns for Sparse Deep Neural Networks [20.374784902476318]
  Pruning, as a method to introduce zeros to model weights, has shown to be an effective method to provide good trade-offs between model accuracy and computation efficiency. Some modern processors are equipped with fast on-chip scratchpad memories and gather/scatter engines that perform indirect load and store operations on such memories. In this work, we propose a set of novel sparse patterns, named gather-scatter (GS) patterns, to utilize the scratchpad memories and gather/scatter engines to speed up neural network inferences.
  arXiv  Detail & Related papers  (2021-12-20T22:55:45Z)
• MoEfication: Conditional Computation of Transformer Models for Efficient Inference [66.56994436947441]
  Transformer-based pre-trained language models can achieve superior performance on most NLP tasks due to large parameter capacity, but also lead to huge computation cost. We explore to accelerate large-model inference by conditional computation based on the sparse activation phenomenon. We propose to transform a large model into its mixture-of-experts (MoE) version with equal model size, namely MoEfication.
  arXiv  Detail & Related papers  (2021-10-05T02:14:38Z)
• Effective Model Sparsification by Scheduled Grow-and-Prune Methods [73.03533268740605]
  We propose a novel scheduled grow-and-prune (GaP) methodology without pre-training the dense models. Experiments have shown that such models can match or beat the quality of highly optimized dense models at 80% sparsity on a variety of tasks.
  arXiv  Detail & Related papers  (2021-06-18T01:03:13Z)
• The Right Tool for the Job: Matching Model and Instance Complexities [62.95183777679024]
  As NLP models become larger, executing a trained model requires significant computational resources incurring monetary and environmental costs. We propose a modification to contextual representation fine-tuning which, during inference, allows for an early (and fast) "exit" We test our proposed modification on five different datasets in two tasks: three text classification datasets and two natural language inference benchmarks.
  arXiv  Detail & Related papers  (2020-04-16T04:28:08Z)
• Training with Quantization Noise for Extreme Model Compression [57.51832088938618]
  We tackle the problem of producing compact models, maximizing their accuracy for a given model size. A standard solution is to train networks with Quantization Aware Training, where the weights are quantized during training and the gradients approximated with the Straight-Through Estimator. In this paper, we extend this approach to work beyond int8 fixed-point quantization with extreme compression methods.
  arXiv  Detail & Related papers  (2020-04-15T20:10:53Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.