[논문 리뷰] Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+

Shuqi Lu, Zhifeng Gao외 3명|arXiv (Cornell University)|2023. 03. 16.

Machine Learning in Materials Science인용 수 6

한 줄 요약

Uni-Mol+는 1D/2D 입력에서 저비용으로 3D 분자 구형을 생성한 후, 이중 트랙 트랜스포머 모델을 사용하여 DFT 최적화된 평형 구형으로 반복적으로 개선함으로써 양자화학적(QC) 성질을 예측하는 딥러닝 프레임워크를 제안한다. 이는 PCQM4MV2 및 OC20 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성하며, OC20 테스트 세트에서 0.3745 eV의 MAE와 11.29%의 EwT를 기록하여 이전 방법들을 크게 앞서며 성능을 뛰어나게 한다.

ABSTRACT

Recent developments in deep learning have made remarkable progress in speeding up the prediction of quantum chemical (QC) properties by removing the need for expensive electronic structure calculations like density functional theory. However, previous methods learned from 1D SMILES sequences or 2D molecular graphs failed to achieve high accuracy as QC properties primarily depend on the 3D equilibrium conformations optimized by electronic structure methods, far different from the sequence-type and graph-type data. In this paper, we propose a novel approach called Uni-Mol+ to tackle this challenge. Uni-Mol+ first generates a raw 3D molecule conformation from inexpensive methods such as RDKit. Then, the raw conformation is iteratively updated to its target DFT equilibrium conformation using neural networks, and the learned conformation will be used to predict the QC properties. To effectively learn this update process towards the equilibrium conformation, we introduce a two-track Transformer model backbone and train it with the QC property prediction task. We also design a novel approach to guide the model's training process. Our extensive benchmarking results demonstrate that the proposed Uni-Mol+ significantly improves the accuracy of QC property prediction in various datasets. We have made the code and model publicly available at \url{https://github.com/dptech-corp/Uni-Mol}.

연구 동기 및 목표

• 정확한 DFT 계산을 위해 필요로 하는 3D 평형 구형과 1D/2D 분자 표현 간의 불일치로 인해 발생하는 양자화학적(QC) 성질 예측의 격차를 해소하기 위해.
• SMILES나 2D 그래프에서 직접 QC 성질을 예측하는 데에 한계가 있는 기존 딥러닝 모델들이 구형 의존성을 제대로 반영하지 못해 정확도에 영향을 주는 문제를 해결하기 위해.
• 신경망을 사용하여 RDKit 등으로 생성된 원시 3D 구형을 DFT 평형 구형으로 반복적으로 개선하는 방법을 개발함으로써 고정밀도 성질 예측을 가능하게 하기 위해.
• 분포 이탈 문제를 완화하기 위해 훈련 중에 가짜 궤적 샘플링 전략을 도입하여 분포 이탈에 대한 일반화 능력과 강건성을 향상시키기 위해.
• 딥러닝을 활용한 고속 스크리닝을 위한 확장 가능한 종단 간 프레임워크를 구축하여 고비용 DFT 계산에 대한 의존도를 줄이기 위해.

제안 방법

• RDKit 등의 저비용 도구를 사용하여 1D SMILES 또는 2D 분자 그래프에서 초기 3D 분자 구형을 생성한다.
• 원자 표현을 위한 트랙과 쌍 표현을 위한 트랙을 가진 이중 트랙 트랜스포머 백본을 활용하며, 둘 다 3D 공간적 및 그래프 위치 인코딩으로 강화된다.
• 원시 구형을 DFT 최적화 평형 구형으로 향해 학습된 신경망을 사용하여 3D 좌표를 반복적으로 업데이트한다.
• 가짜 선형 궤적을 통해 RDKit로 생성된 구형과 DFT 최적화된 구형 사이의 중간 구형을 샘플링하여 보조 학습 전략을 도입함으로써 개선 과정을 감독한다.
• 쌍 표현을 개선하기 위해 외적 곱, 삼각형 업데이트 및 피드포워드 네트워크를 반복 적용하는 다중 헤드 어텐션 메커니즘을 통합한다.
• 대규모 쌍화된 원시 구형과 DFT 구형 데이터셋을 활용하여 L1 손실을 사용한 QC 성질 예측과 구형 개선을 위한 노이즈 제거 목적의 손실을 함께 사용하여 종단 간으로 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 딥러닝 모델이 원시 3D 구형을 DFT 최적화 평형 구조로 향해 개선하는 방식을 학습함으로써, 양자화학적 성질 예측에서 최신 기술 수준의 정확도를 달성할 수 있는가?
• RQ2 직접 1D/2D 입력을 모델링하는 것과 비교해, 구형 개선 과정을 반복적 업데이트의 시퀀스로 모델링하는 것이 예측 정확도를 어떻게 향상시키는가?
• RQ3 훈련 중에 가짜 궤적과 중간 구형을 샘플링하는 전략이 일반화 능력 향상과 분포 이탈 감소에 얼마나 기여하는가?
• RQ4 강화된 쌍 표현과 공간 인코딩을 갖춘 이중 트랙 트랜스포머 아키텍처가 3D 분자 기하학 작업에서 표준 GNN 또는 트랜스포머보다 우수한 성능을 내는가?
• RQ5 제안된 방법이 OC20 벤치마크와 같은 외부 분포(Out-of-Distribution, OOD) 분자 데이터셋에 효과적으로 일반화될 수 있는가?

주요 결과

• Uni-Mol+는 OC20 IS2RE 테스트 세트에서 0.3745 eV의 평균 절대 오차(MAE)를 기록하여 이전 SOTA 방법(DRFormer)의 0.3865 eV보다 유의미하게 높은 성능을 보였다.
• PCQM4MV2 벤치마크에서 Uni-Mol+는 모든 카테고리에서 가장 낮은 MAE를 기록했으며, 내부 데이터(ID)의 경우 0.3795 eV, OOD 흡착(OOD Ads.)의 경우 0.4526 eV, OOD 양면(OOD Both)의 경우 0.4021 eV를 기록했다.
• OC20 테스트 세트에서 Uni-Mol+는 11.29%의 EwT(Energy Weighted Top-1 Accuracy)를 달성하여 이전 SOTA(DRFormer, 9.18%)를 2个百分点 이상 앞서며 우수한 성능을 보였다.
• 절단 실험을 통해 이중 트랙 트랜스포머 백본과 가짜 궤적 샘플링 전략이 성능 향상에 필수적임을 확인했으며, 둘 중 하나를 제거할 경우 정확도가 크게 떨어졌다.
• 모델는 외부 분포 데이터에서 강력한 강건성을 보였으며, OC20 및 PCQM4MV2 벤치마크의 모든 OOD 카테고리에서 가장 높은 EwT 점수를 기록했다.
• 이 방법은 내부 데이터뿐 아니라 도전적인 외부 분포 설정에서도 최신 기술 수준의 성능을 달성하여 강력한 일반화 능력을 보여주었다.