[FederatedLearning][Post-v1] Production-Grade Vertical Federated Learning (VFL) with Split Learning

[ooples]

Summary

Implement a complete vertical federated learning stack. VFL is fundamentally different from the existing horizontal FL framework: instead of each client having the same features for different samples, each party holds different features for the same entities. This requires split neural networks, secure gradient exchange, entity alignment (PSI from #538), label privacy protection, and handling missing features.

Depends on: #538 (PSI for entity alignment)
New facade method: ConfigureVerticalFederatedLearning() (separate from horizontal ConfigureFederatedLearning)

Motivation

The entire existing FL framework is horizontal-only -- all 16 aggregation strategies assume clients have the same model architecture and average parameters. VFL requires a fundamentally different coordination model:

• No model averaging: Instead, parties coordinate forward passes through a split neural network
• Feature partitioning: Party A has features 1-10, Party B has features 11-20, Party C holds labels
• Secure gradient exchange: Gradients flow between parties, must be protected
• Entity alignment: Before training, parties must find shared entities via PSI
• Missing features: In production, not all parties have data for all entities
• Label privacy: The label holder information must be protected from feature holders

Example: Bank (income, credit score) + Hospital (diagnoses, prescriptions) + Retailer (purchase history) jointly predict loan default -- each party keeps its data private, only model activations flow between parties.

Latest Research (2024-2026)

• VFLAIR (ICLR 2025): VFL with limited aligned samples
• VertiBench (ICLR 2024): Benchmark for VFL feature distribution diversity
• VFL with Missing Features (2024): Missing features during training AND inference
• ACM Computing Survey: VFL (2025): Comprehensive VFL survey
• VFL Structured Literature Review (2025): Open challenges and practical deployment
• VFL Certified Unlearning (2025): GDPR-compliant data removal from VFL
• VFL Primal-Dual Unlearning (2025): Sample and label unlearning in VFL
• Neural Network VFL with Server (2024): Server-integrated VFL framework
• FL Production Deployment (2026): Only 5.2% of FL research reaches production; FATE is production-ready VFL

Implementation Plan

Interfaces (4 files)

File	Purpose
src/FederatedLearning/Vertical/IVerticalFederatedTrainer.cs	VFL training orchestration: coordinate split forward/backward passes
src/FederatedLearning/Vertical/IVerticalParty.cs	Represents one party: holds local features, computes local embeddings
src/FederatedLearning/Vertical/ISplitModel.cs	Split neural network: local bottom models + shared top model
src/FederatedLearning/Vertical/ILabelProtector.cs	Protect label holder gradients from feature parties

Classes (10 files)

File	Purpose
src/FederatedLearning/Vertical/VerticalFederatedTrainer.cs	Main VFL coordinator: entity alignment -> split training -> secure gradient exchange
src/FederatedLearning/Vertical/VerticalPartyClient.cs	Feature-holding party: computes local embeddings from local features
src/FederatedLearning/Vertical/VerticalPartyLabelHolder.cs	Label-holding party: computes loss, generates top-model gradients, applies label DP
src/FederatedLearning/Vertical/SplitNeuralNetwork.cs	Split NN: each party runs bottom model, outputs concat/sum/attend, top model on coordinator
src/FederatedLearning/Vertical/VerticalDataPartitioner.cs	Partition features across parties (by column groups, by domain, random)
src/FederatedLearning/Vertical/SecureGradientExchange.cs	Encrypted gradient passing between parties using HE or secret sharing
src/FederatedLearning/Vertical/MissingFeatureHandler.cs	Handle missing feature blocks: zero imputation, mean imputation, learned imputation, skip
src/FederatedLearning/Vertical/LabelDifferentialPrivacy.cs	DP noise on label holder gradients to prevent feature parties from inferring labels
src/FederatedLearning/Vertical/VerticalFederatedUnlearner.cs	GDPR-compliant removal of entity data from VFL models (certified unlearning)
src/FederatedLearning/Vertical/VerticalFederatedBenchmark.cs	VertiBench-style benchmarking for VFL implementations

Options and Enums (8 files)

File	Purpose
src/Models/Options/VerticalFederatedLearningOptions.cs	Top-level VFL config: party definitions, PSI options, split model config
src/Models/Options/SplitModelOptions.cs	Cut layer selection (manual, auto-optimal, balanced-compute), embedding dimension
src/Models/Options/MissingFeatureOptions.cs	Imputation strategy, alignment threshold, minimum overlap ratio
src/Models/Options/VflUnlearningOptions.cs	Unlearning method, certification level, verification
src/Models/Options/VflAggregationMode.cs	Enum: Concatenation, Sum, Attention, Gating
src/Models/Options/MissingFeatureStrategy.cs	Enum: Zero, Mean, Learned, Skip
src/Models/Options/SplitPointStrategy.cs	Enum: Manual, AutoOptimal, BalancedCompute
src/Models/Options/VflUnlearningMethod.cs	Enum: Retraining, GradientAscent, PrimalDual, Certified

Facade Integration

New facade method on AiModelBuilder (VFL is a different paradigm from horizontal FL):

// New method on AiModelBuilder (new partial class file: AiModelBuilder.VerticalFL.cs)
public IAiModelBuilder<T, TInput, TOutput> ConfigureVerticalFederatedLearning(
    VerticalFederatedLearningOptions options,
    IVerticalParty<T>[] parties = null,
    ILabelProtector<T> labelProtector = null)

Usage:

builder.ConfigureVerticalFederatedLearning(new VerticalFederatedLearningOptions
{
    EntityAlignment = new PsiOptions { Protocol = PsiProtocol.ObliviousTransfer },
    SplitModel = new SplitModelOptions
    {
        AggregationMode = VflAggregationMode.Concatenation,
        SplitPoint = SplitPointStrategy.AutoOptimal,
        EmbeddingDimension = 64
    },
    MissingFeatures = new MissingFeatureOptions
    {
        Strategy = MissingFeatureStrategy.Mean,
        MinimumOverlapRatio = 0.5
    },
    Unlearning = new VflUnlearningOptions
    {
        Method = VflUnlearningMethod.Certified,
        Enabled = false
    }
});

Also needs:

• New backing fields in AiModelBuilder (in a new partial class file)
• New method signatures in IAiModelBuilder interface
• New VFL training path in Build() method (separate from horizontal FL path)

Acceptance Criteria

• IVerticalFederatedTrainer coordinating split forward/backward passes
• IVerticalParty representing feature-holding and label-holding parties
• ISplitModel split neural network with local bottom + shared top
• VerticalFederatedTrainer full orchestration (PSI -> train -> evaluate)
• SecureGradientExchange with HE or SS protection
• MissingFeatureHandler with zero, mean, learned, and skip strategies
• LabelDifferentialPrivacy protecting label holder
• VerticalFederatedUnlearner for GDPR compliance
• VerticalDataPartitioner for feature distribution
• ConfigureVerticalFederatedLearning() method on AiModelBuilder
• IAiModelBuilder interface updated with new method signature
• PSI integration from [FederatedLearning][Post-v1] Private Set Intersection (PSI) for Vertical FL Entity Alignment #538 for entity alignment
• Builds on both net10.0 and net471
• Unit tests with synthetic vertically-partitioned data
• XML docs with beginners sections

Estimated: ~22 files

[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

• [FederatedLearning][Post-v1] TEE-Assisted Confidential Computing for Federated Learning #537
• [FederatedLearning][Post-v1] Master tracking issue for FL v2 enhancements #536
• [Phase 3] Implement Federated Learning Framework #398

👤 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 a production-grade Vertical Federated Learning (VFL) system that enables multiple parties to collaboratively train machine learning models on vertically partitioned data (where each party holds different features for the same samples) while preserving privacy through cryptographic protocols.

Observations

The codebase contains a mature distributed training framework (src/DistributedTraining/) with communication backends, collective operations, and synchronization patterns designed for centralized data parallelism. It includes neural network abstractions with layer-based composition, comprehensive model compression and quantization capabilities, deployment infrastructure with runtime management and telemetry, and feature selection utilities. However, it lacks any federated learning implementation, privacy-preserving protocols (differential privacy, homomorphic encryption, secure multi-party computation), or multi-party orchestration mechanisms. VFL will be built as a new system that selectively leverages the existing neural network architecture and communication patterns while introducing entirely new privacy and coordination layers.

Assumptions

Assumption 1: Privacy Protocol Selection

Options Considered:

• Full homomorphic encryption (FHE) for maximum security with high computational overhead
• Partially homomorphic encryption (Paillier/SEAL) with secure multi-party computation for balanced security and performance
• Differential privacy only with no cryptographic guarantees
• Trusted execution environments (TEE) like Intel SGX

Chosen Option: Partially homomorphic encryption (Paillier/SEAL) with secure multi-party computation for balanced security and performance

Rationale: Use Microsoft SEAL for partially homomorphic encryption combined with custom secure multi-party computation protocols and differential privacy. This provides strong cryptographic guarantees while remaining practical for production systems. SEAL is production-grade, actively maintained, has .NET bindings, and allows efficient gradient aggregation. SMPC and DP add complementary layers of protection.

Assumption 2: Coordinator Architecture

Options Considered:

• Centralized dedicated coordinator (third-party server manages all parties)
• Active party as coordinator (label-holding party orchestrates training)
• Fully decentralized peer-to-peer coordination
• Hybrid approach with optional coordinator mode

Chosen Option: Active party as coordinator (label-holding party orchestrates training)

Rationale: Implement active party as coordinator with optional dedicated coordinator mode. This aligns with standard VFL literature where the active party (typically the label holder) naturally coordinates training. Supporting a dedicated coordinator mode provides flexibility for scenarios where no party should have orchestration privileges.

Assumption 3: Integration with Existing Distributed Training

Options Considered:

• Extend ShardedModel and IShardedOptimizer to support feature sharding
• Create entirely new VFL namespace separate from distributed training
• Use composition to wrap existing distributed training with VFL protocols

Chosen Option: Create entirely new VFL namespace separate from distributed training

Rationale: Create new src/VerticalFederatedLearning/ namespace with dedicated abstractions. VFL has fundamentally different semantics (feature partitioning vs. data partitioning, multi-party trust model, split learning architecture) that don't map cleanly to centralized distributed training patterns. Clean separation prevents force-fitting and maintains architectural clarity.

💡 User Tips

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

Plan

Phase 1: Core VFL Abstractions and Split Learning Architecture

Establish the foundational abstractions for vertical federated learning and implement the split neural network architecture that enables parties to train on partitioned features. This phase creates the domain model and basic training mechanics without privacy protocols.

Task 1: Define VFL Core Interfaces and Domain Model

Create the fundamental abstractions that define parties, training sessions, and the split learning contract.

• Define IVFLParty<T> interface representing a training participant with methods for computing embeddings, receiving gradients, and managing local model state
• Create VFLPartyRole enum distinguishing between active party (label holder), passive parties (feature holders), and optional coordinator
• Implement VFLTrainingSession<T> class managing the lifecycle of a multi-party training session with participant registration and session state
• Define IBottomModel<T> and ITopModel<T> interfaces for split learning, where bottom models compute feature embeddings and top models aggregate embeddings for final prediction
• Create VFLModelConfig<T> class specifying feature assignments per party, model architecture for each party's bottom model, and top model architecture
• Add interfaces in src/Interfaces/ for public contracts and create new directory src/VerticalFederatedLearning/Core/ for VFL-specific implementations

Task 2: Implement Split Neural Network Architecture

Build the neural network components that enable split learning across parties.

• Create BottomModel<T> class in src/VerticalFederatedLearning/Models/ that wraps existing NeuralNetwork to compute embeddings from assigned features, with methods for forward pass to embedding space and backward pass receiving top model gradients
• Implement TopModel<T> class aggregating embeddings from all bottom models using concatenation or learned aggregation layers, computing final predictions and loss, and backpropagating gradients to bottom models
• Create SplitNeuralNetwork<T> coordinating the complete split learning workflow: orchestrating forward passes across all parties, aggregating embeddings at top model, computing loss and backpropagating gradients, and managing synchronized parameter updates
• Extend existing LayerBase<T> in src/NeuralNetworks/Layers/ to add EmbeddingAggregationLayer<T> supporting concatenation, weighted sum, and attention-based aggregation strategies
• Implement partial forward pass support in split models that stops at embedding layer rather than computing final prediction

Task 3: Create Feature Partitioning Utilities

Develop utilities for partitioning features across parties and managing feature-based data distribution.

• Create FeaturePartitioner<T> class in src/VerticalFederatedLearning/Data/ with methods for partitioning feature indices across parties based on configuration, validating partitions for completeness and non-overlap, and generating party-specific feature masks
• Implement PartitionStrategy enum supporting manual assignment, random partitioning, correlation-based grouping, and balanced feature distribution
• Extend InMemoryDataLoader in src/Data/Loaders/ to support column-based access patterns with methods like GetFeatureSubset(int[] featureIndices) and CreatePartitionedView(FeaturePartition)
• Create VFLDataLoader<T> that wraps existing data loaders and provides party-specific views of data with only assigned features visible
• Add validation utilities ensuring all parties have the same sample IDs and sample counts match across parties

Task 4: Implement Basic Gradient Exchange Protocol

Create the communication layer for exchanging embeddings and gradients between parties without privacy protocols initially.

• Define IVFLCommunicationProtocol<T> interface in src/Interfaces/ with methods for sending embeddings to coordinator, broadcasting gradients to parties, and synchronizing training state
• Implement PlainTextVFLProtocol<T> in src/VerticalFederatedLearning/Communication/ using existing ICommunicationBackend for message passing but adapted for party-to-party rather than rank-to-rank communication
• Create EmbeddingMessage<T> and GradientMessage<T> classes for structured communication including sender identification, target recipient, and payload data
• Add message routing logic in protocol implementation that handles active party as hub receiving embeddings and broadcasting gradients
• Implement synchronization barriers ensuring all parties complete forward/backward passes before proceeding

🤖 Prompt for AI agents

Implement foundational VFL abstractions and split learning architecture. Create
the domain model for vertical federated learning with multi-party training.

Create core interfaces and enums:
- `IVFLParty<T>` interface with methods for computing embeddings, receiving
gradients, and managing local model state
- `VFLPartyRole` enum with values: ActiveParty, PassiveParty, Coordinator
- `IBottomModel<T>` and `ITopModel<T>` interfaces for split learning components
- `IVFLCommunicationProtocol<T>` interface for embedding/gradient exchange

Implement core classes:
- `VFLTrainingSession<T>` managing multi-party session lifecycle with
participant registration
- `VFLModelConfig<T>` specifying feature assignments and model architectures per
party
- `BottomModel<T>` wrapping existing NeuralNetwork for computing feature
embeddings
- `TopModel<T>` aggregating embeddings and computing final predictions
- `SplitNeuralNetwork<T>` orchestrating the complete split learning workflow
- `EmbeddingAggregationLayer<T>` extending LayerBase with concatenation,
weighted sum, and attention strategies

Create data partitioning utilities:
- `FeaturePartitioner<T>` with partition validation and feature mask generation
- `PartitionStrategy` enum: Manual, Random, CorrelationBased, Balanced
- `VFLDataLoader<T>` providing party-specific feature views
- Extend `InMemoryDataLoader` with `GetFeatureSubset()` and
`CreatePartitionedView()` methods

Implement basic communication protocol:
- `PlainTextVFLProtocol<T>` adapting existing ICommunicationBackend for
party-to-party messaging
- `EmbeddingMessage<T>` and `GradientMessage<T>` with sender/recipient/payload
- Message routing with active party as hub
- Synchronization barriers for forward/backward pass coordination

Add directories: `src/VerticalFederatedLearning/Core/`,
`src/VerticalFederatedLearning/Models/`, `src/VerticalFederatedLearning/Data/`,
`src/VerticalFederatedLearning/Communication/`

Phase 2: Privacy-Preserving Protocols and Secure Aggregation

Add cryptographic protocols that protect sensitive data during VFL training, including entity alignment, homomorphic encryption, and secure multi-party computation.

Task 1: Implement Private Set Intersection (PSI) for Entity Alignment

Build the entity alignment component that identifies common samples across parties without revealing non-overlapping samples.

• Create PrivateSetIntersection<T> class in src/VerticalFederatedLearning/Privacy/PSI/ implementing Diffie-Hellman-based PSI protocol for two-party case and extending to multi-party PSI using coordinator-mediated approach
• Implement EntityAlignmentService<T> orchestrating PSI across all VFL parties, collecting encrypted sample identifiers, computing intersections, and distributing aligned sample indices
• Add SampleIdentifier class supporting multiple ID types (hash, encrypted ID, pseudonymized ID) and collision detection
• Create alignment result caching to avoid re-running PSI on every training session with cache invalidation when party data changes
• Implement alignment verification protocol ensuring all parties see the same aligned sample set

Task 2: Integrate Microsoft SEAL for Homomorphic Encryption

Add homomorphic encryption capabilities for secure gradient aggregation using Microsoft SEAL library.

• Add Microsoft SEAL NuGet package reference to project and create wrapper classes in src/VerticalFederatedLearning/Privacy/HomomorphicEncryption/
• Implement SEALCryptoProvider<T> managing SEAL contexts, encryption parameters (polynomial modulus, coefficient modulus), and key generation/distribution for parties
• Create EncryptedTensor<T> class wrapping SEAL ciphertexts with tensor shape metadata and supporting addition operations on encrypted tensors (used for gradient aggregation)
• Implement HomomorphicGradientAggregator<T> that encrypts gradients at each party, aggregates encrypted gradients using homomorphic addition, and decrypts aggregated result only at coordinator
• Add batching support for encrypting/decrypting large gradient tensors to manage memory and computational overhead

Task 3: Implement Secure Multi-Party Computation Protocols

Create SMPC protocols for scenarios where HE overhead is prohibitive, using additive secret sharing.

• Create SecretSharing<T> class in src/VerticalFederatedLearning/Privacy/SMPC/ implementing additive secret sharing scheme that splits values into random shares summing to original value
• Implement ThreePartyComputation<T> protocol supporting secure multiplication and comparison operations using three-party computation patterns from literature
• Create SecureBeaverTriples<T> service generating and distributing pre-computed Beaver triples for efficient multiplication
• Implement SMPCGradientAggregator<T> as alternative to HE-based aggregation, splitting gradients into shares, computing aggregation on shares, and reconstructing results
• Add protocol selection logic in IVFLCommunicationProtocol allowing runtime choice between HE and SMPC based on performance requirements

Task 4: Add Differential Privacy Mechanisms

Implement differential privacy to protect against gradient leakage attacks and provide provable privacy guarantees.

• Create DifferentialPrivacy<T> class in src/VerticalFederatedLearning/Privacy/DP/ implementing Gaussian mechanism for adding calibrated noise to gradients and tracking privacy budget (epsilon, delta)
• Implement PrivacyAccountant<T> tracking cumulative privacy loss across training iterations using advanced composition theorems and providing privacy budget warnings
• Add DPConfig class specifying noise multiplier, clipping threshold for gradient clipping, and target epsilon/delta values
• Implement gradient clipping in bottom models before adding noise to bound sensitivity and ensure valid DP guarantees
• Create DP-aware learning rate adjustment that compensates for noise injection while maintaining convergence

Task 5: Create Secure Communication Protocol Integration

Combine all privacy mechanisms into a unified secure VFL protocol.

• Create SecureVFLProtocol<T> in src/VerticalFederatedLearning/Communication/ composing PSI, HE/SMPC, and DP mechanisms into cohesive protocol
• Implement protocol negotiation allowing parties to agree on security parameters (HE parameters, DP budget, SMPC vs. HE)
• Add SecurityLevel enum (NoPrivacy, DPOnly, HEOnly, SMPC, FullStack) controlling which privacy mechanisms are active
• Implement secure channel establishment using TLS for transport security on top of cryptographic protocols
• Create security audit logging tracking all cryptographic operations, privacy budget consumption, and protocol violations

🤖 Prompt for AI agents

Implement privacy-preserving cryptographic protocols for secure VFL training.
Add entity alignment, homomorphic encryption, SMPC, and differential privacy
mechanisms.

Implement Private Set Intersection:
- `PrivateSetIntersection<T>` with Diffie-Hellman PSI for two-party and
coordinator-mediated multi-party
- `EntityAlignmentService<T>` orchestrating PSI across parties
- `SampleIdentifier` class supporting hash/encrypted/pseudonymized IDs with
collision detection
- Alignment result caching with invalidation logic
- Alignment verification protocol

Integrate Microsoft SEAL:
- Add Microsoft.SEAL NuGet package
- `SEALCryptoProvider<T>` managing contexts, parameters, and key distribution
- `EncryptedTensor<T>` wrapping SEAL ciphertexts with tensor shape metadata
- `HomomorphicGradientAggregator<T>` for encrypted gradient aggregation
- Batching support for large tensor encryption/decryption

Implement SMPC protocols:
- `SecretSharing<T>` with additive secret sharing scheme
- `ThreePartyComputation<T>` supporting secure multiplication and comparison
- `SecureBeaverTriples<T>` service for pre-computed triples
- `SMPCGradientAggregator<T>` as HE alternative
- Protocol selection logic for runtime HE vs SMPC choice

Add Differential Privacy:
- `DifferentialPrivacy<T>` with Gaussian mechanism and privacy budget tracking
- `PrivacyAccountant<T>` tracking cumulative privacy loss with advanced
composition
- `DPConfig` class with noise multiplier, clipping threshold, epsilon/delta
- Gradient clipping in bottom models
- DP-aware learning rate adjustment

Create unified secure protocol:
- `SecureVFLProtocol<T>` composing PSI, HE/SMPC, and DP
- Protocol negotiation for security parameter agreement
- `SecurityLevel` enum: NoPrivacy, DPOnly, HEOnly, SMPC, FullStack
- TLS secure channel establishment
- Security audit logging

Add directories: `src/VerticalFederatedLearning/Privacy/PSI/`,
`src/VerticalFederatedLearning/Privacy/HomomorphicEncryption/`,
`src/VerticalFederatedLearning/Privacy/SMPC/`,
`src/VerticalFederatedLearning/Privacy/DP/`

Phase 3: VFL Coordinator and Multi-Party Orchestration

Build the coordinator service that orchestrates training across parties, manages session state, and handles fault tolerance.

Task 1: Implement VFL Coordinator Service

Create the central orchestration service managing the complete VFL training lifecycle.

• Create VFLCoordinatorService<T> in src/VerticalFederatedLearning/Coordinator/ managing party registration, session initialization, and training round coordination
• Implement PartyRegistry<T> maintaining registry of connected parties with their capabilities (features, model architecture, supported protocols) and connection status
• Add VFLTrainingState<T> class tracking current training epoch, aggregated loss metrics, and party-specific states (current round, pending operations)
• Implement round coordination logic that broadcasts training start signals, collects embeddings from all passive parties, aggregates at top model, distributes gradients, and synchronizes parameter updates
• Create coordinator discovery protocol allowing parties to locate coordinator endpoint through configuration or service discovery

Task 2: Build Party-to-Coordinator Communication

Develop the communication infrastructure for parties to interact with the coordinator.

• Define IVFLPartyClient<T> interface in src/Interfaces/ for party-side operations: registering with coordinator, receiving training instructions, and uploading embeddings/gradients
• Implement VFLPartyClient<T> in src/VerticalFederatedLearning/Communication/ using HTTP/gRPC for reliable messaging with retry logic and timeout handling
• Create message protocol supporting RegisterParty, StartTrainingRound, UploadEmbeddings, ReceiveGradients, ReportProgress, LeaveSession operations
• Add connection management with heartbeat mechanism, automatic reconnection on failure, and graceful degradation when parties disconnect
• Implement request queuing and flow control preventing coordinator overload when many parties communicate simultaneously

Task 3: Implement Training Session Management

Create session lifecycle management handling multi-party training from initialization through completion.

• Create SessionManager<T> in src/VerticalFederatedLearning/Coordinator/ managing concurrent training sessions with unique session IDs, participant lists per session, and session lifecycle states
• Implement session initialization protocol that performs entity alignment (PSI), negotiates security parameters, and validates party configurations before training starts
• Add training loop orchestration that coordinates multiple epochs, batched training within epochs, and validation rounds between epochs
• Create session teardown logic that finalizes model state, distributes final models to parties, and releases resources
• Implement session recovery mechanisms allowing sessions to resume after coordinator restart using persistent session state

Task 4: Add Fault Tolerance and Checkpointing

Build reliability features ensuring training can recover from party or coordinator failures.

• Extend existing ICheckpointableModel in src/Interfaces/ with IVFLCheckpointable interface supporting multi-party checkpoint coordination
• Create VFLCheckpointManager<T> in src/VerticalFederatedLearning/Coordinator/ coordinating checkpoint creation across all parties, storing coordinator state and aggregated metrics, and managing checkpoint versioning
• Implement party failure detection through heartbeat monitoring and training timeout detection
• Add training continuation logic that reloads checkpoint state, reconnects available parties, and either waits for failed parties or excludes them based on policy
• Create checkpoint synchronization protocol ensuring all parties save checkpoints at same training iteration

Task 5: Implement Asynchronous Training Support

Add asynchronous VFL training mode where parties can train at different speeds.

• Create AsyncVFLCoordinator<T> in src/VerticalFederatedLearning/Coordinator/ supporting asynchronous gradient aggregation where coordinator aggregates gradients as they arrive without waiting for all parties
• Implement staleness-aware aggregation weighting recent gradients higher than stale gradients from slow parties
• Add dynamic party selection allowing coordinator to choose subset of fastest parties per round while ensuring all parties participate over time
• Create fairness metrics tracking participation rate per party and adjusting selection probability to prevent party starvation
• Implement convergence detection for asynchronous mode using gradient norm monitoring and validation loss tracking

🤖 Prompt for AI agents

Build coordinator service for multi-party orchestration, session management, and
fault tolerance. Implement synchronous and asynchronous training coordination.

Implement coordinator service:
- `VFLCoordinatorService<T>` managing party registration, session
initialization, and training rounds
- `PartyRegistry<T>` tracking connected parties with capabilities and connection
status
- `VFLTrainingState<T>` tracking epochs, loss metrics, and party-specific states
- Round coordination logic: broadcast signals, collect embeddings, aggregate,
distribute gradients
- Coordinator discovery protocol

Build party-coordinator communication:
- `IVFLPartyClient<T>` interface for party-side operations
- `VFLPartyClient<T>` using HTTP/gRPC with retry logic and timeout handling
- Message protocol: RegisterParty, StartTrainingRound, UploadEmbeddings,
ReceiveGradients, ReportProgress, LeaveSession
- Connection management with heartbeat, reconnection, and graceful degradation
- Request queuing and flow control

Implement session management:
- `SessionManager<T>` managing concurrent sessions with IDs and lifecycle states
- Session initialization with PSI, security negotiation, and validation
- Training loop orchestration for epochs, batches, and validation rounds
- Session teardown with model finalization and resource release
- Session recovery using persistent state

Add fault tolerance:
- Extend `ICheckpointableModel` with `IVFLCheckpointable` interface
- `VFLCheckpointManager<T>` coordinating multi-party checkpoints
- Party failure detection via heartbeat and timeout monitoring
- Training continuation logic with checkpoint reload and party reconnection
- Checkpoint synchronization protocol

Implement asynchronous training:
- `AsyncVFLCoordinator<T>` with asynchronous gradient aggregation
- Staleness-aware aggregation weighting
- Dynamic party selection ensuring fair participation
- Fairness metrics and selection probability adjustment
- Convergence detection via gradient norm and validation loss

Add directory: `src/VerticalFederatedLearning/Coordinator/`

Phase 4: Feature Engineering and VFL Data Pipeline

Develop data handling utilities specific to VFL scenarios, including columnar operations and feature-based partitioning strategies.

Task 1: Create VFL-Specific Data Loaders

Build data loaders that handle columnar access patterns and party-specific feature views.

• Create VFLDataLoaderBase<T> in src/Data/Loaders/ extending existing DataLoaderBase<T> with columnar access methods and feature partition awareness
• Implement ColumnarInMemoryDataLoader<T> storing data in column-major format for efficient feature-subset extraction
• Add VFLBatchGenerator<T> coordinating batch creation across parties ensuring consistent sample ordering, synchronized batch boundaries, and aligned feature extraction
• Create StreamingVFLDataLoader<T> for large datasets that cannot fit in memory, streaming data with coordinated buffering across parties
• Implement data validation utilities ensuring feature consistency (matching samples across parties) and detecting data quality issues (missing values in features)

Task 2: Implement Feature Partitioning Strategies

Develop intelligent strategies for distributing features across parties beyond manual assignment.

• Create AutomaticFeaturePartitioner<T> in src/VerticalFederatedLearning/Data/ implementing correlation-based clustering that groups correlated features, balanced partitioning distributing features evenly across parties, and information-theoretic partitioning maximizing information gain per party
• Implement FeatureImportancePartitioner<T> using existing feature selectors from src/FeatureSelectors/ to assign important features strategically
• Add PrivacyAwarePartitioner<T> considering feature sensitivity (personally identifiable information) and regulatory constraints when assigning features
• Create partitioning validation checking feature overlap detection and ensuring minimum feature count per party
• Implement partition export/import supporting JSON configuration files for reproducible partitioning

Task 3: Integrate with Existing Feature Selection

Extend existing feature selection utilities to work with VFL's distributed feature sets.

• Modify FeatureSelectorBase<T> in src/FeatureSelectors/ to support distributed feature selection where feature importance is computed locally at each party and aggregated at coordinator
• Create VFLFeatureSelector<T> wrapper coordinating feature selection across parties and handling communication of selection results
• Implement privacy-preserving feature importance aggregation using secure aggregation protocols
• Add party-aware feature elimination allowing recursive feature elimination to work across party boundaries
• Create feature selection result serialization for sharing selected features across parties

Task 4: Build Sample Alignment Pipeline

Create utilities for aligning samples across parties before training begins.

• Create SampleAlignmentPipeline<T> in src/VerticalFederatedLearning/Data/ orchestrating complete alignment workflow from sample ID collection through aligned dataset creation
• Implement SampleIDResolver<T> supporting multiple ID formats (integer keys, string IDs, composite keys) and ID normalization/hashing
• Add alignment quality metrics tracking overlap percentage, duplicate detection, and alignment confidence scores
• Create aligned dataset export generating party-specific datasets with aligned indices and supporting incremental alignment when new samples arrive
• Implement alignment caching and versioning to avoid re-aligning static datasets

🤖 Prompt for AI agents

Develop VFL-specific data handling with columnar loaders, intelligent feature
partitioning, and sample alignment utilities.

Create VFL data loaders:
- `VFLDataLoaderBase<T>` extending DataLoaderBase with columnar access and
feature partition awareness
- `ColumnarInMemoryDataLoader<T>` with column-major storage for efficient
feature-subset extraction
- `VFLBatchGenerator<T>` coordinating batch creation with consistent ordering
and synchronized boundaries
- `StreamingVFLDataLoader<T>` for large datasets with coordinated buffering
- Data validation utilities for feature consistency and quality checks

Implement feature partitioning strategies:
- `AutomaticFeaturePartitioner<T>` with correlation-based clustering, balanced
partitioning, and information-theoretic partitioning
- `FeatureImportancePartitioner<T>` using existing feature selectors for
strategic assignment
- `PrivacyAwarePartitioner<T>` considering feature sensitivity and regulatory
constraints
- Partitioning validation for overlap detection and minimum feature counts
- Partition export/import with JSON configuration support

Integrate with feature selection:
- Modify `FeatureSelectorBase<T>` to support distributed feature selection
- `VFLFeatureSelector<T>` wrapper coordinating selection across parties
- Privacy-preserving feature importance aggregation
- Party-aware recursive feature elimination
- Feature selection result serialization

Build sample alignment pipeline:
- `SampleAlignmentPipeline<T>` orchestrating complete alignment workflow
- `SampleIDResolver<T>` supporting multiple ID formats with
normalization/hashing
- Alignment quality metrics: overlap percentage, duplicates, confidence scores
- Aligned dataset export with incremental alignment support
- Alignment caching and versioning

Extend `src/Data/Loaders/` and `src/VerticalFederatedLearning/Data/`, modify
`src/FeatureSelectors/FeatureSelectorBase<T>`

Phase 5: Production Features and Performance Optimization

Add production-grade features including compression, monitoring, and integration with existing deployment infrastructure.

Task 1: Implement Gradient and Embedding Compression

Apply model compression techniques to reduce communication overhead in VFL.

• Extend existing compression utilities from src/ModelCompression/ to work with VFL embeddings and gradients
• Create VFLCompressionPipeline<T> in src/VerticalFederatedLearning/Optimization/ applying gradient quantization, sparsification, and Huffman encoding
• Implement EmbeddingCompressor<T> using product quantization or dimensionality reduction to compress embeddings before transmission
• Add adaptive compression adjusting compression level based on network bandwidth and compute resources
• Create decompression utilities ensuring lossless reconstruction when needed for security protocols

Task 2: Add VFL Telemetry and Monitoring

Integrate VFL training metrics with existing deployment telemetry system.

• Extend TelemetryCollector in src/Deployment/Runtime/ with VFL-specific metrics including per-party training time, communication overhead, and gradient staleness
• Create VFLMetricsCollector<T> in src/VerticalFederatedLearning/Monitoring/ tracking privacy budget consumption, entity alignment quality, and cryptographic operation overhead
• Implement real-time dashboard metrics exposing party availability, training progress, and bottleneck detection
• Add alerting for privacy violations, party failures, and convergence issues
• Create metrics export supporting Prometheus, statsd, or custom telemetry backends

Task 3: Integrate with Existing Deployment Runtime

Connect VFL training with the existing deployment and serving infrastructure.

• Extend DeploymentRuntime<T> in src/Deployment/Runtime/ to support VFL model deployment with party-specific model versioning and coordinated deployment across parties
• Create VFLModelExporter<T> in src/VerticalFederatedLearning/Deployment/ exporting bottom and top models separately with encryption support for model protection
• Implement VFL inference pipeline coordinating prediction across parties at serving time
• Add A/B testing support for VFL models comparing different privacy configurations or party compositions
• Create model registry tracking VFL model lineage and party contributions

Task 4: Build Configuration and Management Tools

Develop utilities for configuring and managing VFL deployments.

• Create VFLConfigurationBuilder<T> in src/VerticalFederatedLearning/Configuration/ providing fluent API for configuring parties, privacy settings, and training parameters
• Implement configuration validation ensuring security parameters are compatible and party configurations are consistent
• Add configuration templates for common scenarios (two-party VFL, multi-party with coordinator, privacy-first configuration)
• Create CLI tools for VFL management including session start/stop, party registration, and status monitoring
• Implement configuration export/import supporting YAML/JSON formats for reproducible deployments

Task 5: Add VFL-Specific Documentation and Examples

Create comprehensive documentation and example implementations for VFL usage.

• Create docs/VerticalFederatedLearning.md documenting VFL architecture, usage patterns, privacy guarantees, and best practices
• Add code examples in examples/VFL/ demonstrating two-party VFL, multi-party with coordinator, privacy protocol configuration, and production deployment
• Create API documentation for all VFL interfaces and public classes
• Add troubleshooting guide covering common issues (alignment failures, privacy budget exhaustion, party disconnection)
• Create performance tuning guide documenting compression strategies, protocol selection, and optimization techniques

🤖 Prompt for AI agents

Add production-grade features: compression, telemetry, deployment integration,
configuration tools, and comprehensive documentation.

Implement compression:
- Extend existing compression utilities from `src/ModelCompression/` for VFL
embeddings and gradients
- `VFLCompressionPipeline<T>` with quantization, sparsification, and Huffman
encoding
- `EmbeddingCompressor<T>` using product quantization or dimensionality
reduction
- Adaptive compression based on bandwidth and compute resources
- Decompression utilities with lossless reconstruction support

Add telemetry and monitoring:
- Extend `TelemetryCollector` with VFL-specific metrics: per-party training
time, communication overhead, gradient staleness
- `VFLMetricsCollector<T>` tracking privacy budget, alignment quality, and
crypto overhead
- Real-time dashboard metrics for party availability, progress, and bottlenecks
- Alerting for privacy violations, failures, and convergence issues
- Metrics export for Prometheus, statsd, or custom backends

Integrate with deployment:
- Extend `DeploymentRuntime<T>` for VFL model deployment with party-specific
versioning
- `VFLModelExporter<T>` exporting bottom/top models separately with encryption
- VFL inference pipeline coordinating multi-party predictions
- A/B testing for privacy configurations and party compositions
- Model registry tracking lineage and party contributions

Build configuration tools:
- `VFLConfigurationBuilder<T>` with fluent API for parties, privacy, and
training parameters
- Configuration validation for compatible security parameters
- Configuration templates: two-party, multi-party with coordinator,
privacy-first
- CLI tools for session management, party registration, and monitoring
- Configuration export/import with YAML/JSON support

Create documentation:
- `docs/VerticalFederatedLearning.md` with architecture, usage patterns, privacy
guarantees, and best practices
- Code examples in `examples/VFL/`: two-party, multi-party, privacy
configuration, production deployment
- API documentation for all VFL interfaces and public classes
- Troubleshooting guide for alignment failures, budget exhaustion, party
disconnection
- Performance tuning guide for compression, protocol selection, and optimization

Add directories: `src/VerticalFederatedLearning/Optimization/`,
`src/VerticalFederatedLearning/Monitoring/`,
`src/VerticalFederatedLearning/Deployment/`,
`src/VerticalFederatedLearning/Configuration/`, `examples/VFL/`

---

🚀 Next Steps

🤖 All AI agent prompts combined

Task: 1

Implement foundational VFL abstractions and split learning architecture. Create
the domain model for vertical federated learning with multi-party training.

Create core interfaces and enums:
- `IVFLParty<T>` interface with methods for computing embeddings, receiving
gradients, and managing local model state
- `VFLPartyRole` enum with values: ActiveParty, PassiveParty, Coordinator
- `IBottomModel<T>` and `ITopModel<T>` interfaces for split learning components
- `IVFLCommunicationProtocol<T>` interface for embedding/gradient exchange

Implement core classes:
- `VFLTrainingSession<T>` managing multi-party session lifecycle with
participant registration
- `VFLModelConfig<T>` specifying feature assignments and model architectures per
party
- `BottomModel<T>` wrapping existing NeuralNetwork for computing feature
embeddings
- `TopModel<T>` aggregating embeddings and computing final predictions
- `SplitNeuralNetwork<T>` orchestrating the complete split learning workflow
- `EmbeddingAggregationLayer<T>` extending LayerBase with concatenation,
weighted sum, and attention strategies

Create data partitioning utilities:
- `FeaturePartitioner<T>` with partition validation and feature mask generation
- `PartitionStrategy` enum: Manual, Random, CorrelationBased, Balanced
- `VFLDataLoader<T>` providing party-specific feature views
- Extend `InMemoryDataLoader` with `GetFeatureSubset()` and
`CreatePartitionedView()` methods

Implement basic communication protocol:
- `PlainTextVFLProtocol<T>` adapting existing ICommunicationBackend for
party-to-party messaging
- `EmbeddingMessage<T>` and `GradientMessage<T>` with sender/recipient/payload
- Message routing with active party as hub
- Synchronization barriers for forward/backward pass coordination

Add directories: `src/VerticalFederatedLearning/Core/`,
`src/VerticalFederatedLearning/Models/`, `src/VerticalFederatedLearning/Data/`,
`src/VerticalFederatedLearning/Communication/`
===============================================================================

Task: 2

Implement privacy-preserving cryptographic protocols for secure VFL training.
Add entity alignment, homomorphic encryption, SMPC, and differential privacy
mechanisms.

Implement Private Set Intersection:
- `PrivateSetIntersection<T>` with Diffie-Hellman PSI for two-party and
coordinator-mediated multi-party
- `EntityAlignmentService<T>` orchestrating PSI across parties
- `SampleIdentifier` class supporting hash/encrypted/pseudonymized IDs with
collision detection
- Alignment result caching with invalidation logic
- Alignment verification protocol

Integrate Microsoft SEAL:
- Add Microsoft.SEAL NuGet package
- `SEALCryptoProvider<T>` managing contexts, parameters, and key distribution
- `EncryptedTensor<T>` wrapping SEAL ciphertexts with tensor shape metadata
- `HomomorphicGradientAggregator<T>` for encrypted gradient aggregation
- Batching support for large tensor encryption/decryption

Implement SMPC protocols:
- `SecretSharing<T>` with additive secret sharing scheme
- `ThreePartyComputation<T>` supporting secure multiplication and comparison
- `SecureBeaverTriples<T>` service for pre-computed triples
- `SMPCGradientAggregator<T>` as HE alternative
- Protocol selection logic for runtime HE vs SMPC choice

Add Differential Privacy:
- `DifferentialPrivacy<T>` with Gaussian mechanism and privacy budget tracking
- `PrivacyAccountant<T>` tracking cumulative privacy loss with advanced
composition
- `DPConfig` class with noise multiplier, clipping threshold, epsilon/delta
- Gradient clipping in bottom models
- DP-aware learning rate adjustment

Create unified secure protocol:
- `SecureVFLProtocol<T>` composing PSI, HE/SMPC, and DP
- Protocol negotiation for security parameter agreement
- `SecurityLevel` enum: NoPrivacy, DPOnly, HEOnly, SMPC, FullStack
- TLS secure channel establishment
- Security audit logging

Add directories: `src/VerticalFederatedLearning/Privacy/PSI/`,
`src/VerticalFederatedLearning/Privacy/HomomorphicEncryption/`,
`src/VerticalFederatedLearning/Privacy/SMPC/`,
`src/VerticalFederatedLearning/Privacy/DP/`
===============================================================================

Task: 3

Build coordinator service for multi-party orchestration, session management, and
fault tolerance. Implement synchronous and asynchronous training coordination.

Implement coordinator service:
- `VFLCoordinatorService<T>` managing party registration, session
initialization, and training rounds
- `PartyRegistry<T>` tracking connected parties with capabilities and connection
status
- `VFLTrainingState<T>` tracking epochs, loss metrics, and party-specific states
- Round coordination logic: broadcast signals, collect embeddings, aggregate,
distribute gradients
- Coordinator discovery protocol

Build party-coordinator communication:
- `IVFLPartyClient<T>` interface for party-side operations
- `VFLPartyClient<T>` using HTTP/gRPC with retry logic and timeout handling
- Message protocol: RegisterParty, StartTrainingRound, UploadEmbeddings,
ReceiveGradients, ReportProgress, LeaveSession
- Connection management with heartbeat, reconnection, and graceful degradation
- Request queuing and flow control

Implement session management:
- `SessionManager<T>` managing concurrent sessions with IDs and lifecycle states
- Session initialization with PSI, security negotiation, and validation
- Training loop orchestration for epochs, batches, and validation rounds
- Session teardown with model finalization and resource release
- Session recovery using persistent state

Add fault tolerance:
- Extend `ICheckpointableModel` with `IVFLCheckpointable` interface
- `VFLCheckpointManager<T>` coordinating multi-party checkpoints
- Party failure detection via heartbeat and timeout monitoring
- Training continuation logic with checkpoint reload and party reconnection
- Checkpoint synchronization protocol

Implement asynchronous training:
- `AsyncVFLCoordinator<T>` with asynchronous gradient aggregation
- Staleness-aware aggregation weighting
- Dynamic party selection ensuring fair participation
- Fairness metrics and selection probability adjustment
- Convergence detection via gradient norm and validation loss

Add directory: `src/VerticalFederatedLearning/Coordinator/`
===============================================================================

Task: 4

Develop VFL-specific data handling with columnar loaders, intelligent feature
partitioning, and sample alignment utilities.

Create VFL data loaders:
- `VFLDataLoaderBase<T>` extending DataLoaderBase with columnar access and
feature partition awareness
- `ColumnarInMemoryDataLoader<T>` with column-major storage for efficient
feature-subset extraction
- `VFLBatchGenerator<T>` coordinating batch creation with consistent ordering
and synchronized boundaries
- `StreamingVFLDataLoader<T>` for large datasets with coordinated buffering
- Data validation utilities for feature consistency and quality checks

Implement feature partitioning strategies:
- `AutomaticFeaturePartitioner<T>` with correlation-based clustering, balanced
partitioning, and information-theoretic partitioning
- `FeatureImportancePartitioner<T>` using existing feature selectors for
strategic assignment
- `PrivacyAwarePartitioner<T>` considering feature sensitivity and regulatory
constraints
- Partitioning validation for overlap detection and minimum feature counts
- Partition export/import with JSON configuration support

Integrate with feature selection:
- Modify `FeatureSelectorBase<T>` to support distributed feature selection
- `VFLFeatureSelector<T>` wrapper coordinating selection across parties
- Privacy-preserving feature importance aggregation
- Party-aware recursive feature elimination
- Feature selection result serialization

Build sample alignment pipeline:
- `SampleAlignmentPipeline<T>` orchestrating complete alignment workflow
- `SampleIDResolver<T>` supporting multiple ID formats with
normalization/hashing
- Alignment quality metrics: overlap percentage, duplicates, confidence scores
- Aligned dataset export with incremental alignment support
- Alignment caching and versioning

Extend `src/Data/Loaders/` and `src/VerticalFederatedLearning/Data/`, modify
`src/FeatureSelectors/FeatureSelectorBase<T>`
===============================================================================

Task: 5

Add production-grade features: compression, telemetry, deployment integration,
configuration tools, and comprehensive documentation.

Implement compression:
- Extend existing compression utilities from `src/ModelCompression/` for VFL
embeddings and gradients
- `VFLCompressionPipeline<T>` with quantization, sparsification, and Huffman
encoding
- `EmbeddingCompressor<T>` using product quantization or dimensionality
reduction
- Adaptive compression based on bandwidth and compute resources
- Decompression utilities with lossless reconstruction support

Add telemetry and monitoring:
- Extend `TelemetryCollector` with VFL-specific metrics: per-party training
time, communication overhead, gradient staleness
- `VFLMetricsCollector<T>` tracking privacy budget, alignment quality, and
crypto overhead
- Real-time dashboard metrics for party availability, progress, and bottlenecks
- Alerting for privacy violations, failures, and convergence issues
- Metrics export for Prometheus, statsd, or custom backends

Integrate with deployment:
- Extend `DeploymentRuntime<T>` for VFL model deployment with party-specific
versioning
- `VFLModelExporter<T>` exporting bottom/top models separately with encryption
- VFL inference pipeline coordinating multi-party predictions
- A/B testing for privacy configurations and party compositions
- Model registry tracking lineage and party contributions

Build configuration tools:
- `VFLConfigurationBuilder<T>` with fluent API for parties, privacy, and
training parameters
- Configuration validation for compatible security parameters
- Configuration templates: two-party, multi-party with coordinator,
privacy-first
- CLI tools for session management, party registration, and monitoring
- Configuration export/import with YAML/JSON support

Create documentation:
- `docs/VerticalFederatedLearning.md` with architecture, usage patterns, privacy
guarantees, and best practices
- Code examples in `examples/VFL/`: two-party, multi-party, privacy
configuration, production deployment
- API documentation for all VFL interfaces and public classes
- Troubleshooting guide for alignment failures, budget exhaustion, party
disconnection
- Performance tuning guide for compression, protocol selection, and optimization

Add directories: `src/VerticalFederatedLearning/Optimization/`,
`src/VerticalFederatedLearning/Monitoring/`,
`src/VerticalFederatedLearning/Deployment/`,
`src/VerticalFederatedLearning/Configuration/`, `examples/VFL/`

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