2025-11-24T07:55:17.096511

Multitask finetuning and acceleration of chemical pretrained models for small molecule drug property prediction

Adrian, Chung, Boyd et al.
Chemical pretrained models, sometimes referred to as foundation models, are receiving considerable interest for drug discovery applications. The general chemical knowledge extracted from self-supervised training has the potential to improve predictions for critical drug discovery endpoints, including on-target potency and ADMET properties. Multi-task learning has previously been successfully leveraged to improve predictive models. Here, we show that enabling multitasking in finetuning of chemical pretrained graph neural network models such as Kinetic GROVER Multi-Task (KERMT), an enhanced version of the GROVER model, and Knowledge-guided Pre-training of Graph Transformer (KGPT) significantly improves performance over non-pretrained graph neural network models. Surprisingly, we find that the performance improvement from finetuning KERMT in a multitask manner is most significant at larger data sizes. Additionally, we publish two multitask ADMET data splits to enable more accurate benchmarking of multitask deep learning methods for drug property prediction. Finally, we provide an accelerated implementation of the KERMT model on GitHub, unlocking large-scale pretraining, finetuning, and inference in industrial drug discovery workflows.
academic

Multitask finetuning and acceleration of chemical pretrained models for small molecule drug property prediction

基本信息

  • 论文ID: 2510.12719
  • 标题: Multitask finetuning and acceleration of chemical pretrained models for small molecule drug property prediction
  • 作者: Matthew Adrian, Yunsie Chung, Kevin Boyd, Saee Paliwal, Srimukh Prasad Veccham, Alan C. Cheng
  • 机构: Merck & Co., Inc. 和 NVIDIA BioNeMo
  • 分类: cs.LG (机器学习), q-bio.QM (定量生物学方法)
  • 发表时间: 2025年10月14日
  • 论文链接: https://arxiv.org/abs/2510.12719v1

摘要

化学预训练模型(也称为基础模型)在药物发现应用中受到广泛关注。通过自监督训练提取的通用化学知识有潜力改善关键药物发现终点的预测，包括靶点效力和ADMET性质。本研究表明，在化学预训练图神经网络模型(如KERMT和KPGT)的微调中启用多任务学习显著提升了相比非预训练图神经网络模型的性能。令人惊讶的是，KERMT多任务微调的性能提升在更大数据规模下最为显著。此外，作者发布了两个多任务ADMET数据集分割，并提供了KERMT模型的加速实现。

研究背景与动机

核心问题

  1. 数据稀缺性挑战: 在药物发现中，特别是靶点效力预测等任务中，标注数据通常有限(10¹到10⁶个分子)，而整个化学空间估计约有10⁶⁰个分子
  2. 传统方法局限: 监督学习的图神经网络在小数据场景下性能有限，通常需要依赖随机森林等经典方法
  3. 多任务学习潜力: ADMET性质之间存在相关性，为多任务学习提供了机会，但尚未在化学预训练模型微调中得到充分探索

研究动机

  • 利用大规模未标注化学数据进行预训练，学习通用化学知识和模式
  • 探索多任务学习在化学预训练模型微调中的潜力
  • 解决工业级药物发现流程中的计算效率问题

核心贡献

  1. 首次系统性研究: 在化学预训练模型微调中引入多任务学习方法
  2. KERMT模型增强: 提出GROVER的增强版本，支持分布式预训练和加速推理
  3. 反直觉发现: 证明KERMT在大数据规模下表现更优，挑战了预训练模型主要在小数据场景有优势的假设
  4. 基准数据集: 发布两个多任务ADMET数据集分割，促进方法对比评估
  5. 工程优化: 提供加速实现，支持工业级应用

方法详解

任务定义

输入: 分子的SMILES字符串或分子图表示 输出: 多个ADMET性质或靶点效力的预测值 目标: 通过多任务学习提升化学预训练模型在药物性质预测任务上的性能

模型架构

KERMT (Kinetic GROVER Multi-Task)

  • 基础架构: 基于GROVER的图变换器模型
  • 预训练任务:
    • 节点/边级分类: 从节点/边嵌入识别k-hop局部子图
    • 图级多标签分类: 从图嵌入识别分子中存在的官能团
  • 参数规模: ~5100万参数(base版本)
  • 预训练数据: 1100万化合物(来自ZINC15和ChEMBL)

KPGT (Knowledge-guided Pre-training of Graph Transformer)

  • 特色: 使用分子线图表示 + 知识节点(K-node)
  • 预训练任务:
    • 预测被掩码的节点和K-node特征
    • 预测RDKit指纹
    • 预测200个分子描述符
  • 参数规模: ~1亿参数
  • 预训练数据: 200万分子(ChEMBL29)

多任务微调策略

  • 单任务微调: 仅更新编码器和前馈网络权重预测单一性质
  • 多任务微调: 前馈网络输出n个值对应n个性质，编码器权重同时更新

技术创新点

  1. 分布式预训练: 使用PyTorch DDP实现多GPU并行预训练，8个GPU达到86%扩展效率
  2. 加速推理: 集成cuik-molmaker包，实现2.2倍微调加速和2.9倍推理加速
  3. 自动超参优化: 集成Optuna进行超参数搜索
  4. 内存优化: 动态生成分子图和描述符，减少34%内存使用

实验设置

数据集

内部数据集(Merck)

  • ADMET数据: 30个终点，800,733个化合物(截至2024年)
  • 靶点效力: Target 1 (744化合物)，Target 2 (1,163化合物)
  • 分割方式: 80-20时间分割(2018年4月为分界点)

公开数据集

  • 文献ADMET数据: 25个终点，114,112个化合物
  • Biogen数据集: 6个终点，3,521个化合物
  • BindingDB: EGFR (9,462化合物)，BTK (9,337化合物)等
  • 分割方式: 基于PCA降维Morgan指纹的聚类分割

评价指标

  • 主要指标: Pearson r²相关系数
  • 辅助指标: 决定系数R²，平均绝对误差(MAE)，均方根误差(RMSE)
  • 分类评估: 分类富集图，评估高效力分子的正确分类率

对比方法

  • 基线: Chemprop (D-MPNN)
  • 预训练模型: MoLFormer, KPGT, KERMT
  • 评估模式: 单任务(ST)和多任务(MT)变体

实验结果

主要结果

内部ADMET数据性能

在Merck内部数据的时间分割测试中：

  • KERMT MT: 在5个关键终点中表现最佳或并列最佳
  • 性能提升: 相比Chemprop MT在30个终点中的18个表现更优
  • 平均改进: Pearson r²提升0.02(相比Chemprop)和0.04(相比KPGT)

具体结果(Pearson r²):

  • Papp: KERMT MT (0.712) vs Chemprop MT (0.657)
  • EPSA: KERMT MT (0.822) vs Chemprop MT (0.805)
  • Fu,p human: KERMT MT (0.666) vs Chemprop MT (0.641)

公开数据集表现

  • 公开ADMET数据: KPGT表现更优(9/25终点最佳)，KERMT MT仅3/25最佳
  • Biogen数据: 由于样本量小，结果置信度较低
  • 数据规模依赖: KERMT在大数据集(>10k样本)表现更好，KPGT在小数据集(<3k样本)更优

数据规模分析

关键发现：KERMT的优势在大数据规模下更明显

  • 临界点: 训练集>60k数据点时，KERMT显著优于Chemprop
  • 参数规模影响: KERMT(5100万参数)比Chemprop(500万参数)更容易在小数据上过拟合
  • 多任务收益: 随着任务数量增加(1→30个任务)，KERMT性能持续提升

化学空间泛化能力

通过Tanimoto相似性分析：

  • 一致性优势: KERMT在所有相似性区间(0.35-0.7)都优于Chemprop
  • 泛化能力: 虽然不是专门针对低相似性化合物，但整体泛化能力更强
  • 环肽预测: 在环肽子集上两个模型性能相当(Pearson r² = 0.36)

预训练数据影响

使用内部数据重新预训练的实验显示：

  • 有限改进: 即使预训练数据更相似下游任务，性能提升有限
  • 环肽表现: 基础KERMT模型在环肽任务上仍优于内部预训练模型(5/12 vs 1/12任务)
  • 启示: 需要改进预训练任务以更好地捕获相关信息

相关工作

化学预训练模型

  • GROVER: 图变换器，使用原子和键消息传递
  • MoLFormer: 基于SMILES的语言模型，旋转位置编码
  • KPGT: 知识引导的图变换器，整合分子描述符

多任务学习

  • 传统应用: 主要用于从头训练的深度学习架构
  • 本文贡献: 首次系统性地将多任务学习应用于化学预训练模型微调

结论与讨论

主要结论

  1. 多任务微调有效性: KERMT多任务微调显著提升性能，特别是在大数据场景
  2. 数据规模依赖性: 挑战了预训练模型主要在小数据有优势的传统观点
  3. 模型选择指导: 推荐大中型数据使用KERMT MT，小数据使用KPGT ST
  4. 工程可行性: 加速实现使工业级应用成为可能

局限性

  1. 预训练任务优化: 当前预训练任务可能未充分捕获下游任务相关信息
  2. 环肽预测: 在环肽等特殊分子类型上改进有限
  3. 数据集差异: 内部和公开数据集结果不一致，影响方法泛化性评估
  4. 计算资源需求: 大参数模型需要更多计算资源

未来方向

  1. 预训练任务改进: 设计更适合下游多任务学习的预训练目标
  2. 模块化微调: 研究部分冻结编码器对不同数据规模的影响
  3. 跨模态扩展: 探索蛋白质-小分子共同预训练
  4. 基准数据集: 开发更多高质量的多任务基准

深度评价

优点

  1. 实用价值高: 直接解决工业药物发现中的实际问题
  2. 实验全面: 涵盖多个数据集、模型和评估维度
  3. 反直觉发现: 挑战领域假设，提供新的洞察
  4. 工程贡献: 提供完整的开源实现和加速优化
  5. 数据贡献: 发布标准化的多任务基准数据集

不足

  1. 理论分析不足: 缺乏对为什么KERMT在大数据下表现更好的深入理论解释
  2. 预训练策略: 未充分探索针对多任务学习优化的预训练方法
  3. 统计显著性: 部分结果的统计显著性分析可以更严格
  4. 计算成本分析: 缺乏详细的计算成本对比分析

影响力

  1. 学术影响: 为化学信息学和多任务学习交叉领域提供重要参考
  2. 工业应用: 直接适用于制药公司的ADMET预测流程
  3. 开源贡献: 代码和数据的开放促进领域发展
  4. 方法论贡献: 建立了化学预训练模型评估的新标准

适用场景

  1. 大型制药公司: 拥有大规模ADMET数据的组织
  2. 多性质优化: 需要同时预测多个分子性质的场景
  3. 工业流程: 需要高效推理的生产环境
  4. 研究基准: 作为多任务化学性质预测的标准基线

参考文献

论文引用了47篇重要文献，涵盖：

  • 化学预训练模型的基础工作(GROVER, MoLFormer, KPGT)
  • ADMET预测的经典方法和数据集
  • 多任务学习的理论基础
  • 分子表示学习和图神经网络
  • 药物发现中机器学习应用的综述

总体评价: 这是一篇高质量的应用研究论文，在理论贡献、实验验证和工程实现方面都有显著价值。特别是其反直觉的发现和完整的开源贡献，对推动化学信息学领域的发展具有重要意义。