Enhancing Decision-Making in Windows PE Malware Classification During Dataset Shifts with Uncertainty Estimation

TL;DR

This paper improves Windows PE malware classification robustness during dataset shifts by integrating ensemble-based uncertainty estimates with conformal prediction, significantly reducing misclassification acceptance rates in challenging scenarios.

Contribution

It introduces a novel ensemble-based uncertainty estimation method within conformal prediction to enhance malware classifier reliability under dataset shifts.

Findings

Ensemble-based uncertainty estimates reduce incorrect acceptance rates by ~30%.

The method maintains competitive correct acceptance rates under dataset shifts.

Effective in detecting packed malware with distributional differences.

Abstract

Artificial intelligence techniques have achieved strong performance in classifying Windows Portable Executable (PE) malware, but their reliability often degrades under dataset shifts, leading to misclassifications with severe security consequences. To address this, we enhance an existing LightGBM (LGBM) malware detector by integrating Neural Networks (NN), PriorNet, and Neural Network Ensembles, evaluated across three benchmark datasets: EMBER, BODMAS, and UCSB. The UCSB dataset, composed mainly of packed malware, introduces a substantial distributional shift relative to EMBER and BODMAS, making it a challenging testbed for robustness. We study uncertainty-aware decision strategies, including probability thresholding, PriorNet, ensemble-derived estimates, and Inductive Conformal Evaluation (ICE). Our main contribution is the use of ensemble-based uncertainty estimates as Non-Conformity Measures within ICE, combined with a novel threshold optimisation method. On the UCSB dataset, where the shift is most severe, the state-of-the-art probability-based ICE (SOTA) yields an incorrect acceptance rate (IA%) of 22.8%. In contrast, our method reduces this to 16% a relative reduction of about 30% while maintaining competitive correct acceptance rates (CA%). These results demonstrate that integrating ensemble-based uncertainty with conformal prediction provides a more reliable safeguard against misclassifications under extreme dataset shifts, particularly in the presence of packed malware, thereby offering practical benefits for real-world security operations.