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: This work continues on earlier advances in how positional and relational encoding is used by the self attention primitive of Transformers, in particular in the context of source-code understanding. Prior work has biased self attention to encode relative relational information between inputs (e.g., edges or distance on some graph), and to separate the self-attention term due to position from that due to input.  For the former, TPTrans runs the path between any two tokens (specifically, the path in the code Parse Tree between two leaves -- the authors call this the "relative path") through an RNN, and uses the resulting embedding to bias the self attention between those tokens.  For the latter, TPTrans-$\alpha$ runs the path from the parse-tree root to a token through an RNN (the "absolute path"), and uses that to construct a separate positional self-attention term.  The results show that both TPTrans and TPTrans-$\alpha$ outperform prior work (most notably, Code Transformer), but TPTrans seems to dominate TPTrans-$\alpha$ with a smaller set of trainable parameters.