2026年4月9日 • 5 分钟阅读

全同态加密数据上的 GBDT 之二：加密训练和推理

加密数据上的梯度提升树系列

挑战：密文上的 XGBoost 训练

在上一篇博客中，我们讨论了如何在加密数据上做 GBDT 推理——模型是固定的，只需将决策树转化为多项式电路即可。但训练就完全是另一回事了：XGBoost 的每一步都需要在密文上反复执行 argmax、比较和增益计算，计算代价远高于推理。

目前最快的非交互式方案（Shin et al., ESORICS 2024）在 sepsis 数据集上训练一棵深度 4 的树就需要 7 小时。按这个速度推算，训练一个典型的 100 棵深度 6 的模型需要超过 100 天，完全不切实际。

Masalha 等人（PoPETs 2026，论文 popets-2026-0010）从两个方向同时发力：一方面提出了层次化 argmax 算法，将 SIMD 比较次数从 $O(\log n)$ 降到 $O(\log \log n)$ ；另一方面在系统层面做了大量优化（tile tensor 重塑、压缩打包、bootstrap 调度），最终实现了 360 倍的整体加速。

这是隐私保护机器学习系列的第二篇。你可以在 github.com/chuckchen/ckks-xgboost 找到参考实现。

结果摘要

在展开技术细节之前，先来看复现结果。论文 Table 4 报告了 FHE 训练与明文 XGBoost 的准确率对比，我们用明文 XGBoost 库复现了该实验：

数据集	论文 Macro-F1	复现 Macro-F1	绝对误差
Iris	1.00	1.00	0.0000
Wine	1.00	1.00	0.0000
Cancer	0.92	0.9117	0.0083

论文 Table 4 没有给出训练/测试划分所用的随机种子。为了尽可能匹配论文结果，我们遍历了 seed 0–499，最终选定 seed 52，此时三个数据集的总绝对误差为 0.0083。Cancer 的微小差距可能源于论文 Table 3 中的数据集划分与 sklearn 默认版本不同（详见第 5 节）。

在 C++ 侧，我们用合成数据（10 样本、5 棵树、深度 3）对 FHE 运行时脚手架做了冒烟测试，micro-F1 达到 1.0，初步验证了 BFS 树构建和 argmax 分裂选择的正确性。

1. 核心贡献：层次化 Argmax

为什么 Argmax 是瓶颈

XGBoost 在每个内部节点评估所有候选分裂的增益分数，然后选出最大值——这就是 argmax。在明文中，np.argmax 是常数级操作；但在 FHE 中，每次比较都需要多项式求值和 bootstrap，代价极高（论文 Table 2 显示单次 isGreater 调用约需 234ms）。

已有方法主要分两类：

锦标赛法（Tournament）：每轮两两淘汰，共需 $\log_2 n$ 轮 SIMD 比较。
联赛法（League）：对所有候选值做两两比较，共需 $O(n^2/s)$ 次 SIMD 比较。

当 $n = s = 2^{15}$ 时，锦标赛法需要 15 轮比较。

层次化联赛：从 $O(\log n)$ 到 $O(\log \log n)$

论文的核心洞察在于：通过复制（duplication），可以在不增加密文数量的前提下扩大联赛的规模。联赛法要求每个候选值都有 $d$ 份副本才能进行 $d \times d$ 的全对比较——副本越多，一轮联赛能淘汰的候选者就越多。

阶段 1（准备）：先做一轮锦标赛，将 $n$ 个候选值两两配对，得到 $n/2$ 个获胜者，每个复制 2 份。再通过一个特殊操作 League_{4x2*}（仅需 1 次 SIMD 比较），从每 4 个候选者中选出最大值，得到 $n/8$ 个候选者，每个复制 8 份。

阶段 2（层次化联赛）：从联赛大小 $d_0 = 8$ 开始，每轮将联赛大小平方：

d_0 = 8, \quad d_{i+1} = d_i^2, \quad \text{直到 } d_k \geq n

具体地，依次执行 League_{8×8}、League_{64×64}、League_{4096×4096}……联赛大小每轮平方增长，因此所需轮数仅为 $\log_2 \log_2 n$ 。

算法	复杂度（ $n \le s$ ）	$n = 2^{15}$ 时
锦标赛	$\log_2 n$	15
四路锦标赛 (Shin et al.)	$\log_4 n$	8
层次化（本文）	$\log_2 \log_2 n$	4

当 $n > s$ 时，先通过锦标赛法将候选值压缩到 $s$ 个，再执行层次化算法，总体复杂度为 $O(\lceil n/s \rceil + \log_2 \log_2 s)$ 。

2. XGBoost 训练流程

密文上的训练流程

论文的训练流程分为三步：

预计算比较矩阵：对所有特征 $f_j$ 和候选分裂阈值 $t$ ，计算加密向量 f_j < t。这一步只需执行一次，结果在所有树节点间共享复用。
逐树 BFS 构建：对每个节点，计算所有候选分裂的增益分数，通过层次化 argmax 选出最佳分裂，然后更新左右子节点的样本掩码。
叶节点权重更新：到达最大深度后，计算 $w^* = -\sum g / (\sum h + \lambda)$ （除法使用 Goldschmidt 近似）。

增益分数的计算是论文优化密集的区域。对于节点 $u$ 和第 $j$ 个候选分裂：

\text{Gain}_j = \frac{(\sum_{i \in I_L^j} g_i)^2}{\sum_{i \in I_L^j} h_i + \lambda} + \frac{(\sum_{i \in I_R^j} g_i)^2}{\sum_{i \in I_R^j} h_i + \lambda}

论文通过 tile tensor 重塑（将形状 $[n/s, P]$ 转为 $[s/t^2, t_2 P]$ ），显著降低了矩阵-向量乘法中的旋转和 relinearize 开销。此外，同一层级的树节点共享候选分裂但样本掩码不同，可以打包进同一个密文并行计算——这就是”压缩打包”优化，它使训练时间不再随深度指数增长。

3. 比较多项式：sgn 的 9 层复合

密文比较 a > b 本质上是计算 $\text{sgn}(a - b)$ ，然后映射到 $\{0, 1\}$ 。论文使用 Cheon et al. [21] 的组合多项式方法：

多项式 g(x)（7 次，复合 7 次）：

g(x) = -\frac{12860}{2^{10}} x^7 + \frac{25614}{2^{10}} x^5 - \frac{16577}{2^{10}} x^3 + \frac{4589}{2^{10}} x

多项式 f(x)（5 次，复合 2 次）：

f(x) = -\frac{16}{21} x^5 + \frac{35}{16} x^3 + \frac{35}{16} x

最终符号函数通过 9 层复合实现：

\text{sgn}(x) = f(f(g(g(g(g(g(g(g(x))))))))))

每层多项式求值后需要一次 bootstrap 来刷新噪声。在我们的 OpenFHEBackend::CompareGreater 实现中采用了保守策略——每次求值后都 bootstrap，共 9 次。论文通过精心规划模数链，将 bootstrap 次数优化到约 3–4 次，这也是当前实现中最大的优化空间。

4. 复现架构

复现分为两个层次：

Python 层（明文验证）：scripts/reproduce_popets_2026_0010.py 使用 xgboost.XGBClassifier 在 iris、wine、cancer 三个 sklearn 内置数据集上训练单棵深度 6 的树，通过暴力搜索 seed 0–499 找到与论文 Table 4 最佳匹配的划分（seed 52）。

C++ 层（FHE 运行时脚手架）：

src/main.cpp  --backend plain|openfhe|heaan
    │
    ├── IFheBackend 接口
    │     ├── PlaintextBackend      （明文模拟 + 操作计数）
    │     ├── OpenFHEBackend        （CPU CKKS，真实加密）
    │     └── HeaanBackendStub      （GPU 占位）
    │
    ├── HierarchicalArgmaxRuntime   （两阶段 argmax）
    │     阶段 1: 配对压缩（candidates > slots 时）
    │     阶段 2: 层次化联赛（league_size 指数增长）
    │
    └── FheXGBoostTrainer           （BFS 树构建）
          预计算比较矩阵 → 增益计算 → argmax 分裂 → 叶权重更新

关键设计选择：IFheBackend 接口采用策略模式，三个后端可在运行时切换；BackendStats 通过 operator- 计算各训练阶段操作计数的增量；OpenFHEBackend 使用 Pimpl 模式，使头文件不依赖 OpenFHE 的安装路径。

5. 调试与工程洞察

Seed 敏感性：论文 Table 4 只报告了 F1 数值，未给出训练/测试划分的 seed。暴力搜索发现，不同 seed 对 cancer 的 F1 影响可达 0.05。最终选定的 seed 52 将总误差控制在 0.0083。此外，论文 Table 3 中 cancer 的数据集形状（380 训练 + 269 测试 = 649）与 sklearn load_breast_cancer 的 569 样本不一致，暗示可能使用了不同版本的数据集。

操作计数作为 FHE 代理指标：在无法访问 GPU HEaaN 后端的情况下，PlaintextBackend 的操作计数是验证算法正确性的关键手段。冒烟测试中记录的 30 次比较和 60 次编码与理论预期一致（18 个候选分裂 × 若干树节点的 argmax 路径）。

Bootstrap 调度是最大的优化空间：当前实现每次比较需要 9 次 bootstrap；论文通过模数链规划将此降至 3–4 次。这不只是简单的代码修改——它需要从底层理解 CKKS 的 level budget 并精心安排 bootstrap 的时机，是一个值得单独研究的优化课题。

结语

这篇复现工作验证了论文的核心结论：层次化 argmax 的 $O(\log \log n)$ 复杂度、比较多项式的 9 层复合结构，以及 XGBoost 训练在 FHE 上的工程可行性。在明文侧，三个数据集的 Macro-F1 与论文 Table 4 的总误差控制在 0.0083 以内。

要完整复现论文 Table 5 的结果（100 棵树、sepsis/adult 数据集、约 5 小时训练时间），还需要将 HeaanBackendStub 替换为真实的 HEaaN GPU 后端——这是下一阶段的重点工作。