arXiv:2505.12460cs.LG2025-05被引 1

用k-means替代传统分箱法,提升梯度提升树的预测精度。

A Case for Library-Level k-Means Binning in Histogram Gradient-Boosted Trees

  • 以分位数分箱初始化k-means,动态聚类优化分箱边界。
  • 在18个回归数据集上,3个显著提升,最高降55%均方误差。
  • 适合资源受限场景,尤其32-64分箱的GPU部署,可作默认策略。

现代梯度提升决策树(GBDTs)通过直方图分箱加速分裂点查找,将复杂度从O(N log N)降至O(N),但主流分位数分箱策略可能忽略关键边界值。本文提出用以分位数分箱初始化的k-means离散化方法替代分位数分箱,并证明:对任意L-Lipschitz函数,k-means能最大化单个分箱内所有值视为等价时的最坏情况解释方差。在33个OpenML数据集及控制模态、偏度和分箱预算的合成数据上测试,18个回归任务中,k-means在5%显著性水平下无统计损失,且在三个案例中表现更优——尤其在极端偏斜数据集上实现55%的均方误差下降;尽管其平均倒数排名(MRR)略低(0.65 vs 0.72)。15个分类任务中两者统计持平(MRR 0.70 vs 0.68,差距≤0.2个百分点)。合成实验显示,误差降低显著——通常超20%,在异常值增大或分箱预算下降时可达90%。k-means在额外切分价值低时与全枚举分割误差相当,同时恢复了分位数分箱遗漏的关键分裂点。因此,我们建议在库级别内置bin_method=k-means选项,特别适用于回归任务及32-64分箱的GPU场景——因其为“安全默认”,潜力巨大,仅增加一次可缓存的开销(约3.5秒/特征,处理1000万行数据,单个Apple M1线程)。

原文摘要 · Abstract (English)

Modern Gradient Boosted Decision Trees (GBDTs) accelerate split finding with histogram-based binning, which reduces complexity from $O(N\log N)$ to $O(N)$ by aggregating gradients into fixed-size bins. However, the predominant quantile binning strategy - designed to distribute data points evenly among bins -- may overlook critical boundary values that could enhance predictive performance. In this work, we consider a novel approach that replaces quantile binning with a $k$-means discretizer initialized with quantile bins, and justify the swap with a proof showing how, for any $L$-Lipschitz function, k-means maximizes the worst-case explained variance of Y obtained when treating all values in a given bin as equivalent. We test this swap against quantile and uniform binning on 33 OpenML datasets plus synthetics that control for modality, skew, and bin budget. Across 18 regression datasets, k-means shows no statistically significant losses at the 5% level and wins in three cases-most strikingly a 55% MSE drop on one particularly skewed dataset-even though k-means' mean reciprocal rank (MRR) is slightly lower (0.65 vs 0.72). On the 15 classification datasets the two methods are statistically tied (MRR 0.70 vs 0.68) with gaps $\leq$0.2 pp. Synthetic experiments confirm consistently large MSE gains - typically >20% and rising to 90% as outlier magnitude increases or bin budget drops. We find that k-means keeps error on par with exhaustive (no-binning) splitting when extra cuts add little value, yet still recovers key split points that quantile overlooks. As such, we advocate for a built-in bin_method=k-means flag, especially in regression tasks and in tight-budget settings such as the 32-64-bin GPU regime - because it is a "safe default" with large upside, yet adds only a one-off, cacheable overhead ($\approx$ 3.5s per feature to bin 10M rows on one Apple M1 thread).

GBDT分箱优化机器学习高效计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。