[논문 리뷰] Unified 2D and 3D Pre-Training of Molecular Representations

J. Zhu, Yingce Xia외 6명|arXiv (Cornell University)|2022. 07. 14.

Machine Learning in Materials Science인용 수 5

한 줄 요약

이 논문은 분자의 구조 표현 학습을 위한 통합된 2D 및 3D 사전 훈련 프레임워크를 제안하며, 그래프 신경망을 통해 원자 좌표와 원자 간 거리를 동시에 인코딩합니다. 이 모델은 11개의 분자 성질 예측 작업 중 10개에서 최신 기술 수준(SOTA) 성능을 달성하였고, 3D 구형 구조 생성 성능도 크게 향상시켰으며, 2D 전용 작업에서는 평균 8.3%의 정확도 향상을 보였고, 구형 구조 생성 벤치마크에서도 뛰어난 성능을 보였습니다.

ABSTRACT

Molecular representation learning has attracted much attention recently. A molecule can be viewed as a 2D graph with nodes/atoms connected by edges/bonds, and can also be represented by a 3D conformation with 3-dimensional coordinates of all atoms. We note that most previous work handles 2D and 3D information separately, while jointly leveraging these two sources may foster a more informative representation. In this work, we explore this appealing idea and propose a new representation learning method based on a unified 2D and 3D pre-training. Atom coordinates and interatomic distances are encoded and then fused with atomic representations through graph neural networks. The model is pre-trained on three tasks: reconstruction of masked atoms and coordinates, 3D conformation generation conditioned on 2D graph, and 2D graph generation conditioned on 3D conformation. We evaluate our method on 11 downstream molecular property prediction tasks: 7 with 2D information only and 4 with both 2D and 3D information. Our method achieves state-of-the-art results on 10 tasks, and the average improvement on 2D-only tasks is 8.3%. Our method also achieves significant improvement on two 3D conformation generation tasks.

연구 동기 및 목표

• 기존 방법들이 2D 분자 그래프와 3D 구형 구조를 별개로 다루는 데서 비롯되는 한계를 해결하기 위해, 이 둘은 상호 보완적인 성질을 지닌다.
• 단일 모델 내에서 2D 및 3D 분자 표현을 동시에 학습할 수 있는 통합된 딥 러닝 프레임워크를 개발하기 위해.
• 다양한 자기 지도 학습 작업을 기반으로 사전 훈련하여 분자 성질 예측 및 3D 구형 구조 생성 성능을 향상시키기 위해.
• 모델의 강건성과 일반화 능력을 향상시키기 위해 순열 불변성과 회전-이동 불변성을 훈련 목표에 통합하기 위해.

제안 방법

• 모델는 원자 표현과 함께 원자 좌표와 원자 간 거리를 피드포워드 네트워크를 통해 인코딩한 후, 그래프 신경망을 통해 융합합니다.
• 세 가지 사전 훈련 작업을 적용합니다: 마스킹된 원자 및 좌표 복원, 2D 그래프에서의 3D 구형 구조 생성, 3D 구조에서의 2D 그래프 생성.
• 대칭적인 부분 구조를 다루기 위해 순열 불변 손실을 사용하여 원자 교환 시 훈련 안정성에 영향을 주지 않도록 합니다.
• 3D 구형 구조 생성에 대해 회전-이동 불변 손실을 적용하여 공간 변환에 대해 동등성(equivariance)을 보장합니다.
• PCQM4Mv2 데이터셋을 기반으로 사전 훈련을 수행하며, 이 데이터셋은 분자에 대해 2D 그래프와 3D 구조를 모두 제공합니다.
• 사전 훈련된 모델은 분자 성질 예측 및 구형 구조 생성과 같은 후행 작업에서 전이 학습을 통해 미세 조정됩니다.

실험 결과

연구 질문

• RQ1 2D 및 3D 분자 표현의 공동 사전 훈련이 후행 분자 성질 예측 성능을 향상시킬 수 있는가?
• RQ2 손실 함수에 순열 불변성을 통합할 경우, 대칭적인 분자 부분 구조에서의 모델 강건성과 성능에 어떤 영향을 미치는가?
• RQ3 2D 및 3D를 통합한 사전 훈련이 2D 또는 3D 단독 지도 학습 방법에 비해 3D 구형 구조 생성 성능에 얼마나 기여하는가?
• RQ4 각 사전 훈련 목표(복원, 생성)가 최종 모델 성능에 기여하는 정도는 어느 정도인가?
• RQ5 사이클 일致성과 마스킹된 원자 복원을 활용하여 3D 구조가 없는 분자에 대해서도 통합 모델이 일반화 가능한가?

주요 결과

• 제안된 방법은 11개의 분자 성질 예측 작업 중 10개에서 최신 기술 수준 성능을 달성하였으며, 2D 전용 작업에서는 평균 8.3%의 향상을 보였습니다.
• GEOM-QM9 및 GEOM-Drugs 데이터셋에서, 구형 구조 생성 성능이 모든 베이스라인을 초월하여 최고의 커버리지(COV) 및 매칭(MAT) 점수를 기록했습니다.
• 절단 실험을 통해 세 가지 사전 훈련 목표를 동시에 사용할 경우 최고의 성능을 달성함을 확인하였으며, 단일 목표 훈련 대비 MAT 점수는 최대 15% 향상되었습니다.
• 순열 불변 손실은 대칭적인 부분 구조에서 성능 향상을 크게 개선하였으며, 절단 실험에서 예측 오차를 평균 20% 감소시켰습니다.
• 초기 좌표를 무작위로 설정하고 마스킹 복원 및 사이클 일치성을 활용함으로써 3D 구조가 없는 분자에 대해서도 잘 일반화됨을 확인했습니다.
• 통합된 사전 훈련 프레임워크는 3D 구형 구조 생성에 대해 강력한 제로샷 전이 성능를 보이며, 모델의 강건성과 일반화 능력을 입증했습니다.