[논문 리뷰] Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations

Anuroop Sriram, Abhishek Das외 3명|arXiv (Cornell University)|2022. 03. 18.

Machine Learning in Materials Science인용 수 9

한 줄 요약

이 논문은 입력 그래프를 여러 GPU에 분산시켜 억억 개 파rameter를 가진 그래프 신경망(GNN)을 훈련할 수 있도록 하는 그래프 병렬화(Graph Parallelism) 방법을 소개한다. 이를 통해 원자 시뮬레이션을 위한 고차 상호작용(예: 원자 3중주)을 모델링할 수 있으며, Open Catalyst 2020(OC20) 데이터셋에서 최신 기준(SOTA) 성능을 달성하여 힘의 평균 절대 오차(MAE)를 15% 향상시키고 AFbT를 21% 향상시켰다.

ABSTRACT

Recent progress in Graph Neural Networks (GNNs) for modeling atomic simulations has the potential to revolutionize catalyst discovery, which is a key step in making progress towards the energy breakthroughs needed to combat climate change. However, the GNNs that have proven most effective for this task are memory intensive as they model higher-order interactions in the graphs such as those between triplets or quadruplets of atoms, making it challenging to scale these models. In this paper, we introduce Graph Parallelism, a method to distribute input graphs across multiple GPUs, enabling us to train very large GNNs with hundreds of millions or billions of parameters. We empirically evaluate our method by scaling up the number of parameters of the recently proposed DimeNet++ and GemNet models by over an order of magnitude. On the large-scale Open Catalyst 2020 (OC20) dataset, these graph-parallelized models lead to relative improvements of 1) 15% on the force MAE metric for the S2EF task and 2) 21% on the AFbT metric for the IS2RS task, establishing new state-of-the-art results.

연구 동기 및 목표

• 원자 시뮬레이션을 위한 수십억 개 파arameter를 가진 대규모 GNN을 훈련하는 데 도전하는 것, 특히 원자 3중주, 4중주 등 고차 상호작용을 모델링하는 데 초점 맞추기.
• 수백만 개의 중간 크기 그래프로 구성된 데이터셋에 적합하지 않은 기존 GNN 훈련 방법의 메모리 및 계산 제약을 극복하는 것.
• 대규모 Open Catalyst 2020(OC20) 데이터셋에서 DimeNet++ 및 GemNet-T와 같은 최신 기준 GNN의 확장 가능한 훈련을 가능하게 하는 것.
• 그래프 병렬화를 통해 모델 크기를 확장하면 원자 힘과 안정 구조 예측 성능에 상당한 향상이 이루어지는지 입증하는 것.
• 확장된 그래프 네트워크(EGN) 프레임워크 내에서 등방성 및 에너지 보존 모델을 포함한 다양한 메시지 전달 GNN 아키텍처에 일반화 가능한 프레임워크 제공

제안 방법

• 입력 그래프를 여러 GPU에 나누어 분산함으로써 단일 장치에 담을 수 없는 대규모 GNN 훈련을 가능하게 하는 그래프 병렬화 기법을 제안한다.
• 고차 상호작용(예: 결합 각도)을 명시적으로 모델링하기 위해 그래프 네트워크(GN) 프레임워크를 확장된 그래프 네트워크(EGN) 프레임워크로 확장한다.
• 고차 상호작용의 계산을 GPU 간에 분할하여 메모리 및 계산 자원을 장치 간에 효율적으로 분배한다.
• 각 GPU가 입력 그래프의 부분 집합을 처리하면서도 고차 항목에 대한 전체 그래프 연결성을 유지하는 데이터 병렬 유사 전략을 사용한다.
• GPipe와 같은 기존 모델 병렬화 기법과 통합하여 매우 큰 GNN의 가로 방향 및 세로 방향 확장 모두를 가능하게 한다.
• 그래프 분할과 모델 병렬화를 융합한 하이브리드 훈련 전략을 통해 GNN의 파arameter 수와 상호작용의 복잡도를 모두 확장한다.

실험 결과

연구 질문

• RQ1 고차 상호작용으로 인한 높은 메모리 및 계산 요구량에도 불구하고, 억억 개 파arameter를 가진 GNN을 원자 시뮬레이션에 효과적으로 훈련시킬 수 있는가?
• RQ2 수백만 개의 작은 분자 그래프로 구성된 분자 데이터셋에서 기존 분산 GNN 훈련 방법과 비교해 그래프 병렬화의 확장성 및 성능은 어떻게 되는가?
• RQ3 그래프 병렬화를 통해 훈련된 더 큰 GNN은 힘의 MAE 및 AFbT와 같은 핵심 원자 시뮬레이션 벤치마크에서 예측 정확도를 얼마나 향상시키는가?
• RQ4 그래프 병렬화는 등방성 및 에너지 보존을 모델링하는 다양한 GNN 아키텍처에 일반화될 수 있는가?
• RQ5 이러한 대규모 모델 훈련의 환경적 영향은 무엇이며, 전산 자원을 전략적으로 선택함으로써 이를 어떻게 완화할 수 있는가?

주요 결과

• 그래프 병렬화를 통해 DimeNet++ 및 GemNet-T 모델의 파arameter 수가 1억 개 이상으로 증가하여 이전 모델 대비 10배 이상 확장되었다.
• 그래프 병렬화된 GemNet-XL 모델은 S2EF 작업에서 기존 최신 기준 대비 힘의 MAE를 15% 상대적으로 향상시켰다.
• 동일한 모델은 IS2RS 작업에서 AFbT 지표에서도 21% 상대적으로 향상시켜 OC20 데이터셋에서 새로운 최신 기준을 수립했다.
• 이 방법은 확장된 그래프 네트워크(EGN) 프레임워크 내에서 등방성 및 에너지 보존 모델을 포함한 다양한 메시지 전달 GNN 아키텍처에 일반화 가능하다.
• Tesla V100 32GB GPU를 사용한 클라우드 GPU에서 GemNet-XL 모델을 훈련한 결과, 예상 이산화탄소 배출량은 3490~8052kg CO2eq.로, 로스앤젤레스에서 뉴욕으로의 왕복 항공편 16번에 해당한다.
• 저자들은 환경 영향을 줄이기 위해 탄소 상쇄 계약을 체결한 컴퓨팅 자원을 선택할 것을 권장하며, 본 연구의 훈련은 100% 상쇄되었다고 밝혔다.