# Towards Faster and More Compact Foundation Models for Molecular Property Prediction

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 49 references that are in the map.
- First posted on arXiv 28 Apr 2025.

## Builds on (most cited first, 26 of 26)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [A Simple and Effective Pruning Approach for Large Language Models](https://mlsanity.com/paper/2306.11695/) (2023)
- [LLM-Pruner: On the Structural Pruning of Large Language Models](https://mlsanity.com/paper/2305.11627/) (2023)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [A Survey on Knowledge Distillation of Large Language Models](https://mlsanity.com/paper/2402.13116/) (2024)
- [The Open Catalyst 2020 (OC20) Dataset and Community Challenges](https://mlsanity.com/paper/2010.09990/) (2020)
- [The Unreasonable Ineffectiveness of the Deeper Layers](https://mlsanity.com/paper/2403.17887/) (2024)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [Pre-training via Denoising for Molecular Property Prediction](https://mlsanity.com/paper/2206.00133/) (2022)
- [The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials](https://mlsanity.com/paper/2205.06643/) (2022)
- [The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalysts](https://mlsanity.com/paper/2206.08917/) (2022)
- [SPICE, A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials](https://mlsanity.com/paper/2209.10702/) (2022)
- [GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets](https://mlsanity.com/paper/2204.02782/) (2022)
- [LoRAPrune: Structured Pruning Meets Low-Rank Parameter-Efficient Fine-Tuning](https://mlsanity.com/paper/2305.18403/) (2023)
- [From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction](https://mlsanity.com/paper/2310.16802/) (2023)
- [Knowledge Distillation on Graphs: A Survey](https://mlsanity.com/paper/2302.00219/) (2023)
- [Transition1x -- a Dataset for Building Generalizable Reactive Machine Learning Potentials](https://mlsanity.com/paper/2207.12858/) (2022)
- [Comprehensive Graph Gradual Pruning for Sparse Training in Graph Neural Networks](https://mlsanity.com/paper/2207.08629/) (2022)
- [The Open DAC 2023 Dataset and Challenges for Sorbent Discovery in Direct Air Capture](https://mlsanity.com/paper/2311.00341/) (2023)
- [Accelerating Molecular Graph Neural Networks via Knowledge Distillation](https://mlsanity.com/paper/2306.14818/) (2023)

## Built on by (most cited first, 1 of 1)

- [Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning](https://mlsanity.com/paper/2605.08885/) (2026)

## Abstract (arXiv)

Advancements in machine learning for molecular property prediction have improved accuracy but at the expense of higher computational cost and longer training times. Recently, the Joint Multi-domain Pre-training (JMP) foundation model has demonstrated strong performance across various downstream tasks with reduced training time over previous models. Despite JMP's advantages, fine-tuning it on molecular datasets ranging from small-scale to large-scale requires considerable time and computational resources. In this work, we investigate strategies to enhance efficiency by reducing model size while preserving performance. To better understand the model's efficiency, we analyze the layer contributions of JMP and find that later interaction blocks provide diminishing returns, suggesting an opportunity for model compression. We explore block reduction strategies by pruning the pre-trained model and evaluating its impact on efficiency and accuracy during fine-tuning. Our analysis reveals that removing two interaction blocks results in a minimal performance drop, reducing the model size by 32% while increasing inference throughput by 1.3x. These results suggest that JMP-L is over-parameterized and that a smaller, more efficient variant can achieve comparable performance with lower computational cost. Our study provides insights for developing lighter, faster, and more scalable foundation models for molecular and materials discovery. The code is publicly available at: https://github.com/Yasir-Ghunaim/efficient-jmp.