[论文解读] Learning Program Representations with a Tree-Structured Transformer

Wenhan Wang, Kechi Zhang+5 more|arXiv (Cornell University)|Aug 18, 2022

Software Engineering Research被引用 6

一句话总结

该论文提出 Tree-Transformer，一种新颖的树状结构 Transformer 模型，通过利用多头注意力机制捕捉父子节点与兄弟节点之间的依赖关系，并结合双向（自底向上与自顶向下）信息传播，学习程序表示。该方法在程序分类、错误定位和类型推断任务中显著优于现有的基于树和基于图的方法，有效建模长距离依赖关系并学习丰富的节点级表示。

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.

研究动机与目标

• 为解决现有基于树和基于图的模型在捕捉长距离依赖关系以及学习代码有效节点级表示方面的局限性。
• 克服传统递归神经网络中单向信息流的问题，该问题阻碍了叶节点和低层节点的上下文学习。
• 保留并利用程序的有序结构，特别是兄弟节点的顺序，该特性常被如 GNN 等排列不变模型所忽略。
• 保持原始语法树结构，无需复杂的预处理或图增强，降低开销并保持结构保真度。
• 开发一种基于 Transformer 的架构，以有效建模分层且有序的程序结构，提升程序表示学习效果。

提出的方法

• Tree-Transformer 采用专为树结构设计的多头自注意力机制，显式建模父子节点对与兄弟节点对之间的依赖关系。
• 引入一种双向信息传播策略，使用两个并行的 Transformer 单元：一个自底向上单元从叶节点向根节点聚合上下文信息，一个自顶向下单元将根节点的上下文信息分发至其后代节点。
• 通过位置编码整合，捕捉兄弟节点的顺序，使模型能够区分树中不同结构位置。
• 通过注意力融合递归聚合子节点表示来学习节点表示，使上下文信息能够在树中流动。
• 该架构具有递归性和树状结构，支持从抽象语法树（AST）端到端学习向量表示。
• 模型通过标准监督目标在下游任务（如程序分类、错误定位和类型推断）上进行训练。

Figure 1: The overall pipeline of Tree-Transformer.

实验结果

研究问题

• RQ1 基于 Transformer 的架构能否有效建模程序语法树中的长距离依赖关系，特别是跨远距离节点的依赖？
• RQ2 与单向方法相比，双向信息传播（自底向上与自顶向下）是否能提升节点级表示学习效果？
• RQ3 通过位置编码引入兄弟节点顺序，是否能增强模型对代码语法结构差异的区分能力？
• RQ4 Tree-Transformer 在下游程序理解任务中与最先进基于树和基于图的模型相比表现如何？
• RQ5 该模型是否能在无需复杂图构建或结构预处理的情况下保持性能？

主要发现

• Tree-Transformer 在程序分类任务上达到最先进性能，在多个基准测试中优于现有的基于树和基于图的方法。
• 该模型显著提升了节点级预测任务（如错误操作符定位与修复）的性能，表明其在节点级别具有优越的表示学习能力。
• 双向传播机制有效实现了长距离依赖建模，这在需要跨远距离代码元素进行上下文推理的任务中表现尤为明显。
• 通过位置编码整合兄弟节点顺序，在对语法顺序敏感的任务中带来了可测量的性能提升。
• Tree-Transformer 在无需额外图构建或预处理的情况下保持了强大性能，保留了原始 AST 结构并降低了计算开销。
• 消融实验确认，双向传播与多头注意力组件对模型性能均至关重要。

Figure 2: An example of existing tree-structured and graph neural networks for learning on abstract syntax trees.