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 ,與我們聯系。
?