[논문 리뷰] Supervised Pretraining for Molecular Force Fields and Properties Prediction

Xiang Gao, Weihao Gao외 4명|arXiv (Cornell University)|2022. 11. 23.

Computational Drug Discovery Methods인용 수 5

한 줄 요약

이 논문은 3D 기하학적 구조와 원자 전하를 입력으로, 분자 에너지를 타깃 레이블로 사용하여 8600만 개의 분자에서 그래프 신경망의 지도 학습 전훈을 제안한다. 전이 학습을 통해 분자 성질 및 힘장 예측 작업에서의 후행 성능을 크게 향상시키며, 선형 프로빙은 특히 구조 관련 작업에 대해 명시적 지도 학습 없이도 풍부한 구조적 표현을 학습하고 있음을 보여준다.

ABSTRACT

Machine learning approaches have become popular for molecular modeling tasks, including molecular force fields and properties prediction. Traditional supervised learning methods suffer from scarcity of labeled data for particular tasks, motivating the use of large-scale dataset for other relevant tasks. We propose to pretrain neural networks on a dataset of 86 millions of molecules with atom charges and 3D geometries as inputs and molecular energies as labels. Experiments show that, compared to training from scratch, fine-tuning the pretrained model can significantly improve the performance for seven molecular property prediction tasks and two force field tasks. We also demonstrate that the learned representations from the pretrained model contain adequate information about molecular structures, by showing that linear probing of the representations can predict many molecular information including atom types, interatomic distances, class of molecular scaffolds, and existence of molecular fragments. Our results show that supervised pretraining is a promising research direction in molecular modeling

연구 동기 및 목표

• 전훈을 위한 대규모 레이블링 데이터셋을 활용하여 분자 기계 학습에서의 데이터 부족 문제를 해결한다.
• 전훈 단계에서의 전이 학습을 통해 분자 성질 및 힘장 예측 작업의 일반화 성능을 향상시킨다.
• 분자 에너지 기반 지도 학습이 다양한 후행 작업에 유용한 구조 정보를 암묵적으로 학습할 수 있는지 조사한다.
• 물리적으로 타당하고 정확한 타깃(분자 에너지)을 기반으로 한 전훈이 무작위 또는 자기지도 학습 전훈보다 후행 성능을 더 잘 향상시킨다는 것을 입증한다.

제안 방법

• 3D 좌표와 원자 전하를 입력으로, 분자 에너지를 지도 타깃으로 사용하는 8600만 개의 분자 데이터셋에서 메시지 전파 그래프 신경망(EGNN)을 전훈한다.
• 표준 EGNN 아키텍처를 사용하여 주목성 기반 메시지 전파를 통해 분자 그래프를 인코딩하고, 최종 리더아웃 레이어를 통해 에너지를 예측한다.
• 전훈 중에 에너지의 원자 좌표에 대한 기울기의 L2 노름을 최소화하는 힘 정규화 손실을 적용하여, 매끄러운 힘 예측을 유도한다.
• 제한된 레이블 데이터를 가진 후행 분자 성질 및 힘장 예측 작업에서 전훈된 모델을 미세조정한다.
• 학습된 표현의 능력을 평가하기 위해 선형 프로빙을 수행하여 원자 유형, 원자 간 거리, 분자 스케일드, 기능 그룹 예측 능력을 평가한다.
• 모든 선형 프로빙 작업에 음의 로그 우도 손실을 사용하여 학습된 표현의 질을 평가한다.

실험 결과

연구 질문

• RQ1 분자 에너지 기반 지도 학습 전훈이 분자 성질 및 힘장 예측 작업의 후행 성능을 향상시킬 수 있는가?
• RQ2 물리적으로 타당하고 대규모의 타깃(분자 에너지)을 기반으로 한 전훈이 자기지도 학습 또는 작업과 무관한 전훈보다 더 나은 일반화 성능을 낼 수 있는가?
• RQ3 에너지 전훈에서 학습된 표현이 원자 유형, 원자 간 거리, 기능 그룹과 같은 분자 구조 정보를 어느 정도 암묵적으로 포함하고 있는가?
• RQ4 전훈된 모델이 알려지지 않은 분자 스케일드에 대해 제로샷 예측을 일반화할 수 있는가?

주요 결과

• 전훈된 모델을 미세조정하면, 무작위 초기화로 학습하는 것보다 일곱 개인자 성질 예측 작업과 두 개인자 힘장 예측 작업에서 성능이 크게 향상된다.
• 선형 프로빙 결과, 전훈된 모델의 표현은 랜덤 가중치를 가진 모델보다 원자 유형, 원자 간 거리, 분자 스케일드, 85개의 기능 그룹을 훨씬 낮은 오차로 예측할 수 있다.
• 깊은 층을 거쳐도 원자 유형 정보를 유지하는 것으로 보아, 효과적인 메시지 전파와 표현 학습이 이루어지고 있음을 시사한다.
• 힘 정규화 기법이 힘장 예측에 대한 일반화 성능을 향상시켜, 정확한 기울기를 요구하는 작업에 더 적합한 모델을 만들었다.
• 전훈된 모델은 여러 후행 작업에서 새로운 최고 성능을 기록하였으며, 알려지지 않은 분자 스케일드에 대한 제로샷 일반화 테스트에서도 성공적으로 작동했다.
• 결과적으로, 자기지도 학습 전훈의 대안으로 분자 에너지 기반 지도 학습 전훈이 유망하며, 더 나은 후행 성능을 얻고 암묵적으로 구조 정보를 캡처한다는 점을 입증한다.