[논문 리뷰] Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets

Dominique Beaini, Shenyang Huang외 33명|arXiv (Cornell University)|2023. 01. 01.

Machine Learning in Materials Science인용 수 6

한 줄 요약

이 논문은 분자 기계학습에서 기초 모델을 훈련하기 위해 필수적인 대규모, 레이블이 부여된 다중 작업 분자 데이터셋의 새로운 가족인 ToyMix, LargeMix, UltraLarge를 소개한다. 이 데이터셋은 약 1억 개의 분자와 3,000개의 작업을 통해 130억 개 이상의 양자 및 생물학적 성질 레이블을 포함한다. 저자들은 효율적인 다중 작업 및 다중 수준 그래프 학습을 위한 전용 딥러닝 라이브러리인 Graphium을 제시하며, 다양한 감독 데이터를 기반으로 한 사전 훈련이 자원이 부족한 생물학적 작업에서 성능을 향상시킨다는 것을 입증함으로써, 분자 AI 분야에서 기초 모델의 실현 가능성을 뒷받침한다.

ABSTRACT

Recently, pre-trained foundation models have enabled significant advancements in multiple fields. In molecular machine learning, however, where datasets are often hand-curated, and hence typically small, the lack of datasets with labeled features, and codebases to manage those datasets, has hindered the development of foundation models. In this work, we present seven novel datasets categorized by size into three distinct categories: ToyMix, LargeMix and UltraLarge. These datasets push the boundaries in both the scale and the diversity of supervised labels for molecular learning. They cover nearly 100 million molecules and over 3000 sparsely defined tasks, totaling more than 13 billion individual labels of both quantum and biological nature. In comparison, our datasets contain 300 times more data points than the widely used OGB-LSC PCQM4Mv2 dataset, and 13 times more than the quantum-only QM1B dataset. In addition, to support the development of foundational models based on our proposed datasets, we present the Graphium graph machine learning library which simplifies the process of building and training molecular machine learning models for multi-task and multi-level molecular datasets. Finally, we present a range of baseline results as a starting point of multi-task and multi-level training on these datasets. Empirically, we observe that performance on low-resource biological datasets show improvement by also training on large amounts of quantum data. This indicates that there may be potential in multi-task and multi-level training of a foundation model and fine-tuning it to resource-constrained downstream tasks.

연구 동기 및 목표

• 분자 기계학습에서 기초 모델을 훈련하기 위해 필수적인 대규모, 레이블이 부여된 다중 작업 분자 데이터셋의 부족 문제를 해결하기 위해.
• 양자역학 및 실험실 생물학적 분석에서 유래한 감독형 다중 모odal 레이블을 도입하여 자기지도 학습 사전 훈련의 한계를 극복하기 위해.
• Graphium 딥러닝 라이브러리의 개발을 통해 거대하고 이질적인 분자 데이터셋에 대한 효율적인 훈련을 가능하게 하기 위해.
• 다양한 분자 성질에 대한 다중 작업 및 다중 수준 사전 훈련이 자원이 부족한 후행 작업에서의 성능 향상에 기여한다는 강력한 기초 모델과 경험적 증거를 확립하기 위해.

제안 방법

• DFT(B3LYP 등) 및 반경험적(예: PM6) 방법을 통해 계산된 양자역학적(QM) 성질을 통합하여 기존 분자 데이터셋을 정제하고 강화하였다.
• 고속 스크리닝 분석에서 수집한 생물학적 활성 레이블, 포함하여 농도-반응 프ofile, 유전자 발현, 독성 데이터를 활용해 다중 수준(노드 수준 및 그래프 수준)의 레이블을 생성하였다.
• 소규모에서 PubChem의 거의 완전한 커버리지까지 포함하는 세 가지 데이터셋 유형인 ToyMix, LargeMix, UltraLarge를 설계하여 총 130.4억 개의 레이블을 확보하였다.
• 다중 작업 및 다중 수준 그래프 학습에 최적화된 PyTorch 기반 라이브러리인 Graphium을 개발하였으며, 혼합 정밀도 훈련, 모델 파ipelinig, 분산 추론을 지원한다.
• 메시지 전달 신경망과 트랜스포머를 사용하여 기초 모델을 구현하고, 전체 데이터셋 계층에서 훈련하여 전이 학습 성능을 평가하였다.
• 정보량을 극대화하고 효과적인 사전 훈련을 가능하게 하기 위해, 양자 및 생물학적 작업에 대해 회귀 및 분류 목표를 혼합하여 사용하였다.

실험 결과

연구 질문

• RQ1 양자 및 생물학적 레이블을 모두 포함한 대규모, 다중 작업, 다중 수준 분자 데이터셋이 분자 기계학습에서 기초 모델의 효과적인 사전 훈련을 가능하게 할 수 있는가?
• RQ2 자기지도 학습 또는 단일 작업 사전 훈련과 비교해 볼 때, 다양한 감독 데이터를 기반으로 한 사전 훈련이 자원이 부족한 생물학적 성질 예측 작업에서 성능을 향상시키는가?
• RQ3 다중 작업 및 다중 수준 훈련 목표가 분자 모델링 작업 전반에 걸쳐 일반화 및 전이 가능성에 얼마나 기여하는가?
• RQ4 OGB-LSC 및 QM1B와 같은 기존 벤치마크와 비교해 볼 때, 제안된 데이터셋의 데이터 양과 레이블의 풍부함은 어떤가?
• RQ5 GPT-2와 같은 기초 NLP 모델에서 사용된 사전 훈련 데이터 규모에 근접한 130억 개의 레이블을 포함하는 제안된 데이터셋은 기초 모델의 사전 훈련 용량을 충분히 제공하는가?

주요 결과

• 제안된 데이터셋은 널리 사용되는 OGB-LSC PCQM4Mv2 데이터셋보다 300배 이상 많은 데이터 포인트를 포함하며, 양자 전용 QM1B 데이터셋보다는 13배 많다.
• 이 데이터셋은 거의 1억 개의 분자를 포함하고 있으며, 3,000개의 희박하게 정의된 작업을 통해 총 130억 개 이상의 개별 레이블(양자 및 생물학적 성질 모두 포함)을 포함한다.
• 기초 결과에 따르면, 사전 훈련 시 대량의 양자 데이터를 포함할 경우, 자원이 부족한 생물학적 데이터셋에 대한 미세조정 성능이 크게 향상되며, 이는 강력한 전이 학습 잠재력이 있음을 시사한다.
• Graphium 라이브러리는 대규모 다중 작업 데이터셋에 대한 효율적인 훈련을 가능하게 하며, 다양한 가속기에서 혼합 정밀도 및 분산 훈련을 지원한다.
• 제안된 데이터셋의 레이블 수(130.4억 개)는 GPT-2와 같은 기초 NLP 모델에서 사용된 사전 훈련 데이터 규모에 가까워, 분자 표현 학습을 위한 유사한 사전 훈련 능력을 갖춘 것으로 보인다.
• 양자 및 생물학적 성질의 동시 모델링은 모델의 일반화 능력을 향상시키며, 다양한 감독 사전 훈련이 효과적인 분자 기초 모델을 구축하는 데 핵심적이라는 가설을 뒷받침한다.