字节跳动技术团队

PubSub-VFL——基于发布订阅架构的高效异构垂直联邦学习框架

研究背景:垂直联邦学习面临的异构挑战

随着数据隐私保护法规的日益严格和跨机构数据协作需求的不断增长,垂直联邦学习(Vertical Federated Learning, VFL)作为一种重要的隐私保护机器学习范式受到广泛关注。在 VFL 中,不同参与方拥有相同用户的不同特征维度,通过协作训练可以构建更强大的机器学习模型,朴素的 VFL 框架如下图所示。

Image

基于分割学习的朴素 VFL 框架

然而,现有的 VFL 系统在实际部署中面临着严峻的异构性挑战:

  • 资源异构性:不同参与方的计算资源、网络带宽、存储能力存在显著差异

  • 数据异构性:各方的数据分布、数据量、特征维度可能差异巨大

  • 系统异构性:参与方使用不同的硬件平台、操作系统和软件环境

传统的参数服务器(Parameter Server, PS)架构虽然在数据并行场景下表现出色,但在处理 VFL 中的异构环境时存在明显不足:

  • 同步等待开销大:系统性能受限于最慢的参与方,导致整体训练效率低下

  • 通信模式固化:点对点的通信方式缺乏灵活性,难以适应动态的网络环境

  • 资源利用率低:无法根据各方的实际资源情况进行自适应调度和优化

核心创新:PubSub-VFL 框架设计

为解决上述挑战,字节跳动研究团队提出了 PubSub-VFL 框架,其核心创新在于将发布订阅(Pub/Sub)架构与参数服务器架构相结合,构建了一个层次化的异步通信机制。系统总体架构如下图所示。

Image

PubSub-VFL 系统架构

双层异步协同

PubSub-VFL 采用创新的双层异步架构设计:

  • 参与方间(Inter-party)异步:采用发布/订阅(Pub/Sub)架构达成参与方之间的异步通信,借助消息代理(Message Broker)对发布者与订阅者进行解耦。此方法使训练器(Worker)能够专注于本地训练,而无需预先保证数据标识符(ID)的严格对齐,有效地将训练过程与数据 ID 匹配任务相分离,进而实现异步操作,消除等待延迟,提升系统并发性。

  • 参与方内(Intra-party)异步:在 Pub/Sub 架构的基础上,我们实现了跨参与方的异步通信,有效地消除了 Worker 端的等待延迟。为了进一步提高每个参与方内部(即参数服务器(PS)与其 Worker 之间)的计算效率,我们通过一个内部异步机制扩展了这一设计。然而,如果控制不当,这种分层异步可能会阻碍模型收敛。为了解决这个问题,我们提出了一种动态半异步机制,该机制根据训练反馈自适应地调节同步间隔。具体来说,随着模型接近目标准确率,同步间隔会减小,从而在计算速度和收敛稳定性之间取得平衡。

传统 PS 架构的局限

  • 同步等待开销大

  • 通信模式固化

  • 难以处理异构环境

  • 资源利用率低

PubSub-VFL 的优势

  • 层次化异步通信

  • 灵活的消息传递

  • 自适应资源调度

  • 高效异构处理

下图具体形象化展示了 PubSub 对比传统 PS 在性能和架构上的优越性。

Image

传统异步框架与 Pub/Sub 异步框架

安全系统建模

朴素 VFL 方案因未考虑参与方资源与数据异质性,常导致计算负载失衡、资源利用率低,且集中式资源调度会泄露隐私。为此,PubSub-VFL 提出基于系统画像的建模方法,通过刻画双方硬件能力(CPU 核心数)、模型特性(前向 / 反向传播耗时),精准量化计算与通信延迟。

其核心是构建双方计算延迟公式:Image(主动方前向传播耗时)、Image(被动方前向传播耗时),结合通信延迟Image,实现负载均衡的量化分析,为后续参数优化奠定基础,同时避免隐私泄露。

最佳参数规划

传统 VFL 依赖经验设置超参数(worker 数量、批大小),难以在资源约束(内存上限等)下最小化迭代耗时。

PubSub-VFL 在系统规划阶段,将超参数优化转化为带约束的极小化问题,以

Image

为目标,该公式旨在最小化两个时间组合中的较大值,其中Image、Image等分别代表不同计算过程的时间 。

结合内存约束

Image

通过两个分式计算结果取最小值的方式,确定最大可用内存Image。

最后,通过动态规划遍历离散决策空间,高效找到最优参数配置,解决了异质性场景下资源浪费与训练延迟问题。

数据隐私保护

为了增强隐私保护能力,PubSub-VFL 集成了高斯差分隐私(Gaussian Differential Privacy, GDP)协议:

  • 嵌入信息保护:对传输的嵌入向量添加高斯噪声,防止嵌入反转攻击(Embedding Inversion Attack, EIA)

  • 理论收敛保证:研究团队从理论上证明了即使在差分隐私保护下,PubSub-VFL 仍能保证稳定收敛

  • 隐私-效用平衡:通过自适应噪声调节机制,在保护隐私的同时最大化模型性能

实验验证:显著的性能提升

实验设置

  1. 实验环境:为评估我们 PubSub-VFL 系统的性能,我们在五个数据集上进行了广泛的实验。所有实验均使用 Python 3.9 和 PyTorch 1.12 开发,并在配备 Intel (R) Xeon (R) Gold 6530(64 核 CPU)的服务器上进行评估。

  2. 数据集:我们在四个公共基准数据集上评估了 PubSub-VFL,涵盖了回归和分类任务,以及一个大规模的合成数据集。对于回归任务,我们使用了 Energy 和 Blog 数据集。对于分类任务,我们采用了 Bank 和 Credit 数据集。为了评估可扩展性,我们使用 scikit-learn 生成了一个包含 100 万样本和 500 个特征的合成数据集。

  3. 模型:对于顶部模型,我们使用了一个两层的多层感知器(MLP)。对于底部模型,我们使用了两种不同大小的模型,即一个 10 层的 MLP 和一个 ResNet,这可以验证 PubSub-VFL 在不同模型大小下的性能。

  4. 基线:我们采用了以下基线:1)纯 VFL;2)带 PS 的 VFL;3)异步 VFL(AVFL);4)带 PS 的异步 VFL。

性能与资源利用率

PubSub-VFL 在计算和通信效率方面表现出色。与表现最佳的基线 AVFL-PS 相比,PubSub-VFL 的运行时间减少了7倍,CPU 利用率提高了35%。这些收益源于“工人”空闲时间的减少和并行性的提高。此外,分层异步机制提高了收敛效率,导致通信成本低于其他方法。

PubSub-VFL 方案在模型精度方面表现优异。在多个基准数据集上的实验结果显示,PubSub-VFL 的精度与现有最先进方法相当甚至更高,也证明了该方案在不同任务和模型规模下的鲁棒性。

数据集

度量

VFL

VFL-PS

AVFL

AVFL-PS

PubSub-VFL

Energy

RMSE

84.58

84.44

85.41

85.39

85.64

Blog

RMSE

23.20

23.12

23.38

23.45

22.34

Bank

AUC

94.54

94.13

94.12

94.16

96.54

Credit

AUC

81.90

81.34

80.83

80.34

82.34

Synthetic

AUC

91.27

91.31

90.97

91.21

92.87

在资源和数据异构场景中,PubSub-VFL 同样表现出色。当 CPU 核心比例为50:14时,PubSub-VFL 的 CPU 利用率仍高达87.42%,而 AVFL-PS 仅为42.12%。这表明 PubSub-VFL 能有效平衡计算效率,减少运行时间并保持高 CPU 利用率。

参数敏感性分析与消融分析

实验还验证了系统各组件的有效性:

  • 工作节点数量:在4-16个节点范围内,性能随节点数增加而提升,但边际效应递减

  • 批处理大小:存在最优批处理大小,过大或过小都会影响效率

  • 隐私预算:隐私保护强度与模型性能之间存在可调节的权衡关系

通过消融研究,我们发现等待截止时间和内部半异步机制对 PubSub-VFL 的性能至关重要,它们有效地平衡了同步和异步,以减轻梯度陈旧性并确保及时更新。移除这些机制会导致显著的性能下降。

技术亮点与创新贡献

PubSub-VFL 框架的主要技术贡献包括:

  1. 架构创新

首次将 Pub/Sub 架构引入 VFL,实现了层次化异步通信机制,有效解决了异构环境下的同步等待问题

  1. 系统优化

提出系统分析规划驱动的超参数优化方法,能够根据实际资源情况自动调整系统配置

  1. 隐私保护

集成高斯差分隐私协议,在保护隐私的同时提供理论收敛保证

应用前景与行业意义

PubSub-VFL 框架在多个领域具有广阔的应用前景:

  • 金融科技领域:银行、保险公司等金融机构可以在不共享敏感客户数据的前提下,联合训练风险评估和欺诈检测模型,提升金融服务的智能化水平。

  • 医疗健康领域:不同医院和研究机构可以协作训练疾病诊断和药物发现模型,在保护患者隐私的同时推动医疗 AI 的发展。

  • 智能制造领域:制造企业可以与供应链伙伴联合优化生产流程和质量控制,在保护商业机密的前提下实现协同创新。

  • 广告营销领域:广告平台和品牌方可以在保护用户隐私的前提下,联合训练更精准的推荐和定向广告模型。

总结与展望

PubSub-VFL 框架通过引入发布订阅架构和系统分析规划机制,成功解决了垂直联邦学习在异构环境下面临的效率和隐私挑战。该研究不仅在理论上提供了新的架构设计思路,更在实践中展现了显著的性能优势,为联邦学习技术的产业化应用奠定了重要基础。

目前,本文所提及的工作成果(《PubSub-VFL: Towards Efficient Two-Party Split Learning in Heterogeneous Environments via Publisher/Subscriber Architecture》)被机器学习顶级会议 NeurIPS 2025 正式接收。该研究提出了 PubSub-VFL 框架,创新性地将发布订阅(Publisher/Subscriber)架构引入垂直联邦学习,在异构环境下显著提升了双方分割学习的计算效率,为联邦学习在实际部署中面临的资源异构和数据异构问题提供了高效解决方案。