[논문 리뷰] SpreadGNN: Serverless Multi-task Federated Learning for Graph Neural Networks

Chaoyang He, Emir Ceyani외 3명|arXiv (Cornell University)|2021. 06. 04.

Advanced Graph Neural Networks참고 문헌 68인용 수 29

한 줄 요약

SpreadGNN은 부분 라벨이 있는 그래프 신경망의 서버리스, 다중 작업 연합 학습을 가능하게 하며, 분산형 주기적 평균 SGD 및 작업 관계 정규화를 사용하여 비 IID 분자 그래프 데이터셋에서 중앙 집중식 FL 기준선을 능가합니다.

ABSTRACT

Graph Neural Networks (GNNs) are the first choice methods for graph machine learning problems thanks to their ability to learn state-of-the-art level representations from graph-structured data. However, centralizing a massive amount of real-world graph data for GNN training is prohibitive due to user-side privacy concerns, regulation restrictions, and commercial competition. Federated Learning is the de-facto standard for collaborative training of machine learning models over many distributed edge devices without the need for centralization. Nevertheless, training graph neural networks in a federated setting is vaguely defined and brings statistical and systems challenges. This work proposes SpreadGNN, a novel multi-task federated training framework capable of operating in the presence of partial labels and absence of a central server for the first time in the literature. SpreadGNN extends federated multi-task learning to realistic serverless settings for GNNs, and utilizes a novel optimization algorithm with a convergence guarantee, Decentralized Periodic Averaging SGD (DPA-SGD), to solve decentralized multi-task learning problems. We empirically demonstrate the efficacy of our framework on a variety of non-I.I.D. distributed graph-level molecular property prediction datasets with partial labels. Our results show that SpreadGNN outperforms GNN models trained over a central server-dependent federated learning system, even in constrained topologies. The source code is publicly available at https://github.com/FedML-AI/SpreadGNN

연구 동기 및 목표

• 분자 그래프의 중앙 데이터 공유를 저해하는 프라이버시/규제 제약 해결.
• 클라이언트 간 부분 라벨을 처리하는 GNN용 서버리스 연합 다중 작업 학습 프레임워크 개발.
• 분산 네트워크에 적합한 수렴 보장 최적화 방법 제공.
• 비 IID 부분 라벨 설정에서 서버리스 SpreadGNN가 중앙 서버의 FedAvg를 능가함을 입증.

제안 방법

• 부분 라벨이 있는 그래프 신경망에 대해 중앙화 공식(FedGMTL)을 통해 Federated Multi-Task Learning을 확장한 다음 중앙 서버를 제거하여 SpreadGNN를 얻습니다.
• 클라이언트가 로컬 SGD를 수행하고 τ번 반복마다 이웃과 동기화하는 Decentralized Periodic Averaging SGD(DPA-SGD)를 도입합니다.
• 클라이언트별 작업 공분산 행렬 Ωk와 전역 작업 관계 정규화 항 Tr(Φ_task Ω−1 Φ_task^T)를 사용하여 클라이언트 간 작업 정보를 공유합니다.
• 네트워크 가중치 Wk와 Ωk에 대한 교대 최적화를 제안하고 이웃 간 Ωk를 통신을 통해 정렬합니다(부록의 Algorithm 1).
• 비볼록 설정에서 평균화된 모델이 수렴하는 조건을 보여주는 DPA-SGD에 대한 수렴 분석(정리 1)을 제공합니다.
• MoleculeNet 데이터셋(SIDER, Tox21, MUV, QM8)을 비 IID 부분 라벨 분할과 두 가지 GNN 백본(GraphSAGE, GAT)으로 평가합니다.

실험 결과

연구 질문

• RQ1 서버리스이고 비휘발적으로 분산화된 연합 다중 작업 학습 프레임워크가 부분 라벨이 있는 경우 GNN을 효과적으로 학습할 수 있는가?
• RQ2 작업 관계 정규화를 갖춘 DPA-SGD가 비 IID 분자 그래프 설정에서 중앙 집중식 FedAvg보다 성능을 개선하는가?
• RQ3 SpreadGNN이 GNN 아키텍처에 독립적이며 다양한 클라이언트 연결성/토폴로지에 강건한가?
• RQ4 토폴로지, 이웃의 크기, 통신 주기가 학습 성능과 수렴에 어떤 영향을 미치는가?
• RQ5 클라이언트가 로컬 라벨에 없는 작업을 교차-클라이언트 작업 관계를 통해 예측하는 것이 가능한가?

주요 결과

• SpreadGNN은 모든 클라이언트가 통신할 수 있는 경우 중앙 서버가 없었음에도 분자 특성 예측 과제에서 FedAvg보다 성능이 앞섰습니다.
• SpreadGNN은 서버 보조 FedGMTL보다 경쟁력 있거나 더 나은 정확도를 달성하며, 현실적인 서버리스 토폴로지에서 중앙 서버 성능에 근접할 수 있습니다.
• 프레임워크는 GNN 모델의 선택(GraphSAGE 또는 GAT)에 독립적이며, 제한된 클라이언트 연결성에서도 효과적임을 보여줍니다.
• 작업 관계 정규화(Ω)와 분산 평균화를 통해 부분적으로 라벨이 있는 작업들 간 학습이 가능해져 비 IID 설정에서 성능이 향상됩니다.
• 수렴 분석(정리 1)은 DPA-SGD가 비볼록, 분산 FL 설정에서 수렴하는 조건을 제공합니다.