用于多任务ADME性质预测的概率对比预训练

# 概率对比预训练用于多任务ADME性质预测 来源: https://arxiv.org/html/2606.11508 Yifan Xue NVIDIA 圣克拉拉, CA 95051 evax@nvidia\.com &Srimukh Prasad Veccham NVIDIA 圣克拉拉, CA 95051 sveccham@nvidia\.com &Saee Paliwal NVIDIA 圣克拉拉, CA 95051 saeep@nvidia\.com &Tyler Shimko NVIDIA 圣克拉拉, CA 95051 tshimko@nvidia\.com &Micha Livne NVIDIA 圣克拉拉, CA 95051 mlivne@nvidia\.com ###### 摘要 准确预测吸收、分布、代谢和排泄（ADME）性质对药物发现至关重要，但由于ADME终点噪声大、相互依赖且数据有限，这仍然具有挑战性。我们提出了一个分子图-Transformer预训练框架，该框架结合了化学特定的自监督学习和对比互信息机器学习（cMIM）。我们的方法将分子图编码为潜在变量，从图派生的潜在码重建SMILES字符串，并用领域特定的自监督化学任务增强对比目标。我们不将这些任务视为带有单独调优损失权重的辅助正则化器，而是将重建、对比区分和化学特定监督构建为单个概率潜在变量目标中具有单位权重的对数概率因子。对于微调，我们提出了一个多任务GNN读出架构，带有任务特定的多层感知器头，在减轻负迁移并改进异质、非线性任务关系建模的同时，保留共享表示学习。在Biogen、ExpansionRX和ChEMBL-MT基准测试中，与KERMT基线相比，所提出的Contrastive KERMT预训练策略始终改善了这些多任务ADME基准测试的下游预测，在Biogen数据集上提升7.6%，在ExpansionRX上提升9.9%，在ChEMBL-MT上提升9.5%（终点的平均值），差异具有统计学显著性。我们进一步表明，将ADME邻近分子添加到预训练语料库中可以改善迁移，并且对比组件锐化了化学上有意义的潜在邻域。这些结果表明，cMIM通过向KERMT的局部化学特定自监督添加全局潜在邻域塑造来改善ADME表示学习，并且在受控预训练和所述微调协议下，组合比单独任一组件迁移得更好。 ## 1 引言 吸收、分布、代谢和排泄（ADME）性质是药物发现的核心，因为它们决定了一个有活性的分子是否能成为可行的治疗候选物：它必须达到适当的暴露量，持续适当的时间，并避免不利的药代动力学行为（Balani等，2005（https://arxiv.org/html/2606.11508#bib.bib35）；Pellegatti，2012（https://arxiv.org/html/2606.11508#bib.bib36））。在实践中，ADME优化与安全性评估密切相关，因为不良的药代动力学可能会增加毒性风险，而有前景的化合物最终必须平衡暴露量、疗效和安全性。测量这些性质需要物理化学分析、体外实验和体内研究的组合，这些方法成本高、速度慢，并且通常仅适用于化学空间中有限的区域。因此，药物发现项目越来越依赖计算机模型在实验测试之前对化合物进行优先级排序（Lombardo等，2017（https://arxiv.org/html/2606.11508#bib.bib37）；Cáceres01112020；Beckers等，2023（https://arxiv.org/html/2606.11508#bib.bib38））。预训练的分子图模型自然适合这种场景，但实际的ADME预测仍然面临三个相互关联的挑战。首先，学习的潜在邻域应在化学上有意义：在表示空间中接近的分子在结构和性质空间中也应接近。其次，下游ADME数据集是多任务、噪声大且不平衡的：多任务预测可以通过共享表示对数据有限的终点进行正则化，但单个共享预测头可能无法拟合终点特定的非线性，或者在任务相关性弱时导致负迁移（Xu等，2017（https://arxiv.org/html/2606.11508#bib.bib53））。第三，用于预训练的无标签分子可能与用于微调的ADME分析有很大不同，这使得语料库设计成为迁移的重要部分。现有的分子GNN、自监督目标、对比方法和多任务公式解决了这一图景的部分问题；我们将在第5节（https://arxiv.org/html/2606.11508#S5）中回顾它们。我们引入了Contrastive KERMT，一个用于ADME预测的图-Transformer预训练和微调框架。预训练目标将对比互信息机器学习（cMIM）适应于分子：编码分子图，重建其SMILES表示，并使用批内对比项塑造潜在空间。我们的主要方法学变化是将化学特定的自监督任务，包括KERMT/GROVER风格的原子、键和官能团预测（Adrian等，2025（https://arxiv.org/html/2606.11508#bib.bib9）），作为相同概率潜在变量目标中的观测变量。因此，这些任务是具有单位权重的对数概率因子，而不是带有单独调优损失权重的辅助正则化器。对于下游预测，我们将预训练主干与任务特定的MLP头结合，这样每个终点可以学习其自身的后期转换，同时仍然共享共同的分子表示。我们的贡献有三点：（i）我们提出了cMIM的概率扩展，将全局潜在邻域塑造与KERMT风格的化学特定自监督结合在一个单一潜在变量目标中，避免了对辅助正则化权重进行额外的超参数搜索；（ii）我们引入了用于ADME微调的任务特定多层感知器头，每个终点特定的头只通过其对应的任务损失进行更新；以及（iii）我们表明Contrastive KERMT在Biogen、ExpansionRX和ChEMBL-MT上改善了下游ADME预测，并且将ADME对齐的分子添加到预训练语料库中进一步改善了迁移。 ## 2 方法 #### 图到SMILES预训练 我们使用对比互信息机器学习（Livne，2025（https://arxiv.org/html/2606.11508#bib.bib3））在ADME微调之前塑造分子图Transformer的潜在空间。对于分子$i$，令$\boldsymbol{x}_i = (g_i, s_i)$表示底层分子以及两个等价的分子视图：其2D分子图$g_i$和其标准SMILES字符串$s_i$（Weininger，1988（https://arxiv.org/html/2606.11508#bib.bib17））。我们将$g_i$和$s_i$用作同一分子$\boldsymbol{x}_i$可互换的标识符，同时在图到SMILES模型中为它们分配不同的角色。编码器是KERMT图-Transformer主干，并定义了分子级别潜在码上的变分分布$q_\theta(\boldsymbol{z} \mid g_i)$。解码器是一个自回归SMILES Transformer，并定义了$p_\theta(s_i \mid \boldsymbol{z}_i)$。因此，与标准图自编码器不同，该模型对图表示进行编码并重建分子字符串。这种图到SMILES结构鼓励潜在码保留跨分子视图共享的信息，而不是过拟合到单一输入表示。MIM和cMIM图模型的示意图概览见附录中的图4（https://arxiv.org/html/2606.11508#A1.F4）。 #### MIM目标 原始MIM目标学习一个潜在变量模型，该模型最大化输入与潜在码之间的互信息，同时鼓励结构化的潜在聚类（Livne等，2019（https://arxiv.org/html/2606.11508#bib.bib1），2020（https://arxiv.org/html/2606.11508#bib.bib2））。这里我们使用A-MIM，一个非对称变体，其中训练期间仅对后验进行采样，而不对先验进行采样。给定一个批次$\mathcal{B} = \{\boldsymbol{x}_i = (g_i, s_i)\}_{i=1}^B$和潜在样本$\boldsymbol{z}_i \sim q_\theta(\boldsymbol{z} \mid g_i)$，图到SMILES A-MIM损失为：
$$\hat{\mathcal{L}}_{\text{A-MIM}} = -\frac{1}{B}\sum_{i=1}^B\left[\log p_\theta(s_i \mid \boldsymbol{z}_i) + \frac{1}{2}\left(\log q_\theta(\boldsymbol{z}_i \mid g_i) + \log p(\boldsymbol{z}_i)\right)\right],$$
其中$p(\boldsymbol{z})$是标准正态先验。第一项奖励从图派生的潜在码重建SMILES视图。对称的潜在密度项使采样码在编码器和先验下都保持高似然，这正则化了表示并鼓励结构化的潜在分布。SMILES似然以自回归方式评估为$\log p_\theta(s_i \mid \boldsymbol{z}_i) = \sum_{u=1}^{|s_i|} \log p_\theta(s_{i,u} \mid s_{i,<u}, \boldsymbol{z}_i)$。我们使用挑战组织者提供的时间分割，使用药物发现活动中较早合成的分子来训练和验证模型。时间分割提供了比随机或按类留出分割更现实的预测模型前瞻性性能估计（Sheridan，2013（https://arxiv.org/html/2606.11508#bib.bib28））。我们使用活动中后来合成的分子作为测试集，反映真实世界设置中的回顾性性能。 #### ChEMBL-MT。ChEMBL多任务（ChEMBL-MT）数据集是从ChEMBL数据库（Mendez等，2019（https://arxiv.org/html/2606.11508#bib.bib8））由Adrian等人（Adrian等，2025（https://arxiv.org/html/2606.11508#bib.bib9））整理而成。它包含25个终点，共114K个分子。在本工作考虑的所有终点中，ChEMBL-MT中的hERG抑制是唯一的毒性终点；其余终点均为ADME分析。我们使用Adrian等人发布的Taylor-Butina聚类分割（Butina，1999（https://arxiv.org/html/2606.11508#bib.bib21））。在我们的基准测试中，除hERG外的所有终点都属于ADME类别；hERG被包含为一个毒性终点。Biogen和ExpansionRX数据集是单一来源数据集，每个终点在非常相似的实验条件下测量。然而，ChEMBL-MT数据集由来自多个来源的终点聚合而成，实验在不同条件下进行，跨越很长的时间段。这导致了无法完全解释的批次效应。所有三个数据集都适用于多任务学习。通过在这三个具有各自独特特征的数据集上测试我们的模型，我们试图全面且现实地表示ADME预测问题。表4（https://arxiv.org/html/2606.11508#A3.T4）总结了三个下游数据集的大小、终点数量、来源类型、分割和许可证。 表4：用于评估微调性能的三个下游数据集和分割的摘要。 ## 附录 D 附加架构图 我们提供了三种模型变体的详细架构图：KERMT（图5（https://arxiv.org/html/2606.11508#A4.F5））、仅cMIM（图6（https://arxiv.org/html/2606.11508#A4.F6））和Contrastive KERMT（图7（https://arxiv.org/html/2606.11508#A4.F7））。请参阅图注。 图5：KERMT变体的架构。编码器（消息传递+Transformer）之后直接连接词汇预测头和损失计算器；不使用SMILES解码器或cMIM损失。 请参阅图注。 图6：仅cMIM的架构。编码器之后连接一个读出层、潜在分布头和SMILES Transformer解码器；cMIM损失根据潜在码和从解码器获得的重建损失进行计算。不使用词汇预测头。 请参阅图注。 图7：Contrastive KERMT的架构，它结合了KERMT和仅cMIM的所有组件：编码器、词汇预测头以及带cMIM损失的SMILES解码器都在单个预训练目标中处于活跃状态。 表5：三种模型变体的比较。所有变体共享相同的KERMT编码器；仅cMIM和Contrastive KERMT添加了SMILES解码器，而Contrastive KERMT额外保留了词汇预测头。参数数量针对11M分子预训练配置报告；Contrastive KERMT的总参数在表1（https://arxiv.org/html/2606.11508#S3.T1）中的增强配置中波动$\pm 0.1$M，因为SMILES、原子和键词汇表针对每个预训练语料库重建。 ## 附录 E 实现细节 Contrastive KERMT实现基于Adrian等人（Adrian等，2025（https://arxiv.org/html/2606.11508#bib.bib9））描述的KERMT架构。为便于审阅，我们在补充材料中提供了一个匿名化的代码包；该代码包将在发表后去匿名化。 #### 预训练配置。所有预训练运行使用即时计算的自监督目标，cMIM和Contrastive KERMT运行额外使用批内负样本用于对比项；不需要正对分子增强。对于具有对比组件的变体，cMIM温度固定为$\tau=0.1$。所有预训练变体共享相同的KERMT编码器主干和优化器设置，并使用随机种子0；变体在预训练目标和仅预训练模块上有所不同。训练预算随语料库大小变化：11M基础配置，包括所有增强的11M变体，训练100个epoch，其中20个预热epoch；208M基础运行训练6个epoch，其中2个预热epoch；208M基础+Biogen运行训练4个epoch，其中2个预热epoch。对于cMIM变体，后验头预测对角高斯分布$q_\theta(\boldsymbol{z}_i \mid g_i) = \mathcal{N}(\mu_i, \operatorname{diag}(\sigma_i^2))$。使用重参数化技巧抽取潜在样本。为数值稳定性，方差被裁剪到低于$10^{-6}$。后验头是一个仅预训练模块，在下游微调前被丢弃。损失聚合遵循公式(4)（https://arxiv.org/html/2606.11508#S2.E4）中的概率因子。SMILES重建使用教师强制的自回归负对数似然，使用字符级分词器，对每个分子的token求和。高斯后验和先验对数密度对潜在维度求和。KERMT原子、键和官能团目标如KERMT中计算，原子和键头在进入联合目标前进行平均。预训练目标是自监督的，可以即时对任何分子计算。对于下游多任务微调，损失仅针对观测到的分析标签进行计算。 #### 硬件和计算。所有预训练运行使用一个节点，配备8块NVIDIA A100 GPU，位于内部计算集群上，使用同步数据并行训练（NCCL后端）。每个GPU的批次大小为128，所有变体和语料库的有效全局批次大小为1,024个分子每次优化步骤。每次运行的墙上时间成本约为：11M基础配置（100个epoch）3周，208M基础+Biogen配置6周。