Integrating Tree Path in Transformer for Code Representation

1 Introduction . Machine learning for source code aims at building models that learn semantic embedding of programs . The initial representation of source code relied on sequential models adopted from natural language processing , such as n-gram language model [ Hindle et al. , 2016 , Hellendoorn and Devanbu , 2017 ] , Recurrent Neural Networks ( RNNs ) [ Wei et al. , 2019 ] and so on . However , source code is more logical than natural languages , rich in structured information such as Abstract Syntax Tree ( AST ) . Therefore , these previous works struggle to capture the structural complexity of source code . Some works leverage the program AST to model source code structure and linearize the graph by traversing it . Code2seq [ Alon et al. , 2018 ] represents source code as a set of pairwise paths over AST where each path is compressed to a vector using LSTMs [ Hochreiter and Schmidhuber , 1997 ] . Code2seq obtains state-of-the-art for code representation using only AST information , demonstrating the effectiveness of path encoding . [ Kim et al. , 2020 ] leveraged another kind of path from the leaf node to the root of AST by traversing up its ancestors and then coupled the representation of paths with the source code . On the other hand , several works also leverage structured graph neural networks for modelling source code directly . [ Allamanis et al. , 2017 ] used Gated Graph Neural Network ( GGNN ) to embed the semantically meaningful relationships in the code . [ Zhou et al. , 2019 ] proposed a novel graph neural network with composite programming representation for vulnerability identification . However , GNNs ∗Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . typically rely on synchronous message passing , requiring message passing iterations to aggregate information [ Allamanis et al. , 2017 , Fernandes et al. , 2018 ] . Besides , primarily for computational reasons , GNNs for programs usually compute few message-passing iterations . For these reasons , the representation of GNNs tends to be local and struggle to leverage long-range interactions . On the contrary , Transformer [ Vaswani et al. , 2017 ] allows global information flow at each step but lacks the ability to model the structural complexity of source code . Recently , some works have explored introducing structural inductive bias in Transformer to access the global structural representation of source codes . [ Hellendoorn et al. , 2019 ] proposed Graph Relation Embedding Attention Transformer ( GREAT ) , which extends [ Shaw et al. , 2018 ] and biases Transformer with relational information from graph edge types . After that , [ Zügner et al. , 2021 ] proposed Code Transformer based on XLNet [ Yang et al. , 2019 ] , which computes pairwise distances on AST and integrate multiple relations into the attention module . In this paper , we pursue the research line of combining Transformer with additional structure information of source code . Our first starting point is the structural model of Code2seq [ Alon et al. , 2018 ] , which obtains state-of-the-art of code summarization using only pairwise path information in AST . However , Code2seq lacks the modelling of context , which leads us to explore combining path representation with source code context . The other starting point is Code Transformer [ Zügner et al. , 2021 ] , which counts node distances cross AST to capture source code structure . However , the different nodes combinations of paths contain plenty of structure information , which is overlooked by only encoding distances . To overcome the drawback shown in previous works , we adopt the ideas of encoding path to represent source code and integrate them to bias the attention module of Transformer . Specifically , we encode the path between source code tokens across AST and the path from leaf node to tree root for each token . Then we draw inspiration from both relative and absolute positional encoding methods in the NLP field and modify them to integrate path encodings into Transformer . These paths introduce inductive bias into the attention module of Transformer , powering it to know the structure of source codes . For clarity , we name the first path as the relative path and the last one as the absolute path ( see Fig1 ) , since much similarities between path and position on relative or absolute encoding exist . Intuitively , the relative path represents the structural relationship between tokens and shows specific patterns of code block such as For loop , etc . Meanwhile , the absolute path complements the structural information for each token and reveals the program behaviour on it , such as Assignment and Call , etc . Both two kinds of paths contain plenty of structural information of code syntax . Actually , lots of works have introduced these paths into many novel models [ Kim et al. , 2020 , Alon et al. , 2018 , 2020 ] . However , to our best knowledge , the feature relationship between these two paths for code representation learning is still not completely studied . In this paper , we integrate these paths into the unified Transformer framework to analyze each effectiveness and their relationship . The detailed empirical study also leads to our novel code representation model TPTrans , which means encoding the Tree Path into the Transformer . We show the effectiveness of our approaches on the code summarization task across four different languages datasets , in which the model predicts a method ’ s name based on its body . Our model finally outperforms solid baselines and obtains state-of-the-art on most datasets . The contributions of this paper are summarized as follows : 1 . We propose TPTRANS which integrates path encoding in Transformer , powers it to know the structure of source code , and our model significantly outperforms existing baselines . 2 . We present the empirical study to investigate the relationship between two kinds of path encoding methods proposed in previous works , shedding light on this line of future work . 2 Related work . Representation learning for source code The growing availability of open-source code creates opportunities for machine learning from big code . Early research works learned language model on raw text data for code [ Dam et al. , 2016 , Wang et al. , 2016 , Allamanis et al. , 2016 , Iyer et al. , 2016 ] , providing evidence for the naturalness assumption [ Hindle et al. , 2016 ] . These models treat codes as sequences rather than trees and struggle to model the structural complexity of source code . After that , some works leveraged structural information to model source code . [ Mou et al. , 2016 ] proposed tree-based convolutional neural networks to represent source code . [ Alon et al. , 2018 , 2019 ] treated source code as weighted combination of pairwise paths in abstract syntax trees . Several works also use structured graphical models to represent structural aspects of source code directly . For example , [ Allamanis et al. , 2017 ] proposed GGNN to represent program graphs consisting of AST together with control-flow and data-flow edges . After that , [ Hellendoorn et al. , 2019 , Zügner et al. , 2021 ] proposed to bias the self-attention computation of Transformer given the underlying graph structure . In this paper , we explore integrating path representation of syntax tree into Transformer , powering it to capture the structure of code snippets . Structural representation of languages It has been proved that modelling the language structure helps to represent the compositional aspects of natural language [ Socher et al. , 2011 ] and improves the generalization of models [ Chen et al. , 2016 ] . Several works [ Shen et al. , 2018 , Shiv and Quirk , 2019 ] proposed new architectures for modeling tree structures , including in the domain of natural languages processing [ Wang et al. , 2019 ] and code representation learning [ Hellendoorn et al. , 2019 , Zügner et al. , 2021 ] . We pursue this research line and show the benefits of integrating tree structure as inductive bias . Positional encoding for Transformer The absolute positional encoding used in Transformer has been proved inefficient to model the order information of sequence . [ Shaw et al. , 2018 ] firstly proposed to leverage relative positional encoding to Transformer . After that , Transformer-XL [ Dai et al. , 2019 ] re-parameterized the self attention of Transformer to integrate relative positional encoding and T5 [ Raffel et al. , 2019 ] simplified the vector representation of relative positions in [ Shaw et al. , 2018 ] to scalars . Recently , several works also tried to enhance positional encodings in Transformer to integrate structural inductive bias . For example , [ Shiv and Quirk , 2019 ] extended the sequential positional encoding to tree-based positional encoding in Transformer , and [ Zügner et al. , 2021 ] modified the relative positional encoding of Transformer-XL to incorporate multiple relations in tree structural . In this paper , we adopt the idea to modify positional encodings of Transformer and introduce structural bias of path information in the syntax tree for self-attention computing . 3 Approaches . Self-attention is one of the key components of Transformer . The attention module can be formulated as querying a dictionary with key-value pairs . The matrix form equation for attention with a single head is Attention ( Q , K , V ) = softmax ( QKT√ d ) V , ( 1 ) where d is the dimension of the hidden representations , and Q ( Query ) , K ( Key ) , V ( Value ) are specified as the hidden representations of the previous layer . For simplicity , we omit the index of layer , denote x = ( x1 , x2 · · · , xn ) and z = ( z1 , z2 · · · , zn ) as the input and output of self-attention module in the same layer respectively , where n is the sequence length . Then Eq.1 can be rewritten as zi = n∑ j=1 exp ( αij ) ∑n j′=1 exp ( αij′ ) ( xjW V ) , where αij = 1√ d ( xiW Q ) ( xjW K ) T , ( 2 ) where WQ ∈ Rd×dk , WK ∈ Rd×dk , WV ∈ Rd×dv and we set dk = dv = d. Eq.2 is oblivious to structured input since it views input as unordered vector sets . In the NLP field , several works [ Shaw et al. , 2018 , Ke et al. , 2020 , Raffel et al. , 2019 , Dai et al. , 2019 ] propose different ways to bias Transformer towards sequential inputs by incorporating positional information into the self-attention module . Generally , there are two categories of positional encoding methods for self-attention : absolute positional encoding and relative positional encoding . The original Transformer [ Vaswani et al. , 2017 ] use absolute positional encoding to represent positions . Given a sentence , a real-valued vector pi ∈ Rd will be added to the word embedding wi at position i , and the sum vector will be used as the input to Transformer . [ Ke et al. , 2020 ] further proposed TUPE to untie the correlations between absolute positions and words and directly model the relationship between a pair of words or positions using different projection matrices , that is αij = 1√ 2d ( xiW Q ) ( xjW K ) T + 1√ 2d ( piU Q ) ( pjU K ) T , ( 3 ) where UQ , UK ∈ Rd×d are projection matrices for positional embedding and scaling term √ 2d is used to retain the magnitude of αij . Using different pi helps Transformer distinguish the words at different positions i . However , the absolute positional encoding is not effective in capturing the relative word orders . Therefore , [ Shaw et al. , 2018 ] proposed the relative positional encoding , that is zi = n∑ j=1 exp ( αij ) ∑n j′=1 exp ( αij′ ) ( xjW V + aVij ) , where αij = 1√ d ( xiW Q ) ( xjW K + aKij ) T ( 4 ) where aVij , a K ij ∈ Rd are learnable parameters and can be viewed as the relative position embedding between position i and j. Eq.4 extends self-attention to consider the pairwise positional relationship and models it as a directed edge between head and tail word entities [ Shaw et al. , 2018 ] , which means the tail entity should be close to the embedding of the head entity plus vector depending on the relationship [ Bordes et al. , 2013 ] . We draw inspiration from the ideas of encoding absolute and relative positional information and modify Eq.3 and Eq.4 to incorporate path information in the syntax tree of source code .

Summary: The paper proposes an Transformer-based architecture for embedding source code snippets that integrates encoding of absolute and relative AST paths for all code tokens into self-attention. Following up on prior work in relative position encoding, they map AST paths into the position encoding framework, modify the self-attention computation accordingly, and evaluate on method summarization. The approach is compared to recent Transformer-based baselines on the standard method summarization benchmark, as well as to code2seq, which pioneered path-only source code encoding. The results improve performance significantly on almost all benchmarks.