[논문 리뷰] Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching

Shengchao Liu, Hongyu Guo외 1명|arXiv (Cornell University)|2022. 06. 27.

Machine Learning in Materials ScienceMaterials Science인용 수 17

한 줄 요약

이 논문은 좌표 노이즈 제거를 통한 잠재 에너지 표면 모델링을 통해 강력한 분자 표현을 학습하기 위해 SE(3)-불변 노이즈 제거 거리 매칭을 활용하는 새로운 3D 분자 기하학 사전학습 프레임워크인 GeoSSL-DDM을 제안한다. 스코어 매칭을 통해 좌표 노이즈 제거를 쌍별 거리 노이즈 제거로 변환함으로써, 22개의 하류 기하 분자 예측 작업에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Molecular representation pretraining is critical in various applications for drug and material discovery due to the limited number of labeled molecules, and most existing work focuses on pretraining on 2D molecular graphs. However, the power of pretraining on 3D geometric structures has been less explored. This is owing to the difficulty of finding a sufficient proxy task that can empower the pretraining to effectively extract essential features from the geometric structures. Motivated by the dynamic nature of 3D molecules, where the continuous motion of a molecule in the 3D Euclidean space forms a smooth potential energy surface, we propose GeoSSL, a 3D coordinate denoising pretraining framework to model such an energy landscape. Further by leveraging an SE(3)-invariant score matching method, we propose GeoSSL-DDM in which the coordinate denoising proxy task is effectively boiled down to denoising the pairwise atomic distances in a molecule. Our comprehensive experiments confirm the effectiveness and robustness of our proposed method.

연구 동기 및 목표

• 정확한 분자 성질 예측에 필수적인 3D 분자 기하학 데이터에 대한 효과적인 사전학습 방법의 부족을 해결하기 위해.
• 잠재 에너지 표면(PES)에서 분자의 동적이고 노이즈가 많은 성질을 반영하는 자기지도 학습 대체 과제를 개발하기 위해.
• 직접 3D 좌표를 노이즈 제거하는 데 어려움이 있는 문제를 해결하기 위해, SE(3)-불변 스코어 매칭을 활용해 쌍별 원자 거리의 노이즈 제거로 과제를 전환하기 위해.
• 3D 공간에서의 회전 및 이동 변형에 대해 불변인 강력한 분자 표현 학습을 가능하게 하기 위해.
• 다양한 3D 분자 예측 벤치마크에서 제안된 방법의 우수성을 실증적으로 검증하기 위해.

제안 방법

• 청결한 분자 기하학과 그 노이즈가 가미된 변형 간의 상호정보를 최대화하는 일반적인 기하 자기지도 학습 프레임워크인 GeoSSL을 제안한다.
• 좌표 노이즈 제거를 쌍별 원자 거리 노이즈 제거로 재구성함으로써 회전 및 이동 불변성을 보장하는 SE(3)-불변 스코어 매칭 방법인 GeoSSL-DDM을 도입한다.
• NCSN에 영감을 얻어 다중 레이어 노이즈 제거 스코어 매칭 목표를 사용하여 다중 단계의 노이즈 제거를 가능하게 한다.
• 진짜 데이터 분포를 근사하고 청결한 기하학과 노이즈가 가미된 기하학 간의 상호정보를 최대화하기 위해 에너지 기반 모델(EBM)을 사용한다.
• 그래프 신경망(GNNs)을 활용해 분자 구조를 인코딩하고 노이즈가 제거된 거리를 예측하며, 스코어 매칭을 통해 손실을 최적화한다.
• 최적화를 가능하게 하기 위해 상호정보의 하한을 적용하여, 노이즈가 가미된 입력에서 청결한 기하학을 재구성하도록 모델을 효과적으로 훈련시킨다.

실험 결과

연구 질문

• RQ1 3D 분자 기하학에 기반한 자기지도 학습 프레임워크가 기존 2D 그래프 기반 방법보다 하류 분자 예측 과제에서 더 우수한 성능을 낼 수 있는가?
• RQ2 직접 3D 좌표 노이즈 제거보다 쌍별 원자 거리 노이즈 제거가 3D 분자 표현 학습에 더 효과적이고 실현 가능할 수 있는가?
• RQ3 노이즈 제거 레이어의 수가 3D 분자 사전학습에서 모델 성능과 내구성에 어떤 영향을 미치는가?
• RQ4 SE(3)-불변 스코어 매칭이 분자의 매끄러운 잠재 에너지 표면을 효과적으로 포착하면서도 회전 및 이동에 대해 불변성을 유지할 수 있는가?
• RQ5 제안된 방법이 힘 예측 및 에너지 회귀와 같은 다양한 분자 데이터셋과 하류 과제에 일반화 가능한가?

주요 결과

• GeoSSL-DDM은 22개의 하류 기하 분자 예측 과제에서 최신 기술 수준의 성능을 달성하며, 2D 그래프 기반 및 다른 3D 방법을 포함한 9개의 사전학습 베이스라인을 모두 능가한다.
• MD17 데이터셋에서 GeoSSL-DDM은 힘 예측 과제에서 평균 절대 오차 1.637 ± 0.01을 기록하여 비교된 모든 방법 중에서 가장 낮은 수준을 기록했다.
• 일부 힘 예측 과제에서 지도 학습 기반 베이스라인 대비 평균 절대 오차를 최대 20%까지 감소시켜, 자기지도 사전학습에서의 강력한 일반화 능력을 입증했다.
• 실험 결과, 노이즈 제거 레이어 수를 늘릴수록 성능 향상이 이루어지며(특히 L > 1일 경우), 다단계 노이즈 제거가 단일 레이어 접근보다 유의미한 이점을 제공함을 확인했다.
• GeoSSL-DDM은 평가된 모든 과제에서 동일한 작업 [76]을 능가했으며, 특히 더 깊은 노이즈 제거 네트워크를 사용할 경우 그 우수성이 두드러졌다. 이는 다중 레이어 스코어 매칭의 장점을 보여준다.
• 제거 실험 결과, 좌표 직접 노이즈 제거보다 쌍별 거리 노이즈 제거가 더 효과적이며, SE(3) 불변성이 강력한 내구성과 성능 향상에 기여한다는 점을 확인했다.