利用大规模无监督学习提升药物分子表示

 人工智能咨询培训老师叶梓 转载标明出处

在人工智能驱动的药物设计和发现领域，获取具有信息量的分子表示是一个至关重要的前提。近年来，研究者们将分子抽象为图，并利用图神经网络（GNNs）进行分子表示学习，展现出了巨大的潜力。然而，实际应用中GNNs面临着两个主要问题：一是用于监督训练的标记分子数据不足；二是模型对新合成分子的泛化能力较差。

为了解决这些问题，腾讯AI Lab的研究人员提出了一种名为GROVER（Graph Representation frOm self-superVised mEssage passing tRansformer）的新框架。该框架通过在节点、边和图级别上设计精心的无监督学习任务，能够从大量未标记的分子数据中学习丰富的结构和语义信息。

GROVER预训练框架

在大规模无标记数据上进行预训练，GROVER模型能够学习分子的丰富结构和语义信息。该模型基于Transformer构建，使用特制的图神经网络（GNNs）作为自注意力机制的构建块。这种设计使得模型不仅能够捕获图数据中的结构信息，还能在节点和边的消息传递路径上实现信息流动。

GROVER由两个模块组成：节点GNN转换器（node GNN transformer）和边GNN转换器。为了便于说明，本文只详细描述节点GNN转换器（简称node GTransformer），其结构如图1所示。

主要组件包括（图1）：

• 多头注意力机制（Multi-Head Attention）：允许模型在处理查询（Q）、键（K）和值（V）时，能够从多个角度学习信息。
• 层归一化（LayerNorm）：有助于加快训练速度，提高模型的稳定性。
• 前馈神经网络（Feed Forward）：在每个注意力层之后使用，为模型提供额外的非线性能力。
• 节点嵌入（Node Embed）：将节点的原始特征转换为嵌入表示。
• Aggregate2Node和Aggregate2Edge：分别用于将节点嵌入聚合到节点消息和边消息。
• 串联（Concat）：将来自不同层的信息连接起来，提供更丰富的表示。
• 动态消息传递网络（DyMPN）：一种新颖的消息传递机制，通过随机选择消息传递的跳数来增强模型的泛化能力。

GROVER设计的自监督任务如图2所示包括节点级的上下文属性预测任务和图级别的图式预测任务。

• 节点/边级别的上下文属性预测：通过遮蔽节点或边的局部子图，模型需要预测目标节点或边的上下文属性。
• 图级别的模式预测：利用领域知识，例如分子中的官能团，模型需要预测这些模式在分子中的出现。

动态消息传递网络（dyMPN）是GROVER的核心创新之一。传统的GNN消息传递过程需要预先指定每层的迭代次数和每次迭代中的跳数。GROVER引入了一种动态策略，即在每个训练周期中，根据某种随机分布选择每层的跳数。这种策略不仅提高了模型的泛化能力，还避免了过平滑问题。

自监督任务的设计对于预训练模型的成功至关重要。GROVER提出了两种自监督任务：节点/边级别的上下文属性预测和图级别的模式预测。

• 上下文属性预测：通过定义局部子图中的节点/边的统计属性，模型能够预测目标节点/边的上下文属性。图3以一个节点的局部子图为例，展示了如何定义节点的上下文属性。

• 图级别的模式预测：利用RDKit等专业软件检测分子中的官能团等模式，并将这些模式的出现作为预测目标。

预训练完成后，GROVER模型可以用于各种下游任务，如节点分类、链接预测和分子属性预测。通过微调，模型可以在特定任务上达到更好的性能。

想要掌握如何将大模型的力量发挥到极致吗？叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具（限时免费）。

1小时实战课程，您将学习到如何轻松上手并有效利用 Llama Factory 来微调您的模型，以发挥其最大潜力。

优快云教学平台录播地址：https://edu.youkuaiyun.com/course/detail/39987 

实验

实验包括预训练数据的收集、微调任务和数据集的准备、基线比较、以及下游任务的结果分析。

研究者们从ZINC15和Chembl数据集中收集了1100万个未标记的分子样本来预训练GROVER模型。为了模型选择，他们随机划分了10%的未标记分子作为验证集。

为了全面评估GROVER在下游任务上的性能，研究者们在MoleculeNet上的11个基准数据集上进行了实验，这些数据集涵盖了量子力学、物理化学、生物物理学和生理学等领域的各种目标。研究者们采用了更为现实的脚手架分割方法，以8:1:1的比例分配训练/验证/测试集，并进行了三次独立的随机种子脚手架分割实验，报告了平均值和标准差。

研究者们将GROVER与MoleculeNet中的10个流行基线模型和几个最新技术（STOA）方法进行了全面比较。表1中展示了所有模型在所有数据集上的性能比较。其中，标有绿色的方法是预训练方法。

在预训练阶段，研究者们为上下文属性预测任务设置了上下文半径k=1，以提取上下文属性字典，并分别获取了2518个和2686个不同的节点和边上下文属性作为节点和边的标签。每个分子图中随机遮蔽15%的节点和边标签进行预测。对于图级别模式预测任务，使用RDKit提取了85个功能团作为分子的模式，并将模式的标签表示为独热向量。为了评估模型大小的影响，研究者们预训练了两个不同隐藏层大小的GROVER模型：GROVERbase和GROVERlarge，同时保持所有其他超参数不变。

研究者们使用验证损失来选择最佳模型。对于每个训练过程，模型训练了100个周期。对于超参数，研究者们对每个数据集的验证集进行了随机搜索，并报告了最佳结果。

表1记录了所有模型在所有数据集上的整体结果。结果显示，GROVER模型在所有数据集上一致地实现了最佳性能，并且在大多数数据集上都有很大的提升。整体相对提升在所有数据集上是6.1%，在分类任务上是2.2%，在回归任务上是10.8%。这一显著的提升验证了GROVER预训练模型对分子属性预测任务的有效性。

为了研究自监督策略的贡献，研究者们比较了预训练的GROVER和未预训练的GROVER在分类数据集上的性能。结果显示，没有预训练的GROVER性能较差。自监督预训练平均AUC提高了3.8%，这证实了自监督预训练策略可以学习到隐含的领域知识，并增强下游任务的预测性能。

为了验证GTransformer的表达能力，研究者们实现了基于GIN和MPNN的不同骨干网络，并使用了一个包含600K未标记分子的玩具数据集进行了预训练。如图4所示，使用GTransformer骨干网络的GROVER在训练和验证中都优于GIN和MPNN，这再次证实了GTransformer的有效性。

为了证明GTransformer和dyMPN背后的理念，研究者们实现了两个变体：GROVER w/o dyMPN和GROVER w/o GTrans。如图5所示，GROVER w/o GTrans在训练和验证中表现最差，这意味着简单地结合GNN和Transformer并不能增强GNN的表达能力。dyMPN通过在训练过程中引入随机性，虽然稍微损害了训练损失，但验证损失变得更好。因此，dyMPN通过为每个消息传递步骤随机化接受域，为GROVER带来了更好的泛化能力。

https://arxiv.org/pdf/2007.02835v2

GitHub - tencent-ailab/grover: This is a Pytorch implementation of the paper: Self-Supervised Graph Transformer on Large-Scale Molecular Data

关注视频号随时了解前沿论文：