Malware Classification using Machine Learning | by Arpan Das | Towards Data Science

If you love to explore large and challenging data sets, then probably you should give Microsoft Malware Classification a try. Before diving deep in to the problem let’s take few points on what can you expect to learn from this:
How to handle large scale data? Total train data set consist of 200 GB data out of which 50 GB of data is .bytes files and 150 GB of data is .asm files.What feature engineering should you do? If till now you have been working only on text and image data, this will surely boost your intuitions on feature engineering.How to approach a machine learning problem without domain knowledge? In the words of the 1st place winners of this challenge: “Cross Validation is more trust worthy than domain knowledge”.
How to handle large scale data? Total train data set consist of 200 GB data out of which 50 GB of data is .bytes files and 150 GB of data is .asm files.
What feature engineering should you do? If till now you have been working only on text and image data, this will surely boost your intuitions on feature engineering.
How to approach a machine learning problem without domain knowledge? In the words of the 1st place winners of this challenge: “Cross Validation is more trust worthy than domain knowledge”.
Business Problem
In the past few years, the malware industry has grown very rapidly that, the syndicates invest heavily in technologies to evade traditional protection, forcing the anti-malware groups/communities to build more robust software to detect and terminate these attacks. The major part of protecting a computer system from a malware attack is to identify whether a given piece of file/software is a malware.
Machine Learning Problem, KPI and constraints
We can map the business problem to a multi-class classification problem, where we need to predict the class for each given byte files among nine categories (Ramnit, Lollipop, Kelihos_ver3, Vundo, Simda,Tracur, Kelihos_ver1, Obfuscator.ACY, Gatak).
KPI: multi-class log loss, confusion matrix
Constrains: We need to provide the class probability, wrongly classified class labels should be penalized(that’s why log loss has been chosen as KPI) and there should some latency bound.
Handling the Large Size
Keeping the scale of data in mind, It is advised to go for cloud platforms like GCP, AWS or IBM Cloud. We have completed this case study in GCP with Linux based 1vCPU box with 30 GB RAM and 500 GB HDD. Follow this link to know how you can download the data directly in GCP from jupyter notebook itself with the help of CurlWidget. After downloading the data the next challenge will be to extract these files using python as the python support for 7z files are limited. Go for a package called ‘p7zip’ which provides a Linux command line tool to extract 7z files.
# command to install the package in Linuxsudo apt-get install p7zip-full# Sample Command to extract the data from terminal7za x yourfile.tar.7z
The total data set is consists of 10,868 .bytes files and 10,868 asm files total 21,736 files. Byte files are a combination of 256 hex numbers (decimal value ranging between 0 to 255) and a special character(??). Asm files are assembly files obtained using smart compilers.The data set is unbalanced.
The total data set is consists of 10,868 .bytes files and 10,868 asm files total 21,736 files. Byte files are a combination of 256 hex numbers (decimal value ranging between 0 to 255) and a special character(??). Asm files are assembly files obtained using smart compilers.
The data set is unbalanced.
The good thing of this case study is, it demands a lot of patience and experiments to obtain the useful features. We have tried out the following features based on our intuitions and previous works on this field
n-gram features of byte files: We have tried uni-gram, bi-gram and tri-gram features. Though the performance with these features is good but the main challenge lies in handling their large dimensionality. The number of features with uni-gram is 256, with bi-gram it is 65531 and leaving the number of tri-gram features for your exercise. Even with the 30 GB RAM it is not possible to load the total data with 65531 features, so we went for SVD in chunk. First we tried out elbow method with random samples and 3000 features to check the number of components needed to preserve the variance and we found out 1000 as a good number.
n-gram features of byte files: We have tried uni-gram, bi-gram and tri-gram features. Though the performance with these features is good but the main challenge lies in handling their large dimensionality. The number of features with uni-gram is 256, with bi-gram it is 65531 and leaving the number of tri-gram features for your exercise. Even with the 30 GB RAM it is not possible to load the total data with 65531 features, so we went for SVD in chunk. First we tried out elbow method with random samples and 3000 features to check the number of components needed to preserve the variance and we found out 1000 as a good number.
2. Size of byte and asm files: Size of bytes and asm files turned out be be good features. Let’s take a look at the box plots of file sizes
3. Size of byte array: As we mentioned earlier,the byte files are consists of hexadecimal digits. For each files we have taken an array and added each elements in to the array and used the length of these arrays as a feature.
4. Segment count and n-grams of asm files: asm files consists of various segments like prefixes(ex. HEADER,TEXT), opcodes(ex. jmp, hlt), keywords(ex. .dll, std::), registers(ex. edx,epi) etc. We have taken the count of segments as a feature. We have also considered uni-gram, bi-gram and tri-gram of asm files as feature.
5. Pixel intensity of byte and asm features: This is probably the most innovative features of the case this case study. We have borrowed the concept from the 1st place winner solution.As we mentioned earlier the value of each elements in byte files is between 0 to 255. We can easily convert these byte files along with the asm files in to gray scale images. It is observed that the first 800 pixel densities of asm files have significant impact over the final result.
Let’s take a look how a random asm file looks after the conversion
By now you must have observed the dimension of the features and get a sense of the dimensionality of the final features for training. To cope up with the problem we have used random forest to select the important features. Number of important features is a hyper-parameter. We have tried various numbers based of our experiences and obviously there are scope of improvements. Let’s look at a sample code on this and how we have selected best tri-gram asm features. We have also tried t-SNE for multivariate analysis.
We have tried various linear machine learning models like KNN, multi-class logistic regression and ensembles like random forest, XGBoost and LightGBM with various combination of features. We have trained the models with single features, and a combination of byte features and asm features. We have also tried out various CNN architectures over the images obtained from byte and asm files but the obtained result was poorer compared to fine tuned machine learning models. As we did not have any domain knowledge on malware industry we have completely relied on cross validation. As the number of features and files were high, we have chosen RandomizedCV for cross validation and there are scope of improvements. Another question may come to your mind that why we have used both XGBoost and LightGBM? Actually we have started with XGBoost but found out that the training latency of LightGBM is much lower than XGBoost along with similar or better performance. We started with a random model then added the complexity of the model gradually after keeping an eye on the performance.
Finally a use case does not worth anything if it does not provide a satisfactory result. Here let’s list down only the result we obtained after calibrating our final LightGBM model.