N EURAL M ALWARE C ONTROL WITH D EEP R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Antimalware products are a key component in detecting malware attacks, and their engines typically execute unknown programs in a sandbox prior to running them on the native operating system.
Files cannot be scanned indeﬁnitely so the engine employs heuristics to determine when to halt execution.
Previous research has investigated analyzing the sequence of system calls generated during this emulation process to predict if an unknown ﬁle is malicious, but these models require the emulation to be stopped after executing a ﬁxed number of events from the beginning of the ﬁle.
Also, these classiﬁers are not accurate enough to halt emulation in the middle of the ﬁle on their own.
In this paper, we propose a novel algorithm which overcomes this limitation and learns the best time to halt the ﬁle’s execution based on deep reinforcement learning (DRL).
Because the new DRL-based system continues to emulate the unknown ﬁle until it can make a conﬁdent decision to stop, it prevents attackers from avoiding detection by initiating malicious activity after a ﬁxed number of system calls.
Results show that the proposed malware execution control model automatically halts emulation for 91.3% of the ﬁles earlier than heuristics employed by the engine.
Furthermore, classifying the ﬁles at that time improves the true positive rate by 61.5 % , at a false positive rate of 1 % , compared to a baseline classiﬁer.
I NTRODUCTION
Malicious software, or malware , is a serious threat to computer users.
As a ﬁrst line of defense, users and organizations rely on commercial antimalware ( i.e. , antivirus) products to detect malware on their computers, and the antimalware engine is a key component of these malware detection products.
Prior to allowing an unknown ﬁle to be executed on the native operating system, the antimalware engine often tries to detect malware using two main approaches.
First, static analysis employs malware “signatures” ( i.e. , rules) to scan the unknown ﬁle to search for malicious byte sequences in the ﬁle without execution.
Next, the engine utilizes one form of dynamic analysis called emulation to execute the ﬁle in a lightweight sandbox.
Lightweight emulation does not analyze the unknown ﬁle in a full virtual machine (VM).
Instead, the emulator mimics the response of a typical operating system.
If the engine can detect malicious behavior during emulation, the antimalware system blocks the ﬁle from being executed on the native operating system and alerts the user that the ﬁle they are trying to install is malicious.
As a result, the user’s computer is not infected.
Previous dynamic analysis research has focused on analyzing the sequence of system application programming interface (API) calls made by the unknown ﬁle during emulation.
Typically, the authors propose a recurrent, deep learning model to discriminate between the behavior of malicious and benign ﬁles.
Pascanu et al.
(2015) used a recurrent neural network (RNN), or an echo state network (ESN), in combination with either a logistic regression classiﬁer or a multi-layer perceptron (MLP) to detect malware.
Athiwaratkun & Stokes (2017) replace the RNN with a long short-term memory (LSTM) recurrent network or a gated recurrent unit (GRU), and they also propose a character-level convolutional neural network (CNN) to predict if an unknown ﬁle is malicious.
A CNN followed by an LSTM is proposed for this task by Kolosnjaji et al.
In all these solutions, the authors consider a ﬁxed-length input buffer containing the events executed from the beginning of the ﬁle.
The length of this pre-deﬁned window, ν , varies depending on the study with ν ∈ { 50 , 100 , 200 , 65000 } .
In this paper, we propose a novel algorithm, based on deep reinforcement learning (DRL), to overcome this limitation and learn the best time to halt the engine’s emulation to predict whether the unknown ﬁle is malicious or benign.
DRL has been used previously to create adversarial samples to attack a malware classiﬁer (Anderson et al.
Initially, we tried to train a classiﬁer to halt emulation based on the previous t f
( e.g. , 200) behavioral events, but the accuracy of this classiﬁer was not sufﬁcient to effectively halt the emulation.
To the best of our knowledge, this is the ﬁrst paper to propose using deep reinforcement learning to protect users from malware.
This DRL-based neural network, combined with an event classiﬁer and a ﬁle classiﬁer, learns whether to halt emulation after enough state information has been observed or to continue emulation if more events are needed to make a highly conﬁdent prediction.
Unlike previously proposed solutions, the DRL algorithm allows the engine to decide when to stop on a per ﬁle basis.
Results from analyzing a collection of malware and benign ﬁles demonstrates a signiﬁcant improvement in the early stopping of the execution of the ﬁle.
The DRL-based system halts execution of 91.3% of the ﬁles earlier than heuristics used by the production antimalware engine.
When the execution is stopped by the DRL model, the true positive detection rate exhibits a relative increase of 61.5% at a false positive rate of 1.0% compared to the best performing baseline model proposed in Athiwaratkun & Stokes (2017).
Our contributions include the following: 1) we propose a deep reinforcement learning-based system which predicts when to stop emulating an unknown ﬁle, 2) we show that the proposed system signiﬁcantly outperforms several recent neural malware classiﬁcation systems and 3) we provide a theorem and prove that the proposed DRL-based model outperforms these earlier baseline models for malware classiﬁcation.
B ACKGROUND
Engine System Call Events.
The original data for our research was collected by scanning a large collection of Windows Portable Executable (PE) ﬁles with the anonymized for submission production antimalware engine.
The engine collects behavioral events, e t , observed during emulation.
Most of the events are associated with APIs invoked during execution, but other behavioral events, such as unexpected instructions or constructs, are also captured.
In our data, the engine records 114 event types, e t ∈ { 0 , · · · , 113 } , ranging from ﬁle IO, registry APIs, networking APIs, thread or process creation and control, inter-process communication, timing, and debugging APIs.
The behavioral events are logged using a special version of the antimalware engine, and the ﬁles are labeled with
L ∈ { 0 , 1 } where 1 corresponds to a malware ﬁle and 0 indicates that the ﬁle is benign.
Additional details on emulation and scanning can be found in Appendix A, and the threat model is given in
Appendix B.
Conventional Reinforcement Learning.
Conventional Reinforcement learning is normally formulated as a stochastic Markov Decision Process (MDP).
There are four main components in a standard reinforcement learning structure, including States, Actions, Rewards and Policy.
Each of them play a different role in formulating the RL environment.
A general interpretation is that reinforcement learning is a technique to help an agent learn what is the best action and policy to take such that its expected rewards/penalties can be maximized/minimized under a stochastic MDP environment (Sutton & Barto (1998)).
Conventional reinforcement learning deﬁnitions are provided in Appendix C.
N EURAL M ALWARE C ONTROL AND I MPROVED F ILE C LASSIFICATION
An overview of the proposed DRL-based system is shown in Figure 1, and it has two main components: the execution control model and the improved inference model .
An unknown ﬁle is emulated by the antimalware engine, and this generates a sequence of behavioral ﬁle events, E .
The execution control model processes E and is responsible for controlling the ﬁle’s execution.
If the execution control model can make a conﬁdent decision that the ﬁle is either malicious or benign, the execution is halted.
As it is received, each individual event e t is ﬁrst processed by an event classiﬁer which makes a prediction, y e,t , indicating whether or not the most recent event history includes malicious activity.
Initially, we tried halting the emulation based solely on the event classiﬁer’s output, but the classiﬁer’s accuracy was not sufﬁcient to accomplish this task and motivated the need for deep reinforcement learning .
Even though it is a weak signal, y e,t is used to construct a reward signal for the DRL model, which then produces the execution control signal, h t , indicating if the ﬁle execution should be halted or allowed to continue.
[Figure] [Figure]
in Figure 1. In both ﬁgures, the input is e t , and the outputs are
The primary purpose of the DRL model is to better control the ﬁle’s execution.
However, we also found that it can be used to signiﬁcantly improve the overall classiﬁcation of an unknown ﬁle.
This improved prediction, y RL,t , which indicates whether the ﬁle is malicious or benign, is generated by the improved inference model .
The improved inference model boosts the weak predictions from the event classiﬁer , y e,t , based on h t and the output of a baseline ﬁle classiﬁer which offers an initial estimate of the probability, y f , that the ﬁle is malicious based on the initial 200 events generated by the ﬁle.
D EEP E XECUTION C ONTROL
The details of the execution control model in Figure 1 are depicted in Figure 2.
The input is e t , and the outputs are h t and y e,t in both ﬁgures.
As each event is received, it is inserted into the event queue , a ﬁrst in, ﬁrst out (FIFO) queue.
The event classiﬁer then predicts y e,t for the most recent subsequence stored in the event queue.
Since the output layer of the event classiﬁer is a sigmoid function, y e,t is the probability that the most recent subsequence of behavioral events corresponds to malicious activity.
The DRL model depends upon its states, actions, and rewards.
The state, s t , includes information related to all the events received up to and including the most recent event.
For each e t , the event classiﬁer’s prediction y e,t is used as part of the DRL model’s reward function, r t .
The actions for the DRL model include continuing and halting ﬁle execution.
Based in part on y e,t and s t , the DRL model generates separate Q values which are the estimated expected discounted rewards associated with these actions.
The Q value signals are noisy and cannot be used directly.
The action state model ﬁlters the Q values to generate the halting signal h t , which is then used by the antimalware engine to stop the ﬁle’s execution.
Event Classiﬁcation .
The recurrent model structure which is used for the event classiﬁer, and later for the baseline ﬁle classiﬁer, is shown in Figure 3.
For each new event, the event classiﬁer makes a prediction, y e,t , that the behavior associated with the most recent t f behavioral events is malicious.
This event subsequence is stored in the event queue.
Each event in the event queue is input to an embedding layer, and the result is then input to a recurrent layer.
We use a recurrent neural network (RNN) for the recurrent layer.
As proposed in Pascanu et al.
, the recurrent layer’s hidden state is input to a max-pool layer which is able to better detect malicious activity within the subsequence.
We next construct a sparse binary, feature vector consisting of a bag of words (BOW) representation of the event subsequence (114), the ﬁnal hidden state of the recurrent layer which is the recurrent layer’s embedding (1500), and the output of the max-pool layer which is the max-pool embedding (1500).
This feature vector is then input to a shallow neural network in the classiﬁer layer.
The output layer of the neural network is a sigmoid function.
Thus, y e,t is the probability that this most recent event history contains malicious activity.
Deep Reinforcement Learning.
We next present a deep reinforcement learning model to control the antimalware engine’s execution of an unknown ﬁle.
The ﬁrst task is to choose the type of reinforcement learning model for our problem.
The main feature in our problem is a very large state space together with a small action space consisting of two actions A ∈ { continue, halt } .
Considering the small action space, we prefer a value-based reinforcement learning technique which [Figure] [Figure] compares the value functions of the actions directly, instead of learning another policy estimator to ﬁnd the best policy as in policy gradient (Silver et al.
; Levine et al.
or actor-critic-based approaches (Wawrzy´nski (2009); Wawrzy´nski & Tanwani (2013); Lillicrap et al.
The action-value function, Q ∗ ( s t , a t ) = max π E
[ R t | a t , s t , π ] , is the expected reward of taking action a t at state s t following the policy π .
To calculate this value using a conventional value-based approach, it is necessary to store all the Q values in a table for all the state-action pairs encountered during training, which is not feasible if the state space or action space is large.
One approach to overcome this difﬁculty is by training a nonlinear approximator Q ( s t , a t | θ t ) , such as a deep neural network, to estimate Q ∗ ( s t , a t ) at time step t .
However, these types of nonlinear estimators tend to be unstable in practical applications since convergence is not guaranteed.
To address this issue, Mnih et al.
(2013; 2015) recently proposed using a replay buffer in the deep Q network (DQN) which demonstrates better convergence properties.
Since our problem has a very large state space, we also use a DQN as our DRL model structure in this paper.
In a DQN-based DRL structure, the deep neural network action-value function estimator Q ( s t , a t | θ t ) is normally deﬁned at state s t as Q ( s t , a t | θ t ) ∼ Q ∗ ( s t , a t ) by taking the state s t as the input of the neural network, where θ t represents the neural network parameters.
We next describe the state design, action design, reward design, and training for the DRL model using experience replay.
Design of States.
The DRL model’s state, which is illustrated in Figure 5, contains three parts: the position ( i.e. , index) of the current event in the ﬁle ρ t , the current event ID, and the histogram of all the previous events.
The current event position in the ﬁle will be used later to deﬁne the reward r t for the deep reinforcement learning model.
We initially tried to use a one-hot encoding of the event ID, but found that using the event ID directly provides slightly better performance while reducing the size of the state.
The event ID histogram captures the history of the events which have been observed so far.
[Figure]
Design of Actions.
The agent ( i.e. , antimalware engine) can perform two types of actions a t : continue , which is labeled as C , and halt , labeled as H .
The action indicates whether the agent should continue or halt the execution of the ﬁle.
In our deep reinforcement learning model, the selected action a t for state s t is inferred from the output of the neural network.
As shown in the
Figure 2, the outputs of the deep neural network are the estimated action value function Q C =
Q ∗ ( s t , a t = C ) = max π E
[ R t | a t = C, s t , π ] for action C at state s t and Q H =
Q ∗ ( s t , a t = H ) = max π E [ R t | a t = H, s t , π ] for action H at state s t . By comparing the two Q ∗ ( s t , a t ) values for actions H and C , the action with the larger Q value is selected and performed. Design of Rewards.
The reward r t at each state s t is designed based on two criteria:
We prefer for the DRL network to learn to halt emulation as quickly as possible.
Therefore, shorter emulation sequence lengths are assigned a higher reward and longer sequence lengths are given a smaller reward.
The closer an event prediction is to the true label of the ﬁle, the larger the reward should be given at that state.
Based on the above two criteria, the reward is deﬁned as
[Equation]
where y e,t is the event-based prediction generated by the most recent t f = 200 events, and L ∈ { 0 , 1 } is deﬁned as the true label of the training ﬁle.
The decay factor β is chosen experimentally, and ρ t is the position of current event in the ﬁle.
DRL Training.
To train this neural network-based estimator, we use an l 2 loss function deﬁned as L ( θ t ) =
E s t [( ˆ Q ( s t , a t | θ t ) − Q ( s t , a t | θ t )) 2 ] where ˆ Q ( s t , a t | θ t ) is an estimate of Q ( s t , a t | θ t ) .
ˆ
Q ( s t , a t | θ t ) is computed using the current state reward r t together with its neighbors’ estimations from the neural network in an iterative manner, i.e. , ˆ Q ( s t , a t | θ t ) = r t + [Equation] +
[Equation] t + [Equation] | θ t ) where s t +1 are the neighbors of s t , and a t +1 are the corresponding actions generated by the neural network.
In Mnih et al.
(2013), experience replay is used to train the DRL model.
Experience replay helps to alleviate the potential issues of non-stationary distributions and correlated data and is performed by randomly sampling the state pairs.
Whenever one stochastic step is taken by the agent, the current state s t , obtained reward r t , action taken a t and next state s t +1 are combined as one agent experience set ( s t , r t , a t , s t +1 ) , and pushed into the replay memory queue M .
The reinforcement learning updates are performed in minibatches of size B RL , drawing from the replay memory randomly.
The algorithm for training the DRL model using experience replay is provided in Appendix D.
We tested several different stochastic gradient descent optimization methods for training the DRL model and found that adadelta (Zeiler (2012)) performed best.
Furthermore, the convergence of DRL is not always guaranteed.
To help with the convergence, the sum of r t + γ should be within the range of [0,1].
It is important to note that we ﬁrst train the event classiﬁer, as well as the baseline ﬁle classiﬁer in the improved inference model , in isolation prior to training the DRL model – the system is not trained in an end-to-end fashion.
Thus, y e,t in the reward r t is generated by the pre-trained event classiﬁer, and the reward function has the same value for the same event sequence.
Otherwise, the DRL’s reward function can become non-stationary.
Action State Model.
The ﬁnal block in Figure 2 is the action state model which generates the halting signal h t .
The Q value signal is noisy and cannot be used directly to compute h t .
The halting signal is a binary signal which is initialized to have a value of 0 and remains 0 for each e t until Q H > Q C for K consecutive events.
At that point, the value of h t transitions to 1 and continues to maintain that value if any additional events are processed.
I MPROVED I NFERENCE
The purpose of the improved inference model in Figure 1 is used to generate a better DRL-based prediction, y RL,t , of whether or not an unknown ﬁle is malicious.
The details of improved inference model are depicted in Figure 4.
This model has three inputs, namely the events generated by the ﬁle ( e t ), the most recent event predictions ( y e,t ), and the execution control signal ( h t ).
The most recent K values of y e,t are stored in the event prediction queue , which is another FIFO queue.
After h t signals that the ﬁle execution has been halted, the improved ﬁle score model evaluates the event predictions in the queue to generate y RL,t .
The individual y e,t values, stored in the event prediction queue , are noisy and can be difﬁcult to analyze.
In some cases, only setting y RL,t to be the most recent y e,t value can lead to an incorrect prediction that the ﬁle is malicious.
To overcome this issue, the improved inference model also employs an additional baseline ﬁle classiﬁer to improve the accuracy of y RL,t .
To accomplish this task, the initial t f ( e.g. , 200) events generated by the ﬁle are stored in the initial events buffer , and these are processed by the baseline ﬁle classiﬁer to produce an initial prediction that the ﬁle is malicious, y f .
Baseline File Classiﬁer.
The baseline ﬁle classiﬁer also utilizes the structure shown in Figure 3 and follows Athiwaratkun & Stokes (2017).
The input event subsequence in Figure 3 corresponds to the ﬁrst t f events for each ﬁle stored in the initial events buffer (Figure 4).
Here, t f is the same value which denotes the length of the event classiﬁer’s input event queue .
An LSTM is used for the recurrent layer, and the classiﬁer layer uses logistic regression for the ﬁle’s prediction y f .
Similar to the event classiﬁer, y f is the initial estimate of the probability that the ﬁle is malicious based on the initial behavior of the ﬁle.
Improved File Score.
We can now combine y f with the y e,t history stored in the event prediction queue to compute the ﬁnal improved ﬁle classiﬁer score, y RL,t .
Since the initial estimates of the y e,t are noisy, we process the most recent K event predictions from the event prediction queue.
Formally if h t is equal to 1,
[Equation]
where y e,i is the event classiﬁer’s prediction at step i , and y RL,t is the improved inference model’s output, i.e. , the prediction probability that the unknown ﬁle is malicious.
It is important to provide mathematical insight why (2) provides better performance compared to the baseline model which only uses a deep learning-based classiﬁer to process the ﬁrst t f events.
The intuition is that our new model provides a closer prediction to the true label ( e.g. , malware or benign) compared to a single deep learning ﬁle classiﬁer.
Theorem 1.
Let E s tf
[ R t f ] be the expected reward obtained at the state s t f corresponding to event e t f , where t f is the ﬁxed number of events used to train the baseline ﬁle classiﬁer and is also the number of the most recent events used to train the event classiﬁer.
E s ∗ [ R ∗ ] is the expected reward obtained at the halting state s ∗ where s ∗ is selected by the improved inference model when h t transitions from 0 to 1.
If all models are fully trained and the decay factor β deﬁned in reward is small enough, E s tf
[ R t f ] ≤ E s ∗ [ R ∗ ] and | y ∗ − L | ≤ 1 2 where y ∗ ∈ [0 , 1] is the predicted probability of a ﬁle being malicious using the improved inference model, and L is the true label of the given ﬁle.
Then,
[Equation]
where y f is the predicted malware probability using a baseline ﬁle classiﬁer.
This theorem is proved in Appendix E.
Evasion of adversarial learning-based attacks is another important aspect to consider.
Evasion of the model discussed in Appendix F.
E XPERIMENTAL R ESULTS
We next present the results for the proposed neural malware control model.
We ﬁrst describe the datasets which were collected for our experiments.
Next, we present the experimental setup.
We then evaluate how quickly the DRL-based model halts the execution of a ﬁle.
Finally, we compare the ﬁnal prediction that the ﬁle is malicious or benign to the results from several baseline ﬁle classiﬁers.
Datasets.
We collected a dataset of 75 thousand ﬁles which had been evaluated by a production antimalware engine.
These ﬁles were equally split between the malware and benign classes.
First, we discarded any ﬁles whose event sequences were shared between these two classes and which contained less than 50 events.
Furthermore, we ensured that all the event sequences in the datasets were distinct.
This requirement ensured that we did not overﬁt to one particular set of events.
We then split the overall dataset into separate training, validation, and test sets with 50, 10, and 15 thousand ﬁles, respectively.
Again, we maintained an equal split between the two classes for each of the individual datasets.
The event and ﬁle classiﬁers were trained with the training and validation sets, while the DRL-model was trained with 2000 ﬁles from the training set.
The results presented below are based on evaluating the model on the hold out test set. [Figure] [Figure]
Setup.
We implemented the proposed neural malware control model using Keras (Chollet et al.
(2015)) with Theano (Al-Rfou et al.
as the backend deep learning framework.
Several hyperparameters were tuned on smaller datasets.
The decay factor in the DRL model’s reward function is γ = 0 .
The DRL model uses 3 hidden layers and is trained with a minibatch size
B RL = 50 .
The replay memory is initialized with a size µ = 50000 .
How does the DRL-based model’s stopping performance compare to the antimalware engine’s heuristics?
The ﬁles in our dataset were collected by a production antimalware engine, and the number of events recorded for each ﬁle represents the performance of the heuristics employed by the engine to halt emulation.
Thus, by measuring how often the DRL-based model halts execution prior to reaching the end of the ﬁle, we can compare the performance between our model and the engine’s heuristics.
In cases where the DRL model reaches the end of the ﬁle without halting execution, we can infer that the proposed model was not conﬁdent enough to make a decision, and the DRL based-model would have continued to execute the ﬁle.
The results of this evaluation are presented in Table 1 and depend on two values: the number of training ﬁles ( i.e. , epochs) N and the number of consecutive events where Q H > Q C denoted by K .
The fraction of ﬁles where the DRL-model halts execution before the end of the ﬁle, α , is computed as: α =
(Total number of early halted ﬁles)/(Total number of ﬁles).
We make two observations from the results presented in the table.
First, the percentage of ﬁles whose execution is halted by the DRL model earlier than engine’s heuristics continues to increase as the number of training ﬁle N increases.
Better training allows the engine to halt execution earlier.
Second, the percentage of ﬁles which are halted early decreases with K .
The value of K is a proxy for the DRL model’s conﬁdence in the decision to halt the ﬁle’s execution.
It is not surprising that the execution of fewer ﬁles is halted early as we require more conﬁdence ( i.e. , higher value of
K ) in the decision.
Even so, the results show that over 91% of the ﬁles in the test set are halted early compared to the engine’s heuristics after training with only 2000 ﬁles.
This indicates that the engine’s heuristics may be overly cautious when emulating a ﬁle.
In addition, requiring less time for scanning an individual ﬁle leads to better performance when scanning all the ﬁles on the hard drive.
A histogram which indicates the percentage of events which were executed before the DRL model halted execution is provided in Appendix G.
[Table]
The fraction of ﬁles, α , in % where emulation is halted earlier by the proposed deep reinforcement learning model compared to heuristics used by the antimalware engine.
Can DRL improve ﬁle classiﬁcation?
While the results in Table 1 indicate that emulation of the majority of the ﬁles can be stopped earlier than the heuristics employed in the engine, it is important to understand how early halting affects the detection performance.
To measure this, we ﬁrst compute the receiver operating characteristic (ROC) curves for a range of models in Figure 6 for the baseline ﬁle classiﬁer depicted in Figure 3.
The models include LSTM Athiwaratkun & Stokes (2017), RNN Pascanu et al. (2015), gated recurrent unit (GRU) Athiwaratkun & Stokes (2017), convolutional neural network (CNN) Kolosnjaji et al.
(2016) and a simple, single hidden layer, feedforward neural network.
We also include the echo state network (ESN) counterparts for the LSTM and GRU.
All models use logistic regression (LR) for the classiﬁer, because these slightly outperformed a shallow neural network for this dataset.
All models use max-pooling as shown in Figure 3.
These results indicate that the performance of none of the models we investigated dominated all of the other models.
In particular, the ESN version of the GRU offers better performance and low false positive rates (FPRs) while the LSTM outperforms all other models above an FPR ≥ 1 . 2% .
We next compare the two best performing baseline ﬁle classiﬁers to the proposed DRL-based models for K ∈ { 10 , 15 , 20 } in Figure 7.
The ﬁgure clearly indicates that all the DRL-based models offer signiﬁcantly better performance compared to the baseline ﬁle classiﬁers.
In particular, the DRL-based model with K = 20 offers a relative improvement of 61.5 % for the true positive rate (TPR) at an
FPR of 1 % compared to the GRU-ESN-based baseline ﬁle classiﬁer.
The relative improvement of the
TPR is 65.7 % at an FPR of 1 % for the LSTM.
R ELATED W ORK
Deep Reinforcement Learning.
Conventional reinforcement learning has been widely studied in the ﬁelds of machine learning and system control for over three decades (Sutton (1984);
Williams (1992); Littman (1994); Kaelbling et al.
; Sutton & Barto (1998)).
Recently, Mnih et al.
(2013; 2015) successfully applied deep neural network-based Q-learning (DQN) to playing a series of Atari games, by using a replay buffer to improve the system’s convergence.
Also, Silver et al.
(2016; 2017) developed novel algorithms, by applying reinforcement learning to Monte Carlo tree search, to play the Go game and beat human Go masters.
Progress has also been made in improving value-based (Van Hasselt et al.
; Wang et al.
policy gradient (Schulman et al.
(2015); Duan et al.
; Gu et al.
(2017); O’Donoghue et al.
(2017)), and actor-critic (Lillicrap et al.
; Mnih et al.
deep reinforcement learning algorithms, in order to ﬁnd better policies more efﬁciently and to deal with a continuous action space.
Despite the success of applying DRL to different types of games, no existing research has studied malware control and behavioral classiﬁcation using reinforcement learning.
The proposed neural malware control model is the ﬁrst to apply a deep reinforcement learning-based algorithm for these tasks.
Deep Learning for Malware Classiﬁcation.
A number of authors have proposed using DNNs for malware classiﬁcation tasks.
Kephart (1994) employed a shallow neural network in the ﬁrst paper on malware classiﬁcation.
Dahl et al.
ﬁrst investigated the use of a DNN for malware classiﬁcation for dynamic analysis.
Huang & Stokes (2016) proposed a DNN with multitask learning for dynamic analysis where the ﬁrst task was binary ( i.e. , malware versus benign) and the second task was malware family classiﬁcation.
Saxe & Berlin (2015) proposed a DNN for the static analysis of malware.
A separate line of research has investigated using recurrent models for malware classiﬁcation.
Pascanu et al.
ﬁrst proposed using recurrent neural networks and echo state networks for classifying malware sequences.
Athiwaratkun & Stokes (2017) instead proposed an LSTM, a GRU, and a character-level CNN for the sequence classiﬁcation.
Kolosnjaji et al.
(2016) used a CNN followed by an LSTM for this task.
C ONCLUSION
We present a novel, neural malware control model which learns when to halt the execution of an unknown ﬁle based on deep reinforcement learning.
This model is the ﬁrst to use deep reinforcement learning to protect customers from malware.
Fast scanning is an important feature for users when installing software, reading emails with attachments, and searching a hard drive for ﬁles which were maliciously dropped during a drive-by download.
Our results indicate that the proposed model halts execution earlier than a production antimalware engine for more than 91% of the ﬁles in the test set.
More importantly, we show a relative improvement of over 61% in the true positive detection rate of malware at a false positive rate of 1% compared to a number of baseline malware classiﬁers reported in the literature.
Thus, the proposed model offers signiﬁcantly better protection with less delay.
A PPENDIX A: E MULATION AND S CANNING D ETAILS
In some cases, unknown ﬁles are evaluated in a protected sandbox environment to determine if they are malicious or benign prior to running them on the native operating system.
This method of testing unknown ﬁles is called dynamic analysis, and there are two main modes of performing dynamic analysis: virtualization and emulation.
Virtualization involves running an unknown ﬁle in an instrumented version of a virtual machine (VM) which implements a full operating system.
Performing analysis in a VM is the preferred method of analysis because it can fully implement all of the functionality required by the unknown ﬁle.
Virtualization is often used in backend security services to protect customers.
Advanced email services scan email attachments before allowing them to be displayed in the user’s inbox.
Antimalware companies also scan unknown ﬁles in full VMs in order to create new signatures and labeled malware ﬁles for training machine learning classiﬁers.
Emulation is often performed in antimalware engines running on the endpoint computer or cell phone.
Emulation is similar to virtualization in that it tries to induce an unknown malware ﬁle into revealing its malicious behavior.
However, emulation tends to be more restricted because it needs to run much faster and consume less resources compared to the execution of the ﬁle in a virtual machine.
Thus, emulation can also be run in backend services to quickly analyze signiﬁcantly larger numbers of unknown ﬁles compared to virtualization.
If the antimalware engine does not detect that a ﬁle is malicious during emulation, it is then allowed to be installed and executed on the native operating system.
The data analyzed in this study was generated using emulation in a production antimalware engine.
This data was generated in a production emulation scanning environment of a major antimalware computer, instead of by client computers located in the wild.
The behavior revealed in our data is similar to that which would be encountered by the ﬁle running on the client computer’s operating system except for one main difference.
The antimalware engine does not allow the ﬁles to have internet access because this may lead to the infection of other computers in the network.
For emulation or virtualization, the behavioral event sequence corresponding to the system APIs called by the ﬁle are typically recorded.
Attackers often use polymorphic tactics to avoid detection.
In polymorphism, they rearrange or rewrite their code in different ways which appear to be different but accomplish the same task.
To deal with polymorphism, our antimalware engine maps multiple low-level API calls into a single high-level event.
For example, the attacker may use a user mode API (CreateFile), a kernel mode API (ZwCreateFile), or the C++ API (ofstream::open) to create a ﬁle.
All of these events are mapped into the same high-level FileCreate event.
As a result, the 114 behavior events represent many individual low-level API calls.
The labels utilized in this study correspond to production labels used by our antimalware product partners to train malware classiﬁers that identify malware targeting Windows computers.
Benign ﬁles are identiﬁed as those which are known to be safe.
For example, these might be ﬁles belonging to software products which are purchased by users ( e.g. Microsoft Ofﬁce) or downloaded from the internet from sites which are known to be legitimate ( e.g. Adobe Acrobat Reader, Google Chrome).
Other benign ﬁles are determined to be safe by professional analysts.
Labels for malware ﬁles are also generated by manual inspection by professional analysts. In addition, all unknown ﬁles received by the company are scanned by over 20 additional production antimalware engines. If eight or more of these anti-malware engines detect that a ﬁle is malicious, these ﬁles will also be determined to be malicious. A PPENDIX B: T HREAT M ODEL
We next specify the threat model which deﬁnes the assumptions which are made about the proposed detection system and the attacker.
For our system, we require that the computer has not been previously infected, and the trusted computing boundary includes the user account, operating system, and antimalware detection system.
Otherwise, the attacker can alter the detection system’s results and avoid detection.
We also assume that the behavior observed in the antimalware engine’s emulator without internet access is similar to that observed when the ﬁle is executed on the native operating system of the user’s computer.
Some malware tries to identify whether or not it is being emulated or virtualized.
If the malware detects it is being emulated or virtualized, it may disable any further malicious activity ( i e. , cloaking).
Based on our current datasets, our system may fail to detect any malware which successfully employs these cloaking mechanisms.
A PPENDIX C: C ONVENTIONAL R EINFORCEMENT L EARNING D EFINITIONS
In this appendix, we provide deﬁnitions of the elements in conventional reinforcement learning.
Agent and States ( s t ).
An agent interacts with its environment by moving from the current state s t at time t to another state s t +1 at time t
+ 1 .
Each state is normally deﬁned by the useful information from the interaction between an agent and its environment.
Actions ( a t ).
By taking an action a t at state s t , an agent can transfer from its current state to any of its connected neighbors at its next state s t +1 with different probabilities, since the agent can only arrive at one of its neighbors at t + 1 . 3.
Rewards ( r t ).
The agent receives reward r t at time t .
The discounted reward R t is deﬁned as R t = (cid:80) ∞ t
= t 0 γ t − t 0 r t where γ is the discount factor with { 0 ≤ γ ≤ 1 } , and t 0 indicates the starting time step.
After reaching a state, an agent obtains the expected discounted reward ( E [ R t | a t , s t ] ) by considering the policies from the current state s t to its neighbors s t +1 and so on.
The expected discounted reward includes both the pre-deﬁned reward at state s t and the accumulated discounted rewards to be obtained in the future by taking a speciﬁc action a t . 4.
Policy ( π ).
A policy π is a mapping from states to actions.
There are three main types of reinforcement learning: value-based, policy gradient and actor-critic.
In this paper, we focus on a value-based algorithm called Q-learning given the small action space in our problem.
The optimal action-value function in conventional Q-learning is deﬁned as Q ∗ , which is the maximum expected reward obtained by selecting the best policy π at state s t
[Equation]
A PPENDIX D: A DRL A LGORITHM USING E XPERIENCE R EPLAY
The detailed explanation of the DRL algorithm using experience replay in section 4.2 is provided in
Algorithm 1.
A PPENDIX E: P ROOF OF T HEOREM
We present a detailed proof of Theorem 1 in this appendix.
Proof.
We begin by assuming that both the baseline ﬁle classiﬁer and event classiﬁer are fully trained.
By deﬁnition, the expected reward E s tf
[ R t f ] obtained at the state s t f corresponding to event e t f can be calculated as [Equation]
[Equation]
) where t f is the ﬁxed number of events processed by the baseline ﬁle classiﬁer, p t f is the probability of remaining at state s t f , and E s tf +
[ R t f +1 ] represents the expected reward to be obtained by leaving s t f for its neighboring state s t
f +1 including the t f + 1 event.
After substituting the reward (1),
[Equation]
(6) where y e,t f is the predicted probability generated by the event classiﬁer using events [ e 1 , · · · , e t f ] (assuming the ﬁle has more than t f events).
Similarly, E s tf +
[ R t f +1 ] can be written as
[Equation]
(7) where C t f +1 =
(1 − p t f +1 )
E s tf +
[ R t f +2 ] .
By choosing a small value of β , we may approximate that e − βρ tf ≈ e − βρ tf +1 since ρ t f +1 =
ρ t f + 1 .
The event classiﬁer should yield similar predicted probabilities for state s t f and s t
f +1 in a given ﬁle 1: Epochs: N ←2: Minibatch Size: B RL ←3: Decay Factor: β ← 0 .4: Initialize a replay memory M with size µ ← 50000 , DRL model with 3 layers 5: for n=1 → N do 6: Time step in state space: t ←7:
Randomly select an initial state s t 8: while ! End of File do 9: Q ( s t , a t | θ t ) ← DRL( s t ) 10: a ∗ t = arg max a t Q ( s t , a t | θ t ) 11: Perform action a ∗ t , generating next state s t + 12:
Push tuple ( s t , r t , a ∗ t , s t +1 ) into replay memory M 13: for b=1 → B RL do 14:
Randomly select a tuple m from M 15: s t ← m (0) , r t ← m (1) , s t +1 ← m (3) 16:
Q ( s t , a t | θ t ) ← DRL( s t ) 17:
Q ( s t +1 , a t +1 | θ t ) ← DRL( s t +1 ) 18:
Input y e,t from Event Classiﬁer 19: r t ← (0 .
5 − | y e,t − L | ) × e − βρ t 20:
Update ˆ Q ( s t , a t | θ t ) 21:
Update the network by minimizing loss L ( θ t ) 22: end for 23: t ← t +24: end while 25: end for if the classiﬁer is fully trained, i.e. y e,t f ≈ y e,t f +1 .
Using these two approximations and substituting (7) into (6), (6) can be rewritten as
[Equation]
Since there are only two possible actions (halt, continue), the probability of remaining at a particular state p t at step t is either 0 or 1.
If the system halts before t f events, p t f = p t f +1 = 1 .
Similarly, if the system continues until at least t f + 1 events, p t f = p t
f +1 = 0 .
Assuming the system does not transition from continue to halt at event t
= t f , p t f
+ (1 − p t f ) p t
f +1 = 2 p t f − p 2 t f =
p t f .
Under this weak assumption, (8) can be rewritten as
[Equation]
Similarly, E s ∗ [ R ∗ ] , the expected rewards obtained at the halting state s ∗ chosen by the DRL-based model, can be calculated as
[Equation]
where p ∗ is the probability of remaining at state s ∗ , y ∗ is the predicted probability generated by the event classiﬁer using the most recent t f events before the halting state s ∗ , ρ ∗ is the position of the event contained in state s ∗ , and E s ∗(cid:48)
[ R ∗ (cid:48) ] is the expected reward at state s ∗ (cid:48) which is the next state after s ∗ .
Since s ∗ is an absorbing ( i.e. , terminal) state, p ∗ = 1 .
Thus,
[Equation] [Equation]
and the expected reward at s ∗ is larger or equal than those at any other states
[Equation]
Substituting (9) and (11) into (12),
[Equation]
Recalling p t f is equal to either 0 or 1 and ﬁrst considering the case when p t f is equal to 1, (13) can be further rewritten as
[Equation]
Recall that the decay factor β is deﬁned to be positive.
If ρ t f < ρ ∗ , then 0 < e β ( ρ tf − ρ ∗ ) < 1 , and (14) can be rewritten as
[Equation]
If ρ t f > ρ ∗ , then e β ( ρ tf − ρ ∗ ) > 1 , and (14) can be rewritten as
[Equation]
By choosing small enough β such that e β ( ρ tf − ρ ∗ ) → 1 , the second term in ( 16 ) is approaching zero, then | y e,t f − L | ≥ | y ∗ − L | .
To prove it strictly, choose a β satisfying
[Equation]
By substituting (17) into (16), we can derive that | y e,t f − L | ≥ | y ∗ − L | .
Given that | y ∗ − L | ≤ 1 2 , (17) can provide an upper bound of β as
[Equation]
The proceeding derivation was for the case when p t f = 1 .
Following (5) for the case when p t f = 0 , we have E s tf
[ R t f ] =
E s tf +
[ R t f +1 ] .
Similarly we have E s tf +
[ R t f +1 ] =
E s tf +
[ R t f +2 ] since p t f +1 = 0 , and so on.
Thus, E s tf
[ R t f ] =
E s tf + n [ R t f + n ] where s t f + n is the ﬁrst halting state after s t f , such that p t f + n is equal to 1.
Then we can substitute t f by t f + n in our proof in order to show the same result that | y e,t f − L | ≥ | y ∗ − L | for the case when p t f = 0 .
While we have now shown that | y e,t f − L | ≥ | y ∗ − L | for the event classiﬁer, we need to consider the prediction probability of the baseline ﬁle classiﬁer.
The predicted probability y e,t f generated by the event classiﬁer using the ﬁrst t f events is an approximation of y f , which is the predicted probability generated by the baseline ﬁle classiﬁer also using the ﬁrst t f events, i.e. y e,t f ≈ y f , hence we have
[Equation]
It can be concluded that the DRL predicted probability y ∗ is closer to the true label than the baseline ﬁle classiﬁer’s predicted probability y f .
A PPENDIX F: E VASION
Recently, researchers have begun to investigate adversarial attacks on deep reinforcement learning models.
Since the proposed DRL-based system would be run on a user’s computer or cell phone, it is particularly vulnerable to adversarial attacks.
The basic strategy is to destabilize the DRL network by adding a perturbation to a state using a second adversarial agent (Lin et al.
; Pinto et al.
Huang et al.
This research considers two main scenarios:
The attacker knows both the architecture and parameters of a trained neural network policy π .
Thus, they can compute a perturbation directly based on the policy’s parameters.
The attack is a black-box problem assuming the adversary has no access to the target policy, the training algorithm or the model’s parameters.
One possible defense against adversarial attacks on the client is to either run the DRL model, or the entire antimalware engine, in a secure enclave, such as SGX (Ohrimenko et al.
Doing so prevents the attacker from directly observing the model structure and parameters, and the attacker must resort to a black-box type of attack (Huang et al.
Thus, the only way to build an adversarial agent is through the transferability across training algorithms, by guessing the target policy π from another policy trained using a different dataset.
Protecting a DRL-model running outside of a secure enclave from adversarial attacks is an open research topic.
Some promising approaches include the ensemble defense, feature squeezing or building a detection model to discover attacks (Tramer et al.
; Zantedeschi et al.
; Ilyas et al.
A PPENDIX G: E ARLY S TOPPING D ISTRIBUTION
In Table 1 in Section 6, we report that with K = 20 and N = 2000 , the execution control model halts the execution of 91.3% of the ﬁles earlier than the heuristics used by the antimalware engine.
In Figure 8, we provide the histogram indicating the distribution of the percentage of events which are executed by the engine prior to the execution control model halting its execution.
This ﬁgure indicates that execution is halted much earlier than the heuristics for the majority of the ﬁles.