EMA-FS: Accelerating GBDT Training via Gain-Informed Feature Screening

Abstract: Gradient Boosted Decision Trees (GBDT), exemplified by LightGBM, spend a dominant fraction of training time -- typically 65-70% -- constructing per-feature histograms. Existing approaches such as random feature subsampling (feature_fraction) discard features without regard for their predictive utility. We propose EMA-based Feature Screening (EMA-FS), an algorithm-level optimization that maintains an exponential moving average (EMA) of per-feature split gains across boosting iterations and, after a short warmup, restricts histogram construction to the top-K features ranked by historical gain. Unlike random subsampling, EMA-FS is informed: it retains high-gain features while screening out low-gain ones. Operating at the per-tree level, it preserves full compatibility with LightGBM's histogram subtraction trick, requiring no changes to core routines. We evaluate EMA-FS on datasets spanning financial fraud detection, advertising click-through prediction, industrial quality control, and synthetic benchmarks, with feature dimensionalities from 29 to 968. On dense, moderate-to-high-dimensional data it achieves significant speedups: 2.61x on a 500-feature synthetic benchmark and 1.45x on the 432-feature IEEE-CIS Fraud dataset at 30% retention. At 70% retention it improves AUC by 0.11 points while delivering a 1.34x speedup. On extremely sparse data (Bosch, >90% missing) it yields no speedup, as LightGBM's sparse bin optimization already bypasses empty values. We further introduce Stochastic EMA-FS (S-EMA-FS), which replaces deterministic top-K selection with gain-weighted random sampling controlled by a concentration parameter beta, unifying deterministic EMA-FS (beta -> infinity) and random subsampling (beta = 0) in one framework. Both are implemented in ~120 lines of C++ across all six LightGBM tree learners and are fully backward-compatible.

• The paper introduces EMA-FS, which uses an exponential moving average of split gains to selectively screen features and reduce computational overhead in histogram construction.
• It demonstrates significant speedups—up to 2.61×—on dense, high-dimensional datasets with minimal or even beneficial impacts on AUC, highlighting effective regularization.
• The method is lightweight, backward-compatible with LightGBM, and extends to a stochastic variant (S-EMA-FS) that balances deterministic screening with random feature sampling.

EMA-FS: Gain-Informed Feature Screening for Efficient GBDT Training

Background and Motivation

Gradient Boosted Decision Trees (GBDT) are the prevailing architecture for tabular data learning, with frameworks such as LightGBM, XGBoost, and CatBoost converging on optimizations like histogram-based split finding for computational efficiency. Despite their optimization, profiling of LightGBM’s training reveals that histogram construction dominates the wall-clock time—typically accounting for 65-70% of training time—due to its memory-bandwidth-bound nature. Prior attempts at micro-optimization have delivered negligible returns because the core bottleneck is the number, not the speed, of histogram constructions.

The common remedy, random feature subsampling (e.g., LightGBM’s feature_fraction), indiscriminately discards features without regard to their utility, often leading to unnecessary accuracy degradation when aggressively subsampling. The potential for further acceleration thus requires a mechanism to selectively exclude only features of persistently low predictive value.

EMA-FS: Algorithm and Generalization

The paper introduces EMA-based Feature Screening (EMA-FS), which maintains an exponential moving average (EMA) of the observed split gains per feature across trees. After a warmup phase where all features participate, subsequent trees are constructed considering only the top-$K$ features as ranked by the EMA of their historical gains.

Key aspects of the design include:

• Per-tree feature screening to remain compatible with LightGBM’s histogram subtraction trick.
• EMA update after each tree, with a smoothing parameter $\alpha$ dictating the adaptation to changing signal.
• Maximum gain aggregation across nodes within each tree to ensure that strongly interacting features are not underweighted due to localized effects.
• Fully deterministic execution (absent of additional randomness), aiding reproducibility.

Furthermore, the stochastic generalization S-EMA-FS extends EMA-FS by sampling features per tree in a gain-weighted random fashion, with a concentration parameter $\beta$ interpolating between random feature subsampling ($\beta=0$) and deterministic EMA-FS ($\beta \to \infty$).

Empirical Evaluation

The empirical analysis spans five public datasets and two large-scale synthetic fraud detection benchmarks, offering a comprehensive view over feature counts, sparsity, and domain specificity.

Quantitative Results

• Synthetic benchmark (500 dense features): At $r=0.3$ (30% retention), EMA-FS achieves a 2.61× speedup with a marginal reduction in AUC (0.94 vs 0.96); at $r=0.7$, AUC actually improves (+0.11 points), indicating a regularization benefit from screening out low-utility features.
• IEEE-CIS Fraud Detection (432 features): EMA-FS achieves 1.45× speedup at 30% feature retention with an AUC decrease of 0.046; at higher retention, the AUC decrease is minimized. Notably, random feature fraction is slower than baseline due to overhead incurred by LightGBM’s column sampling logic at high feature counts.
• Bosch (968 features, >90% missing): EMA-FS provides no measurable speedup; LightGBM’s internal sparse bin optimization is already efficient for such data, and most features are computationally negligible.
• Low-dimensional datasets (e.g., CreditCard, Criteo): Very limited benefit due to insufficient screening headroom.

S-EMA-FS bridges deterministic screening and random subsampling. With appropriate $\beta$ and $r$, S-EMA-FS achieves simultaneous speedup and accuracy improvement—for example, on large-scale fraud data, S-EMA-FS at $r=0.3$ and $\alpha$0 delivers a 2.4× speedup while improving AUC slightly over baseline.

Analysis of Speedup and Conditions for Applicability

The primary determinant of EMA-FS effectiveness is the presence of dense, moderately to highly dimensional features (100–500+), where a nontrivial fraction of features are noise or otherwise locally uninformative. When feature sparsity is extreme (as in Bosch), or in datasets with uniformly high feature utility, screening is ineffective or may degrade performance.

EMA-FS yields maximal benefit where:

• The signal is concentrated in a subset of features.
• Feature importance is stable across boosting iterations.
• Training employs data- or feature-parallel learners—in these settings, EMA-FS reduces both compute and communication overheads.

Implementation and Integration

The EMA-FS implementation requires only ~120 lines of C++, is backward-compatible with LightGBM across all learner types, and is governed by six parameters (enabling/disabling, feature ratio, warmup, $\alpha$1, stochastic mode, and $\alpha$2). The screening logic is lightweight compared to the cost of dense histogram construction. EMA-FS composes naturally with other LightGBM optimizations (e.g., GOSS).

Practical and Theoretical Implications

By focusing on gain-informed screening, EMA-FS and S-EMA-FS offer a principled and low-overhead approach to balancing accuracy and training efficiency in GBDT. For practitioners, these methods provide substantial training acceleration without requiring disruptive code modifications or extensive tuning, and sometimes even deliver accuracy gains via implicit regularization. Theoretically, gain-weighted stochastic screening introduces an importance sampling interpretation in the feature domain, unifying the ideas of diversity and exploitation.

The strong separation of CPU and communication cost in distributed settings suggests that the method’s utility will trend upward as distributed GBDT training becomes more ubiquitous in production ML pipelines, especially in domains (e.g., financial risk control) where tabular datasets with hundreds of engineered features and frequent retraining cycles are the norm.

Future Directions

• Adaptive feature ratio: Dynamically adjusting the screening threshold $\alpha$3 based on the observed gain distribution.
• Sparse-aware screening: Deeper integration with sparse bin linear cost estimation.
• Joint data-feature sampling: Composing with instance-level samplers such as GOSS.
• Partial reentry: Smart mechanisms for periodically reevaluating low-gain features for potential resurgence in predictive value.

Conclusion

EMA-FS and its stochastic extension S-EMA-FS address the longstanding inefficiency in GBDT histogram construction by leveraging historical split gain statistics for feature screening. These methods are empirically validated to provide substantial, parameterizable speedups (up to 2.6× on 500-dense-feature benchmarks) with negligible or even negative (regularizing) impact on predictive performance for dense, mid- to high-dimensional tabular data. Their practical design, low implementation burden, and backward compatibility position them as immediately usable optimizations for GBDT pipelines, particularly in high-value applications such as fraud detection and risk modeling.