Learning Program Representations with a Tree-Structured Transformer

TL;DR

This paper introduces Tree-Transformer, a recursive neural network with multi-head attention and bi-directional propagation, effectively capturing long-range dependencies in program trees for better code understanding and prediction tasks.

Contribution

The paper presents a novel Tree-Transformer model that improves program representation learning by modeling sibling and parent-children dependencies with attention and bidirectional information flow.

Findings

Outperforms existing tree-based approaches in prediction tasks

Effectively captures long-range dependencies in program trees

Enhances node-level and tree-level program understanding

Abstract

Learning vector representations for programs is a critical step in applying deep learning techniques for program understanding tasks. Various neural network models are proposed to learn from tree-structured program representations, e.g., abstract syntax tree (AST) and concrete syntax tree (CST). However, most neural architectures either fail to capture long-range dependencies which are ubiquitous in programs, or cannot learn effective representations for syntax tree nodes, making them incapable of performing the node-level prediction tasks, e.g., bug localization. In this paper, we propose Tree-Transformer, a novel recursive tree-structured neural network to learn the vector representations for source codes. We propose a multi-head attention mechanism to model the dependency between siblings and parent-children node pairs. Moreover, we propose a bi-directional propagation strategy to allow node information passing in two directions, bottom-up and top-down along trees. In this way, Tree-Transformer can learn the information of the node features as well as the global contextual information. The extensive experimental results show that our Tree-Transformer significantly outperforms the existing tree-based and graph-based program representation learning approaches in both the tree-level and node-level prediction tasks.