[論文レビュー] 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に基づく手法が、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手法とバックボーンの選定に役立つインサイトを提供すること。
• 再現性のあるUQ評価を支援する、ユーザーフレンドリーで拡張可能なベンチマークフレームワークを確立すること。

提案手法

• MUBen は、6 種類のUQ手法（決定的予測、MCドロップアウト、SWAG、BBP、SGLD、アンサンブル）を、3 種類のバックボーンモデル（ChemBERTa、GROVER、Uni-Mol）に適用し、5 つの分子データセットで評価した。
• ベンチマークは、SMILES、分子グラフ、2次元トポロジカル表現といった分子記述子を、異なるバックボーンの入力として用いた。
• 表現学習が不確実性キャリブレーションに与える影響を評価するため、固定済みおよび微調整済みのバックボーン重みを用いてファインチューニングを実施した。
• 評価指標には、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を適用した場合、RMSE が 0.9837 ± 0.0001 にまで低下し、他のすべての手法と比較して優れた回帰性能を示した。
• QM7 では、SGLD および BBP 手法が、特に Uni-Mol および GROVER バックボーンと組み合わせると、NLL および CE が高くなり、過信またはキャリブレーション不良を示した。
• SWAG および MC ドロップアウトは中程度の性能を示し、SWAG は ChemBERTa（NLL: 5.4107 ± 0.0177）および GROVER（NLL: 5.3670 ± 0.0446）で競争力のある結果を得た。
• 固定されたバックボーン重みは、特にアンサンブルおよびSGLD手法において、より安定した不確実性推定をもたらした。これは、微調整がキャリブレーションを劣化させる可能性があることを示唆している。
• Uni-Mol に BBP UQ を適用した場合、QM7 で最高の NLL（6.0873 ± 0.0082）を記録し、不確実性キャリブレーションが著しく劣化していた。一方、アンサンブルバージョンでは NLL が 5.7658 ± 0.0082 に改善され、性能の向上が確認された。

(a) Classification MRR $\uparrow$