论文状态:已完成

BatmanNet: Bi-branch Masked Graph Transformer Autoencoder for Molecular Representation

发表:2022/11/25
原文链接PDF 下载
价格:0.100000
价格:0.100000
已有 2 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文提出了一种新颖的双分支掩蔽图 Transformer 自编码器——BatmanNet,用于有效的分子表示学习。该模型通过简单的自监督策略,同时捕捉分子的局部和全局信息,成功重构掩蔽图中的缺失节点和边,在多个药物发现任务中取得了最先进的结果,展示了其在生物医药领域的巨大潜力。

摘要

Although substantial efforts have been made using graph neural networks (GNNs) for AI-driven drug discovery (AIDD), effective molecular representation learning remains an open challenge, especially in the case of insufficient labeled molecules. Recent studies suggest that big GNN models pre-trained by self-supervised learning on unlabeled datasets enable better transfer performance in downstream molecular property prediction tasks. However, the approaches in these studies require multiple complex self-supervised tasks and large-scale datasets, which are time-consuming, computationally expensive, and difficult to pre-train end-to-end. Here, we design a simple yet effective self-supervised strategy to simultaneously learn local and global information about molecules, and further propose a novel bi-branch masked graph transformer autoencoder (BatmanNet) to learn molecular representations. BatmanNet features two tailored complementary and asymmetric graph autoencoders to reconstruct the missing nodes and edges, respectively, from a masked molecular graph. With this design, BatmanNet can effectively capture the underlying structure and semantic information of molecules, thus improving the performance of molecular representation. BatmanNet achieves state-of-the-art results for multiple drug discovery tasks, including molecular properties prediction, drug-drug interaction, and drug-target interaction, on 13 benchmark datasets, demonstrating its great potential and superiority in molecular representation learning.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

BatmanNet: Bi-branch Masked Graph Transformer Autoencoder for Molecular Representation

1.2. 作者

Zhen Wang, Zheng Feng, Yanjun Li, Bowen Li, Yongrui Wang, Chulin Sha, Min He, Xiaolin Li

1.3. 发表期刊/会议

该论文以预印本 (preprint) 形式发表在 arXiv 上,其发布时间为 2022-11-25T09:44:28.000Z。arXiv 是一个开放获取的论文预印本库,在计算机科学、物理学、数学等领域具有广泛的影响力,许多重要研究成果在正式同行评审发表前会先在此发布。

1.4. 发表年份

2022年

1.5. 摘要

尽管在人工智能驱动的药物发现 (AI-driven Drug Discovery, AIDD) 领域,图神经网络 (Graph Neural Networks, GNNs) 取得了显著进展,但有效的分子表示学习 (molecular representation learning) 仍然是一个开放的挑战,尤其是在标记分子 (labeled molecules) 不足的情况下。最近的研究表明,在未标记数据集上通过自监督学习 (self-supervised learning) 预训练的大型 GNN 模型,在下游分子性质预测 (molecular property prediction) 任务中能够实现更好的迁移性能 (transfer performance)。然而,这些方法通常需要多个复杂的自监督任务和大规模数据集,这导致预训练过程耗时、计算成本高昂且难以端到端 (end-to-end) 进行。

本文提出了一种简单而有效的自监督策略,可以同时学习分子的局部 (local) 和全局 (global) 信息,并进一步提出了一种新颖的双分支掩蔽图 Transformer 自编码器 (bi-branch masked graph transformer autoencoder)——BatmanNet 来学习分子表示。BatmanNet 具有两个量身定制的、互补且不对称的图自编码器 (graph autoencoders),分别用于从掩蔽的分子图中重建缺失的节点 (nodes) 和边 (edges)。通过这种设计,BatmanNet 能够有效地捕获分子的潜在结构和语义信息,从而提高了分子表示的性能。BatmanNet 在13个基准数据集 (benchmark datasets) 上的多个药物发现任务(包括分子性质预测、药物-药物相互作用 (drug-drug interaction, DDI) 和药物-靶点相互作用 (drug-target interaction, DTI))中取得了最先进 (state-of-the-art, SOTA) 的结果,展示了其在分子表示学习方面的巨大潜力和优越性。

1.6. 原文链接

原文链接:https://arxiv.org/abs/2211.13979v3 PDF 链接:https://arxiv.org/pdf/2211.13979v3.pdf 发布状态:预印本 (v3 版本)

2. 整体概括

2.1. 研究背景与动机

人工智能驱动的药物发现 (AIDD) 是一个备受关注的研究领域,其中有效的分子表示学习 (molecular representation learning) 在各种下游任务(如分子性质预测、药物-药物相互作用 (DDI) 预测和药物-靶点相互作用 (DTI) 预测)中扮演着关键角色。图神经网络 (GNNs) 因其在建模图结构数据方面的优势,在分子表示学习中展现出巨大潜力,将原子视为节点,键视为边。

然而,现有方法存在以下主要局限性:

  • 标记数据稀缺: 生物化学领域与小分子相关的任务特定标记数据 (task-specific labeled data) 稀缺,主要是因为通过湿实验室实验获取高质量分子性质标签的成本高昂且耗时。这导致在有限数据集上监督训练深度 GNNs 容易出现过拟合 (overfitting) 问题。
  • 自监督学习的复杂性: 尽管自监督学习 (self-supervised learning) 已被提出用于在未标记数据上预训练大型 GNNs 以提高下游任务性能,但现有方法通常面临以下挑战:
    • 预训练任务复杂: 许多先前研究需要构建多种复杂的预训练任务来学习分子的局部和全局信息,这些任务往往需要引入额外的领域知识 (domain knowledge),如基序 (motifs)、子图 (subgraphs)、原子距离矩阵、分子描述符和指纹 (fingerprints)。这种复杂性以及额外知识的引入使得整个预训练过程难以有效地端到端完成。

    • 计算复杂度和模型规模大: 许多基于 Transformer 的分子预训练模型直接在整个分子图上编码信息,导致计算复杂度极高和模型规模庞大。这不仅耗时、计算昂贵,还需要大量的训练数据集和计算资源,这对于小型研究团队或企业来说可能难以获得。

      为了解决这些挑战,本文旨在设计一种简单、高效的分子自监督学习策略,并提出一种新型模型架构,以减轻上述问题,显著提高分子表示学习的有效性和效率。

2.2. 核心贡献/主要发现

本文通过提出BatmanNet模型及其自监督预训练策略,做出了以下核心贡献:

  • 新颖的自监督预训练策略: 提出了一种简单而强大的自监督预训练策略,通过以高比例(60%)同时掩蔽分子图中的节点和边,并通过基于图的自编码器 (graph-based autoencoder) 框架重建缺失部分。这种策略无需构建多层级的复杂任务或引入额外的领域化学知识,即可有效、自动地学习分子的局部和全局信息,以端到端的方式获取富有表现力的结构和语义知识。
  • 双分支不对称图 Transformer 自编码器架构: 开发了一种简单、有效且可扩展的双分支掩蔽图 Transformer 自编码器 (BatmanNet)。它采用不对称的编码器-解码器 (encoder-decoder) 设计:编码器仅在分子图的可见子集(未掩蔽部分)上操作,而解码器从学习到的表示和掩蔽标记 (mask tokens) 重建分子图,且解码器架构更轻量级。这种设计显著减少了预训练过程的计算量、整体预训练时间和内存消耗。
  • 最先进的性能: BatmanNet 在13个广泛使用的基准数据集上,在多个药物发现任务(包括分子性质预测、药物-药物相互作用 (DDI) 预测和药物-靶点相互作用 (DTI) 预测)中取得了最先进 (SOTA) 的结果。
  • 可视化验证: 通过可视化预训练后的分子表示,证明BatmanNet能够有效区分有效分子和无效分子,表明模型学习到了有意义的结构有效性信息。
  • 效率和可扩展性: 相比于现有 SOTA 模型,BatmanNet 在使用更少训练数据和模型参数的情况下,实现了可比较甚至更优的性能,验证了其卓越的效率和可扩展性。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 图神经网络 (Graph Neural Networks, GNNs)

GNNs 是一类专门设计用于处理图结构数据的神经网络。在分子表示学习中,GNNs 将分子建模为图,其中原子是节点 (nodes),化学键是边 (edges)。GNNs 的核心机制是消息传递 (message passing)邻域聚合 (neighborhood aggregation)。每个节点通过迭代地聚合其邻居节点和边的信息来更新其隐藏表示 (hidden representation)。经过 KK 层消息传递,节点的隐藏表示将包含其 KK 跳邻域 (K-hop neighborhood) 的结构信息。

3.1.2. Transformer 架构

Transformer 是一种基于注意力机制 (attention mechanism) 的深度学习模型架构,最初为自然语言处理 (Natural Language Processing, NLP) 任务设计。其核心组件是多头注意力机制 (Multi-head attention mechanism),允许模型在处理序列数据时同时关注不同位置的信息。Transformer 模型避免了循环神经网络 (Recurrent Neural Networks, RNNs) 和卷积神经网络 (Convolutional Neural Networks, CNNs) 的一些局限性,特别是在捕获长距离依赖 (long-range dependencies) 方面表现出色。在图学习领域,Transformer 也被适配用于处理图数据,通常通过将图结构信息融入注意力计算中。

3.1.3. 自监督学习 (Self-supervised Learning, SSL)

自监督学习是一种机器学习范式,它利用数据自身的信息来生成监督信号,从而训练模型。与传统的监督学习 (supervised learning) 依赖于人类标注的标签不同,自监督学习通过设计“前置任务 (pretext tasks)”来让模型从大规模无标签数据中学习有用的表示。在分子图领域,常见的自监督任务包括掩蔽节点/边预测、图对比学习 (graph contrastive learning) 等。学习到的表示随后可以迁移到下游的监督任务中,通常能提高性能。

3.1.4. 自编码器 (Autoencoder)

自编码器是一种神经网络模型,旨在学习输入数据的有效编码 (encoding) 或表示 (representation)。它由两部分组成:

  • 编码器 (Encoder): 将输入数据映射到低维的潜在空间 (latent space) 表示。
  • 解码器 (Decoder): 从潜在空间表示重建原始输入数据。 自编码器通过最小化输入和重建输出之间的差异来训练。在本文中,BatmanNet 被设计为一个自编码器,通过重建被掩蔽的分子图来学习表示。

3.2. 前人工作与技术演进

3.2.1. 分子表示学习的演进

  • 早期特征工程: 早期方法依赖于固定分子表示,如分子描述符 (molecular descriptors) 和指纹 (fingerprints) (例如 ECFP [50])。这些方法需要复杂的特征工程,并且存在向量稀疏性问题。
  • 深度学习驱动:
    • 序列模型: 随着深度学习的发展,一些研究开始使用分子简化分子线性输入规范系统 (SMILES) 字符串作为输入,并采用基于循环神经网络 (RNN-based) 的模型生成分子表示 (例如 Mol2Vec [54], N-GRAM [43])。还有研究使用掩蔽语言建模 (masked language modeling) 预训练 BERT 风格的模型 (例如 SMILES-BERT [14, 15])。然而,SMILES 表示在捕获分子结构信息和相似性方面存在局限性,并且可能生成无效分子。
    • 图神经网络 (GNNs): 近年来,GNNs 因其在建模图结构数据方面的优势被广泛应用于学习分子图表示。例如,GraphConv [52]、Weave [31]、SchNet [32] 探索了使用图卷积网络编码分子图。MPNN [37] 及其变体 DMPNN [39]、MGCN [40] 扩展了消息传递框架来建模键相互作用。AttentiveFP [35] 则利用图注意力网络 (Graph Attention Network, GAT) 的扩展来学习聚合权重。然而,传统的 GNNs 的消息传递机制主要聚合局部信息,难以捕获分子内的长距离依赖 (long-range dependencies)。

3.2.2. 分子图上的自监督学习 (SSL)

自监督学习在计算机视觉 (Computer Vision, CV) 和自然语言处理 (NLP) 等领域取得了丰硕成果,并逐渐成为分子图领域的核心方向。

  • 基于 2D 拓扑结构的方法:

    • 许多方法从分子 2D 拓扑结构预训练模型 (例如 pre-trainGNN [17], GraphMAE [47])。
    • 一些工作,如 MPG [13] 和 GROVER [18],采用 Transformer 风格的架构来预训练分子图。这些方法通常需要预定义和提取基序 (motifs) 或子图 (subgraphs) 作为自监督任务的预测目标。
    • KPGT [19] 则引入额外知识,并将其随机掩蔽一部分作为重建目标。
    • GMAE [44] 仅掩蔽节点,而 MGAE [45] 仅掩蔽边。GraphMAE [47] 则用描述符替换掩蔽节点。
    • 挑战: 这些方法普遍面临预训练任务复杂、需要额外领域知识、模型参数多和计算复杂度高等问题。
  • 基于 3D 几何结构的方法:

    • 一些方法利用原子空间位置的 3D 几何图,使用几何 GNN 模型 (例如 GEM [20]) 进行预训练。
    • 挑战: 这些方法通常需要大量分子数据进行预训练,且架构复杂,在数据稀疏时泛化能力可能受限。

3.3. 差异化分析

BatmanNet 相较于相关工作的核心区别和创新点体现在以下几个方面:

  • 自监督策略的简洁性与有效性:
    • 区别于多任务和领域知识依赖: 现有方法(如 GROVERMPGKPGTGEM)常需要构建多种复杂的预训练任务(如基序预测、子图预测、原子距离矩阵预测)并引入额外的领域知识。BatmanNet 则采用单一的、简单的双分支图掩蔽和重建任务,无需预定义基序、子图或额外信息,使其更直接、直观且易于实现。
    • 高掩蔽比例的挑战性: BatmanNet 随机掩蔽高比例(60%)的节点和边,使得每个节点/边更有可能同时缺失邻居节点和边。为了重建,模型必须学习更丰富的局部上下文信息。这种高比例掩蔽使得任务比其他使用较小图或基序作为目标的方法更具挑战性,从而促使模型学习更强大的表示。
  • 不对称编码器-解码器架构的高效性:
    • 计算和内存优化: 大多数基于 Transformer 的预训练模型(如 GROVERMPG)直接在整个分子图上编码信息,导致高计算复杂度和大型模型参数。BatmanNet 采用不对称的 Transformer 风格自编码器,编码器仅处理分子图的可见(未掩蔽)部分,而仅在轻量级解码器阶段才处理完整的分子图。这种设计显著减少了模型的参数数量和计算复杂度,从而提高了预训练效率和内存利用率。
    • 双分支互补性:GMAE 仅掩蔽节点或 MGAE 仅掩蔽边不同,BatmanNet 构建了一个双分支互补自编码器,节点分支和边分支分别执行节点掩蔽和边掩蔽,从而增强了模型的表达能力,更全面地捕获分子信息。
  • 端到端学习能力: BatmanNet 的设计使其能够以端到端的方式有效地、自动地学习分子的局部和全局信息,避免了传统方法中手动定义复杂特征或任务的需要。

4. 方法论

本文提出了一种新颖的双分支掩蔽图 Transformer 自编码器 (BatmanNet),用于分子表示学习,并结合了一种简单而有效的自监督预训练策略。

4.1. 方法原理

BatmanNet 的核心思想是利用掩蔽自编码器 (masked autoencoder) 的思想在分子图上进行自监督学习。它通过随机掩蔽分子图中的高比例节点和边,然后训练模型去重建这些缺失的部分。这种重建任务强制模型学习分子图的深层结构和语义信息,包括局部上下文和全局拓扑。为了更全面地捕获分子信息,BatmanNet 采用双分支 (bi-branch) 架构,一个分支专注于节点信息,另一个分支专注于边信息。此外,借鉴 MAE 的成功经验,BatmanNet 采用不对称的编码器-解码器 (asymmetric encoder-decoder) 设计,编码器仅处理可见(未掩蔽)的输入,而轻量级解码器负责从编码器的输出和掩蔽标记中重建完整的图,从而提高计算效率。

4.2. 核心方法详解 (逐层深入)

4.2.1. 预备知识:图神经网络 (GNNs)

GNNs 通过消息传递机制更新节点的表示。对于 GNN 的第 kk 层,节点 vv 的消息聚合 mv(k)\mathbf{m}_v^{(k)} 和隐藏表示 hv(k)\mathbf{h}_v^{(k)} 更新如下: mv(k)=AGG(k)({(hv(k1),hu(k1),euv)uNv}) \mathbf{m}_v^{(k)} = \mathrm{AGG}^{(k)} \left( \left\{ \left( \mathbf{h}_v^{(k-1)} , \mathbf{h}_u^{(k-1)} , \mathbf{e}_{uv} \right) \mid u \in \mathcal{N}_v \right\} \right) hv(k)=σ(W(k)mv(k)+b(k)) \mathbf{h}_v^{(k)} = \sigma \left( \mathbf{W}^{(k)} \mathbf{m}_v^{(k)} + \mathbf{b}^{(k)} \right) 其中:

  • mv(k)\mathbf{m}_v^{(k)}:聚合到的消息 (aggregated message)。
  • hv(k)\mathbf{h}_v^{(k)}:节点 vv 在第 kk 层的表示。
  • hu(k1)\mathbf{h}_u^{(k-1)}:邻居节点 uu 在第 k-1 层的表示。
  • euv\mathbf{e}_{uv}:边 (u, v) 的表示。
  • σ()\sigma(\cdot):激活函数。
  • W(k)\mathbf{W}^{(k)}:第 kk 层的权重矩阵。
  • b(k)\mathbf{b}^{(k)}:第 kk 层的偏置向量。
  • Nv\mathcal{N}_v:节点 vv 的邻居节点集合。
  • AGG(k)()\mathrm{AGG}^{(k)}(\cdot):第 kk 层的邻域聚合函数。 初始时,节点 vv 的表示 hv(0)\mathbf{h}_v^{(0)} 初始化为其原始特征 XvX_v。在经过 KK 次迭代后,通过一个 READOUT 函数得到整个图的表示 hG\mathbf{h}_GhG=READOUT({hv(K)vV}) \mathbf{h}_G = \operatorname{READOUT} \left( \left\{ \mathbf{h}_v^{(K)} \mid v \in \mathcal{V} \right\} \right) 其中 V\mathcal{V} 是所有节点(原子)的集合。

4.2.2. 预备知识:多头注意力机制 (Multi-head attention mechanism)

多头注意力机制是 Transformer 的核心组成部分,它包含多个并行的注意力层。对于查询 (queries) Q\mathbf{Q}、键 (keys) K\mathbf{K} 和值 (values) V\mathbf{V},缩放点积注意力 (scaled dot-product attention) 计算如下: Attention(Q,K,V)=softmax(QKTdk)V { \mathrm{Attention} } ( \mathbf { Q } , \mathbf { K } , \mathbf { V } ) = { \mathrm { softmax } } \left( { \frac { \mathbf { Q } \mathbf { K } ^ { T } } { { \sqrt { d _ { k } } } } } \right) \mathbf { V } 其中 dkd_k 是键的维度。多头注意力机制将多个注意力层的输出拼接并线性投影: ultiHead˙(Q,K,V)=Concat(head1,,headh)WO \mathbf { \dot { \mathbf { u l t i H e a d } } } ( \mathbf { Q } , \mathbf { K } , \mathbf { V } ) = \mathbf { C o n c a t } \left( \mathrm { head } _ { 1 } , \dots , \mathrm { head } _ { h } \right) \mathbf { W } ^ { O } 每个头 headi\mathrm{head}_i 的计算方式为: headi=Attention(QWiQ,KWiK,VWiV), \begin{array} { r } { \mathbf { h e a d } _ { i } = \mathbf { A t t e n t i o n } \left( \mathbf { Q } \mathbf { W } _ { i } ^ { \mathbf { Q } } , \mathbf { K } \mathbf { W } _ { i } ^ { \mathbf { K } } , \mathbf { V } \mathbf { W } _ { i } ^ { \mathbf { V } } \right) , } \end{array} 其中 WiQ,WiK,WiV\mathbf{W}_i^{\mathbf{Q}}, \mathbf{W}_i^{\mathbf{K}}, \mathbf{W}_i^{\mathbf{V}} 是第 ii 个头的投影权重矩阵。

4.2.3. BatmanNet 框架概述

BatmanNet 是一个双分支模型,包含一个节点分支 (node branch) 和一个边分支 (edge branch)。每个分支都专注于从输入图学习节点或边的嵌入。它采用 Transformer 风格的不对称编码器-解码器架构:

  • 编码器 (Encoder): 接收分子图的部分可见信号(未掩蔽的节点和边),并将其嵌入到节点或边的潜在表示中。

  • 解码器 (Decoder): 接收编码器输出的节点和边潜在表示,以及掩蔽标记,然后重建原始分子。解码器比编码器更轻量,仅在预训练阶段用于分子重建任务。在下游任务中,只使用编码器生成分子表示。

    对于一个分子,节点集合表示为 VV,边集合表示为 EE

  • 节点图 (Node Graph) GN=(V,E)G_N = (V, E) 原子 vVv \in V 是节点,键 (u,v)E(u, v) \in E 是边。

  • 边图 (Edge Graph) GEG_E 是节点图的主要对偶 (primary dual),描述原始图中的相邻边,以相似方式在边之间进行消息传递。 初始节点特征为 XvX_v,初始边特征为 euve_{uv}。节点图 GNG_N 和边图 GEG_E 分别作为 BatmanNet 节点分支和边分支的输入。

下图(原文 Figure 2)展示了BatmanNet的整体框架:

该图像是比支分masked图变压器自编码器BatmanNet的示意图。该图展示了节点和边嵌入的过程,包括特征重排序、聚合以及多头注意力机制的使用,突出BatmanNet编码器和解码器的结构设计与功能。 该图像是比支分masked图变压器自编码器BatmanNet的示意图。该图展示了节点和边嵌入的过程,包括特征重排序、聚合以及多头注意力机制的使用,突出BatmanNet编码器和解码器的结构设计与功能。

4.2.4. 编码器和解码器的详细设计

编码器和解码器分别由 NN 层和 MM 层(其中 MNM \ll N)相同的 GNN-Attention 块堆叠而成。每个 GNN-Attention 块采用双层信息提取框架,其中 GNN 用于提取局部信息,注意力层用于捕获全局信息。

在一个 GNN-Attention 块中,GNN 作为信息提取网络的第一层,对输入图执行消息传递操作以提取局部信息,输出学习到的嵌入。随后,采用多头注意力层来捕获图的全局信息。具体而言,一个 GNN-Attention 块包含三个 GNNs,即 GQ()\mathbf{G}_\mathbf{Q}(\cdot)GK()\mathbf{G}_\mathbf{K}(\cdot)GV()\mathbf{G}_\mathbf{V}(\cdot),它们学习查询 Q\mathbf{Q}、键 K\mathbf{K} 和值 V\mathbf{V} 的嵌入: Q=GQ(H),K=GK(H),V=GV(H), \begin{array} { r } { \mathbf { Q } = \mathbf { G } _ { \mathbf { Q } } ( \mathbf { H } ) , } \\ { \mathbf { K } = \mathbf { G } _ { \mathbf { K } } ( \mathbf { H } ) , } \\ { \mathbf { V } = \mathbf { G } _ { \mathbf { V } } ( \mathbf { H } ) , } \end{array} 其中 HRn×d\mathbf{H} \in \mathbb{R}^{n \times d}nn 个节点的隐藏表示矩阵,嵌入大小为 dd。然后应用上述多头注意力公式 (4), (5), (6) 得到 GNN-Attention 块的最终输出。

编码器:

  • 在编码器开始时,使用线性投影层,并添加位置嵌入 (positional embeddings) 来保留未掩蔽节点和边的位置信息。这里采用 Transformer [48] 中提出的绝对正弦位置编码 (absolute sinusoidal positional encoding)。节点和边在输入图中的位置在掩蔽前通过 RDkit 进行索引。
  • 通过将原始图和其对偶图分别输入到编码器的两个分支,得到聚合的节点嵌入 mv\mathbf{m}_v 和边嵌入 mvw\mathbf{m}_{vw}mv=uNvhu, \mathbf { m } _ { v } = \sum _ { u \in \mathcal { N } _ { v } } \mathbf { h } _ { u } , mvw=uNvwhuv. \mathbf { m } _ { v w } = \sum _ { u \in \mathcal { N } _ { v } \setminus w } \mathbf { h } _ { u v } . 其中 Nv\mathcal{N}_v 是节点 vv 的邻居集合,Nvw\mathcal{N}_v \setminus w 表示移除节点 ww 后的邻居集合。
  • 为了克服消息传递阶段的梯度消失 (vanishing gradient) 和过平滑 (over-smoothing) 问题,添加了从节点和边的初始特征到 mv\mathbf{m}_vmvw\mathbf{m}_{vw} 的长距离残差连接 (long-range residual connections)。
  • 最后,应用一个前馈网络 (Feed Forward) 和层归一化 (LayerNorm) 得到未掩蔽节点嵌入和边嵌入作为编码器的最终输出。

解码器:

  • 在解码器开始时,首先使用一个特征重排序层 (Feature Reordering layer)。该层将编码器输出的未掩蔽节点和边嵌入,以及被移除节点和边的掩蔽标记拼接起来,并通过添加相应的位置嵌入来恢复它们在原始输入图中的顺序。
  • 之后,解码器使用与编码器相同但更轻量级的 Transformer 风格架构,以获得节点嵌入和边嵌入。

4.2.5. 重建目标和预训练策略

4.2.5.1. 双分支图掩蔽 (Bi-branch graph masking) 预训练策略

该策略通过单一的预测任务,即双分支图掩蔽和重建,来同时学习分子图的节点和图级别的局部和全局信息。

  • 高比例随机掩蔽: 给定一个分子图,该方法在模型的节点分支和边分支中,分别随机掩蔽高比例(例如 60%)的节点和边。编码器仅在剩余的未掩蔽节点和边上操作。
  • 有向掩蔽 (directed masking): 考虑到 GNNs 中的消息传递过程是有向的,对边的随机掩蔽采用有向掩蔽方案,即移除 (u, v) 不意味着 (v, u) 也被移除。为了区分 (u, v)(v, u),将起始节点(头节点)的特征添加到边的初始特征中。
  • 有效性:
    • 节点级学习: 高比例的随机掩蔽使得每个节点/边很有可能同时缺失邻居节点和边。为了重建这些缺失部分,每个节点和边嵌入必须学习其局部的上下文信息,并且这种方法移除了子图规模和形状的限制,从而促进了超出 kk 跳范围和有限形状的局部上下文信息的捕获。

    • 图级学习: 从随机掩蔽后剩余的节点和边预测整个图,这比其他通常以较小图或基序为目标来学习全局图信息的自监督预训练任务更具挑战性。这种挑战性任务需要一个更强大的模型,具有更大的容量来学习高质量的节点和边嵌入,从而捕获局部和全局级别的分子信息。

      下图(原文 Figure 1)展示了BatmanNet的预训练,通过从潜在表示和掩蔽标记重建原始分子。

      该图像是示意图,展示了BatmanNet模型中节点和边的掩蔽过程。上半部分描述了节点掩蔽策略,通过编码器将初始节点嵌入转化为节点嵌入。下半部分展示了边掩蔽策略,初始边嵌入通过编码器处理得到边嵌入。 该图像是示意图,展示了BatmanNet模型中节点和边的掩蔽过程。上半部分描述了节点掩蔽策略,通过编码器将初始节点嵌入转化为节点嵌入。下半部分展示了边掩蔽策略,初始边嵌入通过编码器处理得到边嵌入。

4.2.5.2. 重建目标

BatmanNet 的节点分支和边分支通过预测掩蔽节点和边的所有特征来重建分子。节点(原子)和边(键)的特征参考补充材料中表 S1。 每个解码器输出后都连接一个线性层,其输出维度设置为原子(节点分支)或键(边分支)特征的总大小。节点和边的重建任务都涉及高维多标签预测,这可以缓解 [18] 中讨论的歧义问题,即使用有限数量的原子或边类型作为节点/边级别的预训练目标可能导致歧义。预训练损失仅在掩蔽标记上计算,类似于 MAE [41]。

最终的预训练损失 Lpretrain\mathcal{L}_{\mathrm{pre-train}} 定义为: Lpretrain=Lnode+Ledge, { \mathcal { L } } _ { \mathrm { p r e - t r a i n } } = { \mathcal { L } } _ { \mathrm { n o d e } } + { \mathcal { L } } _ { \mathrm { e d g e } } , 其中:

  • Lnode\mathcal{L}_{\mathrm{node}} 是节点分支的损失函数。

  • Ledge\mathcal{L}_{\mathrm{edge}} 是边分支的损失函数。

    节点损失 Lnode\mathcal{L}_{\mathrm{node}} 计算如下: Lnode=vVmaskLce(pv,yv), \mathcal { L } _ { \mathrm { n o d e } } = \sum _ { v \in \mathcal { V } _ { \mathrm { m a s k } } } \mathcal { L } _ { \mathrm { c e } } \left( \pmb { p } _ { v } , \pmb { y } _ { v } \right) , 其中:

  • Vmask\mathcal{V}_{\mathrm{mask}}:被掩蔽的节点集合。

  • Lce\mathcal{L}_{\mathrm{ce}}:交叉熵损失 (cross entropy loss)。

  • pv\pmb{p}_v:节点 vv 的预测特征。

  • yv\pmb{y}_v:节点 vv 对应的真实 (ground-truth) 特征。

    边损失 Ledge\mathcal{L}_{\mathrm{edge}} 计算如下: Ledge=(u,v)EmaskLce(p(u,v),y(u,v)), \mathcal { L } _ { \mathrm { e d g e } } = \sum _ { ( u , v ) \in \mathcal { E } _ { \mathrm { m a s k } } } \mathcal { L } _ { \mathrm { c e } } \left( p _ { ( u , v ) } , \pmb { y } _ { ( u , v ) } \right) , 其中:

  • Emask\mathcal{E}_{\mathrm{mask}}:被掩蔽的边集合。

  • p(u,v)p_{(u, v)}:边 (u, v) 的预测特征。

  • y(u,v)\pmb{y}_{(u, v)}:边 (u, v) 对应的真实特征。

4.2.6. 微调 (Fine-tuning) 细节

在下游任务中,只使用 BatmanNet 的编码器。与预训练时输入不完整分子不同,下游任务的输入是完整的、未掩蔽的分子。 经过 NNGNN-Attention 块后,BatmanNet 编码器的两个分支执行节点聚合 (Node Aggregation),分别产生两个节点表示 mvnodebranch\mathbf{m}_v^{\mathrm{node-branch}}mvedgebranch\mathbf{m}_v^{\mathrm{edge-branch}}mvnodebranch=uNvhu, \mathbf { m } _ { v } ^ { \mathrm { n o d e - b r a n c h } } = \sum _ { u \in \mathcal { N } _ { v } } \overline { { \mathbf { h } } } _ { u } , mvedgebranch=uNvwhuv, \mathbf { m } _ { v } ^ { \mathrm { e d g e - b r a n c h } } = \sum _ { u \in \mathcal { N } _ { v } \setminus w } \overline { { \mathbf { h } } } _ { u v } , 其中 hu\overline{\mathbf{h}}_uhuv\overline{\mathbf{h}}_{uv} 是节点分支和边分支的 GNN-Attention 块的隐藏状态。接着,应用一个长距离残差连接,将 mvnodebranch\mathbf{m}_v^{\mathrm{node-branch}}mvedgebranch\mathbf{m}_v^{\mathrm{edge-branch}} 分别与初始节点特征和边特征拼接起来。最后,通过前馈层 (Feed Forward layers) 和层归一化 (LayerNorm) 转换这两个嵌入,生成用于下游任务的最终嵌入。

给定一个分子 GiG_i 及其对应标签 yi\pmb{y}_iBatmanNet 的编码器从节点分支和边分支分别生成两个节点嵌入 Hinodebranch\mathbf{H}_i^{\mathrm{node-branch}}Hiedgebranch\mathbf{H}_i^{\mathrm{edge-branch}}。 参照 GROVER [18],将这两个节点嵌入输入到一个共享的自注意力 READOUT 函数中,生成两个图级别的嵌入 gnodebranch\pmb{g}^{\mathrm{node-branch}}gedgebranch\pmb{g}^{\mathrm{edge-branch}}。它们的计算方式如下: S=softmax(W2tanh(W1H)), \mathbf { S } = \mathrm { softmax } \left( \mathbf { W } _ { 2 } \operatorname { tanh } \left( \mathbf { W } _ { 1 } \mathbf { H } ^ { \top } \right) \right) , g=Flatten(SH), \begin{array} { r } { \pmb { g } = \mathrm { Flatten } ( \mathbf { S } \mathbf { H } ) , } \end{array} 其中:

  • S\mathbf{S}:注意力分数矩阵。
  • H\mathbf{H}:节点嵌入矩阵(可以是 Hinodebranch\mathbf{H}_i^{\mathrm{node-branch}}Hiedgebranch\mathbf{H}_i^{\mathrm{edge-branch}})。
  • W1Rdattn_hidden×dhidden_size\mathbf{W}_1 \in \mathbb{R}^{d_{\mathrm{attn\_hidden}} \times d_{\mathrm{hidden\_size}}}W2Rdattn_out×dattn_hidden\mathbf{W}_2 \in \mathbb{R}^{d_{\mathrm{attn\_out}} \times d_{\mathrm{attn\_hidden}}} 是两个权重矩阵。
  • dattn_hiddend_{\mathrm{attn\_hidden}}:注意力隐藏层维度。
  • dhidden_sized_{\mathrm{hidden\_size}}:隐藏状态维度。
  • dattn_outd_{\mathrm{attn\_out}}:注意力输出维度。
  • Flatten()\mathrm{Flatten}(\cdot):展平操作。

分子性质预测 (Molecular properties prediction): 获得两个图级嵌入 gnodebranch\pmb{g}^{\mathrm{node-branch}}gedgebranch\pmb{g}^{\mathrm{edge-branch}} 后,对两个分支应用前馈层得到预测结果 pinodebranch\pmb{p}_i^{\mathrm{node-branch}}piedgebranch\pmb{p}_i^{\mathrm{edge-branch}}pinodebranch=f(Wgnodebranch+b), \begin{array} { r } { \pmb { p } _ { i } ^ { \mathrm { n o d e \mathrm { - } b r a n c h } } = f \left( W \pmb { g } ^ { \mathrm { n o d e \mathrm { - } b r a n c h } } + b \right) , } \end{array} piedgebranch=f(Wgedgebranch+b), \begin{array} { r } { p _ { i } ^ { \mathrm { e d g e - b r a n c h } } = f \left( W g ^ { \mathrm { e d g e - b r a n c h } } + b \right) , } \end{array} 其中 f()f(\cdot) 是激活函数,WW 是权重矩阵,bb 是偏置向量。

DDI 预测 (DDI prediction): 对于 DDI 任务,输入是两个分子对。它们通过 BatmanNet 编码器被编码成两组图级嵌入 (g1nodebranch,g1edgebranch)(\pmb{g}_1^{\mathrm{node-branch}}, \pmb{g}_1^{\mathrm{edge-branch}})(g2nodebranch,g2edgebranch)(\pmb{g}_2^{\mathrm{node-branch}}, \pmb{g}_2^{\mathrm{edge-branch}})。 预测结果 pinodebranch\pmb{p}_i^{\mathrm{node-branch}}piedgebranch\pmb{p}_i^{\mathrm{edge-branch}} 的计算方式如下: pinodebranch=f(Wgpairnodebranch+b), p _ { i } ^ { \mathrm { n o d e - b r a n c h } } = f \left( W g _ { p a i r } ^ { \mathrm { n o d e - b r a n c h } } + b \right) , piedgebranch=f(Wgpairedgebranch+b), p _ { i } ^ { \mathrm { e d g e - b r a n c h } } = f \left( W g _ { p a i r } ^ { \mathrm { e d g e - b r a n c h } } + b \right) , 其中,pair 级别的图嵌入 gpairnodebranch\pmb{g}_{\mathrm{pair}}^{\mathrm{node-branch}}gpairedgebranch\pmb{g}_{\mathrm{pair}}^{\mathrm{edge-branch}} 是通过拼接两个分子的图嵌入得到的: gpairnodebranch=Concat(g1nodebranch,g2nodebranch), \begin{array} { r } { \pmb { g } _ { p a i r } ^ { \mathrm { n o d e - b r a n c h } } = \mathrm { Concat } \left( \pmb { g } _ { 1 } ^ { \mathrm { n o d e - b r a n c h } } , \pmb { g } _ { 2 } ^ { \mathrm { n o d e - b r a n c h } } \right) , } \end{array} gpairedgebranch=Concat(g1edgebranch,g2edgebranch), g _ { p a i r } ^ { \mathrm { e d g e - b r a n c h } } = \mathrm { Concat } \left( g _ { 1 } ^ { \mathrm { e d g e - b r a n c h } } , g _ { 2 } ^ { \mathrm { e d g e - b r a n c h } } \right) ,

DTI 预测 (DTI prediction): 在 DTI 任务中,作者修改了 Tsubaki et al. [65] 提出的 DTI 预测框架,用 BatmanNet 编码器替换了其化合物编码器(GNN)。蛋白质序列编码器是一个 CNN 模型,它使用注意力机制来导出蛋白质序列表示 yp\pmb{y}_p。给定蛋白质中子序列的隐藏向量集合 S=(s1,s2,...,sn)\mathcal{S} = (s_1, s_2, ..., s_n)yp\pmb{y}_p 的计算方式如下: yp=in(αihi), { \bf \it y } _ { p } = \sum _ { i } ^ { n } \left( \alpha _ { i } h _ { i } \right) , αi=σ(hmThi), \pmb { \alpha } _ { i } = \sigma \left( \pmb { h } _ { m } ^ { T } \pmb { h } _ { i } \right) , hm=f(Wgm+b), \pmb { h } _ { m } = f \left( \pmb { W } \pmb { g } _ { m } + b \right) , hi=f(Wsi+b) \pmb { h } _ { i } = f \left( \pmb { W } s _ { i } + b \right) 其中:

  • yp\pmb{y}_p:蛋白质序列表示。

  • sis_i:蛋白质中第 ii 个子序列的隐藏向量。

  • αi\pmb{\alpha}_i:注意力权重,表示 sis_iyp\pmb{y}_p 的贡献。

  • hm\pmb{h}_m:一个基于分子向量 gm\pmb{g}_m 转换而来的隐藏状态。

  • hi\pmb{h}_i:一个基于子序列隐藏向量 sis_i 转换而来的隐藏状态。

  • σ()\sigma(\cdot):激活函数。

  • W\pmb{W}:权重矩阵。

  • bb:偏置向量。

    与分子性质预测类似,通过 BatmanNet 编码器得到两个分子嵌入 gnodebranch\pmb{g}^{\mathrm{node-branch}}gedgebranch\pmb{g}^{\mathrm{edge-branch}}。然后得到两个蛋白质嵌入 ypnodebranch\pmb{y}_p^{\mathrm{node-branch}}ypedgebranch\pmb{y}_p^{\mathrm{edge-branch}}。最终的预测结果 pinodebranch\pmb{p}_i^{\mathrm{node-branch}}piedgebranch\pmb{p}_i^{\mathrm{edge-branch}} 计算如下: pinodebranch=f(Wypairnodebranch+b), p _ { i } ^ { \mathrm { n o d e - b r a n c h } } = f \left( W y _ { p a i r } ^ { \mathrm { n o d e - b r a n c h } } + b \right) , ypairnodebranch=Concat(gnodebranch,ypnodebranch), { \pmb y } _ { p a i r } ^ { \mathrm { n o d e - b r a n c h } } = \mathrm { Concat } \left( { \pmb g } ^ { \mathrm { n o d e - b r a n c h } } , { \pmb y } _ { p } ^ { \mathrm { n o d e - b r a n c h } } \right) , piedgebranch=f(Wypairedgebranch+b), \begin{array} { r } { p _ { i } ^ { \mathrm { e d g e - b r a n c h } } = f \left( W y _ { p a i r } ^ { \mathrm { e d g e - b r a n c h } } + b \right) , } \end{array} ypairedgebranch=Concat(gedgebranch,ypedgebranch), { \pmb y } _ { p a i r } ^ { \mathrm { e d g e - b r a n c h } } = \mathrm { Concat } \left( { \pmb g } ^ { \mathrm { e d g e - b r a n c h } } , { \pmb y } _ { p } ^ { \mathrm { e d g e - b r a n c h } } \right) , 其中,pair 级别的嵌入 ypairnodebranch\pmb{y}_{\mathrm{pair}}^{\mathrm{node-branch}}ypairedgebranch\pmb{y}_{\mathrm{pair}}^{\mathrm{edge-branch}} 分别通过拼接分子嵌入和蛋白质嵌入得到。

下游任务的最终损失函数 Lfinetune\mathcal{L}_{\mathrm{fine-tune}} 由监督损失 Lsup\mathcal{L}_{\mathrm{sup}} 和分歧损失 (disagreement loss) Ldiss\mathcal{L}_{\mathrm{diss}} 组成。分歧损失旨在训练两个分支的预测结果保持一致: Lfinetune=Lsup+Ldiss, { \mathcal { L } } _ { \mathrm { f i n e - t u n e } } = { \mathcal { L } } _ { \mathrm { s u p } } + { \mathcal { L } } _ { \mathrm { d i s s } } , 其中: Lsup=L(pinodebranch,yi)+L(piedgebranch,yi), \begin{array} { r } { \mathcal { L } _ { \mathrm { s u p } } = \mathcal { L } \left( p _ { i } ^ { \mathrm { n o d e - b r a n c h } } , \pmb { y } _ { i } \right) + \mathcal { L } \left( p _ { i } ^ { \mathrm { e d g e - b r a n c h } } , \pmb { y } _ { i } \right) , } \end{array} Ldiss=pinodebranchpiedgebranch2. \mathcal { L } _ { \mathrm { d i s s } } = \left. p _ { i } ^ { \mathrm { n o d e - b r a n c h } } - p _ { i } ^ { \mathrm { e d g e - b r a n c h } } \right. _ { 2 } .

  • L(,)\mathcal{L}(\cdot, \cdot) 是根据任务类型(分类或回归)选择的监督损失函数。
  • 2\|\cdot\|_2 表示 L2 范数。

5. 实验设置

5.1. 数据集

为了全面评估 BatmanNet 的性能,实验在广泛的下游药物发现任务基准上进行,包括分子性质预测、药物-药物相互作用 (DDI) 预测和药物-靶点相互作用 (DTI) 预测。

5.1.1. 预训练数据集

BatmanNet 的预训练是在 ZINC250K 分子数据集 [55] 上进行的。该数据集包含从 ZINC 数据库 [56] 采样的 250K 个分子。数据集以 9:1 的比例随机划分为训练集和验证集。

5.1.2. 下游任务数据集

论文在以下数据集上进行了分子性质预测、DDI 预测和 DTI 预测任务的评估。所有分子性质预测任务均采用支架分割 (scaffold splitting) 方法将数据集按 8:1:1 的比例划分为训练集、验证集和测试集。这种方法根据分子的结构将数据集分离,能够评估模型在分布外 (out-of-distribution) 的泛化能力。

分子性质预测数据集 (来自 MoleculeNet [60]):

  • BBBP [69]: 记录化合物是否具有穿透血脑屏障 (blood-brain barrier) 的渗透性。
  • SIDER [70]: 一个包含上市药物和不良药物反应 (ADR) 的数据库,分为 27 个系统器官类别。
  • ClinTox [71]: 比较通过 FDA 批准的药物和在临床试验中因毒性被淘汰的药物。
  • BACE [72]: 提供一组人 β\beta-分泌酶 1 (BACE-1) 抑制剂的定量盲测结果。
  • Tox21 [73]: 一个公共数据库,测量化合物对 12 个不同靶点(包括核受体和应激反应)的毒性。
  • ToxCast [74]: 提供基于体外高通量筛选的 8615 种化合物的毒理学数据。
  • FreeSolv [75]: 提供小分子在水中实验和计算的水合自由能。计算值来自使用分子动力学模拟的炼金术自由能计算。
  • ESOL [76]: 一个小型数据集,包含 1128 种化合物的水溶性数据。
  • Lipo [77]: 从 ChEMBL 数据库中整理而来,是药物分子影响膜渗透性和溶解度的重要特征,提供 4200 种化合物的辛醇/水分配系数(pH 7.4 下的 log D)实验结果。
  • QM7 [78]: GDB-13 数据库的一个子集,包含多达七个“重”原子(C、N、O、S)的近 10 亿个稳定且可合成的有机分子。
  • QM8 [79]: 应用于包含多达八个重原子的分子集合(也是 GDB-17 数据库的子集)。它包含计算机生成的量子力学性质。

DDI 预测数据集:

  • BIOSNAP [67]: 包含 1322 种批准药物和 41520 个标记的 DDI,数据来源于药物标签和科学出版物。
  • TWOSIDES [68]: 包含药物组合引起的副作用,共 548 种药物和 48584 对药物-药物相互作用。

DTI 预测数据集:

  • Human 和 C. elegans:Liu et al. 创建,通过系统框架包含了高度可信的化合物-蛋白质对负样本。正样本数据从 DrugBank 4.1 和 Matador 中检索。参照 Tsubaki et al. [65] 和 MPG [13],使用 1:1 比例的正负样本平衡数据集。

5.1.3. 原子和键特征

BatmanNet 使用 RDKit 提取原子和键特征作为 GNN 的输入和 BatmanNet 的重建目标。

以下是原文 [Table S1] 的结果:

Features Size Description
Atom Atom type 23 The atom type (e.g., C, N, O), by atomic number
Number of H 6 The number of bonded hydrogen atoms
Charge 5 The formal charge of the atom
Chirality 4 The chiral-tag of the atom
Is-aromatic 1 Whether the atom is part of an aromatic system or not
Bond Bond type 5 The bond type (e.g., single, double, triple et al.)
Stereo 6 The stereo-configuration of the bond

5.2. 评估指标

论文使用了多种评估指标来衡量模型在不同任务上的性能。

5.2.1. ROC-AUC (Receiver Operating Characteristic - Area Under the Curve)

  • 概念定义: ROC-AUC 是一个用于评估二分类模型性能的指标,它表示 ROC 曲线下方区域的面积。ROC 曲线描绘了在不同分类阈值下,真阳性率 (True Positive Rate, TPR) 与假阳性率 (False Positive Rate, FPR) 之间的关系。AUC 值越高,模型的分类性能越好,因为它表明模型在区分正负样本方面具有更好的能力。
  • 数学公式: AUC=01TPR(FPR1(x))dx \text{AUC} = \int_{0}^{1} \text{TPR}(\text{FPR}^{-1}(x)) dx 其中, TPR=TPTP+FN \text{TPR} = \frac{\text{TP}}{\text{TP} + \text{FN}} FPR=FPFP+TN \text{FPR} = \frac{\text{FP}}{\text{FP} + \text{TN}}
  • 符号解释:
    • TP\text{TP} (True Positive):真阳性,实际为正类且被模型预测为正类的样本数。
    • FN\text{FN} (False Negative):假阴性,实际为正类但被模型预测为负类的样本数。
    • FP\text{FP} (False Positive):假阳性,实际为负类但被模型预测为正类的样本数。
    • TN\text{TN} (True Negative):真阴性,实际为负类且被模型预测为负类的样本数。
    • TPR\text{TPR} (True Positive Rate) 或召回率 (Recall):真阳性率,表示在所有实际正类样本中,被正确预测为正类的比例。
    • FPR\text{FPR} (False Positive Rate):假阳性率,表示在所有实际负类样本中,被错误预测为正类的比例。

5.2.2. RMSE (Root Mean Squared Error)

  • 概念定义: RMSE 是回归任务中常用的评估指标,用于衡量模型预测值与真实值之间的差异。它计算预测误差平方的均值的平方根,因此对大误差的惩罚更大。RMSE 值越小,模型的预测精度越高。
  • 数学公式: RMSE=1ni=1n(yiy^i)2 \text{RMSE} = \sqrt{\frac{1}{n} \sum_{i=1}^{n} (y_i - \hat{y}_i)^2}
  • 符号解释:
    • nn:样本总数。
    • yiy_i:第 ii 个样本的真实值。
    • y^i\hat{y}_i:第 ii 个样本的预测值。

5.2.3. PR-AUC (Precision-Recall Area Under the Curve)

  • 概念定义: PR-AUC 也是一个用于二分类模型性能评估的指标,它表示精确率 (Precision) 与召回率 (Recall) 曲线下方区域的面积。当数据集中正负样本不平衡时,PR-AUC 往往比 ROC-AUC 能更好地反映模型的性能,因为它更关注正样本的预测表现。PR-AUC 值越高,模型在识别正样本方面的性能越好。
  • 数学公式: PR-AUC=01Precision(Recall1(x))dx \text{PR-AUC} = \int_{0}^{1} \text{Precision}(\text{Recall}^{-1}(x)) dx 其中, Precision=TPTP+FP \text{Precision} = \frac{\text{TP}}{\text{TP} + \text{FP}} Recall=TPTP+FN \text{Recall} = \frac{\text{TP}}{\text{TP} + \text{FN}}
  • 符号解释:
    • TP\text{TP} (True Positive):真阳性。
    • FP\text{FP} (False Positive):假阳性。
    • FN\text{FN} (False Negative):假阴性。
    • Precision\text{Precision} (精确率):表示在所有被模型预测为正类的样本中,实际为正类的比例。
    • Recall\text{Recall} (召回率) 或 TPR:表示在所有实际正类样本中,被正确预测为正类的比例。

5.2.4. F1 (F1-score)

  • 概念定义: F1 分数是精确率 (Precision) 和召回率 (Recall) 的调和平均值 (harmonic mean)。它综合考虑了模型的精确率和召回率,当模型在这两者之间取得良好平衡时,F1 分数较高。尤其适用于分类任务中正负样本不平衡的情况。
  • 数学公式: F1=2PrecisionRecallPrecision+Recall \text{F1} = 2 \cdot \frac{\text{Precision} \cdot \text{Recall}}{\text{Precision} + \text{Recall}}
  • 符号解释:
    • Precision\text{Precision} (精确率):如上所述。
    • Recall\text{Recall} (召回率):如上所述。

5.2.5. Precision (精确率)

  • 概念定义: 精确率表示在所有被模型预测为正类的样本中,实际为正类的比例。它衡量了模型预测正类的准确性。高精确率意味着模型在预测为正类时犯的错误较少。
  • 数学公式: Precision=TPTP+FP \text{Precision} = \frac{\text{TP}}{\text{TP} + \text{FP}}
  • 符号解释:
    • TP\text{TP} (True Positive):真阳性。
    • FP\text{FP} (False Positive):假阳性。

5.2.6. Recall (召回率)

  • 概念定义: 召回率表示在所有实际为正类的样本中,被模型正确预测为正类的比例。它衡量了模型识别所有实际正类的能力。高召回率意味着模型能够找出大部分的正类样本。
  • 数学公式: Recall=TPTP+FN \text{Recall} = \frac{\text{TP}}{\text{TP} + \text{FN}}
  • 符号解释:
    • TP\text{TP} (True Positive):真阳性。
    • FN\text{FN} (False Negative):假阴性。

5.3. 对比基线

论文将 BatmanNet 与以下两类共 20 种具有竞争力的基线模型进行了比较:

5.3.1. 无预训练的监督学习模型

这类模型直接在标记数据上进行监督训练,不进行大规模无标签数据的预训练。

  • ECFP [50]: 扩展连通性指纹 (Extended-Connectivity Fingerprints),一种用于分子表征、相似性搜索和结构-活性建模的圆形拓扑指纹。
  • TF_Robust [51]: 一个基于深度神经网络 (DNN) 的多任务框架,以分子指纹作为输入。
  • GraphConv [52]: 一种图卷积模型。
  • Weave [31]: 一种图卷积模型。
  • SchNet [32]: 一种图卷积模型。
  • MPNN [37]: 消息传递神经网络 (Message Passing Neural Network),一种在消息传递过程中考虑边特征的模型。
  • DMPNN [39]: 导向消息传递神经网络 (Directed Message Passing Neural Network),MPNN 的变体。
  • MGCN [40]: 多级图卷积网络 (Multi-level Graph Convolutional Network),一种分层 GNN,用于捕获多级相互作用。
  • AttentiveFP [35]: 注意力指纹 (Attentive Fingerprint),图注意力网络 (GAT) 的扩展。
  • TrimNet [53]: 一种基于图的方法,采用新颖的三元组消息机制来高效学习分子表示。

5.3.2. 有预训练的自监督学习模型

这类模型通过自监督学习在大规模无标签数据上进行预训练,然后微调 (fine-tune) 到下游任务。

  • Mol2Vec [54]: 受 NLP 启发,对分子序列表示进行预训练的模型。
  • N-GRAM [43]: 受 NLP 启发,对分子序列表示进行预训练的模型。
  • SMILES-BERT [14]: 受 NLP 启发,对分子序列表示进行预训练的模型,使用 SMILES 字符串。
  • pre-trainGNN [17]: 一种图基预训练模型,采用各种预训练策略。
  • GraphMAE [47]: 图掩蔽自编码器 (Graph Masked Autoencoder) 用于自监督学习。
  • GROVERbase [18] 和 GROVERlarge [18]: 自监督图 Transformer,基于大规模分子数据。
  • KPGT [19]: 知识引导的图 Transformer 预训练方法,用于分子性质预测。
  • MPG [13]: 有效的自监督框架,用于学习表达性分子全局表示。
  • GEM [20]: 几何增强的分子表示学习模型,具有专门的几何级别自监督学习策略。

5.4. 实验配置

5.4.1. 预训练配置

  • 优化器: Adam [57] 优化器。

  • 学习率调度器: Noam 学习率调度器 [42]。

  • 批量大小 (Batch size): 32。

  • 掩蔽比例 (Masking ratio): BatmanNet 两个分支的掩蔽比例均设置为 0.6。

  • 模型层数: 编码器由 6 层组成,解码器由 2 层组成。

  • 隐藏层大小 (Hidden size): 100。

  • GNN-Attention 块: 每个 GNN-Attention 块包含 3 个 GNN 层和 2 个自注意力头 (self-attention heads)。

  • 模型参数量: 自编码器结构约包含 2.6 百万 (M) 参数。

  • 预训练时间: 在单张 Nvidia RTX3090 GPU 上预训练两天。

    以下是原文 [Table S2] 的结果:

    Hyper-parameter Value Description
    batch_size 32 The input batch_size
    hidden_size 100 The hidden _size of encoder and decoder
    depth 3 The number of GNN layers in GNN-Attention block
    num_enc_mt_block 6 The number of the GNN-Attention block in encoder
    num_dec_mt_block 2 The number of the GNN-Attention block in decoder
    num_attn_head 2 The number of attention heads in the GNN-Attention block
    mask_ratio 0.6 The mask ratio
    init_lr 0.0002 The initial learning rate of Noam learning rate schedular
    max_lr 0.0004 The maximum learning rate of Noam learning rate schedular
    final_lr 0.0001 The final learning rate of Noam learning rate schedular

5.4.2. 微调配置

  • 模型选择: 选择在验证集 (validation set) 上表现最佳的模型。

  • 独立运行: 每个属性预测任务进行三次独立运行,报告 ROC-AUCRMSE 的平均值和标准差。

  • 超参数搜索: 对于每个任务,通过随机搜索 (random search) 尝试不同的超参数组合以找到最佳结果。

    以下是原文 [Table S3] 的结果:

    Hyper-parameter Value Description
    batch_size 32 The input batch size
    ffn_hidden_size 200 The hidden_size of MLP layers
    ffn_num_layer 2 The number of MLP layers
    attn_hidden 200 The hidden_size for the self-attentive readout
    attn_out 2 The number of output heads for the self-attentive readout
    dist_coff 0.1 The coefficient of the disagreement loss
    init_lr max_lr / 10 The initial learning rate of Noam learning rate schedular
    max_lr 0.0001 ~ 0.001 The maximum learning rate of Noam learning rate schedular
    final_lr max_lr / (5-10) The final learning rate of Noam learning rate schedular

5.5. 预训练表示可视化

为了直观观察自监督任务(未进行下游微调)所学到的表示,作者将其投影到二维空间进行可视化。

  • 方法:ZINC 数据集中随机选择 1,500 个具有有效结构的分子,并通过打乱原子特征、改变原子和键的顺序来生成无效分子,引入结构扰动。从预训练的 BatmanNet 最后一层提取嵌入作为分子表示,然后使用 UMAP 算法 [59] 将这些表示映射到二维空间。同时对未预训练的 BatmanNet 模型进行类似分析以进行比较。

    下图(原文 Figure 3)展示了非预训练与预训练分子表示的聚类情况。

    该图像是示意图,展示了非预训练(b)与预训练(a)的分子表示的聚类情况。左侧的图表中,蓝色点代表无效分子,红色点代表有效分子;右侧的图表则显示通过预训练后的分子聚类效果,有效分子的分布更为集中。 该图像是示意图,展示了非预训练(b)与预训练(a)的分子表示的聚类情况。左侧的图表中,蓝色点代表无效分子,红色点代表有效分子;右侧的图表则显示通过预训练后的分子聚类效果,有效分子的分布更为集中。

  • 结果: 预训练的 BatmanNet 表现出增强的区分有效分子和无效分子结构的能力,表明预训练模型能够有效辨别分子的结构有效性。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 分子性质预测

以下是原文 [Table 1] 的结果:

Dataset BACE BBBP ClinTox SIDER Tox21 ToxCast FreeSolv ESOL Lipo
#Molecules 1513 2039 1478 1427 7831 8575 642 1128 4200
#tasks 1 1 2 27 12 617 1 1 1
ECFP [50] 0.861(0.024) 0.783(0.050) 0.673(0.031) 0.630(0.019) 0.760(0.009) 0.615(0.017) 5.275(0.751) 2.359(0.454) 1.188(0.061)
TF_Robust [51] 0.824(0.022) 0.860(0.087) 0.765(0.085) 0.607(0.033) 0.698(0.012) 0.585(0.031) 4.122(0.085) 1.722(0.038) 0.909(0.060)
GraphConv [52] 0.854(0.011) 0.877(0.036) 0.845(0.051) 0.593(0.035) 0.772(0.041) 0.650(0.025) 2.900(0.135) 1.068(0.050) 0.712(0.049)
Weave [31] 0.791(0.008) 0.837(0.065) 0.823(0.023) 0.543(0.034) 0.741(0.044) 0.678(0.024) 2.398(0.250) 1.158(0.055) 0.813(0.042)
SchNet [32] 0.750(0.033) 0.847(0.024) 0.717(0.042) 0.545(0.038) 0.767(0.025) 0.679(0.021) 3.215(0.755) 1.045(0.064) 0.909(0.098)
MPNN [37] 0.815(0.044) 0.913(0.041) 0.879(0.054) 0.595(0.030) 0.808(0.024) 0.691(0.013) 2.185(0.952) 1.167(0.430) 0.672(0.051)
DMPNN [39] 0.852(0.053) 0.919(0.030) 0.897(0.040) 0.632(0.023) 0.826(0.023) 0.718(0.011) 2.177(0.914) 0.980(0.258) 0.653(0.046)
MGCN [40] 0.734(0.030) 0.850(0.064) 0.634(0.042) 0.552(0.018) 0.707(0.016) 0.663(0.009) 3.349(0.097) 1.266(0.147) 1.113(0.041)
AttentiveFP [35] 0.863(0.015) 0.908(0.050) 0.933(0.020) 0.605(0.060) 0.807(0.020) 0.579(0.001) 2.030(0.420) 0.853(0.060) 0.650(0.030)
TrimNet [53] 0.843(0.025) 0.892(0.025) 0.906(0.017) 0.606(0.006) 0.812(0.019) 0.652(0.032) 2.529(0.111) 1.282(0.029) 0.702(0.008)
Mol2Vec [54] 0.841(0.052) 0.876(0.030) 0.828(0.023) 0.601(0.023) 0.805(0.015) 0.690(0.052) 5.752(1.245) 2.358(0.452) 1.178(0.054)
N-GRAM [43] 0.876(0.035) 0.912(0.013) 0.855(0.037) 0.632(0.005) 0.769(0.027) 2.512(0.190) 1.100(0.160) 0.876(0.033)
SMILES-BERT [14] 0.849(0.021) 0.959(0.009) 0.985(0.014) 0.568(0.031) 0.803(0.010) 0.665(0.010) 2.974(0.510) 0.841(0.096) 0.666(0.029)
pre-trainGNN [17] 0.851(0.027) 0.915(0.040) 0.762(0.058) 0.614(0.006) 0.811(0.015) 0.714(0.019)
GraphMAE1 [47] 0.863(0.002) 0.896(0.007) 0.850(0.007) 0.652(0.001) 0.794(0.003) 0.679(0.005) - - -
GROVERbase [18] 0.878(0.016) 0.936(0.008) 0.925(0.013) 0.656(0.023) 0.819(0.020) 0.723(0.010) 1.592(0.072) 0.888(0.116) 0.563(0.030)
GROVERlarge [18] 0.894(0.028) 0.940(0.019) 0.944(0.021) 0.658(0.023) 0.831(0.025) 0.737(0.010) 1.544(0.397) 0.831(0.120) 0.560(0.035)
KPGT [19] 0.855(0.011) 0.908(0.010) 0.946(0.022) 0.649(0.009) 0.848(0.013) 0.746(0.002) 2.121(0.837) 0.803(0.008) 0.600(0.010)
MPG [13] 0.920(0.013) 0.922(0.012) 0.963(0.028) 0.661(0.007) 0.837(0.019) 0.748(0.005) 1.269(0.192) 0.741(0.017) 0.556(0.017)
GEM [20] 0.925(0.010) 0.953(0.007) 0.977(0.019) 0.663(0.014) 0.849(0.003) 0.742(0.004) - - -
BatmanNet 0.928(0.008) 0.946(0.003) 0.926(0.002) 0.676(0.007) 0.855(0.005) 0.756(0.007) 1.174(0.054) 0.736(0.014) 0.578(0.034)

结果分析:

  • 总体性能: BatmanNet 在 9 个数据集中的 6 个上取得了最先进 (SOTA) 的性能。
  • 对比无预训练模型: BatmanNet 在所有数据集上均显著优于分子指纹和无预训练的监督模型,这突显了自监督预训练的优势。
  • 对比自监督模型:
    • 分类任务 (Classification tasks): 在 6 个分类数据集中的 4 个上,BatmanNet 超越了先前的 SOTA 模型 GEM。例如,在 BACESIDERTox21 上取得了最佳的 ROC-AUC
    • 回归任务 (Regression tasks): 在 3 个回归数据集中的 2 个上,BatmanNet 优于先前的 SOTA 模型 MPG。例如,在 FreeSolvESOL 上取得了最低的 RMSE
  • 这些结果有力证明了 BatmanNet 在分子表示学习中的强大能力、有效性和泛化能力。

6.1.2. 效能与效率分析

下图(原文 Figure 4)展示了BatmanNet及一系列先进基线在分子属性预测上的平均AUC与预训练数据集大小的关系。

Fig. 4. Efficacy and Effectiveness Analysis. The figure illustrates the pre-training dataset size and model size for BatmanNet and a series of advanced baselines, along with their average AUC across all classification datasets about molecular property prediction. The horizontal axis represents the dataset size, the vertical axis represents the average AUC, the circle size corresponds to the model size, and the diamond indicates that the size of the model is unknown. 该图像是图表,展示了BatmanNet及一系列先进基线在分子属性预测上的平均AUC与预训练数据集大小的关系。横轴表示数据集大小(百万),纵轴表示平均AUC百分比,图中圆圈的大小对应模型大小,菱形表示模型大小未知。

以下是原文 [Table S4] 的结果:

Model Pre-training Data Size (M) Model Size (M) AVG-AUC /%
GraphMAE 2 - 78.90
GROVERbase 11 40 82.28
GROVERlarge 11 100 83.40
KPGT 2 - 82.53
MPG 11 55 84.18
GEM 20 85.15
BatmanNet 0.25 2.6 84.78

分析: 图 4 和表 S4 比较了 BatmanNet 与先进基线模型的预训练数据集大小、模型大小以及在所有分类数据集上的平均 AUC

  • BatmanNet 使用的预训练数据集规模最小 (0.25M),模型参数量也最小 (2.6M),但其平均 AUC (84.78%) 与之前的 SOTA 模型 GEM (85.15%) 相比具有可比性。
  • 这表明 BatmanNet 在利用更少的训练数据和模型参数的情况下,实现了卓越的效能和效率,解决了现有方法计算成本高昂和模型庞大的问题。

6.1.3. 药物-药物相互作用 (DDI) 预测

以下是原文 [Table 2] 的结果:

Model AUC-ROC PR-AUC F1
LR 0.802(0.001) 0.779(0.001) 0.741(0.002)
Nat.Prot [63] 0.853(0.001) 0.848(0.001) 0.714(0.001)
Mol2Vec [54] 0.879(0.006) 0.861(0.005) 0.798(0.007)
MolVAE [30] 0.892(0.009) 0.877(0.009) 0.788(0.033)
DeepDDI [2] 0.886(0.007) 0.871(0.007) 0.817(0.007)
CASTER [64] 0.910(0.005) 0.887(0.008) 0.843(0.005)
GEM [20] 0.960(0.003) 0.956(0.002) 0.903(0.003)
MPG [13] 0.966(0.004) 0.960(0.004) 0.905(0.008)
BatmanNet 0.972(0.001) 0.966(0.001) 0.916(0.002)

以下是原文 [Table 4] 的结果:

Model Precision Recall F1
DDI_PULearn [66] 0.904 0.824 0.862
GEM [20] 0.928 0.929 0.928
MPG [13] 0.936 0.936 0.936
BatmanNet 0.939 0.939 0.939

结果分析:

  • BIOSNAP 数据集: 在 BIOSNAP 数据集上,BatmanNetAUC-ROC (0.972)、PR-AUC (0.966) 和 F1 (0.916) 三个指标上均超越了之前的 SOTA 方法 MPG
  • TWOSIDES 数据集: 在 TWOSIDES 数据集上,BatmanNetPrecision (0.939)、Recall (0.939) 和 F1 (0.939) 上同样表现最佳,略微优于 MPG
  • 这些结果证明了 BatmanNet 在 DDI 预测任务中的优越性能。

6.1.4. 药物-靶点相互作用 (DTI) 预测

以下是原文 [Table 3] 的结果:

Datasets Model Precision Precision AUC
Human Tsubaki et al. [65] 0.923 0.918 0.970
GEM [20] 0.930 0.930 0.972
MPG [13] 0.952 0.940 0.985
BatmanNet 0.983 0.982 0.998
(Relative improvement) (3.26%) (4.47%) (1.32%)
C.elegans Tsubaki et al. [65] 0.938 0.929 0.978
GEM [20] 0.955 0.954 0.988
MPG [13] 0.954 0.959 0.986
BatmanNet 0.988 0.987 0.999
(Relative improvement) (3.46%) (2.92%) (1.11%)

结果分析:

  • Human 和 C.elegans 数据集: 当使用 BatmanNet 作为化合物编码器时,在 Human 和 C.elegans 数据集上,性能均显著提升。
  • 对比 MPG 与之前最佳模型 MPG 相比,BatmanNetHuman 数据集上的 Precision 分别提高了 3.26% 和 4.47%,AUC 提高了 1.32%。在 C.elegans 数据集上,Precision 分别提高了 3.46% 和 2.92%,AUC 提高了 1.11%。
  • 这些结果表明 BatmanNet 在分子表示学习方面具有强大的迁移学习 (transfer learning) 能力,对于 DTI 预测研究具有广阔前景。

6.2. 消融实验/参数分析

为了深入理解 BatmanNet 框架中各因素对性能的影响,作者在 6 个分类基准数据集上进行了消融研究。

下图(原文 Figure 5)展示了 BatmanNet 模型在不同药物发现任务上的 AUC 结果。

该图像是图表,展示了 BatmanNet 模型在不同药物发现任务上的 AUC 结果。图 (a) 对比了节点分支、边分支和 BatmanNet (双分支) 的性能,显示出 BatmanNet 在多个基准数据集上实现了最佳表现。图 (b) 展示了预训练和非预训练模型在各数据集上的 AUC 变化趋势,图 (c) 则分析了不同掩膜比例对 AVG-AUC 的影响,最高达到 0.848。 该图像是图表,展示了 BatmanNet 模型在不同药物发现任务上的 AUC 结果。图 (a) 对比了节点分支、边分支和 BatmanNet (双分支) 的性能,显示出 BatmanNet 在多个基准数据集上实现了最佳表现。图 (b) 展示了预训练和非预训练模型在各数据集上的 AUC 变化趋势,图 (c) 则分析了不同掩膜比例对 AVG-AUC 的影响,最高达到 0.848。

6.2.1. 双分支信息提取网络的有效性

  • 实验设计: 为了评估双分支设计的影响,作者将 BatmanNet 中的双分支信息提取网络替换为单分支网络(仅节点分支或仅边分支)。所有单分支网络均在与双分支网络相同的条件下进行预训练,且参数量接近 (2.6M)。
  • 结果分析 (Figure 5a): 结果显示,BatmanNet 的双分支设计将平均 AUC 提高了 2.9%(相较于仅使用节点分支)和 3.0%(相较于仅使用边分支)。这表明双分支设计的有效性,它能互补地捕获节点和边的信息,从而提升模型性能。

6.2.2. 自监督预训练策略的影响

  • 实验设计: 比较预训练的 BatmanNet 和未预训练的 BatmanNet 在分子性质预测任务上的分类性能,使用相同的超参数。
  • 结果分析 (Figure 5b): 预训练模型持续优于未预训练模型,平均 AUC 提升了 4.0%。这表明自监督预训练策略有效地捕获了分子丰富的结构和语义信息,从而提升了下游任务的性能。

6.2.3. 不同掩蔽比例的影响

  • 实验设计: 预训练 BatmanNet,掩蔽比例从 10% 到 90% 不等,然后测量所有下游任务的平均 AUC
  • 结果分析 (Figure 5c): 当掩蔽比例设置为 60% 时,模型性能最佳。
    • 较低掩蔽比例: 当掩蔽比例相对较低时,任务不够具有挑战性,模型学习到的信息容量有限。

    • 适中掩蔽比例 (60%): 当使用相对较高的掩蔽比例时,更多节点和边被掩蔽,导致预训练任务更具挑战性。这为剩余节点和边提供了更大的容量,使其在嵌入中捕获分子信息。

    • 过高掩蔽比例: 当掩蔽比例高于 60% 时,剩余的节点和边信息不足以恢复完整的图,导致学习到的嵌入质量开始下降。

      以下是原文 [Table S6] 的结果:

      Ratio BBBP SIDER ClinTox BACE Tox21 ToxCast Avg
      0.1 0.923(0.032) 0.662(0.015) 0.905(0.028) 0.913(0.007) 0.843(0.014) 0.739(0.011) 0.831
      0.2 0.929(0.027) 0.667(0.003) 0.912(0.012) 0.915(0.006) 0.845(0.009) 0.745(0.009) 0.836
      0.3 0.933(0.018) 0.668(0.003) 0.918(0.025) 0.919(0.013) 0.848(0.017) 0.749(0.007) 0.839
      0.4 0.940(0.011) 0.671(0.006) 0.920(0.028) 0.920(0.014) 0.850(0.014) 0.750(0.009) 0.842
      0.5 0.943(0.019) 0.675(0.004) 0.925(0.025) 0.925(0.014) 0.851(0.013) 0.753(0.008) 0.845
      0.6 0.946(0.007) 0.676(0.004) 0.926(0.015) 0.928(0.015) 0.855(0.013) 0.756(0.009) 0.848
      0.7 0.942(0.008) 0.674(0.004) 0.926(0.011) 0.924(0.016) 0.848(0.012) 0.751(0.007) 0.844
      0.8 0.938(0.012) 0.673(0.004) 0.921(0.288) 0.918(0.016) 0.844(0.014) 0.750(0.008) 0.841
      0.9 0.932(0.020) 0.668(0.005) 0.910(0.022) 0.915(0.015) 0.840(0.015) 0.746(0.011) 0.835

下图(原文 Figure S1)展示了掩蔽比例对各个基准数据集的影响。

Fig. S1. The influence of the masking ratio on each benchmark dataset. 该图像是一个图表,展示了不同遮罩比例对多个基准数据集(如 BACE、BBBP 和 ClinTox 等)AUC 值的影响。每个子图中,AUC 值与遮罩比例呈现出不同的趋势,部分数据集在高遮罩比例下表现更好。

这些结果表明,将掩蔽比例设置为 60% 能够实现最佳的预测性能,并且 BatmanNet 在各种数据集上的性能表现出一致性。

6.3. 额外实验

在遵循 GEM [20] 相同实验设置的分子性质预测额外实验中,BatmanNet 在 11 个数据集中的 7 个上取得了 SOTA 性能,所有数据集上的总体相对改进为 1.1%(分类任务上为 1.8%,回归任务上为 0.4%)。

以下是原文 [Table S5] 的结果:

Dataset BACE BBBP Clin Tox SIDER Tox21 ToxCast HIV MUV Avg
#molecules 1513 2039 1478 1427 7831 8575 41127 93087 -
#tasks 1 1 2 27 12 617 1 17 -
D-MPNN 0.809(0.006) 0.710(0.003) 0.906(0.006) 0.570(0.007) 0.759(0.007) 0.655(0.003) 0.771(0.005) 0.786(0.014) 0.746
AttentiveFP 0.784(0.022) 0.643(0.018) 0.847(0.003) 0.606(0.032) 0.761(0.005) 0.637(0.002) 0.757(0.014) 0.766(0.015) 0.735
N-GramRF 0.779(0.015) 0.697(0.006) 0.775(0.040) 0.668(0.007) 0.743(0.004) 0.772(0.001) 0.769(0.007) -
N-GramXGB 0.791(0.013) 0.691(0.008) 0.875(0.027) 0.655(0.007) 0.758(0.009) - 0.787(0.004) 0.748(0.002) -
PretrainGNN 0.845(0.007) 0.687(0.013) 0.726(0.015) 0.627(0.008) 0.781(0.006) 0.657(0.006) 0.799(0.007) 0.813(0.021) 0.742
GROVERbase 0.826(0.007) 0.700(0.001) 0.812(0.030) 0.648(0.006) 0.743(0.001) 0.654(0.004) 0.625(0.009) 0.673(0.018) 0.710
GROVERlarge 0.810(0.014) 0.695(0.001) 0.762(0.037) 0.654(0.001) 0.735(0.001) 0.653(0.005) 0.682(0.011) 0.673(0.018) 0.708
GraphMAE 0.831(0.009) 0.720(0.006) 0.823(0.012) 0.603(0.011) 0.755(0.006) 0.641(0.003) 0.772(0.010) 0.763(0.024) 0.739
GEM 0.856(0.011) 0.724(0.004) 0.901(0.013) 0.672(0.004) 0.781(0.001) 0.692(0.004) 0.806(0.009) 0.817(0.005) 0.781
BatmanNet 0.861(0.028) 0.838(0.005) 0.897(0.012) 0.659(0.003) 0.792(0.003) 0.718(0.007) 0.812(0.009) 0.784(0.014) 0.795
Methods
Model
#molecules ESOL
#tasks 1128 1
D-MPNN AttentiveFP 1.050(0.008) 1.272
N-Gram_RF 0.877(0.029) 1.224
1.074(0.107) 1.083(0.082) 1.525
N-Gram_XGB 1.100(0.006) 2.739
PretrainGNN 0.983(0.090) 1.534
GROVERbase GROVERlarge 0.895(0.017) 1.325
GEM 0.798(0.029) 1.330
1.112
BatmanNet 0.792(0.013) 1.108

7. 总结与思考

7.1. 结论总结

本文针对人工智能驱动的药物发现 (AIDD) 领域中分子表示学习的两个主要挑战——现有自监督学习方法预训练任务复杂性和高计算复杂度与模型规模——提出了一个新颖的自监督学习框架。 核心贡献包括:

  • BatmanNet 模型: 提出了一种双分支掩蔽图 Transformer 自编码器 (BatmanNet),它采用不对称的编码器-解码器设计,显著提高了模型的学习效率和内存利用率。

  • 创新自监督策略: 设计了一种简单而有效的双分支图掩蔽自监督学习策略,通过高比例(60%)同时掩蔽节点和边并进行重建,无需额外的领域知识,即可有效地学习分子的局部和全局信息。

  • 卓越性能: BatmanNet 在分子性质预测、药物-药物相互作用 (DDI) 预测和药物-靶点相互作用 (DTI) 预测等多个药物发现任务的 13 个基准数据集上,持续超越了现有最先进 (SOTA) 方法。

    总而言之,该工作强调了设计有效预训练任务对于分子表示学习的重要性,并展示了一种基于自编码任务的简单、可扩展且领域无关的方法的有效性。

7.2. 局限性与未来工作

作者指出了 BatmanNet 存在的几个潜在改进方向:

  • 三维结构信息的整合: 当前方法主要关注分子的平面拓扑结构,忽略了三维结构信息。未来工作可以通过将三维结构细节整合到节点和边特征中,使模型能够更全面地理解分子特征,尤其是在药物-靶点相互作用等涉及三维结构的任务中提升性能。
  • 预训练数据集规模的扩展: 受限于计算资源,当前模型和数据集规模较小。进一步将该方法扩展到更大的预训练数据集,可以评估在更大规模下能实现的性能提升。
  • 缓解数据驱动方法的固有偏差: 当前完全数据驱动的方法可能存在潜在的数据偏差。未来的努力将集中于探索有效策略,将领域知识整合到当前数据驱动的流程中,以缓解这种偏差并进一步提高模型性能。

7.3. 个人启发与批判

7.3.1. 个人启发

这篇论文提供了一些重要的启发:

  • 自监督学习的潜力: 论文再次强调了自监督学习在数据稀缺领域(如药物发现)的巨大潜力。通过设计巧妙的前置任务,可以从大量无标签数据中学习到有用的、可迁移的表示。
  • 简洁与高效的设计: BatmanNet 的成功表明,在深度学习模型设计中,简洁性和效率并不意味着牺牲性能。不对称的编码器-解码器架构、简单的双分支掩蔽策略,相比于复杂的、多任务的预训练方法,反而能取得更好的效果,这对于资源有限的研究者来说尤其有价值。
  • “挑战性”任务的设计: 论文通过高比例的掩蔽来增加预训练任务的挑战性,迫使模型学习更丰富的上下文信息,这是一个值得借鉴的设计原则。它启发我们,预训练任务的设计不应仅仅是“有意义”,更要能“挑战”模型去学习更深层次的特征。
  • 领域无关性: BatmanNet 的图掩蔽和重建机制相对通用,不 сильно依赖于特定的化学领域知识(如基序定义),这意味着其核心思想可能可以迁移到其他图结构数据表示学习领域,例如社交网络分析、知识图谱嵌入等。

7.3.2. 潜在问题、未经验证的假设或可改进的地方

  • 三维信息缺失的权衡: 尽管作者提到了这是一个未来方向,但当前模型仅考虑二维拓扑结构,可能在某些高度依赖分子构象(如药物-靶点结合位点识别)的任务中存在固有的局限性。三维结构信息的整合并非易事,需要考虑计算成本和数据可用性。
  • 掩蔽比例的敏感性: 消融实验显示掩蔽比例在 60% 达到最佳,但具体如何选择这个比例,以及它是否对所有类型的分子和所有下游任务都普遍适用,可能需要更深入的理论分析或更广泛的实验验证。不同性质的分子或任务可能对掩蔽比例有不同的最优值。
  • 分歧损失的有效性: 微调阶段引入了分歧损失 (disagreement loss),以使节点分支和边分支的预测保持一致。虽然这在实践中可能有效,但其理论依据和对模型泛化能力的具体影响值得更深入的探讨。例如,如果两个分支在某些情况下确实需要学习不同的特征,强制一致性是否会带来负面影响?
  • 模型可解释性: 作为一个基于 Transformer 和 GNN 的复杂模型,BatmanNet 的可解释性可能仍是一个挑战。如何理解模型学习到的分子表示具体对应了哪些化学特性或结构模式,对于药物发现的科学家而言至关重要。未来的工作可以尝试结合可解释性方法来提高模型的透明度。
  • 计算效率与扩展性: 尽管论文强调了其效率优势,但对于更大规模的分子数据集(远超 ZINC250K)和更复杂的分子,其轻量级解码器和不对称架构是否仍能保持SOTA性能和计算优势,需要更进一步的验证。特别是在与 GROVERlarge 等具有更多参数的模型进行比较时,BatmanNet 在某些任务上仍有差距,这表明模型容量和数据规模之间的平衡仍需探索。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。