NEURAL MALWARE CONTROL WITH DEEP REINFORCEMENT LEARNING

1 INTRODUCTION . Malicious software , or malware , is a serious threat to computer users . As a first 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 file 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 file to search for malicious byte sequences in the file without execution . Next , the engine utilizes one form of dynamic analysis called emulation to execute the file in a lightweight sandbox . Lightweight emulation does not analyze the unknown file 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 file from being executed on the native operating system and alerts the user that the file 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 file during emulation . Typically , the authors propose a recurrent , deep learning model to discriminate between the behavior of malicious and benign files . Pascanu et al . ( 2015 ) used a recurrent neural network ( RNN ) , or an echo state network ( ESN ) , in combination with either a logistic regression classifier 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 file is malicious . A CNN followed by an LSTM is proposed for this task by Kolosnjaji et al . ( 2016 ) . In all these solutions , the authors consider a fixed-length input buffer containing the events executed from the beginning of the file . The length of this pre-defined 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 file is malicious or benign . DRL has been used previously to create adversarial samples to attack a malware classifier ( Anderson et al . ( 2018 ) ) . Initially , we tried to train a classifier to halt emulation based on the previous tf ( e.g. , 200 ) behavioral events , but the accuracy of this classifier was not sufficient to effectively halt the emulation . To the best of our knowledge , this is the first paper to propose using deep reinforcement learning to protect users from malware . This DRL-based neural network , combined with an event classifier and a file classifier , 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 confident prediction . Unlike previously proposed solutions , the DRL algorithm allows the engine to decide when to stop on a per file basis . Results from analyzing a collection of malware and benign files demonstrates a significant improvement in the early stopping of the execution of the file . The DRL-based system halts execution of 91.3 % of the files 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 file , 2 ) we show that the proposed system significantly outperforms several recent neural malware classification systems and 3 ) we provide a theorem and prove that the proposed DRL-based model outperforms these earlier baseline models for malware classification . 2 BACKGROUND . Engine System Call Events . The original data for our research was collected by scanning a large collection of Windows Portable Executable ( PE ) files with the anonymized for submission production antimalware engine . The engine collects behavioral events , et , 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 , et ∈ { 0 , · · · , 113 } , ranging from file 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 files are labeled with L ∈ { 0 , 1 } where 1 corresponds to a malware file and 0 indicates that the file 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 definitions are provided in Appendix C . 3 NEURAL MALWARE CONTROL AND IMPROVED FILE CLASSIFICATION . 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 file is emulated by the antimalware engine , and this generates a sequence of behavioral file events , E. The execution control model processes E and is responsible for controlling the file ’ s execution . If the execution control model can make a confident decision that the file is either malicious or benign , the execution is halted . As it is received , each individual event et is first processed by an event classifier which makes a prediction , ye , t , indicating whether or not the most recent event history includes malicious activity . Initially , we tried halting the emulation based solely on the event classifier ’ s output , but the classifier ’ s accuracy was not sufficient to accomplish this task and motivated the need for deep reinforcement learning . Even though it is a weak signal , ye , t is used to construct a reward signal for the DRL model , which then produces the execution control signal , ht , indicating if the file execution should be halted or allowed to continue . The primary purpose of the DRL model is to better control the file ’ s execution . However , we also found that it can be used to significantly improve the overall classification of an unknown file . This improved prediction , yRL , t , which indicates whether the file is malicious or benign , is generated by the improved inference model . The improved inference model boosts the weak predictions from the event classifier , ye , t , based on ht and the output of a baseline file classifier which offers an initial estimate of the probability , yf , that the file is malicious based on the initial 200 events generated by the file . 4 DEEP EXECUTION CONTROL . The details of the execution control model in Figure 1 are depicted in Figure 2 . The input is et , and the outputs are ht and ye , t in both figures . As each event is received , it is inserted into the event queue , a first in , first out ( FIFO ) queue . The event classifier then predicts ye , t for the most recent subsequence stored in the event queue . Since the output layer of the event classifier is a sigmoid function , ye , 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 , st , includes information related to all the events received up to and including the most recent event . For each et , the event classifier ’ s prediction ye , t is used as part of the DRL model ’ s reward function , rt . The actions for the DRL model include continuing and halting file execution . Based in part on ye , t and st , 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 can not be used directly . The action state model filters the Q values to generate the halting signal ht , which is then used by the antimalware engine to stop the file ’ s execution . Event Classification . The recurrent model structure which is used for the event classifier , and later for the baseline file classifier , is shown in Figure 3 . For each new event , the event classifier makes a prediction , ye , t , that the behavior associated with the most recent tf 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 . ( 2015 ) , 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 final 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 classifier layer . The output layer of the neural network is a sigmoid function . Thus , ye , 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 file . The first 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 compares the value functions of the actions directly , instead of learning another policy estimator to find the best policy as in policy gradient ( Silver et al . ( 2014 ) ; Levine et al . ( 2016 ) ) or actor-critic-based approaches ( Wawrzyński ( 2009 ) ; Wawrzyński & Tanwani ( 2013 ) ; Lillicrap et al . ( 2015 ) ) . The action-value function , Q∗ ( st , at ) = maxπ E [ Rt|at , st , π ] , is the expected reward of taking action at at state st 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 difficulty is by training a nonlinear approximator Q ( st , at|θt ) , such as a deep neural network , to estimate Q∗ ( st , at ) 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 ( st , at|θt ) is normally defined at state st as Q ( st , at|θt ) ∼ Q∗ ( st , at ) by taking the state st 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 file ρt , the current event ID , and the histogram of all the previous events . The current event position in the file will be used later to define the reward rt 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 . Design of Actions . The agent ( i.e. , antimalware engine ) can perform two types of actions at : 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 file . In our deep reinforcement learning model , the selected action at for state st 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 QC = Q∗ ( st , at = C ) = maxπ E [ Rt|at = C , st , π ] for action C at state st and QH = Q∗ ( st , at = H ) = maxπ E [ Rt|at = H , st , π ] for action H at state st. By comparing the two Q∗ ( st , at ) values for actions H and C , the action with the larger Q value is selected and performed . Design of Rewards . The reward rt at each state st is designed based on two criteria : 1 . 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 . 2 . The closer an event prediction is to the true label of the file , the larger the reward should be given at that state . Based on the above two criteria , the reward is defined as rt = ( 0.5− |ye , t − L| ) e−βρt ( 1 ) where ye , t is the event-based prediction generated by the most recent tf = 200 events , and L ∈ { 0 , 1 } is defined as the true label of the training file . The decay factor β is chosen experimentally , and ρt is the position of current event in the file . DRL Training . To train this neural network-based estimator , we use an l2 loss function defined as L ( θt ) = Est [ ( Q̂ ( st , at|θt ) − Q ( st , at|θt ) ) 2 ] where Q̂ ( st , at|θt ) is an estimate of Q ( st , at|θt ) . Q̂ ( st , at|θt ) is computed using the current state reward rt together with its neighbors ’ estimations from the neural network in an iterative manner , i.e. , Q̂ ( st , at|θt ) = rt + γmaxat+1 Q ( st+1 , at+1|θt ) where st+1 are the neighbors of st , and at+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 st , obtained reward rt , action taken at and next state st+1 are combined as one agent experience set ( st , rt , at , st+1 ) , and pushed into the replay memory queue M . The reinforcement learning updates are performed in minibatches of sizeBRL , 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 rt + γ should be within the range of [ 0,1 ] . It is important to note that we first train the event classifier , as well as the baseline file classifier 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 , ye , t in the reward rt is generated by the pre-trained event classifier , 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 final block in Figure 2 is the action state model which generates the halting signal ht . The Q value signal is noisy and can not be used directly to compute ht . The halting signal is a binary signal which is initialized to have a value of 0 and remains 0 for each et until QH > QC for K consecutive events . At that point , the value of ht transitions to 1 and continues to maintain that value if any additional events are processed .

Summary: This paper attempts to train a predictor of whether software is malware. Previous studies have emulated potential malware for a fixed number of executed instructions, which risks both false negatives (haven’t yet reached the dangerous payload) and false positives (malware signal may be lost amidst too many other operations). This paper proposes using deep reinforcement learning over a limited action space: continue executing a program or halt, combined with an “event classifier” which predicts whether individual parts of the program consist of malware. The inputs at each time step are one of 114 high level “events” which correspond to related API invocations (e.g. multiple functions for creating a file). One limitation seems to be that their dataset is limited only to events considered by a "production malware engine", so their evaluation is limited only to the benefit of early stopping (rather than continuing longer than the baseline malware engine). They evaluate a variety of recurrent neural networks for classifying malware and show that all significantly underperform the “production antimalware engine”. Integrating the event classifier within an adaptive execution control, trained by DQN, improves significantly over the RNN methods.