[论文解读] Contrastive Learning for Source Code with Structural and Functional Properties

Yangruibo Ding, Luca Buratti+4 more|arXiv (Cornell University)|Sep 29, 2021

Software Engineering Research参考文献 31被引用 8

一句话总结

BOOST 是一种用于源代码的自监督对比学习模型，通过利用代码的结构和功能特性来提升代码理解能力。通过生成功能等价但结构不同的代码，以及功能不同但语法相似的代码，它通过对比学习和一种新颖的节点类型掩码语言建模目标学习鲁棒表示，实现了在显著更少的预训练数据下的最先进性能。

ABSTRACT

Pre-trained transformer models have recently shown promises for understanding the source code. Most existing works expect to understand code from the textual features and limited structural knowledge of code. However, the program functionalities sometimes cannot be fully revealed by the code sequence, even with structure information. Programs can contain very different tokens and structures while sharing the same functionality, but changing only one or a few code tokens can introduce unexpected or malicious program behaviors while preserving the syntax and most tokens. In this work, we present BOOST, a novel self-supervised model to focus pre-training based on the characteristics of source code. We first employ automated, structure-guided code transformation algorithms that generate (i.) functionally equivalent code that looks drastically different from the original one, and (ii.) textually and syntactically very similar code that is functionally distinct from the original. We train our model in a way that brings the functionally equivalent code closer and distinct code further through a contrastive learning objective. To encode the structure information, we introduce a new node-type masked language model objective that helps the model learn about structural context. We pre-train BOOST with a much smaller dataset than the state-of-the-art models, but our small models can still match or outperform these large models in code understanding and generation tasks.

研究动机与目标

• 解决现有代码模型过度依赖文本和语法特征、常忽略功能语义的局限性。
• 通过自动代码转换引入功能等价性和语法相似性，改进代码表征学习。
• 通过引入更高效的数据自监督训练方法，减少对大规模预训练数据集的依赖。
• 通过一种新颖的节点类型掩码语言建模目标，增强代码模型的结构感知能力。

提出的方法

• 通过自动化的、以结构为导向的代码转换算法生成两类增强样本：功能等价但结构不同的代码，以及语法相似但功能不同的代码。
• 应用对比学习目标，使功能等价代码对的表示更接近，同时使功能不同但相似的代码对的表示更远离。
• 引入一种节点类型掩码语言建模目标，通过在抽象语法树（AST）中掩码特定节点类型，提升模型对代码结构上下文的理解能力。
• 使用对比损失进行模型预训练，以促进等价代码变体之间的语义对齐，同时保留恶意或非预期行为之间的差异。
• 该方法设计为数据高效，即使在比最先进模型更小的预训练数据集上，也能实现优异性能。
• 最终模型在下游代码理解与生成任务上进行微调，展现出优于或相当的性能，且所需预训练数据更少。

实验结果

研究问题

• RQ1 当显式建模功能等价性和语法相似性时，自监督对比学习是否能提升代码表征？
• RQ2 节点类型掩码语言建模目标在增强代码模型结构感知能力方面有多有效？
• RQ3 在多小的预训练数据集上，性能可与或优于更大规模模型在代码理解任务中的表现？
• RQ4 自动化代码转换技术能否生成有意义的正样本和负样本来增强代码表征的对比学习？
• RQ5 结合功能与结构特性是否能带来更鲁棒且泛化能力更强的代码表征？

主要发现

• 尽管预训练数据量显著少于现有模型，BOOST 在代码理解与生成任务上仍实现了最先进性能。
• 该模型在多个代码相关基准测试中优于更大的最先进模型，证明了其数据高效训练方法的有效性。
• 引入节点类型掩码语言建模目标显著提升了模型捕捉代码表征中结构上下文的能力。
• 通过功能等价与功能不同的代码对进行对比学习，可生成更鲁棒且语义更清晰的表征。
• 自动化代码转换流水线成功生成了高质量的正样本与负样本，有效增强了对比学习效果。
• BOOST 的性能在与预训练数据量大得多的模型对比中仍具竞争力或更优，凸显其数据效率优势。