解密Angel PowerFL联邦学习平台中的纵向GBDT算法

导语:  GBDT（或XGBoost）算法是一种十分流行的树集成学习算法，不但是数据科学竞赛的常胜工具，在工业界的具体业务场景也有广泛的落地场景。然而，近年来用户隐私数据保护条例逐渐完善，“数据孤岛”逐渐形成，不但数据难以收集，不同公司或团队之间的数据也难以共享，这直接影响着机器学习模型的效果。为了应对这个问题，联邦学习技术逐渐进入人们的视线。本文聚焦腾讯自研的联邦学习平台Angel PowerFL中纵向联邦GBDT算法实现，介绍纵向联邦GBDT算法的原理和流程，并讲解相关的优化技术。

梯度提升决策树算法（Gradient Boosting Decision Tree, GBDT）由于其高准确率和可解释性，被广泛应用于分类、回归、排序等各位问题中。目前有许多有名的GBDT算法实现，如XGBoost、LightGBM、CatBoost等。然而，这些主流实现专注于传统的机器学习场景，要求所有的训练数据都是可以访问到的。近年来，人们对数据隐私的保护逐渐注重起来，相关的法律法规也在逐渐完善，一方面，用户的数据变得更难收集，另一方面，不同公司的数据直接用于商业合作可能会导致法律纠纷。为了应对这种“数据孤岛”的现象，联邦学习逐渐成为了一个被广泛研究并应用的话题。联邦学习可以在保护数据隐私的前提下，联合多个数据源来进行机器学习模型训练，从而提高模型的精度。

得益于GBDT算法的强大能力，联邦GBDT算法不仅在学术界得到广泛研究，也受到了工业界的青睐。然而，在实际应用中，业界的联邦学习框架性能不尽人意，可支持的数据集规模也比较小。为了解决这种供需矛盾，腾讯自研的联邦学习平台Angel PowerFL（以下简称PowerFL）针对纵向联邦GBDT算法的瓶颈进行了一系列的分析与优化，并提出一种高效的纵向联邦GBDT算法实现。与业界开源框架相比，PowerFL在公开数据集上可达到18.9倍的训练性能提升，并可以支持更大规模的数据集。

在本文，笔者首先简要地对GBDT算法和纵向联邦GBDT算法进行介绍，然后讲述PowerFL中纵向联邦GBDT算法的实现与优化技巧，最后进行一系列的实验对比。

GBDT算法简介

在本章节中，笔者简要介绍GBDT算法的核心思想与关键步骤。

GBDT示例

GBDT算法使用Boosting策略，通过依次学习多棵决策树（基模型）来不断地提升模型能力。上图是一个GBDT模型的示例，GBDT算法的预测输出为所有决策树的预测值之和，以分类问题为例，一个GBDT模型的预测值为

其中是树的数量，是学习率，是第棵树对样本的预测结果。

在GBDT所采取的Boosting策略中，第棵决策树学习的是前棵决策树的预测值的残差。对于第棵决策树，GBDT算法会根据前棵决策树的预测值来计算一阶梯度和二阶梯度，然后第棵决策树会根据梯度来进行训练，并达到不断拟合残差的目标。

基于梯度直方图的最优分裂点找寻

梯度直方图是GBDT算法中的核心数据结构，几乎所有主流的GBDT实现均采用基于梯度直方图的方法来训练一颗决策树。上图展示了该算法如何寻找一个特征维度上的最优分裂点，具体包括以下几个步骤。

1. 选举候选分裂点：该步骤通常在数据预处理阶段进行。在读入训练数据后，GBDT算法会对每个特征选举一些候选的分裂点，常用的方法是选择该特征上所有特征值的分位数。例如，可以使用一种名为Quantile Sketch的数据结构，对每个特征进行近似的分位数查询。
2. 构建梯度直方图：对于一个树结点，扫描树结点上的所有训练样本，根据其特征值，将样本梯度累加到对应的直方图桶中。在直方图构建完成后，直方图上的每个桶是梯度的汇总，如年龄在的范围内且落入该树结点的所有样本的梯度之和。
3. 寻找最优分裂点：对于每个特征的一阶和二阶梯度直方图，进行从左到右的遍历，根据GBDT算法的分裂准则，计算每个候选分裂点的增益，选择所有候选分裂点中，分裂增益最高的一个，作为该树结点的最优分裂点。

纵向联邦GBDT算法

纵向联邦是跨企业合作中十分常见的联邦学习模式，以两方为例，有且仅有一方拥有机器学习的目标label信息，称为Guest方，而没有label信息的则称为Host方。通过纵向联邦学习，Guest方可以借助Host方的特征数据，提高机器学习模型的能力，同时又能保护各个参与方的数据隐私。

如前文所述，GBDT算法中的核心步骤是梯度直方图的构建，为了寻找最优分裂点，首要条件是需要让Host方根据己方的特征构建梯度直方图。然而，由于梯度是根据label进行计算的，往往蕴涵着label的相关信息。以分类任务中常用的logistic loss为例，其一阶梯度为，不难看出，对于正样本，其梯度恒负，对于负样本，其梯度恒正。因此，直接发送梯度至Host方是会导致Guest方label信息的泄漏。

纵向联邦GBDT算法伪代码

为了解决这个问题，联邦GBDT算法需要对样本梯度进行加密保护。进一步的，由于梯度直方图的构建只含有加法运算，所有满足加法同态的加密协议（见下文）均为可选的方案。上图展示了基于同态加密的纵向联邦GBDT算法伪代码，其中参与方之间的发送和接收分别使用红色和蓝色字体高亮，主要流程可以总结为：

1. 梯度加密：对于一棵新的决策树，Guest方根据当前预测值和label进行梯度计算，并将梯度加密后发送至Host方。
2. 构建梯度直方图：对于决策树中的树结点，双方根据自己的特征数据建立梯度直方图，Host方将梯度直方图发送给Guest方。
3. 寻找最优分裂点：Guest方对Host方的梯度直方图进行解密，并根据分裂增益计算公式，寻找双方的全局最优分裂点；若最优分裂点属于Host方，Guest方需要将分裂的ID返回给Host方进行解析。
4. 分裂树结点：拥有最优分裂点的一方，对该树结点上的样本进行分裂划分，并将分裂划分结果发送给对方，用于更新样本-树结点之间的索引。
5. 更新预测值：Guest方根据权重计算公式计算叶子结点的预测权重值，并更新样本预测值；Host方无法得知叶子结点的权重。

隐私分析

笔者对上述流程中，针对双方交互中传输的信息，进行隐私安全分析。

• 加密梯度：Host方接收加密后的样本梯度，这一步的隐私性由加密算法来进行保证，如PowerFL使用了十分经典的Paillier加密算法，其安全性是可靠的。
• 梯度直方图：Guest方接收Host方梯度直方图后，会解密复原直方图桶中的值。由于梯度直方图是刻画特征维度的信息，Guest方是否能从中获取到Host方的信息呢？答案是否定的。每个直方图桶代表了一个特征区间内（如年龄在）的样本梯度总和，即，其中代表落入这个区间的样本集合。通常来说，样本的个数远远大于直方图桶的个数，想要利用反推出是一个NP难的问题。此外，Host方每个分裂点ID对应的特征值范围是没有暴露给Guest方的，因此，Guest方并不清楚每个直方图桶对应的含义，这进一步保护了Host方的数据安全。
• 最优分裂点：一方面，在Guest发送最优分裂点给Host方之前，消去了分裂增益这一信息，只保留候选分裂的ID信息，因此Host方无法从中学习到任何与模型有关的信息；另一方面，如上所述，Host方每个候选分裂点对应的特征是没有暴露给Guest方的，从而也不会暴露与特征相关的元信息。
• 样本分裂划分：在分裂树结点时，由于最优分裂点对应的特征值被有且仅有一方拥有，因此需要告知另一方每个样本被划分至左子结点或右子结点，以达到样本索引的一致性。若样本被划分至同一个子结点，另一方只能得知样本在某个特征上均大于（或均小于）某个阈值，而无法确定样本在该特征上的具体值，甚至无法确定该特征的具体含义以及阈值的具体大小。同理，若样本被划分至不同的子结点，另一方只能得知样本在某个特征上可以被某个阈值所区分，而具体的特征信息仍然是不可知的。

上述的纵向联邦GBDT算法协议，不但可以保证参与双方的数据安全，而且可以实现无损的计算，因此，联邦训练的模型，与将所有数据放到同一个数据中心再进行传统训练，其模型精度是完全一致的，这在本文的实验环节中也得到了验证。

Paillier同态加密协议

Paillier加密协议是最常用的同态加密协议之一，满足加同态性质。Paillier同态加密协议支持以下操作：

• 加密：
• 解密：
• 同态加法：
• 数乘操作：

PowerFL和FATE均支持Paillier加密协议，并在纵向联邦GBDT实现中应用该协议来进行样本梯度的加密，直方图的同态加法，以及直方图的解密。

系统架构与实现

接下来笔者介绍一下腾讯自研联邦学习平台PowerFL中的纵向联邦GBDT实现。

系统架构

系统架构

上图是该实现的系统设计架构，包含了Guest方和Host方的独立计算，以及跨公网的交互。该系统中有三种角色，分别为：

• Scheduler：负责掌控整体的训练流程，向Worker节点分发任务，并收集模型训练的元信息，也起到保存GBDT模型的作用。
• Worker：每个Worker节点拥有一部分数据集（按样本划分，data parallel），并根据该数据集执行Scheduler分发的任务，如构建梯度直方图，更新样本-树结点索引等。值得注意的是，Guest方和Host方的样本划分是在Worker级别对齐，例如，若Guest方的Worker #1拥有100条训练样本，则Host方的Worker #1有且仅有该100条训练样本。这种Worker级别的样本对齐，对于训练中的双方交互起到很大的便利作用。
• Channel：由于Guest方和Host方一般代表着不同的合作企业，其任务是运行在不同的数据中心的。对于Scheduler和Worker，通常是IDC集群中统一管理的计算容器（如Yarn，Kubernetes等），为了避免受到攻击，这些机器往往是无法直接连接外网的。因此，PowerFL在少量堡垒网关机器上部署消息队列，双方网关机通过token验证进行消息的传输，而Scheduler和Worker则通过订阅不同的topic来进行通信。

系统实现

腾讯自研的联邦学习平台PowerFL对易用性、高效性与可扩展性进行了充分的考虑：

• 在每个参与方内部（Scheduler和Worker）由Apache Spark作为计算引擎，可以更方便地与其他任务流进行对接。
• 使用Apache Pulsar作为跨公网传输的消息队列，可以支撑大量的网络传输，可扩展性好，并支持断点自动续传等功能。
• 使用C++实现了一个高效的Paillier密文运算库，并使用OpenMP进行多线程支持，性能远超已有开源实现。

瓶颈分析与优化

在实际运行中，笔者观察到，存在一些十分耗时的操作，这些操作不但会导致性能的下降，而且会使得计算资源使用率降低。因此，在本章节，笔者对纵向联邦GBDT算法进行复杂度分析，并提出多项优化技术。

复杂度分析

在联邦GBDT算法中，出于对隐私安全的保护，引入了同态加密操作，包括加密、解密和密文加法操作。由于密文的运算操作比普通浮点数的运算耗时长数十甚至数百倍，也进而成为了训练过程中的主要瓶颈。因此，笔者着重对密文操作的复杂度进行分析。

为了方便读者理解，这里首先定义一些数学标识：

接下来，笔者对加密、同态加法、解密三种密文运算的复杂度进行分析：

• 加密：为了训练一颗决策树，纵向联邦GBDT算法需要对所有参与训练的样本进行梯度加密，因此，每棵树的加密复杂度为，该操作仅Guest方进行。
• 同态加法：构建梯度直方图需要扫描树结点上的所有样本的非零特征值，对于决策树的每一层，样本总量不变，因此，每层的同态加法复杂度为，该操作仅Host方进行。
• 解密：对于每个树结点，需要对所有特征的梯度直方图进行解密，而梯度直方图的大小与候选分裂点个数相关，因此，每个树结点的解密复杂度为，该操作仅Guest方进行。

优化1:喷砂式梯度加密

分析发现，对于每棵决策树，由于样本量通常比较巨大，样本梯度的加密操作往往耗时比较长。不仅如此，在短时间内发送全量的加密梯度，对于网关机来说也是一个很大的压力——由于公网带宽是受限的，传输加密梯度的耗时不可忽视。所以，在实际运行中，由于样本梯度的加密、传输比较耗时，Host方处于一个长时间的闲置状态（如下图的上半部分所示），造成计算资源的严重浪费。

梯度加密、发送、以及根结点梯度直方图构建的甘特图

显然，Host方闲置等待的主要原因是决策树根结点包含了全量的样本，使得创建梯度直方图需要使用全量的梯度。然而，笔者分析发现，梯度直方图的创建是一个梯度累加的过程，这个过程并不需要不同样本之间的协同计算。受到这个启发，笔者提出了喷砂式梯度加密策略。如上图的下半部分所示，该策略的主要思想是将全量样本分为多个小批量进行处理。每次Guest方对一批样本的梯度进行加密和发送，Host方接收到一批样本后，马上应用于根结点的梯度直方图构建中，并由后台线程继续下一批加密梯度的接收。该策略的优点可以汇总如下：

• 计算与计算并行、计算与通信并行——Guest方的加密操作与Host方的同态加法操作可以进行并行，Guest方（Host方）的密文发送（密文接收）与加密操作（同态加法操作）可以进行并行。
• 一般来说，用于做消息队列的网关机数量少于用于计算的Worker数量，因此，网关机的总I/O带宽是小于所有Worker的总I/O带宽的，对加密梯度进行分批发送，可以缓解网关机的带宽压力。
• 在原始的算法中，对于Host方来说，处理根结点的复杂度为，在进行优化后，假设样本量足够多，或分批足够细，渐进复杂度降低为。

在实际的训练中，笔者发现，在大部分情况下，加密和传输梯度的耗时可以被完全覆盖，不再成为瓶颈之一。

优化2:梯度直方图密文打包压缩

有加密操作就势必有解密操作。通过之前的分析，解密操作只与特征维度以及候选分裂点个数相关，与样本个数无关。尽管看似并非瓶颈，然而存在三个至关重要的问题：

• 对于决策树模型来说，树结点的个数与树的深度是成指数关系的——第层上的树结点个数为，因此，解密操作与传输梯度直方图的耗时随着树深度的增长而指数级增长；然而，同态加法的复杂度并不会随着树深的变化而变化。因此，对于较深树结点的训练，解密和通信的开销将会占很大的一部分比例。
• 通过在上一个优化点中，加密、梯度传输、同态加法可以很好的进行并行覆盖然而，梯度直方图则并非如此——梯度直方图的传输需要等待所有累加（同态加法）操作结束后才能开始，此外，同一层上树结点的梯度直方图通常是同时计算的，即使是不同树结点之间也难以做到计算与通信并行。
• 在Paillier加密中，由于解密操作涉及求解中国剩余定理等复杂的运算，单次解密的耗时通常是单次加密或同态加法的数十、甚至上百倍。

基于以上三点，加速梯度直方图的公网传输以及解密是十分必要的。

Paillier密文打包压缩 为了同时降低公网传输和解密的耗时，笔者提出了一种基于多项式计算的密文打包压缩算法，该算法的主要思想是将多个密文压缩为一个，在解密时只需对压缩后的密文进行解密。若将个密文打包压缩为一个，则跨公网的传输量和解密的操作次数均降低至原来的，是一个十分可观的提升。

在使用Paillier加密算法对浮点数进行加密前，首先需要将其编码为对应的大整数，随后对进行加密得到密文。熟悉低精度训练的读者应该了解，在机器学习训练中，研究者之所以可以使用更少的精度来表示浮点数，是由于模型或梯度中的数值通常在一个较小的范围内。类似的，由于的范围较小，在编码后，大整数也在一个较小的范围内。然而，为了进行高质量的加密，Paillier加密需要在一个很大的空间内，如在4096-bit的空间内，对大整数进行多次幂模运算。毫无疑问，这极大地造成了数据空间的膨胀。因此，笔者提出问题：是否可以通过多项式计算，对多个密文放缩操作，以填满膨胀的数据空间呢？

受到这个启发，笔者介绍一种通过多项式计算，将多个密文压缩为单个密文，并支持解密与解压的算法，具体公式如下：

其中分别表示密文空间中的同态加法和数乘运算，是足够大的放缩因子，如。该算法利用了密文空间的膨胀性质，将多个密文通过放缩拼接在一起，起到降低密文总大小和减少解密次数的作用。

应用于GBDT算法 在整数编码中，由于非负数和负数的编码范围不同，若同时对正数和负数进行放缩拼接，毫无疑问会让压缩率大打折扣。然而，在训练过程中，一阶梯度有正有负（二阶梯度非负），要如何保证压缩算法的正确性呢？

笔者分析发现，尽管样本梯度有正负，单个样本梯度是可以有范围限制的：对于分类任务使用的logistic loss，其一阶梯度范围在之间；对于回归任务使用的RMSE loss，可以对梯度大小进行约束（类似于L1正则化）。因此，每个直方图桶的大小，也必然是有范围的，即当前树结点的样本数量和单个梯度最大值的乘积。于是，PowerFL在压缩前对直方图桶通过加法进行偏移，使其保持恒正，在解压后再通过减法偏移恢复。由于这个偏移量是Guest和Host双方均已知的，因此不会造成任何信息泄露。

在实际运行中，PowerFL默认将32个密文压缩到一起，从而以较低的开销获取近32倍的跨公网传输和解密性能的提升。

优化3:计算加速技巧

除了对加密和解密进行优化，PowerFL对同态加法操作也进行了许多优化，包括以下几点：

• 高效的样本采样：除了传统的均匀采样算法，PowerFL还支持Gradient-based One Side Sampling (GOSS)采样算法，该算法可以根据样本梯度值进行更精准的采样。实验发现，对于许多数据集，每棵决策树使用50%的采样率就可以达到与全量样本训练接近的模型精度。
• 直方图作差：若两个树结点为兄弟结点，其样本集合互斥，且并集等于父结点的样本集合，因此兄弟树结点的梯度直方图之和也等于父结点的梯度直方图。PowerFL通过缓存父结点的梯度直方图，并在兄弟结点中选择样本量较小的一者进行梯度直方图构建，最后通过作差的方式得到样本量较大的结点的梯度直方图。该优化可以降低至少一半的计算开销。
• 批量JNI调用：在构建直方图时，使用一个缓存区对待累加的梯度进行缓存，直到缓存区存满再调用JNI同态加法操作，以减少JNI调用开销。
• 多线程加速：在构建直方图时，将训练样本划分为多个工作集，使用多线程加速。

实验对比

在本章节，笔者对PowerFL中的纵向联邦GBDT实现进行性能测试，这里主要考量的是腾讯自研平台相比FATE的训练效率提升，以及纵向联邦训练与非联邦训练的模型精度比较，并在最后介绍PowerFL在已有业务场景中的性能指标。

小规模数据集训练性能

首先，笔者使用两个小规模的公开数据集，对比PowerFL和FATE (SecureBoost) 的训练性能。使用的数据集，实验环境，以及超参数设置如下所示

数据集描述：

实验环境：Guest方和Host方各一个Worker，16 cores，30GB内存。

超参数设置：20棵决策树，学习率0.1，树最大深度6，样本采样率100%

下表展示了PowerFL和FATE在两个小规模数据集上的训练性能，笔者也使用XGBoost进行了模型指标对比，其中，XGBoost-Guest指的是仅用Guest方数据训练。从实验结果可以看出，

• PowerFL相比FATE，在两个小规模数据集上分别可以达到17.8倍和18.9倍的效率提升。
• PowerFL训练得到的模型，其精度与使用全量数据进行训练是相当的，均优于仅用Guest方数据训练的情况。

较大规模数据集训练性能

接下来，笔者使用三个较大规模的公开数据集进行实验。使用的数据集，实验环境，以及超参数设置如下所示

数据集描述：

实验环境：Guest方和Host方各8个Workers，8 cores，30GB内存。

超参数设置：20棵决策树，学习率0.1，树最大深度6，样本采样率20%（对于XGBoost我们使用100%采样率）

下表展示了在这三个数据集上的实验结果。

• FATE在这三个数据集上运行失败或无法在合理时间内返回结果，体现出PowerFL更强大的处理能力。
• 得益于高效的样本采样算法，PowerFL在使用低采样率的情况下，同样能达到和XGBoost全量训练相近的结果，大幅度降低了计算量。
• 纵向联邦训练和使用全量数据训练得到AUC相近，均高于仅用Guest方数据训练的情况，这充分体现了纵向联邦GBDT算法的无损性。

实际落地业务场景

最后，笔者简单介绍一下PowerFL在实际业务场景中的性能指标。如下表所示，PowerFL在可支持千万级别样本的训练场景，以及亿级别样本的预测场景，并在可观的运行时间内完成任务。

总结

本文介绍了腾讯自研联邦学习平台PowerFL中的纵向联邦GBDT算法实现，并对其中的运行瓶颈和优化技巧进行了简要的介绍。相比于开源联邦学习平台FATE，PowerFL在小规模数据集上可以达到18.9倍的性能提升，并支持更大规模的数据集，并达到与非联邦训练相同的模型精度。

参考文献

1. Jerome Friedman et al. Additive logistic regression: a statistical view of boosting. Annals of statistics 2000.
2. Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. SIGKDD 2016.
3. Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, et al. Lightgbm: A highly efficient gradient boosting decision tree. NeurIPS 2017.
4. Liudmila Prokhorenkova et al. CatBoost: unbiased boosting with categorical features. NeurIPS 2018.
5. Michael Greenwald, Sanjeev Khanna, et al. Space-efficient online computation of quantile summaries. SIGMOD 2001.
6. Kewei Cheng et al. Secureboost: A lossless federated learning framework. arXiv preprint arXiv:1901.08755.
7. Pascal Paillier. Public-key cryptosystems based on composite degree residuosity classes. International conference on the theory and applications of cryptographic techniques 1999.
8. Christine Jost et al. Encryption Performance Improvements of the Paillier Cryptosystem. IACR Cryptol. ePrint Arch. 2015.