SlowGuess/ACL_2020_OCR huggingface dataset:
Ruichu Cai $^{1}$ , Zhihao Liang $^{1}$ , Boyan Xu $^{1*}$ , Zijian Li $^{1}$ , Yuexing Hao $^{2}$ , Yao Chen $^{3}$
<sup>1</sup> School of Computer Science, Guangdong University of Technology, China
2 Rutgers University New Brunswick, USA
<sup>3</sup> Advanced Digital Sciences Center, Singapore
cairuichu@gmail.com, zhihaolzh95@gmail.com, hpakyim@gmail.com,
leizigin@gmail.com, yh599@scarletmail.rutgers.edu, yao.chen@adsc-create.edu.sg
Existing leading code comment generation approaches with the structure-to-sequence framework ignores the type information of the interpretation of the code, e.g., operator, string, etc. However, introducing the type information into the existing framework is non-trivial due to the hierarchical dependence among the type information. In order to address the issues above, we propose a Type Auxiliary Guiding encoder-decoder framework for the code comment generation task which considers the source code as an N-ary tree with type information associated with each node. Specifically, our framework is featured with a Type-associated Encoder and a Type-restricted Decoder which enables adaptive summarization of the source code. We further propose a hierarchical reinforcement learning method to resolve the training difficulties of our proposed framework. Extensive evaluations demonstrate the state-of-the-art performance of our framework with both the auto-evaluated metrics and case studies.
The comment for the programming code is critical for software development, which is crucial to the further maintenance of the project codebase with significant improvement of the readability (Aggarwal et al., 2002; Tenny, 1988). Code comment generation aims to automatically transform program code into natural language with the help of deep learning technologies to boost the efficiency of the code development.
Existing leading approaches address the code comment generation task under the structure-to-sequence (Struct2Seq) framework with an encoder-decoder manner by taking advantage of the inherent structural properties of the code. For instance, existing solutions leverage the syntactic structure of abstract syntax trees (AST) or parse trees from
(a) Struct2Seq example
Figure 1: Comment generation frameworks. Different types are denoted as different colors and shapes in (b).
(b) TAG example
source code have shown significant improvement to the quality of the generated comments (Liang and Zhu, 2018; Alon et al., 2018; Hu et al., 2018; Wan et al., 2018); Solutions representing source code as graphs have also shown high-quality comment generation abilities by taking advantage of extracting the structural information of the codes (Xu et al., 2018a,b; Fernandes et al., 2018).
Although promising results were reported, we observe that the information of the node type in the code is not considered in these aforementioned Struct2Seq based solutions. The lack of such essential information lead to the following common limitations: 1) Losing the accuracy for encoding the source code with the same structure but has different types. As shown in Fig. 1(a), a Tree-LSTM (Tai et al., 2015) encoder is illustrated to extract the structural information, the two subtrees of the code 'Select' and 'Compare' in the dashed box have the same structure but different types, with the ignorance of the type information, the traditional encoders illustrate the same set of neural network parameters to encode the tree, which leads to an inaccurate generation of the comment. 2) Losing both the efficiency and accuracy for searching the large vocabulary in the decoding procedure,
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 291-301 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
especially for the out-of-vocabulary (OOV) words that exist in the source code but not in the target dictionary. As shown in the Fig. 1(a), missing the type of 'ACL' node usually results in an unknown word 'UNK' in the generated comments. Thus, the key to tackle these limitations is efficiently utilizing the node type information in the encoder-decoder framework.
To well utilize the type information, we propose a Type Auxiliary Guiding (TAG) encoder-decoder framework. As shown in Fig. 1(b), in the encoding phase, we devise a Type-associated encoder to encode the type information in the encoding of the N-ary tree. In the decoding phase, we facilitate the generation of the comments with the help of type information in a two-stage process naming operation selection and word selection to reduce the searching space for the comment output and avoid the out-of-vocabulary situation. Considering that there is no ground-truth labels for the operation selection results in the two-stage generation process, we further devised a Hierarchical Reinforcement Learning (HRL) method to resolve the training of our framework. Our proposed framework makes the following contributions:
- An adaptive Type-associated encoder which can summarize the information according to the node type;
- A Type-restricted decoder with a two-stage process to reduce the search space for the code comment generation;
- A hierarchical reinforcement learning approach that jointly optimizes the operation selection and word selection stages.
Code comment generation frameworks generate natural language from source code snippets, e.g. SQL, lambda-calculus expression and other programming languages. As a specified natural language generation task, the mainstream approaches could be categorized into textual based method and structure-based method.
The textual-based method is the most straightforward solution which only considers the sequential text information of the source code. For instance, Movshovitz-Attias and Cohen (2013) uses topic models and n-grams to predict comments with given source code snippets; Iyer et al. (2016) presents a language model Code-NN using LSTM networks with attention to generate descriptions
about C# and SQL; Allamanis et al. (2016) predicts summarization of code snippets using a convolutional attention network; Wong and Mooney (2007) presents a learning system to generate sentences from lambda-calculus expressions by inverting semantic parser into statistical machine translation methods.
The structure-based methods take the structure information into consideration and outperform the textual-based methods. Alon et al. (2018) processes a code snippet into the set of compositional paths in its AST and uses attention mechanism to select the relevant paths during the decoding. Hu et al. (2018) presents a Neural Machine Translation based model which takes AST node sequences as input and captures the structure and semantic of Java codes. Wan et al. (2018) combines the syntactic level representation with lexical level representation by adopting a tree-to-sequence (Eriguchi et al., 2016) based model. Xu et al. (2018b) considers a SQL query as a directed graph and adopts a graph-to-sequence model to encode the global structure information.
Copying mechanism is utilized to address the OOV issues in the natural language generation tasks by reusing parts of the inputs instead of selecting words from the target vocabulary. See et al. (2017) presents a hybrid pointer-generator network by introducing pointer network (Vinyals et al., 2015) into a standard sequence-to-sequence (Seq2Seq) model for abstractive text summarization. COPYNET from Gu et al. (2016) incorporates the conventional copying mechanism into Seq2Seq model and selectively copy input segments to the output sequence. In addition, Ling et al. (2016) uses the copying mechanism to copy strings from the code.
Our targeted task is considered as the opposite process of natural language to programming code (NL-to-code) task. So some of the NL-to-code solutions are also taken as our references. Dong and Lapata (2016) distinguishes types of nodes in the logical form by whether nodes have child nodes. Yin and Neubig (2017); Rabinovich et al. (2017); Xu et al. (2018a) take the types of AST nodes into account and generate the corresponding programming codes. Cai et al. (2018) borrows the idea of Automata theory and considers the specific types of SQL grammar in Backus-Naur form (BNF) and generates accurate SQL queries with the help of it.
Inspired by the methods considering the type
Figure 2: TAG Encoder and Decoder framework.
information of the code, our solution differs from the existing method with a Type-associated Encoder that encodes the type information during the substructure summarization and a Type-restricted Decoder that can reduce search space for the code comment generation. In addition, two improvements are developed according to our objectives. First, we design a type-restricted copying mechanism to reduce the difficulty of extracting complex grammar structure from the source code. Second, we use a hierarchical reinforcement learning methods to train the model in our framework to learn to select from either copy or other actions, the details will be presented in Section 3.
We first make the necessary definition and formulation for the input data and the code comment generation problem for our Type Auxiliary Guiding (TAG) encoder-decoder framework.
Definition 1 Token-type-tree. Token-type-tree  $T_{x,\tau}$  represents the source code with the node set  $V$ , which is a rooted  $N$ -ary tree. And  $V = \{v_1, v_2, \dots, v_{|V|}\}$  denotes a partial order nodes set satisfying  $v_1 \preceq v_2 \preceq \dots, \preceq v_{|V|}$ . Let internal node  $v_j = \{x_j, \tau_j\}$ , where  $x_j$  denotes the token sequence and  $\tau_j$  denotes a type from grammar type set  $\mathcal{T}$ .
Token-type-tree can be easily constructed from token information of the original source code and type information of its AST or parse tree. According to Definition 1, we formulate the code comment generation task as follows.
**Formulation 1 Code Comment Generation with Token-type-tree as the Input.** Let  $S$  denote training dataset and labeled sample  $(T_{x,\tau}, \mathbf{y}) \in S$ , where  $T_{x,\tau}$  is the input token-type-tree,  $\mathbf{y} = (y_1, y_2, \dots, y_M)$  is the ground truth comment with  $M$  words. The task of code comment generation is to design a model which takes the unlabeled sample  $T_{x,\tau}$  as input and predicts the output as its comment, denoted as  $\mathbf{y}$ .
Our framework follows the encoder-decoder manner, and consists of the revised two major components, namely the Type-associated Encoder and Type-restricted Decoder. As shown in Fig. 2.
The Type-associated Encoder, as shown in Fig. 2, recursively takes the token-type-tree  $T_{x,\tau}$  as input, and maintains the semantic information of the source code in the hidden states. Instead of using the same parameter sets to learn the whole token-type-tree, Type-associated Encoder utilizes multiple sets of parameters to learn the different type of nodes. The parameters of the cells are adaptively invoked according to the type of the current node during the processing of the input token-type-tree. Such a procedure enables the structured semantic representation to contain the type information of the source code.
The Type-restricted Decoder, as shown in the right part of Figure 2, takes the original toke-type-tree  $T_{x,\tau}$  and its semantic representation from encoder as input and generates the corresponding comment. Different from conventional decoders which generate output only based on the target dictionary, our Type-restricted Decoder considers both
input code to the encoder and target dictionary as the source of output. Attention mechanism is employed to compute an attention vector which is used to generate the output words through a two-stage process: (1) Determine either to copy from the original token-type-tree or to generate from the current hidden state according to the distribution of the operation. (2) If the copying operation is selected, the words are copied from the selected node from the token-type-tree  $T_{x,\tau}$  with restricted types; otherwise, the candidate word will be selected from the target dictionary. The above two-stage process is guided by the type which is extracted from the hidden state of encoder with the help of attention mechanism. Such a process enables adaptive switching between copying and generation processes, and not only reduces the search space of the generation process but also addresses the OOV problem with the copying mechanism.
Although the proposed framework provides an efficient solution with the utilization of the type information in the code, training obstacles are raised accordingly: (1) No training labels are provided for the operation selection stage. (2) There is a mismatch between the evaluation metric and the objective function. Thus, we further devised an HRL method to train our TAG model. In the HRL training, the TAG model feeds back the evaluation metric as the learning reward to train the two-stage sampling process without relying on the ground-truth label of operation selection stage.
4 Type-associated Encoder
The encoder network aims to learn a semantic representation of the input source code. The key challenge is to provide distinct summarization for the sub-trees with the same structure but different semantics. As shown in the Type-associated Encoder in Fig. 1, the blue and red dashed blocks have the same 3-ary substructure. The sub-tree in the blue box shares the same sub-structure with the tree in the red box, which is usually falsely processed by the same cell in a vanilla Tree-LSTM. By introducing the type information, the semantics of the two subtrees are distinguished from each other.
Our proposed Type-associated Encoder is designed as a variant  $N$ -ary Tree-LSTM. Instead of directly inputting type information as features into the encoder for learning, we integrate the type information as the index of the learning parameter sets of the encoder network. More specifically, differ
ent sets of parameters are defined through different types, which provides a more detailed summarization of the input. As is shown in Fig. 1(b), the two sub-trees in our proposed Type-associated Encoder are distinguished by the type information. The tree contains  $N$  ordered child nodes, which are indexed from 1 to  $N$ . For the  $j$ -th node, the hidden state and memory cell of its  $k$ -th child node is denoted as  $\pmb{h}_{jk}$  and  $\pmb{c}_{jk}$ , respectively. In order to effectively capture the type information, we set  $\pmb{W}_{\tau_j}$  and  $\pmb{b}_{\tau_j}$  to be the weight and bias of the  $j$ -th node, and  $\pmb{U}_{\tau_{jk}}$  be the weight of the  $k$ -th child of the  $j$ -th node. The transition equation of the variant  $N$ -ary Tree-LSTM is shown as follow:
$$
\pmb {i} _ {j} = \sigma \left(\pmb {W} _ {\tau_ {j}} ^ {(\pmb {i})} \phi (x _ {j}) + \sum_ {l = 1} ^ {N} \pmb {U} _ {\tau_ {j l}} ^ {(\pmb {i})} \pmb {h} _ {j l} + \pmb {b} _ {\tau_ {j}} ^ {(\pmb {i})}\right), \quad (1)
$$
$$
\pmb {f} _ {j k} = \sigma \left(\pmb {W} _ {\tau_ {j k}} ^ {(\pmb {f})} \phi (x _ {j}) + \sum_ {l = 1} ^ {N} \pmb {U} _ {\tau_ {j l, k}} ^ {(\pmb {f})} \pmb {h} _ {j l} + \pmb {b} _ {\tau_ {j k}} ^ {(\pmb {f})}\right), \quad (2)
$$
$$
\pmb {o} _ {j} = \sigma \left(\pmb {W} _ {\tau_ {j}} ^ {(\pmb {o})} \phi (x _ {j}) + \sum_ {l = 1} ^ {N} \pmb {U} _ {\tau_ {j l}} ^ {(\pmb {o})} h _ {j l} + \pmb {b} _ {\tau_ {j}} ^ {(\pmb {o})}\right), \quad (3)
$$
$$
\boldsymbol {u} _ {j} = \tanh  \left(\boldsymbol {W} _ {\tau_ {j}} ^ {(\boldsymbol {u})} \phi (x _ {j}) + \sum_ {l = 1} ^ {N} \boldsymbol {U} _ {\tau_ {j l}} ^ {(\boldsymbol {u})} h _ {j l} + \boldsymbol {b} _ {\tau_ {j}} ^ {(\boldsymbol {u})}\right), \tag {4}
$$
$$
\boldsymbol {c} _ {j} = \boldsymbol {i} _ {j} \odot \boldsymbol {u} _ {j} + \sum_ {l = 1} ^ {N} \boldsymbol {f} _ {j l} \odot \boldsymbol {c} _ {j l}, \tag {5}
$$
$$
\boldsymbol {h} _ {j} = \boldsymbol {o} _ {j} \odot \tanh  \left(\boldsymbol {c} _ {j}\right), \tag {6}
$$
We employ the forget gate (Tai et al., 2015) for the Tree-LSTM, the parameters for the  $k$ -th child of the  $j$ -th node's is denoted as  $\pmb{f}_{jk}$ .  $\pmb{U}_{\tau_{jl,k}}$  is used to represent the weight of the type for the  $l$ -th child of the  $j$ -th node in the  $k$ -th forget gate. The major difference between our variants and the traditional Tree-LSTM is that the parameter set  $(W_{\tau}, U_{\tau}, b_{\tau})$  are specified for each type  $\tau$ .
5 Type-restricted Decoder
Following with the Type-associated Encoder, we propose a Type-restricted Decoder for the decoding phase, which incorporates the type information into its two-stage generation process. First of all, an attention mechanism is adopted in the decoding phase which takes hidden states from the encoder as input and generates the attention vector. The resulted attention vector is used as input to the following two-stage process, named operation selection stage and word selection stage, respectively. The operation selection stage selects between generation operation and copying operation
for the following word selection stage. If the generation operation is selected, the predicted word will be generated from the targeted dictionary. If the copying operation is selected, then a type-restricted copying mechanism is enabled to restrict the search space by masking down the illegal grammar types. Furthermore, a copying decay strategy is illustrated to solve the issue of repetitively focusing on specific nodes caused by the attention mechanism. The details of each part are given below.
Attention Mechanism: The encoder extracts the semantic representation as the hidden state of the rooted nodes, denoted as  $\pmb{h}_r$ , which are used to initialize the hidden state of the decoder,  $\pmb{z}_0 \gets \pmb{h}_r$ . At time step  $m$ , given output  $y_{m-1}$  and the hidden state of the decoder  $z_{m-1}$  at last time step  $m-1$ , the hidden state  $z_m$  is recursively calculated by the LSTM cells in the decoder,
$$
\boldsymbol {z} _ {m} = L S T M \left(\boldsymbol {z} _ {m - 1}, y _ {m - 1}\right). \tag {7}
$$
The attention vector  $\pmb{q}$  is calculate with:
$$
\alpha_ {m j} = \frac {\exp \left(\boldsymbol {h} _ {j} ^ {\top} \boldsymbol {z} _ {m}\right)}{\sum_ {j = 1} ^ {| V _ {x} |} \exp \left(\boldsymbol {h} _ {j} ^ {\top} \boldsymbol {z} _ {m}\right)},
$$
$$
\widetilde {\boldsymbol {q} _ {m}} = \sum_ {j = 1} ^ {| V _ {x} |} \alpha_ {m j} \boldsymbol {h} _ {j}, \tag {8}
$$
$$
\pmb {q} _ {m} = \tanh \left(\pmb {W} _ {q} \left[ \widetilde {\pmb {q}}, \pmb {z} _ {m} \right]\right),
$$
where  $W_{q}$  is the parameters of the attention mechanism. The attention vector contains the token and type information, which is further facilitated in the following operation selection and word selection stages.
Operation Selection Stage: Operation Selection Stage determines either using the copying operation or the generation operation to select the words based on the attention vector and hidden states from the encoder. Specifically, given the attention vector  $\mathbf{q}_m$  at time step  $m$ , Operation Selection Stage estimates the conditional probabilities as the distribution of the operation  $p(\hat{a}_m|\hat{y}_{<m};T_{x,\tau})$ , where  $\hat{a}_m\in \{0,1\}$  and 0 and 1 represents the copy and the generation operations, respectively. A fully connected layer followed by a softmax is implemented to compute the distribution of the operations.
$$
p (\hat {a} _ {m} | \hat {y} _ {<   m}; T _ {x, \tau}) = s o f t m a x (\boldsymbol {W} _ {s} \boldsymbol {q} _ {m}), \quad (9)
$$
The  $W_{s}$  in the Eq. 9 is the trainable parameters. Since there is no ground-truth label for operation
selection, we employ an HRL method to jointly train the operation selection stage and the following stage, the details are provided in Section 6.
Word Selection Stage: Word Selection Stage also contains two branches. The selection between them is determined by the previous stage. If the generation operation is selected in the Operation Selectoin Stage, the attention vector will be fed into a softmax layer to predict the distribution of the target word, formulated as
$$
p \left(y _ {m} \mid \hat {a} _ {m} = 1, \hat {y} _ {<   m}; T _ {x, \tau}\right) = \text {s o f t m a x} \left(\boldsymbol {W} _ {g} \boldsymbol {q} _ {m}\right), \tag {10}
$$
where  $W_{g}$  is the trainable parameters of the output layer. Otherwise, if the copy operation is selected, we employ the dot-product score function to calculate score vector  $\pmb{s}_m$  of the hidden state of the node and the attention vector. Similarly, score vector  $\pmb{s}_m$  will be fed into a softmax layer to predict the distribution of the input word, noted as:
$$
\pmb {s} _ {m} = \left[ \pmb {h} _ {1}, \pmb {h} _ {2}, \dots , \pmb {h} _ {| V _ {x} |} \right] ^ {\top} \pmb {q} _ {m}
$$
$$
p \left(y _ {m} \mid \hat {a} _ {m} = 0; \hat {y} _ {<   m}; T _ {x, \tau}\right) = \text {s o f t m a x} \left(\boldsymbol {s} _ {m}\right). \tag {11}
$$
One step further, to filter out the illegally copied candidates, we involve a grammar-type based mask vector  $\pmb{d}_m \in \mathbb{R}^{|V_x|}$  at each decoding step  $m$ . Each dimension of  $\pmb{d}_m$  corresponds to each node of the token-type-tree. If the mask of the node in token-type-tree indicates the node should be filtered out, then the corresponding dimension is set as negative infinite. Otherwise, it is set to 0. Thus, the restricted copying stage is formulated as
$$
p \left(y _ {m} \mid \hat {a} _ {m} = 0, \hat {y} _ {<   m}; T _ {x, \tau}\right) = \text {s o f t m a x} \left(\boldsymbol {s} _ {m} + \boldsymbol {d} _ {m}\right). \tag {12}
$$
The word distribution of the two branches is represented with a softmax over input words or target dictionary words in Eq. 10 and Eq. 12. At each time step, the word with the highest probability in the word distribution will be selected.
Copying Decay Strategy: Similar to the conventional copying mechanism, we also use the attention vector as a pointer to guide the copying process. The type-restricted copying mechanism tends to pay more attention to specific nodes, resulting in the ignorance of other available nodes, which makes certain copied tokens repeatedly active in a short distance in a single generated text, lead to a great redundancy of the content.
So we design a Copying Decay Strategy to smoothly penalize certain probabilities of outstand
ingly copied nodes. We define a copy time-based decay rate  $\lambda_{mi}$  for the  $i$ -th tree node  $x_{i}$  in the  $m$ -th decoding step. If one node is copied in time step  $m$ , its decay rate is initialized as 1. In the next time step  $m + 1$ , it is scaled by a coefficient  $\gamma \in (0,1)$ :
$$
\lambda_ {m + 1, i} = \gamma \lambda_ {m, i} \tag {13}
$$
The overall formulation for the Type-restricted Decoder is:
$$
\begin{array}{l} p \left(y _ {m} \mid \hat {a} _ {m} = 0, \hat {y} _ {<   m}; T _ {x, \tau}\right) = \tag {14} \\ \operatorname {s o f t m a x} \left(\boldsymbol {s} _ {m} + \boldsymbol {d} _ {m}\right) \odot (1 - \boldsymbol {\lambda} _ {m}) \\ \end{array}
$$
6 Hierarchical Reinforcement Learning
There remain two challenges to train our proposed framework, which are 1) the lack of ground truth label for the operation selection stage and 2) the mismatch between the evaluation metric and objective function. Although it is possible to train our framework by using the maximum likelihood estimation (MLE) method which constructs pseudo-labels or marginalize all the operations in the operation selection stage (Jia and Liang, 2016; Gu et al., 2016), the loss-evaluation mismatch between MLE loss for training and non-differentiable evaluation metrics for testing lead to inconsistent results (Keneshloo et al., 2019; Ranzato et al., 2015). To address these issues, we propose a Hierarchical Reinforcement Learning method to train the operation selection stage and word selection stage jointly.
We set the objective of the HRL as maximizing the expectation of the reward  $R(\hat{\pmb{y}},\pmb {y})$  between the predicted sequence  $\hat{\pmb{y}}$  and the ground-truth sequence  $\pmb{y}$ , denoted as  $L_{r}$ . It could be formulated as a function of the input tuple  $\{T_{x,\tau},\pmb {y}\}$  as,