[논문 리뷰] Accurate transition state generation with an object-aware equivariant elementary reaction diffusion model

Chenru Duan, Yuanqi Du외 2명|arXiv (Cornell University)|2023. 04. 12.

Machine Learning in Materials Science인용 수 4

한 줄 요약

이 논문은 반응물과 생성물 구조를 활용하여 초당 3D 전이 상태(TS) 구조를 정확하게 생성하는 새로운 객체 인식 기반 SE(3)-동치성 확산 모델인 OA-ReactDiff를 소개한다. 이 모델은 진짜 TS 구조에서 중앙값 RMSD가 0.08 Å이며, 신뢰도 기반 추천 시스템을 통해 가장 도전적인 반응의 14%만 최적화함으로써 장벽 예측 정확도를 2.6 kcal/mol 이내로 유지한다.

ABSTRACT

Code and model release for OA-ReactDiff, a generalizied Denoising Diffusion Probabilistic Model (DDPM) for generating sets of 3D structures (i.e., reactant, transition state, and product) in chemical reactions. OA-ReactDiff reduces the transition state search from hours/days to seconds and complemented conventional intuition-based reaction exploration with generative AI solutions that can discover surprising "unintended" reactions. A codebase under continuous development can be found at https://github.com/chenruduan/OAReactDiff. Update: Upload the checkpoint for the confidence model.

연구 동기 및 목표

• 기존에 수시간에 걸쳐 DFT 계산이 필요한 반응 메커니즘 규명에서의 전이 상태(TS) 탐색 계산적 병목 현상을 해결하기 위해.
• 반응물, 전이 상태, 생성물과 같은 다객체 시스템에서 순열, 이동, 회전 대칭성을 모두 유지하는 기계학습 모델을 개발하기 위해.
• 비용이 많이 드는 DFT 기반 TS 최적화에 대한 의존도를 줄이기 위해 고정확도의 TS 구조를 신속하게 생성하고, 불확실성 정량화를 통해 가장 신뢰할 수 있는 예측을 식별하기 위해.
• 인간 또는 계산 자원의 간섭을 최소화하면서도 대규모의 미지의 메커니즘을 가진 반응 네트워크를 스케일러블하게 구축하기 위해.

제안 방법

• 모델는 반응물, TS, 생성물 조각 간의 순열, 이동, 회전 대칭성을 유지하는 객체 인식 기반 SE(3) 동치성 그래프 신경망(GNN)을 사용한다.
• 2차 다항식 노이즈 스케줄과 L_simple 손실 함수를 사용하는 확산 프로세스를 적용하며, 9,000개의 반응으로 구성된 Transition1x15 데이터셋에서 훈련한다.
• 원자 순서 매핑이나 조각 정렬이 필요 없이 단일 순방향 전파에서 반응물, TS, 생성물의 공동 3D 구조를 생성한다.
• 생성된 TS 구조를 '좋음' 또는 '나쁨'으로 레이블링하기 위해 RMSD 임계값(<0.2 Å)을 사용하는 이진 분류기로 구성된 신뢰도 스코어링 모델을 훈련하여 불확실성 인식 선택을 가능하게 한다.
• 최종 예측 파이프라인은 신뢰도 스코어를 사용해 다수의 생성 결과 중 가장 신뢰도 높은 TS 후보를 순위 매기고 선택한다.
• V100 GPU를 사용하여 5,000개의 확산 스텝, 96개의 반경 기저 함수, 196개의 은닉 채널을 사용하며, 화학적으로 다양한 대규모 데이터셋에서 엔드 투 엔드로 훈련된다.

실험 결과

연구 질문

• RQ1 객체 인식 기반 SE(3) 동치성 확산 모델은 반복적인 DFT 최적화 없이도 초등 반응에 대해 정확한 3D 전이 상태 구조를 생성할 수 있는가?
• RQ2 이러한 모델이 다수의 분자 조각 간에 순열, 이동, 회전 대칭성을 얼마나 잘 유지할 수 있는가?
• RQ3 신뢰도 기반 순위 매기기 시스템은 정확도를 유지하면서도 필요한 DFT 최적화 수를 줄일 수 있는가?
• RQ4 정확도, 속도, 계산 비용 측면에서 기존의 TS 탐색 방법과 비교해 모델의 성능은 어떠한가?

주요 결과

• 모델는 진짜 TS에서 중앙값 RMSD가 0.08 Å로 매우 높은 구조적 정확도를 달성한다.
• 신뢰도 기반 샘플링을 통해 평균 RMSD는 0.13 Å로 감소하여 불확실성 정량화를 통한 신뢰성 향상을 입증한다.
• 단일 GPU에서 반응당 6초 미만으로 TS 구조를 예측하여 기존 DFT 기반 방법 대비 100배 빠른 속도를 기록한다.
• 신뢰도 모델을 통해 DFT 최적화를 안내함으로써, 장벽 예측 정확도 2.6 kcal/mol을 확보하기 위해 가장 도전적인 반응의 14%만 전체 DFT 평가를 거치면 된다.
• 검증용 테스트 세트(1,073개의 반응, TS 구조에 겹침 없음)에서도 강한 성능을 보여, 새로운 반응에 대한 일반화 능력이 뛰어나다.
• 신뢰도 모델은 신뢰할 수 있는 불확실성 정량화를 달성하였으며, 예측 확률과 실제 RMSD 성능 간 강한 상관관계를 보였다.