[논문 리뷰] Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

Rui Jiao, Jiaqi Han외 3명|arXiv (Cornell University)|2022. 07. 18.

Machine Learning in Materials Science인용 수 5

한 줄 요약

이 논문은 3D 분자의 기하학적 정보를 효과적으로 활용하기 위해 E(3)-등변 메시지 전파와 물리학에 영감을 받은 사전 훈련 작업(예: 리만-가우시안 노이즈 제거를 통한 등가 힘 예측 및 노이즈 스케일 예측)을 활용하는 자기지도 학습 사전 훈련 프레임워크인 3D-EMGP를 제안한다. 이 방법은 MD17 및 QM9 벤치마크에서 최신 기술을 초월하며, 3D 좌표 증강이 적용된 2D 사전 훈련 베이스라인조차도 앞서는 성능을 기록한다.

ABSTRACT

Pretraining molecular representation models without labels is fundamental to various applications. Conventional methods mainly process 2D molecular graphs and focus solely on 2D tasks, making their pretrained models incapable of characterizing 3D geometry and thus defective for downstream 3D tasks. In this work, we tackle 3D molecular pretraining in a complete and novel sense. In particular, we first propose to adopt an equivariant energy-based model as the backbone for pretraining, which enjoys the merits of fulfilling the symmetry of 3D space. Then we develop a node-level pretraining loss for force prediction, where we further exploit the Riemann-Gaussian distribution to ensure the loss to be E(3)-invariant, enabling more robustness. Moreover, a graph-level noise scale prediction task is also leveraged to further promote the eventual performance. We evaluate our model pretrained from a large-scale 3D dataset GEOM-QM9 on two challenging 3D benchmarks: MD17 and QM9. Experimental results demonstrate the efficacy of our method against current state-of-the-art pretraining approaches, and verify the validity of our design for each proposed component.

연구 동기 및 목표

• 분자의 동역학 및 에너지 예측과 같은 작업에서 필수적인 3D 기하학적 정보를 포착하지 못하는 기존 2D 분자 사전 훈련 방법의 한계를 해결하기 위해.
• E(3)-등변 모델과 3D 인지 사전 훈련 목표를 활용하여 3D 분자 기하학을 완전히 활용하는 자기지도 학습 사전 훈련 프레임워크를 개발하기 위해.
• 사전 훈련 손실에서 대칭성 불변성(E(3)-불변성)을 보장하여 3D 공간 변환에 대한 강건성과 일반화 능력을 향상시키기 위해.
• 다양한 3D 백본과 벤치마크를 대상으로 추론 및 일반화 연구를 통해 제안된 구성 요소의 효과를 평가하기 위해.

제안 방법

• 3D 회전, 이동, 반사에 대해 불변성을 확보하기 위해 등변 GNN(e.g., EGNN) 기반의 E(3)-등변 메시지 전파 백본을 사용하여 원자 수준의 힘을 예측한다.
• 리만-가우시안 분포를 사용하여 노이즈 제거 손실로 재구성된 원자 수준의 사전 훈련 작업인 힘 예측 기반 작업을 도입하여 E(3)-불변성을 유지한다.
• 정규화된 3D 분자 구조와 노이즈가 첨가된 쌍에 대해 적용된 노이즈 스케일을 예측하는 그래프 수준의 대체 작업을 설계하여 강건성과 일반화 능력을 향상시킨다.
• 두 사전 훈련 목표를 통합하여 3D-EMGP라는 단일 프레임워크를 구축하고, 3D 구조를 포함한 대규모 GEOM-QM9 데이터셋에서 훈련한다.
• 에너지 기반 모델링 접근 방식을 사용: 풀링된 노드 표현에서 에너지를 예측하고, 힘을 에너지의 기울기로 계산하여 전역 일관성을 향상시킨다.
• 에너지 투영 헤드를 추가하여 에너지 분포도의 시각화 및 평가를 위해 사전 훈련된 모델을 다운스트림 3D 작업에 피니팅하여 적용한다.

실험 결과

연구 질문

• RQ1 3D 분자 그래프를 위한 자기지도 학습 사전 훈련 프레임워크는 기존 2D 사전 훈련 방법보다 3D 벤치마크 작업에서 더 우수한 성능을 내는가?
• RQ2 특히 리만-가우시안 노이즈 제거를 통해 사전 훈련 손실에 E(3)-불변성을 강제 적용할 경우, 모델의 강건성과 성능에 어떤 영향을 미치는가?
• RQ3 원자 수준의 힘 예측과 그래프 수준의 노이즈 스케일 예측을 결합한 사전 훈련 작업이 최종 성능에 기여하는 정도는 어떠한가?
• RQ4 에너지 기반 모델링 전략(예측된 에너지의 기울기로 힘을 계산하는 방식)은 직접적인 힘 예측보다 더 높은 성능을 내는가?
• RQ5 제안된 사전 훈련 프레임워크는 다양한 3D GNN 백본에 대해 얼마나 일반화 가능한가?

주요 결과

• 3D-EMGP는 MD17 힘 예측 벤치마크에서 테스트 MAE 0.0876을 기록하여 다음으로 우수한 베이스라인(0.0968)을 크게 앞서며, 3D 좌표 증강이 적용된 2D 사전 훈련 방법조차도 뛰어넘는 성능을 보였다.
• 아블레이션 연구에서 리만-가우시안 노이즈 제거 기법을 제거할 경우 MAE가 0.0912로 증가하여, E(3)-불변성과 강건성을 유지하는 데서 이 기법의 핵심적 역할을 입증했다.
• 그래프 수준의 노이즈 스케일 예측 작업은 의미 있는 기여를 하며, 이를 제거할 경우 MAE가 0.0876에서 0.0905로 증가하여 일반화 능력 향상에 기여하는 것으로 나타났다.
• 에너지 기반 힘 예측 전략(예측된 에너지의 기울기로 힘 계산)은 직접 힘 예측(0.0914)보다 낮은 MAE(0.0876)를 기록하여 전역 기하 패턴을 더 잘 포착하는 것으로 확인되었다.
• 제안된 방법은 다양한 백본에 대해 잘 일반화되며, SchNet 및 TorchMD-ET에 3D-EMGP를 적용했을 때도 일관되게 성능 향상이 관찰되어 광범위한 적용 가능성을 입증했다.
• 에너지 분포도의 시각화 결과, 3D-EMGP는 구조 주변에 매끄럽고 안정된 에너지 표면을 생성하는 반면, 2D 사전 훈련 방법은 거친 에너지 분포를 보이며, 물리적 에너지 분포와의 일치도가 더 높다는 것을 시사했다.