FedGrAINS: Personalized SubGraph Federated Learning with Adaptive Neighbor Sampling

TL;DR

FedGrAINS introduces an adaptive, sampling-based regularization for personalized federated learning on subgraphs, utilizing generative flow networks to improve GNN training across heterogeneous client data.

Contribution

The paper proposes FedGrAINS, a novel method using GFlowNets for adaptive node importance evaluation, enhancing personalized subgraph federated learning.

Findings

FedGrAINS improves federated GNN performance over baselines.

Adaptive sampling based on GFlowNets enhances training efficiency.

Regularization with FedGrAINS yields more accurate personalized models.

Abstract

Graphs are crucial for modeling relational and biological data. As datasets grow larger in real-world scenarios, the risk of exposing sensitive information increases, making privacy-preserving training methods like federated learning (FL) essential to ensure data security and compliance with privacy regulations. Recently proposed personalized subgraph FL methods have become the de-facto standard for training personalized Graph Neural Networks (GNNs) in a federated manner while dealing with the missing links across clients' subgraphs due to privacy restrictions. However, personalized subgraph FL faces significant challenges due to the heterogeneity in client subgraphs, such as degree distributions among the nodes, which complicate federated training of graph models. To address these challenges, we propose \textit{FedGrAINS}, a novel data-adaptive and sampling-based regularization method for subgraph FL. FedGrAINS leverages generative flow networks (GFlowNets) to evaluate node importance concerning clients' tasks, dynamically adjusting the message-passing step in clients' GNNs. This adaptation reflects task-optimized sampling aligned with a trajectory balance objective. Experimental results demonstrate that the inclusion of \textit{FedGrAINS} as a regularizer consistently improves the FL performance compared to baselines that do not leverage such regularization.