DRMD: Deep Reinforcement Learning for Malware Detection under Concept Drift

Malware poses a significant and ever-evolving threat to mobile devices, personal computers, and enterprise systems. Every day, tens of thousands of new Android and Windows applications appear, far outpacing the capacity for manual analysis. Automated malware detection using machine learning approaches is therefore crucial to ensure user security.

A core challenge of malware detection arises from the scale of new samples that emerge daily and the limited capacity for manual review. Onwuzurike et al. estimated that an upper bound of 10,000 applications are submitted to the Google Play Store each day, while Miller et al. assumed the manual review capacity of an average company at only 80 samples per day. This significant gap necessitates the use of machine learning (ML) to accurately and consistently classify new applications based on previously labeled data.

A key limitation of traditional Supervised Learning (SL) classifiers is their reliance on the assumption that data follows an Independent and Identically Distributed (IID) structure. However, this assumption does not hold in the domain of Android malware, due to the continuous evolution of both benign applications (goodware) and malware over time. Goodware evolves naturally as software developers introduce new features, while malware adapts to better mimic goodware and evade detection. This continual evolution leads to a divergence between training and testing data that expands over time, a phenomenon known as concept drift.

Approaches such as active learning (AL) to periodically retrain on a subset of new samples and classification with rejection to selectively abstain from classification can mitigate performance degradation caused by evolving threats. However, manual labeling in the malware domain is resource-intensive. Therefore, AL techniques must carefully balance selecting the most informative samples while minimizing labeling costs.

We present the first approach in this domain that reformulates malware detection as a one-step MDP, or contextual bandit, where each episode corresponds to a single sample, thus avoiding spurious dependencies between states.

DRMD is a DRL-based approach that, in its most complete form, unifies classification, AL, and rejection in a single malware detection pipeline. DRMD first formulates the Malware Detection MDP (MD-MDP), then uses a deep neural network to learn an effective detection policy. Concretely, MD-MDP extends ICMDP with corrected episode definitions and integrates the mechanisms necessary for effective malware detection in practice.

The application of DRL to detection tasks has been shown to improve performance in more complex domains with class imbalance. However, the use of DRL in malware detection has seen little work likely due to the drift and class imbalance of the domain. A policy for malware detection should consider the larger decision-making process of classification, rejection, and AL; thus allowing for the optimization of performance to be rejection-aware. Finally, the agent must adapt to concept drift within the domain constraints. As a result, DRL policies must move beyond the stationary perspectives of prior problem domains. Our MD-MDP and agent design achieve these considerations through a one-step episodic formulation that leverages an expanded reward function for temporal robustness and the integration of rejection.

• Episodes. In our formulation, each episode consists of a single step: classifying or rejecting one sample.
• States. For episode e, the sample x_e consists of d features, where F ⊆ ℝ^d represents the feature space, hence x_e ∈ F. The state consists of the features of the current sample; therefore, the state space and state are defined as S = F and s = x_e, respectively.
• Actions. A chosen action a ∈ A represents the decision of the agent for the current state s. The action space is A = {a_gw, a_mw, a_rej}, where a_gw classifies the sample as goodware, a_mw as malware and a_rej rejects the sample, abstaining from classification. Active learning is integrated by retraining on samples for which the agent selected a_rej.

The rewards need to encapsulate the dynamics of classifications and rejection. The three underlying components of classification in the domain are the accuracy of the predictions, the natural imbalance in the distribution of the samples, and the evolution of samples over time.

The classification (R_clf) and complete DRMD (R_cr) reward functions are defined as:

Using our formulation, we create two types of policies for evaluation. First, a Classification-Only Policy (π_clf) using the classification actions (a_gw and a_mw), and reward R_clf(s, a). Second, a Classification-Rejection Policy (π_cr) that uses all actions, and reward R_cr(s, a).

We use Proximal Policy Optimization (PPO), modified for one-step episodes, to train the DRMD agent. Additionally, PPO has seen successful application to several other security domains such as automated cyber defense, vulnerability discovery, and adversarial machine learning. The on-policy nature of PPO and its use of rollout updates facilitate learning on the most recent samples compared to experience replay buffers, used in off-policy methods, which would be more susceptible to concept drift.

While we have formulated the MD-MDP as a one-step MDP or contextual bandit problem, existing deep contextual bandits (DCBs) are not well-suited to our setting. First, DCBs assume stationary sub-Gaussian (bounded) reward noise, an assumption violated in malware domains due to concept drift. Furthermore, DRMD significantly outperforms DCBs, NeuralTS and NeuralUCB, over all settings for the classification–only policy with an average ∆Base of 9.77.

The agent employs separate actor and critic neural networks, each consisting of four hidden layers with 512 neurons per layer, using LeakyReLU activations and dropout for regularization. During training, each epoch consists of one episode per training sample using the current policy. The collected states, actions, and rewards are then used to compute the clipped policy and value losses to update the actor and critic networks over mini batches. When new samples become available after the initial training, the agent will be fine-tuned on a sliding window of the newest samples (instead of on all samples) to focus updates on adapting to changes in the domain.

Lin et al. formulated ICMDP, forming states as samples, actions as predicted labels, rewarding +1 for correct and -1 for incorrect classification of minority class samples with majority class sample rewards down-scaled according to the class distribution. SINNER builds on the ICMDP framework and extends its reward function to the multi-class setting, employing a Dueling Double DQN to classify the family of Windows PE malware samples. Although these approaches improve performance by harnessing DRL, they do not consider concept drift, a core problem in malware detection.

Prior work has applied the ICMDP to several domains such as malware family classification. However, the formulation of ICMDP assumes state transitions between classification samples, thus treating independent samples as sequentially dependent. To overcome this, we cast malware detection as a one-step MDP (aka contextual bandit).

Therefore, our MD-MDP formulation ensures that policy updates do not introduce spurious correlations as a result of artificial dependencies between independent samples.

To isolate the difference from MD-MDP compared to ICMDP, we train a version of the DRMD agent using the transitions and reward structure of ICMDP. Concretely, episodes continue until either a) malware is misclassified or b) the agent runs out of training samples. Otherwise, states and actions are the same as in the MD-MDP. As a result, the next states are based on the next sample from the dataset. Rewards are structured as in R_acc, however, rewards from goodware samples are scaled down by σ̂. For fairness, we use the same PPO architecture and hyperparameters to train both our agent and the ICMDP agent.

Overall, the MD-MDP outperforms ICMDP in 97 out of 100 settings, with 45 settings being statistically significant, and achieves an average ∆IC of 1.94.

Since concept drift in malware detection is inherently time-dependent, time-aware evaluations are essential to assess classifier performance in real-world scenarios. A valid time-aware evaluation must adhere to the following constraints. (C1) Training samples must always precede testing samples. (C2) All testing samples must be drawn from the same time window. (C3) The malware distribution in the testing set must reflect a realistic distribution. Violating these constraints can introduce biases that artificially inflate the performance of the classifier, leading to unrealistic expectations in practical deployments.

We use two Android malware datasets in our evaluation: Transcendent and Hypercube. Transcendent provides an established Android malware dataset, while Hypercube is an updated dataset with newer malware samples. Transcendent consists of 259,230 applications collected between 2014 and 2018. Hypercube contains 159,839 applications between 2021 and 2023. Both datasets contain a malware distribution of ≈ 10% according to the spatial constraint of Tesseract, and the samples have been labeled using a standard VirusTotal detection threshold (VTT = 4).

To ensure that the performance of DRMD is independent of a single feature representation, we evaluate it using two different Android malware feature spaces for each dataset. We use feature spaces that enable us to evaluate our approach across varying amounts of information per sample.

In order to evaluate the performance of detectors, we use the Area Under Time (AUT) metric, introduced by Pendlebury et al., and defined as:

where: P(X_k) is the performance on samples X_k using metric P, with P being the F1 score in our paper.

To compare our novel MD-MDP with the ICMDP, we evaluate the two model designs using the same (our) DRL agent architecture, thereby isolating the effect of this factor on the performance. The difference in performance is denoted as ∆IC. Moreover, we compare DRMD with a version of our model architecture that uses Supervised Learning (SL) and standard cross-entropy loss, which we call SL-DRMD.

In addition to SL-DRMD, for baselines, we also evaluate commonly used classifiers in Android malware detection: Drebin, DeepDrebin, and Ramda. The Drebin classifier is a linear support vector machine (SVM) originally used for its corresponding feature space and serves as a foundational benchmark due to its simplicity. DeepDrebin applies a multilayer perceptron (MLP). Finally, the Ramda classifier, developed for its corresponding feature space, uses a variational autoencoder (VAE) and a MLP to classify samples. Together, Drebin, DeepDrebin, Ramda, and SL-DRMD are evaluated across five random seeds on the four dataset-feature-space pairs and the difference between DRMD and the best performing baseline, for that setting, is reported as ∆Base.

This evaluation aims not only to investigate the base performance of DRMD but also to assess the impact that including rejection and AL (both monthly and integrated) has on model performance over time. We evaluate DRMD using five seeds, following the work of Alam et al., for each combination of dataset, feature space, and configuration, totaling 172 settings. Our results demonstrate that DRMD, both classification-only and classification-rejection policies, consistently outperform considered baselines, achieving an average AUT improvement of 8.66 and 10.90, respectively.

First, we consider DRMD trained to learn a Classification–Only policy (π_clf), without the rejection action. Initially, when there is no AL or rejection, DRMD surpasses the best baselines by 6.94 on average across five runs and all dataset-feature-space pairs. These initial results are promising, yet do not consider a complete malware detection pipeline. Thus, we consider malware detection pipeline configurations using monthly budgets for the uncertainty-based AL and rejection (Rej).

Monthly Rejection. Here, we consider how integrating uncertainty-based monthly rejection affects performance when no AL is present. DRMD shows statistically significant improvements in performance over the best baseline in 15 of the 16 settings, with an average ∆Base of 15.01. Increasing the number of uncertain samples rejected leads to increased performance for both AUT and ∆Base. This means that the samples DRMD has high certainty on lead to more correct classifications than the baselines.

Monthly Active Learning. In the AL setting, both the DRMD policy and the baselines use monthly uncertainty-based AL to select samples for retraining at each testing month. DRMD demonstrates statistically significant improvements in 9 of the 16 settings (where AL > 0 and Rej = 0), achieving an average ∆Base of 2.52. As in the rejection-only setting, both DRMD and the baselines use the same AL mechanism and budgets. Note that the diminishing returns of AL as performance improves coupled with the initial performance gains of DRMD contribute to the decreasing ∆Base as AL rates increase.

Monthly Rejection & Active Learning. DRMD has statistically significant improvements in 52 of the 64 settings, achieving up to 29.07 AUT improvement (AL = 400 and Rej = 400), and an average ∆Base of 8.71. Thus, when the rejection of samples at high risk of misclassification and the cost-sensitive drift adaptation of AL are applied together on DRMD, this performs better than when only using either individually, highlighting how our approach can yield improvements in real world settings.

Extending from the Classification-only policy, the DRMD Classification-Rejection policy (π_Σ) integrates the rejection action, a_rej. Importantly, this action enables the agent to abstain from making a decision on a sample if there is insufficient information, high uncertainty, or concept drift. Note how, previously, rejections occurred only at the end of each month; however, the rejection action provides real-time abstention as the samples are seen, thus opening up the potential for real-time detection systems in future work. To maintain the capability for such decision making while facilitating AL, the samples resulting from the rejection action are used for retraining. To ensure consistency in the number of rejections, the seed-wise average rejections (rounded up) of DRMD is used as the budget for baseline rejections.

Integrated Rejection (IR). The IR setting introduces the rejection action without any additional AL on rejected samples. The results show statistically significant improvements over all baselines in 9 of the 12 settings with an average ∆Base of 16.15. This shows that the agent can effectively learn a policy that rejects samples likely to be misclassified. When comparing the two feature spaces (Ramda and Drebin), we see a disparity in the average number of rejections. Specifically, policies trained on Ramda reject more samples than those trained on Drebin. We hypothesize that the reduction of information encapsulated by Ramda results in greater prediction uncertainty compared to the Drebin feature space. Increasing R_cost intuitively reduces the number of rejected samples with only a minimal effect on ∆Base. This implies that higher R_cost values lead DRMD to be more selective in rejection.

IR with Active Learning (IRAL). To go beyond the mitigation of concept drift and allow self-adaptation, we include the samples rejected in each period for retraining DRMD by updating the fine-tuning sample sliding window. Compared with the best baseline for each setting, DRMD shows statistically significant improvements in 7 of the 12 settings with an average ∆Base of 11.60. In Hypercube-Drebin and Transcendent-Drebin, retraining on rejected samples leads to higher confidence in classification actions. As R_cost increases, this leads to a reduction in the number of rejected and, thus, retraining samples in subsequent periods causing the observed drop in performance.

IR with Augmented Active Learning (IRAAL). Augmented AL extends the rejection-based sampling used in IRAL with uncertainty sampling so that, given a predetermined budget for AL, the model samples up to the predetermined labeling budget for retraining if the agent was conservative in rejections, and downsamples to the budget if the agent is rejecting more samples. The results show statistically significant performance gains over the baseline approaches in 33 of 48 settings, with an average ∆Base of 10.75. Highlighting that the augmented AL budget improves stability and performance compared to the IRAL results.

In Table 3 we present a cumulative component breakdown of the AUT performance of DRMD across the four dataset–feature-space pairs.

Starting from the basic DRMD (single 128-neuron layer, ±1 reward, no rejection), performance is uniformly low. Introducing temporal and malware reward scaling immediately increases AUT across all settings, highlighting the importance of incorporating both the spatial (malware distribution) and temporal aspects of the domain into the reward function. Increasing the capacity of the model, both by the number of hidden layers and by the number of their neurons, provides marginal improvements.

Adding the integrated reject action without reward R_rej leads to the model not rejecting any samples and to negligible performance differences. Crucially, rewarding rejected outcomes provides a large performance increase as it allows the agent to learn to balance abstention and misclassification risk, with the rejection cost further refining this trade-off. Subsequently, integrated rejection with AL enables the agent to adapt to drift over time and coupled with the sliding retraining window provides performance increases in all but the Transcendent-Drebin setting where DRMD becomes too conservative in rejection. By rectifying this, augmented active learning yields a increase in AUT performance across most settings.

In this work, we focus on the use of uncertainty-based AL and rejection. This paper aims to establish a foundational basis for DRL-based malware detection integrating AL and rejection.

Future work could extend the MD-MDP and DRMD to include additional aspects of the malware detection pipeline, such as feature retrieval, as done in the image domain. Additionally, extending the rejection reward design to integrate sources of information, such as conformal evaluation, explanation, or contrastive learning, provides a potentially valuable next step. Together these directions for future work represent the path forward to a DRL policy that unifies the complete Android malware detection problem.

We introduce a new one-step MDP formulation for malware detection, MD-MDP, which enables the unification of DRL, cost-aware rejection, and AL to counteract concept drift in temporally split evaluations. We use MD-MDP to train the DRMD agent. Our results demonstrate that across two Android malware datasets and two feature spaces, DRMD outperforms the baselines of Drebin, DeepDrebin, Ramda, and SL-DRMD (SL version of DRMD) by an average of 8.66 AUT using the classification-only policy. Furthermore, we show that DRMD learns effective strategies that enable real-time detection, rejection, and AL leading to an average 10.90 AUT increase over the best baselines for all settings using the classification-rejection policy.