IoT - FedMalDetect: Federated Learning Based Malware Detection for IoT Edge Devices

Abstract

Rapid adoption of the Internet of Things (IoT) devices has transformed industries by enabling automation and seamless connectivity, but it also introduces significant security challenges, particularly for malware threats. To address these challenges, numerous AI-based solutions have been proposed to enhance IoT device security. Additionally, federated learning (FL) based solutions have been proposed to enhance security and data privacy guarantees. However, these methods require access to labeled data, but data annotation is an arduous task in security. Further, storing such a large pool of labeled data in memory-constrained IoT devices is difficult. We propose 10T-FedMaIDetect, a novel semi-supervised FL framework for dynamic malware detection in IoT edge devices based on network traces to address these issues. Our approach is a two-stage training framework that combines unsupervised federated learning with supervised fine-tuning. In the first stage, clients train feature learning models independently, and the server aggregates them to build a global model. In the second stage, this global model is paired with a classifier and fine-tuned on a publicly available labeled dataset of mal ware. The resulting model is then used as the deployable detection system at the IoT devices. Our solution achieves comparable detection performance with a ROC-AUC of 0.96 and a PR-AUC of 0.943, with the existing semi-supervised FL method, FedMSE, on the N-BaIoT dataset without requiring any client-side labels. Additionally, we deployed our framework on a real-world IoT testbed to evaluate its deployment feasibility and observed comparable detection performance with reduced memory and CPU resources.