Towards Faster and More Compact Foundation Models for Molecular Property Prediction

Yasir M. Ghunaim, Andrés Villa, Gergo Ignacz, Gyorgy Szekely, Motasem Alfarra, Bernard Ghanem

Submission Track: Full Paper

Submission Category: AI-Guided Design

Keywords: Molecular Property Prediction, Efficient Models

TL;DR: We propose an efficient block reduction strategy for foundation models in molecular property prediction, demonstrating that pruning later interaction blocks reduces computational cost while maintaining comparable performance.

Abstract: Advancements in machine learning for molecular property prediction have improved accuracy but at the cost of increased complexity and longer training times. The recent Joint Multi-domain Pre-training (JMP) foundation model has demonstrated strong performance across various downstream tasks while reducing training time. However, fine-tuning on small-scale datasets remains time consuming, and larger datasets with more training samples pose even greater challenges. In this work, we investigate strategies to enhance efficiency by reducing model size while preserving performance. Through an analysis of layer contributions in JMP, we find that later interaction blocks provide diminishing returns, suggesting opportunities for model simplification. We explore block reduction strategies, where we prune the pre-trained model before fine-tuning, and assess their impact on efficiency and accuracy. Our findings reveal that removing two interaction blocks results in minimal performance drop, reducing model size by 32\% while increasing inference throughput by 1.3×. This confirms that JMP-L is over-parameterized, and a smaller, more efficient variant can achieve comparable performance at a 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: github.com/Yasir-Ghunaim/efficient-jmp.

AI4Mat Journal Track: Yes

Submission Number: 65