A SHARPNESS-AWARE FEDERATED LEARNING FRAMEWORK WITH MC-DROPOUT FOR CNN-BASED MALWARE CLASSIFICATION TO MITIGATE OVERFITTING AND MODEL UNCERTAINTY IN HETEROGENEOUS ENVIRONMENTS
Abstract
Machine learning techniques are not yet widely applicable in the field of malware detection, especially because of data heterogeneity, privacy issues, and the overfitting problem of the model. The presented paper introduces an innovative federated learning framework that combines CNN-based models with Sharpness Aware Minimization (SAM) and Monte Carlo Dropout (MC-Dropout) that overcome these issues. Particularly, the framework is effective operating in decentralized and non-IID settings since it encourages flat minima in the loss landscape, thus boosting model generalization and robustness. Also, predictive uncertainty can be measured due to the utilization of MC-Dropout, a crucial feature when it comes to enhancing the reliability of a decision in the field of cybersecurity. Experimental findings confirm that the suggested method remains considerably superior to the traditional techniques, having high levels of accuracy (up to 99.86%) and precision (99.87%) along with the meaningful estimation of uncertainties. The prospect of unifying the concepts of sharpness-aware optimization and uncertainty quantification under a federated architecture presents a potential solution towards having reliable and privacy-preserved malware detection in heterogenous operational environments.
