Table of Contents
Fetching ...

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.

FedGrAINS: Personalized SubGraph Federated Learning with Adaptive Neighbor Sampling

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 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.
Paper Structure (18 sections, 7 equations, 1 figure, 8 tables, 2 algorithms)

This paper contains 18 sections, 7 equations, 1 figure, 8 tables, 2 algorithms.

Figures (1)

  • Figure 1: The illustration of the proposed joint training scheme of 2-layer GNN and GFlowNet for our personalized subgraph FL framework, FedGrAINS from the client side. For any node $v$ in the client's graph, we first estimate the important neighbors with GFlowNet, then sample $k$ nodes using the Gumbel-Max trick huijben2022review-gumbel. After choosing the important nodes for two layers, we input the sampled subgraph to the GNN. After obtaining our novel loss, we backpropagate through the classification and the trajectory balance losses, with the reward function defined as the cross-entropy loss at a specific round. With classification loss as a reward function, GFlowNet adapts itself to select and sample neighbors vital for optimal task performance.