[论文解读] Code Prediction by Feeding Trees to Transformers

Seohyun Kim, Jinman Zhao+2 more|arXiv (Cornell University)|Mar 30, 2020

Software Engineering Research参考文献 66被引用 32

一句话总结

本论文通过利用 AST 结构的 Transformer 模型提升下一个 token 的代码预测，提出 SeqTrans、PathTrans、TravTrans，以及 TravTrans+，在 Python 代码数据集上显著提升比前方法的准确性。TravTrans 取得最强结果，在 py150 上实现最高约 58% 的 mean reciprocal rank (MRR)，比前沿方法高出 14–18 个百分点。

ABSTRACT

We advance the state-of-the-art in the accuracy of code prediction (next token prediction) used in autocomplete systems. First, we report that using the recently proposed Transformer architecture even out-of-the-box outperforms previous neural and non-neural systems for code prediction. We then show that by making the Transformer architecture aware of the syntactic structure of code, we further increase the margin by which a Transformer-based system outperforms previous systems. With this, it outperforms the accuracy of an RNN-based system (similar to Hellendoorn et al. 2018) by 18.3%, the Deep3 system (Raychev et al 2016) by 14.1%, and an adaptation of Code2Seq (Alon et al., 2018) for code prediction by 14.4%. We present in the paper several ways of communicating the code structure to the Transformer, which is fundamentally built for processing sequence data. We provide a comprehensive experimental evaluation of our proposal, along with alternative design choices, on a standard Python dataset, as well as on a Facebook internal Python corpus. Our code and data preparation pipeline will be available in open source.

研究动机与目标

• 推动并改进在 IDE 自动完成中对下一个 token 的代码预测。
• 评估将 AST 结构集成到 Transformer 模型中对预测准确性的影响。
• 将基于 Transformer 的、具树结构感知的模型与先前的非 Transformer 与基于 AST 的方法进行比较。
• 在标准与内部 Python 数据集上进行评估以展示泛化能力。

提出的方法

• 将部分程序表示为记号序列（SeqTrans）或基于 AST 的结构（PathTrans、TravTrans）以输入到 Transformer。
• 引入 PathTrans：将从叶节点到根的根路径表示进行嵌入，并与叶节点嵶入结合以输入到 Transformer。
• 引入 TravTrans：将 AST 的 DFS 顺序记号输入到 Transformer 以融合结构信息。
• 引入 TravTrans+：通过 AttnTreeRel 中的树关系矩阵 R，在注意力中增加显式的树关系路径以增强 TravTrans。
• 在 py150 和 Facebook 内部 Python 数据集上，与强基线（SeqRNN、Deep3、Code2Seq）进行比较。
• 使用 mean reciprocal rank (MRR) 作为对叶节点 token 的下一个 token 预测的评估指标。

实验结果

研究问题

• RQ1 Transformer 是否在 Python 代码预测的下一个 token 上优于基于 RNN 的模型？
• RQ2 利用 AST 结构的 Transformer 模型是否优于以前的基于 AST 的代码预测方法？
• RQ3 SeqTrans、PathTrans、TravTrans 与 TravTrans+ 在不同数据集上的准确性对比如何？
• RQ4 树关系增强（TravTrans+）是否相对于 TravTrans 具有额外的收益？
• RQ5 显著性分析可提供哪些关于 TravTrans 在预测时如何关注树结构的洞见？

主要发现

• TravTrans 实现最佳整体性能，在 py150 上达到 58.0% 的 MRR，并显著高于此前的方法。
• SeqTrans（基线 Transformer）在不同数据集和任务中达到 54.9–50.1% 的 MRR，优于 SeqRNN（36.6–23.8% MRR）。
• PathTrans 与 TravTrans 相对于非 Transformer 的 AST 基线（Deep3、Code2Seq）有所提升，TravTrans 在 py150 上相比 Deep3 和 Code2Seq 的 MRR 提升了 14.1–14.4 个百分点。
• 在 Facebook 内部的 Python 数据集上，基于 Transformer 的模型对基线显示出类似的相对优势。
• 显著性分析表明 TravTrans 在预测 token 时关注最相关的树区域，特别是聚焦于父节点。