[논문 리뷰] FAENet: Frame Averaging Equivariant GNN for Materials Modeling

Alexandre Duval, Victor Schmidt외 5명|arXiv (Cornell University)|2023. 04. 28.

Machine Learning in Materials Science인용 수 7

한 줄 요약

FAENet는 기하 데이터의 확률적 기준 프레임 평균화(SFA)를 통해 구조적 제약 없이 E(3)-동변성을 달성하는 새로운 경량 그래프 신경망을 제안한다. 이 방법은 OC20 및 QM9 데이터셋에서 최신 기술 수준의 정확도와 계산 효율성을 확보하며, 이전의 동변성 모델보다 확장성과 대칭성 유지 측면에서 뛰어나다.

ABSTRACT

Applications of machine learning techniques for materials modeling typically involve functions known to be equivariant or invariant to specific symmetries. While graph neural networks (GNNs) have proven successful in such tasks, they enforce symmetries via the model architecture, which often reduces their expressivity, scalability and comprehensibility. In this paper, we introduce (1) a flexible framework relying on stochastic frame-averaging (SFA) to make any model E(3)-equivariant or invariant through data transformations. (2) FAENet: a simple, fast and expressive GNN, optimized for SFA, that processes geometric information without any symmetrypreserving design constraints. We prove the validity of our method theoretically and empirically demonstrate its superior accuracy and computational scalability in materials modeling on the OC20 dataset (S2EF, IS2RE) as well as common molecular modeling tasks (QM9, QM7-X). A package implementation is available at https://faenet.readthedocs.io.

연구 동기 및 목표

• 복잡한 아키텍처 설계를 통해 대칭성을 강제하는 기존의 동변성 GNN이 겪는 표현력과 확장성의 한계를 해결하기 위해.
• 모델 아키텍처가 아닌 데이터 변환을 통해 E(3) 대칭성을 유지하는 유연한 프레임워크를 개발하기 위해.
• 기하 정보를 자유롭게 처리하면서도 이론적·실증적 대칭성 불변성을 유지하는 경량이면서 고성능의 GNN—FAENet를 설계하기 위해.
• SFA 기반의 대칭성 강제 방식이 데이터 증강 및 대칭성 강제 없음과 비교하여 재료 성능 예측에서의 모델 정확도와 대칭성 유지 측면에서 뛰어나다는 것을 실증적으로 검증하기 위해.

제안 방법

• E(3)-동변성을 강제하기 위해 원자 구성을 표준 기준 프레임으로 투영하는 대칭성 유지 데이터 증강 기법으로 확률적 기준 프레임 평균화(SFA)를 제안한다.
• 대칭성 유지 아키텍처 제약 없이 상대 원자 위치를 처리하는 GNN 아키텍처인 FAENet를 도입한다. 이는 SFA를 통해 불변성을 확보한다.
• 기준 프레임 평균화를 통해 입력 그래프를 표준 좌표계로 변환하여, 임의의 모델이 회전, 반사, 이동에 대해 동변성이 되도록 한다.
• 일반화 및 대칭성에 대한 강건성을 향상시키기 위해 SFA에서 다중 기준 프레임을 확률적으로 샘플링하는 미분 가능 기준 프레임 선택 메커니즘을 도입한다.
• 상대 위치를 에지 특성으로 사용하는 기하 메시지 전파 기법을 활용하여 표현력을 극대화하면서도, 데이터 사전 처리를 통해 대칭성을 유지한다.
• 다양한 데이터셋을 대상으로 회전 불변성, 반사 불변성, 힘의 동변성 등의 지표를 사용하여 대칭성 유지 정도를 측정함으로써 방법의 타당성을 검증한다.

Figure 1: Overview of FAENet architecture making use of the Stochastic Frame Averaging framework. FAENet takes an input sample $D$ and makes an invariant $\Phi^{\text{{invar}}}_{{\mathcal{F}}}$ (e.g. energy) and/or equivariant prediction $\Phi^{\text{{equiv}}}_{{\mathcal{F}}}$ (e.g. atomic forces) b

실험 결과

연구 질문

• RQ1 대칭성 유지 아키텍처 설계 없이 데이터 변환을 통해 GNN에서 E(3)-동변성을 효과적으로 강제할 수 있는가?
• RQ2 SFA 기반의 대칭성 강제 방식이 데이터 증강 및 대칭성 강제 없음과 비교하여 모델 정확도와 대칭성 유지 측면에서 어떻게 다른가?
• RQ3 FAENet와 같은 경량 GNN이 재료 성능 예측 과제에서 복잡한 대칭성 제약 아키텍처보다 더 높은 정확도와 확장성을 달성할 수 있는가?
• RQ4 전체 기준 프레임 평균화 또는 표준 데이터 증강과 비교해 SFA는 대칭성 일반화를 얼마나 향상시키는가?
• RQ5 FAENet는 고도로 수정된 아키텍처 없이도 고체 결정 및 분자와 같은 다양한 재료 데이터셋에 효과적으로 일반화될 수 있는가?

주요 결과

• S2EF 데이터셋에서 FAENet와 SFA를 사용할 경우 테스트 MAE가 464 meV에 도달하여 SchNet(918 meV) 및 ForceNet(655 meV)를 능가하며, 뛰어난 정확도와 확장성을 입증한다.
• SFA는 S2EF에서 회전 불변성 오차를 6.94 meV로 줄였고, 반사 불변성 오차는 6.96 meV로 낮췄다. 이는 데이터 증강(8.88 및 9.12 meV) 및 No-FA(26.81 및 27.61 meV)보다 유의미하게 뛰어나다.
• IS2RE에서 FAENet와 SFA를 사용할 경우 MAE가 569 meV를 기록하며, ForceNet 대비 12% 향상, SchNet 대비 13% 향상되었고, 강력한 대칭성 유지 성능도 유지했다.
• SE(3)-SFA 변종은 전체 FA보다 略적으로 더 나은 지표를 기록하여, 적은 수의 기준 프레임에서 학습함으로써 일반화 능력 향상이 가능하며 대칭성 손실 없이도 가능함을 시사한다.
• No-FA 및 No-Train는 대칭성 지표가 유의미하게 열등하다(예: 51.1 meV의 회전 불변성 오차), 이는 대칭성 강제가 견고한 성능을 위해 필수적임을 입증한다.
• SFA는 특히 3D 대칭성 하에서 데이터 증강보다 대칭성 유지 및 예측 정확도에서 뛰어나며, S2EF에서 회전 동변성 오차가 60% 감소함을 보였다.

Figure 2: OC-20 Direct IS2RE performance / scalability trade-off. FAENet achieves the best MAE across methods while being much faster than the majority of baselines.