Learn to Reverse DNNs from AI Programs Automatically

1 Introduction

Over the last few years, deploying Deep Neural Networks (DNNs) on edge devices is becoming a popular trend for giant AI providers. The AI providers compile the private high-quality DNN models into stand-alone programs [Li et al., 2020; Chen et al., 2018a; Rotem et al., 2018; Cyphers et al., 2018] and would like to sell them to other companies, organizations, and governments with a license fee. This edge-deployment situation further exacerbates the risk of model leakage. The DNN model leakage will result in the following three consequences. First, the potential loss of intellectual property. DNNs parameters are typically trained from TB datasets with high training costs. For example, training a DNN on the google cloud TPU will cost more than $400K [Devlin et al., 2018; Raffel et al., 2019; Zhu et al., 2021]. Thus, DNNs architecture and parameters are precious intellectual property of the developers. Second, the potential infringement of privacy. DNNs parameters have been proven to remember sensitive information (e.g. SSN) in the DNNs training data [Carlini et al., 2020]. Therefore, the leakage of DNN parameters will result in the privacy leakage of training data. Last but not least, the potential risk of being attacked by adversaries. DNNs are vulnerable to adversarial attacks [Carlini and Wagner, 2017]. If the attackers obtain the parameters, they could easily compute DNNs gradients to launch white-box attacks. For Example, reversing Apple's image hash model, NeuralHash 1 , will enable attackers to break the content protection provided by macOS. Therefore, conducting penetration testing to protect DNNs' privacy is essential.

Quantifying model leakage risk of on-device DNNs automatically is a challenging task because of the following two reasons: (1) Reversing DNN architectures requires understanding the semantics of the machine code of each DNN layer. Machine code is incomprehensible to humans; unfortunately, modern DNN usually contains hundreds of layers, and each layer corresponds to a function that needs to infer the semantic. (2) The cross-platform deployment of DNN models increase the difficulty of understanding the semantics of the code. DNNs may be deployed on different hardware platforms. The machine code on different hardware platforms have different syntax. Therefore, understanding the semantics of machine code on different hardware platforms requires different domain-specific knowledge. Fig. 1 shows the challenge in reversing DNNs, the left figure shows the computational graph of Resnet-18, reversing ResNet-18 requires inferring the semantic of each node in the graph, and the right figure shows different assembly instructions syntax on different hardware platforms.

In this paper, we propose NNReverse, the first static model stealing attack to reverse DNNs from AI programs automatically. Different from existing model extraction attacks [Jagielski et al., 2020; Oh et al., 2019; Tram`er et al., 2016], which approximate the victim DNNs functionality rather than reversing DNN architectures [Krishna et al., 2019]. Our approach does not train a substitute model to approximate input-output pairs from the victimDNNs. Instead, our approach learns and infer the semantic of machine code for different hardware platforms automatically.

1 Apples-neuralhash-algorithm-has-been-reverse-engineered

Our intuition is that although DNN architectures are complicated, the basic types (i.e. conv, dense) of DNN layers are enumerable. Thus, it is possible to reverse DNN architectures by inferring each DNN layer type. Furthermore, because different DNN layers are implemented as different binary functions, we can infer DNN layer type by inferring semantic corresponding binary functions. Specifically, our attack includes two phases: offline phase and online phase. In the offline phase, we build a database to collect the binary of each optimized kernel layer, then we train a semantic representation model on the decompiled functions. The semantic representation model can be trained with machine code from different hardware platforms, thus, NNReverse could reverse DNNs for different hardware platforms. Although some semantic representation approaches for NLP fields have been proposed, they are not suitable for assembly functions because they ignore the topology semantic in assembly functions. We propose a new embedding algorithm to overcome this limitation and represent assembly functions more precisely, which combines syntax semantic representation and topology semantic. In the online phase, we decompile the victim DNN executable and collect the binaries. By the pre-collected knowledge in the database, we reconstruct the whole DNN model. We summarize our contribution as follows:

• Characterizing: We characterize the model leakage risk of on-device DNNs. Specifically, we show that ondevice DNNs could be reversed easily.
• Approach: We propose a new fine-grained embedding approach for representing assembly functions, which combines syntax representation and topology structure representation.
• Evaluation: We implement NNReverse and evaluate NNReverse on two datasets. The evaluation results show NNReverse can achieve higher accuracy in searching similar semantic functions than baseline methods, and we apply NNReverse to reverse real-world DNNs as a case study. Results show that NNReverse can reverse the DNNs without accuracy loss.
• Dataset: We open-source our decompiled assembly function dataset 2 , including more than 80k functions compiled for four hardware platforms. Our dataset is the first dataset for DNN reverse engineering, contributing to future research on on-device DNN security.

2 Background &RelatedWork

2.1 DNNDeployment

DL compilers are used to generate portable AI programs from high-level DL frameworks such as Pytorch, TensorFlow, and Caffe. The generated AI programs are advanced optimized (e.g. Layer Fusion) to match the target hardware devices. The working process of DL compiler is shown in Fig. 2. DL compilers accept DL models from highlevel DL frameworks as input and first translate DL models into high-level intermediate representation (IR). The highlevel IR (as known as graph IR) is a computational graph, where nodes represent the atomic DL operators (Convolution, Pooling, Convolution+ReLu, etc.) or tensors, and the edges represent the direction of dataflow. Then DL compilers translate each atomic DL operator into low-level IR and generate the executable code by optimizing the computation and memory access for the target hardware devices. DL compilers enable running DL models on different hardware devices such as X86, ARM, NVIDIA, and TPU. For more details about DL compilers, we refer readers to the existing work [Li et al., 2020; Chen et al., 2018a; Rotem et al., 2018; Cyphers et al., 2018]. In this paper, we consider TVM [Chen et al., 2018b], a widely-used DL compiler to deploy DNNs. Given a DNNs from high-level frameworks (e.g. PyTorch), TVMwill output three files: graph file, parameter file, and binary file. Notice that, These files are necessary for compiled AI programs to run. There are similar functionality files for other DL compilers. For example, Apple Core ML applies the file formats .espresso.net and .espresso.weights to store computational graphs and weights.

2 https://drive.google.com/file/d/NNReverse

2.2 Model Extraction Attacks

As shown in Table 1, a number of attack algorithms [Jagielski et al., 2020; Oh et al., 2019; Tram`er et al., 2016; Zhu et al., 2021] have been proposed to extract DNNs. We summarize the existing attacks into two categories based on their extraction fidelity. The first type of existing attack is functionality attack, they treat the victimmodel as a label oracle and train another model to approximate the input-output pair from the victim model. These attacks can not steal the DNN architecture and the DNN parameters, instead, they just clone the functionality of the victimmodel. The second type of existing attack is steal attack. The goal of this attack is to reverse the actual DNN architecture or parameters. For example, [Yan et al., 2020] introduced a cache-based side-channel attack to steal DNN architectures, [Roberts et al., 2019] assume the model architecture and the output of the softmax layer is known, and replicate the model parameters by feeding noise inputs. Different from the existing steal attack, we propose an infer-based steal attack which requires no assumption about the environment in which the DNNmodel is running.

2.3 Binary Code Analysis

In this section, we introduce some binary analysis techniques for reverse engineering. Traditional binary analysis techniques often extract hand-crafted features by domain experts

to match similar functions and infer binary function semantics. Recently, motivated by the representation learning in NLP fields, researchers propose to apply learning-based approaches for binary code analysis. Specifically, researchers propose to learn a function representation that is supposed to encode the binary function syntax and semantics in low dimensional vectors, known as function embeddings. For example, inspired by Doc2Vec, Asm2Vec proposes to model binary code as documents and design a novel learning algorithm to embed binary functions. However, the learning architectures adopted in existing approaches are limited in modeling long-range dependencies (based on Doc2Vec). The binary code for the DNN operators is long-range dependent because of the loop operations. For example, the four-level loops in the Conv2D layer. Unlike existing approaches, our approach models both instruction and topology semantics of binary code, thus representing the semantics more precisely.

3 NNReverse

3.1 ProblemFormulation

In this paper, we consider an AI model privatization deployment scenario, where the DNN model developers compile their private DNN models into stand-alone AI programs and sell the programs to third-party customers with subscription or perpetual licensing. The considered scenario is realistic because many AI models have been deployed on mobile devices (e.g., Apple NeuralHash, Mobile Translation App, AIoT)and this assumption is widely used in existing work. The endusers can physically access the AI programs, especially the machine code (executable code) of the AI programs. The goal of malicious attackers is to reverse the DNNs from the AI program. Specifically, they can reverse the architecture of the DNNs and reconstruct the DNNs with high-level frameworks (i.e. Pytorch) to compute the DNNs gradients to launch evasion attacks. Formally, our goal includes: (i) inferring the kernel operator type of each binary function in the victim AIprograms; (ii) reconstructing whole DNNs with high-level DL frameworks. Reversing on-device DNNs have a severe real-world impact because DNNs parameters are precious intellectual property, and reverse engineering can bridge the gap between white-box and black-box attacks for DNNs (e.g., bypass Apple NeuralHash). Furthermore, understanding the semantic intricacies of deployed binary code makes automatically reversing on-device DNNs difficult.

3.2 Design Overview

As shown in Fig. 3, our approach includes offline and online steps. In the offline step, we collect a set of DNNs with known architectures and use TVM [Chen et al., 2018b], a widely-used DL compiler, to compile them into AI programs including a binary file for each program. We keep the debugging information for each binary file to keep a record of the kernel operator type of each binary function. We then train a semantic-representation model with these binary files. In the online step, given a victim AI program, we first parse the AI program and decompile the parsed stripped binary file into a list of assembly functions. Next, we feed each decompiled assembly function into the trained semantic representation model to estimate the vectors. Finally, we compare the estimated vectors with the vectors in our offline database to infer the kernel operator type. We could reverse the DNNs arch after inferring the kernel operator type and rebuild it with PyTorch by combing the weight files.

3.3 Semantic RepresentationModel

In this section, we propose an embedding model to represent assembly functions as continuously distributed vectors. As shown in Fig. 4, an assembly function can be represented as a control flow graph (CFG), where nodes are blocks consisting of sequential assembly instructions and edges are control flow paths. Our embedding paradigm, in particular, fuses both text representation and structure representation from the assembly functions.

Text Embedding. We first model the assembly functions as sequential assembly instructions and seek to learn a function ϕt(·), whose input is a sequential assembly instruction and output is a N dimensional vector. We embed the semantics of assembly instructions with more fine-grained modeling rather than treating each instruction as a token. Specifically, for each assembly instruction, we differentiate the opcode and operands in the instruction. Then we concatenate the vector of opcode and operands as the representation of the instruction. For instruction with more than one operands, we use the average vector instead. It can be formulated as:

Where ins represents an assembly instruction, Ω(·) and Θ(·) are the functions to parse opcode and operands respectively, and E(·) is the basic embedding layer that transfer tokens to corresponding numeric vectors. Take the assembly instruction ins = push, rbx as example, then Ω(ins) = push and Θ(ins) = rbx. Because the operands may be a memory address (e.g. 0x73ff) or a immediate value (e.g. 0x16, 0x1024). To eliminate the issue of out of vocabulary (OOV), we propose a blur strategy. Specifically, for memory address value and immediate value, we compare the value with 1000. If the value is less than 1000, we keep the original value, otherwise, we use token large to replace the original value.

After modeling assembly instruction, we seek to represent the semantic of the assembly function. Specifically, we treat each instruction I(ins) as a "word" and apply a similar idea with doc2vec. Formally, our training objective can be formulated as Equation 2. Where M is the number of the training assembly functions, fj is the j th assembly function, T(·) is the text representation, Ci is the neighbor instruction of Ii(ins) and ϕt is the text embedding under learned.

Algorithm 1 Structure Embedding Algorithm

Input: G. {A set of CFGs topology structures}

Input: D. {Maximum degree of rooted subgraphs}

• 1: Initialization: Sample ϕs fromRG×N
• 2: for gi ∈ G do
• 3: for n ∈ NODE(gi) do
• 4: for d = 0 to D do
• 5: subg = ComputeWLSubgraph(n, gi, d)
• 6: ϕs = ϕs - ∂-logP(subg | ϕs(gi)) ∂ϕs
• 7: end for
• 8: end for
• 9: end for

Structure Embedding. Given a set of CFG topology structures, we intend to learn a distributed representation for every CFGs' topology structure. Formally, we represent each CFG as a graph G = (N, E), where N are assembly blocks and E are connections between blocks. Our goal is to learn a embedding function ϕs(·), whose input is g ∈ G and output is a N dimensional vector. Our intuition is that two graphs are more topology similar if these two graphs have more topology similar sub-graphs. Different from the semantic embedding that we apply random-walk paths as representation, applying sub-graphs will keep higher-order substructure and non-linear substructure information. As shown in Algorithm 1. Our algorithm require a set of graphs G as the training corpus. In each training iteration, we compute the WeisfeilerLehman sub-graphs [Shervashidze et al., 2011] of each node in each graph gi, then we try to maximize the possibility that the subgraphs appear under the whole graph and update our structure embedding function ϕs(·). We repeat this process iteratively until the maximum epochs is reached.

Feature Fusion. Finally, we combine the vectors from text embedding and structure embedding to obtain the final function embedding. Equation 3 shows our fusion process, where f is the assembly function, T(f) and S(f) are the text representation and the structure representation of f respectively, ϕt and ϕs are the text embedding function and structure embedding function, and || is the vector concatenation.

Binary Matching. After the model is well trained, we obtain the pre-trained database U = {(uj , namej)}, where uj and namej are the vector and the kernel type of the j th function in the database. During the inference step, given a binary function B, we first feed B into our model to get a vector v. we then apply Eq.(4) to assign B a known kernel type.

3.4 DNNReconstruction

After training a semantic representation model ϕbinary. We propose a automatic algorithm to rebuild the victim DNNs. As shown in Algorithm 2, the inputs of our algorithm are the victim AI program and a trained semantic representation model. We first parse the program to collect the graph files (CG) and collect all kernel nodes N in CG based on whether the node has an edge point. While list N is not empty, we select a node c from N based on whether c's all precursors have been inferred (line 3). For the selected node c, we find the corresponding binary function f (line 4) and apply our semantic representation model to infer the kernel type k of this function (line 6). After that, we compute the output of node c and remove it from N. We repeat this process iteratively until N is empty. Finally, we return the computed tensor of the output node in CG. By doing so, we reconstruct the DNNs from the AI programs.

Algorithm 2 DNNs Reconstruction Algorithm

Input: VicProg. {VictimAI program }

Input: ϕbinary {Semantic representation model}

• 1: CG = ParseFile(VicProg)
• 2: N = CollectKernelNodes(CG)
• 3: while N is not empty do
• 4: c = SelectNode(N) {select a node whose all precursors are computed}
• 5: f = FindBinaryFunction(c)
• 6: k = ϕbinary(f) {Infer kernel type based on Eq. 4}
• 7: c.type = k {Assign kernel to node c}
• 8: Compute node c
• 9: Remove c fromN
• 10: end while

Output: Return computed tensor of the output node in CG.

4 Evaluation

We answer following questions through empirical evaluation.

• RQ1: How accurate of proposed representation model?
• RQ2: Can our algorithm 2 rebuild DNNs correctly?
• RQ3: Is our approach sensitive to hyper-parameters.
• RQ4: What's the contribution of each component?

4.1 Experimental Setup

Datasets. We first download the open source models from TorchVision and HuggingFace as our dataset. Our model dataset includes vision models (e.g. VGG-11, etc.) and textual models (e.g. BertSentenceClassification, etc.). We then use TVM, a popular deep learning compiler [Chen et al., 2018b] to compile deep neural networks into AI programs. We consider four different deployment platforms (e.g., Android, Intel, ARM, and Aarch64), and we set the compilation optimization level from 0 to 4. For the executable files of the compiled programs, we use radare2 to disassemble the binary functions into assembly instructions. We split 80% of the data as our training dataset and the rest 20% for testing. The statistic of our dataset can be found in Table 4.

Comparison Baselines. We compare our proposed method with six unsupervised semantic representation methods: word2vec, doc2vec, and asm2vec. For word2vec, we use the average vectors of all instructions in the assembly function to represent the semantic. For word2vec and doc2vec, we consider both Skip-Gram and CBOW models. For a fair comparison, for each method, we set the dimensions of the embedding as 100.

4.2 (RQ1) Results for Binary Function Searching

We follow existing work [Alon et al., 2018; Alon et al., 2019] to apply Semantic Precision, Semantic Recall and Semantic F-1 score as our metrics. Table 3 summarizes the main results. From the results, we could observe that: the proposed method significantly improves the binary function mapping accuracy across two datasets on all hardware platforms, The observation demonstrates our embedding technique's success in enhancing the prediction capacity in the binary mapping problem.

4.3 (RQ2) Reconstruct DNNs

In this section, we evaluate the effectiveness of our reconstruction algorithm (i.e. Algorithm 2). Because a learningbased technique cannot guarantee soundness results, we choose DNNs with high kernel operator similarity scores as examples. The selected model architectures are listed in Table 5. We use the accuracy loss as our evaluation metric. Formally, we feed random inputs into the original and reconstructed models. Then, the accuracy loss is defined as the accuracy difference between the accuracy of the original model and the reconstructed model. Table 5 lists some examples of

our successful reversed DNN models.

We perform the sensitivity analysis of our proposed techniques on the Textual dataset. (a) embedding dimension: We first conduct experiments to understand how different embedding dimensions affect our technique. Specifically, we configure the embedding dimension as 100, 200, 300 and measure F-1 under different settings. (b) context length: We then configure the context size as 3, 5, 7 and train different models. The results of different embedding configure are listed in Fig. 5. From the results, we observe that NNReverse can keep a stable F1 scores on different embedding dimension and context length settings, which indicates that NNReverse is not sensitive to embedding dimension and context length.

4.5 (RQ4) Ablation Studies

Can fine-grained modeling eliminate Out of Vocabulary (OOV)? We first conduct experiments to explore whether our fine-grained instruction embedding can eliminate the issue of OOV. We model each assembly instruction as a "word" and measure the OOV ratio, and then we measure the OOV ratio of our instruction embedding. The OOV results are summarized in Table 6. We observe that directly modeling each instruction as "word" will result in a high OOV ratio. This is because assembly instruction may contain the memory address, and the value of the address is had to predict.

The performance of each representation component? In this study, we compare our embedding method to solo text embedding and solo topological embedding to investigate the efficacy of embedding combinations. From the results in Table 2, we observe that: (1) for most cases, 20 out of 30, combing text embedding and structure embedding can increase binary mapping accuracy. (2) structure embedding performs the poorest in all scenarios, implying that simply applying structure information is insufficient to infer function semantics. CFG structures, on the other hand, are complementary to text, and combining them yields better results.

5 Discussion &Conclusion

In this research, we present a technique for automatically reversing DNNs from deployed AI applications. In particular, we propose a novel embedding technique for capturing the semantics of assembly functions. The experiments demonstrate that NNReverse is possible to correctly reverse DNNs.

Acknowledgments

This work was supported by NSF grant CCF-2146443.