[논문 리뷰] MUBen: Benchmarking the Uncertainty of Molecular Representation Models

Yinghao Li, Lingkai Kong외 5명|arXiv (Cornell University)|2023. 06. 14.

Computational Drug Discovery MethodsComputer Science인용 수 3

한 줄 요약

MUBen은 다양한 백본 모델과 불확실성 추정 기법을 사용하여 여러 분자의 성질 예측 작업에서 최신 분자 표현 모델의 불확실성 정량화(UQ) 방법을 종합적으로 평가하는 벤치마크이다. 이는 백본 선택과 UQ 방법 선택이 校정(calibration)과 불확실성 추정에 중대한 영향을 미친다는 것을 드러내며, 앙상블 및 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 방법과 백본을 선택하는 데 실질적인 통찰을 제공하기 위해.
• 재현 가능성을 보장하며 확장성이 뛰어난 사용자 우호적인 확장 가능한 벤치마크 프레임워크를 구축하기 위해.

제안 방법

• MUBen은 6종의 UQ 방법(정적 예측, MC Dropout, 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 방법은 높은 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)을 기록하여 불확실성 추정의 열악한 校정을 보였으나, 앙상블 버전은 NLL 5.7658 ± 0.0082를 기록하여 성능 향상을 보였다.

(a) Classification MRR $\uparrow$