[论文解读] A comprehensive study on the prediction reliability of graph neural networks for virtual screening

Soojung Yang, Kyung Hoon Lee+1 more|arXiv (Cornell University)|Mar 17, 2020

Machine Learning in Materials Science被引用 5

一句话总结

本文对图神经网络（GNNs）在虚拟筛选中预测可靠性的全面分析，强调了输出的概率解释。在数据稀缺和类别不平衡条件下，评估了模型架构、正则化和损失函数的影响，发现适当的正则化和调整α权重的焦点损失能显著提高可靠性与筛选成功率，尤其在类别不平衡的情况下。

ABSTRACT

Prediction models based on deep neural networks are increasingly gaining attention for fast and accurate virtual screening systems. For decision makings in virtual screening, researchers find it useful to interpret an output of classification system as probability, since such interpretation allows them to filter out more desirable compounds. However, probabilistic interpretation cannot be correct for models that hold over-parameterization problems or inappropriate regularizations, leading to unreliable prediction and decision making. In this regard, we concern the reliability of neural prediction models on molecular properties, especially when models are trained with sparse data points and imbalanced distributions. This work aims to propose guidelines for training reliable models, we thus provide methodological details and ablation studies on the following train principles. We investigate the effects of model architectures, regularization methods, and loss functions on the prediction performance and reliability of classification results. Moreover, we evaluate prediction reliability of models on virtual screening scenario. Our result highlights that correct choice of regularization and inference methods is evidently important to achieve high success rate, especially in data imbalanced situation. All experiments were performed under a single unified model implementation to alleviate external randomness in model training and to enable precise comparison of results.

研究动机与目标

• 为解决深度学习模型在虚拟筛选中因数据稀缺和类别不平衡导致的不可靠性（如过度自信和校准性差）问题。
• 评估模型架构、正则化和损失函数对分子性质分类中预测可靠性的影响。
• 为训练生成可信概率估计的GNN提供可操作的指导建议，以支持化合物筛选。
• 不仅通过准确率，还通过真实世界虚拟筛选场景中的校准性和置信度可靠性来评估模型性能。
• 推荐针对高精度或高召回率虚拟筛选应用的最优训练策略——特别是正则化和损失函数的选择。

提出的方法

• 本研究在所有实验中采用统一的GNN实现，以最小化训练随机性并确保公平比较。
• 在相同的训练和评估协议下，评估了多种GNN架构，包括GCN、GIN和GraphSAGE。
• 系统分析了权重衰减、Dropout和蒙特卡洛Dropout（MC-DO）等正则化技术对模型校准的影响。
• 对比了二元交叉熵、焦点损失（可调α和γ）以及加权交叉熵等损失函数在性能和可靠性方面的表现。
• 通过校准曲线、期望校准误差（ECE）和熵直方图量化预测可靠性，以评估置信度与准确度的一致性。
• 评估在类别不平衡的数据集上进行分子性质分类任务，模拟真实虚拟筛选条件。

实验结果

研究问题

• RQ1 不同GNN架构在分子性质分类中如何影响概率预测的可靠性？
• RQ2 在数据稀缺和类别不平衡条件下，正则化（尤其是权重衰减和MC-DO）在多大程度上能改善模型校准？
• RQ3 各种损失函数，特别是具有不同α和γ值的焦点损失，如何影响预测可靠性和筛选表现？
• RQ4 当数据稀疏且不平衡时，GNN的概率输出是否可以被信任用于虚拟筛选？
• RQ5 哪些训练和推理策略能在保持可靠置信度估计的同时，实现最高的虚拟筛选成功率？

主要发现

• 适当的正则化，尤其是蒙特卡洛Dropout（MC-DO），能显著改善模型校准并减少预测过度自信。
• 在类别不平衡数据设置中，采用α ≠ 0.5且γ > 0.0的焦点损失优于标准交叉熵，提升了可靠性和筛选成功率。
• 使用α ≠ 0.5的加权交叉熵在处理类别不平衡方面非常有效，提供了更好的精确率与召回率之间的权衡。
• 焦点损失中较高的γ值会因损失的固有类别敏感性而降低预测可靠性，导致对少数类的过拟合。
• 采用适当正则化和损失函数训练的模型能实现更高的期望校准误差（ECE）降低，并更好地实现置信度与准确度的一致性。
• 本研究建议避免对置信度进行惩罚，而应通过在损失函数中使用定制化的α权重，以实现针对特定虚拟筛选目标的高精度或高召回率模型。