[논문 리뷰] Self-Supervised Graph Transformer on Large-Scale Molecular Data

Yu Rong, Yatao Bian외 5명|arXiv (Cornell University)|2020. 06. 18.

Computational Drug Discovery Methods참고 문헌 62인용 수 415

한 줄 요약

GROVER는 1,000만 개의 라벨이 없는 분자에서 자기지도학습 작업으로 그래프 트랜스포머를 사전 학습하고, 미세 조정을 거친 후 MoleculeNet의 11개 벤치마크에서 평균 약 6%의 큰 향상을 달성합니다.

ABSTRACT

How to obtain informative representations of molecules is a crucial prerequisite in AI-driven drug design and discovery. Recent researches abstract molecules as graphs and employ Graph Neural Networks (GNNs) for molecular representation learning. Nevertheless, two issues impede the usage of GNNs in real scenarios: (1) insufficient labeled molecules for supervised training; (2) poor generalization capability to new-synthesized molecules. To address them both, we propose a novel framework, GROVER, which stands for Graph Representation frOm self-superVised mEssage passing tRansformer. With carefully designed self-supervised tasks in node-, edge- and graph-level, GROVER can learn rich structural and semantic information of molecules from enormous unlabelled molecular data. Rather, to encode such complex information, GROVER integrates Message Passing Networks into the Transformer-style architecture to deliver a class of more expressive encoders of molecules. The flexibility of GROVER allows it to be trained efficiently on large-scale molecular dataset without requiring any supervision, thus being immunized to the two issues mentioned above. We pre-train GROVER with 100 million parameters on 10 million unlabelled molecules -- the biggest GNN and the largest training dataset in molecular representation learning. We then leverage the pre-trained GROVER for molecular property prediction followed by task-specific fine-tuning, where we observe a huge improvement (more than 6% on average) from current state-of-the-art methods on 11 challenging benchmarks. The insights we gained are that well-designed self-supervision losses and largely-expressive pre-trained models enjoy the significant potential on performance boosting.

연구 동기 및 목표

• 제한된 라벨 데이터와 큰 화학 공간 하에서 강력한 분자 표현의 필요성에 대한 동기를 제시한다.
• 일반화를 향상시키기 위한 분자 그래프의 자기지도 사전 학습 프레임워크를 제안한다.
• 그래프 인식 어텐션과 동적 메시지 전달을 갖춘 트랜스포머 기반 인코더를 설계한다.
• 대규모의 라벨 없는 데이터에서의 사전 학습이 다운스트림 분자 특성 예측을 향상시킴을 보인다.

제안 방법

• 자체 지도성 메시지 전달 트랜스포머에서 노드/에지 GNN 트랜스포머로 구성된 GROVER를 도입한다.
• GNN 기반의 질의/키/값을 Transformer 인코더에 전체 노드에 대해 공급하는 이중 수준 정보 추출을 사용한다.
• 일반화 향상을 위해 임의의 도약 수를 갖는 동적 메시지 전달(dyMPN)을 구현한다.
• 노드/에지 맥락화된 속성 예측을 노드/에지 수준의 자기지도학습으로 설계한다.
• RDKit로 탐지된 모티프를 다중 레이블 타깃으로 사용하는 그래프 수준 모티프 예측을 도입한다.
• ZINC15, ChEMBL의 1,100만 분자에서 100M 파라미터 모델로 250개의 GPU에 걸쳐 11M 라벨 없는 분자를 사전 학습한다.

실험 결과

연구 질문

• RQ1 대규모의 라벨이 없는 분자 그래프에서의 자기지도 사전 학습이 미세 조정 후 다운스트림 속성 예측을 향상시키는가?
• RQ2 그래프 인식 트랜스포머 인코더와 동적 메시지 전달이 전통적인 GNN보다 더 우수한 표현을 제공하는가?
• RQ3 맥락 인식 노드/에지 및 모티프 기반 그래프 수준 사전 텍스트 작업이 성능과 일반화에 어떤 영향을 미치는가?
• RQ4 GROVER가 MoleculeNet 벤치마크에서 모델 규모와 학습 데이터와 함께 확장될 때의 성능은 어떠한가?

주요 결과

• GROVER는 11개 데이터셋 모두에서 일관되게 최상의 성능을 달성하며 평균 상대 개선은 6.1%(분류에서 2.2%, 회귀에서 10.8%)이다.
• GROVER_large는 모든 데이터셋에서 최첨단 베이스라인을 능가하며, GROVER_base는 8/11 데이터셋에서 우수하다.
• 자기지도 학습 사전 학습은 분류 과제에서 사전 학습 없음에 비해 평균 AUC가 3.8% 증가하도록 제공하며, 특히 작은 데이터셋에 도움을 준다.
• 요소 분해 실험에서 GROVER의 GTransformer 백본은 GIN 및 MPNN 백본보다 표현력이 높음을 확인하며, dyMPN은 학습 손실의 약간의 영향에도 불구하고 일반화를 개선한다.
• 소형 라벨 체제에서 GROVER는 SOTA 대비 23.9%의 상대 향상을 보이며 뚜렷한 이익을 얻는다.