[FederatedLearning][Post-v1] Federated learning over graphs (secure protocols)

[ooples]

Out of scope for FL v1.\n\nDetails: docs/FederatedLearning_OutOfScope_Backlog.md\n\nGoal: Implement graph federated learning with specialized secure protocols for graph-structured data.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

✅ Successfully generated implementation plan.

💡 Use @coderabbitai <feedback> to replan

Examples

• Example 1
• Example 2

🔗 Similar Issues

Related Issues

• [Phase 3] Implement Federated Learning Framework #398
• [FederatedLearning][Post-v1] Full on-device confidential training in iOS/Android TEEs #537
• [FederatedLearning][Post-v1] Create post-v1 backlog tracking issue #536

👤 Suggested Assignees

• ooples

---

🧪 Issue enrichment is currently in early access.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

[coderabbitai]

Implementation Plan

Approach

Implement federated learning capabilities for Graph Neural Networks with specialized secure protocols that address graph-specific privacy concerns including structural privacy, neighborhood information leakage, and topology heterogeneity. The implementation follows a hybrid approach by creating core federated learning abstractions as a new module while leveraging existing communication backends and optimizer patterns.

Observations

The codebase contains comprehensive Graph Neural Network implementations (GraphNeuralNetwork, GraphAttentionNetwork, GraphSAGENetwork) with supporting graph data structures (GraphData, graph layers) and loaders. A mature distributed training framework exists with sharded optimizers, multiple communication backends (NCCL, MPI, Gloo), and gradient compression utilities. The optimizer architecture follows an interface-driven design with OptimizerBase and GradientBasedOptimizerBase as extension points. No federated learning infrastructure currently exists, but the distributed training patterns provide a strong foundation. Privacy mechanisms (differential privacy, secure aggregation, homomorphic encryption) are not yet implemented.

Assumptions

Assumption 1: Federated Learning vs. Distributed Training Separation

Options Considered:

• Extend existing ShardedOptimizerBase directly for federated scenarios
• Create separate src/FederatedLearning module with dedicated abstractions
• Hybrid approach with base FL infrastructure that leverages existing communication backends

Chosen Option: Hybrid approach with base FL infrastructure that leverages existing communication backends

Rationale: Hybrid approach - create core federated learning abstractions (client, server, aggregator) as a new module while leveraging existing communication backends and optimizer patterns. This provides clear conceptual separation between federated and distributed training while maximizing code reuse. The pattern aligns with the codebase's preference for interface-driven design and composition (similar to how GradientCompressionOptimizer wraps other optimizers).

Assumption 2: Graph Partitioning Strategy Default

Options Considered:

• Community-based partitioning (Metis/Louvain) for natural clustering
• Random node sampling for simplicity and baseline
• Hybrid approach with pluggable partitioning strategies

Chosen Option: Hybrid approach with pluggable partitioning strategies

Rationale: Hybrid approach with community-based partitioning as the default. Implement an IGraphPartitioner interface that allows pluggable strategies, defaulting to Louvain (which is simpler than Metis and produces quality communities). This aligns with the codebase's extensibility patterns (IGraphDataLoader, ICommunicationBackend) and provides a practical default while enabling future experimentation with other strategies.

Assumption 3: Privacy Mechanism Integration Point

Options Considered:

• Directly in federated optimizer during aggregation
• As decorators wrapping communication backends
• As middleware layer between optimizer and communication

Chosen Option: As middleware layer between optimizer and communication

Rationale: Middleware approach - implement privacy mechanisms as composable components that wrap federated aggregators. This follows the decorator pattern already present in GradientCompressionOptimizer and allows privacy mechanisms to be independently tested, composed (e.g., both differential privacy AND secure aggregation), and optionally enabled. Create IPrivacyMechanism interface that federated aggregators can optionally use.

💡 User Tips

Regenerate the plan with different choices with @coderabbitai <feedback>.

Plan

Phase 1: Core Federated Learning Infrastructure

Establish foundational abstractions for federated learning that separate client-side training from server-side aggregation, leveraging existing communication backends while introducing federated-specific coordination patterns.

Task 1: Define Federated Learning Abstractions

Create core interfaces and base classes that define the federated learning lifecycle.

• Define IFederatedClient<T> interface in src/Interfaces/ with methods for local training, model retrieval, and parameter submission
• Define IFederatedServer<T> interface in src/Interfaces/ for client registration, round coordination, and model distribution
• Define IModelAggregator<T> interface in src/Interfaces/ for aggregating client models with pluggable aggregation strategies (FedAvg, FedProx, weighted averaging)
• Create base class FederatedCoordinator<T> in src/FederatedLearning/ that manages training rounds, client selection, and convergence tracking
• Add FederatedTrainingConfiguration<T> in src/FederatedLearning/ to hold federated-specific settings (number of clients, rounds, local epochs, client selection strategy)

Task 2: Implement Base Federated Optimizer

Create optimizer that coordinates federated training rounds and integrates with existing optimizer architecture.

• Create FederatedOptimizerBase<T, TInput, TOutput> in src/FederatedLearning/ extending GradientBasedOptimizerBase
• Implement federated training loop: select clients, distribute model, collect updates, aggregate, update global model
• Add methods for client model collection and global model distribution using existing serialization patterns from src/Serialization/
• Integrate with existing communication backends by wrapping ICommunicationBackend for federated scenarios
• Support both synchronous (all clients per round) and asynchronous (subset of clients) federated training modes

Task 3: Create Federated Aggregation Strategies

Implement standard federated aggregation algorithms that combine client models.

• Create FederatedAveraging<T> aggregator in src/FederatedLearning/Aggregators/ implementing weighted parameter averaging (FedAvg)
• Create FederatedProximal<T> aggregator in src/FederatedLearning/Aggregators/ implementing FedProx with proximal term for handling non-IID data
• Add aggregation utilities for handling client dropout and partial participation
• Integrate with existing AllReduce patterns from distributed training for efficient parameter averaging
• Support aggregation metrics (client contribution weights, convergence tracking)

🤖 Prompt for AI agents

This phase establishes the core federated learning infrastructure by creating
abstractions for clients, servers, and aggregators that work with the existing
optimizer architecture.

Create IFederatedClient<T>, IFederatedServer<T>, and IModelAggregator<T>
interfaces in src/Interfaces/
Add FederatedCoordinator<T> base class and FederatedTrainingConfiguration<T> in
src/FederatedLearning/

Implement FederatedOptimizerBase<T, TInput, TOutput> extending
GradientBasedOptimizerBase
Add federated training loop with client selection, model distribution, update
collection, and aggregation
Integrate with ICommunicationBackend and serialization patterns
Support both synchronous and asynchronous training modes

Create FederatedAveraging<T> and FederatedProximal<T> aggregators in
src/FederatedLearning/Aggregators/
Implement weighted averaging, proximal term handling, client dropout support
Add aggregation metrics and integrate with AllReduce patterns

Phase 2: Graph Partitioning and Federated Data Loading

Implement graph-specific data partitioning strategies that distribute graph data across federated clients while maintaining structural integrity and handling missing neighbor problems.

Task 1: Implement Graph Partitioning Utilities

Create partitioning algorithms that divide graphs into client subgraphs.

• Create IGraphPartitioner<T> interface in src/Interfaces/ defining partitioning contract (Partition method returning client-to-node assignments)
• Implement LouvainPartitioner<T> in src/Data/Graph/Partitioning/ using community detection for natural graph clustering
• Implement RandomNodePartitioner<T> in src/Data/Graph/Partitioning/ for baseline random partitioning
• Implement DegreeBasedPartitioner<T> in src/Data/Graph/Partitioning/ distributing nodes based on degree distribution to balance computation
• Add partitioning quality metrics (edge cut, modularity, balance) for evaluating partition strategies
• Ensure partitioners work with existing GraphData structures from src/Data/Structures/

Task 2: Create Federated Graph Data Loader

Extend graph data loading to support federated subgraph extraction with handling for missing neighbors.

• Create FederatedGraphDataLoader<T> in src/Data/Graph/ extending GraphDataLoaderBase
• Implement subgraph extraction that creates local GraphData instances per client from partitions
• Add missing neighbor handling: create placeholder nodes for cross-client edges, track missing neighbor IDs
• Implement overlapping node strategies where boundary nodes are replicated to maintain message passing
• Support edge sampling for link prediction tasks that respects client boundaries
• Integrate with existing CitationNetworkLoader, MolecularDatasetLoader, and OGBDatasetLoader patterns for federated scenarios

Task 3: Implement Graph Sampling Strategies

Create subgraph sampling mechanisms that can be used during federated training rounds.

• Create ISubgraphSampler<T> interface in src/Interfaces/ for dynamic subgraph sampling
• Implement NeighborhoodSampler<T> in src/Data/Graph/Sampling/ for K-hop neighborhood sampling around anchor nodes
• Implement RandomWalkSampler<T> in src/Data/Graph/Sampling/ for sampling via random walks
• Add label-guided sampling that balances class distributions across client subgraphs
• Support incremental sampling for efficient multi-round federated training
• Enable sampling to work with both disjoint and overlapping subgraph settings

🤖 Prompt for AI agents

This phase implements graph-specific data partitioning and loading to distribute
graph data across federated clients while handling structural integrity
challenges.

Create IGraphPartitioner<T> interface in src/Interfaces/ with Partition method
Implement LouvainPartitioner<T>, RandomNodePartitioner<T>, and
DegreeBasedPartitioner<T> in src/Data/Graph/Partitioning/
Add partitioning quality metrics (edge cut, modularity, balance)
Ensure compatibility with GraphData<T> structures

Create FederatedGraphDataLoader<T> extending GraphDataLoaderBase in
src/Data/Graph/
Implement subgraph extraction creating per-client GraphData instances
Add missing neighbor handling with placeholder nodes and tracking
Implement overlapping node strategies and edge sampling for link prediction
Integrate with existing data loader patterns

Create ISubgraphSampler<T> interface in src/Interfaces/
Implement NeighborhoodSampler<T> and RandomWalkSampler<T> in
src/Data/Graph/Sampling/
Add label-guided sampling and incremental sampling support
Enable both disjoint and overlapping subgraph modes

Phase 3: Privacy-Preserving Mechanisms for Graph Federated Learning

Implement differential privacy and secure aggregation protocols specialized for graph-structured data, addressing structural privacy concerns and neighborhood information leakage.

Task 1: Implement Differential Privacy for GNNs

Create differential privacy mechanisms adapted for graph neural network message passing operations.

• Create IDifferentialPrivacy<T> interface in src/Interfaces/ defining privacy budget, noise injection, and sensitivity calculation methods
• Implement GraphAggregationPerturbation<T> (GAP) in src/FederatedLearning/Privacy/ that adds calibrated noise to GNN aggregation steps for edge-level and node-level privacy
• Add sensitivity calculator for GNN operations accounting for message passing and neighborhood aggregation
• Implement privacy budget tracker using moments accountant technique that accumulates privacy loss across training rounds
• Add gradient clipping specifically for graph layers before noise addition to bound sensitivity
• Create PrivacyConfiguration<T> in src/FederatedLearning/Privacy/ to manage privacy budgets (epsilon, delta), noise scale, clipping thresholds

Task 2: Implement Secure Aggregation Protocols

Create cryptographic protocols for securely aggregating client models without revealing individual contributions.

• Create ISecureAggregationProtocol<T> interface in src/Interfaces/ for secure aggregation operations
• Implement SecretSharingAggregator<T> in src/FederatedLearning/Security/ using Shamir secret sharing for threshold-based aggregation
• Implement AdditiveSecretSharing<T> in src/FederatedLearning/Security/ for lightweight pairwise masking
• Add mask generation and cancellation mechanisms that ensure individual client models remain hidden
• Integrate with Microsoft SEAL library for homomorphic encryption support (create wrapper in src/FederatedLearning/Security/Encryption/)
• Support client dropout handling in secure aggregation (reconstruct secrets when clients fail)

Task 3: Create Privacy-Aware Federated Aggregator

Combine privacy mechanisms with federated aggregation as composable middleware.

• Create PrivacyAwareFederatedAggregator<T> in src/FederatedLearning/Aggregators/ that wraps base aggregators with privacy mechanisms
• Implement middleware pattern allowing composition of multiple privacy mechanisms (DP + secure aggregation)
• Add privacy budget enforcement that rejects training rounds when budget exhausted
• Create privacy audit logging that tracks privacy expenditure across federated rounds
• Support per-client and global privacy accounting for heterogeneous privacy requirements
• Ensure privacy mechanisms are independently testable and can be disabled for non-privacy scenarios

🤖 Prompt for AI agents

This phase implements privacy-preserving mechanisms specialized for graph
federated learning, including differential privacy and secure aggregation
protocols.

Create IDifferentialPrivacy<T> interface in src/Interfaces/ with privacy budget,
noise injection, and sensitivity methods
Implement GraphAggregationPerturbation<T> in src/FederatedLearning/Privacy/ with
calibrated noise for GNN operations
Add sensitivity calculator for message passing and neighborhood aggregation
Implement privacy budget tracker using moments accountant
Add gradient clipping for graph layers
Create PrivacyConfiguration<T> for managing privacy budgets and thresholds

Create ISecureAggregationProtocol<T> interface in src/Interfaces/
Implement SecretSharingAggregator<T> and AdditiveSecretSharing<T> in
src/FederatedLearning/Security/
Add mask generation and cancellation mechanisms
Create Microsoft SEAL wrapper in src/FederatedLearning/Security/Encryption/
Support client dropout handling with secret reconstruction

Create PrivacyAwareFederatedAggregator<T> in src/FederatedLearning/Aggregators/
Implement middleware pattern for composing multiple privacy mechanisms
Add privacy budget enforcement and audit logging
Support per-client and global privacy accounting
Ensure independent testability and optional disabling

Phase 4: Federated Graph Neural Network Training Integration

Integrate federated learning with existing GNN implementations, handling graph-specific challenges like topology heterogeneity and missing neighbors during message passing.

Task 1: Create Federated Graph Optimizer

Implement optimizer that coordinates federated training specifically for graph neural networks.

• Create FederatedGraphOptimizer<T, TInput, TOutput> in src/FederatedLearning/ extending FederatedOptimizerBase
• Override UpdateParameters to handle graph-specific aggregation accounting for topology differences across clients
• Implement missing neighbor handling during aggregation: propagate placeholder embeddings or use learned proxies
• Add support for heterogeneous graph aggregation where clients have different graph structures
• Integrate with existing GNN training patterns from src/NeuralNetworks/ (GraphNeuralNetwork, GraphAttentionNetwork, GraphSAGENetwork)
• Support graph-level and node-level federated learning tasks

Task 2: Adapt GNN Models for Federated Training

Extend existing GNN implementations to support federated training scenarios.

• Create IFederatedGraphModel<T> interface in src/Interfaces/ defining federated-specific methods (GetSubmodelParameters, SetGlobalParameters, HandleMissingNeighbors)
• Update GraphNeuralNetwork.cs, GraphAttentionNetwork.cs, GraphSAGENetwork.cs in src/NeuralNetworks/ to implement federated interface
• Add missing neighbor embeddings that can be updated during federated rounds
• Implement cross-client message passing simulation for overlapping nodes
• Support auxiliary losses for graph smoothness regularization that work in federated settings
• Ensure backward compatibility with existing non-federated GNN training

Task 3: Implement Federated Graph Training Coordinator

Create high-level orchestrator that manages end-to-end federated graph learning.

• Create FederatedGraphTrainingCoordinator<T> in src/FederatedLearning/ that coordinates all phases: partitioning, client training, aggregation, privacy
• Implement client selection strategies (random, weighted by graph size, importance sampling)
• Add convergence detection mechanisms accounting for graph non-IID challenges
• Support evaluation on global test graphs that span multiple client subgraphs
• Integrate compression from src/ModelCompression/ to reduce communication overhead for graph model updates
• Create training loop that alternates between local GNN training and federated aggregation rounds
• Add support for both node classification, graph classification, and link prediction federated tasks

🤖 Prompt for AI agents

This phase integrates federated learning with existing GNN implementations,
handling graph-specific challenges and creating a complete training pipeline.

Create FederatedGraphOptimizer<T, TInput, TOutput> extending
FederatedOptimizerBase in src/FederatedLearning/
Override UpdateParameters for graph-specific aggregation with topology
differences
Implement missing neighbor handling with placeholder embeddings or learned
proxies
Add heterogeneous graph aggregation support
Integrate with existing GNN implementations (GraphNeuralNetwork,
GraphAttentionNetwork, GraphSAGENetwork)
Support graph-level and node-level federated tasks

Create IFederatedGraphModel<T> interface in src/Interfaces/ with
GetSubmodelParameters, SetGlobalParameters, HandleMissingNeighbors methods
Update GraphNeuralNetwork.cs, GraphAttentionNetwork.cs, GraphSAGENetwork.cs to
implement federated interface
Add missing neighbor embeddings and cross-client message passing simulation
Support auxiliary losses for graph smoothness regularization
Ensure backward compatibility with non-federated training

Create FederatedGraphTrainingCoordinator<T> in src/FederatedLearning/
Implement client selection strategies (random, weighted, importance sampling)
Add convergence detection for graph non-IID challenges
Support global test graph evaluation
Integrate model compression for communication efficiency
Create training loop alternating between local training and federated
aggregation
Support node classification, graph classification, and link prediction tasks

---

🚀 Next Steps

🤖 All AI agent prompts combined

Task: 1

This phase establishes the core federated learning infrastructure by creating
abstractions for clients, servers, and aggregators that work with the existing
optimizer architecture.

Create IFederatedClient<T>, IFederatedServer<T>, and IModelAggregator<T>
interfaces in src/Interfaces/
Add FederatedCoordinator<T> base class and FederatedTrainingConfiguration<T> in
src/FederatedLearning/

Implement FederatedOptimizerBase<T, TInput, TOutput> extending
GradientBasedOptimizerBase
Add federated training loop with client selection, model distribution, update
collection, and aggregation
Integrate with ICommunicationBackend and serialization patterns
Support both synchronous and asynchronous training modes

Create FederatedAveraging<T> and FederatedProximal<T> aggregators in
src/FederatedLearning/Aggregators/
Implement weighted averaging, proximal term handling, client dropout support
Add aggregation metrics and integrate with AllReduce patterns
===============================================================================

Task: 2

This phase implements graph-specific data partitioning and loading to distribute
graph data across federated clients while handling structural integrity
challenges.

Create IGraphPartitioner<T> interface in src/Interfaces/ with Partition method
Implement LouvainPartitioner<T>, RandomNodePartitioner<T>, and
DegreeBasedPartitioner<T> in src/Data/Graph/Partitioning/
Add partitioning quality metrics (edge cut, modularity, balance)
Ensure compatibility with GraphData<T> structures

Create FederatedGraphDataLoader<T> extending GraphDataLoaderBase in
src/Data/Graph/
Implement subgraph extraction creating per-client GraphData instances
Add missing neighbor handling with placeholder nodes and tracking
Implement overlapping node strategies and edge sampling for link prediction
Integrate with existing data loader patterns

Create ISubgraphSampler<T> interface in src/Interfaces/
Implement NeighborhoodSampler<T> and RandomWalkSampler<T> in
src/Data/Graph/Sampling/
Add label-guided sampling and incremental sampling support
Enable both disjoint and overlapping subgraph modes
===============================================================================

Task: 3

This phase implements privacy-preserving mechanisms specialized for graph
federated learning, including differential privacy and secure aggregation
protocols.

Create IDifferentialPrivacy<T> interface in src/Interfaces/ with privacy budget,
noise injection, and sensitivity methods
Implement GraphAggregationPerturbation<T> in src/FederatedLearning/Privacy/ with
calibrated noise for GNN operations
Add sensitivity calculator for message passing and neighborhood aggregation
Implement privacy budget tracker using moments accountant
Add gradient clipping for graph layers
Create PrivacyConfiguration<T> for managing privacy budgets and thresholds

Create ISecureAggregationProtocol<T> interface in src/Interfaces/
Implement SecretSharingAggregator<T> and AdditiveSecretSharing<T> in
src/FederatedLearning/Security/
Add mask generation and cancellation mechanisms
Create Microsoft SEAL wrapper in src/FederatedLearning/Security/Encryption/
Support client dropout handling with secret reconstruction

Create PrivacyAwareFederatedAggregator<T> in src/FederatedLearning/Aggregators/
Implement middleware pattern for composing multiple privacy mechanisms
Add privacy budget enforcement and audit logging
Support per-client and global privacy accounting
Ensure independent testability and optional disabling
===============================================================================

Task: 4

This phase integrates federated learning with existing GNN implementations,
handling graph-specific challenges and creating a complete training pipeline.

Create FederatedGraphOptimizer<T, TInput, TOutput> extending
FederatedOptimizerBase in src/FederatedLearning/
Override UpdateParameters for graph-specific aggregation with topology
differences
Implement missing neighbor handling with placeholder embeddings or learned
proxies
Add heterogeneous graph aggregation support
Integrate with existing GNN implementations (GraphNeuralNetwork,
GraphAttentionNetwork, GraphSAGENetwork)
Support graph-level and node-level federated tasks

Create IFederatedGraphModel<T> interface in src/Interfaces/ with
GetSubmodelParameters, SetGlobalParameters, HandleMissingNeighbors methods
Update GraphNeuralNetwork.cs, GraphAttentionNetwork.cs, GraphSAGENetwork.cs to
implement federated interface
Add missing neighbor embeddings and cross-client message passing simulation
Support auxiliary losses for graph smoothness regularization
Ensure backward compatibility with non-federated training

Create FederatedGraphTrainingCoordinator<T> in src/FederatedLearning/
Implement client selection strategies (random, weighted, importance sampling)
Add convergence detection for graph non-IID challenges
Support global test graph evaluation
Integrate model compression for communication efficiency
Create training loop alternating between local training and federated
aggregation
Support node classification, graph classification, and link prediction tasks

💡 Iterate on the plan with: @coderabbitai <feedback>

Example Feedback
- @coderabbitai You can skip phase 3. Add a simple unit test case for phase 2.
- @coderabbitai For assumption 1 go ahead with option 3 and replan.