[论文解读] A benchmark study on reliable molecular supervised learning via Bayesian learning

Doyeong Hwang, Grace Lee+3 more|arXiv (Cornell University)|Jun 12, 2020

Computational Drug Discovery Methods参考文献 28被引用 4

一句话总结

本文对用于提升分子图神经网络（GNN）预测可靠性的贝叶斯学习方法进行了基准测试，表明诸如随机权重平均（SWA）和SWAG等技术可在多种GNN架构和分子性质预测任务中实现校准良好、具备不确定性感知的预测。研究结果表明，贝叶斯方法可降低模型的过度自信，提升虚拟筛选的成功率，提供可信的预测概率。

ABSTRACT

Virtual screening aims to find desirable compounds from chemical library by using computational methods. For this purpose with machine learning, model outputs that can be interpreted as predictive probability will be beneficial, in that a high prediction score corresponds to high probability of correctness. In this work, we present a study on the prediction performance and reliability of graph neural networks trained with the recently proposed Bayesian learning algorithms. Our work shows that Bayesian learning algorithms allow well-calibrated predictions for various GNN architectures and classification tasks. Also, we show the implications of reliable predictions on virtual screening, where Bayesian learning may lead to higher success in finding hit compounds.

研究动机与目标

• 评估图神经网络中使用贝叶斯深度学习进行分子性质预测的可靠性。
• 解决神经网络在虚拟筛选中对分布外（OOD）化合物预测过度自信的问题。
• 评估贝叶斯推理方法是否能改善基于GNN的分子分类中的校准性与不确定性估计。
• 为可靠的分子机器学习提供实用基准，尤其针对药物发现等紧急应用场景。
• 发布代码以支持可复现性，并推动计算化学与虚拟筛选领域的广泛应用。

提出的方法

• 本研究采用期望校准误差（ECE）定量评估预测可靠性，衡量预测置信度与实际正确率之间的差异。
• 通过近似推理方法（MAP、深度集成、MC-Dropout、SGLD、SWA、SWAG）应用贝叶斯学习，以估计模型权重的后验分布。
• 采用图神经网络（GIN、GAT、GatedGCN等）作为主干模型，输入特征来自分子图及节点/边表示。
• 对于贝叶斯边际化，通过多次采样权重的预测结果进行平均，利用预测概率的方差实现不确定性估计。
• 模型在四个MoleculeNet数据集（BACE、BBBP、HIV和Tox21）上使用骨架划分的数据分割（80:10:10）进行训练，各方法的超参数独立调优。
• SWA和SWAG采用循环学习率调度，并在150–250个周期内进行权重平均，SWAG的后验协方差按1.0缩放。

实验结果

研究问题

• RQ1 贝叶斯学习方法是否能在多种架构和任务中提升基于GNN的分子性质预测器的校准性？
• RQ2 不同贝叶斯推理技术在分子分类基准上的预测可靠性与性能表现如何比较？
• RQ3 校准良好的预测在多大程度上通过减少假阳性和假阴性，提升了虚拟筛选的成功率？
• RQ4 不确定性估计对真正例、假阳性、真阴性和假阴性预测行为的影响如何？
• RQ5 SWA和SWAG相较于其他贝叶斯方法，在实现高准确率与低ECE方面表现如何？

主要发现

• SWA及其变体SWAG在所有GNN架构和分子性质预测任务中均持续实现了最低的期望校准误差（ECE），表明其具有更优的校准性能。
• 使用贝叶斯学习方法显著降低了过度自信，预测概率与实际正确率高度一致，尤其在分布外（OOD）样本上表现更优。
• 直方图分析显示，真正例通常被赋予高预测概率，而假阳性则多以低置信度预测——表明不确定性量化具有可靠性。
• 采用SWA和SWAG的模型在BACE和BBBP数据集上实现了AUROC（>0.85）和F1分数（>0.75）的高分，同时ECE值低于0.05，展现出高性能与高可靠性兼具。
• 研究发现MC-Dropout和深度集成也改善了校准性，但其表现不如SWA和SWAG稳定，尤其在类别不平衡的Tox21数据集上。
• 在https://github.com/AITRICS/mol_reliable_gnn发布的代码库支持可复现的基准测试，并可支持药物发现等紧急应用场景，如针对COVID-19的研究。