[FederatedLearning][Post-v1] PSI for vertical FL entity alignment

[ooples]

Out of scope for FL v1.\n\nDetails: docs/FederatedLearning_OutOfScope_Backlog.md\n\nGoal: Explore/implement private set intersection workflows to align entities across parties for vertical federated learning.

[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] 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 Private Set Intersection (PSI) workflows to align entities across parties for vertical federated learning, enabling secure identification of common data samples without revealing non-overlapping records.

Observations

Codebase Summary

AiDotNet is a comprehensive machine learning framework with mature distributed training infrastructure (25+ files in src/DistributedTraining/) supporting data-parallel training via DDP, FSDP, and ZeRO strategies. The system features pluggable communication backends (ICommunicationBackend<T> with MPI, NCCL, Gloo, and in-memory implementations), layer-based neural network abstractions, and serialization infrastructure. However, the codebase currently lacks federated learning abstractions, privacy-preserving primitives, and external entity ID handling—all required for vertical FL with PSI.

Assumptions

Assumption 1: PSI Protocol Selection

Options Considered:

• DDH-based PSI with Bloom filters (better for asymmetric set sizes, moderate security, complex implementation)
• OPRF-based PSI (balanced performance, strong security, simpler to implement with existing crypto libraries)
• Homomorphic encryption-based PSI using Microsoft SEAL (strongest security, highest performance cost, most complex)

Chosen Option: OPRF-based PSI (balanced performance, strong security, simpler to implement with existing crypto libraries)

Rationale: Start with OPRF-based PSI (Naor-Pinkas style) as the initial implementation. This protocol provides strong security guarantees, reasonable performance for moderate dataset sizes, and can be implemented using .NET's existing cryptographic primitives in System.Security.Cryptography. The modular design will allow adding DDH-based or HE-based protocols as alternatives in the future. Developers will appreciate the simpler implementation that doesn't require heavy cryptographic library dependencies initially.

Assumption 2: Entity ID Representation

Options Considered:

• Extend existing loaders to accept external ID mappings in-memory (simple, no persistence)
• Create separate entity ID abstraction layer with persistent mapping (more robust, supports larger datasets)
• Embed IDs directly in data structures like Matrix/Vector (invasive, breaks existing API)

Chosen Option: Create separate entity ID abstraction layer with persistent mapping (more robust, supports larger datasets)

Rationale: Create a lightweight IEntityIdMapper abstraction that maintains a mapping between external entity IDs (strings) and internal row indices (integers). This keeps the core tensor/matrix operations unchanged while providing a clean integration point for PSI output consumption. The mapper will be optional—existing data loaders continue working without it, but VFL scenarios explicitly provide entity mappings.

Assumption 3: VFL Coordination Architecture

Options Considered:

• Fork distributed training abstractions specifically for VFL (complete separation, clearer semantics)
• Extend existing ICommunicationBackend and sharding configuration for VFL (reuses infrastructure, risk of confusion)
• Create parallel VFL-specific communication layer using existing backend as transport (best of both worlds)

Chosen Option: Create parallel VFL-specific communication layer using existing backend as transport (best of both worlds)

Rationale: Introduce a VerticalFederatedCoordinator that uses ICommunicationBackend as the transport layer but implements VFL-specific semantics (feature aggregation, split-learning patterns, encrypted gradient exchange). This reuses the proven communication patterns while avoiding contamination of the data-parallel abstractions with incompatible VFL concepts. Developers familiar with the distributed training code will recognize the communication primitives while understanding VFL's distinct coordination requirements.

💡 User Tips

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

Plan

Phase 1: PSI Protocol Foundation

Build the cryptographic protocol layer supporting private set intersection between parties. This phase establishes the core PSI algorithms using OPRF-based approaches with .NET's System.Security.Cryptography primitives.

Task 1: Create PSI Protocol Abstractions

Define interfaces and base types for PSI protocol implementations to support multiple algorithm variants.

• Create src/PrivacyPreserving/PSI/IPSIProtocol.cs interface defining ComputeIntersectionAsync() with party role parameters (client/server)
• Add PSIProtocolBase abstract class in same directory handling common validation, timeout management, and error handling
• Define PSIConfiguration class capturing protocol parameters (hash function selection, security parameters, communication timeouts)
• Introduce PSIResult record type containing intersection indices, metadata (protocol timing, communication overhead), and party-specific views

Task 2: Implement OPRF-Based PSI Protocol

Build concrete PSI implementation using oblivious pseudorandom function pattern with cryptographic hashing.

• Create src/PrivacyPreserving/PSI/OPRFBasedPSI.cs implementing IPSIProtocol
• Use System.Security.Cryptography.HMACSHA256 for PRF construction with server-side keying
• Implement client-side blinding/unblinding using multiplicative group operations
• Add server-side PRF evaluation without learning client inputs
• Implement intersection determination by comparing PRF outputs with constant-time equality checks

Task 3: Build PSI Communication Primitives

Create message serialization and protocol handshake coordination leveraging existing communication backends.

• Create src/PrivacyPreserving/PSI/Messages/ directory for protocol message types
• Add PSIProtocolMessage base class with message type enumeration and serialization contracts
• Implement specific message types: ClientBlindedSet, ServerPRFEvaluations, ClientIntersectionRequest
• Create PSIProtocolCoordinator wrapping ICommunicationBackend<T> to handle tagged message routing for PSI handshakes
• Use message tags to multiplex PSI traffic from other distributed operations on shared communication channel

🤖 Prompt for AI agents

Build the cryptographic PSI protocol foundation using OPRF-based approach. This
establishes secure set intersection capabilities between parties.

Create PSI protocol abstractions:
- Define `IPSIProtocol` interface with `ComputeIntersectionAsync()` method
accepting party roles
- Create `PSIProtocolBase` abstract class with validation, timeout, and error
handling
- Add `PSIConfiguration` class for protocol parameters
- Create `PSIResult` record containing intersection indices and metadata

Implement OPRF-based PSI:
- Create `OPRFBasedPSI` class implementing `IPSIProtocol`
- Use `HMACSHA256` for PRF with server-side keying
- Implement client blinding/unblinding with multiplicative group operations
- Add server PRF evaluation preserving input privacy
- Implement constant-time intersection determination

Build communication primitives:
- Create message types directory under PSI
- Define `PSIProtocolMessage` base with serialization
- Implement specific messages: `ClientBlindedSet`, `ServerPRFEvaluations`,
`ClientIntersectionRequest`
- Create `PSIProtocolCoordinator` wrapping `ICommunicationBackend<T>` for
message routing
- Add message tagging for PSI traffic multiplexing

Phase 2: Entity Alignment Integration

Integrate PSI outputs into data loading pipeline by introducing entity ID abstractions and alignment-aware loaders. This phase bridges cryptographic PSI results with the existing tensor-based training data structures.

Task 1: Introduce Entity ID Mapping Abstraction

Create lightweight mapping layer maintaining correspondence between external entity identifiers and internal tensor row indices.

• Create src/Data/EntityAlignment/IEntityIdMapper.cs interface with methods: AddMapping(string externalId, int rowIndex), GetRowIndex(string externalId), GetAlignedIndices(IEnumerable<string> intersectionIds)
• Implement InMemoryEntityIdMapper using Dictionary<string, int> for bidirectional lookup
• Add thread-safe operations supporting concurrent query during data loading
• Include serialization support extending existing JSON converter patterns in src/Serialization/

Task 2: Create PSI-Aware Data Loaders

Extend existing data loader abstractions to consume PSI intersection results and filter datasets to aligned entities.

• Create src/Data/Loaders/AlignedDataLoader.cs base class accepting IEntityIdMapper and PSIResult in constructor
• Override batch iteration to only yield rows corresponding to intersection indices
• Maintain correspondence between feature matrix rows and entity IDs during sampling
• Add CreateFromPSI() factory method constructing aligned loaders from PSI protocol output
• Ensure compatibility with existing episodic sampling patterns in EpisodicDataLoaderBase for meta-learning scenarios

Task 3: Implement Entity ID Preprocessing

Create utilities for extracting and preprocessing entity identifiers from raw datasets before PSI execution.

• Create src/DataProcessor/EntityIdExtractor.cs handling common ID formats (CSV columns, database keys, composite identifiers)
• Add normalization and hashing helpers ensuring consistent ID representation across parties
• Implement Bloom filter-based preprocessing for large-scale cardinality reduction before PSI
• Provide validation utilities checking ID uniqueness and format consistency

🤖 Prompt for AI agents

Integrate PSI results with data loading by creating entity ID mapping and
alignment-aware loaders. This bridges PSI output with tensor-based training
data.

Create entity ID mapping:
- Define `IEntityIdMapper` interface with `AddMapping()`, `GetRowIndex()`, and
`GetAlignedIndices()` methods
- Implement `InMemoryEntityIdMapper` using bidirectional dictionary lookup
- Add thread-safe concurrent operations
- Extend existing serialization patterns for persistence

Build PSI-aware data loaders:
- Create `AlignedDataLoader` base class accepting `IEntityIdMapper` and
`PSIResult`
- Override batch iteration to filter rows by intersection indices
- Maintain entity ID to feature matrix row correspondence
- Add `CreateFromPSI()` factory method
- Ensure compatibility with `EpisodicDataLoaderBase` patterns

Implement entity preprocessing:
- Create `EntityIdExtractor` for common ID formats (CSV, database, composite)
- Add normalization and hashing for consistent representation
- Implement Bloom filter preprocessing for cardinality reduction
- Provide validation for ID uniqueness and format consistency

Phase 3: VFL Coordination Layer

Establish vertical federated learning coordination patterns extending existing distributed training infrastructure for feature-partitioned scenarios. This phase implements the orchestration required for multi-party collaborative training.

Task 1: Create Vertical FL Coordinator

Build orchestrator managing feature aggregation and gradient coordination across parties with different feature subsets.

• Create src/FederatedLearning/Vertical/VerticalFederatedCoordinator.cs using ICommunicationBackend<T> for transport
• Implement CoordinateForwardPass() method aggregating embeddings from feature parties into active party's model
• Add CoordinateBackwardPass() distributing gradients to feature parties for local parameter updates
• Handle party role assignment (active party vs passive parties) based on label ownership
• Use barrier synchronization from existing distributed training patterns to coordinate computation phases

Task 2: Implement Split-Learning Layer Abstractions

Extend neural network layer abstractions to support vertical feature partitioning and secure embedding aggregation.

• Create src/NeuralNetworks/Layers/SplitInputLayer.cs accepting multiple feature sources from different parties
• Add SecureEmbeddingAggregator layer combining embeddings with privacy-preserving aggregation
• Extend LayerBase.Forward() contract to support receiving encrypted intermediate activations
• Implement gradient splitting in backward pass routing party-specific gradients to respective feature owners
• Leverage existing SplitLayer.cs composition patterns but add VFL-specific encryption hooks

Task 3: Build VFL Training Configuration

Create configuration abstractions for vertical federated learning scenarios extending existing sharding configuration patterns.

• Create src/FederatedLearning/Vertical/VFLConfiguration.cs holding party topology (active vs passive), feature ownership mapping, and encryption parameters
• Add PartyRole enumeration distinguishing label owner (active) from feature providers (passive)
• Include FeaturePartitionMap describing which features belong to each party
• Extend ShardingConfiguration<T> pattern with VFL-specific parameters while maintaining communication backend compatibility
• Provide validation ensuring exactly one active party and consistent feature coverage across passive parties

🤖 Prompt for AI agents

Build VFL coordination layer for multi-party collaborative training with feature
partitioning. This extends distributed training infrastructure for vertical
federated learning.

Create vertical FL coordinator:
- Create `VerticalFederatedCoordinator` using `ICommunicationBackend<T>` for
transport
- Implement `CoordinateForwardPass()` aggregating embeddings from feature
parties
- Add `CoordinateBackwardPass()` distributing gradients to feature parties
- Handle party role assignment (active vs passive) based on label ownership
- Use barrier synchronization for computation phase coordination

Implement split-learning layers:
- Create `SplitInputLayer` accepting multiple feature sources from different
parties
- Add `SecureEmbeddingAggregator` for privacy-preserving embedding combination
- Extend `LayerBase.Forward()` to support encrypted intermediate activations
- Implement gradient splitting in backward pass for party-specific routing
- Leverage `SplitLayer.cs` patterns with VFL encryption hooks

Build VFL configuration:
- Create `VFLConfiguration` with party topology, feature ownership, and
encryption parameters
- Add `PartyRole` enumeration (active vs passive)
- Include `FeaturePartitionMap` for feature ownership
- Extend `ShardingConfiguration<T>` pattern with VFL parameters
- Validate single active party and consistent feature coverage

Phase 4: Security and Serialization

Implement encrypted communication patterns and secure state management for VFL training rounds. This phase ensures privacy guarantees during multi-party computation.

Task 1: Add Encrypted Communication Support

Extend communication backend with encryption capabilities for gradient and embedding exchange.

• Create src/PrivacyPreserving/Encryption/IEncryptionScheme.cs interface supporting additive homomorphic operations
• Implement basic Paillier-style homomorphic encryption wrapper using .NET BigInteger in System.Numerics
• Create SecureCommunicationBackend decorator wrapping ICommunicationBackend<T> with automatic encryption/decryption
• Add encryption key management during VFL initialization phase via secure channel establishment
• Support selective encryption—only gradients/embeddings encrypted, not protocol coordination messages

Task 2: Implement VFL Round State Persistence

Build checkpoint and state management for federated learning round progression and recovery.

• Create src/FederatedLearning/Vertical/VFLRoundState.cs capturing round number, participating parties, convergence metrics, and PSI result references
• Extend existing serialization patterns from src/Serialization/ with VFL-specific converters
• Implement SaveRoundState() and LoadRoundState() methods coordinated via barriers ensuring consistent checkpointing
• Add support for resuming interrupted training by reconstructing entity alignment from persisted PSI results
• Include audit logging tracking which parties contributed to each round for compliance requirements

Task 3: Add Differential Privacy Mechanisms

Introduce gradient perturbation capabilities for differential privacy guarantees in VFL scenarios.

• Create src/PrivacyPreserving/DifferentialPrivacy/GaussianMechanism.cs adding calibrated noise to gradient updates
• Implement privacy budget tracking (epsilon, delta) across VFL training rounds
• Add PrivacyAccountant class computing cumulative privacy loss using advanced composition theorems
• Integrate noise addition into VerticalFederatedCoordinator.CoordinateBackwardPass() before gradient transmission
• Provide configuration options for per-party vs global privacy budgets supporting heterogeneous privacy requirements

🤖 Prompt for AI agents

Implement security and state management for VFL training with encrypted
communication and privacy guarantees. This ensures secure multi-party
computation.

Add encrypted communication:
- Define `IEncryptionScheme` interface with additive homomorphic operations
- Implement Paillier-style encryption using `BigInteger`
- Create `SecureCommunicationBackend` decorator wrapping
`ICommunicationBackend<T>` with auto encryption/decryption
- Add encryption key management during VFL initialization
- Support selective encryption (gradients/embeddings only, not coordination
messages)

Implement round state persistence:
- Create `VFLRoundState` capturing round number, parties, metrics, and PSI
references
- Extend serialization patterns with VFL-specific converters
- Implement `SaveRoundState()` and `LoadRoundState()` with barrier coordination
- Support resuming training with entity alignment reconstruction
- Add audit logging for party contributions per round

Add differential privacy:
- Create `GaussianMechanism` for calibrated gradient noise
- Implement privacy budget tracking (epsilon, delta) across rounds
- Add `PrivacyAccountant` computing cumulative privacy loss
- Integrate noise into `CoordinateBackwardPass()` before transmission
- Support per-party and global privacy budget configurations

---

🚀 Next Steps

🤖 All AI agent prompts combined

Task: 1

Build the cryptographic PSI protocol foundation using OPRF-based approach. This
establishes secure set intersection capabilities between parties.

Create PSI protocol abstractions:
- Define `IPSIProtocol` interface with `ComputeIntersectionAsync()` method
accepting party roles
- Create `PSIProtocolBase` abstract class with validation, timeout, and error
handling
- Add `PSIConfiguration` class for protocol parameters
- Create `PSIResult` record containing intersection indices and metadata

Implement OPRF-based PSI:
- Create `OPRFBasedPSI` class implementing `IPSIProtocol`
- Use `HMACSHA256` for PRF with server-side keying
- Implement client blinding/unblinding with multiplicative group operations
- Add server PRF evaluation preserving input privacy
- Implement constant-time intersection determination

Build communication primitives:
- Create message types directory under PSI
- Define `PSIProtocolMessage` base with serialization
- Implement specific messages: `ClientBlindedSet`, `ServerPRFEvaluations`,
`ClientIntersectionRequest`
- Create `PSIProtocolCoordinator` wrapping `ICommunicationBackend<T>` for
message routing
- Add message tagging for PSI traffic multiplexing
===============================================================================

Task: 2

Integrate PSI results with data loading by creating entity ID mapping and
alignment-aware loaders. This bridges PSI output with tensor-based training
data.

Create entity ID mapping:
- Define `IEntityIdMapper` interface with `AddMapping()`, `GetRowIndex()`, and
`GetAlignedIndices()` methods
- Implement `InMemoryEntityIdMapper` using bidirectional dictionary lookup
- Add thread-safe concurrent operations
- Extend existing serialization patterns for persistence

Build PSI-aware data loaders:
- Create `AlignedDataLoader` base class accepting `IEntityIdMapper` and
`PSIResult`
- Override batch iteration to filter rows by intersection indices
- Maintain entity ID to feature matrix row correspondence
- Add `CreateFromPSI()` factory method
- Ensure compatibility with `EpisodicDataLoaderBase` patterns

Implement entity preprocessing:
- Create `EntityIdExtractor` for common ID formats (CSV, database, composite)
- Add normalization and hashing for consistent representation
- Implement Bloom filter preprocessing for cardinality reduction
- Provide validation for ID uniqueness and format consistency
===============================================================================

Task: 3

Build VFL coordination layer for multi-party collaborative training with feature
partitioning. This extends distributed training infrastructure for vertical
federated learning.

Create vertical FL coordinator:
- Create `VerticalFederatedCoordinator` using `ICommunicationBackend<T>` for
transport
- Implement `CoordinateForwardPass()` aggregating embeddings from feature
parties
- Add `CoordinateBackwardPass()` distributing gradients to feature parties
- Handle party role assignment (active vs passive) based on label ownership
- Use barrier synchronization for computation phase coordination

Implement split-learning layers:
- Create `SplitInputLayer` accepting multiple feature sources from different
parties
- Add `SecureEmbeddingAggregator` for privacy-preserving embedding combination
- Extend `LayerBase.Forward()` to support encrypted intermediate activations
- Implement gradient splitting in backward pass for party-specific routing
- Leverage `SplitLayer.cs` patterns with VFL encryption hooks

Build VFL configuration:
- Create `VFLConfiguration` with party topology, feature ownership, and
encryption parameters
- Add `PartyRole` enumeration (active vs passive)
- Include `FeaturePartitionMap` for feature ownership
- Extend `ShardingConfiguration<T>` pattern with VFL parameters
- Validate single active party and consistent feature coverage
===============================================================================

Task: 4

Implement security and state management for VFL training with encrypted
communication and privacy guarantees. This ensures secure multi-party
computation.

Add encrypted communication:
- Define `IEncryptionScheme` interface with additive homomorphic operations
- Implement Paillier-style encryption using `BigInteger`
- Create `SecureCommunicationBackend` decorator wrapping
`ICommunicationBackend<T>` with auto encryption/decryption
- Add encryption key management during VFL initialization
- Support selective encryption (gradients/embeddings only, not coordination
messages)

Implement round state persistence:
- Create `VFLRoundState` capturing round number, parties, metrics, and PSI
references
- Extend serialization patterns with VFL-specific converters
- Implement `SaveRoundState()` and `LoadRoundState()` with barrier coordination
- Support resuming training with entity alignment reconstruction
- Add audit logging for party contributions per round

Add differential privacy:
- Create `GaussianMechanism` for calibrated gradient noise
- Implement privacy budget tracking (epsilon, delta) across rounds
- Add `PrivacyAccountant` computing cumulative privacy loss
- Integrate noise into `CoordinateBackwardPass()` before transmission
- Support per-party and global privacy budget configurations

💡 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.