[논문 리뷰] ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction

Seyone Chithrananda, Gabriel Grand외 1명|arXiv (Cornell University)|2020. 10. 19.

Computational Drug Discovery Methods참고 문헌 29인용 수 378

한 줄 요약

ChemBERTa는 77M PubChem SMILES에 대해 마스킹 언어 모델링으로 트랜스포머를 사전 학습하고 MoleculeNet에서 하류 성능을 평가하여 대규모 전 pretraining의 이점과 PubChem-77M 데이터 세트를 제공합니다.

ABSTRACT

GNNs and chemical fingerprints are the predominant approaches to representing molecules for property prediction. However, in NLP, transformers have become the de-facto standard for representation learning thanks to their strong downstream task transfer. In parallel, the software ecosystem around transformers is maturing rapidly, with libraries like HuggingFace and BertViz enabling streamlined training and introspection. In this work, we make one of the first attempts to systematically evaluate transformers on molecular property prediction tasks via our ChemBERTa model. ChemBERTa scales well with pretraining dataset size, offering competitive downstream performance on MoleculeNet and useful attention-based visualization modalities. Our results suggest that transformers offer a promising avenue of future work for molecular representation learning and property prediction. To facilitate these efforts, we release a curated dataset of 77M SMILES from PubChem suitable for large-scale self-supervised pretraining.

연구 동기 및 목표

• 트랜스포머 기반 사전 학습이 분자 특성 예측 성능에 미치는 영향을 시연한다.
• 사전 학습 데이터세트 크기, 토크나이제이션, 문자열 표현이 하류 작업에 미치는 영향을 평가한다.
• 자체 지도 사전 학습을 위한 대규모 공개 SMILES 데이터세트(PubChem 77M)를 제공한다.
• SMILES와 SELFIES 표현 및 두 토크나이저(BPE와 SmilesTokenizer)의 비교를 이 맥락에서 수행한다.

제안 방법

• 분자 문자열에 대해 RoBERTa 스타일 트랜스포머(12 계층, 12 어텐션 헤드)를 적응시킨다.
• 77M PubChem SMILES에서 15% 마스킹 및 시퀀스 길이 512로 MLM으로 사전 학습한다.
• 토크나이저(BPE vs SmilesTokenizer) 및 표현(SMILES vs SELFIES)을 실험한다.
• MoleculeNet 분류 과제에서 사전 학습된 모델을 80/10/10 분할 및 ROC-AUC를 기준으로 조기 종료를 사용해 미세 조정한다.
• 선정된 과제에서 Chemprop 기준선(D-MPNN, RF, SVM)과의 성능 비교 및 사전 학습 크기에 따른 스케일링 분석을 수행한다.
• BertViz를 사용해 주의(attention) 시각화를 제공하여 모델 헤드와 화학적 특징 간의 관련성을 설명한다.

실험 결과

연구 질문

• RQ1 MLM 사전 학습이 MoleculeNet 과제에서 분자 특성 예측을 개선하는가?
• RQ2 사전 학습 데이터 크기(100K–10M)가 하류 성능에 어떤 영향을 주는가?
• RQ3 토크나이즈 전략(BPE vs SmilesTokenizer) 및 문자열 표현(SMILES vs SELFIES)이 결과에 유의미한 영향을 미치는가?
• RQ4 주의 시각화가 트랜스포머 헤드의 화학 관련 패턴을 드러내는가?

주요 결과

작업	ChemBERTa 10M ROC-AUC	ChemBERTa 10M PRC-AUC	D-MPNN ROC-AUC	D-MPNN PRC-AUC	RF ROC-AUC	RF PRC-AUC	SVM ROC-AUC	SVM PRC-AUC
BBBP	0.643	0.620	0.708	0.697	0.681	0.692	0.702	0.724
ClinTox (CT_TOX)	0.733	0.975	0.906	0.993	0.693	0.968	0.833	0.986
HIV	0.622	0.119	0.752	0.152	0.780	0.383	0.763	0.364
Tox21 (SR-p53)	0.728	0.207	0.688	0.429	0.724	0.335	0.708	0.345

• ChemBERTa는 상위 Chemprop 기준선에 근접하지만 특정 MoleculeNet 과제에서 능가하지는 못한다.
• 100K에서 10M으로 사전 학습 크기를 확장하면 BBBP, ClinTox, Tox21에서 평균 ROC-AUC 증가(+0.110) 및 PRC-AUC 증가(+0.059)의 체계적인 이점을 보인다.
• 10M-PubChem 사전 학습에서 ChemBERTa는 BBBP ROC-AUC 0.643 및 PRC-AUC 0.620; ClinTox ROC-AUC 0.733 및 PRC-AUC 0.975; HIV ROC-AUC 0.622 및 PRC-AUC 0.119; Tox21 ROC-AUC 0.728 및 PRC-AUC 0.207를 달성한다.
• SmilesTokenizer가 한 과제에서 BPE보다 약간 우수했다(PRC-AUC +0.015 on Tox21 SR-p53).
• SELFIES 표현은 SMILES에 비해 Tox21 SR-p53에서 실질적인 하류 성능 차이를 보이지 않았다.
• BertViz를 이용한 주의 시각화는 일부 헤드가 기능 그룹과 방향족 고리 등에 선택적으로 주의를 기울이며 괄호 유사 구조를 추적하는 헤드도 있음을 보여주었다.