Self-MCKD: Enhancing the Effectiveness and Efficiency of Knowledge Transfer in Malware Classification

1. Introduction

With the rapid development of the Internet, computers have become an integral part of people’s daily lives. However, this growth has also led to a gradual increase in the number of malware, which has become one of the major threats to computer security. Some hackers exploit vulnerabilities in computers to steal data or obtain private information without user’s permission. As a representative example, the WannaCry malware caused billions of dollars in damages to millions of computer devices[29].

To address such threats in a timely and effective manner, many countermeasures have been proposed. Traditionally, many studies have focused on malware classification methods using signature-based approach, relying on statistical methods. However, as malware continues to evolve with new variants, these methods have become less effective, leading to a decrease in classification performance.

In response, some studies have shifted towards AI-based methods, which have shown significant promise in improving malware classification performance[30]. Until recently, AI-based malware classification methods have focused on enhancing performance by increasing model size and depth. However, these methods lead to a substantial rise in computational complexity and the number of parameters, significantly raising computational cost during the training process. Moreover, the maintenance cost of AI-based malware classification methods also increases, as frequent retraining and transfer learning are required to keep up with evolving malware variants. These challenges not only increase the training time of AI-based malware classification methods but also increase the inference time. For example, such challenges present practical limitations when applying AI-based malware classification methods to IoT networks in CPS industrial, which is a resource-constrained environment[32].

To address these challenges, there is a growing need for research on lightweight techniques for AI-based malware classification methods. The existing lightweight techniques for AI-based malware classification methods are grouped into two categories of architectures, i.e., model compaction techniques and model optimization techniques. Model compaction techniques simplify the model structure by replacing layers with more efficient ones for lightweight AI-based malware classification methods. However, these methods suffer from computational complexity and long inference time. On the other hand, model optimization techniques reduce computational complexity and inference time by optimizing computational complexity of AI-based malware classification methods. In particular, knowledge distillation, which is a representative model optimization technique, has been actively studied as an efficient lightweight technique for AI-based malware classification methods.

Knowledge distillation allows small and shallow AI-based malware classification methods, as the student model, to mimic large and deep AI-based malware classification methods, as the teacher model, through knowledge transfer. Unfortunately, previous studies on knowledge distillation techniques for AI-based malware classification methods still face issues of computational complexity and inefficient knowledge transfer. Also, they require large and deep AI-based malware classification methods for the teacher model, which leads to high computational cost during the training process.

To solve these challenges, we propose an efficient knowledge distillation technique for AI-based malware classification methods, called Self-MCKD. As shown in Figure 1, different from traditional knowledge distillation techniques for AI-based malware classification methods, which transfer combined output logits composed of target class and non-target classes, Self-MCKD transfers output logits that are separated into them. With the separation of output logits, Self-MCKD enables efficient knowledge transfer by assigning weighted importance to the target class and non-target classes. Also, Self-MCKD only uses small and shallow methods as both the teacher and student models to overcome the need for using large and deep methods as the teacher model in knowledge distillation techniques for AI-based malware classification methods. Since Self-MCKD enables efficient knowledge transfer, it not only enhances the performance of small and shallow AI-based malware classification methods but also enables them to surpass large and deep methods. Also, it accelerates convergence speed and reduces the training time of small and shallow AI-based malware classification methods.

Main contributions of this paper can be summarized as follows:

• We propose Self-MCKD, a novel knowledge distillation technique for AI-based malware classification methods. Self-MCKD separates the output logits into the target class and non-target classes, assigning a weighted importance factor that maximizes knowledge transfer effectiveness.
• To enhance knowledge transfer efficiency, we introduce small and shallow AI-based malware classification methods as both the teacher and student models in Self-MCKD.
• From the experimental results using representative benchmark datasets, we show that Self-MCKD significantly improves the performance of a small and shallow baseline model, even surpassing that of a large and deep baseline model. Also, we show that Self-MCKD outperforms previous knowledge distillation techniques in both effectiveness and efficiency. Specifically, with the same FLOPs, MACs, and parameters, Self-MCKD achieves improvements in accuracy, recall, precision, and F1-score by 0.2% to 1.4%, while reducing training time by 20 seconds.

2. Related Works

In this section, we first overview previous AI-based malware classification methods and lightweight techniques for AI-based malware classification methods, then we introduce the most effective knowledge distillation-based technique for AI-based malware classification methods.

2.1. AI-Based Malware Classification Methods

In the field of malware classification, many previous methods have utilized executable files represented by opcodes or assembly codes [18]. However, these methods faced several limitations, such as the difficulty of analyzing obfuscated code, the time-consuming nature of static analysis, and the high resource demands of dynamic analysis[19]. To address these limitations, several studies have applied AI-based methods, specifically machine learning (ML) and deep learning (DL), after converting executable files into gray-scale images [17]. Nataraj et al. were the first to apply the K-Nearest Neighbors (KNN) algorithm, an ML technique, to executable files converted into gray-scale images for malware classification [1]. To achieve improved performance over KNN algorithm, Hardy et al. proposed an SVM-based malware classification method [20]. Also, Drew et al. leveraged modern gene classification tools [21], and Ahmadi et al. leveraged XGBoost techniques [22], to progressively improve the performance of malware classification with more advanced ML models. However, as the variety of malware variants increased, the complexity of malware analysis also increased. Moreover, this revealed certain limitations of ML-based malware classification methods, especially their inability to automatically extract and learn nuanced features of malware variants.

To address limitations of ML-based malware classification methods, DL-based malware classification methods, which can automatically extract nuanced features, have emerged. In particular, convolutional neural network (CNN)-based malware classification methods have gained prominence due to their superior image analysis capabilities when applied to executable files converted into gray-scale images. Daniel et al. [2] applied a shallow, simple CNN model for malware classification using gray-scale images of malware. They also showed that DL-based malware classification methods outperform previous ML-based malware classification methods. Following this study, Mahmoud et al. hypothesized that deeper CNN architectures could further improve the performance of malware classification [3]. Based on this assumption, they enhanced the performance of DL-based malware classification methods using the VGG16 model. To achieve even better classification performance, Pratikkumar et al. proposed a migration learning framework for malware classification, utilizing more advanced models such as ResNet52 and VGG19 models [4].

2.2. Lightweight Techniques for DL-Based Malware Classification Methods

As DL-based malware classification methods have become increasingly complex and deeper, the associated hardware burden has also increased. Also, since this has led to longer training and inference times, several recent studies have focused on developing lightweight DL-based malware classification methods. In this section, we summarize and discuss key studies related to the lightweight techniques for DL-based malware classification methods. Early lightweight techniques for DL-based malware classification methods focused on simplifying the model structure, also known as model compaction techniques. Specifically, Fathurrahman et al. [23] proposed a lightweight malware classification method by replacing the deep structures of the CNN model with residual layers and global average pooling. This approach outperformed the lightweight CNN-based malware classification method proposed by Su et al. [24], offering both reduced complexity and improved performance. However, model compaction techniques, which simply replace layer structures, still suffer from computational complexity and long inference time [5,6,10].

To address these challenges, efficient lightweight techniques, also known as model optimization techniques, have emerged. These techniques, such as weight pruning [7], model quantization [8], and knowledge distillation [9], reduced computational complexity and inference time in DL-based malware classification methods. In particular, knowledge distillation has been actively studied as an efficient lightweight technique for DL-based malware classification methods. Zhi et al. proposed a knowledge distillation-based technique for recurrent neural network (RNN)-based malware classification methods [11]. Specifically, they enhanced the performance of the RNN-based malware classification method (student model), using a transformer-based malware classification method, as the teacher model. Xia et al. [12] also proposed a lightweight technique for multi-layer perceptron (MLP)-based malware classification methods using knowledge distillation. Different from Zhi et al.’s technique, which uses different families of models for the teacher and student models, Xia et al.’s technique uses the same family of MLP models for both. By using the same MLP family of models, they improved the performance of the student MLP-based malware classification method to a level comparable to the teacher MLP-based model. Building on this foundation, recent studies have further combined effective learning methods with knowledge distillation, such as the two-dimensional Fourier transform and federated learning[33,34].

Previous knowledge distillation-based techniques for DL-based malware classification methods showed that small and shallow student models could achieve performance comparable to large and deep teacher models. However, despite the promising performance of these techniques, they still suffered from computational complexity and inefficient knowledge transfer. Moreover, many of these techniques relied on large and deep DL-based methods as the teacher model, which incurred high computational cost during training process. These problems motivate us to seek a new knowledge distillation-based technique for DL-based malware classification methods, which enables efficient knowledge transfer with minimal computational cost and complexity.

3. Preliminary

3.1. Knowledge Distillation

Knowledge distillation was first introduced by Bucila et al. for model compression [13] and gained popularity after Hinton et al. generalized it [9]. Knowledge distillation is a model optimization technique that allows the small and shallow student model to mimic the performance of a large and deep teacher model. Specifically, the student model S is trained to mimic the performance of the teacher model T by leveraging the teacher’s “Dark Knowledge”. Here, “Dark Knowledge” refers to the softened class probabilities predicted by the teacher model, which are distilled into the student model. By utilizing this “Dark Knowledge”, the student model improves its performance by learning not only from the original training data but also by incorporating the nuanced knowledge provided by the teacher’s predictions. The process of distilling this “Dark Knowledge”, which is called knowledge distillation process($KD$), can be expressed as follows:

$$KD=\sum _{x_i\in X}\mathrm{KL}\left(softmax\left({\displaystyle \frac{f_T\left(x_i\right)}{\tau }}\right),softmax\left({\displaystyle \frac{f_S\left(x_i\right)}{\tau }}\right)\right),$$

(1)

where X represents the training dataset, $x_i$ represents individual data points, $f_T\left(x_i\right)$ and $f_S\left(x_i\right)$ represent the class probabilities predicted by the teacher and student models, respectively. Additionally, $\tau$ is the temperature parameter used to soften the class probabilities. The Kullback-Leibler(KL) divergence is used as the loss function to measure the difference between the softened outputs of the teacher and student models. By minimizing KL divergence, the student model is trained to mimic the behavior of the teacher model, effectively absorbing the teacher’s “Dark Knowledge” and enhancing its own classification performance. Knowledge distillation can be implemented in two main ways: logit-based and feature-based approaches. In Figure 2, we show the differences between these two knowledge distillation approaches.

As shown in Figure 2 (a), the feature-based approach, which was first proposed by Romero et al. [31], distills knowledge to the student model from the teacher model’s hint layers. This approach focuses on training the student to mimic the output of the intermediate layers. Specifically, the feature-based approach transfers the rich knowledge of the teacher to the student, enabling the student to achieve performance similar to that of the teacher. However, since the knowledge distilled from the hint layers has high computational complexity, the feature-based approach requires significant training time for the student model. Also, the feature-based approach is inefficient for lightweight techniques of DL-based malware classification methods because DL-based malware classification methods require frequent training and transfer learning.

On the other hand, the logit-based approach distills knowledge to the student model from the teacher model’s output layer, which produces softened class probabilities as shown in Figure 2 (b). Different from the feature-based approach, this approach focuses on training the student to mimic the output logits of the teacher. Also, since the logit-based approach has low computational complexity, it can be a solution to the long training time problem of the feature-based approach. For example, traditional knowledge distillation, which is a representative logit-based approach, enables the student model to reduce the training time compared to the feature-based approach. Unfortunately, the existing logit-based approaches suffer from inefficient knowledge transfer, which limits the performance of the student model. Therefore, we propose a new logit-based approach, which can transfer knowledge more efficiently to the student model compared to the existing logit-based approaches.

4. Proposed Method

In this section, we overview the operation of Self-MCKD in detail. First, we introduce the difference between traditional knowledge distillation and Self-MCKD, then we reformulate traditional knowledge distillation for Self-MCKD. Next, we describe the network architecture in Self-MCKD used for both the teacher and student models.

4.1. Revisiting Traditional Knowledge Distillation

Traditional knowledge distillation transfers combined knowledge of target and non-target classes to the student model. However, this approach suffers from inefficient knowledge transfer, as combined knowledge is often emphasized in an unbalanced manner, reducing its effectiveness [14]. Also, traditional knowledge distillation tends to use high-performing large and deep teacher models, which require high computational cost during the training process. In the field of malware classification, effective knowledge transfer is essential for better generalization against evolving malware variants. However, traditional knowledge distillation has limitations in this domain due to its unbalanced knowledge transfer. Moreover, its reliance on large and deep teacher models incurs high computational costs, making frequent retraining challenging and limiting adaptability to newly emerging malware threats.

To address these problems in the field of malware classification, we propose a novel knowledge distillation technique called Self-MCKD. Different from traditional knowledge distillation, Self-MCKD transfers the knowledge of target and non-target classes independently. Also, Self-MCKD does not rely on a large and deep teacher model. Instead, it utilizes a pre-trained small and shallow student model to serve as the teacher, transferring knowledge to the untrained student model during the training process.

4.2. Notations for Self-MCKD

In this section, we define two types of output probabilities—target malware class probabilities and non-target malware class probabilities—to reformulate knowledge transfer process of traditional knowledge distillation. When a malware sample of the i-th class is given, the output probabilities of the malware classification model are defined as $\mathbf{p}=[p_1,p_2,p_3,\dots ,p_i,\dots ,p_M]\in {\mathbb{R}}^{1\times M}$, where $p_i$ represents the output probabilities of the i-th malware class, and M is the number of malware families. Here, $p_i$ is calculated as follows:

$$p_i={\displaystyle \frac{exp\left(z_i\right)}{{\sum }_{j=1}^{M}exp\left(z_j\right)}},$$

(2)

where $z_i$ is the output logit for the i-th malware class in the malware classification model.

Next, to separate the output probabilities related to the target malware class and non-target malware classes, we define binary probabilities as $\mathbf{bm}=[p_{tm},p_{ntm}]\in {\mathbb{R}}^{1\times 2}$. Here, $p_{tm}$ and $p_{ntm}$ can be expressed into:

$$p_{tm}={\displaystyle \frac{exp\left(z_{tm}\right)}{{\sum }_{j=1}^{M}exp\left(z_j\right)}},p_{ntm}={\displaystyle \frac{{\sum }_{k=1,k\ne tm}^{M}exp\left(z_k\right)}{{\sum }_{j=1}^{M}exp\left(z_j\right)}},$$

(3)

where $p_{tm}$ represents the probabilities of the target malware class, and $p_{ntm}$ represents the probabilities of all the other non-target malware classes. Meanwhile, we defined $\widehat{\mathbf{p}}=[{\widehat{p}}_1,\dots ,{\widehat{p}}_{i-1},{\widehat{p}}_{i+1},\dots ,{\widehat{p}}_{M-1}]\in {\mathbb{R}}^{1\times (M-1)}$ that means the output probabilities of non-target malware classes, i.e., excluding the i-th malware class. Here, ${\widehat{p}}_i$ is calculated as follows:

$${\widehat{p}}_i={\displaystyle \frac{exp\left(z_i\right)}{{\sum }_{j=1,j\ne tm}^{M}exp\left(z_j\right)}}.$$

(4)

4.3. Reformulating Knowledge Distillation for Self-MCKD

Based on above notations, we reformulate the knowledge distillation process($KD$) from Equation (1) as follow:

$$\begin{array}{cc}\hfill KD& =KL\left(p^{\mathcal{T}}\right|\left|p^{\mathcal{S}}\right)\hfill \\ & =p_{tm}^{\mathcal{T}}log\left({\displaystyle \frac{p_{tm}^{\mathcal{T}}}{p_{tm}^{\mathcal{S}}}}\right)+\sum _{i=1,i\ne tm}^M{p}_i^{\mathcal{T}}log\left({\displaystyle \frac{p_i^{\mathcal{T}}}{p_i^{\mathcal{S}}}}\right),\hfill \end{array}$$

(5)

where $p^{\mathcal{T}}$ represents the output probabilities of the teacher model and $p^{\mathcal{S}}$ represents the output probabilities of the student model.

Different from the loss function of traditional knowledge distillation, which aims to simply reduce the overall difference in probabilities between the teacher and student models, Self-MCKD aims to reduce the difference for both the target and non-target malware classes separately between the teacher and student models. Also, since Self-MCKD replaces the large and deep teacher model with a pre-trained small and shallow student model, $p^{\mathcal{T}}$ can be replaced with $p^{t\mathcal{S}}$. Finally, considering that $\widehat{p}$ equals $p_i/p_{ntm}$, we can reformulate Equation (5) as follows:

$$\begin{array}{cc}\hfill KD& =p_{tm}^{t\mathcal{S}}log\left({\displaystyle \frac{p_{tm}^{t\mathcal{S}}}{p_{tm}^{\mathcal{S}}}}\right)+\sum _{i=1,i\ne tm}^M{p}_{ntm}^{t\mathcal{S}}{\widehat{p}}_i^{t\mathcal{S}}\left(log\left({\displaystyle \frac{{\widehat{p}}_i^{t\mathcal{S}}}{{\widehat{p}}_i^{\mathcal{S}}}}\right)+log\left({\displaystyle \frac{p_{ntm}^{t\mathcal{S}}}{p_{ntm}^{\mathcal{S}}}}\right)\right)\hfill \\ \hfill & =p_{tm}^{t\mathcal{S}}log\left({\displaystyle \frac{p_{tm}^{t\mathcal{S}}}{p_{tm}^{\mathcal{S}}}}\right)+p_{ntm}^{t\mathcal{S}}log\left({\displaystyle \frac{p_{ntm}^{t\mathcal{S}}}{p_{ntm}^{\mathcal{S}}}}\right)+p_{ntm}^{t\mathcal{S}}\sum _{i=1,i\ne tm}^M{\widehat{p}}_i^{t\mathcal{S}}log\left({\displaystyle \frac{{\widehat{p}}_i^{t\mathcal{S}}}{{\widehat{p}}_i^{\mathcal{S}}}}\right)\hfill \\ \hfill & =KL({\mathbf{bm}}^{t\mathcal{S}}\left|\right|{\mathbf{bm}}^{\mathcal{S}})+p_{ntm}^{t\mathcal{S}}KL({\widehat{\mathbf{p}}}^{t\mathcal{S}}\left|\right|{\widehat{\mathbf{p}}}^{\mathcal{S}}).\hfill \end{array}$$

(6)

After reformulating the Equation (6), it can be written as follows:

$$KD=KL({\mathbf{bm}}^{t\mathcal{S}}\left|\right|{\mathbf{bm}}^S)+p_{ntm}^{t\mathcal{S}}KL({\widehat{\mathbf{p}}}^{t\mathcal{S}}\left|\right|{\widehat{\mathbf{p}}}^S).$$

(7)

As shown in Equation (7), the loss function of traditional knowledge distillation can be reformulated as a weighted sum of two terms. Here, $KL\left({\mathbf{bm}}^{t\mathcal{S}}\right|\left|{\mathbf{bm}}^S\right)$ represents the knowledge of the target malware class between the teacher and student models, which we refer to as (Target Malware Class Knowledge Distillation) TMKD. On the other hand, $KL\left({\widehat{\mathbf{p}}}^{t\mathcal{S}}\right|\left|{\widehat{\mathbf{p}}}^S\right)$ represents the knowledge of the non-target malware classes between the teacher and student models, which we refer to as (Non-Target Malware Class Knowledge Distillation) NMKD. These separated knowledges of Self-MCKD are shown in Figure 3. According to Zhao et al.’s work [14], the $p_{ntm}^{t\mathcal{S}}$ factor, derived from traditional knowledge distillation, enforces an emphasis on non-target malware classes. This enforced emphasis prevents from autonomously adjusting the importance of target and non-target malware classes. As a result, it restricts the effectiveness of knowledge transfer to student models. To address the restrictive effect of the $p_{ntm}^{t\mathcal{S}}$ factor, we introduce a factor $\beta$. Also, to proper emphasis on TMKD, we introduce a factor $\alpha$. Consequently, these factors enable Self-MCKD to emphasize TMKD and NMKD, respectively. By applying these factors and making the necessary substitutions to Equation (7), the loss function of Self-MCKD can be defined as Equation (8) and calculated using Algorithm 1.

$$\mathbf{Self}-\mathbf{MCKD}=\alpha \mathbf{TMKD}+\beta \mathbf{NMKD}.$$

(8)

5. Evaluation Results

To demonstrate the effectiveness and efficiency of the proposed Self-MCKD in the field of malware classification, we evaluated its performance across various datasets. Specifically, we aimed to evaluate the lightweighting effect of Self-MCKD when applied to representative DL-based malware classification methods in different dataset environments. To this end, we tried to answer the following questions:

• How do factors $\alpha$ of TMKD and $\beta$ of NMKD influence on the performance of Self-MCKD?
• Does Self-MCKD show good performance on small and shallow malware classification model under various malware datasets?
• Does Self-MCKD show the better performance than the other knowledge distillation method?

5.1. Experimental Environments

In this section, we describe experimental environments including dataset, model configuration, evaluation metric, and so on.

Dataset: We conducted the experiments using the Malimg dataset(Malimg) [28] and the Microsoft Malware Classification Challenge Dataset(MMCC) [27], both of which are representative datasets for malware classification. The Malimg, originally compiled by Nataraj et al. [1], consists of 9,339 malware samples represented as grayscale images, corresponding to 25 malware families. These families include “Adialer”, “Allaple”, “Lolyda”, “C2LOP”, and “Swizzor”, and several others. The MMCC dataset, provided by Microsoft Malware Classification Challenge, consists of 21,741 malware samples categorized into 9 malware families, namely “Ramni”, “Lollipop”, “Kelihos_ver3”, “Vundo”, “Simda”, “Tracur”, “Kelihos_ver1”, “Obfuscator.ACY”, and “Gatak”. As with Malimg, the distribution of malware samples across classes in the MMCC dataset is not uniform, with some families having significantly more samples than others. In our experiments, we randomly select 90% of malware samples from each family for training and used the remaining 10% for testing.

Model Configuration: To demonstrate the effectiveness and efficiency of the proposed Self-MCKD, we selected representative DL-based malware classification methods, such as ResNet and VGGNet. In previous malware classification studies [3,35], ResNet and VGGNet achieved promising classification performance. Thus, we used ResNet110 and VGG16 models as large and deep baseline models for malware classification. Also, to apply small and shallow baseline models in the proposed Self-MCKD, we selected much lighter homogeneous models, such as ResNet32 and VGG8 models. For the implementation of these models, we followed the original ResNet and VGGNet architecture [16,35] without modifications. Also, we set these models’ output layer to 9 or 25, depending on the number of malware families in the dataset, and used the stochastic gradient descent(SGD) optimization function to train them. Lastly, we set hyperparameters, such as batch size, epoch, learning rate, momentum, and weight decay to 8, 25, 0.001, 0.9, and 0.0005, respectively.

Evaluation Metric: To evaluate the performance of Self-MCKD, we measured accuracy, recall, precision, and F1-score under different malware datasets. These metrics are derived from the confusion matrix. In the confusion matrix, True Positive (TP) represents correctly predicted samples of the target malware class, while True Negative (TN) represents correctly predicted samples of the non-target malware class. False Positive (FP) represents non-target malware class samples that are incorrectly predicted as belonging to the target malware class. On the other hand, False Negative (FN) represents target malware class samples incorrectly predicted as belonging to the non-target malware class. Based on the confusion matrix, accuracy, recall, precision, and F1-score are calculated as follows:

$$Accuracy={\displaystyle \frac{TP+TN}{FP+TP+FN+TN}}.$$

(9)

$$Recall={\displaystyle \frac{TP}{TP+FN}}.$$

(10)

$$Precision={\displaystyle \frac{TP}{TP+FP}}.$$

(11)

$$F1-score={\displaystyle \frac{2·Precision·Recall}{Precision+Recall}}.$$

(12)

Also, to consider the computational efficiency and complexity of malware classification baseline models, we measured floating point operations (FLOPs), multiply-accumulate operations (MACs), and the number of model parameters. Here, FLOPs refer to the total number of floating-point operations required during a forward pass and MACs represent the number of multiply-accumulate operations. such metrics reflect the computational cost of the convolutional and fully connected layers. In knowledge distillation, both the student and teacher models contribute to the FLOPs and MACs during the inference process. However, the number of parameters, which represents the total trainable parameters of malware classification baseline models, is measured only for the student model during the training process. We conducted experiments five times and reported the average results to ensure a reliable evaluation based on these metrics.

Implementation Environment: We implemented malware classification baseline models, knowledge distillation methods using PyTorch version 2.2.2 and Python version 3.11.9. For efficient experiments, the performance was measured on Windows OS with PyTorch-CUDA version 12.1, Intel Core i9-12900KF processor, NVIDIA RTX 3060 GPU, and 64 GB of memory.

5.2. Experimental Analysis

5.2.1. How Do Factors $\alpha$ of TMKD and $\beta$ of NMKD Influence on the Performance of Self-MCKD?

To evaluate the impact of the factors $\alpha$ and $\beta$ on the performance of Self-MCKD, we measured the accuracy of the malware classification model while adjusting these factors. Simultaneously, we aimed to verify that these factors provide better performance compared to the $p_{ntm}^{t\mathcal{S}}$ factor derived from the traditional knowledge distillation reformulation. To determine the appropriate values for $\alpha$ and $\beta$, we conducted experiments by setting their sum to 1 to ensure stability during the training process and prevent degradation of knowledge transfer effectiveness. Also, we varied the $\alpha$ and $\beta$ values from (0.1, 0.9) to (0.9, 0.1).

In Table 1, we show the experimental results of the VGG8 model applied to Self-MCKD using MMCC and Malimg datasets. According to Table 1, it is observed that too much emphasis on NMKD by increasing the factor $\beta$ reduced the effectiveness of knowledge transfer. For example, in the experimental results on both datasets, while the VGG8 model with traditional knowledge distillation achieved an accuracy of 98.44% and 99.36%, the VGG8 model with Self-MCKD achieved an accuracy of 98.07% and 99.25%, when we set $\alpha$ and $\beta$ at 0.2 and 0.8, respectively. Also, we observed that too much emphasis on TMKD by increasing the factor $\alpha$ also negatively affected knowledge transfer. For example, in the experimental results on MMCC dataset, the VGG8 model with Self-MCKD, where $\alpha$ and $\beta$ were set to 0.9 and 0.1, achieved an accuracy of 98.25%, which is lower than an accuracy of 98.44% achieved by the VGG8 model with traditional knowledge distillation.

On the other hand, when balancing between TMKD and NMKD, we observed that Self-MCKD transferred knowledge more effectively than traditional knowledge distillation. Specifically, when we set both $\alpha$ and $\beta$ to 0.5, the VGG8 model with Self-MCKD achieved accuracy of 98.52% and 99.46%, respectively, which are higher than that of the VGG8 model with traditional knowledge distillation. Also, it is observed that a proper balance between TMKD and NMKD maximizes the effectiveness of knowledge transfer. For example, the VGG8 model with Self-MCKD, with $\alpha$ and $\beta$ were set to 0.7 and 0.3, achieved an accuracy of 98.62% and 99.57%, respectively, which outperformed the VGG8 model with traditional knowledge distillation.

5.2.2. Does Self-MCKD Shows Good Performance on Small and Shallow Malware Classification Model Under Various Malware Datasets?

To show the effectiveness and efficiency of Self-MCKD on small and shallow malware classification models under various datasets, we measured the computational complexity and accuracy of ResNet32 and VGG8 models, both with and without Self-MCKD using MMCC and Malimg datasets. Simultaneously, we aimed to show that the ResNet32 and VGG8 models with Self-MCKD outperforms larger and deeper malware classification models by comparing their performance to ResNet110 and VGG16 baseline models. Specifically, we measured FLOPs, MACs and the number of model parameters to assess computational complexity, while evaluating performance in terms of accuracy and training time. In Table 2, we show the experimental results for ResNet and VGGNet under MMCC and Malimg datasets.

From the experimental results, we first observed that the large and deep malware classification models outperformed small and shallow malware classification models under various datasets. For example, on the MMCC dataset, ResNet32 and VGG8 baseline models achieved accuracies of 96.87% and 97.70%, respectively, which are lower than accuracies of 97.79% and 98.25% achieved by the ResNet110 and VGG16 baseline models. Also, on the Malimg dataset, the ResNet110 and VGG16 baseline models achieved accuracies of 99.36% and 99.25%, respectively, while the ResNet32 and VGG8 baseline models both achieved an accuracy of 99.04%. However, deeper models such as the ResNet110 and VGG16 baseline models cost approximately four times more in terms of FLOPs, MACs, and parameters compared to shallower models such as ResNet32 and VGG8. Also, these models required more than twice the training time.

Second, it is observed that the small and shallow malware classification models with Self-MCKD achieved better performance than the large and deep malware classification models, while using only minimal additional computational cost. For example, on the MMCC dataset, while the VGG16 baseline model achieved an accuracy of 98.25% using 627.04M FLOPs, 313.52M MACs, and 14.77M parameters, the VGG8 model with Self-MCKD achieved an accuracy of 98.62% with only 272.02M FLOPs, 136.01M MACs, and 3.96M parameters. Also, while the VGG16 baseline model took 91 minutes and 8 seconds to train, the VGG8 model with Self-MCKD required only 28 minutes and 19 seconds, which is more than three times faster. Similarly, the ResNet110 baseline model required 508.42M FLOPs, 254.21M MACs, and 1.74M parameters to achieve an accuracy of 97.79%. Its performance is higher than the ResNet32 baseline model, which achieved an accuracy of 96.87% with 138.89M FLOPs, 69.45M MACs, and 0.47M parameters. However, the ResNet32 model with Self-MCKD outperformed the ResNet110 baseline model with 277.78M FLOPs, 138.90M MACs, and 0.47M parameters. For the Malimg dataset, the ResNet32 and VGG8 models with Self-MCKD achieved comparable performance to the ResNet110 and VGG16 baseline models. For example, while the ResNet110 and VGG16 baseline models achieved accuracies of 99.36% and 99.25%, respectively, the ResNet32 and VGG8 models with Self-MCKD achieved accuracies of 99.25% and 99.57%, respectively.

Third, we observed that the small and shallow malware classification models with Self-MCKD effectively transfers knowledge from the teacher to student models. In other words, Self-MCKD not only improved the classification accuracy of the student model but also enhanced its convergence speed. Specifically, the accuracy of the ResNet32 and VGG8 baseline models improved from 96.87% and 97.70% to 98.53% and 98.62%, respectively, using Self-MCKD on the MMCC dataset. Similarly, both the ResNet32 and VGG8 baseline models achieved an accuracy of 99.04% without Self-MCKD. On the other hand, for both models with Self-MCKD, the average accuracy increased by approximately 0.3%. Such results show that the ResNet32 and VGG8 models with Self-MCKD maintained the number of parameters while reducing the training time by approximately 1.5 times compared to the ResNet32 and VGG8 baseline models.

5.2.3. Does Self-MCKD show the better performance than the other knowledge distillation methods?

To compare the classification performance of Self-MCKD with other logit-based knowledge transfer methods, we conducted experiments using Vanilla KD and MLKD. We measured the accuracy, precision, recall, F1-score, and training time of Vanilla KD, MLKD and Self-MCKD. Here, Vanilla KD refers to a traditional knowledge distillation method that transfers knowledge, combining both target and non-target classes, from the teacher model to the student model. Also, MLKD refers to Multi-Level Logit Knowledge Distillation (MLKD) [36], a state-of-the-art method that distills knowledge at multiple logit levels. Unlike these approaches, Self-MCKD transfers knowledge of target and non-target classes independently. To compare the effectiveness of knowledge transfer, we used the same small and shallow malware classification models as the teacher and the student model in Vanilla KD and MLKD, as with Self-MCKD. For all methods, including Vanilla KD, MLKD, and Self-MCKD, we set the temperature factor to 4, following common practices in knowledge distillation experiments [9,14,36].

As shown in Table 3, we observed that Vanilla KD and MLKD also improved small and shallow malware classification models. For example, on the MMCC dataset, the VGG8 model with Vanilla KD achieved an accuracy, recall, precision, and F1-score of 98.44%, 98.34%, 98.36%, and 98.29%, respectively, while the VGG8 baseline model achieved 97.70%, 97.70%, 97.67%, and 97.68%, respectively. Also, for ResNet32, MLKD achieved an accuracy, recall, precision, and F1-score of 97.70%, 97.70%, 97.74%, and 97.57%, respectively, while the ResNet32 baseline model achieved an accuracy, recall, precision, and F1-score of 96.87%, 97.42%, 97.39%, and 97.39%, respectively. Similarly, for the Malimg dataset, it is observed that the VGG8 model with Vanilla KD and MLKD improved its malware classification accuracy, recall, precision, and F1-score by approximately 0.3%. Especially, MLKD outperformed Vanilla KD in recall, precision, and F1-score across both datasets. However, although Vanilla KD and MLKD reduced the training time by approximately 1.5 times compared to the baseline, MLKD required slightly longer training time than Vanilla KD.

Even though Vanilla KD and MLKD improved the performance of small and shallow malware classification models, Self-MCKD outperformed both methods in all metrics. For example, on the MMCC dataset, the ResNet32 and VGG8 models with Self-MCKD surpassed the accuracy, recall, precision, and F1-score of same models with Vanilla KD and MLKD by about 0.2% to 0.8%. Similarly, on the Malimg dataset, the ResNet32 model with Self-MCKD achieved higher accuracy, recall, precision, and F1-score of same models with Vanilla KD and MLKD by about 0.4%. Notably, Self-MCKD showed shorter training times compared to baseline and other knowledge distillation methods.

In Figure 4, we showed ROC curves along with AUC scores for each method on the MMCC dataset. As shown for class index 8 (e.g., Gatak), indicated by the light green line in Figure 4, Self-MCKD outperformed other methods achieving an AUC score exceeding 0.4. In Figure 5, we also visualized the T-SNE results for each method on the Malimg dataset. Compared to the baseline and other knowledge distillation methods, Self-MCKD produced more distinct and well-separated representations.

6. Conclusions

In this paper, we showed that traditional knowledge distillation in AI-based malware classification methods reduces the effectiveness of knowledge transfer by combining the knowledge of the target malware class and non-target classes. Also, this approach often relies on a high-performing, large and deep teacher model, which requires high computational cost during the training process, further reducing the efficiency of knowledge transfer. To address these problems, we proposed a novel knowledge distillation technique, called Self-MCKD. Its main idea is to transfer two types of knowledge, NMKD and TMKD, by separating the output logits of the target class and non-target classes. Also, Self-MCKD utilized a pre-trained small and shallow student model to serve as the teacher model. From the experimental results under various malware datasets, we showed that the separation of output logits in Self-MCKD, consisting of NMKD and TMKD increased the effectiveness of malware knowledge transfer. Specifically, Self-MCKD outperformed traditional knowledge distillation in all evaluation metrics. Furthermore, we showed that the small and shallow malware classification methods with Self-MCKD achieved better performance than the large and deep malware classification methods, while requiring only minimal additional computational cost. Also, we showed that Self-MCKD improved the efficiency of malware knowledge transfer by using small and shallow methods as both the teacher and student models. From such experimental results, we believe that Self-MCKD can enhance the malware knowledge in terms of the effectiveness and efficiency of malware classification better than existing knowledge distillation techniques. In future work, we will validate Self-MCKD under scenarios with highly imbalanced class distributions to assess its robustness and adaptability in real-world malware classification tasks.