專為決策樹打造,新加坡國立大學&清華大學聯合提出快速安全的聯邦學習新系統
聯邦學習是機器學習中一個非常火熱的領域,指多方在不傳遞數據的情況下共同訓練模型。隨著聯邦學習的發展,聯邦學習系統也層出不窮,例如 FATE, FedML, PaddleFL, TensorFlow-Federated 等等。然而,大部分聯邦學習系統不支持樹模型的聯邦學習訓練。相比于神經網絡,樹模型具有訓練快,可解釋性強,適合表格型數據的特點。樹模型在金融,醫療,互聯網等領域有廣泛的應用場景,比如用來做廣告推薦、股票預測等等。
決策樹的代表性模型為 Gradient Boosting Decision Tree (GBDT)。由于一棵樹的預測能力有限,GBDT 通過 boosting 的方法串行訓練多棵樹,通過每棵樹用來擬合當前預測值和標簽值的殘差的方式,最終達到一個好的預測效果。代表性的 GBDT 系統有 XGBoost, LightGBM, CatBoost, ThunderGBM,其中 XGBoost 多次被 KDD cup 的冠軍隊伍所使用。然而,這些系統都不支持聯邦學習場景下的 GBDT 訓練。近日,來自新加坡國立大學和清華大學的研究者提出了一種專注于訓練樹模型的聯邦學習新系統 FedTree。
- 論文地址:https://github.com/Xtra-Computing/FedTree/blob/main/FedTree_draft_paper.pdf
- 項目地址:https://github.com/Xtra-Computing/FedTree
FedTree 系統介紹FedTree 架構圖如圖 1 所示,共有 5 個模塊: 接口,環境,框架,隱私保護以及模型。
圖 1: FedTree 系統架構圖
接口:FedTree 支持兩種接口:命令行接口和 Python 接口。用戶只需給定參數(參與方數量,聯邦場景等),就可以用一行命令來運行 FedTree 進行訓練。FedTree 的 Python 接口和 scikit-learn 兼容,可以調用 fit() 和 predict() 來進行訓練和預測。
環境:FedTree 支持在單機上模擬部署聯邦學習,和在多機上部署分布式聯邦學習。在單機環境上,FedTree 支持將數據進行劃分成多個子數據集,每個子數據集作為一個參與方進行訓練。在多機環境上,FedTree 支持將每個機器作為一個參與方,機器之間 通過 gRPC 進行通信。同時,除了 CPU 以外,FedTree 支持使用 GPU 來加速訓練。
框架:FedTree 支持橫向和縱向聯邦學習場景下 GBDT 的訓練。橫向場景下,不同參與方有著不同的訓練樣本和相同的特征空間。縱向場景下,不同參與方有著不同的特征空間和相同的訓練樣本。為了保證性能,在這兩種情境下,多有參與方都共同參與每個節點的訓練。除此之外,FedTree 還支持集成學習,參與方并行地訓練樹之后再進行聚合,以此減少參與方之間的通信開銷。
隱私:由于訓練過程中傳遞的梯度可能會泄露訓練數據的信息,FedTree 提供不同的隱私保護方法來進一步保護梯度信息,包括同態加密 (HE) 和安全聚合 (SA)。同時,FedTree 提供了差分隱私來保護最終訓練得到的模型。
模型:以訓練一棵樹為基礎,FedTree 通過 boosting/bagging 的方法支持訓練 GBDT/random forest。通過設置不同的損失函數,FedTree 訓練得到的模型支持多種任務,包括分類和回歸。
實驗表 1 總結了不同系統在 a9a, breast 和 credit 上的 AUC 和 abalone 上的 RMSE,FedTree 的模型效果和用所有數據訓練 GBDT(XGBoost, ThunderGBM)以及 FATE 中的 SecureBoost (SBT) 幾乎一致。而且,隱私保護策略 SA 和 HE 并不會影響模型性能。
表 1: 不同系統的模型效果比較
表 2 總結了不同系統每棵樹的訓練時間(單位:秒),可以看到 FedTree 相比于 FATE 快了很多,在橫向聯邦學習場景下能達到 100 倍以上的加速比。
表 2: 不同系統每棵樹的訓練時間比較
更多研究細節請參考 FedTree 論文原文。