Rethinking Positional Encoding in Tree Transformer for Code Representation (2022)

Abstract: Transformers are now widely used in code representation, and several recent works further develop tree Transformers to capture the syntactic structure in source code. Specifically, novel tree positional encodings have been proposed to incorporate inductive bias into Transformer.In this work, we propose a novel tree Transformer encoding node positions based on our new description method for tree structures.Technically, local and global soft bias shown in previous works is both introduced as positional encodings of our Transformer model.Our model finally outperforms strong baselines on code summarization and completion tasks across two languages, demonstrating our model’s effectiveness.Besides, extensive experiments and ablation study shows that combining both local and global paradigms is still helpful in improving model performance. We release our code at https://github.com/AwdHanPeng/TreeTransformer .

## introduction machine learning for source code aims to learn the semantic embedding of programs. due to the format similarity between code and text @xcite , transformers @xcite are now widely used in code representation @xcite @xcite . unlike natural language, source code is more logical and has rich structures such as abstract syntax trees (ast). therefore, one research topic of code intelligence is representing the syntax tree of code. several recent works proposed novel tree-based transformers by defining the position of each node to handle tree structure @xcite . in this work, we pursue the research line of tree transformer for learning code ast. in transformer, positional encoding is crucial to exploit potential structures of data (such as code * corresponding authors asts or graphs) as other components are entirely position-invariant. recently, a growing research trend is adapting transformer to more complex structured data than plain text by modifying positional encoding, not only in language processing @xcite but also in graph representation learning field @xcite . as for the structure-aware transformers, the key step of encoding positions is to find a proper description for input structure, which means abstracting the physical structure of data into a suitable mathematical form. for example, the vanilla transformer @xcite regards the potential text order in natural languages as the arrangement of natural numbers, while the graph transformers @xcite treat the positional relationship between graph nodes as the adjacent matrix or laplacian further. intuitively, a good description should be information lossless from which the whole structure of the original data could be precisely reconstructed. recently, several tree-based transformers incorporated with advanced positional encoding are presented to process the code syntax tree. @xcite represented the position of each node using sibling orders of all nodes that existed in its path to the root, while @xcite defined the relative distance between two nodes as the traversing up and down steps along the shortest path connecting them. however, some tree structures are overlooked by these previous approaches: the first method assumes the regular tree with a fixed number of node children, and the second ignores the sibling feature in traversing. in this paper, we present a new description method for node positions from which the corresponding tree can be rebuilt without ambiguity. specifically, the position of each node is recursively described as a list including multiple 2d coordinates. the coordinate list of each node first inherits from its parent and then includes a new 2d coordinate of itself, where the first dimension indicates its sibling order and the second is the total child number of its parent. we incorporate the proposed description method into transformer, powering it to capture tree structures. technically, a growing popular approach is to encode structure as soft inductive bias in positional encodings of transformer, in which attention between all nodes is allowed rather than the strict aspect of message passing. to the best of our knowledge, the soft bias in previous works is usually introduced either in local or global. the local methods integrate the structure relation as one-hop edges only for two adjacent nodes @xcite , so each node knows its multi-hop subgraph only by stacking model layers. in global approaches, inductive bias is injected into the attention between any nodes regardless of whether adjacent in trees or graphs @xcite , in which structures can percolate fully across graphs in a single layer @xcite . the local and global methods show expressiveness in previous works, but the relationship between them is still not completely studied to the best of our knowledge. in this paper, we propose a new tree transformer that integrates our tree description in local and global, exploring the interaction between local and global bias. our model finally outperforms solid baselines and obtains stateof-the-art in code summarization and completion tasks across two different language datasets. besides, the ablation results show that both global and local methods are powerful, and the combination improves model performance further. the contributions of this paper are summarized as follows: ## related work representation learning for source code the availability of big code shows opportunities for representation learning of programs. traditionally, code intelligence designers have relied predominantly on structure or context. early research works relied on raw text data for code snippets @xcite @xcite , mainly focusing on context and struggle to capture code structure. after that, a growing active research topic is to represent the syntax tree structure of code. @xcite proposed tree-based convolutional neural networks and @xcite treated codes as weighted combination of pairwise paths in ast. @xcite proposed a custom positional encoding to extend transformers to treestructured data. @xcite defines the relative distance on the tree as the shorted path between nodes consisting of up and down steps. our work pursues the research line to model code trees, powering transformer to learning ast by integrating tree positional encodings. besides, several works also explored leveraging different code representations jointly, including context, ast structure and other code graphs. @xcite proposed ggnn to represent program graphs consisting of ast with control-flow and data-flow. @xcite @xcite proposed to learn structure and context together by introducing bias in the self-attention of code context with the underlying tree structure. ## approach in this section, we present our tree transformer in two parts. we first show the novel two-dimensional description for trees, by which each node's position is converted as a coordinate list. after that, we embed position from the description for each node and then integrate position encoding in the selfattention of transformer in local and global. ## experiment setup we focus on two tasks of code representation learning: code summarization and completion. code summarization is one of the most popular tasks in which the function name is predicted given a function body. the method body typically forms complete logical units, and the function name tends to be precisely descriptive. therefore, code summarization is widely used as benchmark by several previous works @xcite @xcite . the code completion is another useful benchmark, and we mainly focus on the completion for each tree node @xcite @xcite . in this task, ast is linearized in depth-first order for all baselines. each node is then predicted given the partial tree built on all the previous nodes in depth-first order. in both summarization and completion tasks, we use the python150k and javascript150k datasets foot_0 preprocessed by @xcite . both datasets consist of program files from github foot_1 and are widely used to evaluate code represen-tation models. to avoid biased results, duplicate files are removed by the duplication list provided by @xcite , and identical code files and functions are further filtered out by @xcite . both datasets are split into training/validation/testing sets with 60%/6.7%/33.3% based on github usernames. in this paper, we mainly compare our model with tree-based transformer models. before feeding trees into all compared transformers, we linearize asts in depth-first order and convert them as node sequences. note that each node in ast is associated with a type, but not all nodes have values (such as non-leaf nodes in trees). therefore, a @xmath0. the type and value embeddings are added as input node features fed into models. for all compared models in both tasks, node tokens are not split into subtokens for computational efficiency following @xcite . ## conclusion we propose a new tree transformer encoding position for each node based on our novel twodimensional description of tree structures. technically, our model introduces soft bias as the positional encoding of transformer in global and local ways. our model finally outperforms strong baselines on code summarization and completion tasks across two different languages, highlighting the effectiveness of our approach.