Centroid-Based Learning for Malware Detection and Novel Family Identification

While building on existing centroid methods, we believe our work introduces novel adaptations and innovations tailored to the specific challenges of malware detection. Primarily, our key innovation is the Graph Centroid Model (GCM). Our centroid layer is able to locate class representatives, eliminating the need for labor-intensive manual creation. Additionally, because it is integrated within a neural network, it eliminates the need for k-NN clustering outside of the network. This departure from traditional methods signifies a practical advancement in efficiency and accuracy, especially in the context of malware classification. Our proposed mechanism, which relies on the distance between new data points and existing class representatives, introduces computational advantages. By replacing k-nearest neighbors (k-nn) with this distance calculation, we achieve a more efficient and scalable solution, demonstrating our commitment to addressing real-world computational challenges in malware detection. The symbiotic relationship between Centroids, CFGs, and GNNs in our approach allows for the identification of emerging trends and patterns, making it a valuable tool for zero-day malware detection.

While acknowledging that complete elimination of the OOD problem is challenging, our approach focuses on minimizing semantically ambiguous regions. The Centroid Layer introduces a concrete representation for each class, enhancing the model's ability to assign distinct regions for different malware families. This contributes to a reduction in the instances where diverse samples are scored similarly.

In the context of malware detection, reducing the OOD problem is crucial for ensuring the robustness of security systems. We will provide a more in-depth discussion on how our approach translates these technical improvements into practical benefits, emphasizing the enhanced reliability of our model in distinguishing between emerging threats and known malware families.

We will highlight the novelty of the centroid net to better suit the ICLR audience:

Optimizing Centroid Net with Classification Loss

When a Centroid Net is optimized with a standard classification loss, it results in an effective classifier.

We consider a Centroid Net as defined by the function $F(x) = y_{i^*}$, where $i^* = \arg\min_{j} ||f(x) - c_j||$. Here, $f$ represents the neural network that generates feature representations, $c_j$ are the centroids, and $y_{i^*}$ is the predicted label.

The optimization process aims to minimize the classification loss, typically measured using a loss function such as cross-entropy, which is what we use in our methodology. This loss is defined as $L = -\sum_{k}y_k \log(p_k),$ where $y_k$ is the ground truth label and $p_k$ is the predicted probability for class $k$. In the context of the Centroid Net, $p_k$ corresponds to the probability of selecting centroid $k$ based on the feature representations generated by the GCN.

Let us denote the classification loss as $L_{\text{class}}$. The optimization process involves updating the centroids $c_j$ and the neural network weights to minimize $L_{\text{class}}$. Mathematically, we have:

$

\min_{c_j, \text{weights}} L_{\text{class}}(F(x), y),

$

where $$y$$ represents the ground truth labels.

As optimization progresses, the centroids $c_j$ adapt to the specific feature representations generated by the neural network $f$. This adaptation occurs because the loss encourages the network to produce feature representations that effectively discriminate between different classes.

As the optimization converges, the Centroid Net assigns inputs to centroids that are close in feature space. This means that inputs with similar feature representations will be assigned to the same or nearby centroids, leading to effective classification. This behavior aligns with the optimization objective of minimizing the classification loss.

Optimizing the Centroid Net with a standard classification loss function allows us to adapt centroids and neural network features to more effectively discriminate between classes. Therefore, the Centroid Net serves as a trainable component of a larger neural network.

Lastly, we will include the references you provided, thank you for the suggestions.

Responses to Weaknesses Thank you very much for your comments on our work. 1. Hashemi et al. implements a graph embedding technique based on OpCodes, representing a lower-level abstraction in capturing malware behavior. In contrast, our work integrates (CFGs) and the centroid model directly into the Graph Neural Network (GNN). Our use of CFGs allows for a higher-level abstraction, capturing more nuanced and semantically rich patterns in malware execution. The utilization of Control Flow Graphs adds a unique layer of sophistication to our approach, enabling the extraction of intricate, context-aware features that are essential for accurate malware classification. Furthermore, our Graph Centroid Model (GCM), embedded within a Graph Neural Network (GNN), dynamically creates representatives in the embedding space. This dynamic nature empowers our model not only to efficiently classify known malware families but also to adapt to emerging threats, making it highly responsive to the evolving landscape of malware.

While CFGs have already shown promise in malware classification, our work introduces a new methodology, Centroid Nets, specifically designed for detecting out-of-distribution (OOD) new-family threats. The CFGs, processed by our Centroid-enabled GNN (GCM), capture intricate patterns, enabling effective classification of known families and the identification of emerging threats. A critical distinction lies in our innovative use of centroid models, providing a dynamic and adaptive mechanism for classifying malware families. Importantly, our workflow for constructing CFGs involves an open-source tool, providing an accessible and cost-free alternative to the often-expensive IDAPro. This democratization of tools is crucial for fostering collaboration and innovation within the academic community, allowing researchers easy access to the analysis of malware families.

We will be revising our paper with distinct and explicit wording to emphasize the novelty of our Centroid model and its contributions. We appreciate the opportunity to clarify these points and further highlight the uniqueness of our approach.

2. We will ensure that Table 1 includes a more comprehensive comparison with recent approaches for new class family detection, providing a clearer picture of the effectiveness of our algorithm. We will also use center loss in our follow-ups. We will certainly include more baseline analysis in our future version of this work.

3. Malware datasets often exhibit complexities, including diverse variations within a single family. The visualization serves as a valuable insight into the nuanced nature of malware distribution. Our model is designed to handle such complexities by leveraging the embeddings and centroid-based approach, which, as shown in our experiments, enhances the classification performance, especially in the context of novel and evolving malware families. In fact, Figure 5 may also illustrate that some families are mislabeled and could provide information in malware families that actually contain multiple separate types of malware. This could be a useful feature of our work for malware analysts.

4. We appreciate the suggestion regarding the inclusion of a benign software dataset. Our focus is primarily on malware family discovery and classification rather than binary detection. In a real-world scenario, our model can be part of a larger cybersecurity system where distinguishing between benign and malicious software is handled by complementary components. However, we understand the importance of a holistic cybersecurity solution, and we hope to explore malware detection for our update of this work.

5. We appreciate your feedback on the presentation and formatting of the paper. In the upcoming update, we will address the concerns you raised: We will remove repetitions of the equation and move the dataset subsection to outside the Methods section for a more logical structure. In the revision, we will provide concise and informative captions for all figures, and we will move the training loss figures in the supplementary material rather than the main text to enhance readability.

Questions

1. Our dataset consists of CFGs that we retrieved from samples of various malware families, examples which include berbew and sfone. For each experiment, the training set comprises samples of known families used to train and fine-tune the model. The testing set, on the other hand, contains samples from both known families (for evaluating model performance on familiar threats) and completely new, unseen families (for evaluating the model's capability in detecting novel threats). We ensured a diverse representation of families to provide a comprehensive evaluation of our approach.
2. We will remove these two values in our update.

Thank you for recognizing the clarity in our paper's logic and the importance of the research topic. We acknowledge the concern and agree that clarity is crucial. In the revised manuscript, we will provide a more detailed and explicit explanation of the design of Centroid Nets. This will include a step-by-step breakdown of the model architecture, training process, and proofs of correctness. We have included some of what we will add below. We appreciate this feedback. In the revised version, we will include a comprehensive comparison with recent state-of-the-art methods for detecting new class families. This will offer a more thorough evaluation of our approach against other OOD methods, rather than other GNN methods. Our use-case, malware datasets, would contain both new and old families in any given validation set. We would tune epsilon to the malware families that we do have in both the training and validation set.

Algorithm 1: Centroid Partitioning for Malware Family Discovery

1. Input: samples - set of data samples, threshold - distance threshold

2. Output: centroids - set of centroids representing malware families

3. Initialize an empty set of centroids.

4. Choose a random sample from samples and add it to centroids.

5. for each sample in samples do

   1. Compute distances (distances) from sample to all centroids.
   2. Find the minimum distance (min_distance) from distances.
   3. if min_distance > threshold then
      1. Add sample to the set of centroids.
   4. end if

6. end for

7. return centroids.

Algorithm 2: Compute Distances

1. Input: sample - a data sample, centroids - set of centroids

2. Output: distances - list of distances from sample to each centroid

3. Initialize an empty list distances.

4. for each centroid in centroids do

   1. Compute the distance (distance) between sample and centroid.
   2. Add distance to the list of distances.

5. end for

6. return distances.

The CentroidPartitioning procedure takes in a set of samples and a threshold distance value. It starts by randomly selecting one sample from the set as the first centroid. It then iterates over each sample and computes the distance to all existing centroids using the ComputeDistances procedure. If the minimum distance to any centroid is greater than the threshold, the sample is added as a new centroid.

The ComputeDistances procedure takes in a sample and a set of centroids. It iterates over each centroid and computes the distance between the sample and the centroid using a distance metric such as Euclidean distance or Manhattan distance. It returns a list of distances between the sample and each centroid.

Centroid Proof of Differentiability

To prove the differentiability of $L_j$ with respect to $f(x)$, we will calculate its gradient, $\nabla L_j$.

First, we express $L_j$ in terms of the individual components of $f(x)$ and $c_j$ (where $c_j$ are centroids):

$

L_j = \sum_{i=1}^m (f(x)_i - c_j)_i^2,

$

where $(f(x)_i - c_j)_i$ denotes the $i$-th component of the vectors $f(x)$ and $c_j$.

Now, we compute the partial derivative of $L_j$ with respect to the $k$-th component of $f(x)$:

$

\frac{\partial L_j}{\partial f(x)_k} &= \frac{\partial}{\partial f(x)k} \sum{i=1}^m (f(x)_i - c_j)i^2 \ &= \sum{i=1}^m \frac{\partial}{\partial f(x)_k} (f(x)_i - c_j)i^2 \ &= 2 \sum{i=1}^m (f(x)_i - c_j)i \delta{ik},

$

where $\delta_{ik}$ is the Identity matrix, which is 1 when $i = k$ and 0 otherwise.

We combine these partial derivatives to form the gradient vector $\nabla L_j$:

$

\nabla L_j &= \left(\frac{\partial L_j}{\partial f(x)_1}, \frac{\partial L_j}{\partial f(x)_2}, \ldots, \frac{\partial L_j}{\partial f(x)_m}\right) \ &= 2 \left((f(x)_1 - c_j)_1, (f(x)_2 - c_j)_2, \ldots, (f(x)_m - c_j)_m\right) \ &= 2 \left(f(x)_1 - c_j)_1, f(x)_2 - c_j)_2, \ldots, (f(x)_m - c_j)_m\right) \ &= 2 (f(x) - c_j).

$

Therefore, the gradient $\nabla L_j$ of $L_j$ with respect to $f(x)$ is $2 (f(x) - c_j)$.