[论文解读] Code Completion by Modeling Flattened Abstract Syntax Trees as Graphs

Yanlin Wang, Hui Li|arXiv (Cornell University)|Mar 17, 2021

Software Engineering Research被引用 5

一句话总结

该论文提出CCAG，一种新颖的代码补全模型，通过将扁平化的抽象语法树（AST）表示为图结构，以捕捉代码中的顺序性、重复性和结构性模式。通过使用带有多任务学习和基于不确定性的损失加权机制的AST图注意力模块，CCAG在代码补全基准测试中达到最先进性能，在值预测和类型预测任务中均优于现有方法。

ABSTRACT

Code completion has become an essential component of integrated development environments. Contemporary code completion methods rely on the abstract syntax tree (AST) to generate syntactically correct code. However, they cannot fully capture the sequential and repetitive patterns of writing code and the structural information of the AST. To alleviate these problems, we propose a new code completion approach named CCAG, which models the flattened sequence of a partial AST as an AST graph. CCAG uses our proposed AST Graph Attention Block to capture different dependencies in the AST graph for representation learning in code completion. The sub-tasks of code completion are optimized via multi-task learning in CCAG, and the task balance is automatically achieved using uncertainty without the need to tune task weights. The experimental results show that CCAG has superior performance than state-of-the-art approaches and it is able to provide intelligent code completion.

研究动机与目标

• 解决现有代码补全方法在捕捉AST中顺序性、重复性和结构性模式方面的局限性。
• 通过将扁平化的AST表示为图结构以保留顺序和结构依赖关系，从而提升代码补全性能。
• 在无需手动调整超参数的情况下，实现多任务学习中值预测与类型预测任务的自动平衡。
• 通过专用的图注意力机制捕捉多样化依赖关系，增强表征学习能力。

提出的方法

• CCAG使用前序深度优先遍历将部分AST表示为图结构，以保留顺序和结构信息。
• 提出AST图注意力模块（ASTGab），其采用三个不同的注意力层，以捕捉AST节点间的不同类型依赖关系。
• 在ASTGab中引入残差连接，以支持深层堆叠并提升表征学习效果。
• 模型采用多任务学习进行值预测和类型预测，结合基于不确定性的损失加权机制，以自动平衡两项任务。
• 基于不确定性的损失公式使模型能够通过梯度下降端到端学习任务权重，避免手动调参。
• 联合损失函数将值预测和类型预测损失相结合，引入可学习的温度参数，并通过正则化防止过拟合。

实验结果

研究问题

• RQ1 将扁平化的AST建模为图结构是否能通过捕捉顺序和结构性模式来提升代码补全性能？
• RQ2 基于图的注意力机制在AST中学习多样化依赖关系方面，对代码补全的效能如何？
• RQ3 在多任务学习中，基于不确定性的损失加权是否优于手动任务加权？
• RQ4 代码中的重复性模式在多大程度上影响基于AST的代码补全模型的有效性？
• RQ5 统一的图表示是否能在代码补全任务中超越基于序列或纯图结构的AST建模方法？

主要发现

• CCAG在基准代码补全数据集上达到最先进性能，超越现有SOTA方法。
• 通过利用图结构AST表示，模型显著提升了值预测和类型预测的准确率。
• 基于不确定性的损失加权机制能有效平衡多任务学习，且无需手动调整超参数。
• 结合残差连接的AST图注意力模块可实现更深、更有效的表征学习，适用于代码补全任务。
• 实证结果证实，捕捉代码中的顺序性和重复性模式可增强模型泛化能力与预测质量。