[Paper Review] An Ensemble of Pre-trained Transformer Models For Imbalanced Multiclass Malware Classification

Ferhat Demirkıran, Aykut Çayır+2 more|arXiv (Cornell University)|Dec 25, 2021

Advanced Malware Detection Techniques4 citations

TL;DR

This paper proposes a bagging-based Random Transformer Forest (RTF) ensemble model combining pre-trained BERT and CANINE transformer architectures to improve multiclass malware classification on highly imbalanced datasets using API call sequences. The RTF model achieves state-of-the-art F1-score of 0.6149 on the Catak dataset, outperforming individual models and prior methods by leveraging attention mechanisms and ensemble learning to capture complex sequence relationships in malware behavior.

ABSTRACT

Classification of malware families is crucial for a comprehensive understanding of how they can infect devices, computers, or systems. Thus, malware identification enables security researchers and incident responders to take precautions against malware and accelerate mitigation. API call sequences made by malware are widely utilized features by machine and deep learning models for malware classification as these sequences represent the behavior of malware. However, traditional machine and deep learning models remain incapable of capturing sequence relationships between API calls. On the other hand, the transformer-based models process sequences as a whole and learn relationships between API calls due to multi-head attention mechanisms and positional embeddings. Our experiments demonstrate that the transformer model with one transformer block layer surpassed the widely used base architecture, LSTM. Moreover, BERT or CANINE, pre-trained transformer models, outperformed in classifying highly imbalanced malware families according to evaluation metrics, F1-score, and AUC score. Furthermore, the proposed bagging-based random transformer forest (RTF), an ensemble of BERT or CANINE, has reached the state-of-the-art evaluation scores on three out of four datasets, particularly state-of-the-art F1-score of 0.6149 on one of the commonly used benchmark dataset.

Motivation & Objective

• To address the challenge of highly imbalanced multiclass malware classification using behavioral features from API call sequences.
• To evaluate the effectiveness of pre-trained transformer models like BERT and CANINE in capturing long-range dependencies in API call sequences for malware family identification.
• To develop an ensemble model that improves generalization and performance on rare malware families in imbalanced datasets.
• To compare inference and training efficiency of transformer-based models against traditional LSTM and single-layer Transformer models.
• To demonstrate the first application of CANINE in malware classification and validate its performance through ensemble learning.

Proposed method

• The method employs a bagging-based ensemble framework called Random Transformer Forest (RTF), where multiple instances of pre-trained BERT or CANINE models are trained on random subsets of the training data.
• Each RTF model uses tokenization-free input representation of API call sequences, preserving sequence order and enabling attention mechanisms to model dependencies across calls.
• The ensemble aggregates predictions via majority voting to improve robustness and reduce variance, especially on minority malware families.
• Pre-processing steps, including sequence length normalization and data augmentation, are applied to enhance performance on the Catak dataset with dynamic API calls.
• Evaluation is conducted using F1-score and AUC as primary metrics due to class imbalance, with training and inference times measured for efficiency analysis.
• The model leverages multi-head attention and positional embeddings inherent in BERT and CANINE to model long-range relationships in API call sequences more effectively than RNNs or single-layer Transformers.

Figure 1 : Transformer model architecture.

Experimental results

Research questions

• RQ1 RQ.1: What are the suitable classification metrics for imbalanced multiclass malware classification, and how do F1-score and AUC compare in evaluating model performance?
• RQ2 RQ.2: How do pre-trained transformer models like BERT and CANINE perform compared to traditional models (e.g., LSTM) and single-layer Transformers on imbalanced malware datasets?
• RQ3 RQ.3: Can an ensemble of pre-trained transformers, such as a Random Transformer Forest (RTF), achieve state-of-the-art performance on benchmark malware datasets?
• RQ4 RQ.4: What is the trade-off between inference speed and classification accuracy when using large pre-trained models like BERT and CANINE?
• RQ5 RQ.5: How do pre-processing techniques affect the performance of transformer models on dynamic API call sequences, particularly on the Catak dataset?

Key findings

• The RTF ensemble model achieves a state-of-the-art F1-score of 0.6149 on the Catak dataset, outperforming all baseline models and prior methods on this benchmark.
• BERT and CANINE pre-trained models significantly outperform a single-layer Transformer and LSTM models in terms of F1-score and AUC across all four datasets, especially on rare malware families.
• The CANINE model is successfully applied for the first time in malware classification and achieves state-of-the-art results on the VirusShare and VirusSample datasets, demonstrating its effectiveness for static API call sequences.
• Training time for BERT and CANINE is substantially higher than for LSTM and single-layer Transformers, with BERT taking 141.02 minutes on the Oliveira dataset, but this does not impact real-time response since models are pre-trained.
• Inference time for RTF is approximately 4.5 seconds per sample on average, which is acceptable for AV scanners, and remains within a reasonable range despite the use of large pre-trained models.
• The RTF model achieves state-of-the-art performance on three out of four datasets, confirming the effectiveness of ensemble learning with pre-trained transformers for imbalanced malware classification.

Figure 2 : The outputs of pre-processing steps.