CODEFUSION:代码生成领域的创新预训练模型 [译]

Mukul Singh Microsoft 印度德里 &José Cambronero Sumit Gulwani Vu Le Microsoft 美国雷德蒙 &Carina Negreanu Microsoft Research 英国剑桥 &Gust Verbruggen Microsoft 比利时科尔贝根

摘要

试想一个只能修改他们最新写的那一行代码的程序员,他们需要反复从头开始写代码,直到写对为止——这听起来是不是很费劲?传统的自然语言转代码的自回归模型也面临着类似的困境,一旦生成了代码,就很难再回头修改。为此,我们推出了 CODEFUSION,这是一个突破性的预训练扩散代码生成模型。它能够在自然语言的引导下,反复优化整段代码,直到生成高质量的程序代码。我们在 Bash、Python 以及 Microsoft Excel 的条件格式规则转换任务上对 CODEFUSION 进行了全面测试。测试结果令人振奋:尽管 CODEFUSION 只有 75M 的参数,但它在 top-1 准确率上与最先进且参数量高达 350M–175B 的自回归系统不相上下,而在 top-3 和 top-5 准确率上更是超越对手,彰显了它在保证代码质量的同时,还能提供丰富多样的代码生成选项的卓越能力。

1 引言

自回归代码生成模型,如 Wang et al. (2021)、Brown et al. (2020)、Scholak et al. (2021)、Feng et al. (2020) 和 Fried et al. (2022) 在生成代码的过程中,一旦生成了某个部分,就很难回头再去修改。这种方式可能导致生成的代码变化不大,缺乏多样性 Lin et al. (2023)。为了解决这个问题,之前的研究尝试了一些新的解码策略,如 grouped beam search Vijayakumar et al. (2018) 和 nucleus sampling Holtzman et al. (2019),来增加生成代码的多样性。

近来,扩散模型在图像生成领域取得了卓越的表现 Dhariwal and Nichol (2021),并且已经被扩展到文本生成领域,表现出能够生成更加多样化的文本 Li et al. (2022); Lin et al. (2023)。这些模型通过嵌入层将离散的 tokens 转化为连续的嵌入表示,并在其中添加高斯噪声,再通过预测将其去噪,模拟扩散过程。在最后一步,模型选择与去噪后的嵌入最接近的词汇 token。但在代码生成的场景下,由于代码的语法和语义限制较强,这种方法可能会生成无效的代码。

为了解决这个问题,我们提出了 CODEFUSION,一个自然语言到代码的模型,结合了编解码器架构 Raffel et al. (2020) 和扩散过程。首先,编码器将自然语言指令转换为连续表示,作为扩散模型的条件输入。然后,扩散模型对加入高斯噪声的表示进行去噪,得到去噪后的嵌入。最后,我们将这些嵌入送入一个 transformer 解码器,得到代码 tokens 的概率分布,并从中选择概率最高的 token 作为生成结果。

为了预训练 CODEFUSION 进行代码生成任务,我们将连续段落去噪 (CPD) 任务,原先用于文本 Lin et al. (2023),扩展到代码领域。具体来说,我们只对代码中的标识符和关键字添加噪声。通过这种方式,模型学会了在有噪声的环境中理解和生成代码,尤其是在理解代码中关键元素(如变量名、函数名和控制流关键字)的关系方面。

我们的实验结果表明,与传统的自回归模型相比,CODEFUSION 能够生成更多样化的代码(参见表 2),并且在保持语法正确性方面表现优异,比 Genie(一种基于扩散过程的文本生成模型)生成的语法正确的代码多出 48.5%(在三种编程语言上的平均结果,参见表 3)。

我们在三种不同的编程语言中对 CODEFUSION 这个从自然语言到代码的转换工具进行了评估,这三种语言分别是:Python(由 Yin 等人在 2018 年提出),Bash(由 Lin 等人在 2018 年提出),以及 Microsoft Excel 中的条件格式规则(由 Singh 等人在 2022 年提出)。我们的测试结果表明,尽管 CODEFUSION 的参数量仅为 75M,但它在 top-1 的结果上与参数量达到 350M 到 175B 的顶尖系统相当或甚至更好。在 top-3 和 top-5 的评比中,CODEFUSION 的表现超过了所有其他基准模型。

通过这项工作,我们取得了以下成果:

  1. 我们开创性地提出了 CODEFUSION,这是第一个基于扩散理论的自然语言到代码转换模型。
  2. 我们将连续段落去噪技术(CPD)应用到代码转换中,并证明了这种方法能够显著提升 CODEFUSION 的性能。
  3. 我们对 CODEFUSION 在三种不同编程语言中的自然语言到代码的转换任务进行了全面评估,并将其与自回归代码模型和文本扩散模型进行了对比。

2 相关研究

NL-to-code(自然语言转代码)是一个既热门又充满挑战的任务,涉及从一种序列转换到另一种序列。在以往的研究中,人们采用了诸如 RNNs (由 Brunner 和 Stockinger 在 2021 年提出) 和语义解析器 (由 Lee 等人在 2021 年提出) 等技术来解决这个问题。Vaswani 等人在 2017 年提出的基于注意力机制的变压器模型也被引入到代码领域,产生了一系列的模型,包括仅有编码器的 CodeBERT (由 Feng 等人在 2020 年提出)、有编码器和解码器的 T5 (由 Raffel 等人在 2020 年提出)、以及仅有解码器的 GPT-3 (由 Brown 等人在 2020 年提出)。近期,更大规模的基于变压器的模型,如 CodeGen (由 Nijkamp 等人在 2023 年提出) 和 TransCoder (由 Sun 等人在 2023 年提出),展现出了令人振奋的成果。而指令调优技术也被应用于提升模型性能,具体例子包括 StarCoder (由 Li 等人在 2023 年提出) 和 WizardCoder (由 Luo 等人在 2023 年提出)。这些模型能够以自回归和非自回归的方式进行解码 (由 Su 等人在 2021 年提出)。最新的一些综述文章 (如 Xu 和 Zhu、Niu 等人在 2022 和 2023 年的工作) 也更深入地探讨了这些模型。

扩散模型,一种流行的机器学习模型,已经成功应用于无监督任务和基于文本条件的图像生成任务。近期,这种方法也被扩展到了处理离散领域的文本数据上。

3 方法

本节详细介绍了 CODEFUSION 的架构及其训练和推理过程。如图 1 所示,CODEFUSION 的架构包括编码器 (E)、去噪器 (N) 和解码器 (D) 三个主要组件。

架构图
架构图

图 1: CODEFUSION 的架构图,展示了编码器(E)、去噪器(N)和解码器(D)单元。

3.1 架构

CODEFUSION 这个系统的任务是将自然语言的输入语句 s={s1,s2,,sk}s=\{s_{1},s_{2},\cdots,s_{k}\} 转换成相对应的代码片段 y^={y^1,y^2,,y^s}\hat{y}=\{\hat{y}_{1},\hat{y}_{2},\cdots,\hat{y}_{s}\},这个过程中输入输出的数据都会被统一格式化到一个固定的大小,即 nn 维。在这个系统中,有三个关键的组件,都是基于 Transformer 技术构建的,包括编码器 EE、去噪器 NN 和解码器 DD,此外还有一个分类头 HH

首先,编码器 EE 的作用是把分词后的自然语言输入转换成向量的形式,即 Es=E(s)={e1,e2,,en}E_{s}=E(s)=\{e_{1},e_{2},\cdots,e_{n}\}

接着,在给定编码后的语句 EsE_{s} 和某一时刻 tt 的条件下,去噪器 NN 负责预测并清除嵌入在程序中的噪声 ϵt\epsilon_{t},从而得到去噪后的程序嵌入 x^0=N(xt,t,Es)\hat{x}^{0}=N(x^{t},t,E_{s})。值得注意的是,NN 采用了一种特殊的 Transformer 结构,能够在 xtx^{t}EsE_{s} 之间进行交叉注意,并对 xtx^{t} 实现完全的自注意。

在将去噪后的嵌入数据转换回具体的代码 token 之前,我们需要通过解码器 DD 进一步处理,这一次系统会将注意力放在 x^0\hat{x}^{0}EsE_{s} 上,并对 x^0\hat{x}^{0} 实施全面的自注意,从而得到最终的隐藏表示 Ds={d1,d2,,dn}=D(x0,Es)D_{s}=\{d_{1},d_{2},\cdots,d_{n}\}=D(x^{0},E_{s})。这种处理方式与之前文本扩散的方法有所不同,后者在生成 token 时是各自独立的,而我们的方法能够确保在生成每一个维度(did_{i})时都能充分考虑到其他维度的信息。

最后,通过分类头 HH,我们将隐藏表示 DsD_{s} 转换成具体的代码 token,并计算出每个 token 的概率分布 p(ydi)p(y|d_{i})。在这个过程中,我们并不会在所有可能的 token 中进行搜索,而是直接选择概率最大的那个,即 y^i=argmaxyp(ydi)\hat{y}_{i}=\arg\max_{y}p(y|d_{i}),作为我们的输出结果。这样,我们就能够根据自然语言的输入,高效准确地得到对应的代码片段。

3.2 训练

我们通过两个步骤对 CODEFUSION 进行训练:首先在代码片段上对去噪器和解码器进行无监督的预训练,接着在(用户语句,代码片段)这样的数据对上对编码器、去噪器和解码器进行有监督的微调。借鉴了之前在文本处理领域的成果,我们使用了一个可训练的嵌入层 LL,将代码片段 yy 转换成一个连续的向量空间,在这个空间中我们可以在不同的时间点 tt 加入或去除噪声 ϵt\epsilon_{t}

我们受到了以往在文本处理上应用扩散模型的启发,并且将 genie Lin 等人在 2023 年提出的损失函数调整应用到了 CODEFUSION,加入了解码器中隐藏状态 DsD_{s} 的影响。在每一个时间点 tt,损失函数被定义为:

Lt=ϵ^tϵt+DsL(y)\displaystyle\mathcal{L}_{t}=\|\hat{\epsilon}_{t}-\epsilon_{t}\|+\|D_{s}-L(y)%\|-\log p(y|D_{s})

这个损失函数包括三个部分:

  1. 我们通过最小化预测噪声 ϵ^t\hat{\epsilon}_{t} 和实际噪声 ϵt\epsilon_{t} 之间的差异来训练去噪器 NN
  2. 我们通过最小化解码器的隐藏状态 DsD_{s} 和嵌入后的代码之间的差异来共同训练解码器 DD 和嵌入层 LL
  3. 我们利用标准的交叉熵损失函数,比较分类头根据 DsD_{s} 预测出的代码标记和实际的代码片段 yy,以此来进行训练。

我们的模型由三大部分构成:去噪器、解码器和分类头。利用损失函数,我们可以对其进行训练,使其达到我们所期望的效果。

在代码片段的大量数据中,我们对去噪器(记为 NN)和解码器(记为 DD)进行了预训练。在这个过程中,我们主要进行了两项任务:一是无监督的代码生成,二是我们改良版的连续段落去噪(简称 CPD)任务,具体来说就是 Lin et al. 在 2023 年的一篇论文(链接:2023)中提到的一种方法,这种方法专门针对代码,并且只对目标语言中的标识符或内置关键词相关的标记进行掩盖。预训练阶段,我们会在这两个任务之间进行随机选择。

无论是预训练还是后续的微调阶段,我们都会使用 Lt\mathcal{L}_{t} 这个标准进行评估。需要注意的是,在预训练阶段,并没有涉及到自然语言的处理,因此去噪器 NN 没有接收到 EsE_{s} 这个输入。在无监督代码生成的任务中,我们用高斯噪声来代替 EsE_{s},而在 CPD 任务中,EsE_{s} 是通过将经过掩盖的代码 yy 传递给编码器 EE 计算得到的。

3.3 推断阶段

在进行预测的时候,我们首先利用高斯噪声来设置初始值 xtx^{t},接着在 TT 个时间步长内,按照预先设定的规则逐渐减少噪声,最终得到 x^0\hat{x}^{0}(详情见 Ho et al.,2020 的研究)。值得注意的是,在这个去噪的过程中,我们并没有使用解码器。只有在迭代完成后,我们才运用解码器来生成最终的预测代码 y^\hat{y}。为了得到更准确的结果,我们还会对 y^\hat{y} 进行一番筛选,只保留那些在 pad 标记之前出现的标记。

系统描述Python(使用 CodeBERT)Bash(基于模板)CF 规则(执行)
系统模型参数数量top-1top-3top-5top-1top-3top-5top-1top-3top-5
T5t5-large770M80.482.384.867.168.970.371.173.474.6
CodeT5codet5-large770M80.583.185.067.669.370.572.775.375.8
GPT-3text-davinci-003175B82.583.785.866.967.768.470.372.472.8
ChatGPTgpt-3.5-turbo20B80.682.583.966.166.967.870.873.174.5
StarCoderstarcoder15.5B79.282.084.164.565.366.570.672.874.5
CodeT5+codet5p-16b16B79.682.184.565.766.167.270.572.974.3
CodeGencodegen-350m350M80.181.883.767.269.270.371.473.775.0
Diffusion-LMCustom50M70.474.376.559.461.662.062.465.568.2
GenieCustom93M73.277.180.360.061.562.362.966.868.7
CODEFUSIONCustom75M80.786.390.366.770.272.072.876.778.5

表 1: 我们将 CODEFUSION 与其他基线方法在自然语言到代码生成任务上进行了比较,覆盖了 Python、Bash 和 CF 规则三个领域。我们报告了 top-1、top-3 和 top-5 的预测精度。其中,模型列显示了底层基础模型的名称,参数数量列显示了模型的参数数量。我们在括号中注明了每种语言对应的评价指标。

4 评估设置

我们简要介绍了训练过程、对比的基线方法、评估基准和使用的评价指标,并在附录中提供了关于训练和基线方法的更多详细信息。

4.1 基准测试

我们使用 CODEFUSION 对几种复杂程度各异的编程语言进行了自然语言到代码(NL-to-code)的任务评估,这些语言包括 Python、Bash 以及 Microsoft Excel 中的条件格式化(CF)规则。其中,Python 任务采用的是 CoNaLa 数据集,该数据集由 Yin et al. (2018) 提供,包含了一系列复杂的、由多条语句组成的 StackOverflow 中的 Python 代码片段及其对应的自然语言问题。Bash 任务采用的数据集由 Lin et al. (2018) 提供,包括了一些复杂的单行 Bash 命令及其自然语言描述。CF 任务则使用由 Singh et al. (2022) 提供的数据集,该数据集包含了一些简单的单行 Excel CF 规则及其自然语言描述。这些基准测试数据集和预训练数据(详见下一节)都已经公开提供。11https://github.com/microsoft/prose-benchmarks/tree/main/CodeFusion

4.2 训练

在我们的实验中,我们将编码器(E)实例化为一个预训练的 CodeT5 编码器(嵌入维度为 512),去噪器(N)实例化为一个由 10 层 Transformer 块组成的模型,解码器(D)实例化为一个由 6 层 Transformer 解码器层组成的模型,分类头(H)则是一个单层的全连接层。

在训练和预训练阶段,我们采用了由 Wu et al. (2023) 提供的一个包含 1200 个扩散步骤的平方根噪声计划。我们使用了由 Wang et al. (2021) 提供的 CodeT5 的分词器和词汇表,并将目标代码长度设置为 128 个令牌。我们使用了不带权重衰减的 AdamW 优化器(Loshchilov and Hutter, 2019)以及 5e-4 的学习率。

我们在代码片段上对扩散模型和解码器模型进行了预训练。对于 Excel,我们使用了一个公共语料库,其中包含了 450,000 条条件格式化规则(Singh et al., 2022)。对于 Python 和 Bash,我们从 GitHub 笔记本和带有 python、bash 和 powershell 标签的 StackOverflow 帖子中提取了代码片段,使用了正则表达式提取器进行检测(Lin et al., 2018)。

4.3 对比模型

我们使用了一系列 Transformer 和文本扩散模型作为对比模型:T5(Raffel et al., 2020)、GPT-3(text-davinci-003,Brown et al., 2020)、ChatGPT(gpt-3.5-turbo,OpenAI, 2023)、CodeT5(Wang et al., 2021)、StarCoder(Li et al., 2023a)、CodeT5+(Wang et al., 2023)、CodeGen(Nijkamp et al., 2023)、Diffusion-LM(Li et al., 2022)和 GENIE(Lin et al., 2023)。其中,GPT-3、ChatGPT、StarCoder 和 CodeT5+ 采用了在上下文中学习的设置,并使用了五个根据自然语言表达和 SentenceBERT 相似度动态选取的示例(Reimers and Gurevych, 2019)。其他模型在我们的数据集上进行了微调。Diffusion-LM 和 GENIE 也在我们的语料库上进行了预训练。

4.4 评估指标

我们采用了数据集中提供的一种模板匹配方式来评估 Bash 代码的生成质量,这种方法涵盖了一些基础的标准化处理。对于 Python 代码,我们选择了 CodeBERTScore(Zhou et al.,2023)作为评估工具,该工具已经证明在非执行型代码匹配方面具有很高的准确度。而对于 CF 代码,我们使用了执行匹配的方法(Singh et al.,2022),即通过在数据列上执行一条规则,并将结果与预期输出进行比对来进行评估。

在评估代码生成的多样性方面,我们主要关注三个方面:1)生成代码中不同 token n-grams 的数量与 token 总数的比值(参考 Vijayakumar et al.,2018);2)CodeBERT 编码的成对相似性的统计数据(参考 Perlitz et al.,2023);3)字符串编辑距离的成对比较的统计数据(同样参考 Perlitz et al.,2023)。

5 评估结果分析

在这部分,我们主要探讨了三个问题:Q1. CODEFUSION 是否能够生成既正确又多样化的代码?Q2. 不同设计决策对性能有何影响?Q3. 在模型的扩散步骤中,潜在表示是如何变化的?

表 2 展示了对于 Python、Bash 和 CF 规则,CODEFUSION 与其他基准模型在生成的前 5 个代码方案中多样性的比较。我们从三个维度进行了评估:token-level n-grams 的多样性,CodeBERT 编码的成对相似性,以及字符串编辑距离的统计数据。

模型N-grams(↑ 更好)嵌入相似度(↓ 更好)编辑距离(↑ 更好)
1234最小值最大值平均值最小值最大值平均值
T50.290.450.540.60.930.990.970.080.580.36
CodeT50.270.410.460.510.940.990.970.080.560.32
GPT-30.240.360.420.480.970.990.990.000.280.21
StarCoder0.320.480.560.610.890.980.960.110.610.39
CodeGen0.280.410.490.550.940.990.970.090.520.34
CODEFUSION0.530.650.740.810.830.970.910.210.830.57

表 2:对于 Python、Bash 和 CF 规则,CODEFUSION 与其他基线模型在生成的前 5 个代码方案中多样性的对比。我们从 token-level n-grams 的多样性,CodeBERT 编码的成对相似性,以及字符串编辑距离的统计数据三个维度进行了评估。

5.1 性能与多样性 (Q1)

表格 1 对 CODEFUSION 和其他基准模型在不同设置下的性能进行了总结。

在单一最佳选择(top-1)的情况下,CODEFUSION 的表现非常出色,甚至超过了一些体量更大的自回归模型。以 Python 为例,只有体量更庞大的 GPT-3 (175B) 能胜过 CODEFUSION (75M)。而在前三和前五的选择中,CODEFUSION 更是超越了所有对比模型,这验证了之前的研究,即高性能的自回归模型在生成多样性上往往会有所牺牲 Poesia et al., 2022

在表格 2 中,我们看到在各个基准任务上,CODEFUSION 和一系列自回归模型(如 T5, CodeT5, StarCoder, CodeGen, GPT-3)在前五个选择的多样性上的对比。结果显示,CODEFUSION 生成的代码更为多样。

和 CODEFUSION 类似,其他的扩散方法(如 Diffusion-LM 和 genie)在前三和前五的选择中相较于单一选择有所改进。但由于生成了一些语法错误的代码,它们的表现仍然不如 CODEFUSION。在表格 3 中,我们可以看到 CODEFUSION 和其他扩散模型在生成语法正确代码上的对比。CODEFUSION 生成的代码在语法上更为准确,比 Diffusion-LM 高出 33.8%,比 Genie 高出 26.2%,这是在所有三种编程语言上的平均结果。

系统PythonBashCF 规则
Diffusion-LM19.540.474.3
Genie24.254.278.6
CODEFUSION67.673.494.5

表 3:在 CODEFUSION 和基于文本扩散的模型中,单一最佳选择中语法正确的代码占比。CODEFUSION 生成了更多语法正确的代码。

5.2 消融研究 (Q2)

表格 4 展示了对 CODEFUSION 进行各种修改后的结果。

从结果来看,移除任何一个预训练任务都会对性能造成明显的负面影响,其中代码生成任务的移除会导致性能下降 10.9%,而 CPD 任务的移除会导致性能下降 4.6%,这些数据是基于三种编程语言的平均结果。另外,通过比较用 grounding 和 clamping 方法替代 D 和 H,我们发现在最终确定代码前使用解码器能带来性能上的提升。

模型PythonBashCF 规则
M - 代码生成70.952.364.2
M - CPD 目标76.761.168.2
M 采用 Grounding73.560.563.2
M 采用 Clamping77.163.264.4
M (CODEFUSION)80.766.772.8

表 4:对 CODEFUSION (M) 进行消融研究的结果。“-”表示相应的组件被移除。完整模型在所有情况下都表现最好。

5.3 逐步提升(Q3)

我们深入研究了 CODEFUSION 是如何一步步接近其最终输出结果的。在这项实验中,我们选择在某个时间步 t(取值范围在 0 到 T 之间)停止去噪处理,并生成当前状态下的代码片段。我们测量了每个时间步(以 100 步为单位递增)所得到的标准化字符串编辑距离。如图 2 所示,随着 t 的增加,编辑距离逐渐减小。值得注意的是,CF(CodeFusion)规则的编辑距离下降得更快,这是因为与生成完整程序和 Bash 命令相比,生成 CF 规则相对更为简单。图 3 通过一个实际例子直观地展示了这一点。

平均标准化编辑距离
平均标准化编辑距离

图 2:随着扩散时间步增加,CODEFUSION 生成代码的平均标准化编辑距离。

逐步去噪的过程
逐步去噪的过程

图 3:CODEFUSION 在一个 Python 示例基准测试上进行逐步去噪的过程。

6 总结

我们提出了 CODEFUSION,这是首个利用扩散过程从自然语言生成代码(NL-to-code)的模型。通过使用专注于代码生成的解码器和预训练,CODEFUSION 能够生成语法更准确、更富有多样性的程序,相较于现有的文本扩散模型和自回归代码模型都表现出更强的竞争力。在对 Python、Bash 和 Excel 条件格式规则的代码生成任务中,CODEFUSION 的表现可与最先进的 Transformer 代码生成模型相媲美。

7 局限性

CODEFUSION 的应用范围并不全球化,因为我们仅考虑了英语的自然语言输入。另外,自然语言规范的详细程度不同,描述不够详细的情况可能会导致性能下降。我们虽然支持多种编程语言,但对于更复杂的语言,性能可能会有所下降。我们还注意到,当面临生成较长的代码片段或语法依赖关系较长的程序时,CODEFUSION 的性能会受到影响。由于基于扩散的模型是通过迭代去噪的方式工作,因此随着目标生成长度的增加,推理延迟会显著增加。

附录 A 实施细节

我们首先介绍了进行实验的环境和设置,然后详细说明了每个组件的参数设置和训练过程。

A.1 硬件配置

我们的所有实验都是在 Python(版本 3.8.7)环境下进行的,使用了一台配备 Intel Core i7 处理器(基础频率 1.8 GHz)、4 个 V100 GPU 单元、64 位操作系统和 56 GB 内存的计算机。对于每个数据集,CODEFUSION 的预训练平均耗时 8 小时,微调平均耗时 3 小时。

A.2 训练参数设置

在训练和预训练阶段,我们采用了一个包含 1200 个扩散步骤的平方根噪声计划 Wu et al. (2023)。我们使用了 CodeT5 Wang et al. (2021) 提供的分词器和词汇表。我们设置学习率为 5e-4m,批处理大小为 64,目标长度为 128,并使用了 AdamW 优化器进行训练,其中不包含权重衰减 Loshchilov and Hutter (2019)。

A.3 预训练数据简介

表 5 展示了用于预训练的样本数据,包括样本数量和平均长度(以令牌数计)。从表中我们可以看出,在用于预训练 CODEFUSION 的数据集中,Python、Bash 和条件格式规则三种语言的代码片段数量分别为 56K、35K 和 100K,平均长度分别为 78.4、45.3 和 23.4 个令牌。

语言样本数平均长度
Python56,00078.4
Bash35,00045.3
条件格式规则100,00023.4

表 6 则对 CODEFUSION 进行了详细的性能评估,将其与其他几个基线模型在文本到代码生成任务上的表现进行了对比。从表中我们可以看出,CODEFUSION 在处理 Python、Bash 和条件格式规则时的准确度都相当高,尤其是在 top-1 和 top-3 的准确匹配上表现突出。其中,“模型”列显示了各系统所使用的基础大语言模型 (LLM),而 #P 则表示了模型中的参数数量。

系统模型参数数目Python(top-1/3/5)Bash(top-1/3/5)条件格式规则(top-1/3/5)
T5t5-large770M5.2/6.1/6.713.5/14.7/15.263.1/65.2/67.6
CodeT5codet5-large770M5.5/6.4/7.114.1/14.9/15.563.2/65.7/67.8
..................
CODEFUSIONCustom75M5.1/7.2/9.013.5/15.3/16.463.4/67.6/69.1

通过这些数据,我们可以清楚地了解到 CODEFUSION 在各个方面的性能,以及与其他模型相比的优势和不足。

最后附上的图示展示了去噪过程中的各个阶段,提供了对模型如何处理和改进代码的直观理解。

图 4:这幅图展示了 CODEFUSION 在处理一个 Python 基准测试示例时,从最初的纯噪声逐渐过渡到生成目标代码的各个阶段。可以看出,CODEFUSION 最终成功地完成了这一任务。

附录 B 基线系统的实现细节

在这一部分,我们将介绍一些基准系统的实现详情。

B.1 T5, CODET5 和 CODEGEN

T5 是 CODEFUSION 的基模型,它采用了编解码的架构方式。CodeT5 则是在 T5 的基础上,通过在代码生成和理解任务上进行预训练来加以改进的。我们对 T5 和 CodeT5 进行了训练,使它们能够根据给定的输入生成代码,其输入方式与 CODEFUSION 在微调阶段的输入方式相同(参见第 3.2 节)。这些模型通过标准的交叉熵损失函数和 Adam 优化器进行优化,学习率设置为 1e−4,训练周期为 100。实验结果表明,CodeT5 的性能优于 T5,这主要是因为 CodeT5 在代码生成方面进行了预训练。这些模型是在数据集的训练部分上进行训练,并在测试部分上进行评估的。在解码过程中,我们使用了束搜索方法,束大小设置为 5。

B.2 GPT-3, CHATGPT, STARCODER 和 CODET5+

我们在少量样本的设置下,直接使用了预训练好的 GPT-3(Text DaVinci 003)、ChatGPT(gpt-3.5-turbo)、StarCoder(bigcode/starcoder)和 CodeT5+(codet5p-16b)模型,没有进行进一步的微调。我们在提示信息中给出了 5 个示例。这些示例是根据训练语料中的句子与查询句子的余弦相似度,使用 sentenceBERT Reimers 和 Gurevych(2019)嵌入方法进行选取的。对于 GPT-3 和 ChatGPT,我们测试了从 0 到 1 的温度值(步长为 0.1),并在所有其他参数保持默认值的情况下,报告了最佳结果。对于 StarCoder 和 CodeT5+,我们采用了束搜索方法进行解码,束大小同样设置为 5。

B.3 GENIE

我们按照 Lin 等人在 2023 年的论文中的描述,实现了 Genie 模型。我们将扩散时间步设置为 1200,嵌入维度设置为 256,编码和生成长度均设置为 128,这些参数设置是为了与 CODEFUSION 保持一致。我们还使用了与预训练 CODEFUSION 相同的语料进行预训练。在采样 top-k 候选时,我们使用了 Lin 等人在他们的论文中推荐的推理算法。

B.4 DIFFUSION-LM

我们根据 Li 等人在 2022 年的论文中的描述,实现了 Diffusion-LM 模型。我们将扩散时间步、嵌入维度和编码生成长度的参数设置,与 Genie 模型保持一致。我们利用控制生成分类器,将文本到代码的生成任务转换为一个完成任务,这是通过简单扩展作者在论文中为文本生成提出的填充任务实现的。

附录 C 更多结果

在这个部分,我们展示了从我们的语料库中获取的一些额外的结果和示例。

C.1 精确匹配结果

我们在这里展示了 CODEFUSION 和其他主流模型在标准测试集上的精确匹配准确率。你可以在表 6 中找到具体数据。简单来说,CODEFUSION 在最高准确率方面与基于 transformer 的模型不相上下,并且比其他基于扩散的文本生成方法更为出色。而在前三高准确率方面,CODEFUSION 更是领先于所有对比模型。这与我们在表 1 中看到的结果是一致的,证明了 CODEFUSION 能够为不同任务生成更优质、种类更丰富的候选程序。

图 5: 这是一个 CODEFUSION 处理 Python 基准测试时失败的例子,展示了从纯噪音出发,逐步去噪,直至生成目标代码的过程。最终的生成结果并不正确,因为正确的函数名应该是 os.removedirs,而且这个函数只能用来删除空目录,不符合用户删除含有文件的目录的需求。

C.2 扩散可视化

CODEFUSION 通过迭代去除潜在表示中的噪声,最终构建出目标程序。这个过程可以通过将每个时刻的表示转换为离散的令牌来直观展示。具体的实现方式可以参考 5.3 节。图 4 展示了一个成功案例,其中 CODEFUSION 能够从 Python 基准测试中正确生成代码。图中展示了输入的查询、目标代码,以及 CODEFUSION 在不同时间步 t=1200 的输出结果。可以清晰地看到 CODEFUSION 是如何逐步去噪并最终生成正确代码的。图 5 则展示了一个失败案例,CODEFUSION 未能生成正确的代码。用户的需求是删除一个包含文件的目录树 'folder_name/',但 CODEFUSION 生成的代码中使用了一个不存在的函数 os.removedir,正确的函数名应该是 os.removedirs,并且这个函数只能用来删除空目录,不能满足用户的需求。

图 6: 这个图表展示了 CODEFUSION 在 Python 编程语言上的 Top-1 CodeBERT 分数随着扩散时间步增加的变化情况。可以看出,随着时间步的增加,性能逐渐提升并在 T=1000 时趋于稳定。

C.3 扩散时间步的影响

扩散时间步的数量对生成质量有直接影响,这一点在 Saharia et al. (2022a) 的研究中已经得到证实。我们对 CODEFUSION 进行了实验,探索了在 Python 编程语言上,不同时间步 T={200,400,600,800,1000,1200}T=\{200,400,600,800,1000,1200\} 对模型性能的影响,并将结果以 CodeBERT 分数的形式绘制在图 6 中。从图中可以看出,增加时间步确实能够提升 CODEFUSION 的性能,但性能提升的幅度随着时间步的增加而减小,当时间步达到 1000 时,性能基本稳定。需要注意的是,增加时间步同时也会增加模型的推理延迟和内存需求。

C.4 延迟和内存

扩散模型通常具有复杂的结构,参数众多,因此其延迟和内存需求相对较高,特别是在进行重复采样和序列去噪操作时。CODEFUSION 具有 7500 万个参数,需要 544 MB 的磁盘空间。在标准测试集上,其平均推理延迟为 2318 毫秒,平均 GPU 内存使用量为 928 MB,最大 GPU 内存使用量达到 1042 MB。

附录 D 背景知识

这个直译和意译的结合方式能够帮助理解原文的精确含义,同时又不失去原文的语境和细节。

D.1 基于 Transformer 的序列生成

Transformer 语言模型(由 Vaswani 等人于 2017 年提出)属于条件生成模型的一种,并通过自回归(AR)解码的方式来实现。这种模型能够预测目标词 yt 的出现概率,这一预测是基于已知的条件输入和之前生成的词序列 y1,y2,,yt1y_{1},y_{2},\cdots,y_{t-1}。整个生成序列的概率可以用下面的公式表示:

P(yx)=i=1Np(yiy1:i1;x)\mathcal{P}(y|x)=\prod_{i=1}^{N}p(y_{i}|y_{1:i-1};x)

D.2 扩散模型

扩散过程是一个按照离散时间推进的马尔可夫过程。这个过程从一个初始状态 x0x_{0} 开始,在时间点 t=0t=0 时刻,x0x_{0} 取自原始的数据分布。随着时间的推进,我们逐步给 x0x_{0} 加入一些高斯噪声,这个加入噪声的过程是根据一组预定的方差值 β1,,βT\beta_{1},\cdots,\beta_{T} 来进行的。因为这个过程仅仅是在添加噪声,所以在任何一个时间点 t+1t+1,我们都可以用 xtx_{t} 来表示 xt+1x_{t+1},具体如下所示:

q(xt+1xt)=N(xt+1;1βt+1xt,βt+1I)\small q(x_{t+1}\|x_{t})=\mathcal{N}\left(x_{t+1};\sqrt{1-\beta_{t+1}}x_{t},\beta_{t+1}I\right)

在模型训练的阶段,扩散模型的任务是学习如何进行逆扩散过程,也就是预测当前状态 xtx_{t} 在时间点 tt 时刻的噪声。通过使用这个预测出来的噪声,我们就能够回推出前一个状态 xt1x_{t-1},方法是从 xtx_{t} 中减去这个噪声,并对均值进行重新缩放。所以,从 xtx_{t} 推算出 xt1x_{t-1} 的分布实际上就是一个均值为 μθt1\mu_{\theta}^{t-1}、方差为 σθt12\sigma_{\theta}^{t-1^{2}} 的高斯分布,这里的 θ\theta 是神经网络的参数。

训练扩散模型的目标是尽量减小实际均值 μθt1\mu_{\theta}^{t-1} 和预测均值 μ^θt1\hat{\mu}_{\theta}^{t-1} 之间的均方误差。

D.3 文本生成中的扩散模型

近期的文本生成研究开始涉足扩散模型。Li 等人在 2022 年提出了一种系统,它利用嵌入技术把文本中的离散词元转换成连续的表征,作为目标分布。接着,扩散模型被训练为逐渐从随机高斯噪声中去噪,生成目标分布中的样本。Lin 等人在 2023 年将这一方法扩展到了条件文本生成,通过将输入编码并与扩散状态进行拼接。Wu 等人在 2023 年提出了使用夹持技术来强制模型在每个时间步都紧贴目标分布,从而提升了性能。