[논문 리뷰] GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction

Dominic Masters, Josef Dean외 8명|arXiv (Cornell University)|2022. 11. 18.

Machine Learning in Materials Science인용 수 10

한 줄 요약

GPS++는 PCQM4Mv2 데이터셋에서 분자의 성질 예측을 위해 최적화된 하이브리드 MPNN/Transformer 모델로, 3D 원자 좌표와 노이즈 제거 손실을 통합하여 성능을 향상시켰다. 112개 모델 앙상블과 Graphcore IPU 가속을 사용해 1시간 32분 내에 추론을 완료하고, 테스트 세트의 평균 절대 오차(MAE)가 0.0719로 기록되어 OGB-LSC 2022 챌린지에서 1위를 달성했다.

ABSTRACT

This technical report presents GPS++, the first-place solution to the Open Graph Benchmark Large-Scale Challenge (OGB-LSC 2022) for the PCQM4Mv2 molecular property prediction task. Our approach implements several key principles from the prior literature. At its core our GPS++ method is a hybrid MPNN/Transformer model that incorporates 3D atom positions and an auxiliary denoising task. The effectiveness of GPS++ is demonstrated by achieving 0.0719 mean absolute error on the independent test-challenge PCQM4Mv2 split. Thanks to Graphcore IPU acceleration, GPS++ scales to deep architectures (16 layers), training at 3 minutes per epoch, and large ensemble (112 models), completing the final predictions in 1 hour 32 minutes, well under the 4 hour inference budget allocated. Our implementation is publicly available at: https://github.com/graphcore/ogb-lsc-pcqm4mv2.

연구 동기 및 목표

• 대규모 PCQM4Mv2 데이터셋에서 분자 성질 예측을 위한 고도로 효율적이고 정확한 모델을 개발하기 위해.
• 3D 원자 좌표와 인덕티브 바이어스를 활용하여 그래프 신경망의 일반화 능력을 향상시키고 과도한 스무딩을 줄이기 위해.
• 하드웨어 가속을 통한 훈련 및 추론을 활용해 엄격한 추론 시간 제약(4시간 예산) 내에서 최신 기술 성능을 달성하기 위해.
• 순수한 Transformer 모델에 비해 훨씬 적은 파라미터를 사용하면서도 하이브리드 MPNN/Transformer 아키텍처가 성능을 동등하거나 초월할 수 있음을 보여주기 위해.
• 테스트-챌린지 분할에서 예측의 강건성과 일반화 능력을 극대화하기 위해 112개의 모델 앙상블을 구축하기 위해.

제안 방법

• GPS++는 국소적인 인덕티브 바이어스와 전역 정보 흐름을 균형 있게 유지하기 위해 깊은 메시지 전달 신경망(MPNN)과 편향된 자기주의 기반 메커니즘을 조합한다.
• 거리 특징을 주의 메커니즘에 통합하기 위해 거리 기반의 그룹화된 입력 마스킹 전략을 통해 3D 원자 좌표를 모델에 통합한다.
• 깊은 아키텍처에서의 과도한 스무딩과 표현 학습 향상을 방지하기 위해 사전 훈련 단계에서 노이즈 제거 손실을 적용한다.
• MPNN 레이어는 국소적 메시지 전달을 담당하고, 주의 레이어는 장거리 의존성을 가능하게 하며, 주의는 이웃 원자들 쪽으로 편향된다.
• 훈련 및 추론은 Graphcore의 IPU 하드웨어를 통해 가속되며, 1초당 100개의 에포크를 처리하고, 112개 모델 앙상블의 전체 추론을 1시간 32분 내에 완료한다.
• 입력 특징, 드롭아웃 비율, 마스킹 전략을 다양하게 조정하여 112개의 다양한 모델 앙상블을 구성함으로써 일반화 및 강건성을 향상시킨다.

실험 결과

연구 질문

• RQ1 하이브리드 MPNN/Transformer 아키텍처가 더 적은 파라미터를 사용하면서도 순수한 Transformer 모델을 능가하거나 능가할 수 있는가?
• RQ2 그룹화된 입력 마스킹을 통한 3D 원자 좌표 통합이 모델 성능 향상과 일반화 능력 향상에 얼마나 효과적인가?
• RQ3 노이즈 제거 사전 훈련 목표가 깊은 그래프 신경망에서 분자의 그래프에 대해 과도한 스무딩을 얼마나 줄일 수 있는가?
• RQ4 IPU를 활용한 하드웨어 가속이 깊은 모델과 대규모 앙상블에 대해 엄격한 시간 제약 내에서 대규모 훈련 및 추론을 가능하게 하는가?
• RQ5 초기 설정을 다양하게 조정함으로써 달성된 모델 앙상블의 다양성이 PCQM4Mv2 테스트-챌린지 분할에서 최종 예측 성능에 크게 기여하는가?

주요 결과

• GPS++는 PCQM4Mv2 데이터셋에서 테스트-챌린지 평균 절대 오차(MAE) 0.0719를 기록하여 OGB-LSC 2022 챌린지에서 상위 3위를 차지했다.
• 최종 112개 모델 앙상블은 단일 IPU와 AMD EPYC CPU를 사용해 1시간 32분 내에 추론을 완료했으며, 4시간 예산 이내에 성공적으로 수행되었다.
• 모델는 파라미터 수가 64%에 불과한 4430만 개(6900만 개 대비)로, 최신 기술 성능을 보이는 순수한 Transformer(Luo 등, 2022)와 유사한 성능을 달성했다.
• 그룹화된 마스킹을 통한 3D 거리 특징 통합은 주의 메커니즘의 효과를 크게 향상시켰으며, 아블레이션 실험에서 3D 특징 제거 시 성능 향상이 거의 없음을 확인했다.
• 입력 특징, 드롭아웃, 마스킹 전략을 다양하게 조정한 다양한 구성의 앙상블은 단일 모델의 MAE 0.0755에서 프록시 세트 기준 0.0722로 개선되었고, 최종 테스트 세트에서는 0.0719로 추가로 향상되었다.
• IPU 하드웨어에서 3분당 1 에포크의 속도로 훈련이 가능하여 깊이 있는(16층) 아키텍처의 빠른 반복 및 24시간 이내 훈련이 가능했다.