[논문 리뷰] On the Scalability of GNNs for Molecular Graphs

Maciej Sypetkowski, Frederik Wenkel외 5명|arXiv (Cornell University)|2024. 04. 17.

Machine Learning in Materials Science인용 수 6

한 줄 요약

이 논문은 서로 다른 GNN 아키텍처가 큰 분자 그래프에서 어떻게 확장되는지 연구하고, 더 큰 너비(width), 깊이(depth), 데이터 규모, 다양화된 사전학습이 지속적으로 이득을 준다는 것을 보여주며, MolGPS를 그래프 기반의 기초 모델로 소개하여 많은 작업에서 SOTA를 달성합니다.

ABSTRACT

Scaling deep learning models has been at the heart of recent revolutions in language modelling and image generation. Practitioners have observed a strong relationship between model size, dataset size, and performance. However, structure-based architectures such as Graph Neural Networks (GNNs) are yet to show the benefits of scale mainly due to the lower efficiency of sparse operations, large data requirements, and lack of clarity about the effectiveness of various architectures. We address this drawback of GNNs by studying their scaling behavior. Specifically, we analyze message-passing networks, graph Transformers, and hybrid architectures on the largest public collection of 2D molecular graphs. For the first time, we observe that GNNs benefit tremendously from the increasing scale of depth, width, number of molecules, number of labels, and the diversity in the pretraining datasets. We further demonstrate strong finetuning scaling behavior on 38 highly competitive downstream tasks, outclassing previous large models. This gives rise to MolGPS, a new graph foundation model that allows to navigate the chemical space, outperforming the previous state-of-the-arts on 26 out the 38 downstream tasks. We hope that our work paves the way for an era where foundational GNNs drive pharmaceutical drug discovery.

연구 동기 및 목표

• 폭(width), 깊이(depth), 데이터셋 크기, 라벨 수, 그리고 데이터셋 다양성이 분자 그래프에서 GNN 성능에 어떤 영향을 미치는지 평가합니다.
• 사전학습, 미세조정, 프루빙 설정에서 세 가지 아키텍처(MPNN++, Graph Transformer, GPS++)를 비교합니다.
• 분자용 그래프 기초 모델을 구축하기 위한 효과적인 사전학습 전략과 아키텍처 선택을 결정합니다.
• MolGPS 스타일의 분자 발견을 위한 기초 모델에 대해 실용적인 지침을 제공합니다.

제안 방법

• Graphium 라이브러리를 사용하여 세 가지 아키텍처(MPNN++, Graph Transformer, GPS++)의 확장 동작을 연구합니다.
• 최대 5백만 분자와 수 천 개의 라벨을 가진 대규모 다태스크 데이터셋 LargeMix에서 사전학습합니다.
• TDC, Polaris, MoleculeNet의 38개 다운스트림 작업에서 미세조정 및 프루빙을 평가합니다.
• 그래프 트랜스포머를 강화하기 위해 위치 인코딩(PSE)과 구조적 바이어스를 사용합니다.
• 폭(width), 깊이(depth), 데이터셋 크기, 라벨에 대한 확장을 분석하고 관측된 추세를 거듭제곱 법칙 확장 관계 L_theta ~ (D_C / D)^beta 및 L_theta ~ (theta_C / theta)^alpha에 적합시킵니다.
• L1000 없이 사전학습한 후 MPNN++, Transformer, GPS++의 지문을 결합하여 MolGPS를 제안하고, 그 다음 MLP 헤드를 사용합니다.

실험 결과

연구 질문

• RQ1 분자 GNN이 너비(width), 깊이(depth), 데이터셋 크기, 라벨 다양성이 증가함에 따라 어떻게 확장되는가?
• RQ2 어떤 아키텍처(MPNN++, Graph Transformer, GPS++)가 분자 특성 예측 및 다운스트림 작업에 대해 가장 잘 확장되는가?
• RQ3 38개의 다운스트림 분자 벤치마크에서 성능을 최대화하는 사전학습 및 미세조정/프루빙 전략은 무엇인가?
• RQ4 아키텍처 선택과 사전학습을 하나로 통합하여 작업 전반에서 최첨단 결과를 달성하는 기초 모델(MolGPS)을 만들 수 있는가?

주요 결과

• 모든 모델은 너비(width), 깊이(depth), 데이터셋 크기 증가에 따라 확장성이 좋다.
• MPNN++는 작은/데이터가 부족한 작업에서 매개변수 효율성과 데이터 효율성이 더 높고, Transformer는 데이터가 더 필요하지만 너비의 이점으로 확장 가능하다.
• 하이브리드 GPS++는 저매개변수 영역에서 MPNN++의 이점을 활용하고, 더 큰 데이터 규모에서 Transformer와 유사하게 확장된다.
• 데이터셋 크기(분자 수)를 증가시키면 상당한 성능 향상이 나타나며, GPS++와 Transformer는 양자 작업에서 더 큰 이득을 얻는다.
• 프루빙은 Transformer가 다운스트림 작업에서 너비의 이점을 누리는 반면, 깊이 확장은 혼합된 경향을 보이며, 라벨 확장이 성능을 강하게 증가시킨다.
• MolGPS(MPNN++, Transformer, GPS++의 지문을 결합) 제안은 38개 다운스트림 작업 중 26개에서 SOTA를 달성하고 분자 발견을 위한 통합 그래프 기초 모델을 향상시킨다.
• L1000 작업 없이의 사전학습이 규모에 걸쳐 다운스트림 성능을 개선하고, 다중 지문 프루빙이 결과를 더욱 향상시킨다.