[논문 리뷰] TF-GNN: Graph Neural Networks in TensorFlow

Oleksandr Ferludin, Arno Eigenwillig외 25명|arXiv (Cornell University)|2022. 07. 07.

Advanced Graph Neural Networks인용 수 15

한 줄 요약

TF-GNN는 이질적이고 다중관계적 그래프 데이터를 네 가지 수준의 API 계층을 통해 지원하는 확장 가능한 그래프 신경망(GNN)을 위한 오픈소스 텐서플로 라이브러리이다. 연구자와 개발자가 원시 텐서플로 연산에서 고수준 Keras 모델에 이르기까지 다양한 추상화 수준으로 GNN을 구축, 훈련, 배포할 수 있도록 하며, OGBN-MAG에서 최신 기술 성능을 달성한다. 더 작은 최적화된 MPNN 모델이 더 크고 복잡한 아키텍처를 능가한다.

ABSTRACT

TensorFlow-GNN (TF-GNN) is a scalable library for Graph Neural Networks in TensorFlow. It is designed from the bottom up to support the kinds of rich heterogeneous graph data that occurs in today's information ecosystems. In addition to enabling machine learning researchers and advanced developers, TF-GNN offers low-code solutions to empower the broader developer community in graph learning. Many production models at Google use TF-GNN, and it has been recently released as an open source project. In this paper we describe the TF-GNN data model, its Keras message passing API, and relevant capabilities such as graph sampling and distributed training.

연구 동기 및 목표

• 복잡한 이질적 관계 데이터에서 그래프 신경망을 훈련하기 위한 확장 가능하고 프로덕션 준비가 된 프레임워크의 부족을 해결하기 위해.
• 초보자부터 전문 연구자에 이르기까지 다양한 수준의 기술 능력을 가진 개발자를 위한 계층적 API를 제공하기 위해.
• 분산 훈련, TPU 가속, 다양한 모odal의 사전 학습 모델에의 접근을 가능하게 하기 위해 텐서플로 생태계와 원활하게 통합하기 위해.
• 실제 응용 프로그램에서 GNN의 효율적인 모델 개발, 하이퍼파rameter 튜닝 및 배포를 가능하게 하기 위해.
• 간단하고 잘 튜닝된 GNN 모델이 더 복잡한 아키텍처보다 벤치마크 그래프 학습 작업에서 뛰어난 성능을 낼 수 있음을 입증하기 위해.

제안 방법

• TF-GNN는 네 단계의 API를 구현한다: 이질적 그래프 표현을 위한 데이터 레벨, GNN 레이어를 구축하기 위한 Keras 메시지 전파 API, 메시지 전파 동작을 정의하기 위한 GraphUpdate, 그리고 훈련과 추론을 위한 오케스트레이터.
• 엣지 중심, 노드 중심, 그래프 중심의 GNN 빌딩 블록을 지원하여 복잡한 관계 구조를 민첩하게 모델링할 수 있다.
• 노드 및 엣지 특징이 그래프 이웃 영역을 통해 적용되는 학습 가능한 함수를 통해 업데이트되는 메시지 전파 프레임워크를 사용한다.
• 텐서플로 생태계와 통합되어 분산 훈련, TPU 가속, 추론을 위한 SavedModel 내보내기를 지원한다.
• 빠른 프로토타ип링을 위한 사전 빌드된 구성 요소인 VanillaMPNN, GraphSAGE, GATv2, MultiHeadAttention 등의 모델 라이브러리가 제공된다.
• 자동 최적화를 위한 Vizier 통합을 통해 모델 용량, 드롭아웃, 정규화, 메시지 차원 등의 하이퍼파라미터 튜닝을 지원한다.

실험 결과

연구 질문

• RQ1 텐서플로 위에 구축된 통합적이고 확장 가능한 GNN 프레임워크를 통해 이질적이고 다중관계적 그래프 데이터를 지원할 수 있는가?
• RQ2 계층적 API 설계가 다양한 수준의 머신러닝 전문 지식을 가진 사용자를 효과적으로 지원할 수 있는가?
• RQ3 간단하고 잘 최적화된 MPNN 모델이 대규모 그래프 벤치마크에서 더 복잡한 트랜스포머 기반 아키텍처를 능가할 수 있는가?
• RQ4 텐서플로 생태계와의 통합이 GNN의 훈련 효율성과 프로덕션 준비도를 어느 정도 향상시킬 수 있는가?
• RQ5 자동 하이퍼파라미터 튜닝이 아키텍처의 복잡성 없이도 높은 성능을 내는 GNN 구성 요소를 효과적으로 발견할 수 있는가?

주요 결과

• Vizier를 통한 하이퍼파라미터 최적화로 훈련된 589만 파라미터의 MPNN 모델이 OGBN-MAG에서 검증 정확도 0.5149를 달성했으며, 이는 훨씬 더 큰 2,680만 파라미터의 이질적 그래프 트랜스포머(HGT) 모델을 능가했다.
• Vizier 최적화된 MPNN 모델은 테스트 정확도 0.5027을 기록하여 HGT의 테스트 정확도 0.4982를 초월했으며, 하이퍼파라미터 튜닝이 아키텍처의 복잡성보다 더 효과적임을 입증했다.
• TF-GNN 프레임워크는 과학자들이 여러 팀에서 GNN 모델의 프로덕션 배포를 가능하게 하여 확장성과 사용성의 타당성을 입증했다.
• 엣지 중심 및 그래프 중심 메시지 전파를 지원함으로써, 노드 중심 프레임워크에서 다루기 어려운 모델(예: 그래프 트랜스포머)을 자연스럽게 표현할 수 있었다.
• 텐서플로 생태계와의 통합을 통해 TPU에서의 효율적 훈련과 추론을 위한 SavedModel로의 원활한 내보내기가 가능해져 배포의 부담을 줄였다.
• 모델 라이브러리와 고수준 API는 개발 시간을 크게 단축시켜 GNN의 신속한 실험과 배포를 가능하게 했다.