MSSA: multi-stage semantic-aware neural network for binary code similarity detection

Instruction normalization

In the field of BCSD, when encountering assembly instructions that were not included in the training data, referred to as the Out-of-Vocabulary (OOV) issue, it can significantly impair the detection performance. To address the OOV issue, it is helpful to abstract and standardize the opcodes and operands in assembly instructions. We apply the following strategies to normalize assembly instructions and reduce their vocabulary size:

(1) Keep all opcodes and registers not normalized, just as the way it is.

(2) Replace all literals with const. For example, sub rsp, 07h becomes sub rsp, const.

(3) Retain operations involving +, -, and * in operands.

(4) Replace the operand following the call operation with foo. For example, call sub_401E85 becomes call foo.

(5) Preserve offset values in operands.

(6) Replace occurrences of var or arg in operands with ptr, since they are typically associated with pointers. For example, test[ebp+arg_0],1 becomes test [ebp+ptr], const.

(7) Represent all the other symbols apart from the above six situations as tag.

Intra-block learning

The execution behavior of assembly instructions is influenced by preceding instructions and can affect subsequent ones, creating complex dependencies. This challenges unidirectional models in capturing the complete semantic information. To address this challenge, the Intra-Block Learning stage employs Bidirectional Gated Recurrent Unit (Bi-GRU) (Schuster & Paliwal, 1997) to further learn the semantic relationships between assembly instructions within a single basic block. Bi-GRU processes data in both directions (forward and backward), thereby capturing the semantic flow from front to back and the dependencies from back to front simultaneously. This bidirectional processing enables a comprehensive understanding of the sequence of assembly instructions within a basic block, ensuring that the semantic context is fully represented.

The feature vector of a basic block b_i as derived in the Block Embedding stage can then be transformed to $b_i^{{}^{\prime }}$ after Intra-Block Learning. $b_i^{{}^{\prime }}$ can be formulated as: (3)${\displaystyle b_i^{{}^{\prime }}=Bi\text{\_}GRU\left(b_i\right),b_i^{{}^{\prime }}\in {\mathbb{R}}^{m\times k},}$ where k is the number of neurons in the Bi-GRU and m is the maximum number of instruction sequences contained in a basic block.

Inter-block learning

As a function typically consists of multiple basic blocks, understanding the structural relationships between these basic blocks is essential for comprehensively grasping the function’s overall functionality. The Inter-Block Learning stage aims to thoroughly capture this structural information and integrate it with the semantic information of each basic block as derived from the Intra-Block Learning stage. MSSA employs an iterative network to uncover and integrate structural information among basic blocks. The iterative network processes the adjacency matrix A, as derived in the Pre-Processing phase, through multiple rounds of iteration, enhancing the model’s ability to generate a more robust and nuanced representation of the function’s overall structure and behavior. The iterative network is shown in Fig. 2.

Figure 2A is composed of vertices and edges, where vertices represent basic blocks, and edges represent control dependencies between basic blocks. The iterative network includes T iterations. After T iterations, the initial features of basic blocks are further enriched with both structural and semantic information across all basic blocks. In contrast to Gemini (Xu et al., 2017a), we omitted the addition operation in the output of the iterative network. We contend that simply adding the basic block features to derive function-level features risks losing critical information.

Figure 2B illustrates the computation process for the i-th basic block in the t-th iteration where the block’s initial feature vector $b_i^{{}^{\prime }}$ is transformed into its inter-block vector ${\mu }_i^{\left(t\right)}$. The inputs for this process are: $b_i^{{}^{\prime }}$, the feature vector of basic block i; and the sum of ${\mu }_j^{\left(t-1\right)}$, namely the inter-block vectors of all the basic blocks adjacent to the i-th block in the adjacency matrix A. The computation process can be formulated as: (4)${\displaystyle {\mu }_i^{\left(t\right)}=tanh\left(b_i^{{}^{\prime }}+\sigma \left(\sum _{j\in N\left(i\right)}{\mu }_j^{\left(t-1\right)}\right)\right),}$ where the tanh function is used as the activation function, and σ is a multi-layer fully connected network responsible for computing an embedding vector with more powerful representation capabilities. N(i) represents the set of basic blocks adjacent to the basic block i. σ can be formulated as: (5)${\displaystyle \sigma \left(x\right)=P_1\times ReLu\left(P_2\times \dots ReLU\left(P_u\times x\right)\right),}$ where u represents the embedding depth of each basic block vertex, x represents the input, and P represents a fully connected parameter matrix. Through T layers of iteration, the features of each vertex are propagated to other vertices via the iterative network, ensuring that control flow information is disseminated among all basic blocks.

Define the output layer dimension of the iterative network as l. The iterative network can convert the feature vectors of all basic blocks into the corresponding function’s feature vector f. f can be formulated as: (6)${\displaystyle f=IterNet\left(\left(b_1^{{}^{\prime }},b_2^{{}^{\prime }},\dots ,b_i^{{}^{\prime }},\dots ,b_n^{{}^{\prime }}\right),A\right),f\in {\mathbb{R}}^{n\times m\times l}.}$

Function-Level learning

The high-level semantic meaning of a function extends beyond its constituent basic blocks and is implied throughout the entire structure of the function. The arrangement of all the basic blocks within a function directly reflects the core semantics and logical structure of the function, which are crucial for understanding the function’s overall behavior. Therefore, conducting in-depth learning at the function level is necessary to capture these comprehensive semantics.

MSSA utilizes Bi-GRU with an attention mechanism (Bahdanau, Cho & Bengio, 2014), where Bi-GRU performs bidirectional learning on the function’s feature vector f, and the attention mechanism enhances the model’s ability to focus on key features. This combination enables the model to more accurately identify the critical features that determine the function’s semantics, thereby achieving a deeper semantic capture of the entire function.

The calculation formula for the attention (att) mechanism is as: (7)${\displaystyle att=sum\left(x\ast softmax\left(tanh\left(x\cdot W+b\right)\right)\right),}$ where x represents the input, W represents the weight matrix, and b represents the bias vector.

Let us define the output layer dimension of Bi-GRU as h. The function’s feature vector f is transformed to f′.f′can be formulated as: (8)${\displaystyle f^{{}^{\prime }}=Bi\text{\_}GR{U}_{att}\left(f\right),f^{{}^{\prime }}\in {\mathbb{R}}^h.}$

Baselines

We compared MSSA to the following four baselines.

Gemini (Xu et al., 2017a): A graph embedding network to compute the embedding vector for each node in the ACFG, ultimately forming the embedding vector for the entire ACFG (https://github.com/Yunlongs/Gemini)

Asm2Vec (Ding, Fung & Charland, 2019): An assembly language embedding model based on the Distributed Memory Model of Paragraph Vectors (PV-DM) (Luo et al., 2023) model (https://github.com/oalieno/asm2vec-pytorch).

SAFE (Massarelli et al., 2019a): An attention-based assembly language embedding model that uses an RNN architecture and attention mechanism to generate function embeddings (https://github.com/gadiluna/SAFE).

jTrans (Wang et al., 2022): A jump-aware Transformer-based model for assembly language embedding, stands as one of the state-of-the-art approaches (https://github.com/vul337/jTrans).

Evaluation metrics

Determining whether a function pair is similar or not is a classification task. Therefore, we adopt the Accuracy, Precision, Recall, F1-Score (F1) and Area Under the Curve (AUC) metrics to evaluate the classification performance of BCSD. AUC is the area under the Receiver Operating Characteristic (ROC) curve, where the ROC curve is plotted based on the True Positive Rate (TPR) and False Positive Rate (FPR). The calculation formulas for each indicator can be formulated as: (10)${\displaystyle Accuracy=\frac{TP+TN}{TP+TN+FP+FN},}$ (11)${\displaystyle Precision=\frac{TP}{TP+FP},}$ (12)${\displaystyle Recall=\frac{TP}{TP+FN},}$ (13)${\displaystyle F1=\frac{2\ast Precision\ast Recall}{Precision+Recall},}$ (14)${\displaystyle TPR=\frac{TP}{TP+FN},}$ (15)${\displaystyle FPR=\frac{FP}{FP+TN},}$ where the True Positive (TP) refers to a correctly defined patient sample, True Negative (TN) refers to a correctly defined healthy case, False Positive (FP) refers to a patient who has been incorrectly identified, False Negative (FN) refers to a health case with an incorrect definition.

Besides, Recall@1 and Mean Reciprocal Rank (MRR) are often used to evaluate the retrieval performance of BCSD (Wang et al., 2022; Gu, Shu & Hu, 2022; Bottou, 2010). Rank@1 means whether the true similar function pair has the highest score, namely ranking first, within a given set of functions. MRR calculates the average of the multiplicative inverse of the rank at which the first correct function is retrieved for a given set of functions. The calculation formulas for MRR and Recall@1 indicators can be formulated as: (16)${\displaystyle \text{Recall}@1=\frac{1}{\left|\mathrm{F}\right|}\sum _{{\mathrm{f}}_{\mathrm{i}}\in \mathrm{F}}\tau \left({\text{Rank}}_{{\mathrm{f}}_{\mathrm{i}}^{\mathrm{gt}}}<=1\right),}$ (17)${\displaystyle \mathrm{MRR}=\frac{1}{\left|\mathrm{F}\right|}\sum _{{\mathrm{f}}_{\mathrm{i}}\in \mathrm{F}}\frac{1}{{\text{Rank}}_{{\mathrm{f}}_{\mathrm{i}}^{\mathrm{gt}}}},}$ where F = f₁, f₂, f₃, …f_n represents the function pool, $f_{\mathrm{i}}^{\mathrm{gt}}$ represents the similar functions of the corresponding function f_i, ${\text{Rank}}_{f_{\mathrm{i}}^{\mathrm{gt}}}$ represents its ranking in the retrieval list, τ(⋅) is defined as : (18)${\displaystyle \tau \left(\mathrm{x}\right)=\left\{\begin{array}{c}0,\mathrm{x}=\text{False}\hfill \\ 1,\mathrm{x}=\text{True}\hfill \end{array}.\right.}$

Hyperparameter selection

We used Adaptive Moment Estimation (Adam) (Kingma & Ba, 2014), Stochastic Gradient Descent (SGD) (Dozat, 2016), and Nesterov-accelerated Adaptive Moment Estimation (Nadam) (Hinton & Salakhutdinov, 2006) for optimization. In addition, we conducted multiple experiments on the number of neurons in each network based on empirical values to select the parameters with the best performance.

We ultimately chose the following hyperparameters for MSSA: the word vector length d = 100, the number of basic blocks n = 20, the number of instructions per basic block m = 20, the number of Bi-GRU neurons k = 128, the number of iterations of the network T = 5, the embedding depth of the iterative network u = 3, the output dimension of the iterative network l = 128, the function embedding dimension h = 256, and the training batch size is 1,024.

Classification performance

We conducted experiments to evaluate the classification performance of MSSA as compared to other four baselines using the BinaryCorp-3M dataset.

The results, presented in Table 2, show that MSSA consistently outperforms all the baselines in terms of Accuracy, Recall, F1 and AUC on classification performance, and ranks second only to jTrans regarding Precision. Specifically, MSSA outperforms jTrans by 0.0035 in AUC, surpasses SAFE by 0.0012 in Recall, 0.0018 in Accuracy, and 0.0018 in F1. However, its Precision is only 0.0158 lower than jTrans.

Retrieval performance

We also conducted experiments to evaluate the retrieval performance of MSSA as compared to other baselines using the BinaryCorp-3M dataset. We set the function pool size to 32 and divided BinaryCorp-3M into six groups based on optimization levels: O0-O3, O0-Os, O1-O3, O1-Os, O2-O3, and O2-Os.

The BCSD results on retrieval performance are presented in Tables 3 and 4. The results indicate that the retrieval performance of MSSA is second only to jTrans. In terms of MRR, the gap between MSSA and jTrans is most significant in the O0-O3 group with MSSA being 0.1162 lower. Meanwhile, the gap is minimal in the O2-Os group with MSSA being 0.0181 lower. On average, MSSA’s MRR is 0.0497 lower than jTrans. Similarly, for Recall@1, MSSA exhibits the largest gap with jTrans in the O0-O3, being 0.1772 lower. In the O2-Os group, the gap is minimal, with MSSA being 0.0343 lower. On average, MSSA’s Recall@1 is 0.0811 lower than jTrans.

The results also show that the retrieval performance of MSSA surpasses other baselines. In terms of MRR, MSSA exhibits the most significant advantage over SAFE in the O2-Os group, achieving an improvement of 0.0517. In the O1-O3 group, although the performance gain of MSSA is smaller, it still achieves an increase of 0.0135. Overall, MSSA averages a 0.0273 higher MRR than SAFE. Furthermore, considering the Recall@1 metric, MSSA demonstrates the most notable difference compared to SAFE in the O2-Os group, achieving a boost of 0.0837. While the advantage of MSSA in the O1-O3 group is relatively modest, it still manages to achieve an increase of 0.022. On average, MSSA outperforms SAFE by 0.0409 in the Recall@1.