SlowGuess/NAACL_2022_OCR huggingface dataset:
Xin Wang $^{1\diamond}$  Yasheng Wang $^{2}$  Yao Wan $^{3}$  Jiawei Wang $^{4}$  Pingyi Zhou $^{2}$  Li Li $^{4}$  Hao Wu $^{5}$  Jin Liu $^{1\boxtimes}$
$^{1}$ School of Computer Science, Wuhan University, China  $^{2}$ Huawei Noah's Ark Lab  
 $^{3}$ School of Computer Sci. & Tech., Huazhong University of Science and Technology, China  
 $^{4}$ Faculty of Information Technology, Monash University, Australia  
 $^{5}$ School of Information Science and Engineering, Yunnan University, China  
{xinwang0920, jinliu}@whu.edu.cn
Recent years have witnessed increasing interest in code representation learning, which aims to represent the semantics of source code into distributed vectors. Currently, various works have been proposed to represent the complex semantics of source code from different views, including plain text, Abstract Syntax Tree (AST), and several kinds of code graphs (e.g., Control/Data Flow Graph). However, most of them only consider a single view of source code independently, ignoring the correspondences among different views. In this paper, we propose to integrate different views with the natural-language description of source code into a unified framework with Multi-View contrastive Pre-training, and name our model as CODE-MVP. Specifically, we first extract multiple code views using compiler tools, and learn the complementary information among them under a contrastive learning framework. Inspired by the type checking in compilation, we also design a fine-grained type inference objective in the pre-training. Experiments on three downstream tasks over five datasets demonstrate the superiority of CODE-MVP when compared with several state-of-the-art baselines. For example, we achieve 2.4/2.3/1.1 gain in terms of MRR/MAP/Accuracy metrics on natural language code retrieval, code similarity, and code defect detection tasks, respectively.
Code intelligence that utilizes machine learning techniques to promote the productivity of software developers, has attracted increasing interest in both communities of software engineering and artificial intelligence (Lu et al., 2021; Feng et al., 2020; Wang et al., 2022; Wan et al., 2022a; Wu et al., 2021). To achieve code intelligence, one fundamental task is code representation learning (also
Models	Tokens	AST	Graph	PT
CodeBERT (Feng et al., 2020)	✓	✗	✗	✗
GraphCodeBERT (Guo et al., 2021)	✓	✗	✓	✗
SynCoBERT (Wang et al., 2021)	✓	✓	✗	✗
CodeGPT (Lu et al., 2021)	✓	✗	✗	✗
PLBART (Ahmad et al., 2021)	✓	✗	✗	✗
TreeBERT (Jiang et al., 2021)	✓	✓	✗	✗
ContraCode (Phan et al., 2021)	✓	✗	✗	✓
CoTexT (Phan et al., 2021)	✓	✗	✗	✗
CodeT5 (Wang et al., 2021b)	✓	✗	✗	✗
CODE-MVP (Our work)	✓	✓	✓	✓
Table 1: Comparison with current pre-trained code models. PT: Program Transformation.
known as code embedding), which aims to preserve the semantics of source code in distributed vectors (Alon et al., 2019). It can support various downstream tasks about code intelligence, including code defect detection (Omri and Sinz, 2020; Zhao et al., 2021b,a), code summarization (Wan et al., 2018), code retrieval (Wan et al., 2019), and code clone detection (White et al., 2016).
Current approaches to code representation borrow ideas from the successful deep learning methods in natural language processing, mainly attributed to the naturalness hypothesis in source code (Allamanis et al., 2018). From our investigation, existing approaches mainly represent the source code from different views of code, including code token in plain text (Iyer et al., 2016), Abstract Syntax Tree (AST) (Bui et al., 2021a), and Control/Data Flow Graphs (CFGs/DFGs) of code (Cummins et al., 2020; Wang and Su, 2020). Recently, many attempts have been made to pretrain a masked language model for source code, such as CodeBERT (Feng et al., 2020), GraphCodeBERT (Guo et al., 2021), SynCoBERT (Wang et al., 2021), CodeGPT (Lu et al., 2021), PLBART (Ahmad et al., 2021), CoTexT (Phan et al., 2021), and CodeT5 (Wang et al., 2021b). Table 1 shows the contribution of our work when compared with current pre-trained language models for source code.
$\diamond$  Work conducted during an internship at Huawei Noah's Ark Lab.
Findings of the Association for Computational Linguistics: NAACL 2022, pages 1066 - 1077 July 10-15, 2022 ©2022 Association for Computational Linguistics
Despite much progress in code representation learning, most of them only consider a single view of source code independently, ignoring the consistency among different views (Feng et al., 2020; Lu et al., 2021; Ahmad et al., 2021; Wang et al., 2021b). Usually, a program, accompanied by a corresponding natural-language comment (NL), can be parsed into multiple views, e.g., the source code tokens, AST, and CFG. We argue that these different views contain complementary semantics of the program. For example, the source code tokens (e.g., method name identifiers) and natural-language comments always reveal the lexical semantics of code, while the intermediate structures of code (e.g., AST and CFG) always reveal the syntactic and executive information of code. In addition, a program can also be transformed (or rewritten) into different variants that have equivalent functionality. We think that different variants of the same program reveal the functional information of code. That is, those different program variants with the same functionality are expected to represent the same semantics.
Inspired by the aforementioned insights, this paper proposes a novel CODE-MVP for code representation, which aims to integrate multiple views of the code into a unified framework with multi-view contrastive pre-training. Concretely, we first extract multiple views of code using several compiler tools, and learn the complementary information among them under a multi-view contrastive learning framework. Meanwhile, inspired by the type checking in compilation process, we also introduce fine-grained type inference as an auxiliary task in the pre-training process to encourage the model to learn more fine-grained type information.
To summarize, the contributions of this paper are two-fold: (1) We are the first to represent source code from multiple views, including the code tokens, AST, CFG, and various program equivalents, under a unified multi-view contrastive pre-training framework. Meanwhile, we also introduce an auxiliary task of inferring type annotations for variables. (2) We extensively evaluate CODE-MVP on three program comprehension tasks. Experimental results demonstrate the superiority of CODE-MVP when compared with several state-of-the-art baselines. Specifically, CODE-MVP achieves 2.4/2.3/1.1 gain on MRR/MAP/Accuracy metrics in natural language code retrieval, code similarity, and code defect detection tasks, respectively.
Figure 1: An example of converting a program from source code into machine code in compilation process.
2 Multiple Views of Code
We borrow ideas from the way that computers process the source code in compilation, where a program would be converted into multiple views. Figure 1 shows the process of converting a program from source code to machine code. During this process, the compiler would automatically utilize some program analysis techniques to verify the correctness of source code, including lexical, syntax, and semantic analyses. In the lexical analysis, a program is treated as a sequence of tokens and checked for spelling problems. In the syntax analysis, syntactic rules of programs are defined by the context-free grammar (Javed et al., 2004). Then the program could be parsed as an AST, based on which many program transformation heuristics can be applied to rewrite the program while maintaining the same desired functionality. In the semantic analysis, semantic rules of the program are defined by the attribute grammar (Paakki, 1995). Then the compiler could check the types of code tokens, and a decorated AST could be obtained. After the three stages above, a translator will convert the source code to its Intermediate Representation (IR), which is then considered as the basis for building Control/Data Flow Graphs (CFGs/DFGs) for further optimizations in the static analysis. Finally, the IR of the source code should be converted into machine code to execute through a code generator. Next, we introduce how we extract different views of the source code. Figure 2 illustrates multiple views of source code with an example.
Abstract Syntax Tree (AST). An AST, which is composed of leaf nodes, non-leaf nodes and edges between them, contains rich syntactic structural information of source code. In the AST, an assignment statement  $\mathrm{y} = 0$  can be represented by a non-leaf node assignment that points to three
Figure 2: Multiple views of source code.
leaf nodes  $(0, y, \text{and} =)$ . In this paper, we parse a snippet of source code into an AST using a standard compiler tool tree-sitter. To feed an AST into our model, we apply depth-first traversal to convert it into a sequence of AST tokens (Kim et al., 2021).
Control Flow Graph (CFG). CFG, which represents the execution semantics of the program in the form of a graph, is one intermediate representation of programs. A CFG consists of basic blocks and directed edges between them, where each directed edge reflects the execution order of the two basic blocks in the program. We can easily traverse the CFG along directed edges to parse it into a token sequence, which reveals the execution semantics of the program. In this paper, we use a static analyzer Scalpel² (Li et al., 2022) to construct the CFGs for Python code snippets.
Program Transformation (PT). The program transformation operations aim to produce multiple variants for a given program that satisfy the same desired functionality (Rabin et al., 2020). These different variants of a program can help the model capture functional semantics. In this work, we employ the following program transformation heuristics on ASTs and rewrite one program into another equivalent variant.
- Function and Variable Renaming. We randomly take new names from a set of candidates, such as VAR_i, FUNC_i, to rename the names of variables and functions in a program. This heuristic will not change the AST structure of the program, except for the textual appearance of variable and function names in the AST.
- Loop Exchange. The for and while loops represent the same functionality in a program. We traverse the AST to identify the for and while loop nodes, and replace for loops with while loops or vice versa. We also modify the initialization, condition and afterthought simultaneously.
- Dead Code Insertion. We first traverse the AST to identify several basic blocks (Mendis et al., 2019), and then randomly select a basic block and insert dead code snippets into it. Note that the dead code snippets are predefined and selected from a set of candidates.
3 CODE-MVP
3.1 Tasks and Notations
We define the set of program samples in multiple views (i.e. NL, PL, AST, CFG, PT) as  $S = \{S^1,\dots ,S^m\}$ , where  $m$  represents the number of views,  $s_i^a\in S^a$  represents a program in the view of  $a$ . Given a program, the PL view denotes its textual appearance, the NL view denotes its corresponding natural-language comment, and the PT denotes the variants of this program based on program transformation. The AST and CFG are extracted from a program using several compiler tools. CODE-MVP adopts two forms of input, i.e., single-view input  $x_{i}^{a} = \{<\mathrm{CLS}>,s_{i}^{a}\}$  and dual-view input  $x_{i}^{ab} = \{<\mathrm{CLS}>,s_{i}^{a},<\mathrm{SEP}>,s_{i}^{b}\}$ , where  $a$  and  $b$  denote two different views of the program. Following (Devlin et al., 2019), a special token  $<\mathrm{CLS}>$  is appended at the beginning of each input sequence, and  $<\mathrm{SEP}>$  is used to concatenate two sequences. Subsequently, the representation of  $<\mathrm{CLS}>$  is used to represent the entire sequence, and  $<\mathrm{SEP}>$  is used to split two views of sub-sequences. Given a set of programs with their corresponding multiple
$^{1}$ https://github.com/tree-sitter/tree-sitter
$^{2}$ https://github.com/SMAT-Lab/Scalpel
Figure 3: An illustration of our proposed multi-view contrastive pre-training framework.
views, we aim to learn the code representation by utilizing the mutual information existing in different views. Our intuition is to learn complementary information from multiple views of code by pulling the code under different views together and pushing the dissimilar ones apart.
Figure 3 shows a simple example of our multi-view contrastive pre-training framework. Given a program  $s_i$ , we use the same program to construct a pair of positive samples  $(x_i^a = \{<\mathrm{CLS}>, s_i^a\}$  vs  $x_i^b = \{<\mathrm{CLS}>, s_i^b\})$  in the form of views  $a$  and  $b$ , as described above. We take  $x_i^a$  and  $x_i^b$  as the input of CODE-MVP respectively. The last hidden representations of  $<\mathrm{CLS}>$  tokens in the two inputs can be formulated as  $h_i^a = \mathrm{CODE - MVP}(x_i^a)$  and  $h_i^b = \mathrm{CODE - MVP}(x_i^b)$ . We utilize a projection head (a two-layer MLP) to map hidden representations to a space, i.e.,  $v_i^a = f(h_i^a)$ ,  $v_i^b = f(h_i^b)$ . Then the multi-view contrastive objective can be performed. During the pre-training process, we also design other two pre-training tasks, i.e., fined-grained type inference (FGTI) task and multi-view masked language modeling (MMLM).
3.3 Multi-View Contrastive Learning
We train CODE-MVP with paired data and unpaired data. Paired data refers to those program samples with paired NL, while unpaired data stands for those isolated program samples without paired NL. Next, we explain how we construct positive
and negative samples for these two cases.
Multi-View Positive Sampling. We design Single-View (for paired and unpaired data) and Dual-View (for paired data only, which needs the NL) methods to construct multi-view positive samples for the MVCL objective:
- Single-View. To bridge the gap between different views of a same program, we consider the view of a program  $x_{i}^{a}$  as a positive sample w.r.t another view  $x_{i}^{b}$ . That is,  $(x_{i}^{a} = \{<\mathrm{CLS}>, s_{i}^{a}\}$  vs  $x_{i}^{b} = \{<\mathrm{CLS}>, s_{i}^{b}\})$  forms an inter-view positive pair, since  $x_{i}^{a}$  and  $x_{i}^{b}$  are two different views of a same program  $x_{i}$ .
- Dual-View. There are a total of  $C_m^2$  combinations for two views of a same program. For efficiency, we focus on the features of the program itself, and propose the NL-conditional dual-view contrastive pre-training strategy, freezing the position of NL. Concretely, we construct a NL-conditional interview positive pair by replacing the second view in the input  $\{<\mathrm{CLS}>, s_i^{\mathrm{NL}}, <\mathrm{SEP}>, s_i^a\}$  to be  $\{<\mathrm{CLS}>, s_i^{\mathrm{NL}}, <\mathrm{SEP}>, s_i^b\}$ , where  $\forall a, b \neq \mathrm{NL}$ .
It is worth mentioning that there are many combinations to construct positive pairs. Some combinations are not considered in this work, such as the AST vs PT of the same program, and the CFG vs PT of the same program. Simultaneously, for training efficiency and downstream applications, we comprehensively consider eight combinations. They are (1) single-view: (NL vs PL), (NL vs PT), (PL vs AST), (PL vs CFG), and (PL vs PT); and (2) dual-view: (NL-PL vs NL-AST), (NL-PL vs NL-CFG), and (NL-PL vs NL-PT).
Multi-View Negative Sampling. Since the processes of unpaired data and paired data are similar, here we take the unpaired data as an example. We leverage in mini-batch and cross minibatch sampling strategies (Chen et al., 2020) to construct intra-view and inter-view negative samples, respectively. Given a mini-batch of training data  $b_{1} = [x_{1}^{a},\ldots ,x_{n}^{a}]$  in the view of  $a$  with size  $n$ , we can easily get another positive mini-batch data  $b_{2} = [x_{1}^{b},\dots,x_{n}^{b}]$  in the view of  $b$ , where  $(x_i^a \lor x_i^b)$  denotes an inter-view positive pair. For  $x_{i}^{a}$ , the intra-view negative samples are  $\{x_j^a\} ,\forall i\neq j$ , and the inter-view negative samples are  $\{x_j^b\} ,\forall i\neq j$ . Finally, for each  $x_{i}$ , we can get a set of  $2n - 2$  negative samples.
Figure 4: Pre-training with fine-grained type inference and multi-view masked language modeling.
For an input  $x_{i}^{a}$  with representation  $v_{i}^{a}$  under the view of  $a$ , it has one positive sample  $x_{i}^{b}$  with representation  $v_{i}^{b}$  under the view of  $b$ . It also has a negative sample set  $\mathbf{V}^{-} = \{\pmb{v}_{1}^{-},\dots,\pmb{v}_{2n - 2}^{-}\}$  with size  $2n - 2$ , which consists of two types of negative sample subsets, e.g., intra-view negative sample set  $\mathbf{V}_1^-$  with size  $n - 1$ , where  $\pmb{v}_j^a\in \mathbf{V}_1^-, \forall j\neq i$ , and the inter-view negative sample set  $\mathbf{V}_2^-$  with size  $n - 1$ , where  $\pmb{v}_j^b\in \mathbf{V}_2^-, \forall j\neq i$ . We define the similarity of a pair of samples as the dot product of their representations. Then the loss function for a positive pair  $(x_i^a,x_i^b)$  can be defined as:
$$
l \left(x _ {i} ^ {a}, x _ {i} ^ {b}\right) = - \ln \frac {\exp \left(\boldsymbol {v} _ {i} ^ {a} \cdot \boldsymbol {v} _ {i} ^ {b}\right)}{\exp \left(\boldsymbol {v} _ {i} ^ {a} \cdot \boldsymbol {v} _ {i} ^ {b}\right) + \sum_ {k = 1} ^ {2 n - 2} \exp \left(\boldsymbol {v} _ {i} ^ {a} \cdot \boldsymbol {v} _ {k} ^ {-}\right)}. \tag {1}
$$
We calculate the loss for the same pair twice with order switched, i.e.,  $(x_{i}^{a}, x_{i}^{b})$  is changed to  $(x_{i}^{b}, x_{i}^{a})$  as the dot product with negative samples for  $x_{i}^{a}$  and  $x_{i}^{b}$  are different. Overall, the MVCL loss function is defined as follows:
$$
\mathcal {L} _ {\mathrm {M V C L}} = - \frac {1}{| \mathcal {N} |} \sum_ {i} ^ {| \mathcal {N} |} \left[ l \left(x _ {i} ^ {a}, x _ {i} ^ {b}\right) + l \left(x _ {i} ^ {b}, x _ {i} ^ {a}\right) \right], \tag {2}
$$
where  $\mathcal{N}$  denotes the set of all program samples covering all different views.
3.4 Pre-Training with Type Inference
Figure 4 shows the other two pre-training tasks, including fine-grained type inference and multiview masked language modeling.
Fine-Grained Type Inference. Several previous works (Wang et al., 2021; Wang et al., 2021b) have proven the importance of symbolic properties in programming languages. Two concurrent works, SynCoBERT (Wang et al., 2021) and CodeT5 (Wang et al., 2021b) let the model divide the code token types into identifier or
non-identifier. Inspired by the type checking in compilation process, we propose a fine-grained type inference (FGTI) objective to capture the fine-grained type information of variables (Li et al., 2022; An et al., 2011). First, we parse all source codes into ASTs. Then, we traverse the AST and use the type checker to obtain fine-grained identifier types. We employ BPE tokenizer (Sennrich et al., 2016) to tokenize tokens and let sub-tokens inherit the type information of the token. Finally, we define the loss function as follows:
$$
\mathcal {L} _ {\mathrm {F G T I}} = - \frac {1}{| \mathcal {Z} |} \sum_ {i} ^ {| \mathcal {Z} |} \sum_ {j} ^ {| \mathcal {T} |} Y _ {i j} \log P _ {i j}, \tag {3}
$$
where  $\mathcal{Z}$  denotes the set of all tokens that need to inference types,  $\mathcal{T}$  represents the set of all types contained in the pre-training corpus,  $Y_{ij}$  denotes the label of token  $i$  in type  $j$ , and  $P_{ij}$  denotes the predicted probability of token  $i$  in type  $j$ .
Multi-View Masked Language Modeling. In addition to the multi-view contrastive learning objective and fine-grained type inference objective, we also extend the Masked Language Modeling (MLM) to the multi-view program corpus, named MMLM. Given a data point  $x$ , we randomly select  $15\%$  of tokens in  $x$  and replace them with a special token  $<\text{MASK} >$ , following the same settings in (Devlin et al., 2019). The MMLM objective aims to predict original tokens which are masked out. We calculate the MMLM loss as follows:
$$
\mathcal {L} _ {\mathrm {M M L M}} = - \frac {1}{| \mathcal {M} |} \sum_ {i} ^ {| \mathcal {M} |} \sum_ {j} ^ {| \mathcal {V} |} Y _ {i j} \log P _ {i j}, \tag {4}
$$
where  $\mathcal{M}$  denotes the set of masked tokens,  $\nu$  represents the vocabulary,  $Y_{ij}$  denotes the label of the masked token  $i$  in class  $j$ , and  $P_{ij}$  denotes the predicted probability of token  $i$  in class  $j$ .
3.5 Overall Training Objective
The overall loss function in CODE-MVP is the integration of several components we have defined before.
$$
\mathcal {L} = \mathcal {L} _ {\mathrm {M V C L}} + \mathcal {L} _ {\mathrm {F G T I}} + \mathcal {L} _ {\mathrm {M M L M}} + \lambda \| \Theta \| ^ {2}, \tag {5}
$$
where  $\Theta$  contains all trainable parameters of the model, and  $\lambda$  is the coefficient of  $L_{2}$  regularizer.
Tasks	Datasets	Train	Valid	Test
Natural Language	AdvTest	251K	9.6K	19.2K
Code Retrieval	CosQA	19.6K	0.5K	0.5K
CoNaLa	2.4K	-	0.5K
Code-to-Code Retrieval	Python800	72K	4K	4K
Code Clone Detection	Python800	144K	8K	8K
Code Defect Detection	GREAT	100K	5K	5K
Table 2: Statistics of datasets for downstream tasks.
We conduct experiments to answer the following research questions: (1) How effective is CODE-MVP compared with the state-of-the-art baselines? (2) How do different components and different views affect our CODE-MVP?
4.1 Pre-Training Dataset and Settings
Different programming languages often require different program analyzers. Existing program analysis tools rarely support multiple programming languages and multi-view program transformations. For convenience, we choose Python for our experiments, as it is very popular and used in many projects. We pre-train CODE-MVP on the Python corpus of CodeSearchNet dataset (Husain et al., 2019), which consists of 0.5M bimodal Python functions with their corresponding natural-language comments, as well as 1.1M unimodal Python functions.
CODE-MVP is built on the top of Transformer (Vaswani et al., 2017), and consists of a 12-layer encoder with 768 hidden sizes and 12 attention heads. The pre-training procedure is conducted on 8 NVIDIA V100 GPUs for 600K steps, with each mini-batch containing 128 sequences up to 512 tokens including special tokens. According to the length distribution of samples in the training corpus, we set the lengths of PL/AST/CFG/PT in unpaired data to 512, and set the lengths of NL and PL/AST/CFG/PT in paired data to 96 and 416 respectively. The learning rate of CODE-MVP is set to  $1e-4$  with a linear warm up over the first 30K steps and a linear decay. CODE-MVP is trained with a dropout rate of 0.1 on all layers and attention weights. We initialize the parameters of CODE-MVP by GraphCodeBERT (Guo et al., 2021) and utilize a BPE tokenizer (Sennrich et al., 2016).
4.2 Evaluation Tasks, Datasets and Metrics
We select several program comprehension tasks to evaluate CODE-MVP, including natural language code retrieval, code similarity, and code defect detection. We pre-train CODE-MVP on Python corpus, and choose several public Python datasets to evaluate it, as shown in Table 2.
Natural Language Code Retrieval. This task aims to find the most relevant code snippet from a collection of candidates, given a natural language query. We choose three datasets to evaluate this task, including AdvTest (Lu et al., 2021), CoNaLa (Yin et al., 2018), and CoSQA (Huang et al., 2021). We adopt the Mean Reciprocal Rank (MRR) metric to evaluate the performance of code retrieval. In AdvTest dataset, we set the learning rate as  $5e - 5$ , the batch size as 32, the maximum fine-tuning epoch as 20, the maximum length of both query and code sequence as 256. In CoNaLa and CoSQA datasets, we set the learning rate as  $5e - 5$ , the batch size as 32, the maximum fine-tuning epoch as 30, the maximum length of query and code sequence as 128. In AdvTest and CoSQA datasets, we save the optimal checkpoint on the validation set, and test it on the testing set. In CoNaLa dataset, we report the best results on the testing set.