[论文解读] Structured Multi-task Learning for Molecular Property Prediction

Shengchao Liu, Meng Qu+3 more|arXiv (Cornell University)|Feb 22, 2022

Machine Learning in Materials Science被引用 7

一句话总结

本文提出SGNN-EBM，一种用于分子性质预测的结构化多任务学习框架，通过显式任务关系图提升性能。通过结合用于潜在任务表征学习的态图神经网络（SGNN）与用于结构化输出预测的能量模型（EBM），SGNN-EBM在新构建的包含382项任务的ChEMBL-STRING数据集上实现了最先进性能，通过在潜在空间与输出空间中联合建模任务关系，并利用噪声对比估计实现高效训练，显著优于强基线模型。

ABSTRACT

Multi-task learning for molecular property prediction is becoming increasingly important in drug discovery. However, in contrast to other domains, the performance of multi-task learning in drug discovery is still not satisfying as the number of labeled data for each task is too limited, which calls for additional data to complement the data scarcity. In this paper, we study multi-task learning for molecular property prediction in a novel setting, where a relation graph between tasks is available. We first construct a dataset (ChEMBL-STRING) including around 400 tasks as well as a task relation graph. Then to better utilize such relation graph, we propose a method called SGNN-EBM to systematically investigate the structured task modeling from two perspectives. (1) In the \emph{latent} space, we model the task representations by applying a state graph neural network (SGNN) on the relation graph. (2) In the \emph{output} space, we employ structured prediction with the energy-based model (EBM), which can be efficiently trained through noise-contrastive estimation (NCE) approach. Empirical results justify the effectiveness of SGNN-EBM. Code is available on https://github.com/chao1224/SGNN-EBM.

研究动机与目标

• 通过利用来自领域知识的结构化任务关系，解决分子性质预测中的数据稀缺问题。
• 探究显式任务关系图在药物发现中提升多任务学习性能的机制。
• 开发一个统一框架，在潜在空间与输出空间中同时建模任务关系，以提升泛化能力。
• 构建一个大规模、领域特定的多任务数据集（ChEMBL-STRING），包含382项任务，并基于蛋白质-蛋白质相互作用构建任务关系图。

提出的方法

• 通过整合ChEMBL生物活性检测数据与STRING蛋白质-蛋白质相互作用网络，构建ChEMBL-STRING数据集，基于蛋白质靶点定义382项二分类任务。
• 使用态图神经网络（SGNN）通过在任务关系图上传播特征，学习潜在空间中的任务表征。
• 将多任务预测形式化为结构化预测，采用能量模型（EBM），其中能量函数整合了SGNN学习到的任务表征。
• 通过噪声对比估计（NCE）高效训练EBM，通过区分真实标签配置与噪声样本，优化联合标签分布。
• 探索两种噪声分布策略：固定均匀噪声与基于预训练SGNN模型初始化的自适应噪声。
• 在EBM中使用Gibbs采样进行推理，并对SGNN的层数进行消融研究，以评估其对长程依赖建模能力的影响。

实验结果

研究问题

• RQ1 基于生物知识的显式任务关系图是否能提升分子性质预测中多任务学习的性能？
• RQ2 与仅在潜在空间中建模任务关系相比，同时在潜在空间与输出空间中联合建模任务关系对预测性能有何影响？
• RQ3 在基于NCE的EBM训练中，不同噪声分布选择对下游性能有何影响？
• RQ4 SGNN的深度（层数）如何影响其在关系图中捕捉长程任务依赖的能力？
• RQ5 与独立任务预测相比，使用EBM进行结构化预测在多大程度上提升了性能？

主要发现

• SGNN-EBM在所有评估数据集上均取得最佳平均ROC-AUC性能，持续优于所有强基线多任务学习模型。
• SGNN组件本身即优于单任务学习与标准多任务学习基线，证明了结构化任务表征学习的价值。
• 将SGNN与EBM结合在SGNN-EBM中，相较于仅使用SGNN，实现了稳定且显著的性能提升，验证了结构化输出建模的有效性。
• 在NCE训练中，使用自适应更新的噪声分布（由预训练SGNN模型初始化）优于固定或均匀噪声分布，性能更优。
• 增加SGNN的层数可提升性能至一定程度（L=4），但进一步增加层数会降低性能，可能由于Gibbs采样在捕捉长程依赖方面存在局限。
• 消融研究证实，通过EBM进行结构化预测在所有预训练先验分布下均带来稳定但适中的性能提升，进一步验证了该方法的鲁棒性。