[论文解读] Reliable Graph Neural Networks for Drug Discovery Under Distributional Shift

Kehang Han, Balaji Lakshminarayanan+1 more|arXiv (Cornell University)|Nov 25, 2021

Computational Drug Discovery Methods参考文献 29被引用 15

一句话总结

本文提出 GNN-SNGP，一种通过整合谱归一化与基于距离感知的高斯过程分类器，在分布偏移条件下实现可靠的图神经网络，用于药物发现。该方法在新提出的 CardioTox 基准和三个现有数据集上验证，展现出更优的鲁棒性，显著减少了对分布外分子的过度自信误判，同时保持高准确率。

ABSTRACT

The concern of overconfident mis-predictions under distributional shift demands extensive reliability research on Graph Neural Networks used in critical tasks in drug discovery. Here we first introduce CardioTox, a real-world benchmark on drug cardio-toxicity to facilitate such efforts. Our exploratory study shows overconfident mis-predictions are often distant from training data. That leads us to develop distance-aware GNNs: GNN-SNGP. Through evaluation on CardioTox and three established benchmarks, we demonstrate GNN-SNGP's effectiveness in increasing distance-awareness, reducing overconfident mis-predictions and making better calibrated predictions without sacrificing accuracy performance. Our ablation study further reveals the representation learned by GNN-SNGP improves distance-preservation over its base architecture and is one major factor for improvements.

研究动机与目标

• 解决在药物发现中使用的图神经网络存在的过度自信误判这一关键可靠性问题，特别是在分布偏移条件下。
• 识别出过度自信的误判通常发生在与训练数据结构上相距较远的分子上，表明缺乏对距离的感知能力。
• 开发一种新型图神经网络架构 GNN-SNGP，通过增强表征中的距离保持能力并引入不确定性感知预测，提升模型可靠性。
• 提出 CardioTox，一个包含分布内和分布偏移测试集的真实世界基准，用于在现实药物发现条件下评估模型的可靠性。
• 证明改进的距离感知能力可带来更好的校准效果并减少过度自信的假阴性预测，这对安全的药物开发至关重要。

提出的方法

• 提出 GNN-SNGP，一种端到端可训练的图神经网络，结合谱归一化（SNGP）与基于随机特征的高斯过程（GP）层，用于不确定性估计。
• 在图神经网络的消息传递层中应用谱归一化，以在潜在空间中保持几何距离，增强距离感知能力。
• 用 GP 层替换最后的全连接层，基于与训练数据的距离建模预测不确定性，从而降低对分布外分子的过度自信。
• 使用基于指纹的图距离（Rogers 和 Hahn，2010）量化分子与训练数据的距离，并据此定义基于距离的数据划分以用于评估。
• 使用标准图神经网络损失配合不确定性正则化进行训练，实现在不牺牲准确率的前提下实现可靠的不确定性量化。
• 开展消融实验，以分离 GP 层与保持距离的表征在整体鲁棒性中的独立贡献。

实验结果

研究问题

• RQ1 在药物毒性预测的图神经网络中，对与训练数据结构上相距较远的分子，其过度自信误判的发生程度如何？
• RQ2 在药物发现中，通过增强图神经网络表征的距离感知能力，能否减少在分布偏移下的过度自信误判？
• RQ3 在分布外数据上，GNN-SNGP 与基线图神经网络及其他不确定性估计方法相比，在校准性和鲁棒性方面表现如何？
• RQ4 保持距离的表征与基于 GP 的分类器相结合，是否能带来优于单一组件的不确定性量化效果？
• RQ5 该方法在心脏毒性之外的其他基于图的药物发现任务中是否具有泛化能力？

主要发现

• 在测试-OOD1 分割上，GNN-SNGP 达到最高的 AUROC（0.935）和最低的 ECE（0.002），在校准性和鲁棒性方面显著优于 GNN 基线和其他变体。
• 在最具挑战性的测试-OOD2 分割上，GNN-SNGP 将过度自信假阴性（OFNs%）降低至 0.00%，而 FPs + Dense 基线为 5.71%。
• 消融研究证实，保持距离的 GNN 编码器和 GP 层均对可靠性提升有贡献，其中 GP 层在减少过度自信方面尤为有效。
• GNN-SNGP 在 OOD 分割上 ECE（降至 0.021）和 Brier 得分（0.208）持续改善，表明其不确定性估计更具校准性。
• 该模型在保持高准确率（AUROC > 0.89）的同时减少了过度自信预测，证明可靠性与性能并非相互排斥。
• 该方法具有良好的泛化能力：GNN-SNGP 在 molHIV、BBBP 和 BACE 数据集上均实现一致改进，证实其在心脏毒性之外的更广泛应用潜力。