[논문 리뷰] May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular Conformations

Rui Feng, Qi Zhu외 5명|arXiv (Cornell University)|2023. 08. 24.

Machine Learning in Materials Science인용 수 4

한 줄 요약

이 논문은 3D 분자 구형에 대한 力 중심의 사전 훈련 프레임워크인 ET-OREO를 제안한다. 이는 원자적 힘을 직접 학습함으로써 평형 및 비평형 데이터를 통합한다. 평형 데이터에 대해선 노이즈 제거 및 제로-힘 정규화를, 비평형 데이터에 대해선 직접적인 힘 감독을 사용하여 힘을 통합된 목표로 삼는다. 이로 인해 사전 훈련되지 않은 Equivariant Transformer보다 힘 정확도가 약 3배 높고, 분자 동역학 시뮬레이션에서 NequIP보다 2.45배 빠른 추론 속도를 달성하면서 QM9 성능 예측에서 최신 기준 수준의 성능을 유지한다.

ABSTRACT

Recent works have shown the promise of learning pre-trained models for 3D molecular representation. However, existing pre-training models focus predominantly on equilibrium data and largely overlook off-equilibrium conformations. It is challenging to extend these methods to off-equilibrium data because their training objective relies on assumptions of conformations being the local energy minima. We address this gap by proposing a force-centric pretraining model for 3D molecular conformations covering both equilibrium and off-equilibrium data. For off-equilibrium data, our model learns directly from their atomic forces. For equilibrium data, we introduce zero-force regularization and forced-based denoising techniques to approximate near-equilibrium forces. We obtain a unified pre-trained model for 3D molecular representation with over 15 million diverse conformations. Experiments show that, with our pre-training objective, we increase forces accuracy by around 3 times compared to the un-pre-trained Equivariant Transformer model. By incorporating regularizations on equilibrium data, we solved the problem of unstable MD simulations in vanilla Equivariant Transformers, achieving state-of-the-art simulation performance with 2.45 times faster inference time than NequIP. As a powerful molecular encoder, our pre-trained model achieves on-par performance with state-of-the-art property prediction tasks.

연구 동기 및 목표

• 기존 3D 분자 사전 훈련 모델이 평형 구형에만 집중하고 비평형 데이터에선 실패하는 한계를 해결하기 위해.
• 단일 물리적 기반 목표를 사용해 평형 및 비평형 분자 구형 간의 사전 훈련을 통합하기 위해.
• 다양한 구형 데이터에서 직접 원자적 힘을 학습함으로써 분자 동역학 시뮬레이션의 정확도와 안정성을 향상시키기 위해.
• 하류 작업(예: 성능 예측 및 시뮬레이션)을 위한 3D 분자 표현의 일반화 및 이식 가능성 향상을 위해.
• 소규모, 저비용 폴리머에서 효과적인 사전 훈련을 가능하게 하여 데이터 생성 비용을 절감하기 위해.

제안 방법

• 원자적 힘을 주요 신호로 삼는 力 중심의 사전 훈련 목표를 제안하며, 이는 힘의 물리적 불변성과 분자 간 이식 가능성에 기반한다.
• 비평형 구형에 대해서는 DFT와 같은 백그라운드 방법으로 확보한 원자적 힘을 직접적으로 모델이 예측하도록 감독한다.
• 평형 구형에 대해서는 제로-힘 정규화와 힘 기반 노이즈 제거를 적용하여 근접 평형 힘을 근사하고 훈련을 안정화시킨다.
• 여러 출처에서 유래한 약 1,500만 개의 구형을 포함하는 통합 데이터셋을 사용해 Equivariant Transformer 모델을 훈련한다.
• 사전 훈련된 모델을 하류 작업(분자 동역학 시뮬레이션 및 QM9에서의 양자 성능 예측 포함)의 강력한 인코더로 활용한다.
• 에너지에 의존하지 않고 힘의 일관성에 기반한 노이즈 제거 목표를 평형 데이터에 도입하여 강인성을 향상시킨다.

(b) Potential energy curve during simulations.

실험 결과

연구 질문

• RQ1 원자적 힘을 기반으로 한 통합된 사전 훈련 목표가 평형 및 비평형 3D 분자 구형을 효과적으로 표현할 수 있는가?
• RQ2 에너지 기반 또는 노이즈 제거 기반 방법과 비교해, 힘 중심의 사전 훈련은 분자 동역학 시뮬레이션의 정확도와 안정성에 어떻게 기여하는가?
• RQ3 비평형 구형을 통합함으로써 3D 분자 표현의 일반화 능력이 성능 예측 작업에서 얼마나 향상되는가?
• RQ4 저비용 DFT 데이터로 훈련된 모델이 소규모 분자에서 출발해 큰, 새로운 폴리머로의 외삽에 성공할 수 있는가?
• RQ5 데이터 효율성과 이식 성능 측면에서, 힘 기반 목표는 전통적인 에너지 기반 또는 노이즈 제거 기반 목표와 어떻게 비교되는가?

주요 결과

• ET-OREO는 사전 훈련되지 않은 Equivariant Transformer 모델보다 약 3배 높은 힘 정확도를 달성한다.
• ET-OREO는 NequIP보다 2.45배 빠른 추론 속도를 보이며 안정적인 분자 동역학 시뮬레이션을 가능하게 하며, DFT 기준 힘과 높은 상관관계를 유지한다.
• ET-OREO는 최대 360개 원자를 가진 큰, 새로운 폴리머를 성공적으로 시뮬레이션하여 강력한 일반화 능력과 구형 샘플링 능력을 입증한다.
• ET-OREO는 QM9 HOMO-LUMO 성능 예측에서 최신 기준 수준의 성능을 달성했으며, NoisyNode와 동등하며, 처음부터 훈련된 모델들을 능가한다.
• rich한 비평형 데이터에 접근함에도 불구하고 평형 성능 예측에서 유의미한 정확도 향상이 관찰되지 않아, 이러한 작업은 비평형 데이터 포함으로부터 유의미한 이점을 얻지 못할 수 있음을 시사한다.
• 제로-힘 정규화와 힘 기반 노이즈 제거가 평형 데이터 훈련에서 효과적으로 안정성을 확보하여 기존 Equivariant Transformer에서 관찰된 불안정성 문제를 해결한다.