CAFE-GB: Scalable and Stable Feature Selection for Malware Detection via Chunk-wise Aggregated Gradient Boosting

TL;DR

CAFE-GB tackles stability and scalability challenges in high-dimensional malware feature spaces by aggregating local gradient-boosting feature importances computed on overlapping data chunks to yield a stable global ranking. It identifies a principled fixed feature budget (notably $k=100$) and demonstrates that using the top-$k$ features achieves parity with full-feature baselines on BODMAS and CIC-AndMal2020 while reducing dimensionality by over 95%. The approach yields low inter-feature redundancy, interpretable SHAP explanations, and favorable runtime/memory profiles that reduce downstream classification overhead. These results indicate CAFE-GB as a practical, scalable preprocessing step for large-scale malware detection pipelines across heterogeneous datasets.

Abstract

High-dimensional malware datasets often exhibit feature redundancy, instability, and scalability limitations, which hinder the effectiveness and interpretability of machine learning-based malware detection systems. Although feature selection is commonly employed to mitigate these issues, many existing approaches lack robustness when applied to large-scale and heterogeneous malware data. To address this gap, this paper proposes CAFE-GB (Chunk-wise Aggregated Feature Estimation using Gradient Boosting), a scalable feature selection framework designed to produce stable and globally consistent feature rankings for high-dimensional malware detection. CAFE-GB partitions training data into overlapping chunks, estimates local feature importance using gradient boosting models, and aggregates these estimates to derive a robust global ranking. Feature budget selection is performed separately through a systematic k-selection and stability analysis to balance detection performance and robustness. The proposed framework is evaluated on two large-scale malware datasets: BODMAS and CIC-AndMal2020, representing large and diverse malware feature spaces. Experimental results show that classifiers trained on CAFE-GB -selected features achieve performance parity with full-feature baselines across multiple metrics, including Accuracy, F1-score, MCC, ROC-AUC, and PR-AUC, while reducing feature dimensionality by more than 95\%. Paired Wilcoxon signed-rank tests confirm that this reduction does not introduce statistically significant performance degradation. Additional analyses demonstrate low inter-feature redundancy and improved interpretability through SHAP-based explanations. Runtime and memory profiling further indicate reduced downstream classification overhead. Overall, CAFE-GB provides a stable, interpretable, and scalable feature selection strategy for large-scale malware detection.

Figures (5)

• Figure 1: proposed CAFE-GB featture selection framework.
• Figure 2: Impact of feature budget on BODMAS
• Figure 3: Impact of feature budget on ANDMAL2020
• Figure 4: Correlation heatmap for BODMAS dataset with k=100.
• Figure 5: SHAP summary plots for the LightGBM model on AndMal2020 and BODMAS datasets. The plot shows the impact of each feature on the model's output, with features ordered by their importance. Positive SHAP values indicate a higher likelihood of the positive class (malware), while negative values indicate a lower likelihood.