Towards Faster and More Compact Foundation Models for Molecular Property Prediction

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.