Very Fast Bayesian Additive Regression Trees on GPU¶

会议: NeurIPS 2026
arXiv: 2410.23244
代码: https://github.com/Gattocrucco/bartz
领域: 优化算法 / Optimization
关键词: 贝叶斯可加回归树, GPU加速, JAX, 无分支MCMC, 树集成

一句话总结¶

针对贝叶斯可加回归树（BART）在传统 CPU 上因树结构动态变化和 MCMC 串行采样导致难以处理大规模样本的瓶颈，本文基于 JAX 提出了一种定长二叉堆矩阵表征与无分支（branchless）向量化算法，在保持完整贝叶斯推断的同时在 GPU 上实现最高 200 倍加速，使 BART 的运行时间可与 XGBoost 正面匹敌。

研究背景与动机¶

贝叶斯可加回归树（Bayesian Additive Regression Trees, BART）是统计学与机器学习领域经典的非参数回归方法。它将未知回归函数建模为多个浅层二叉决策树的累加和，并通过先验分布控制单棵树的复杂度，再借助 Metropolis-within-Gibbs 采样从后验分布中抽取树结构与叶节点参数。得益于先验对过拟合的天然正则化，BART 在因果推断比赛（如 ACIC Benchmark）和常规表格数据回归中通常具有优于随机森林及未经精调 boosting 的预测精度，同时天然具备开箱即用的不确定性量化能力。

然而，在产业界和大规模数据竞赛中，BART 远不如以 XGBoost 为代表的梯度提升树（GBDT）流行。核心矛盾在于两者的计算可扩展性差距悬殊：XGBoost 采用贪心策略和高度优化的直方图并行算法，可轻松扩展到上亿样本；而经典 BART（如 dbarts、bartMachine）依赖复杂的 MCMC 采样，每轮迭代都在动态调整树的指针拓扑结构，不仅在多层条件分支（if-else）下无法向量化，更难以利用 GPU 等现代高吞吐并行硬件。现有文献中 BART 分析的最大规模通常仅在数十万样本以内，面对海量数据时耗时极其漫长。

为了兼顾精度与速度，已有工作如 XBART 虽通过牺牲完全贝叶斯性质将运行效率提升了 20-30 倍，但失去了严谨的先验后验框架与灵活扩展性。本文的切入点在于重新审视 GPU 硬件对算法结构的约束：只要为树深度设定合理的显式上限，BART 繁杂的动态树更新与样本叶节点查找就能完全映射为稠密矩阵算子与无分支算术指令。核心 idea：通过固定最大树深度的二叉堆矩阵表达与全局无分支遍历算法，将 BART MCMC 的拓扑提案与统计量归约转化为 JAX 编译的高并发 GPU 算子，在完全保留贝叶斯后验的前提下达成最高 200 倍单核加速。

方法详解¶

整体框架¶

BART 模型将目标变量标量观测 $y_i$ 建模为 $m$ 棵独立回归树输出之和加上高斯噪声： $$y_i = \sum_{j=1}^{m} g(\mathbf{x}_i; T_j, M_j) + \varepsilon_i, \quad \varepsilon_i \overset{\text{i.i.d.}}{\sim} \mathcal{N}(0, \sigma^2)$$ 其中 $T_j$ 代表第 $j$ 棵树的二叉拓扑结构与划分规则（划分特征轴与分割点），$M_j = \{\mu_{j,1}, \ldots, \mu_{j,b_j}\}$ 代表对应叶节点的值。后验推断采用 Metropolis-within-Gibbs 算法轮流更新单棵树结构 $T_j$、叶节点取值 $M_j$ 以及全局误差方差 $\sigma^2$。

在传统实现中，树结构基于动态链表或指针对象存储，样本下推通过递归条件跳转完成。本文算法（在开源 Python 包 bartz 中实现）利用 JAX 框架，将整个过程改造为基于定长张量的流式计算：所有特征预先量化为 1 字节整数索引以节省显存带宽；树拓扑压平为紧凑的二叉堆；通过无分支迭代统一所有样本的下推步数；将可跨树并行的结构提议、样本计数与方差计算先在 GPU 上一并完成，最后仅保留按顺序串行规约叶节点残差与更新全局残差缓存的必要步骤。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入特征量化与张量初始化<br/>特征离散为单字节索引，分配定长堆矩阵与缓存"] --> B["跨树并行独立提案<br/>批量生成所有树的 GROW/PRUNE 候选结构与随机正态数"]
    B --> C["无分支树遍历与叶索引更新<br/>固定深度循环消除分支，更新 n×m 节点归属缓存"]
    C --> D["逐树串行残差规约与 Metropolis 步<br/>跨样本聚集叶内残差并判定接受率，更新连续叶值"]
    D --> E["全局残差增量广播与方差采样<br/>更新缓存残差，基于残差平方和共轭采样误差方差 σ²"]

关键设计¶

1. 定长二叉堆矩阵表征：消除动态指针内存分配与内存碎片 传统 C++/Java 树实现为每个节点分配包含左右子节点指针、切分维度、阈值及元数据的对象（在 BART/dbarts 中单节点开销高达 48 至 64 字节），节点随增删动态分配，内存跳跃严重破坏 GPU 访存局部性。本文通过固定树的最大深度 $D$（默认设为 $D=6$），将每棵树表示为定长的满二叉树二叉堆（Heap）。整个森林由 3 个固定维度的矩阵表示：切分轴矩阵、切分点索引矩阵和叶节点值矩阵。由于最大深度为 $D$ 的树最多具有 $2^D - 1$ 个节点，其中内部决策节点最多 $2^{D-1} - 1$ 个，叶节点最多 $2^{D-1}$ 个，在 32 位浮点与整型下单棵树的全部表征开销仅为 $2^{D+3}$ 字节（$D=6$ 时仅 512 字节）。空节点或非活跃节点直接填充 0 并在计算中保持哑操作，彻底排除了 GPU 上的动态显存分配开销。

2. 无分支树遍历与单字节特征量化：消除控制流分歧并极大压缩显存占用 在 GPU 上执行分支跳转（if-else）会导致同一个 Warp/Thread Block 内的线程产生严重的分歧（branch divergence）。为了将样本寻找所属叶节点的操作完全向量化，本文设计了无分支的定步长遍历循环（Branchless Traverse）。对每一个输入样本，循环严格执行 $D-1$ 步，维护一个布尔标志指示是否已触达叶节点，并通过位运算和算术选择动态计算子节点堆索引： $$\text{child\_index} = 2 \times \text{index} + \mathbb{I}(x_{\text{axis}} \ge \text{split})$$ 若已处于叶节点，则保持索引不变。此外，BART 在连续特征轴上的候选分割点通常不超过 255 个（默认设置为 100 个），本文将输入矩阵 $X$ 的所有浮点特征映射为分割网格索引，每个特征仅消耗 1 个字节（uint8）。配合大小为 $n \times n_{\text{tree}}$ 字节的样本叶节点索引缓存，整个算法的核心内存消耗严格控制在 $n(p + n_{\text{tree}})$ 字节，较 dbarts 的 $\approx 24 n(p + n_{\text{tree}})$ 字节降低了逾 20 倍。

3. 跨树并行解耦与残差规约流水线：最大化 GPU 算力利用 BART 的标准 Gibbs 采样要求在更新第 $j$ 棵树时，必须依赖由其余 $m-1$ 棵树计算得出的当前残差 $R_i = y_i - \sum_{k \ne j} g(\mathbf{x}_i; T_k, M_k)$，呈现天然的树间串行性。为了最大化挖掘 GPU 硬件并发度，本文将每轮 MCMC 严格拆分为“跨树并行阶段”与“树间串行阶段”。在跨树并行阶段，算法同时对全部 $m$ 棵树生成随机 GROW/PRUNE 提案、无分支更新假想最大拓扑下的叶索引分布、预计算落入各叶节点的样本数，并生成后续抽样所需的所有标准正态和均匀随机变量，这些计算与残差取值完全独立。在树间串行阶段，唯一的计算瓶颈在于对样本残差进行按叶索引规约求和（indexed reduce）： $$\sum_{i: \mathbf{x}_i \in \text{Leaf}_l} R_i$$ 由于该操作具有 $O(n)$ 复杂度，当样本量 $n$ 足够大时，GPU 的流多处理器（SM）能被充分打满；在 Metropolis 接受/拒绝并重抽样叶值后，通过广播叶值增量更新全局残差缓存向量，避免全量重新预测。

一个完整示例¶

以单棵树的最大深度限制为 $D=3$、样本量为 $n$ 为例，说明单步无分支遍历与残差更新机制： 1. 二叉堆编号：根节点编号为 1，第 2 层为 2（左）、3（右），第 3 层潜在叶节点为 4、5、6、7。若当前树仅在根节点切分（左侧为叶，右侧在节点 3 继续切分），则节点 2 的切分点值标记为 0（代表叶节点），节点 3 记录具体的特征轴与切分阈值。 2. 定步长遍历：所有样本首先初始化 index = 1, leaf_found = False。 - 循环第 1 步：读取根节点规则，若 $x_i \ge \text{split}_1$ 则 child_index = 3，否则 child_index = 2。更新 index。 - 循环第 2 步：对于落入节点 2 的样本，由于其 split_2 == 0，标志位 leaf_found 翻转为 True，其索引被强制锁定在 2；对于落入节点 3 的样本，leaf_found 仍为 False，根据特征对比继续下推至节点 6 或 7。循环在第 $D-1=2$ 步统一终止，无任何线程提前退出或跳转。 3. 残差累加与更新：各样本依据最终的 index（取值为 2、6 或 7）并行原子累加残差值，仅需一次线性的按索引归约即可完成该树叶节点充分统计量的计算。

损失函数 / 训练策略¶

作为严谨的贝叶斯方法，BART 不依赖梯度下降损失优化，而是基于先验与后验共轭性进行 MCMC 采样： - 树结构先验：深度为 $d$ 的节点成为内部节点的先验概率为 $P_d = \alpha / (1+d)^\beta$（默认 $\alpha=0.95, \beta=2$），天然惩罚深树。 - 叶节点值采样：叶参数先验为独立正态分布 $\mu \sim \mathcal{N}(\mu_\mu, \sigma_\mu^2)$。在给定树结构与残差下，叶节点均值参数的条件后验仍为标准正态分布，可直接解析采样。 - 误差方差采样：先验为逆卡方分布 $\nu\lambda/\sigma^2 \sim \chi^2_\nu$（默认 $\nu=3$），在遍历完所有 $m$ 棵树后，直接根据全样本残差平方和从逆伽马分布中精确抽取新的 $\sigma^2$。 - MCMC 运行设置：默认执行 2000 轮 MCMC 迭代并丢弃前 1000 轮作为 burn-in，采用多链（通常 $\ge 2$ 条）监测后验函数值的一致性。

实验关键数据¶

主实验¶

论文通过合成基准测试评估了本文实现 bartz 在单核 CPU 及 Nvidia A100 GPU 上的单轮 MCMC 耗时，并与目前最快的 C++ 实现 dbarts（单核 M1 Pro CPU）以及 XGBoost（Nvidia L4 GPU）进行端到端对比。在常见低维设置（$p=100, n_{\text{tree}}=200$）下不同样本规模 $n$ 的单轮运行时间如下所示：

样本规模 $n$	dbarts (CPU, ms)	bartz (CPU, ms)	bartz (A100, ms)	XGBoost (L4, ms/fit)	bartz-A100 加速比 (vs bartz-CPU)
$10^3$	0.45	1.80	1.10	2.50	1.6×
$10^4$	4.20	5.10	1.12	6.80	4.6×
$10^5$	41.0	43.0	1.25	38.0	34.4×
$10^6$	420.0	435.0	2.15	280.0	202.3×
$3 \times 10^6$	1280.0	1310.0	6.40	850.0	204.7×

(注：BART 需要约 1000 轮 MCMC 产出后验样本，单轮耗时 2.15 ms 意味着 1000 轮总耗时仅约 2.15 秒；XGBoost 单次拟合耗时约 0.28 秒，若计入超参数网格调优及交叉验证的数十至上百次拟合，BART 在数百万样本下的端到端训练耗时已具备完全竞争优势。)

消融实验¶

为了分析特征维度 $p$ 与树数量 $n_{\text{tree}}$ 对 GPU 显存限制及加速瓶颈的影响，论文对比了固定规模与随样本动态缩放（$n/p=10, n/n_{\text{tree}}=8$）下的显存占用与最大可达加速比：

配置维度模式	参数设置	峰值显存占用模式	显存溢出临界点 ($n$)	最大有效加速比 (vs CPU)	性能特征与瓶颈说明
低维基准 (Low-Low)	$p=100, n_{\text{tree}}=200$	$\approx n(p+n_{\text{tree}})$ 字节	$> 5 \times 10^6$	200×	显存极为充裕，在大 $n$ 时残差规约打满 GPU 并发
高维特征 (High-p, Low-$n_{\text{tree}}$)	$p=n/10, n_{\text{tree}}=200$	受限特征索引矩阵	$\approx 3 \times 10^4$	~30×	特征矩阵显存迅速耗尽，GPU 加速尚未完全释放
大森林 (Low-p, High-$n_{\text{tree}}$)	$p=100, n_{\text{tree}}=n/8$	受限叶索引缓存	$\approx 2.5 \times 10^4$	~25×	节点索引缓存 $n \times n_{\text{tree}}$ 超出单卡容量上限
极端高维 (High-High)	$p=n/10, n_{\text{tree}}=n/8$	双重高阶缓存	$\approx 8 \times 10^3$	~15×	极小 $n$ 下即 OOM，但验证了后验采样一致性

(在测试集规模为 1000 的非线性交互预测验证中，bartz 与 dbarts、BART、bartMachine 的测试集 RMSE 几乎完全重合，证实无分支二叉堆表征未损害任何统计推断精度。)

关键发现¶

GPU 吞吐的拐点效应：当 $n \le 10^4$ 时，由于 GPU 调度开销和向量化利用不充分，bartz-A100 的耗时几乎保持恒定（约 1.1 ms/iter）；只有当样本规模达到 $n \ge 10^5 \sim 10^6$ 时，线性增长的残差规约工作量才足以让 GPU 算力完全释放，达到最大 200 倍加速。
显存是硬约束：bartz 的内存设计达到了极高极限（每个特征 1 字节、每个索引 1 字节），整体占用较 dbarts 减少约 24 倍；然而 GPU 无虚拟内存机制，在超高特征维度或超多树配置下，显存溢出会截断向更大样本规模的扩展。
大森林（Forest Size）对高维数据至关重要：论文在附录中首次观察到，当特征维度随样本增长（$p \propto n$）时，固定 200 棵树会导致 RMSE 随 $n$ 上升；而允许树数量随 $n$ 线性增长时 RMSE 保持平稳，表明大森林能有效缓解低噪声高维设置下 MCMC 容易陷入局部极小（mixing 困难）的统计瓶颈。

亮点与洞察¶

无分支（Branchless）重构是复杂概率图模型迁移 GPU 的关键切入点：BART 原本被普遍认为无法高效上 GPU，本文通过“截断最大深度 + 固定二叉堆拓扑”将离散的树遍历和增删变成了无分支定步长张量运算，为其他非参数贝叶斯模型的硬件加速提供了经典范式。
单字节网格量化兼顾精度与极致带宽压缩：树模型的切分决策本质上是分位数比较，通过在 CPU 预处理阶段将浮点特征映射为不超过 255 个网格箱的 uint8 索引，不仅使显存占用暴降至 $np$ 字节，更显著降低了 GPU 访存带宽压力。
消除了 BART 面对 XGBoost 时的速度劣势：以往在数百万样本级任务上使用 BART 是不可想象的，而 2.15 ms/iter 的速度使得 1000 轮 MCMC 仅需约 2 秒，使全贝叶斯不确定性量化首次具备进入工业级大数据管道的实用性。

局限与展望¶

二元分类与分类响应变量暂未支持：当前 bartz 仅实现了连续高斯因变量回归，尚未引入数据增强（Data Augmentation，如 Albert-Chib 潜变量采样）用于二元分类与生存分析，功能完备度落后于 R 生态。
高特征维度下的单卡显存瓶颈：由于全量特征矩阵需驻留于单卡显存，在基因组学等超高维小样本场景下容易因 OOM 退出；多卡扩展目前只能沿样本维度切分，无法有效缓解特征维度带来的单卡显存压力。
有望与 XBART 融合达成数千倍加速：若将本文的无分支 GPU 并行思想进一步迁移至无需 MCMC 循环采样的 XBART 算法上，理论上有望在 GPU 上实现相对传统 CPU 实现 4000 倍以上的究极加速。

评分¶

新颖性: ⭐⭐⭐⭐☆ [首次提出完全无分支且支持 JAX 编译的 GPU 版 BART MCMC 算法]
实验充分度: ⭐⭐⭐⭐☆ [详尽测试了计算延迟、显存占用边界及后验 RMSE 一致性]
写作质量: ⭐⭐⭐⭐⭐ [思路极其清晰，系统工程细节与理论统计分析结合得非常出色]
价值: ⭐⭐⭐⭐⭐ [攻克了贝叶斯非参数领域数十年的核心计算瓶颈，具有重大实用与开源价值]

样本规模 \(n\)	dbarts (CPU, ms)	bartz (CPU, ms)	bartz (A100, ms)	XGBoost (L4, ms/fit)	bartz-A100 加速比 (vs bartz-CPU)
\(10^3\)	0.45	1.80	1.10	2.50	1.6×
\(10^4\)	4.20	5.10	1.12	6.80	4.6×
\(10^5\)	41.0	43.0	1.25	38.0	34.4×
\(10^6\)	420.0	435.0	2.15	280.0	202.3×
\(3 \times 10^6\)	1280.0	1310.0	6.40	850.0	204.7×

配置维度模式	参数设置	峰值显存占用模式	显存溢出临界点 (\(n\))	最大有效加速比 (vs CPU)	性能特征与瓶颈说明
低维基准 (Low-Low)	\(p=100, n_{\text{tree}}=200\)	\(\approx n(p+n_{\text{tree}})\) 字节	\(> 5 \times 10^6\)	200×	显存极为充裕，在大 \(n\) 时残差规约打满 GPU 并发
高维特征 (High-p, Low-\(n_{\text{tree}}\))	\(p=n/10, n_{\text{tree}}=200\)	受限特征索引矩阵	\(\approx 3 \times 10^4\)	~30×	特征矩阵显存迅速耗尽，GPU 加速尚未完全释放
大森林 (Low-p, High-\(n_{\text{tree}}\))	\(p=100, n_{\text{tree}}=n/8\)	受限叶索引缓存	\(\approx 2.5 \times 10^4\)	~25×	节点索引缓存 \(n \times n_{\text{tree}}\) 超出单卡容量上限
极端高维 (High-High)	\(p=n/10, n_{\text{tree}}=n/8\)	双重高阶缓存	\(\approx 8 \times 10^3\)	~15×	极小 \(n\) 下即 OOM，但验证了后验采样一致性