[论文解读] MUBen: Benchmarking the Uncertainty of Molecular Representation Models

Yinghao Li, Lingkai Kong+5 more|arXiv (Cornell University)|Jun 14, 2023

Computational Drug Discovery MethodsComputer Science被引用 3

一句话总结

MUBen 是一个全面的基准测试，用于评估在多种分子性质预测任务中，基于先进分子表征模型的不确定性量化（UQ）方法。该基准使用多种主干模型和UQ技术，揭示了主干模型选择和UQ方法选择对校准和不确定性估计具有显著影响，其中集成方法和基于SGLD的UQ方法在NLL和CE等关键指标上表现更优。

ABSTRACT

Large molecular representation models pre-trained on massive unlabeled data have shown great success in predicting molecular properties. However, these models may tend to overfit the fine-tuning data, resulting in over-confident predictions on test data that fall outside of the training distribution. To address this issue, uncertainty quantification (UQ) methods can be used to improve the models' calibration of predictions. Although many UQ approaches exist, not all of them lead to improved performance. While some studies have included UQ to improve molecular pre-trained models, the process of selecting suitable backbone and UQ methods for reliable molecular uncertainty estimation remains underexplored. To address this gap, we present MUBen, which evaluates different UQ methods for state-of-the-art backbone molecular representation models to investigate their capabilities. By fine-tuning various backbones using different molecular descriptors as inputs with UQ methods from different categories, we assess the influence of architectural decisions and training strategies. Our study offers insights for selecting UQ for backbone models, which can facilitate research on uncertainty-critical applications in fields such as materials science and drug discovery.

研究动机与目标

• 为解决预训练分子表征模型中不确定性量化（UQ）方法缺乏系统性评估的问题。
• 研究架构选择（主干模型）和训练策略如何影响在多样化分子性质预测任务中UQ方法的表现。
• 为在不确定性敏感的应用（如药物发现和材料科学）中选择最优UQ方法和主干模型提供可操作的见解。
• 建立一个用户友好、可扩展的基准框架，支持分子AI中UQ的可复现评估。

提出的方法

• MUBen 在5个分子数据集上，对3种主干模型（ChemBERTa、GROVER、Uni-Mol）应用6种UQ方法（确定性预测、MC Dropout、SWAG、BBP、SGLD和集成方法）进行评估。
• 该基准使用SMILES、分子图和2D拓扑表示等分子描述符作为不同主干模型的输入。
• 通过固定主干权重和微调主干权重两种方式执行微调，以评估表征学习对不确定性校准的影响。
• 评估指标包括RMSE、MAE、NLL（负对数似然）和CE（分类误差），结果在分类和回归任务中均有报告。
• 该框架以代码形式实现，并发布在GitHub上，以支持可扩展性和可复现性。
• 实验涵盖QM7、Lipophilicity等多样化数据集，覆盖量子、物理和生物性质。

Figure 1: An overview of MUBen with all datasets, backbone models, UQ methods, and metrics enumerated. Elements in green and blue are related to classification and regression tasks, respectively.

实验结果

研究问题

• RQ1 在预训练分子表征模型上应用时，哪些UQ方法能产生最校准的不确定性估计？
• RQ2 不同主干架构（如ChemBERTa、GROVER、Uni-Mol）如何影响UQ方法的表现？
• RQ3 与使用固定权重相比，微调主干权重是否能改善不确定性校准？
• RQ4 UQ方法在不同类型分子性质预测任务（分类与回归）中的表现如何？
• RQ5 架构设计与UQ策略对不确定性估计质量的相对影响是什么？

主要发现

• 集成方法和基于SGLD的UQ在所有数据集中持续实现最低的NLL和CE值，表明其具有更优的校准性和不确定性估计能力。
• 在Lipophilicity数据集中，ChemBERTa搭配集成UQ方法实现了0.9837 ± 0.0001的RMSE，为所有方法中的最低值，表现出优异的回归性能。
• SGLD和BBP方法在QM7上常产生较高的NLL和CE，表明存在过度自信或校准不足，尤其在Uni-Mol和GROVER主干上更为明显。
• SWAG和MC Dropout表现中等，其中SWAG在ChemBERTa（NLL: 5.4107 ± 0.0177）和GROVER（NLL: 5.3670 ± 0.0446）上取得了具有竞争力的结果。
• 固定主干权重通常能产生更稳定的不确定性估计，尤其对集成和SGLD方法而言，表明微调可能损害校准性能。
• Uni-Mol搭配BBP UQ在QM7上表现出最高的NLL（6.0873 ± 0.0082），表明不确定性校准较差，而其集成变体实现了5.7658 ± 0.0082的NLL，性能显著提升。

(a) Classification MRR $\uparrow$