Random CapsNet Forest Model for Imbalanced Malware Type Classification Task

Malware type classification is important as much as malware detection problem because system protection softwares are able to decide their strategies with respect to malware family types. Malware families have different behavior and affects on a computer system. For instance, each malware family uses different resources, files, ports and other components of operating systems. In addition to this, because of trends in technology, new malware types have occurred.

Malware type classification is the most common problem in cybersecurity domain, because strategies of protection systems vary with respect to malware family type. Malware type classification can be done via three different ways such as static, dynamic and image based nataraj2011comparative; abijah2019intelligent; ni2018malware. This paper focuses on image based malware family classification problem. In order to do that, two important datasets have been used and obtained results have been compared other models in the literature.

This paper proposes a new model named random capsule network forest (RCNF) based on boostrap aggregation (bagging) ensemble and capsule network (CapsNet) breiman1996bagging; sabour2017dynamic. The main assumption behind the proposed method is to reduce variance of different CapsNet models (as weak learners) using bagging. In this perspective, the main contributions of this paper can be listed as bellow:

• This paper introduces the first application of CapsNet on malware type classification.

• The first ensemble model of CapsNet is presented in the paper.

• The proposed model uses simple architecture engineering instead of complex convolutional neural network architectures and domain specific feature engineering techniques.

This paper also presents a detailed comparison of proposed model with other models in the literature.

The paper is organized as follows, section 2 presents a literature survey for CapsNet applications and previous malware analysis studies. Datasets used in the paper are described in section 3. Section 4 gives details of inspiring model and the proposed model. Results are investigated in section 5. The paper is concluded by section 6.

Convolutional neural networks (CNNs) are very popular and successful models for image based datasets. As expected in visual static malware analysis, there are many CNN architectures. CapsNet, as a new CNN structure, has been implemented in 2017 sabour2017dynamic. Especially in health domain jimenez2018capsule, CapsNet have many appliations. For instance Afshar et al. afshar2018brain uses CapsNet for brain tumor classification problem like classification of breast cancer histology images iesmantas2018convolutional. Mobiny et al mobiny2018fast creates a fast CapsNet architecture for lung cancer diagnosing. Another important application area of CapsNet is object segmentation. LaLonde et al. lalonde2018capsules uses CapsNet for object segmentation. The most brilliant idea in machine learning is generative adverserial networks (as known as GAN). GANs are based on CNN structures. CapsNet is very useful to make GANs better removing weakest point of CNNs jaiswal2018capsulegan. In other words, these studies show that CapsNet is a promising architecture against to standard CNN.

In the literature, although there are many applications of CapsNets, there is a missing and important area. This area is computer and information security. This gap can be seen easily in a pre-print version of a survey about CapsNets patrick2019capsule. In that manner, this paper aims to a robust malware prediction model based on CapsNet architecture for imbalanced malware dataset. This model is the first application of CapsNets in malware prediction domain.

The base CapsNet and the proposed RCNF models have been tested on two the most known malware datasets called Malimg and Microsoft Malware 2015 (BIG2015) respectively. This section describes these two datasets.

In this section, general capsule networks, base CapsNet architecture for Malimg and base CapsNet architecture for Microsoft Malware 2015 Dataset are described. CapsNet architectures are different for both Malimg and BIG2015 Dataset.

CapsNet and RCNF ensemble model are tested on two different datasets called Malimg and BIG2015. Malimg dataset has been divided into three parts, training, validation and test sets. Training set has 7004 samples, validation set has 1167 samples and test set has 1167 samples. BIG2015 has also been divided into three parts like Malimg dataset. In the experiments of CapsNet and RCNF ensemble model for BIG2015 dataset, training set has 8151 samples, validation set has 1359 samples and test set has 1358 samples.

The First experiment is to obtain performance results of single base CapsNet estimators for each dataset. The second experiment is about performance of RCNF model. Model evaluation has been done in terms of accuracy, F1 measure and number of parameters of deep neural nets. These performance metrics are defined as follows:

where true positive (TP) and false positive (FP) are the number of instances correct and wrong classified as positive, respectively. True negative (TN) and false negative (FN) are the number of instances correct and wrong classified as negative, respectively. Accuracy is the ratio of number of true predictions to all instances in the set as shown in equation (3). F1 is shown as the equation (4) in terms of true positives, false negatives and false positives.

Figure 7 shows confusion matrices for each test part of both datasets. Each confusion matrix (figure 6(a) and 6(b)) implies that a model containing single CapsNet incorrectly predicts rare malware families in both datasets.

Figure 8 is the confusion matrix of RCNF containing 5 base CapsNet models. This confusion matrix shows prediction accuracy of the model for each malware family type in Malimg test set. Class 8, 10, 20 and 21 have been predicted wrongly by the RCNF model. On the other hand, the model has been very successful at correctly predicting other malware types in the test set. This confusion matrix also shows that RCNF is successful at correctly predicting rare malware types in Malimg test set.

In the second experiment, RCNF is tested on BIG2015 dataset. Figure 8 shows the prediction results of RCNF containing 10 base CapsNet for BIG2015 dataset. Class 4 is the rarest malware type in the whole dataset. Training, validation and test sets are stratified, so the class distribution is preserved for each partition. This result shows that RCNF is able to predict the rarest malware type pretty well. Class 0, 1, 2 and 6 are predicted perfectly by RCNF. If the performance of RCNF is compared with the performance of a single CapsNet model, it is easily seen that RCNF is better than a single CapsNet at predicting rare malware families for imbalanced datasets.

Table 3 shows the test performance of the proposed models and others for Malimg dataset. Yue yue2017imbalanced uses a weighted loss function to handle imbalance class distribution problem in Malimg dataset and also uses transfer learning yosinski2014transferable method to classify malware family types. Due to using transfer learning, the architecture has 20M parameters and the model is so huge. Cui et al. cui2018detection use classical machine learning methods such as K-Nearest Neighbor and support vector machines. They have trained these algorithms using GIST and GLCM features which are feature engineering methods for images and have applied resampling to the dataset in order to solve imbalance dataset problem. RCNF does not use weighted loss function or any sampling method to overcome imbalanced dataset problem. Our results are higher than these two methods and results also show that CapsNet and RCNF do not require any method for extra feature engineering in Malimg dataset. A single CapsNet architecture for Malimg dataset has 90592 trainable parameters and RCNF has 452960 trainable paramters, so our proposed methods are reasonably smaller than Yue’s method.

Table 4 compares the test performance of the proposed models and others for BIG2015 dataset. Chen chen2018deep and Khan et al. khan2018analysis use transfer learning architectures for the dataset. Test results of our proposed methods are better than those two models, but our results are very close to Chen chen2018deep in terms of accuracy. Our proposed models are better than Gibert et al. gibert2018end in terms of accuracy, but F1 score of RCNF is very close to this model. Our proposed models are obviously better than models of Cao et al. cao2018efficient, Zhao et al. zhao2019maldeep, Kim et al. kim2018classifying and Kim et al. kim2018detecting. Jung et al. jung2018malware propose a reasonably smaller model than our models in terms of number of parameters, but our model has higher accuracy score than this model. Abijah et al. abijah2019intelligent, Safa et al. safa2019benchmarking, Kebede et al. kebede2017classification and Yan et al. yan2018detecting propose deep learning models whose accuracy scores are close to our proposed models. Data resampling and weighted loss function for imbalanced dataset, transfer learning and data augmentation have never been used to train CapsNet and RCNF models for both BIG2015 and Malimg datasets. In other words, these results are basic performances of these models.

For those tables, last studies using Mailmg and BIG2015 datasets are chosen. In order to compare in a fair way, these models are drawn from visual static malware analysis studies.

This paper introduces a first application of CapsNet on imbalanced malware family type classification task. Moreover, the first ensemble model of CapsNet called RCNF is introduced in this paper. Proposed models do not require any complex feature engineering methods or architecture for deep networks. In order to show that, we used two different malware family type datasets such as Malimg and BIG2015. These datasets are used image based visual static malware classification. Our proposed models can use these datasets directly using raw pixel values.

Datasets in the paper are highly imbalanced in terms of class distribution. CapsNet and RCNF do not use oversampling, undersampling and weighted loss function during training phase. Results show that CapsNet and RCNF are the best models least suffering from imbalanced class distribution among others.

Experiment results show that a single CapsNet model has good performance for both BIG2015 and Malimg datasets. However, we have assumed an ensemble model of CapsNet can help to increase generalization performance and RCNF has better generalization perofrmance results than a single CapsNet model as expected.

As a future work, we are planning to develop a hybrid architecture for malware classification. This hybrid method will be based on CapsNet architeture. This paper shows that CapsNet are useful for imbalanced malware dataset in visual static malware analysis problem domain.