Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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

About

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.

Yan Song• 2026

Related benchmarks

TaskDatasetResultRank
CTR PredictionCriteo
AUC0.7545
346
Industrial Defect DetectionBosch
AUC70.24
8
ClassificationSynthetic
AUC96.14
5
ClassificationFraudDense
AUC0.588
5
Binary ClassificationSynthetic 500 features (test)
Inference Time (s)1.5
5
Fraud DetectionIEEE-CIS Fraud Detection
AUC0.9095
5
ClassificationIEEE-CIS
AUC0.9252
4
ClassificationFraudMixed
AUC0.5868
4
ClassificationCriteo
AUC0.7615
4
Financial Fraud DetectionCredit Card Fraud (test)
AUC76.22
4
Showing 10 of 14 rows

Other info

Follow for update