GenTAL: Generative Denoising Skip-gram Transformer for Unsupervised Binary Code Similarity Detection

1 INTRODUCTION . 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 can not 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 , can not 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 . 2 RELATED WORK . 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 . 3 GENERATIVE DENOISING SKIP-GRAM TRANSFORMER . 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 . 3.1 INSTRUCTION PREPROCESSING , MASKING , AND RAGGED TENSOR . 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 .

Summary: 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).