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: ==+== 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.