GENTAL: GENERATIVE DENOISING SKIP-GRAM TRANSFORMER FOR UNSUPERVISED BINARY CODE SIMILARITY DETECTION
Year: 2021

Abstract:
Binary code similarity detection serves a critical role in cybersecurity. It alleviates the huge manual effort required in the reverse engineering process for malware analysis and vulnerability detection, where often the original source code is not available for analysis. Most of the existing solutions focus on a manual feature engineering process and customized code matching algorithms that are inefficient and inaccurate. Recent deep learning based solutions embed the semantics of binary code into a latent space through supervised contrastive learning. However, one cannot cover all the possible forms in the training set to learn the variance of the same semantics. In this paper, we propose an unsupervised model aiming to learn the intrinsic representation of assembly code semantics. Specifically, we propose a Transformer-based auto-encoder like language model for the low-level assembly code grammar to capture the abstract semantic representation. By coupling a Transformer encoder and a skip-gram style loss design, it can learn a compact representation that is robust against different compilation options. We conduct experiments on four different block-level code similarity tasks. It shows that our method is more robust compared to the state-of-the-art.

heading: 1 INTRODUCTION
text: Reverse engineering is the process of analyzing a given binary program without its source code. It routinely requires experienced analysts and demands a huge amount of manual effort. This process is essential in many critical security problems, such as malware analysis, vulnerability discovery, and Advanced Persistent Threat (APT) tracking. Binary similarity detection is an important solution to reduce the amount of manual effort, by detecting known parts and pieces in the target under investigation. Binary-level similarity detection is more difficult than source-level similarity detection, because most of the semantic-rich literals and structures, such as constants, function names, and variable names, are altered or no longer available in the form of assembly language. The data structures are also lost, due to the compilation process, since the debugging information is typically stripped in commercial-off-the-shelf programs.
Many existing approaches rely on manual feature engineering to model the semantics of assembly code. For example, a fragment of assembly code can be modeled into a numeric vector based on (1) the ratio of algorithmic operations, (2) the ratio of transferal operations, and (3) the ratio of function calls [1, 2, 3, 4]. Alternatively, assembly code can be modeled as word-based n-grams [5]. These approaches cannot capture the rich semantics carried in the assembly code. Some other approaches rely on symbolic constraint solving to measure the logical relationship between each pair of code fragments [6, 7]. However, these methods are computationally expensive and do not scale well.
In comparison, recent deep learning approaches have shown to be more effective and robust to detect similar binary code. Typically, a neural network model is proposed and coupled with a contrastive loss function [8, 9, 10, 11]. The network is trained with limited pairs of assembly code with 0-1 labels. Asm2Vec [12] extends the idea of Word2Vec [13] and PVDM [14], and follows an unsupervised paradigm to learn the function representation.
Although deep learning has been proven to be effective, there are still many practical barriers that prevent the aforementioned approaches from a wide adoptation. First of all, most approaches
except Asm2Vec uses a supervised paradigm and learns the model parameters by directly decreasing the distance between similar assembly code pairs in the training set. While supervised models potentially provide better performance in their trained tasks, they usually suffer when a new task is introduced. For example, [10] learns the tasks for cross-optimization level detection and crossarchitecture detection on ARM and X86-64. However, it does not guarantee cross-compiler and cross-OS binaries performance, thus it lacks robustness. On the other hand, Asm2Vec, which uses an unsupervised approach, is mainly based on PVDM. This limits its stability during the inference, since it needs multiple rounds to accumulate the gradient. Moreover, Word2Vec based models are less context-dependent than the state-of-the-art models such as BERT [15]. They carry less semantics and discriminative power for the downstream tasks. To illustrate this, Fig. 1 shows 4 pieces of compiled binary code in assembly code, all generated from the same source code. The choice of compilers and optimization levels mainly results in performance differences, while the obfuscated code can be generated for various reasons, both benign and malicious. The semantics of the programs are similar or the same, but the code itself can look drastically different from one another. A supervised model may be able to learn to detect similar binaries cross-compiled, but changing the optimization level or adding obfuscation can create unseen patterns, leading to detection failures.
To learn a compact representation of assembly code through an unsupervised language model, one can use BERT and other Transformer-based models, since they have shown their effectiveness in modeling natural language semantics. However, assembly code, despite the fact that it is missing many semantic-rich literals present in the original source code, complies with a more static syntax. The original Transformer architecture, with the flattened sequence of position-embedded subwords, cannot address the position invariance of the assembly code instructions (see Figure 1). The same instructions may be placed in a different place, but contribute to similar program semantics. Additionally, the original masked language model in BERT distribute the memory of reconstructing the masked tokens into hidden layers of different timestamps. This is against our goal, which is to have a single compact vector representation for all code.
With the above observations, we propose GenTAL, a generative denoising skip-gram Transformer for assembly language representation learning. It follows an unsupervised learning paradigm for binary code similarity detection. We model assembly code instructions through the Transformer encoder in a way that their syntax is preserved. Inspired by the denoising autoencoders, we propose to combine the originally time-distributed memory of the masked language model into a single dense vector, and leverage a skip-gram like structure for masked instructions recovery. This allows the model to embed the semantics of the assembly code into a very compact representation for a more effective similarity detection. Our contributions are as follows:
• We propose a new Transformer-based unsupervised language model for assembly code instructions, the form of human-understandable binary code. The model follows the syntax of assembly code and is able to address the instruction position invariance issue. • We proposed to combine skip-gram style reconstruction loss with the masked language model to
condense the originally time-distributed memory into a single compact embedding vector. This design simulates a denoising autoencoder and provides a unified representation of semantics. • We conduct experiments on five different scenarios for code similarity detection and compare our
methods against traditional TFIDF based and state-of-the-art machine learning-based methods. We show that GenTAL is more robust and able to outperform the baselines in all applications.

heading: 2 RELATED WORK
text: Binary Code Similarity. Assembly code can be regarded as a natural language in certain aspects. For this reason, NLP techniques are often used as encoders. Other ML-based binary code similarity detection works use different approaches for representation learning. Supervised learning approaches, such as Gemini [8], Diff [9], [10], and BinDeep [16] all use siamese networks to reduce loss and use cosine distance to compute similarity. Gemini manually extracts block features and feeds them into a graph neural network. Diff feeds raw byte input into CNN for learning function embedding, which lacks the modeling of block semantics. Yu, et al. [10] extend the BERT model for code semantics learning by introducing the same graph prediction task (SGP) and graph classification task (GC). They also train a graph neural network for assembly code representation learning and a CNN on the adjacency matrix for additional order-aware embedding. BinDeep uses Instruction2Vec [17] with LSTM for instruction and blocks embedding to enable sequence-aware modeling. Asm2Vec [12] follows an unsupervised paradigm and uses PVDM for block embedding. Other traditional graph matching methods include BinDiff [18] and Binslayer [19]. These use static graph patterns. As a result, the performance can be severely hindered with any changes in graphs, which often happen with different compiler settings and obfuscation.
Unsupervised Language Models. BERT [15] is the state-of-the-art language model for NLP pretraining based on Transformer [20] architecture. Transformer can learn the contextual and sequential information within a sentence, while also maintaining multiple levels of semantics. It trains in parallel and is thus faster compared to RNN-based models. There are several variations of the original BERT model. ELECTRA [21] builds on top of BERT and adds a discriminator that predicts which tokens are originally masked. Albert [22] achieves similar performance to BERT, while using fewer parameters through factorized embeddings parameterization and cross-layer weight sharing. RoBERTa [23] further up-trains BERT with heavier parameters and discards the next sentence prediction (NSP) task from BERT.
Word2Vec is also an unsupervised learning technique for language models, which uses Skip-gram or Continuous-Bag-Of-Words (CBOW) to learn word embedding based on a fixed length of windowed context. Doc2Vec extends Word2Vec and adds another ID token for document/sentence representation with the Distributed Memory (PVDM) and Distributed Bag-Of-Words (DBOW) variance. Tokenization is also an important task in language models, since the performance can vary significantly depending on the quality of the tokenizers. There are different levels of tokenization, such as character level and subword level. Byte Pair Encoding (BPE) [24, 25] generates subword vocabulary by learning the frequency of characters in a large corpus. Unigram [26] learns the language model by optimizing the word occurrence given a sequence and then builds the vocabulary by sorting the loss of subwords. These methods are used to combat the out-of-vocabulary (OOV) issues, which many large-scale language models can struggle with.

heading: 3 GENERATIVE DENOISING SKIP-GRAM TRANSFORMER
text: In this section, we describe the details of GenTAL with respect to three major components: preprocessing, including instruction masking and encoding; the code fragment encoding using Transformer; and the reconstruction loss through a skip-gram style approach. The overall framework of GenTAL is shown in Fig. 2. Given a sequence of assembly instructions, which is generated from disassembling the binary data into assembly code, we preprocess the assembly code first to extract the instructions as subword sequences and apply masking. Then each subword is mapped into subword embedding and coupled with an instruction-level positional embedding. After being merged as the instruction
level embeddings, we feed them into a Transformer encoder and obtain the CLS step hidden layer output, which is finally position-encoded again to recover the masked instructions. We formally define GenTAL’s goal as to learn an encoding function G that maps a sequence of instructions into a b-dimensional space C → Rb where its semantics will be preserved by optimizing:∑
j ∑ i P (ti|pj , G(C)) (1)
Here C is the sequence of assembly code, pj is the position of the j-th masked instruction, and ti is the i-th subword of that j-th instruction. Simply put, we reconstruct the full instruction rather than individual disconnected subwords or the original tokens.

heading: 3.1 INSTRUCTION PREPROCESSING, MASKING, AND RAGGED TENSOR
text: An assembly instruction contains an operation and operands. We first treat each instruction as a plain text sequence and clean up the assembly code by removing address-specific literals such as loc_DD20. These literals depend on the base address of the binary and can be changed when generated in a slightly different environment. Next, we pre-tokenize the assembly code following the syntax of the assembly language as this can help us alleviate the out-of-vocabulary (OOV) issues. For example, instead of treating [rbg+var_19] as a single token, we break it into [, rbg, +, var, 19]. After, we train an unigram [26, 27] model on the training corpus for subword tokenization. This can further mitigate the OOV issue, especially on long constants. For example, 0x00fffff1 will be broken down as 0x00, ffff, and f1. If any compiler or obfuscator manipulates this constant by shifting or padding, the subwords can still be matched.
Given a sequence of assembly instructions in a linear order, we follow the Masked Language Model used in BERT. However, instead of masking individual subwords or subwords that constitute an original token, we mask a full instruction, since there exists a strong correlation among subwords presented in a single instruction. For example, operand PUSH will very likely be followed by a stack register. XOR is also used quite often with the same register as its two operands. This issue will make the reconstruction task too easy. Instead, we mask the complete instruction.
Unlike many approaches that treat assembly instructions as a flat sequence of words, we keep the original structure and dimension of [blocks, instructions, tokens] as input to preserve the execution logic and carry better semantics. In the past, it was very difficult to model the data structure this way, due to the padding issue where some instructions have a significantly larger number of subwords. Our implementation uses the recent ragged tensor interface in TensorFlow, which allows variable sequence length over the subword dimension without padding. Besides masking, we add a CLS token at the start of the assembly instructions, which acts as a condensed vector for a collective representation. CLS is left out for masking.

heading: 3.2 CODE FRAGMENT ENCODING
text: Next, we obtain subword embeddings and aggregate them into instruction-level representation, before feeding it into the Transformer. Given a subword t of instruction, we map it to its subword embedding, and then couple it with a position embedding. After, we aggregate the position-encoded subword embedding into instruction-level embedding as e. Specifically:
e = 1/m ∗ ∑
t EM(t) + PE1(t) (2)
where m is the instruction’s subword count. The subword-level positional embedding is important in our case, since long strings and constants are broken down into different subwords. For example, the two constants 00f1 and f100 may yield the same set of subwords f1 and 00. With position encoding, the information order can be preserved. After obtaining the embedding for instructions, we feed them into a Transformer model. Suppose the embedding of all the instructions is denoted as E, and e is one of them, we have:
H = TE(E) (3)
Note that in this step, the CLS token is treated as an individual instruction to distinguish it from the other instructions, and the Transformer only runs over the instruction-level embeddings. This design also helps us mitigate the sequence length limit of assembly code. Assembly code fragments are typically much longer than a natural language sentence in terms of subwords count. Now the limit applies to the number of instructions rather than the number of subwords.

heading: 3.3 SKIP-GRAM RECONSTRUCTION LOSS
text: In previous BERT-based methods such as Yu, et al. [10], the MLM task is to predict the masked tokens with the vector outputs from the Transformer corresponding to the time dimension:
ymlmj = Θ(h i j),∀j ∈M (4)
where M is the masked tokens in an assembly instruction and Θ denotes a standard softmax function. The authors also propose other binary classification tasks including next sentence prediction (NSP), same graph prediction task (SGP), and graph classification task (GC). They use the first vector output from the Transformer for classification:
ynsp, ysgp, ygc = σ(h i 0) (5)
where σ denotes a sigmoid function. Although this approach is able to predict the masked token with full exposure to the token embedding, the vector used for the downstream task carries less information.
Similar to the MLM task in BERT, we have replaced the original tokens with a MASK token, swapped it with a random token from the corpus or kept it unchanged. Since we do not fully adopt BERT and only train the MLM task, the blocks are not sampled to be pairs, but individually fed into the network. The training task is to recover the masked or swapped tokens within each instruction using the output from the positional encoding layer. With all the instruction vectors hi from the Transformer encoder,
we take the first one hcls, which is the CLS vector. Using a skip-gram approach, hcls is treated as the paragraph (collectively representing all instructions in our case) representation. In order to model the position information, we duplicate the hcls vector and tile it to the original instruction length Lins to get the matrix hcls. It is concatenated to another positional encoding layer PE2 before the final prediction layer. PE2 encodes the prediction position, which contains the mask locations. This way, we enable hcls to also carry the position information while maintaining a condense learning representation. In particular:
p = PE2(mask)
hp = hcls ⊕ p (6)
where hp is the final prediction vector, encoded with position corresponding to each masked instruction position. Each hpl predicts all m tokens within the instruction l. Formally, we define our MLM task to be:
yl1, ...y l m = Φ(h p l ),∀l ∈ Lins (7)
Φ denotes the sampled softmax [28], a more efficient method compared to the standard softmamx. For multiclass training, we need to train a function F (x, y), which is GenTAL in this case, to compute the logits, which is the relative log probabilities of class y given the context x:
F (xi, y)← log(P (y|x)) +K(x) (8) where K(x) is an arbitrary function that does not depend on y. We use a larger vocabulary size compared to NLP tasks, such as BERT, because of OOV and the training can slow down significantly. Sampled softmax can accelerate the process by reducing the softmax logits calculated for each training example by picking a small set of sampled classes S ⊂ L using a sampling function Q(y|x):
P (Si = S|xi) = ∏ y∈S Q(y|xi) ∏ y∈(L−S) (1−Q(y|xi)) (9)
Then a set of candidates Ci = Si ∪ ti can be created which contains the union of the target class and sampled classes. Through computing posterior probability that y is the target class given xi and Ci, which is P (ti = y|xi, Ci), we can obtain: log(P (ti = y|xi, Ci)) = log(P (y|xi))− log(Q(y|xi)) +K ′(xi, Ci) (10) which are the sampled relative logits that can feed into a softmax classifier. We can further write it as:
logits = F (x, y)− log(Q(x|y)) (11)

heading: 4 EXPERIMENTS
text: Experiments are conducted on a Linux server with 4 Nvidia RTX6000s, an Intel Xeon Gold 5218 CPU, and 300 GB of memory. The software used includes Python 3.8.6 and Tensorflow 2.4.1.
We collect two independent datasets, where one is for unsupervised training and the other is for evaluation. For the training set, we collect malware from MalwareBazaar and Malpedia1 as well as benign system and software programs for Linux and Windows (see Table. 1). Each method is trained on 100,000 code fragments and roughly 1 million assembly instructions. The training set is split into 90% training and 10% validation for hyperparameter tuning. GenTAL is trained multiple times with different random seeds to estimate the variance. The vocabulary size is 80,000 for training the unigram tokenizer through the SentencePiece package. The maximum sequence length is 512, the embedding dimension is 192, the number of attention heads is 6, and the feedforward dimension is 32.
The evaluation dataset is generated by compiling four widely used utility and numeric libraries in different configurations. In practical scenarios, the ground truth similarity is difficult to obtain, especially on the block-level mapping. We present the first verified block-level assembly fragment mapping dataset. Matching on block-level is much difficult than on function-level, since it lacks many details present in the other blocks. Additionally, the function-level mapping used in the other related works cannot guarantee that the assembly code is actually mapping to the same source code, since optimization or obfuscation will mix source code from different functions and break function integrity. We generate mapping pairs based on (1) cross-compiler mapping blocks compiled with GCC and Clang, (2) cross-optimization mapping blocks with different optimizations, and (3) cross-obfuscation mapping original and obfuscated blocks [29].
1https://malpedia.caad.fkie.fraunhofer.de/, and https://bazaar.abuse.ch/

heading: 4.1 COMPARED METHODS AND EVALUATION METRICS
text: We use mean rank (MRR) and precision at 1 (P@1) as the evaluation metrics, since the similarity detection task can be regarded as an information retrieval task. The following list contains the baselines we implement to compare against GenTAL:
TFIDF-str: TFIDF-based matching [30] by treating code as text. IDF is directly estimated on the testing set, so this approach has more edge over other DL-based methods.
TFIDF-mne: same as TFIDF-str, but we only use the operands of the instruction [30]. IDF is directly estimated on the testing set, so this approach has more edge over other DL-based methods.
MLM: BERT-based solution [10] that shows superior results compared to other binary code similarity detection works, such as Gemini [8] and Skip-thought [31], by a large margin, as well as other machine learning solutions. Only the MLM task is included in this variant.
MLM+NSP: Same as MLM, but with the additional NSP task.
MLM+SGP: Same as MLM, but this variant adds the same graph prediction (SGP) task to classify whether the pairs belong to the same graph.
MLM+SFP: Same as MLM, but we introduce an additional task in which to identify if the pairs are extracted from the same function to further add in more supervision.
MLM+GC: Same as MLM, but with the GC task (graph supervision) from Yu, et al [10]. The task includes classification of optimization level and architecture.
MLM+DIS: We implemented ELECTRA [21], where a discriminator is added to MLM to predict which tokens are originally masked, after the generator part of the model has recovered the tokens.
All-no-GC: This variant includes MLM, NSP, SGP, SFP, and DIS, except GC. GC includes supervision and can potentially correlate to the evaluation task.
ALL: We include all previous BERT-based tasks including GC.

heading: 4.2 RESULTS
text: We first evaluate the cross-compiler task between Clang and GCC shown in Table 2. The table on the top shows the evaluation when only the code fragments with the same number of instructions are used. The assembly code will be more similar than in the other situations, such as obfuscation. Therefore all baselines have relatively good performance. GenTAL is able to outperform most other baselines, except a slightly lower MRR for libtomcrypt against MLM+DIS. The evaluation of cross-compiler with different lengths is at the bottom. It is clear that similar pieces of binary code with different lengths can be much harder to detect, due to changed sequences and tokens. Moreover, when using O0, which is unoptimized, the outcomes between different compilers can vary more. We show that BERT-based models struggle to capture the semantics, regardless of the tasks learned, while GenTAL is able to significantly perform better under this circumstance. For the cross-optimization level evaluation shown in Table 3, O2-O3 is shown at the top and O0-O3 is at the bottom. For O2-O3, all methods have again good performance, since O3 is closer to O2 and both optimization levels usually perform the same. Note that all BERT-based models are nearly identical in this task, leading to another confirmation that additional tasks have little contribution to the specific downstream task. For O0-O3, there are gaps in the results, as shown in the bottom table. While GenTAL is able to achieve greater than 0.5 MRR most of the time, BERT-based models struggle to identify similar binary code, due to the larger difference between the two optimization levels. Finally, Table 4 shows the evaluation for the obfuscation task, where GenTAL leads all the other methods. Code obfuscation
is the most impactful similarity detection task, since malware often bypass detection software using it.
Overall, we can show that GenTAL consistently outperforms all baselines under all compiling configurations, with significant margins. We observe that for more difficult tasks, such as crosscompiler with different lengths, O0-O3, and obfuscation, our design still has reasonable performance with the MRR ranging from 0.3 to 0.7. GenTAL only trains with the MLM task. Therefore it also eliminates the need for sampling pairs for training additional tasks, such as NSP and SFP. Moreover, the model parameter is much smaller compared to the BERT-based models, since we use only 1 Transformer encoder and much smaller dimensions in the other layers. The additional training tasks on top of MLM result in little to no performance gain as shown in all tables. In Fig. 3, we plot the training epochs vs. similarity detection evaluation on the fly for the obfuscation task. Even with only learning the MLM task, GenTAL’s evaluation performance increases faster than for all the other baselines. Therefore, we believe that a more condensed vector using skip-gram contains more semantic information and can result in better similarity detection.

heading: 5 CONCLUSION
text: In this paper, we propose GenTAL, a generative approach with denoising skip-gram and Transformer for binary code similarity detection. We improve upon the BERT-based representation learning specifically for assembly code, by introducing the PVDM paradigm to condense the code fragment representation. This design allows us to capture more semantics from assembly code, which is important as often the compiled code can have different lengths and tokens based on configurations, such as optimization levels, compilers, and obfuscations. When evaluating GenTAL, we show that its performance is superior to both BERT-based and TFIDF-based methods for all the different compiler configurations. In practice, our approach is able to train with fewer data and adapt to more scenarios due to its generative nature. The limitation of GenTAL is that the downstream similarity detection performance can be hindered by the OOV problem, especially for heavy numeric libraries with many constants. Although GenTAL achieves better generalizability than the state-of-the-art solutions, there are still areas or tasks, such as other obfuscation methods, which were not evaluated.

Review summaries:
==+== A. Paper summary  This paper proposes a new model architecture to encode assembly code. By using a new reconstruction loss with the new architecture, the paper achieves better fine-tuning results on downstream binary code similarity detection tasks. 

This paper proposes an enhanced transformer for dealing with the generation of assembly code. This is achieved by feeding the model with one instruction at the time (as BERT did), and hence each of these is divided into sub-words in a position-preserving way. The authors shows the efficacy of GenTAL against many other work proposed in literature, showing that their solution is able to better grasp all the possible different level of optimization applied by the compiler (O1 to O3).