PubSub-VFL: Towards Efficient Two-Party Split Learning in Heterogeneous Environments via Publisher/Subscriber Architecture

AI summary

Overview

Research area: Privacy-preserving collaborative machine learning, specifically Two-Party Split Learning, also known as Vertical Federated Learning (VFL), and distributed systems architecture.

Technical level: Intermediate. The paper assumes familiarity with federated learning terminology (top/bottom models, cut layers, embeddings, gradients, parameter servers) and distributed systems concepts (publish/subscribe, asynchronous training). The mathematical formulation is presented compactly.

Scope: The paper proposes PubSub-VFL, a system-and-algorithm design for two-party VFL that combines a Publisher/Subscriber (Pub/Sub) architecture with a parameter server (PS) design, a hierarchical asynchronous mechanism, and a privacy-preserving hyperparameter optimization model, evaluated on five datasets.

What This Paper Is About

Organizations frequently want to jointly train machine learning models on data split across parties (for example, a bank holding transaction records and an insurer holding accident reports for the same customers), but privacy rules such as GDPR and HIPAA prevent pooling the raw data. Two-Party Split Learning lets each side train part of a network and exchange only intermediate results, but the standard design is synchronous and tightly coupled, so workers sit idle waiting for each other, and mismatches in compute resources and feature dimensions between the two parties make the imbalance worse. PubSub-VFL's goal is to restructure this pipeline so that compute runs asynchronously, waiting time drops, and hyperparameters are chosen automatically from system profiles without leaking private data.

Key Contributions

1. A Pub/Sub-based VFL architecture that decouples data ID alignment from model training. The design introduces separate embedding channels and gradient channels, with a unique batch ID attached to each training batch so intermediate results can be matched across parties without forcing workers to synchronize. For n samples and batch size B, the system maintains ⌈n/B⌉ embedding and gradient channels.

2. A hierarchical asynchronous mechanism. Inter-party asynchrony comes from the Pub/Sub layer; intra-party asynchrony (between the parameter server and its workers) comes from an adaptive semi-asynchronous mechanism whose synchronization interval ΔT_t shrinks as the model approaches target accuracy, governed by a tanh-based schedule with initial interval ΔT₀.

3. A privacy-preserving hyperparameter optimization model. The authors formalize computation delays, backward-pass delays, top-model delays, memory constraints, and communication delays (T_emb = E/B_b, T_grad = G/B_b) into an objective that minimizes max(T_A, T_P), then solve for the optimal number of workers (w_a, w_p) and batch size B using a dynamic programming search over discrete candidate values.

4. Theoretical and empirical validation. The paper claims a theoretical analysis showing stable convergence and compatibility with differential privacy, plus experiments on five benchmark datasets showing 2∼7× acceleration and computational resource utilization up to 91.07%.

Main Findings

• Accuracy is preserved or improved. On the five datasets, PubSub-VFL reports Energy RMSE 85.64 (baseline AVFL-PS: 85.39), Blog RMSE 22.34 (AVFL-PS: 23.45), Bank AUC 96.54 (AVFL-PS: 94.16), Credit AUC 82.34 (AVFL-PS: 80.34), and Synthetic AUC 92.87 (AVFL-PS: 91.21).

• Speed and utilization gains on the synthetic dataset. Using B = 256, w_a = 8, and w_p = 10 to reach a target accuracy of 91%, PubSub-VFL achieves a 7× reduction in running time and 35% higher CPU utilization compared to the best-performing baseline, AVFL-PS.

• Robustness under resource heterogeneity. With CPU core ratios of 50:14, 48:16, 40:24, and 36:28 between the active and passive parties, PubSub-VFL holds CPU utilization as high as 87.42% at the 50:14 ratio, while AVFL-PS drops to 42.12%.

• Robustness under data (feature) heterogeneity. With feature size ratios of 50:450, 100:400, 150:350, and 200:300, the paper reports similar trends: reducing the dimension processed by the active party further lowers running time by balancing load between the parties.

• More workers is not always better. With B = 32 on the synthetic dataset, w = 8 gives the best combination: accuracy 92.06%, time 668.11 s, CPU 88.04%, waiting 1.5288 s, communication 888.77 MB. At w = 50 the numbers degrade to 92.21% accuracy, 1661.74 s, 45.12% CPU, 19.843 s waiting, and 2197.27 MB.

• Batch size has a strong effect on efficiency. At B = 256 the system records 92.67% accuracy, 92.54 s, 91.07% CPU, 1.1389 s waiting, and 439.45 MB; at B = 16 accuracy is 91.70% with 987.64 s and 48.64% CPU, and at B = 1024 accuracy is 92.21% with 865.74 s and 52.67% CPU.

• Differential privacy integrates with limited cost. Varying the GDP parameter μ over {0.1, 0.5, 1, 2, 4, 8, 10, +∞} has minimal effect on accuracy and CPU utilization, but noticeably increases communication cost due to added noise slowing convergence. The DP protocol also improves defense against embedding inversion attacks, measured by Attack Success Rate.

Methodology in Plain English

The authors start from a standard VFL setup in which a passive party without labels computes embeddings from its features, sends them to an active party that holds labels and the top model, and receives gradients back. They identify two structural problems: the strict lockstep between the two sides, and the fact that the two sides rarely have matching compute capacity or feature dimensionality.

Their response is architectural. Instead of workers directly coordinating, they route embeddings and gradients through named channels labelled by batch ID, mimicking a message broker. This means a fast worker never has to stall for a slow counterpart. To stop fast workers from flooding channels, each channel buffer holds at most p embeddings and q gradients, timestamped, and discards the oldest under First-In-First-Out when full. A waiting deadline T_ddl (set to 10 s) discards a stalled batch and reassigns it to an available worker pair. Inside each party, a parameter server aggregates worker updates only every ΔT_t epochs, where ΔT_t starts at ΔT₀ = 5 and follows a tanh curve so that synchronization is more frequent early and less frequent later.

Because privacy rules prevent centralized resource scheduling, the authors estimate computation times from observations of a synchronous baseline, fit proportionality constants for forward, backward, and top-model costs, express memory limits as M_A(B) = M_A0 + ρ_A·B^χ and M_P(B) = M_P0 + ρ_P·B^χ, and then run a dynamic programming search over candidate worker counts (w_a ∈ [2,50], w_p ∈ [2,50]) and batch sizes ({16, 32, 64, 128, 256, 512, 1024}) subject to C_a + C_p = 64, picking the configuration that minimizes the slower party's per-iteration time.

Experiments run on Python 3.9 and PyTorch 1.12 on a server with an Intel Xeon Gold 6530 (64-core CPU), with a two-layer MLP top model and either a ten-layer MLP or a ResNet bottom model.

Why This Matters

Impact on research. The paper reframes a well-known bottleneck in vertical federated learning — synchronization stalls between parties — as a systems architecture problem rather than purely an algorithmic one. It also shows that a formal profiling and hyperparameter search can be done under privacy constraints, which is a distinguishing move relative to prior asynchronous VFL work that assumes centralized control.

Real-world applications.

• Bank–insurer collaboration to predict customer credit scores, the motivating example the paper gives explicitly.
• Healthcare analytics where hospitals and other institutions hold different feature sets for the same patients.
• E-commerce collaborations where different firms hold complementary user signals.
• Autonomous driving data partnerships, cited as a sector where data collaboration drives technological advancement.

Industry relevance. The baselines compared against include FATE and PaddleFL, which the paper describes as mature industrial frameworks, and the author affiliations include ByteDance Inc. alongside City University of Hong Kong. The measured outcomes — running time, CPU utilization, waiting time per epoch, and communication cost in MB — are the operational metrics that matter when deploying such systems at scale.

Future Directions

• Beyond two parties. The formulation is explicitly for a two-party (active/passive) setting; whether the Pub/Sub channel design and the dynamic programming planner extend to multi-party VFL is not addressed.
• Tighter pipelining models. The paper notes in a remark that the additive iteration-time formula is a simplification, and that a fully pipelined Pub/Sub system with enough buffering could achieve lower iteration times. Treating pipelining explicitly is left open.
• Quantifying the privacy–communication trade-off. The paper observes that the DP protocol raises communication cost noticeably because of slower convergence, but does not report a full characterization of this trade-off across all datasets.
• Robustness of the estimator. The optimization model relies on proportionality constants and system profiles determined empirically; how sensitive the chosen hyperparameters are to inaccurate profiling, and how performance degrades under staleness or divergence from the semi-asynchronous interval, are open questions. Several supporting results (Tables 5, 6, 7 and Algorithms 1, 2) are deferred to appendices and are not included in the provided content.

Target Audience

This paper is most useful for distributed systems and privacy-preserving machine learning researchers, applied scientists building cross-organization training pipelines, and engineers working on industrial federated learning frameworks such as FATE or PaddleFL who need concrete guidance on worker counts, batch sizing, and asynchronous scheduling. Readers without background in federated learning or parameter-server architectures will need to consult the appendices and background references first.

Authors’ abstract

With the rapid advancement of the digital economy, data collaboration between organizations has become a well-established business model, driving the growth of various industries. However, privacy concerns make direct data sharing impractical. To address this, Two-Party Split Learning (a.k.a. Vertical Federated Learning (VFL)) has emerged as a promising solution for secure collaborative learning. Despite its advantages, this architecture still suffers from low computational resource utilization and training efficiency. Specifically, its synchronous dependency design increases training latency, while resource and data heterogeneity among participants further hinder efficient computation. To overcome these challenges, we propose PubSub-VFL, a novel VFL paradigm with a Publisher/Subscriber architecture optimized for two-party collaborative learning with high computational efficiency. PubSub-VFL leverages the decoupling capabilities of the Pub/Sub architecture and the data parallelism of the parameter server architecture to design a hierarchical asynchronous mechanism, reducing training latency and improving system efficiency. Additionally, to mitigate the training imbalance caused by resource and data heterogeneity, we formalize an optimization problem based on participants' system profiles, enabling the selection of optimal hyperparameters while preserving privacy. We conduct a theoretical analysis to demonstrate that PubSub-VFL achieves stable convergence and is compatible with security protocols such as differential privacy. Extensive case studies on five benchmark datasets further validate its effectiveness, showing that, compared to state-of-the-art baselines, PubSub-VFL not only accelerates training by $2 \sim 7\times$ without compromising accuracy, but also achieves a computational resource utilization rate of up to 91.07%.