FedGrAINS: Personalized SubGraph Federated Learning with Adaptive Neighbor Sampling
Emir Ceyani, Han Xie, Baturalp Buyukates, Carl Yang, Salman Avestimehr
TL;DR
FedGrAINS addresses privacy-preserving learning on distributed graphs with missing inter-subgraph links by introducing a data-adaptive, sampling-based regularizer. It uses Generative Flow Networks (GFlowNets) to estimate layer-wise neighbor importance and to perform task-aligned adaptive sampling of adjacency matrices across $L$ GNN layers, guided by a trajectory balance objective. The framework jointly trains a node classifier and an importance sampler at each client, enabling personalized subgraph FL without additional communication. Experimental results on six real-world datasets show consistent improvements over strong baselines, especially in disjoint subgraph settings with increased client counts, demonstrating its effectiveness against node-degree heterogeneity and missing-link challenges. The approach is plug-and-play with existing subgraph FL methods and has practical implications for privacy-preserving, scalable graph learning.
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.
