Towards Faster and More Compact Foundation Models for Molecular Property Prediction

📄 arXiv: 2504.19538v1 📥 PDF

作者: Yasir Ghunaim, Andrés Villa, Gergo Ignacz, Gyorgy Szekely, Motasem Alfarra, Bernard Ghanem

分类: cs.LG, q-bio.BM

发布日期: 2025-04-28

🔗 代码/项目: GITHUB

💡 一句话要点

通过剪枝JMP模型，实现分子性质预测中更快速紧凑的Foundation模型

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 分子性质预测 Foundation模型 模型剪枝 JMP模型 模型压缩 效率优化 药物发现

📋 核心要点

  1. 现有分子性质预测模型计算成本高、训练时间长，限制了其在大规模数据集上的应用。
  2. 通过分析JMP模型各层贡献，发现后期交互块收益递减，提出剪枝策略以压缩模型。
  3. 实验表明，移除两个交互块能在性能损失最小的情况下，显著减小模型尺寸并提升推理速度。

📝 摘要（中文）

分子性质预测的机器学习方法在提高准确性的同时，也增加了计算成本和训练时间。最近，联合多域预训练(JMP)基础模型在各种下游任务中表现出强大的性能，并且比以前的模型减少了训练时间。尽管JMP具有优势，但在从小规模到大规模的分子数据集上对其进行微调仍然需要相当长的时间和计算资源。本文研究了通过减少模型大小来提高效率同时保持性能的策略。为了更好地理解模型的效率，我们分析了JMP的层贡献，发现后面的交互块提供的收益递减，这表明存在模型压缩的机会。我们通过剪枝预训练模型来探索块减少策略，并评估其对微调期间效率和准确性的影响。我们的分析表明，移除两个交互块会导致性能下降最小，同时将模型大小减少32%，并将推理吞吐量提高1.3倍。这些结果表明JMP-L是过度参数化的，并且更小、更高效的变体可以以更低的计算成本实现相当的性能。我们的研究为开发更轻、更快、更可扩展的分子和材料发现基础模型提供了见解。代码已公开发布在：https://github.com/Yasir-Ghunaim/efficient-jmp。

🔬 方法详解

问题定义：论文旨在解决分子性质预测中，现有Foundation模型（如JMP）计算成本高、模型体积大的问题。这些问题限制了模型在资源受限环境下的部署和应用，也增加了大规模分子数据集上的微调成本。现有方法通常侧重于提高预测精度，而忽略了模型的效率和可扩展性。

核心思路：论文的核心思路是通过模型剪枝，移除JMP模型中冗余的交互块，从而在保持预测性能的同时，显著减小模型尺寸并提高推理速度。作者观察到JMP模型后期交互块的贡献递减，认为这些块存在过度参数化，可以通过剪枝进行压缩。

技术框架：论文的技术框架主要包括以下几个步骤：1) 对预训练的JMP模型进行层贡献分析，确定可剪枝的交互块；2) 实施剪枝策略，移除选定的交互块；3) 在分子性质预测数据集上对剪枝后的模型进行微调；4) 评估剪枝前后模型的性能（预测精度）和效率（模型大小、推理速度）。

关键创新：论文的关键创新在于发现了JMP模型中交互块的冗余性，并提出了一种有效的剪枝策略，能够在显著减小模型尺寸的同时，保持甚至略微提升预测性能。这种方法不同于以往侧重于模型结构设计的优化，而是从模型压缩的角度出发，提高了现有Foundation模型在分子性质预测任务中的实用性。

关键设计：论文的关键设计包括：1) 交互块选择：通过分析各层对预测结果的贡献，选择贡献较小的交互块进行剪枝。具体方法未知，论文中未详细描述如何量化“贡献”；2) 剪枝比例：实验中移除两个交互块，模型大小减少32%；3) 微调策略：使用与原始JMP模型相同的微调策略，以保证公平比较。具体的参数设置、损失函数等细节未在摘要中详细说明。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，通过移除JMP-L模型中的两个交互块，模型大小减少了32%，推理吞吐量提高了1.3倍，而性能下降 minimal。这表明JMP-L模型存在过度参数化，通过剪枝可以获得更高效的模型变体，为分子性质预测任务提供了一种更具性价比的解决方案。

🎯 应用场景

该研究成果可应用于药物发现、材料科学等领域，加速新分子和材料的筛选与设计。通过降低计算成本和提高推理速度，使得Foundation模型能够在资源受限的环境下部署，并应用于更大规模的分子数据集，从而加速相关领域的研发进程。未来，该方法可以推广到其他Foundation模型的压缩，提高其在各个领域的应用效率。

📄 摘要（原文）

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.