通过剪枝减少分子模型参数,提速30%且几乎不损失性能。
Towards Faster and More Compact Foundation Models for Molecular Property Prediction
- 剪掉JMP模型后两层交互块,压缩32%体积。
- 推理速度提升1.3倍,微调性能下降不足1%。
- 适合需要高效部署的分子药物研发场景。
机器学习在分子性质预测中的进步提升了准确性,但伴随更高的计算成本和训练时间。最近的联合多领域预训练(JMP)基础模型在多个下游任务中表现优异,且训练时间较以往更短。然而,在从小规模到大规模分子数据集上微调JMP仍需大量时间和资源。本文研究了通过减小模型规模来提升效率的方法。我们分析了JMP各层的贡献,发现后期交互模块带来的收益递减,提示存在压缩空间。通过剪枝预训练模型并评估其对微调效率与准确率的影响,结果表明移除两个交互块仅导致微小性能下降,使模型尺寸减少32%,推理吞吐量提升1.3倍。这表明JMP-L存在过度参数化,更小高效的变体可在保持相近性能的同时显著降低计算开销。本研究为开发更轻量、快速、可扩展的分子与材料发现基础模型提供了洞见。代码已公开:https://github.com/Yasir-Ghunaim/efficient-jmp。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。