[논문 리뷰] Molecular Geometry-aware Transformer for accurate 3D Atomic System modeling

Zheng Yuan, Yaoyun Zhang외 4명|arXiv (Cornell University)|2023. 02. 02.

Machine Learning in Materials ScienceMaterials Science인용 수 3

한 줄 요약

Moleformer는 상대적 거리와 각도를 기반으로 회전 및 이동 불변성 공간 인코딩을 사용하여 원자, 결합, 비공유 원자 쌍을 기하학적으로 인식하는 트랜스포머 아키텍처를 제안한다. 이는 OC20에서 이완 에너지 예측에서 최고 성능을 기록하고 QM9에서 기존 GNN 및 트랜스포머 모델을 능가하는 경쟁력 있는 성능을 달성하며, 고차 기하 상호작용을 명시적으로 모델링함으로써 달성된다.

ABSTRACT

Molecular dynamic simulations are important in computational physics, chemistry, material, and biology. Machine learning-based methods have shown strong abilities in predicting molecular energy and properties and are much faster than DFT calculations. Molecular energy is at least related to atoms, bonds, bond angles, torsion angles, and nonbonding atom pairs. Previous Transformer models only use atoms as inputs which lack explicit modeling of the aforementioned factors. To alleviate this limitation, we propose Moleformer, a novel Transformer architecture that takes nodes (atoms) and edges (bonds and nonbonding atom pairs) as inputs and models the interactions among them using rotational and translational invariant geometry-aware spatial encoding. Proposed spatial encoding calculates relative position information including distances and angles among nodes and edges. We benchmark Moleformer on OC20 and QM9 datasets, and our model achieves state-of-the-art on the initial state to relaxed energy prediction of OC20 and is very competitive in QM9 on predicting quantum chemical properties compared to other Transformer and Graph Neural Network (GNN) methods which proves the effectiveness of the proposed geometry-aware spatial encoding in Moleformer.

연구 동기 및 목표

• 기존 트랜스포머 모델이 원자만 입력으로 사용하고 결합, 각도, 비공유 상호작용을 명시적으로 모델링하지 못하는 한계를 해결하기 위해.
• 결합 각도와 토판 각도와 같은 기하학적 요소를 주의 메커니즘에 직접 통합하여 분자의 에너지 및 성능 예측 정확도를 향상시키기 위해.
• 주기적 경계 조건 보정 및 데이터 증강을 도입하여 분포 외 분자 시스템에서의 일반화 능력을 향상시키기 위해.
• 결합 길이 예측을 위한 엣지 수준의 보조 손실을 도입하여 기하학적 감독을 통해 에너지 예측 성능 향상을 입증하기 위해.

제안 방법

• Moleformer는 이전 트랜스포머의 표준 원자만의 입력을 초월하여 원자와 엣지(결합 및 비공유 원자 쌍)를 입력으로 사용한다.
• 노드와 엣지 간의 상대적 거리와 각도를 계산하는 기하학적 인식 공간 인코딩을 도입하여 회전 및 이동 불변성을 확보한다.
• 각 레이어에서 노드와 엣지 모두에 대해 자기주의를 적용하여 세 개 이상의 몸체 상호작용(예: 결합 각도, 토판 각도)을 명시적으로 모델링할 수 있도록 한다.
• 주기적 구조를 가진 시스템, 특히 촉매 시뮬레이션에서 정확도를 향상시키기 위해 새로운 주기적 경계 조건(PBC) 보정을 적용한다.
• OC20 데이터셋의 분포 외 광범위한 흡착체에서 일반화 능력을 향상시키기 위해 노이즈가 섞인 노드 데이터 증강을 사용한다.
• 예측된 결합 길이를 보조적으로 감독하기 위해 엣지 수준의 손실을 도입하여 에너지 추정 성능을 향상시킨다.

Figure 1: The overview of Moleformer. (a) A toy example of a catalyst system where \ch CH3 is the adsorbate and the surface is composed of copper atoms. Moleformer prefers edges between adosrbate-adsorbate atoms and adsorbate-surface atoms as model inputs. (b) Architecture of Moleformer: Nodes and s

실험 결과

연구 질문

• RQ1 트랜스포머 아키텍처에서 결합과 비공유 원자 쌍을 명시적으로 모델링하면 원자만을 고려하는 모델에 비해 3D 원자 시스템의 에너지 예측 성능이 향상되는가?
• RQ2 상대적 거리와 각도를 기반으로 한 기하학적 인식 공간 인코딩은 분자의 에너지 및 성능 예측 작업에 어떤 영향을 미치는가?
• RQ3 주기적 경계 조건 보정을 통합할 경우 주기적 구조를 가진 촉매 시스템에서 에너지 예측 성능이 얼마나 향상되는가?
• RQ4 결합 길이 예측을 위한 엣지 수준의 보조 손실을 추가하면 시스템 에너지 정확도 예측 능력이 향상되는가?
• RQ5 특히 OC20 데이터셋의 희귀 흡착체를 고려할 때 Moleformer는 분포 외 분자에 대해 어떻게 일반화되는가?

주요 결과

• Moleformer는 OC20 데이터셋에서 초기 상태에서 이완 상태로의 에너지 예측에서 최고 성능을 기록하였으며, 도메인 내 분할에서 에너지 MAE가 413 meV로 나타났다.
• 주기적 경계 조건 보정이 포함된 경우 IS2RS에서 위치 MAE가 감소하고, IS2RE 분할에서 에너지 MAE는 0.014 eV 감소하고 EWT는 0.7% 향상되었다.
• 노이즈가 섞인 노드 데이터 증강을 적용하지 않은 경우 분포 외 흡착체 분할에서 에너지 MAE가 크게 증가하여 데이터 증강이 일반화 능력을 향상시킨다는 것을 확인하였다.
• 엣지 수준의 보조 손실을 제거할 경우 도메인 내 및 도메인 외 분할 모두에서 에너지 MAE가 증가하여 기하학적 감독을 통한 효과성을 입증하였다.
• QM9 데이터셋에서 Moleformer는 분자의 오비탈 에너지와 같은 양자화학적 성능 예측에서 경쟁적인 결과를 도출하였으며, 최근 최고 성능 모델인 SEGNN 및 EQGAT을 능가하였다.
• 절단 실험 결과 모든 구성 요소—PBC 보정, 데이터 증강, 엣지 수준의 손실—이 성능 향상에 기여하며, 특히 엣지 수준의 손실이 가장 큰 기여를 하였다.

Figure 2: PBC Correction. (a)-(c) An example of a catalyst system where \ch CH3 is the adsorbate and the surface is composed of copper atoms viewed along the <110>orientation. The dotted lines show the borders of the unit cell. Grey-colored small atom is a hydrogen atom that goes across the unit cel