Towards Faster and More Compact Foundation Models for Molecular Property Prediction

TL;DR

This work targets efficiency in molecular property prediction by compressing the JMP-L foundation model. It identifies diminishing returns from deeper interaction blocks and demonstrates that pruning two blocks reduces parameters by 32% and increases inference throughput by ~1.3x with minimal loss in many tasks. By combining block reduction with knowledge distillation during pre-training and fine-tuning, the authors preserve accuracy on both in-distribution and out-of-distribution tasks while substantially cutting compute. The findings show that block-reduced variants, especially with distillation, offer practical routes to faster, more scalable models for molecular and materials discovery. The results hold across multiple architectures and underscore the value of pre-training-aware compression for domain-specific foundation models.

Abstract

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.

Figures (9)

• Figure 1: Block Reduction for Efficient Foundation Models. The top model represents the foundation model JMP-L, where interaction blocks extract features, which are concatenated and processed by FinalMLP before making predictions. The bottom model is its pruned version, constructed by removing low-importance blocks and adjusting FinalMLP. To mitigate performance degradation, we apply both feature distillation (node-to-node and edge-to-edge) and output distillation to transfer knowledge from the original model.
• Figure 1: Block Relevance Analysis for EquiformerV2. This figure illustrates the contribution of each transformer block in EquiformerV2 to the final prediction. The model consists of eight transformer blocks, $f_1$ to $f_8$, with relevance decreasing in the deeper layers. The results suggest that later blocks are strong candidates for pruning to improve efficiency with minimal performance loss.
• Figure 2: Block Relevance Analysis. This table illustrates the contribution of each output block in JMP-L to the final prediction. The first block, $f_1$, represents the embedding output, while $f_2$ to $f_7$ correspond to the six interaction blocks of GemNet-OC. The results show diminishing returns in deeper interaction blocks, suggesting they are strong candidates for pruning to improve efficiency without considerable performance loss.
• Figure 2: Evaluation on downstream tasks using EquiformerV2. We assess the performance of EquiformerV2 across various downstream tasks after applying block reduction (BR) and knowledge distillation (KD). Performance is reported in meV/Å for force targets, meV for the QM9 energy target, and eV for the QMOF band gap target. The teacher model refers to the original EquiformerV2 with 8 transformer blocks.
• Figure 3: Evaluation on downstream tasks. Evaluation of the performance across various downstream tasks using different block reduction strategies: block reduction (BR), block reduction with a randomly initialized MLP (BR/RandomMLP), and block reduction combined with knowledge distillation (BR+KD). Performance is measured in MAE: meV/Å for force targets, meV for the QM9 energy target, and eV for the QMOF band gap target. The original JMP-L model utilizes 6 blocks.