[论文解读] Calibration and generalizability of probabilistic models on low-data chemical datasets with DIONYSUS

Gary Tom, Riley J. Hickman+4 more|arXiv (Cornell University)|Dec 3, 2022

Machine Learning in Materials Science被引用 4

一句话总结

本论文在低数据量化学数据集（≤2000个分子）上评估了概率机器学习模型的校准性和泛化能力，采用多种分子表征方法和模型。结果表明，贝叶斯神经网络（BNNs）和带有概率输出头的图神经网络（GNNGP）在低数据环境下展现出优越的不确定性校准能力和分布外泛化性能，并提出了DIONYSUS框架，以实现小分子性质预测的可复现、开源分析。

ABSTRACT

Deep learning models that leverage large datasets are often the state of the art for modelling molecular properties. When the datasets are smaller (< 2000 molecules), it is not clear that deep learning approaches are the right modelling tool. In this work we perform an extensive study of the calibration and generalizability of probabilistic machine learning models on small chemical datasets. Using different molecular representations and models, we analyse the quality of their predictions and uncertainties in a variety of tasks (binary, regression) and datasets. We also introduce two simulated experiments that evaluate their performance: (1) Bayesian optimization guided molecular design, (2) inference on out-of-distribution data via ablated cluster splits. We offer practical insights into model and feature choice for modelling small chemical datasets, a common scenario in new chemical experiments. We have packaged our analysis into the DIONYSUS repository, which is open sourced to aid in reproducibility and extension to new datasets.

研究动机与目标

• 解决在小样本数据集（<2000个分子）中分子性质预测的不确定性校准与泛化能力缺乏系统性评估的问题，此类情况在早期药物发现中十分常见。
• 研究不同分子表征方法（Morgan指纹、Mordred描述符、基于图的表征）和概率模型（SNGP、GP、BNN、NGBoost、GNNGP）对预测性能和不确定性估计的影响。
• 在实际高风险场景中评估模型性能，例如通过截断聚类分割模拟新型骨架结构的分布外泛化，以及基于贝叶斯优化的分子设计。
• 为低数据量分子机器学习中的模型与特征选择提供可操作的见解，尤其面向数据有限的实验研究人员。
• 开发并发布DIONYSUS开源代码库，以确保研究可复现，并支持未来扩展至新数据集和新模型。

提出的方法

• 在5个多样化的化学数据集（Delaney、Freesolv、BACE、RBioDeg、BBBP）上评估多种概率模型——SNGP、高斯过程（GP）、贝叶斯神经网络（BNN）、NGBoost以及带有概率输出头的图神经网络（GNNGP）。
• 采用三种分子表征策略：Morgan指纹（MFP）、Mordred描述符（Mordred）以及基于图的表征（如SMILES生成的节点/边特征）。
• 使用期望校准误差（ECE）评估模型校准性，使用AUROC（分类任务）和RMSE（回归任务）评估预测性能，并报告95%置信区间。
• 开展两项模拟实验：(1) 基于贝叶斯优化的分子设计（最小化半衰期、自由能），(2) 通过HDBScan聚类分割模拟新型骨架结构的分布外泛化。
• 应用UMAP降维技术可视化化学空间，验证聚类结构的合理性，确保分布外分割具有实际意义。
• 将完整的分析流程以开源形式发布至DIONYSUS代码库，以支持可复现性并便于未来扩展。

实验结果

研究问题

• RQ1 在低数据量化学数据集（<2000个分子）上，不同概率模型（SNGP、GP、BNN、NGBoost、GNNGP）在不确定性校准与泛化能力方面的表现如何？
• RQ2 在多样化的分子任务中，哪种分子表征方法（Morgan指纹、Mordred描述符、基于图的表征）能产生最稳健且校准良好的预测结果？
• RQ3 这些模型在分布外化学结构（尤其是新型分子骨架）上的泛化能力如何？通过截断聚类分割方法进行评估。
• RQ4 在低数据设置下，基于这些模型的贝叶斯优化在多大程度上能高效发现最优分子（如最小化生物降解半衰期或溶剂化自由能）？
• RQ5 模型不确定性校准在高风险分子设计决策中具有何种实际意义？例如预测药物性肝损伤或血脑屏障通透性。

主要发现

• 贝叶斯神经网络（BNNs）和GNNGP在不确定性校准方面表现最佳，其在BBBP数据集上的ECE值分别为0.029和0.041，表明其不确定性估计具有高度可信度。
• 在BACE数据集上，GNNGP实现了AUROC为0.879 ± 0.039，ECE为0.127 ± 0.050，显示出在低数据量二分类任务中的强大性能。
• 在贝叶斯优化实验中，BNN和GP等模型在最小化溶剂化自由能和生物降解半衰期方面持续优于其他模型，且在20–30次迭代内即收敛至最优分子。
• 基于图的表征在不确定性校准和预测性能方面普遍优于Mordred和MFP，尤其在复杂分子任务中表现更优。
• 截断聚类分割协议成功评估了分布外泛化能力，结果表明BNN和GNNGP等模型在新型骨架结构上仍能保持高预测准确性。
• DIONYSUS框架展现出良好的可复现性与可扩展性，可在极低配置开销下实现对多样化数据集与模型配置的一致性评估。