TRACED: Execution-aware Pre-training for Source Code

基本信息

2024 IEEE/ACM 46th International Conference on Software Engineering (ICSE)

博客贡献人

柴进

作者

Fengyi Zhang, Bihuan Chen, Yufei Zhao, Xin Peng

标签

代码表示，预训练模型

摘要

  大多数现有的用于源代码的预训练语言模型集中于学习静态代码文本，或用静态代码结构（抽象语法树、依赖关系图等）来增强。然而，在程序真正执行之前不会完全暴露语义信息。即在不理解程序执行的情况下，静态预训练模型无法全面捕获动态代码属性，例如分支覆盖和运行时变量值，也就影响了它们在代码理解任务（例如检索语义克隆和检测软件漏洞）中的效率。

  因此，为了弥合语言模型的静态性质和程序的动态特性之间的差距，本文作者引入了TRACED，一种执行感知的源代码预训练策略。具体来说，作者用源代码、可执行输入和相应的执行跟踪的组合来预训练代码语言模型。让代码模型在预训练期间能够学习到复杂的执行逻辑，使模型能够静态地估计动态代码属性，而无需在特定于任务的微调期间重复执行代码。

  另外，为了说明文章所提方法的有效性，作者对三个下游任务进行了微调和评估：静态执行估计、克隆检索和漏洞检测。实证结果表明TRACED均表现良好

问题定义

  以往的研究通常使用源代码的静态文本，或添加静态代码结构（抽象语法树、依赖图）进行扩充，再采用自然语言的预训练策略来学习程序表示。作者认为在许多源代码理解任务中，更全面的理解程序行为(program behavior)很重要，因为许多更深层次的程序语义只有在代码执行时才会显现出来。

代码语义克隆检测：程序行为包括确定两段代码在相似的输入下是否表现相似，即使它们的结构明显不同

代码漏洞检测：程序行为包括确定能否执行有潜在问题的代码，以及什么样的值流(value flows)会暴露漏洞

  而现有的代码模型主要被训练来捕获静态代码属性，它们在推理程序行为方面并不有效。

程序状态表示

  本文作者将执行特定时间步长的程序状态定义为当前作用域中每个已定义变量的值集，也就是说程序状态相当于调试器的值映射表。

值量化

  考虑到具体值跨越了广泛的可能值，尤其还存在不同的数据类型，若直接将变量值跟踪为具体值，会导致高维、复杂但稀疏的数据分布。因此，本文作者定义了30个值类别，涵盖了广泛的变量类型，将连续但稀疏的变量值映射到离散的箱中。

执行覆盖表示

  记录执行期间的执行覆盖率，根据哪些行被执行，哪些行没有被执行，并构建执行覆盖率特征供模型学习。

方法

方法架构

主要包括三个阶段：

（1）追踪源代码并设计特征

（2）使用程序追踪进行执行感知预训练

（3）加载预训练的权重并执行特定于任务的微调

第一阶段：追踪和特征工程

  此阶段的目标是为预训练准备数据。

1、第一步是用每个输入执行程序，以生成相应的跟踪。跟踪记录运行时变量值以及执行覆盖率，记录程序的完整执行历史，并显示整个执行过程中程序状态的变化

2、然后，将跟踪中记录的具体运行时值量化到预定义的值范围中，以降低数据的复杂性和稀疏性，使模型更容易学习变量值之间的模式和关系。这些有限的可能输出可以在训练期间用作基本事实标签。

3、另外还创建了程序状态标签和执行覆盖标签

第二阶段：基于追踪的执行感知预训练

  此阶段利用从阶段1获得的预处理样本和标签来执行有监督的预训练。

  本文作者使用基于Transformer Encoder的模型来学习程序跟踪，以提高模型对程序执行的理解。具体地，为了实现产生执行感知代码表示的目标，这篇文章提出了三个预训练任务：

生成源代码

作者认为，理解代码的文本是模型捕捉更复杂信号（如程序执行）的基础。

通过掩码语言建模（MLM）的方法，屏蔽源代码中一定比例的tokens，并训练模型根据周围的上下文来重建被掩码的tokens。

预测程序状态

通过预测在阶段1中生成的程序状态标签，模型可以学习捕获数据流和代码执行的副作用。

预测执行覆盖

通过预测在阶段1中生成的执行覆盖标签，模型可以学习捕获动态控制流，并帮助模型理解程序状态是如何达到和演变的。

阶段3：特定于任务的微调

  此阶段用于验证TRACED模型应用在几个下游任务上的性能

  作者考虑了三个特定的下游任务微调TRACED模型：执行覆盖和运行时变量值预测、克隆检测、漏洞检测。利用TRACED在预训练期间学习的执行信号，静态地推理程序执行过程并完成特定任务，而不需要再次执行程序。

追踪和特征工程

程序状态表示

  类似于实际debug的过程，为了记录程序状态，作者在每一行执行后获取并使用一个value mapping M记录变量的当前值。

程序状态定义

将特定行$l$执行后的程序状态定义为$s(l)$，表示为此时的一组变量值：

s ( l ) = { M ( v , l ) ∣ v ∈ V , l ∈ L } s(l)=\{ M(v,l)|v\in V,l\in L\} s(l)={ M(v,l)∣v∈V,l∈L}

$V$表示所有跟踪变量的集合，$L$是带有源代码的行的集合

【注意】

1、不记录没有可执行代码的行的程序状态，例如图3的第8行。

2、考虑到循环或递归的存在，认为最后出现的值通常足以捕获循环和递归的结果。例如，当调用递归函数时，只有返回变量的最后一个出现的值将被用来完成调用者的后续执行。类似地，循环结束时的最终值将参与将来的执行。因此，使用每行的最后一次执行来确定最终的程序状态，在执行过程中$s(l)$会不断更新，直到执行终止。

量化变量值

  文章作者认为具体的值并不总是必要的，为了降低数据复杂度和增加数据密度，定义了30个类别进行变量值的量化。同时，为了更全面地表示变量值，所提出的量化类别考虑了数据类型、值类型，以及动态运行时的具体值。

构建可学习标签

程序状态标签

  经过“程序状态表示”和“量化变量值”后，会产生一系列程序状态，每个状态由一组量化变量值表示（如图3所示），作者在这些程序状态基础上构建代码模型的可学习特征。具体来说，为可以量化至30个类别中的变量构建标签，并训练模型，使模型在给定源代码表示的情况下预测这些标签。

  每个变量的标签可表示为一个元组：(数据类型、值类型、量化值)，然后为所有可量化的有效变量构建这样的标签，组合所有标签的集合被认为是代码样本的程序状态标签。

执行覆盖率标签

  与程序状态标签保持一致，为变量的每次出现构建了执行覆盖标签。用“Yes”或“No”来表示变量的覆盖情况，其中，在执行行内的变量都会被标记为“Yes”；在未执行行内的变量都会被标记为“No”，同时会给这些变量分配一个“Unknown”量化值，结合变量本身的数据类型和值类型组成一个程序状态标签。示例如下：

TRACED模型

模型架构

  TRACED的主干是一个12层Transformer Encoder，用来学习通用代码表示。在Transformer层之上，堆叠多个多层感知器（MLP）层作为不同任务的预测头。

  在预训练期间，如图4所示，TRACED应用语言模型预测层（即LM层）来预测给定其上下文表示的掩码token，应用程序状态预测层来预测定义的“程序状态标签”，应用执行覆盖层来预测执行覆盖标签。对于特定于任务的微调，加载Transformer层预训练的权重，而预测层由为特定下游任务定制的新的层代替。

执行感知预训练(Execution-aware Pre-training)

模型输入

  每个预训练样本包括可执行程序的源代码和有效的可执行输入，可执行输入和源代码被展平并连接为一个序列。并且为了区分输入和源代码，TRACED使用特殊的[SEP]标记来进行分隔。

可执行输入： E = { e 1 , . . . , e i } E=\{e_1,...,e_i\} E={ e1​,...,ei​