[论文解读] Bayesian Graph Neural Networks for Molecular Property Prediction

George S. Lamb, Brooks Paige|arXiv (Cornell University)|Nov 25, 2020

Machine Learning in Materials Science参考文献 30被引用 9

一句话总结

本论文提出使用不确定性量化方法的贝叶斯图神经网络（GNNs），在QM9数据集上对分子性质预测任务中的消息传递和读出层进行处理。实验表明，与标准深度学习基线相比，完整的贝叶斯GNN显著提升了预测准确性、校准性以及下游分子搜索任务的性能。

ABSTRACT

Graph neural networks for molecular property prediction are frequently underspecified by data and fail to generalise to new scaffolds at test time. A potential solution is Bayesian learning, which can capture our uncertainty in the model parameters. This study benchmarks a set of Bayesian methods applied to a directed MPNN, using the QM9 regression dataset. We find that capturing uncertainty in both readout and message passing parameters yields enhanced predictive accuracy, calibration, and performance on a downstream molecular search task.

研究动机与目标

• 解决GNN在分布外分子骨架上泛化能力差的问题，该问题源于对训练数据的过拟合。
• 提升分子性质预测中的模型校准性和不确定性量化，特别是在回归设置中。
• 评估贝叶斯推理在GNN的消息传递和读出组件中的影响。
• 在真实世界的分子搜索任务中基准化贝叶斯方法，以评估其实际应用价值。
• 提供开源的PyTorch代码，以支持可复现性与进一步研究。

提出的方法

• 采用有向消息传递神经网络（D-MPNN）作为基础GNN架构，用于分子表征学习。
• 应用八种贝叶斯推理方法：MAP、GP、MC Dropout（DropR、DropA）、SWAG、SGLD、BBP和DUN，用于不确定性估计。
• 使用后处理t-分布似然函数以改善校准性，尤其在未显式建模类内不确定性时。
• 通过从后验近似中采样，实现基于集成的预测分布，以估计不确定性。
• 采用变分推理或基于MCMC的方法端到端训练模型，以近似模型参数的后验分布。
• 通过可靠性图评估不确定性，并使用平均排名和校准面积作为12个QM9回归任务中的关键指标。

实验结果

研究问题

• RQ1 在消息传递和读出层中同时应用贝叶斯学习，是否能提升分子性质预测中的预测准确性和校准性？
• RQ2 不同贝叶斯推理方法在QM9回归任务中的不确定性量化和性能表现如何比较？
• RQ3 贝叶斯GNN是否能提升下游分子生成或搜索任务的性能？
• RQ4 建模类内不确定性对校准性的影响如何，特别是在使用高斯分布与t-分布似然函数时？
• RQ5 不确定性量化如何影响模型在未见分子骨架上的泛化能力？

主要发现

• 完整的贝叶斯GNN，尤其是同时在消息传递和读出层中引入不确定性的模型，在QM9数据集上实现了1.95 ± 0.40的最优平均排名，显著优于标准MAP训练方法（4.08 ± 0.16）。
• DropA方法（在整棵GNN中应用MC Dropout）在分子搜索任务中取得了77.52 ± 0.77%的得分，优于MAP（72.22 ± 0.57%）及其他贝叶斯方法。
• 使用后处理t-分布似然函数改善了校准性，减少了使用高斯似然函数时观察到的病理性过度自信现象。
• 采用贝叶斯读出层的模型（如GP、DropR）校准性优于仅在消息传递层中引入贝叶斯方法的模型，表明读出阶段的不确定性对预测置信度更为关键。
• SGLD和BBP方法虽然给出了较高的不确定性估计，但预测准确性较差，表明在该设置下MCMC和变分方法可能无法始终实现良好泛化。
• DUN方法（采用对角加低秩协方差结构）在平均排名（4.40 ± 0.25）和搜索得分（76.02 ± 1.09）方面表现优异，表明其在不确定性估计方面具有强鲁棒性。