Nat Commun｜基于图神经网络的迁移学习在多保真度数据下的分子性质预测

2024年2月26日，英国剑桥大学Pietro Lio教授团队，联合瑞典阿斯利康，在Nature Communications上发表文章Transfer learning with graph neural networks for improved molecular property prediction in the multi-fidelity setting。

作者研究了图神经网络在化学分子的多保真度（multi-fidelity）数据迁移学习方面的潜力，通过利用低保真度数据进行预训练，在小样本的高保真度数据上改进分子特性预测。为了增强从预训练任务中有效学习分子表示的性能，作者提出了一种基于神经网络的分子图自适应读出方法。结果表明，所提出的方法优于现有的迁移学习策略。

背景

化学分子数据集往往具有不同的保真度，规模较大的数据集，其误差也较大，保真度较低。而要获取误差较小、具有高保真度的数据，其时间和经济成本往往较为高昂，这使得高保真度的数据规模往往较小。如何融合不同保真度的数据集，提升分子性质预测的性能，是一个值得研究的问题。

方法

本文提出的框架如图1所示。变分图自编码器（VGAE）用于根据化学键的连通性传播和学习分子中各个原子的表示z1,z2,...，然后将其在邻域上进行聚合（Agg.），得到更新后的原子表示z1*,z2*,...。用于将原子表示聚合成单个分子表示z* 的读出函数是神经网络(自适应聚合器)，符号||表示向量拼接，μ(x)和σ(x)表示VGAE学习到的均值和标准差，“Dense NN”是一个将分子表示映射到分子性质数值的多层感知器。图中所示的四个实验流程列在右上方。

实验1（E1）：用监督信息训练一个低保真度（low-fidelity，LF）模型，得到分子表示z*。实验2（E2）：加入低保真度数据标签。实验3（E3）：加入低保真度模型生成的表示。实验4（E4）：训练具有相同架构的模型来预测高保真值，具有自适应读出的图神经网络在VGAE层固定的情况下，只对读出层进行微调。

图1 基于图神经网络的多保真度数据迁移学习结构图

图神经网络中的标准读出函数(即求和、求平均和求最大值)没有任何参数，因此不适合于迁移学习。本文提出了一种神经网络架构的读出函数，将学习到的节点表示聚合到图嵌入中，也就是在微调阶段，固定图神经网络架构中负责学习有效节点表示的部分，并在小样本下游任务中微调读出层。本文采用集合Transformer读出函数，它保留了标准读出函数的排列不变性特性。由于图可以看作是节点的集合，本文利用这种架构作为图神经网络中的读出函数。

在这个Transformer中，编码器是由多个经典的多头注意力块组成的链，没有位置编码。解码器组件由投影的多头注意力块组成，通过多个自注意力模块链和线性投影块做进一步处理。与典型的基于集合的神经结构只能单独地处理单个元素相比，本文提出的自适应读出可以建模由邻域聚合方案生成的所有节点表示之间的相互作用。该架构的一个特点是图表示的维度可以从节点输出维度和聚合方案中分离出来。

变分图自编码器（VGAE）由一个概率编码器和一个概率解码器组成，与在向量值输入上操作的标准变分自编码器架构相比，VGAE通过图卷积层的叠加得到编码器分量，学习表示隐空间编码高斯分布的参数矩阵μ和σ。此外，该模型通常假设存在自环，即邻接矩阵的对角线由1组成。解码器通过sigmoid函数传递隐变量之间的内积，重构邻接矩阵中的元素。通过优化证据下界损失函数来训练变分图自编码器，该函数可以看作是基于表示学习的重建误差，加上变分分布q(⋅)与先验p(⋅)之间的Kullback-Leibler散度（KL散度）正则化项，其中先验假设为高斯分布，由于图的邻接矩阵通常是稀疏的（即Aij=0的元素远多于Aij=1的元素），所以在训练时通常对Aij=0的样本进行采样，而不是取所有的负样本。

作者通过在变分图自编码器的表示空间上添加前馈分量来扩展这种神经网络，将标签信息引入VGAE的训练优化过程，实现了有监督的变分图自编码器。实际上，从传统的变分自编码器（VAE）过渡到图上的VGAE，然后再过渡到建模分子结构的VGAE，这不是一个平凡（trivial）的过程，原因至少有两个。首先，原始的VGAE只重建图的连通性信息，而不重建节点(原子)特征。这与传统的VAE形成对比，在传统的VAE中，潜在表示直接针对实际输入数据进行优化。因此，重建函数之间的平衡(分别用于连通性和节点特征)是图深度学习中的一个悬而未决的问题。其次，对于分子水平的预测和潜在空间表示等任务，变分图自编码器的读出功能至关重要。标准读出函数会导致类似于完全无监督训练的无信息表示，在迁移学习任务中表现不佳。因此，这里提出的监督变分图自编码器在图表示学习方面也是一个进步，可以用于建模数百万规模的具有挑战性的分子任务。

结果

作者将基于图神经网络的多保真度数据迁移学习与一些具有代表性的方法进行了比较。作者首先针对不同药物的低保真度单剂量（single dose，SD）和高保真度剂量响应（dose response，DR）的pIC50标签进行比较。将未进行迁移学习的非监督VGAE模型作为基线，对比其他模型预测结果的MAE下降的幅度，降幅越大说明模型性能越好。结果如图2所示。其中Label表示仅使用高保真度标签和神经网络读出，也就是未进行低保真度迁移学习的有监督VGAE。LF表示使用低保真度数据迁移学习，并分别采用加和（sum）读出和神经网络（neural）读出。其中Emb, Pred lbl, Hyb lbl分别表示使用表示数据重构误差（也就是无监督的VGAE），使用低保真度模型预测的标签，混合（hybrid）低保真度模型预测和原有的低保真度标签。Tune VGAE表示在加和读出中，微调VGAE，Tune readout表示在神经网络读出中，固定VGAE，只微调读出层。

结果表明，结合了多保真度数据迁移学习、神经网络读出的有监督的VGAE算法的模型（Tune readout）在大多数情况下取得了最优结果。

图2 与其他方法对比

作者设计了消融实验。对于测试集中的标签值超出训练集中的标签值分布（out of distribution，OOD）的情况进行了验证，在OOD和非OOD情况下的学习过程分别对应迁移学习中的归纳式学习（inductive learning）和直推式学习（transductive learning）。OOD是药物发现中的常见现象，是限制实际应用中分子筛选性能的主要瓶颈。

如图3所示，训练集的大小对于归纳式学习的性能有显著影响，在未进行增强时（将未进行迁移学习的非监督VGAE模型作为基线），归纳式学习预测结果的相关系数平方值很低。随着训练集扩大，不同模型的性能均有所提升。无论对于归纳式学习还是直推式学习，神经网络读出的性能均优于加和读出。在使用神经网络读出进行归纳式学习时，即使采用非监督的自编码器，其结果也与有监督的自编码器相当。在使用神经网络读出进行直推式学习时，即使采用非监督的自编码器，其结果也超越了有监督的自编码器。结果表明自适应的神经网络读出策略可以较好地解决OOD问题。

图3 消融实验

作者还进行了案例分析。如图4所示，选取若干种化合物，将其使用均值读出函数、加和读出函数和神经网络读出函数（集合Transformer）的测试MAE进行对比（越低越好，三种方法的MAE最高值归一化为1）。结果表明，加和读出略优于均值读出，神经网络读出显著优于以上两种读出。

图4 案例分析

总结

在本文中，作者研究了在分子设计和药物发现中通常遇到的在小样本数据集上学习有效的分子性质预测模型的问题。本文的重点是展示迁移学习与图神经网络在多保真度数据中的应用。

更具体地说，作者主要关注的是在化学分子数据集中，大规模的低保真度数据和小样本的高保真度数据之间的知识迁移。本文的主要算法贡献在于识别和解决经典图神经网络的缺点，克服将图神经网络用于多保真度数据迁移学习的瓶颈。作者提出了改进方案主要包括以下两点：（1）通过在低保真度数据上训练的模型生成的表示进行迁移，也就是在低保真度上预训练，在高保真度上进行微调。（2）采用具有自适应读出的图神经网络，增强分子图表示学习的性能。实验结果的有效性和通用性展现了图神经网络在多保真度数据迁移学习中的应用前景。

基于分子表示的迁移学习因其高效和广泛的适用性而备受关注，未来，可探索将更多的分子表示学习模型应用于上述迁移学习框架，例如可以提供不确定性估计的高斯过程模型等。而基于变分图自编码器架构在分子生成任务中的广泛应用，本文提出的图神经网络迁移学习框架不仅可用于分子性质预测，也有潜力用于分子生成。另一个方向是，在同一个预训练任务上训练的模型可在分子性质预测的多个不同的微调任务上得到应用。

参考文献

[1] Buterez et al. Transfer learning with graph neural networks for improved molecular property prediction in the multi-fidelity setting. Nat Commun. 2024