[논문 리뷰] ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations

Weihua Hu, Muhammed Shuaibi외 6명|arXiv (Cornell University)|2021. 03. 02.

Machine Learning in Materials Science참고 문헌 53인용 수 24

한 줄 요약

ForceNet은 OC20에서 물리 기반 GNNs보다 원자 힘 예측에서 우수한 성능을 보이고 학습 및 추론 속도에서 상당한 속도 향상을 제공하는 확장 가능한 힘 중심 그래프 신경망이며; ForceNet-large는 계산 비용이 크게 감소한 상태에서 DimeNet++와 동등하거나 이를 능가합니다.

ABSTRACT

With massive amounts of atomic simulation data available, there is a huge opportunity to develop fast and accurate machine learning models to approximate expensive physics-based calculations. The key quantity to estimate is atomic forces, where the state-of-the-art Graph Neural Networks (GNNs) explicitly enforce basic physical constraints such as rotation-covariance. However, to strictly satisfy the physical constraints, existing models have to make tradeoffs between computational efficiency and model expressiveness. Here we explore an alternative approach. By not imposing explicit physical constraints, we can flexibly design expressive models while maintaining their computational efficiency. Physical constraints are implicitly imposed by training the models using physics-based data augmentation. To evaluate the approach, we carefully design a scalable and expressive GNN model, ForceNet, and apply it to OC20 (Chanussot et al., 2020), an unprecedentedly-large dataset of quantum physics calculations. Our proposed ForceNet is able to predict atomic forces more accurately than state-of-the-art physics-based GNNs while being faster both in training and inference. Overall, our promising and counter-intuitive results open up an exciting avenue for future research.

연구 동기 및 목표

• 비용이 큰 양자 물리 계산에서 원자 힘에 대한 빠르고 정확한 근사의 필요성에 대한 동기 부여.
• 아키텍처에서 명시적 물리 제약에 의존하지 않는 유연하고 표현력 있는 힘 중심 GNN 개발.
• 회전 관련 물리 속성을 암묵적으로 강제하기 위해 물리 기반 데이터 증강 활용.
• OC20의 200M+ 구조 데이터세트를 포함한 규모 확장성 및 효율성 조사, 비평형 구조 포함.

제안 방법

• 원자들을 노드로, 엣지들이 3D 상호작용을 조건부 필터 컨볼루션으로 포착하는 인코더–디코더 GNN 아키텍처로 ForceNet 개발.
• 전체 3D 상대 배치와 회전 공변 방향 정보를 인코딩하는 엣지 피처 e_st 사용.
• 원시 엣지 피처를 더 풍부한 표현으로 변환하기 위한 기저 함수 B 도입(구면 조화 함수 버전 포함) 및 엣지/노드 MLP에서의 표현력 있는 비선형성(Swish) 사용.
• 학습 중 회전 데이터 증강을 적용하여 명시적 아키텍처 회전 제약을 강제하지 않고 회전 공변 예측을 유도.
• 에너지 보존 제약을 명시적으로 적용하지 않고 확장성 확보를 위해 최종 노드 임베딩에서 디코더 MLP를 통해 힘 예측을 직접 구성.
• OC20 S2F 태스크에서 SchNet, DimeNet++, DimeNet++-large, GNS와의 비교를 통해 학습/추론 예산을 신중하게 제어.
• 성능에 대한 영향 연구를 위해 항등, 선형+활성화, 가우시안, 삼각함수, 구면 조화 함수의 기저 함수 범위와 Swish 활성화의 효과를 평가.

실험 결과

연구 질문

• RQ1 명시적 회전 제약 없이도 힘 중심 GNN이 비균형 대규모 데이터셋에서 물리 기반 GNN을 능가하는 최첨단 성능을 달성할 수 있는가?
• RQ2 아키텍처 선택(기저 함수, 활성화, 엣지 컨디셔닝) 및 데이터 증강이 힘 예측 정확도와 회전 공변성에 어떻게 영향을 미치는가?
• RQ3 ForceNet의 학습/추론 효율성과 기존 모델 간의 OC20 확장 시 trade-off는 무엇인가?
• RQ4 회전 증강이 비 회전 불변 모델에 어느 정도로 회전 공변 특성을 부여할 수 있는가?

주요 결과

• ForceNet은 유사한 계산 자원에서 OC20의 원자당 힘 예측에서 물리 기반 GNN을 능가하며 학습과 추론이 더 빠릅니다.
• ForceNet-large은 DimeNet++-large와 동등하거나 더 나은 힘 예측을 달성하는 반면 훨씬 적은 계산으로 동작합니다(학습 약 8.2x–18.9x 감소, 추론 약 2.4x–3.0x 빨라짐, 비교에 따라 다름).
• ForceNet–large는 평균 힘 MAE에서 DimeNet++를 능가하고 학습 및 추론이 더 크게 빠르며, 초기 결과는 SchNet 대비 유사한 성능에 도달하는 데 약 6.3x의 계산량 감소, DimeNet++ 대비 최대 약 6.5x 빠른 추론을 시사합니다.
• 회전 증강과 대규모 데이터는 회전하에서의 불안정성을 줄이는 데 효과적이며 ForceNet의 예측을 회전 공변 특성에 더 가깝게 만듭니다.
• 소실 연구에서 표현력 있는 엣지 수준 계산과 Swish 활성화를 포함한 구면 기저의 사용이 ForceNet의 성능에 결정적이며, 이를 제거하면 정확도가 저하됩니다.