[论文解读] TreeCaps: Tree-Based Capsule Networks for Source Code Processing

Nghi D. Q. Bui, Yijun Yu+1 more|arXiv (Cornell University)|Sep 5, 2020

Software Engineering Research参考文献 46被引用 7

一句话总结

TreeCaps 提出了一种新颖的胶囊网络架构，将基于树的卷积神经网络（TBCNN）与胶囊路由相结合，直接从抽象语法树（AST）学习源代码表示，避免了噪声较大的基于图的语义分析。该模型在 Java 和 C/C++ 程序的代码功能分类与函数名预测任务中，实现了最先进的准确率与鲁棒性，优于现有模型如 Code2vec、GGNN 和 GREAT。

ABSTRACT

Recently program learning techniques have been proposed to process source code based on syntactical structures (e.g., Abstract Syntax Trees) and/or semantic information (e.g., Dependency Graphs). Although graphs may be better at capturing various viewpoints of code semantics than trees, constructing graph inputs from code needs static code semantic analysis that may not be accurate and introduces noise during learning. Although syntax trees are precisely defined according to the language grammar and easier to construct and process than graphs, previous tree-based learning techniques have not been able to learn semantic information from trees to achieve better accuracy than graph-based techniques. We propose a new learning technique, named TreeCaps, by fusing together capsule networks with tree-based convolutional neural networks, to achieve learning accuracy higher than existing graph-based techniques while it is based only on trees. TreeCaps introduces novel variable-to-static routing algorithms into the capsule networks to compensate for the loss of previous routing algorithms. Aside from accuracy, we also find that TreeCaps is the most robust to withstand those semantic-preserving program transformations that change code syntax without modifying the semantics. Evaluated on a large number of Java and C/C++ programs, TreeCaps models outperform prior deep learning models of program source code, in terms of both accuracy and robustness for program comprehension tasks such as code functionality classification and function name prediction

研究动机与目标

• 开发一种用于源代码理解的深度学习模型，避免基于图的语义分析带来的噪声与复杂性。
• 通过利用抽象语法树（AST）的精确结构，提升程序理解任务中的学习准确率与鲁棒性。
• 针对将胶囊网络应用于可变大小 AST 时胶囊数量动态变化的挑战，提出新型路由机制。
• 评估模型在语义保持的代码变换下的性能与鲁棒性，此类变换模拟对抗性攻击。

提出的方法

• 将基于树的卷积神经网络（TBCNN）适配用于从 AST 中提取局部节点特征，生成可变大小的胶囊输出。
• 引入主可变胶囊（PVC）层，根据输入 AST 的结构动态生成胶囊数量。
• 提出两种路由算法：(1) 共享权重的动态路由（DRSW），以及 (2) 一种新型的可变到静态（VTS）路由，通过选择激活度最高的胶囊将它们映射到固定数量的二级胶囊。
• 从二级胶囊（SC）层到具有固定胶囊数的最终代码胶囊（CC）层应用动态路由，该胶囊数根据分类任务定制。
• 在 DRSW 中使用共享变换矩阵，在 VTS 中采用选择性注意力机制，以提升路由效率与表征学习能力。
• 采用基于边缘的损失函数，实现代码分类与函数名预测任务的端到端训练。

实验结果

研究问题

• RQ1 胶囊网络架构是否能在不依赖基于图的语义依赖关系的情况下，有效从 AST 中学习？
• RQ2 所提出的 VTS 路由机制在收敛速度与准确率方面相较于标准动态路由表现如何？
• RQ3 TreeCaps 在语义保持的代码变换下，其预测一致性相较于现有模型保持程度如何？
• RQ4 与 SOTA 模型如 Code2vec、GGNN 和 GREAT 相比，TreeCaps 是否在代码功能分类与函数名预测任务中表现更优？
• RQ5 模型的鲁棒性与其在语法改变但语义等价的代码上的泛化能力之间是否存在相关性？

主要发现

• TreeCaps 在代码功能分类与函数名预测任务中，准确率与 F1 分数均高于 Code2vec、ASTNN、TBCNN、GGNN、GREAT 与 GNN-FiLM。
• VTS 路由算法收敛速度更快，验证准确率更高，可能归因于参数更少且路由过程更聚焦。
• 在语义保持的变换下，TreeCaps-VTS 的预测变化百分比（PPC）最低——OJ 数据集仅为 12.3%，表明其为测试中鲁棒性最强的模型。
• 该模型在 Java 与 C/C++ 的多种代码库中均保持高性能，展现出跨编程语言的强大泛化能力。
• 尽管仅使用 AST，TreeCaps 在 F1 分数与训练时间方面均优于 GREAT，表明其效率与表征学习能力更优。
• 模型的鲁棒性表明其在对抗性场景中具有强大潜力，例如检测或抵御代码混淆与对抗性样本。