NeurIPS 2021 | 分布偏移下的用于药物发现的可靠图神经网络

编译 | 程思雨 审稿 | 刘名权

今天给大家介绍一篇由Google Research机构的Kehang Han、Balaji Lakshminarayanan、Jeremiah Liu共同发表的文章:《Reliable Graph Neural Networks for Drug Discovery Under Distributional Shift》。在分布偏移下对过度自信的错误预测的关注，要求我们对用于药物发现中的关键任务的图神经网络进行广泛的可靠性研究。该文章首先介绍了CardioTox,，一个真实世界的药物心脏毒性基准，以促进这方面的努力。作者的一个探索性研究表明，过于自信的错误预测往往与训练数据相距甚远。这进而引导作者开发了距离感知的GNNs: GNN-SNGP。通过对 CardioTox 和三个既定基准的评估，他们证明了 GNN-SNGP 在增加距离感知、减少过度自信的错误预测和在不牺牲精确性能的情况下做出更好的校准预测方面的有效性。作者的消融研究进一步揭示了由GNN-SNGP 学习的表征改进了其基本结构上的距离保存，并且是改进的主因之一。

1介绍

近年来，图神经网络(GNNs)在物理、生物医学和计算化学等科学应用中表现出了显著的性能。一个重要的例子是早期的药物发现，GNNs在关键任务中表现出了希望，如命中发现和责任筛选(即分别预测候选药物分子的结合亲和力和毒性)。

然而，阻碍在实际应用中采用 GNN 的一个关键的可靠性问题是过于自信的错误预测。例如，在责任筛选中，一个过度自信的假阴性(以下简称 OFN)预测会导致 GNN 模型将一个有毒分子标记为安全，从而导致严重后果，将其泄漏到药物开发的下一个阶段。药物发现任务的第二个关键特征进一步加剧了这种担忧: 数据分布的偏移; 药物发现任务通常通过转移到以前在训练数据中没有表现出来的特征空间中的区域来明确地评估新的分子。因此，测试分子与训练数据的特点截然不同，并且可以携带新的毒性信号，这些信号在模型中是前所未见的。在这种情况下，仅有高分布精度的模型是不够的; 在分布偏移下，量化模型的可靠性和设计技术，以提高对过度自信的稳健性变得特别重要的。

目前，缺乏一个药物发现的基准，它针对的是分布偏移下模型可靠性的现实担忧。因此，作者引进了 CardioTox，一个基于现实世界中的药物发现问题的数据基准，它是由来自 ChEMBL、NCATS和 FDA 验证数据库的9千多类药物分子编译而成的。为了评估模型的可靠性，作者生成了额外的分子注释，并提出了新颖的度量方法来围绕负责任的GNN应用，根据真实世界的标准衡量模型。这是他们的第一个贡献。

通过使用CardioTox，他们的第二个贡献是一个关于在分布偏移下过于自信的错误预测背后的根本原因的探索性研究，以及使用原则性建模方法来减轻它。特别是，他们观察到，许多过度自信的错误预测的分子在结构上与训练数据不同(也就是说，它们“远离”于基于分子指纹图距离的训练数据)。这种失效模式表明，提高 GNN 模型的距离意识将是一个有效的解决方案: 一个测试分子，其远离决策边界(例如，一个有毒但新颖的分子，由于其新颖性，可能拥有很少毒性信号定义决策边界)，仍然应该得到低信心，如果它远离训练数据。为此，最近提出的光谱归一化神经高斯过程(SNGP)展示了实现这一目标的具体方法。具体地说，它给特征提取器施加了一个保持距离的正则化(即光谱归一化) ，并用一个距离感知分类器(即随机特征高斯过程)代替稠密输出层。SNGP 在视觉和语言问题方面显示出了很好的稳健性改进。这是首次将距离感知设计原则引入到 GNN 模型中，以提高分子图的可靠性。

总而言之，他们的贡献如下(附录 E 总结了相关工作) :

• Data: 他们将真实世界的药物发现数据集(IID 和分布移位数据集) CardioTox 引入到健壮性社区，以促进图模型的可靠性研究。 他们进一步设计了基于距离的数据分割，以定量测量模型的距离感知。
• Model: 他们开发了一个端到端可训练的 GNN-SNGP 体系结构及其消融版本 GNN-GP。 从实际经验来看，这种方法不仅在准确性(例如 AUROC)方面优于其基本架构，而且在健壮性(例如 ECE)方面也优于其基本架构，特别是在减少过度自信的错误预测方面。 通过以CardioTox的距离为基础的数据分裂测量，GNN-SNGP 在数据移位时的距离感知能力高于基线，这解释了其在过度自信的错误预测减少的能力。 他们的实现和CardioTox数据集将通过 Github 开源。
• Ablation Study: 他们进行了大量的消融研究，证实了鲁棒性的改进来自于距离感知分类器和保持距离的潜在表象。 他们还通过测试其他图形建模领域: molHIV，BBBP 和 BACE 来研究这种方法的可推广性，并得到了一致的结果。

2方法

GNN基线：在这项研究中，作者使用一个普通的消息传递神经网络(MPNN)作为 GNN 的基线。具体地说，消息函数是由全连接层建模而成。

GNN-GP: 提高分类器的距离意识。根据 Liu 等人2020年在视觉和语言模型中提出的高斯过程层，作者将其引入到图域中，并开发了 GNN-GP 模型来增加距离感知。图 S1(附录 A)显示了高层架构的改变。

图1. 通过保持距离的特征提取器(跳过连接和光谱归一化)和距离感知分类器(神经高斯过程层)改进 GNN 结构。

如图1所示，通过随机傅里叶特征生成逼近 GP 核函数，使得高斯过程可以用 GNN 进行端到端的可训练。在推理过程中，每个样本都可以得到 logit 预测和 logit 方差，这两个方差都可以通过均值场近似计算预测概率。

GNN-SNGP: 结合距离保持的特征提取。由于特征提取中的特征塌陷，神经表示不能忠实地保持输入流形中的距离。Liu 等人在2020提出在特征提取中保持输入距离，将光谱归一化(SN)应用于残差网络。因此，结合 SN 和 GP 可以提高模型的总体距离感知能力，有助于面向对象的检测和其他鲁棒性度量。由于普通 MPNN 模型(GNN 基线)没有残差连接，作者通过以下两个更改创建 GNN-SNGP (图1)。首先，作者使用 SN 对消息函数进行全连接层建模；其次，作者通过节点更新函数将残差连接添加到消息传递层。

3CardioTox:分布偏移下的药物心脏毒性

评估指标

作者评估了流行的不确定性度量，如预期校准误差(ECE)，Brier 评分(Brier) ，负对数似然(NLL)。为了评估模型在捕捉图距离和缓解过度自信方面的能力，作者还引入了两个新的指标: 过度自信假否定百分比(OFNs%)：

请注意，OFNs% 可能会被所有预测的不确定性所“欺骗”，因此一个理想的不确定性模型应该在校准指标和 OFNs% 中表现良好。

精确性和鲁棒性能

表1显示 GNN-GP 在 AUROC 和 CardioTox 任务的健壮性指标方面都优于 GNN 基线。对于不连续(Test-IID)和移位测试集(Test-OOD1和Test-OOD2)来说都是如此。GNN-SNGP 在鲁棒性方面显示了额外的增益。如果资源允许，GNN-SNGP的深度集成将进一步提高性能，消除了Test-OOD2中的所有的OFNs。值得注意的是，作者的GNN-SNGP集成在AUROC表现上优于以往最先进的神经模型。

表1: 药物心脏毒性基准的准确性(AUROC) ，健壮性(ECE，Brier，NLL)和过度自信(OFNs)性能。结果平均超过10粒种子。

为什么距离意识有助于减少过分自信的错误预测？

作者观察到很大一部分的OFNs距离训练数据集很远(即60%的OFNs的 Tanimoto distance> 0:65，见图 S2b)。当一个 GNN 模型缺乏距离意识时，这种情况就会发生: 一个有毒但新奇的分子可能远离模型的决策边界的错误一侧，这是由于缺乏已知的在训练数据集中存在的毒性特征。由于不知道到训练数据集的距离，GNN 基线模型倾向于根据到决策边界的距离进行预测，并对这种情况给予很高的置信度。

为此，GNN-GP利用了GP的距离感知能力，能够自然地将距离纳入预测的不确定性中。如表1所示，SN版本的GNN-SNGP在两个数据源位移测试集中获得了最高的距离感知(DA-AUC)，与OFNs缩减具有良好的相关性。图S2b显示了OFNs中距离样本百分比的下降趋势(从GNN基线到GNN- GP再到GNN- SNGP)。总体而言，使用GP可以提高超过80%的基线OFNs的不确定度估计，同时也可以改善校准(图S4)。

额外消融以评估SN和GP的相对贡献

为了理解这两个建模分量(即距离保持特征提取器和距离感知分类器)的相对性能贡献，作者在附录H中进行了广泛的消融研究。他们利用通过GNN 基线、 GNN-GP和 GNN-SNGP (增加距离保持)学习到的潜在表示，将它们提供给距离意识增强的分类器: 全连接层、 gps 层、精确高斯过程分类器(GPC )。表 S2表明，神经表征的距离保持和分类器的距离意识之间存在协同作用。使用最小距离感知分类器(即全连接层) ，AUROC 在增加距离保持的神经表示时下降。使用最具距离感知能力的分类器(即精确 GPC) ，提高神经表征的距离保持能力有利于降低准确性、鲁棒性以及过度自信。作者发现，当两者都在体系结构中(GNN-SNG中嵌入 GPC)时GNN 模型可以达到最好的性能。

现有基准的额外结果

与 CardioTox 中的结果一致，GNN-GP 在三个既定的图形分类基准上也优于 GNN 基准: molHIV，BBBP 和 BACE (见附录 G)。可以对表 S1中的结果进行一些观察。首先，GNN-GP始终能够实现更高的AUROC，改善校准，并减少对所有基准的过于自信的错误预测。其次，与GNN-GP相比，光谱归一化版本GNN-SNGP在精度性能下降有限的情况下，可以进一步提高模型的鲁棒性。

4结论

在这项研究中，作者介绍了GNN-GP和GNN-SNGP，以及来自药物发现设置的新的基准CardioTox。通过对四个数据集的评估，他们证明了它们在减少过度自信的错误预测和在不牺牲精确性的情况下制作更好校准的 GNN 模型的有效性。健壮性的提高似乎来自于距离意识的提高。作者进一步发现，由 SN 和 GP 引入的嵌入空间改善了基本结构的距离保持性，是提高嵌入精度、鲁棒性和过度自信表现的主要因素之一。

未来展望：对其他基础 GNN 架构如 GAT和 PNA进行更广泛的实证研究将是有趣的。另一个有趣的方向是，通过附录H中详细描述的消融研究，找到一种定量的方法来测量学习表征中的距离保持，并通过精确性和稳健性的表现来理解保持和特定任务压缩之间的权衡。

参考资料

Reliable Graph Neural Networks for Drug Discovery.

( https://arxiv.org/abs/2111.12951)