[논문 리뷰] Calibration and generalizability of probabilistic models on low-data chemical datasets with DIONYSUS

Gary Tom, Riley J. Hickman외 4명|arXiv (Cornell University)|2022. 12. 03.

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 다양한 확률적 모델(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 프레임워크는 재현 가능성과 확장성을 입증하였으며, 최소한의 설정 오버헤드로 다양한 데이터셋과 모델 구성에서 일관된 평가를 가능하게 하였다.