[論文レビュー] 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.

研究の動機と目的

• 初期段階のドラッグディスcoveryにおいて一般的な小規模データセット（2000分子未満）における分子物性予測の不確実性キャリブレーションと一般化性能に関する体系的評価の不足に対処する。
• 異なる分子特徴抽出法（Morganフィンガープrint、Mordred記述子、グラフベース）と確率的モデル（SNGP、GP、BNN、NGBoost、GNNGP）が予測精度と不確実性推定に与える影響を調査する。
• ベイジアン最適化を用いた分子設計や、アブレーテッドクラスタースプリットを用いた分布外一般化といった、実用的で高リスクなシナリオにおけるモデル性能を評価する。
• 特に限られたデータしか入手できない実験研究者を想定し、低データ分子機械学習におけるモデルおよび特徴抽出法の選定に役立つ実用的知見を提供する。
• 再現可能性を確保し、新たなデータセットやモデルへの拡張を可能にするために、DIONYSUSオープンソースリポジトリを構築・公開する。

提案手法

• 5つの多様な化学データセット（Delaney, Freesolv, BACE, RBioDeg, BBBP）に対して、SNGP、ガウス過程（GP）、ベイジアンニューラルネットワーク（BNN）、NGBoost、および確率的ヘッドを備えたグラフニューラルネットワーク（GNNGP）の複数の確率的モデルを評価する。
• 3つの分子特徴抽出戦略を用いる：Morganフィンガープリント（MFP）、Mordred記述子（Mordred）、およびSMILESからのノード/エッジ特徴を用いたグラフベース表現。
• 期待キャリブレーション誤差（ECE）を用いたキャリブレーション評価と、AUROC（分類）およびRMSE（回帰）を用いた予測性能評価を行い、95%信頼区間を提示する。
• 2つのシミュレート実験を実施：(1) 分子設計のためのベイジアン最適化（半減期・自由エネルギーの最小化）、(2) HDBScanを用いたクラスタースプリットを用いた分布外一般化の評価。
• 化学空間の可視化とクラスタ構造の妥当性を検証するために、UMAPを用いた次元削減を実施し、意味のある分布外スプリットを確保する。
• 完全な分析パイプラインをオープンソースのDIONYSUSリポジトリとして公開し、再現可能性と将来的な拡張性を支援する。

実験結果

リサーチクエスチョン

• RQ1 異なる確率的モデル（SNGP、GP、BNN、NGBoost、GNNGP）は、低データ化学データセット（2000分子未満）において、キャリブレーションと一般化性能の観点でどのように性能を発揮するか？
• 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フレームワークは再現性とスケーラビリティを示し、最小限のセットアップコストで多様なデータセットとモデル設定に対して一貫した評価を可能にした。