XGBoost 是一种广泛用于表格数据建模的机器学习算法。为了将 XGBoost 模型从单站点学习扩展到多站点协作训练,NVIDIA 开发了 Federated XGBoost ,这是一种用于联邦学习的 XGBoost 插件。它包括用于跨去中心化数据源联合训练 XGBoost 模型的垂直协作设置,以及 水平直方图和树的联合学习 。
在垂直环境下,每个方拥有面向全体人口的部分特征,并且只有一个方持有该标签。标签所有者被称为主动方,而所有其他方都是被动方。在横向设置下,每个方拥有所有特征和标签信息,但仅限部分人口。
此外, NVIDIA FLARE 是一款与领域无关、开源且可扩展的联邦学习 SDK,通过引入处理通信挑战的功能来增强现实世界的联邦学习体验。其中包括多项并发训练作业,以及可能因网络条件而中断的作业。
目前,Federated XGBoost 是在完全相互信任的前提下构建的,这表明除了模型训练之外,任何组织都不想了解更多信息。但在实践中,对于联邦协作而言,诚实而好奇的环境更为逼真。例如,在垂直联合 XGBoost 中,被动方可能有兴趣从主动方发送的梯度中恢复标签信息。在水平联合学习中,服务器或其他客户端可以访问每个客户端的梯度直方图并了解其数据特征。
NVIDIA Flare 2.5.2 和 XGBoost 联邦安全通过保护这些潜在的信息问题,扩展了联邦 XGBoost 的范围。具体来说:
- 安全的联合算法(包括横向和纵向算法)已实施并添加到 XGBoost 库支持的联合方案中,以解决不同假设下的数据安全模式。
- 我们使用插件和处理器接口系统将同态加密 (HE) 功能添加到安全的联邦 XGBoost 工作流,该系统旨在通过 XGBoost 的计算和 NVIDIA Flare 的通信,通过适当的加密和解密在两者之间可靠有效地桥接。
- HE 插件基于 CPU 和 CUDA 加速,可根据硬件和效率要求提供通用适应性。CUDA 插件的速度远快于当前的第三方解决方案。
借助 HE,可以通过密文执行关键的联合计算步骤,并且相关资产 (gradients 和 partial histograms) 经过加密,不会被其他方在计算期间学习。这为用户提供了数据安全保证,这是联邦学习的基本优势之一。
如本文所述,采用联合 XGBoost 的 CUDA 加速同态加密增加了数据隐私安全保护,与第三方解决方案相比,垂直 XGBoost 的速度提升高达 30 倍。
协作模式和安全模式
对于垂直 XGBoost,主动方持有标签,该标签可被视为整个过程中“最有价值的资产”,不应被被动方访问。因此,从模型训练的角度来看,本例中的主动方是“主要贡献者”,担心会将此信息泄露给被动客户端。在这种情况下,安全保护主要针对被动客户端的标签信息。
为保护用于垂直协作的标签信息,在主动方计算每个样本的梯度后,每一轮 XGBoost 都会对梯度进行加密,然后再发送给被动方 (图 1)。在接收到加密的梯度 (密文) 后,这些梯度将根据每个被动方的特定特征分布进行累积。生成的累积直方图将返回给活跃方,然后解密,并由活跃方进一步用于树构建。
对于水平 XGBoost,每方拥有“同等状态”(完整特征和部分人口的标签),而联邦服务器则执行聚合,而不拥有任何数据。因此,在这种情况下,客户端担心信息泄露给服务器以及彼此之间。因此,要保护的信息是每个客户端的本地直方图。
为保护本地直方图以进行水平协作,我们将对直方图进行加密,然后再将其发送到联邦服务器进行聚合。然后,将通过密文执行聚合,并将加密的全局直方图返回给客户端,然后对其进行解密并用于树构建。这样,服务器将无法访问明文直方图,而每个客户端只能在聚合后学习全局直方图,而不是单个本地直方图。
使用适当的 HE 方案进行加密
多个库涵盖各种 HE 方案 (支持和不支持 GPU),因此正确选择最高效的方案以满足特定联合 XGBoost 设置的特定需求非常重要。让我们来看一个示例,假设 N = 5 个参与者,M = 200K 个数据样本总数,J = 30 个特征总数,并且每个特征直方图都有 K = 256 个插槽。根据联合学习应用程序的类型:(垂直或水平应用程序,我们将需要不同的算法。
对于垂直应用,加密目标是单个 g/h 数字,计算是根据它们所属的直方图槽添加加密数字。理论上,由于 g/h 的数量与样本数量相同,因此对于每一轮提升:
- 所需的总加密量为 M * 2 = 400k (g 和 h),并且每次都加密一个数字
- 所需的加密加法总量为 (M – K)*2*J ≈ 12m
在这种情况下,最佳方案选择是 Paillier,因为需要在单个数字上执行加密。使用针对 CKKS 等向量的方案将严重浪费空间。
另一方面,对于水平应用,加密目标是局部直方图 G/H,计算是将局部直方图相加,形成全局直方图。对于每一轮 Boosting:
- 所需的总加密将为 N * 2 = 10 (G 和 H),并且每次都加密长度为 J * K = 7680 的向量
- 所需的加密加法总数为 (N – 1) * 2 = 18
在这种情况下,最佳方案选择是 CKKS,因为它能够一次性处理直方图向量 (例如长度为 7680)。
我们提供仅支持 CPU 和高效 GPU 加速的加密解决方案。
示例结果
通过实施之前在 XGBoost 和 NVIDIA Flare 上描述的工作流,我们使用 信用卡欺诈检测数据集 测试了安全的联邦工作流。结果如下:
垂直学习的 AUC(安全和非安全):
[0] eval-auc:0.90515 train-auc:0.92747
[1] eval-auc:0.90516 train-auc:0.92748
[2] eval-auc:0.90518 train-auc:0.92749
横向学习的 AUC(安全和非安全):
[0] eval-auc:0.89789 train-auc:0.92732
[1] eval-auc:0.89791 train-auc:0.92733
[2] eval-auc:0.89791 train-auc:0.92733
通过将树模型与集中式基准进行比较,我们得出了以下观察结果:
垂直联合学习 (非安全) 的树模型与集中式基准完全相同,例如 PyTorch、pandas、LangChain、Megatron、NVIDIA、cuOpt、Stable Diffusion、Llama、Jetson、Google、Python、Hugging Face、Arm、Anaconda、Siemens、DPU、GPU、Github 等。
垂直联合学习 (安全) 的树结构与集中式基准相同。此外,由于各方拥有不同的特征子集,因此它会在不同的方面生成不同的树记录,并且不应学习其他人拥有的特征的切割信息。
横向联合学习 (安全和非安全) 的树模型与集中式基准不同。这是通过全局数据 (集中) 或本地数据 (水平) 进行的初始特征分位数计算得出的结果。
有关更多详情,请参阅 NVIDIA Flare Secure XGBoost 示例 。
加密方法的效率
为了对我们的解决方案进行基准测试,我们使用具有不同特征的各种数据集进行了实验,包括大小差异 (从小到大) 和特征维度差异 (从少到多)。这些基准测试旨在展示我们算法的可靠性,并突出显示速度和效率方面的显著性能改进。
数据集和数据拆分
我们使用了三个数据集 (涵盖不同的数据大小和特征大小) 来说明它们对加密方法效率的影响。表 1 汇总了数据特征。信用卡欺诈检测数据集标记为 CreditCard , Epsilon 数据集标记为 Epsilon, HIGGS 数据集的子集标记为 HIGGS。
CreditCard | HIGGS | Epsilon | |
数据记录大小 | 284807 | 6200000 | 400000 |
特征大小 | 28 | 28 | 2000 |
训练集大小 | 227845 | 4000 万 | 320000 |
验证集大小 | 56962 | 2200,000 | 80000 |
对于垂直联合学习,我们将训练数据集拆分成两个客户端,每个客户端拥有相同数据记录的不同特征 (表 2)。
特征 | CreditCard | HIGGS | Epsilon |
标签客户端 | 10 | 10 | 799 |
非标签客户端 | 18 | 18 | 1201 |
对于水平联合学习,我们将训练集平均分为三个客户端 (表 3)。
数据记录 | CreditCard | HIGGS | Epsilon |
客户 1 | 75948 | 1333333 | 106666 |
客户端 2 | 75948 | 1333333 | 106666 |
客户端 3 | 75949 | 1333334 | 106668 |
实验结果
我们使用以下参数执行端到端 XGBoost 训练:num_trees = 10、max_depth = 5、max_bin = 256。测试使用 NVIDIA Tesla V100 GPU 和 Intel E5-2698 v4 CPU 执行。图 3 和图 4 显示了时间比较。请注意,模拟在同一台机器上运行,因此联合通信成本可以忽略不计。
安全的垂直联合 XGBoost
我们将 NVIDIA Flare 工作流 CUDA 加速的 Paillier 插件 (称为 GPU 插件) 的时间成本与用于安全垂直联合 XGBoost 的现有第三方开源解决方案的时间成本进行了比较。两者都经过 HE 加密。图 3 显示,根据数据和特征大小的组合,我们的解决方案的速度可提升 4.6 倍到 36 倍。请注意,第三方解决方案仅支持 CPU。
安全的横向联合 XGBoost
对于安全的横向 Federated XGBoost,第三方产品没有包含 HE 的安全解决方案。因此,我们将不使用加密的 NVIDIA Flare 工作流的时间成本与使用 CPU 的 CKKS 加密插件 (称为 CPU 插件) 的时间成本进行比较,以了解数据保护的加密开销。
如图 4 所示,在这种情况下,计算速度明显快于垂直场景 (低几个数量级),因此在如此合理的开销下可能不需要 GPU 加速。仅对于具有非常宽直方图的数据集 (例如 Epsilon),加密开销会更大 (但仍仅为垂直设置的约 5%)。
总结
在本文中,我们展示了 GPU 加速的同态加密如何增强联合 XGBoost 的安全性,从而通过 NVIDIA FLARE 实现保护隐私的横向和纵向联合学习。与联合 XGBoost 的现有工作相比,新功能可提供 1)安全的联合 XGBoost 管道,确保算法级别的数据安全;以及 2)高效的 CUDA 加速解决方案,其速度远快于由 GPU 计算助力的市场上的当前替代方案。这将鼓励对通常使用 XGBoost 的数据安全和学习效率要求较高的领域进行调整,例如金融业中的欺诈检测模型训练。
如需了解更多信息和端到端示例,请访问 GitHub 上的 NVIDIA/NVFlare 。如有疑问或意见,请发送电子邮件至 federatedlearning@nvidia.com ,与我们联系。