跳转至

TreeWalker: Partial Evaluation for Grouped Tree-Ensemble Inference

会议: NeurIPS 2026
arXiv: 2610.03939
代码: https://github.com/Shopify/treewalker/tree/neurips2026-final
领域: 优化/理论
关键词: 决策树集成推理、部分求值、分组预测、GBDT、位掩码分区

一句话总结

本文提出针对分组树集成推理的部分求值引擎 TreeWalker，将组内共享的不变特征作为静态输入、变化特征作为动态输入，每棵树整组仅遍历一次并沿变化节点划分行位掩码，在保持完全精确精度的同时实现数倍至数十倍的推理加速。

研究背景与动机

梯度提升决策树（GBDT，如 LightGBM 与 XGBoost）是处理表格数据预测的主流模型。在众多实际工业预测场景中，模型需要打分的样本往往不是孤立单行，而是以具有强特征共性的“行分组（grouped rows）”形式到达。例如在离散时间生存分析中，单个患者沿时间轴离散化展开为 \(G\) 个时间步，患者基础生理特征在各行恒定不变，仅时间编码变化；在搜索排序与点击率（CTR）预估中，同一搜索 Session 内的用户画像和查询上下文对所有候选商品恒定，仅商品特征变化；在金融风险的假设情景分析（what-if scenario analysis）中，大部分基线变量保持固定，仅微调受冲击的若干特定参数。在这些任务中，组内恒定不变的特征占比通常高达 48% 至 70%。

然而，现有的高性能树推理系统主要集中于两大维度：一是优化“单次访问开销”（Per-visit cost），例如通过 LLVM IR 编译（lleaves）、C 代码生成（tl2cgen）或张量算子转化（Hummingbird）降低节点跳转解释成本；二是优化“单次遍历开销”（Per-walk cost），如 QuickScorer 系列通过叶子位向量交集提前终止遍历。现有所有推理框架均将输入矩阵拆分为独立单行逐行打分，完全忽视了组内样本在特征侧的高度冗余相关性。每一行都必须完整走一遍决策树，导致在恒定特征主干上的前缀遍历被机械重复了 \(G\) 次，造成了极大的算力浪费。

为了打破这种重复计算的瓶颈，本文引入了程序分析与编译领域的在线部分求值（online partial evaluation）思想。将每组输入中跨行共享的恒定特征视为静态绑定输入，将行间变化的特征视为动态输入，把整棵决策树的遍历转化为基于绑定时间折叠（binding-time fold）的在线特化过程。核心 idea：对每个样本组将每棵树仅遍历一次，恒定特征前缀以 \(\mathcal{O}(1)\) 复杂度直接推进，当且仅当遇到变化特征分裂点时对活跃行位掩码进行切分与空子树剪枝，从而将整组样本的遍历次数从 \(G \cdot T\) 大幅削减至逼近 \(T\)。

方法详解

整体框架

TreeWalker 面向明确具备分组语义的预测任务，直接读取标准训练完成的 LightGBM 和 XGBoost 模型，无需任何重训练或结构微调。在输入一组包含 \(G\) 个样本的观测数据（分为恒定特征子集 \(x_c\) 与变化特征子集 \(x_v\)）时，系统首先提取组内所有变化特征并完成谓词位掩码预计算；随后引擎为每一棵树初始化一个覆盖全组的活跃行位掩码（active row bitmask），从根节点出发执行带约束的遍历。

在遍历过程中，若当前节点的分裂特征属于恒定特征，则由于整组样本在该特征上的取值完全相同，判定结果唯一，算法以 \(\mathcal{O}(1)\) 开销直接步进至对应子节点；若遇到变化特征分裂点，算法利用预计算掩码或即时评估将当前活跃行掩码按左右子树判定条件切分为两个不相交的子掩码，并借助未分裂快捷路径（unsplit shortcut）剪除空掩码子树，仅对非空子集继续向下推进；最终到达叶子节点时，将叶子预测权重一次性累加至所有当前置位行的预测结果中。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入分组数据 (G 个样本)<br/>静态恒定特征 xc 与动态变化特征 xv"] --> B["谓词预计算<br/>排序变化特征+双指针扫描生成位掩码"]
    B --> C["树遍历与活跃行掩码初始化<br/>掩码置为全 1 (全组活跃)"]
    C --> D{"节点分裂特征判定"}
    D -->|恒定特征 (静态)| E["恒定步进<br/>O(1) 直接跳转至唯一确定子节点"]
    D -->|变化特征 (动态)| F["掩码切分与未分裂快捷路径<br/>拆分活跃行位掩码并剪除空子树"]
    E --> D
    F -->|仅非空子树| G["重排 DFS 递归推进<br/>轻子节点压栈递归，重子节点尾循环"]
    G --> D
    D -->|到达叶节点| H["叶节点批量累加<br/>对掩码置位样本并行累加叶子权重"]
    H --> I["输出整组 G 个样本最终预测值"]

关键设计

1. 恒定步进与位掩码切分：以单次树遍历替代逐行冗余计算

传统树遍历将每个样本独立送入根节点向下追踪，导致由恒定特征构成的公共路径被重复遍历 \(G\) 次。TreeWalker 为整组 \(G\) 个样本在当前遍历树中维护一个统一的活跃行位掩码（对于 \(G \le 64\) 使用单字 u16/u32/u64，宽组则使用 64 位数组）。从根节点出发，当节点分裂特征为恒定特征时，所有活跃样本严格共享相同的分支走向，算法无需评估逐行状态，直接以 \(\mathcal{O}(1)\) 时间推进至目标子节点；当且仅当遇到变化特征分裂节点 \(u\) 时，算法调用分区函数将当前活跃掩码 \(A(u)\) 拆分为两个互斥的子掩码 \((A_L(u), A_R(u))\)，分别携带各样本分配至左右子树。这一机制将整棵树在恒定特征上的投影子图 \(T_c(\tau, X)\) 压缩为单次共享执行，从根本上消除了输入相关性带来的计算冗余。

2. 未分裂快捷路径与重路径 DFS：消除无效调用并保持缓存局部性

在变化特征节点处进行位掩码切分时，组内样本可能全部满足或全部不满足分裂阈值，导致某一子树的位掩码为空。朴素递归会盲目进入空掩码子树造成冗余开销。TreeWalker 设计了未分裂快捷路径（unsplit shortcut）：一旦检测到 \(m_L = 0\) 或 \(m_R = 0\)，算法直接将节点指针更新为非空子节点并原地循环，彻底避免递归进入空子树，在理论上消除了所有 \(U(\tau, X) = V(\tau, X) - B(\tau, X)\) 次无效子树访问。此外，系统在模型解析阶段采用重路径深度优先遍历（Heavy-Path DFS）对节点进行连续紧凑布局（每个节点占用 16 字节，包含阈值、跳转偏移、特征索引与紧凑标志位）。将训练权重较大（样本量或 Hessian 综合更高）的“重子节点”直接紧邻存放在内存下标 idx + 1 处；在遍历遇到双侧非空分支时，优先递归进入轻子节点，而对重子节点采用尾调用消除（tail-call elimination）原地循环推进，确保高频执行路径在 CPU 缓存中维持高度连续的顺序扫描。

3. 谓词预计算机制：解耦特征比较与全林遍历

在整片森林包含大量决策树时，相同的变化特征切分阈值会在不同树中反复出现。如果每次访问变化节点都重新逐行比较数值，开销将随树数 \(T\) 线性膨胀。TreeWalker 的默认执行路径引入了森林级谓词预计算（Predicate Precomputation）：在解析阶段提取整片森林中每个变化特征 \(f \in F_v\) 的所有去重谓词阈值集合 \(\Pi_f\)；在预测阶段，针对当前组的 \(G\) 个数值仅执行一次快速排序（复杂度 \(\mathcal{O}(G \log G)\)），随后利用双指针归并扫描以均摊 \(\mathcal{O}(1)\) 的成本直接生成每个谓词对应的 \(G\) 位输出掩码。在后续树遍历过程中，变化节点的切分操作退化为简单的按索引查表与按位与（bit-level AND）操作，使得整片森林的谓词供给开销 \(Q_{\text{pre}}(\mathcal{F}, X) = \mathcal{O}\left(\sum_{f \in F_v} (G \log G + |\Pi_f^{\text{num}}| + G \cdot |\Pi_f^{\text{cat}}|)\right)\) 仅在组级别支付一次，与树的总数 \(T\) 完全解耦。

4. 结构化工作量分解定理与理论渐近界

为了精确刻画算法收益，论文在单位开销模型下建立了树集成结构化工作量分解定理（Theorem 1）。对于任意树 \(\tau\) 与大小为 \(G\) 的组 \(X\)，遍历总开销 \(W(\tau, X)\) 严格等于恒定节点访问数 \(C(\tau, X)\)、变化节点访问数 \(V(\tau, X)\)、叶节点写入数 \(G\) 以及谓词掩码供给开销 \(Q(\tau, X)\) 之和，且受恒定投影子树规模 \(|T_c(\tau, X)|\) 有界约束：

\[W(\tau, X) = C(\tau, X) + V(\tau, X) + G + Q(\tau, X) \le |T_c(\tau, X)| + G + Q(\tau, X)\]

在此基础上，推论 1 针对即时追踪评估器（Trace Evaluator，即无预计算版本，\(Q_{\text{trace}} = G \cdot d_v\)）证明：相对于基线逐行完全遍历 \(W_{\text{full}} = G(d + 1)\)，归一化工作量比值满足：

\[\frac{W_{\text{trace}}(\tau, X)}{G(d + 1)} \le \frac{|T_c(\tau, X)|}{G(d + 1)} + \frac{d_v + 1}{d + 1}\]

随着组规模 \(G \to \infty\)，结构项 \(\frac{|T_c|}{G(d+1)} \to 0\)，相对单行完全遍历的渐近加速比收敛至理论极限 \(\frac{1}{1 - \bar{\rho}_G} = \frac{\mathbb{E}[d + 1]}{\mathbb{E}[d_v + 1]}\)（其中 \(d\) 为平均树深，\(d_v\) 为路径上变化特征节点数）。而在默认预计算模式下，单行计算开销在 \(G\) 增大时迅速逼近叶节点写入的物理底线 \(1/L\)。

一个完整示例

考虑一个由 4 个样本组成的行组（\(G=4\)），初始活跃位掩码为 [1111]（二进制 0b1111）。样本共享恒定特征 \(x_1 = 2\)，但具有变化特征 \(x_5 = [2, 4, 7, 9]\)。 1. 根节点判定：根节点划分条件为 \(x_1 \le 3\)。由于 \(x_1\) 为恒定特征且 \(2 \le 3\) 成立，整组 4 个样本无歧义地全部走向左子节点，掩码 [1111] 保持不变，耗时 \(\mathcal{O}(1)\)；相比之下标准逐行推理必须重复执行 4 次浮点比较。 2. 变化节点切分：左子节点划分条件为 \(x_5 \le 5\)。由于 \(x_5\) 为变化特征，查阅预计算掩码或即时评估后，样本 0 和 1 满足条件，样本 2 和 3 不满足。掩码被切分为左子掩码 \(m_L = \text{0b0011}\)（第 0, 1 行）与右子掩码 \(m_R = \text{0b1100}\)（第 2, 3 行）。 3. 未分裂快捷检测：左右掩码均非空，算法未触发未分裂剪枝。根据重路径权重，算法递归向下处理轻子树（携带 \(m_L\)），随后尾循环处理重子树（携带 \(m_R\)）。 4. 叶节点累加：左侧路径到达权重为 \(+0.8\) 的叶节点，遍历活跃位为 0 和 1，对 results[0] 和 results[1] 累加 \(+0.8\)；右侧路径到达权重为 \(-0.4\) 的叶节点，对 results[2] 和 results[3] 累加 \(-0.4\)。

实验关键数据

主实验

评估在 Google Compute Engine 专属实例（Intel Xeon Platinum 8581C 与 Arm Neoverse V2）上运行，覆盖基准配置（\(T=500, L=8\)；SUPPORT 与 FLCHAIN 取生存分析中位数 \(G=16\)，Expedia 采用真实 Session 规模 \(G=5\sim 37\)），对比原生 LightGBM、原生 XGBoost、C 编译代码生成工具 tl2cgen、LLVM 编译工具 lleaves 以及同数据布局的行独立完全遍历（Full Walk）。

数据集 任务类型 组规模 \(G\) TreeWalker 延迟 (\(\mu\text{s}\)) 原生 LightGBM (\(\mu\text{s}\)) 原生 XGBoost (\(\mu\text{s}\)) tl2cgen (\(\mu\text{s}\)) 对 Full Walk 算法加速比 对原生框架系统加速比
SUPPORT 生存分析 16 63.8 453.2 574.6 146.7 2.9× 7.1× (LGBM) / 9.0× (XGB)
FLCHAIN 生存分析 16 72.4 481.5 612.0 166.5 3.0× 6.6× (LGBM) / 8.5× (XGB)
Expedia CTR/排序 经验均值 (13.8) 168.0 487.2 1015.0 198.3 2.7× 2.9× (LGBM) / 6.0× (XGB)
Credit (k=1) 情景分析 128 18.2 - - - 15.2× (Intel) / 9.2× (Arm) 理论界一致 (误差 < 3.7%)

注：在更宽的组规模下（\(G=128\)），TreeWalker 在 SUPPORT 和 FLCHAIN 上的算法加速比进一步攀升至 7.7×–9.1×（Intel）和 5.5×–6.0×（Arm）；在全量 544 个网格单元中，最大端到端系统加速比达到 54×（\(T=2000, L=8, G=128\) 于 SUPPORT 数据集）。

消融实验

论文在 \(T=500, L=8\) 基准配置下，对 TreeWalker 的 8 项核心功能开关进行了逐一禁用的消融测试（报告禁用后相对于完整配置的延迟劣化比例）：

配置 / 禁用开关 SUPPORT (G=16) 延迟增加 FLCHAIN (G=4) 延迟增加 Expedia (Session) 延迟增加 说明与机制分析
完整引擎 (Full) 0.0% (基准) 0.0% (基准) 0.0% (基准) 启用所有静态与动态优化
禁用谓词预计算 (Precompute) +21% ~ +39% +11% +93% ~ +139% 退化为单树逐行即时计算，高基数及深树场景严重掉速
禁用谓词去重 (Predicate Dedup) +28% ~ +34% +23% ~ +27% +77% ~ +82% 无法跨树共享相同切分阈值的预计算位掩码
禁用未分裂快捷路径 (Unsplit Shortcut) +12% ~ +18% +8% ~ +11% +22% ~ +29% 盲目递归进入空位掩码子树，增加额外控制流分支
禁用树重排序 (Tree Ordering) +3% ~ +8% +3% ~ +6% \(\le\) +1% 丧失跨树缓存友好度与前缀公共性
禁用单调提前终止 (Monotone Early-Term) -3% ~ +7% -2% ~ +4% -1% ~ +3% 仅影响即时追踪模式，受分支预测器干扰效果有正有负

关键发现

  • 组规模 \(G\) 的单调增益性：算法加速比随组宽 \(G\) 呈严格单调上升趋势，从 \(G=1\) 时的约 \(1.0\times\) 稳步上升至 \(G=128\) 时的 \(7.7\times \sim 9.1\times\)；在 \(G > 128\) 的超宽分组（最高 1008 行）中，整组处理速度可达同布局单行基线的 \(17.1\times\)。
  • 树深度对收益的敏感度：浅树模型（\(L=2\)）展现出比深树模型（\(L=16\)）更高的相对加速比（\(3.9\times \sim 4.7\times\) 对比 \(1.9\times \sim 2.6\times\)），原因在于浅树中恒定特征前缀占据整条路径的比例显著更高。
  • 数值精度优势：在 64 位浮点（LightGBM）下，TreeWalker 在所有 3020 万行数据上均与 Treelite GTIL 严格保持 0-ULP 比特级一致；在 32 位浮点（XGBoost）下，TreeWalker 采用 64 位累加叶子权重，在全部 117,096 行测试样本中均比 XGBoost 原生 32 位累加更加贴近数学精确参考值，消除了级数求和截断误差。

亮点与洞察

  • 将程序特化编译原则创新引入树模型推理：敏锐捕捉到工业界分组预测任务（生存模型、排序检索、假设检验）中输入侧高度共性的物理事实，首次将程序分析中的部分求值范式实例化为行位掩码的在线下推传播。
  • 正交且可组合的技术维度：现有加速方案聚焦于“单次访问编译”与“单次遍历位向量交集”，而 TreeWalker 独辟蹊径开辟了“遍历次数缩减（Number-of-walks）”这一全新维度，在理论上完全可与硬件向量化（SIMD）和底层即时编译（JIT）无缝级联。
  • 完备严谨的工作量分解理论闭环：不仅提出了工程实现，更给出了具有恒等式意义的单位开销工作量分解定理，并通过硬件计数器严格验证了渐近线收敛性，使理论预测与工程实测高度统一。

局限与展望

  • 依赖预先定义的分组语义边界：当前实现假定输入批次已经具备明确的分组标识与恒定/动态特征边界，不具备从杂乱无序的任意 Batch 中自动聚类发掘共享常量子空间的能力。
  • 极小分组场景的常数开销：在 \(G=1\) 或 \(G=2\) 的极小样本情形下，位掩码分配与谓词预计算的管理开销会抵消路径复用收益，可能导致微弱性能回退（约 1% 至 20%）。
  • 单线程 CPU 绑定与架构拓展：目前核心引擎为针对单线程低延迟 CPU 服务设计，尚未扩展至 GPU 等大规模并行架构；对于部分重叠常量的层级分组（hierarchical subgrouping）尚需更精细的动态掩码追踪机制。

相关工作与启发

  • vs lleaves / tl2cgen: 后两者在解析时（parse-time）将树模型固化编译为 LLVM IR 或 C 动态库，优化的是单次节点访问的指令开销；TreeWalker 则在预测时（predict-time）针对输入组的常量模式特化决策路径，优化的是遍历总次数。两者分别对应外层控制流与内层计算核，具有天然互补性。
  • vs QuickScorer / RapidScorer: QuickScorer 家族通过叶子位向量交集在单棵树内短路遍历，但在深树（\(L \ge 8\)）、存在缺失值路由或高基数类别特征时无法通用支持；TreeWalker 在标准特征切分上进行样本维度的位掩码切分，对复杂缺失值与类别特征具有完整的原生支持。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将部分求值与位掩码分区机制引入决策树集成推理，开辟了缩减遍历次数的正交新维度。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 大数据集、跨 x86 与 Arm 架构、712 个测试单元、5 大外部基线以及严格的硬件计数器定理对账。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论定义清晰，从算法伪代码、数据结构布局到工作量分解定理与推论层层递进，论证严谨自洽。
  • 价值: ⭐⭐⭐⭐⭐ 对于生存分析、实时搜索推荐等重度依赖树模型分组打分的云端服务，具备立竿见影的吞吐提升与计算降本价值。

相关论文