EMA-FS：通过增益信息特征筛选加速GBDT训练

# 通过增益感知的特征筛选加速GBDT训练
来源: https://arxiv.org/html/2606.26337
###### 摘要

梯度提升决策树（GBDT），以LightGBM为例，其训练时间的主要部分花在构建每个特征的直方图上——通常占墙钟时间的65–70%。现有的降低该成本的方法，例如随机特征子采样（feature_fraction），在丢弃特征时不考虑其预测效用。我们提出基于EMA的特征筛选（EMA-FS），这是一种算法层面的优化，它维护每个特征分裂增益的指数移动平均（EMA），并在短暂的预热期后，仅对按历史增益排名前KK的特征进行直方图构建。与随机子采样不同，EMA-FS是*有信息依据的*：它优先保留高增益特征，同时筛选出持续低增益的特征。通过在每棵树级别上操作，EMA-FS完全兼容LightGBM的直方图减法技巧，无需更改核心直方图或分裂查找例程。

我们在五个数据集上评估了EMA-FS，涵盖金融欺诈检测（IEEE-CIS Fraud, Credit Card Fraud）、广告点击率预测（Criteo）、工业质量控制（Bosch）以及合成基准，特征维度从29到968。在具有密集特征且中等至高维度的数据集上，EMA-FS实现了显著加速：在500特征合成基准上达到2.61×\times，在432特征的IEEE-CIS欺诈检测数据集上（30%特征保留率）达到1.45×\times。在合成基准上以70%保留率运行时，EMA-FS在实现1.34×\times加速的同时，AUC提升了0.11个点，这证明了通过噪声特征去除带来的隐式正则化。我们还识别并描述了EMA-FS的局限性：在极度稀疏的数据集（Bosch，968特征，>90%缺失值）上，EMA-FS未提供可衡量的加速，因为LightGBM的稀疏分箱优化已绕过空特征值。

我们进一步引入随机EMA-FS（S-EMA-FS），这是一种泛化方法，用增益加权随机采样替代确定性top-KK选择，并由浓度参数β\\beta 控制。S-EMA-FS将确定性EMA-FS（β→∞\\beta\\to\infty）和随机特征子采样（β=0\\beta=0）统一为单一框架的两个特例，结合了增益信息引导的特征选择与驱动随机森林准确性的集成多样性。

EMA-FS和S-EMA-FS在所有六种LightGBM树学习器类型（串行、GPU、特征并行、数据并行、投票并行和线性树）中以约120行C++实现，完全向后兼容，并通过六个直观参数控制。

关键词：梯度提升，特征选择，直方图构建，LightGBM，训练加速，重要性采样，欺诈检测

## 1 引言

梯度提升决策树（GBDT）仍然是表格数据监督学习的主导方法，在分类和回归基准上持续取得最先进的结果（Grinsztajn et al., 2022 (https://arxiv.org/html/2606.26337#bib.bib4)）。现代GBDT框架——LightGBM（Ke et al., 2017 (https://arxiv.org/html/2606.26337#bib.bib6)）、XGBoost（Chen & Guestrin, 2016 (https://arxiv.org/html/2606.26337#bib.bib2)）和CatBoost（Prokhorenkova et al., 2018 (https://arxiv.org/html/2606.26337#bib.bib8)）——已通过基于直方图的分裂查找、数据并行训练和梯度量化进行了广泛的计算效率优化。

尽管有这些优化，对于涉及数百个特征和数百万样本的大规模应用，训练时间仍然是一个关键瓶颈。对LightGBM训练循环的分析揭示了一个显著的不平衡：直方图构建约占墙钟时间的69%，而分裂查找约占26%，其余5%分布在数据划分和簿记上（表1 (https://arxiv.org/html/2606.26337#S1.T1)）。这种分布在不同数据集大小和硬件配置上保持一致，反映了直方图累积的内存带宽受限性质。

表1：LightGBM在代表性工作负载（80K×500特征，100棵树，4线程）上的训练时间分解。先前尝试加速直方图构建的内循环——包括SIMD向量化、缓存行优化和内存布局变换——仅带来边际改善，因为工作负载已经达到内存带宽饱和。我们初步工作中的八次不同的微优化尝试均未能产生可衡量的加速，这证实了加速的途径不在于使每个直方图更快，而在于*构建更少的直方图*。

LightGBM现有的feature\_fraction参数通过每棵树随机选择一组特征来解决这个问题。然而，随机选择是无信息的：它可能同样丢弃高增益特征和低增益特征，导致在激进的子采样率下不必要的准确性损失。

在本文中，我们提出基于EMA的特征筛选（EMA-FS），一种简单而有效的算法，它使用指数移动平均追踪每个特征的分裂增益，并在短暂的预热后，将直方图构建限制在按历史增益排名的前KK个特征。EMA-FS基于一个关键观察：特征重要性在提升迭代中是*持久的*：在早期树中产生高增益分裂的特征在后期树中往往仍然重要，尤其是在信号结构稳定的数据集中。

我们的贡献如下：

1. 1. **算法**。我们引入EMA-FS，一种增益驱动的特征筛选方法，其直方图构建成本的降低与筛选比例成正比，且无需更改LightGBM的核心直方图或减法例程。
2. 2. **随机泛化**。我们提出S-EMA-FS，用增益加权随机采样替代确定性top-KK选择。S-EMA-FS将EMA-FS和随机特征子采样（feature\_fraction）统一为单一参数化框架的两个端点，结合了增益信息引导的选择与集成多样性。
3. 3. **实现**。我们提供了一个完整的实现，集成到所有六种LightGBM树学习器类型（串行、GPU、特征并行、数据并行、投票并行和线性树）中，约120行C++，通过六个直观参数控制。
4. 4. **综合评价**。我们在五个数据集上进行了实验，涵盖金融欺诈检测、广告和工业质量控制。我们描述了EMA-FS的优势（密集、中等至高维数据）和局限性（稀疏数据、低维数据），并与随机特征子采样进行了直接比较。

## 2 相关工作

##### 基于直方图的GBDT。

LightGBM（Ke et al., 2017 (https://arxiv.org/html/2606.26337#bib.bib6)）为GBDT引入了基于直方图的分裂查找，将每个节点的分裂搜索成本从O(n·F)O(n·F)降低到O(B·F)O(B·F)，其中BB表示分箱数（通常B≪nB≪n）。XGBoost（Chen & Guestrin, 2016 (https://arxiv.org/html/2606.26337#bib.bib2)）采用使用分位数草图近似分裂查找。CatBoost（Prokhorenkova et al., 2018 (https://arxiv.org/html/2606.26337#bib.bib8)）使用有序提升和 oblivious 树。这三个框架在每个树节点上为所有选中的特征构建直方图，使得直方图构建成为宽数据集的主要成本。

##### 特征子采样。

随机特征子采样（XGBoost中的colsample\_bytree，LightGBM中的feature\_fraction）是一种广泛使用的正则化技术，受随机森林（Breiman, 2001 (https://arxiv.org/html/2606.26337#bib.bib1)）启发。通过每棵树随机选择一部分特征，它同时减少了过拟合和训练时间。然而，随机选择是无信息的，没有利用特征重要性的异质性。

##### 特征重要性和选择。

从分裂增益或分裂计数导出的特征重要性分数是GBDT模型的标准输出（Friedman, 2001 (https://arxiv.org/html/2606.26337#bib.bib3)）。这些分数通常*事后*用于模型解释或特征工程。一些工作探索了在集成方法中使用特征重要性进行迭代特征选择（Guyon & Elisseeff, 2003 (https://arxiv.org/html/2606.26337#bib.bib5)），但这些方法在模型层面操作（使用选中的特征重新训练），而不是在提升循环内部。Xu等人（2014 (https://arxiv.org/html/2606.26337#bib.bib9)）提出用于随机森林的基于梯度的特征选择，但该方法需要单独的特征排名步骤。

##### 基于梯度的一侧采样（GOSS）。

LightGBM的GOSS（Ke et al., 2017 (https://arxiv.org/html/2606.26337#bib.bib6)）通过保留大梯度实例并采样小梯度实例来减少用于直方图构建的数据点数量。EMA-FS是互补的：GOSS减少数据维度，而EMA-FS减少特征维度。这两种技术可以结合。

##### 成本效益梯度提升（CEGB）。

Nan等人（2022 (https://arxiv.org/html/2606.26337#bib.bib7)）提出基于获取成本惩罚特征使用。EMA-FS的不同之处在于它使用*观察到的*分裂增益而不是外部指定的成本，并且其目标是训练加速而不是部署时的成本降低。

##### 具有二元掩码的自适应特征选择（AFS-BM）。

Akman等人（2024 (https://arxiv.org/html/2606.26337#bib.bib10)）在GBDT训练期间联合优化二元特征掩码和模型参数，目标是在高维数据上提高预测准确性。EMA-FS在三个方面不同：(i) 目标——我们针对直方图构建成本降低，而非准确性提高；(ii) 机制——我们使用观察到的分裂增益的轻量级EMA，而非联合掩码-参数优化；(iii) 集成点——EMA-FS在LightGBM每棵树特征选择阶段内部操作，无需辅助优化循环，保留了直方图减法技巧，仅增加了约120行C++。这两种方法是互补的：AFS-BM寻求全局信息性特征掩码，而EMA-FS利用每轮增益信号来减少主要计算成本。

##### 用于GBDT的重要性加权特征采样。

与我们的工作更接近的是，Zhou等人（2022 (https://arxiv.org/html/2606.26337#bib.bib11)）提出LGBM-CBFS，这是一种启发式方法，根据累积分裂增益导出的重要性分数在LightGBM中对特征进行采样。LGBM-CBFS针对稀疏高维数据上的准确性和收敛性，而在低速率下统一的feature\_fraction通常无法采样到信息性特征。EMA-FS在三个方面不同：(i) 我们使用增益的*指数移动平均*而非累积和，使得权重能够适应提升后期特征重要性的变化；(ii) 我们的目标是直方图构建成本降低，经验证据表明轻量级掩码过滤避免了feature\_fraction在高维数据上引入的col\_sampler开销（第6.2节 (https://arxiv.org/html/2606.26337#S6.SS2)）；(iii) 我们通过β\\beta参数引入了确定性与随机性的统一，其中LGBM-CBFS的增益加权采样成为其特例之一（β=1\\beta=1且stochastic=true）。

## 3 基于EMA的特征筛选

### 3.1 动机

考虑一个具有FF个特征的数据集。在提升树的每个内部节点处，LightGBM为每个特征构建一个梯度统计直方图，每个叶子需要O(nleaf·F)O(n_{\text{leaf}}·F)次内存访问，其中nleafn_{\text{leaf}}是该叶子中的数据点数。对于具有LL个叶子的树中的所有节点，总成本为O(N·F)O(N·F)，其中NN是训练集大小（每个数据点在每个级别恰好贡献给一个叶子的直方图）。

驱动EMA-FS的关键观察是，在许多实际数据集中，一小部分特征占了分裂增益的绝大部分。如果我们能在训练早期识别出这些高增益特征，并将直方图构建限制在仅这些特征上，我们就可以按比例降低主要训练成本。

### 3.2 算法

EMA-FS维护每个特征的增益估计g=(g1,...,gF)\mathbf{g}=(g_1,\ldots,g_F)，使用每棵树后更新的指数移动平均（EMA）。该算法有两个阶段：

##### 预热阶段（树1,...,W）。

所有FF个特征在每个节点上都进行评估。每棵树tt之后，记录在这棵树所有节点上观察到的每个特征的最大增益g^f(t)\hat{g}_f^{(t)}。EMA更新如下：

gf(t)=α·g^f(t)+(1−α)·gf(t−1),g_f^{(t)}=\alpha·\hat{g}_f^{(t)}+(1-\alpha)·g_f^{(t-1)}, (1)
其中α∈(0,1)\alpha\in(0,1)是平滑参数（本文中我们统一使用α=0.3\alpha=0.3）。

##### 筛选阶段（树W+1,...,T）。

在构建树tt之前，根据其当前EMA值gf(t−1)g_f^{(t-1)}按降序对特征进行排序。选择前K=⌈r·F⌉K=\lceil r·F\rceil个特征，其中r∈(0,1]r\in(0,1]是ema\_fs\_feature\_rate参数。只有选中的特征参与树上所有节点的直方图构建和分裂查找。树构建完成后，使用选中特征的增益更新EMA（未选中特征保留其先前的EMA值，按因子(1−α)(1-\alpha)衰减）。

完整算法见算法1 (https://arxiv.org/html/2606.26337#alg1)。

算法1 基于EMA的特征筛选（EMA-FS）

0: 数据集D\mathcal{D}，具有FF个特征，树数量TT，预热WW，比率rr，EMA参数α\alpha
1: 初始化gf←0g_f\leftarrow 0对所有f∈{1,...,F}f\in\{1,\ldots,F\}
2: 初始化S←{1,...,F}S\leftarrow\{1,\ldots,F\} ⊳\triangleright 选中的特征
3: for t=1t=1 to TT do
4: 重置g^f(t)←0\hat{g}_f^{(t)}\leftarrow 0对所有ff
5: if t>Wt>W then
6: K←max(1,⌈r·F⌉)K\leftarrow\max(1,\lceil r·F\rceil)
7: S←top-K features bygfS\leftarrow\text{top-}K\text{ features by }g_f (递减)
8: endif
9: 只使用SS中的特征构建树tt:
10: for tree tt中的每个节点 do
11: 为SS中的特征构建直方图
12: 在SS中找到最佳分裂
13: 记录: g^f(t)←max(g^f(t),gainf)\hat{g}_f^{(t)}\leftarrow\max(\hat{g}_f^{(t)},\;\text{gain}_f) 对于分裂特征ff
14: endfor
15: 更新EMA: gf←α·g^f(t)+(1−α)·gfg_f\leftarrow\alpha·\hat{g}_f^{(t)}+(1-\alpha)·g_f 对所有ff
16: endfor

### 3.3 设计决策

##### 每棵树特征选择。

EMA-FS在树级别而非节点级别选择特征。这对于与LightGBM的直方图减法技巧的兼容性至关重要，该技巧通过从父直方图中减去较小兄弟节点的直方图来计算兄弟节点的直方图。直方图减法要求两个兄弟节点使用相同的特征集，当选择是每棵树时，这一点得到保证。

##### 用于增益追踪的EMA。

我们使用指数移动平均而不是累积平均值或滑动窗口来追踪特征增益。EMA自然适应训练过程中特征重要性的变化：早期树中的增益对当前估计有很大影响，而在许多轮后其影响衰减。这种衰减率由α\alpha控制：α\alpha越大，EMA对新信息反应越快，但可能导致估计方差更大；α\alpha越小，估计越稳定，但适应变化较慢。我们凭经验选择α=0.3\alpha=0.3，作为响应性和稳定性之间的合理折中。