Paper name
CODEFUSION: A Pre-trained Diffusion Model for Code Generation
Paper Reading Note
Paper URL: https://arxiv.org/abs/2310.17680
TL;DR
- 2023 微软出的文章,提出了 CODEFUSION,这是首个基于 diffusion 的自然语言到代码 (NL-to-code) 生成模型。与自回归模型相比,CODEFUSION 产生了更多样化的代码。
Introduction
背景
-
想象一下,如果有一名代码开发人员只能更改他们编写的最后一行代码,那么在编写一个函数并确保其正确之前,需要从头开始编写多少次呢?
- 作者认为基于自然语言生成代码的自回归模型有类似的局限性:它们不容易重新考虑之前生成的标记(tokens)。这一限制可能导致在相关文本领域生成的多样性降低。为了平衡生成候选项的多样性和质量,先前的工作探索了解码策略,如分组束搜索 (grouped beam search) 或核采样 (nucleus sampling)
-
扩散模型在图像生成中表现出色,最近已被扩展用于生成多样化的文本。这些方法使用嵌入层将离散 token 转换为连续嵌入,在此基础上可以添加和预测高斯噪声,以模拟扩散过程。为了将去噪后的嵌入映射回离散文本,这些方法选择最接近嵌入的词汇 token。在代码领域,token 之间存在许多语法和语义约束,独立地将嵌入映射回 token 可能会生成无效的程序。
本文方案
- 提出了 CODEFUSION,一个自然语言到代码(NL-to-code)的生成模型
- 它结合了编码器-解码器架构(Raffel et al., 2020)与扩散过程。编码器将自然语言映射为连续表示,扩散模型使用这一表示作为额外条件来去噪随机高斯噪声输入。为了生成语法正确的代码,我们随后将去噪后的嵌入输入到 transformer 解码器,通过完整的自注意力和与嵌入话语的交叉注意力,获得代码 token 的概率分布。最后,我们在每个索引选择最高概率的 token。
- 为了预训练 CODEFUSION 进行代码生成,我们将连续段落去噪(CPD)任务扩展到代码领域。具体来说,我们只对代码中对应于标识符或目标语言内置关键词的 token 应用噪声。这个去噪任务使模型能够学习关键代码 token(如变量名、函数名和控制流内置函数)之间的关系。
- 我们发现,与自回归模型相比,CODEFUSION 产生了更多样化的代码(n-gram 比例更高,嵌入相似性更低,编辑距离更高)。CPD 目标,它使模型偏向于学习以语境感知的方式去除噪声,并配以能够访问完整去噪表示的解码器,共同使 CODEFUSION 在与 GENIE 比较时(一个文本扩散模型),生成了 48.5% 更多语法正确的代码(平均跨三种语言)。我们在三种不同语言的自然语言到代码上评估 CODEFUSION:Python,Bash 和 Microsoft Excel 中的条件格式化规则。我们的结果表明,CODEFUSION的(7500万参数)top-1 结果与更大的最新系统(3.5亿–1750亿参数)相当或更好。在 top-3 和 top-5 中,CODEFUSION 的表现优于所有基线。
Methods
架构设计
- CODEFUSION 的输入是自然语言话语 s = s 1 , s 2 , … , s k s={s1,s2,…,sk} s=s1,s2,…,sk,输出是预测的代码片段 y = y 1 , y 2 , … , y s y^={y^1,y^2,…,y^s} y=y1,y2,…,ys。输入和输出都被填充到固定的维度 n。
- CODEFUSION有三个主要的基于 transformer 的组件(一个编码器E,一个去噪器N,一个解码器D)和一个分类头H。
- 编码器(E,预训练的 CodeT5):将分词话语 s 转换为向量表示 E s = E ( s ) = e 1 , e 2 , … , e n Es=E(s)={e1,e2,…,en} Es=E(s)=e1,e2,…,en。
- 去噪器 (N,10层 transformer 模块):在编码的话语 Es 和时间 t 的条件下,去噪器(N)预测并移除噪声 εt, 从噪声程序嵌入 xt 获得预测的去噪程序嵌入
x
0
=
N
(
x
t
,
t
,
E
s
)
x^0=N(xt,t,Es)
x0=N(xt,t,Es)。
- N 是一个带有 xt 和 Es 之间的交叉注意力以及 xt 上的完全自注意力的 transformer 模块。
- 解码器(D,6 层 transformer 模块):应用交叉注意力到 x 0 x^0 x0 和 Es,以及在 x 0 x^0 x0 上的完全自注意力,来计算最终的隐藏表示 D s = d 1 , d 2 , … , d n = D ( x 0 , E s ) Ds={d1,d2,…,dn}=D(x^0,Es) Ds=d1,d2,…,dn=D(x0,Es)。与以前的文本扩散方法不同,在那里标记是独立生成的,完全自注意力允许每个隐藏维度( d i di di)在关于其他维度的完整信息中生成。
- 分类头 (H,单个全连接层):Ds 被投影到实际的代码 token 上,通过一个分类头 H 计算代码 token 的分布 p ( y ∣ d i ) p(y∣di) p(y∣di)。我们不执行对这些 token 的搜索,并为每个 i 选择 y i = a r g m a x y p ( y ∣ d i ) y^i=argmax_{y}p(y∣di) yi=argmaxyp(y∣di)。
训练
-
两阶段训练:
- 在代码片段上无监督预训练去噪器和解码器
- 在(自然语言,代码片段)上有段度微调编码器、去噪器和解码器
-
根据之前在文本扩散上的工作,我们使用一个可训练的嵌入层 L 来将代码片段 y 嵌入到一个连续的空间,在时间步 t 我们可以添加(并移除)噪声 εt。我们从之前在文本扩散上的工作中获得灵感,并通过结合来自解码器的隐藏表示 Ds 将来自 GENIE 的损失适应到 CODEFUSION。在时间步 t,损失计算为
- 我们通过最小化预测噪声 ε ^ t \hat{ε}_{t} ε^t 与实际噪声 ε t ε_{t} εt 之间的误差来训练去噪器 N。
- 我们通过最小化隐藏表示 Ds 与嵌入代码之间的误差来训练解码器 D 和嵌入层 L。
- 我们在分类头的输出上应用标准交叉熵损失,该头根据 Ds 产生预测的代码标记,并与真实代码片段 y 进行比较。损失函数允许我们使用扩散目标训练模型的三个主要组件(去噪器、解码器和分类头)。
-
预训练阶段的两个任务
- 无监督的代码生成
- 对代码的连续段落去噪(CPD):这个特定于代码的CPD任务只掩盖与目标语言的标识符或内置关键词相关的标记
-
预训练和微调任务都使用 Lt 损失。因为在预训练中没有自然语言话语,所以去噪器N没有输入Es。在无监督代码生成任务中,Es被每个去噪时间步骤采样的高斯噪声所替代。在CPD任务中,Es通过将被掩盖的代码y通过编码器E计算得出。
推理
- 在推理过程中,用高斯噪声初始化 xt 并在 T 时间步骤上迭代移除(由调度器确定的)比例的噪声,以获得 x ^ 0 \hat{x}_0 x^0。在这个迭代去噪过程中,不使用解码器。在这个迭代过程之后,解码器产生最终预测的代码 y ^ \hat{y} y^。我们对 y ^ \hat{y} y^ 进行后处理,选择到第一个填充 token 为止的 token。
Experiments
评测指标
- 准确度评测:
- bash 生成:模板匹配
- python:CodeBERTScore,该评分已被证明是一个高质量的非执行基代码匹配指标。
- CF:通过在数据列上执行规则并与预期输出比较
- 评估生成的多样性
- 测量生成代码中不同标记 n-gram 的数量除以标记数
- CodeBERT 编码对的成对相似性的摘要统计数据
- 成对字符串编辑距离的摘要统计数据
定量指标
- 引起众多关注的表格,透露了 gpt-3.5-turbo 的参数量为 20B,不过作者已经将该论文撤稿,同时强调这个参数量信息不可靠。总结了 CODEFUSION 及其基准在 top-1、top-3 和 top-5 设置中的性能。对于 Python,只有 GPT-3 (175B) 的表现优于 CODEFUSION (75M)。在 top-3 和 top-5 中,CODEFUSION 超越了所有基线,这与之前的观察一致,即自回归模型在高 top-1 性能的情况下牺牲了生成的多样性
多样性
- 本文方法的多样性更佳
生成代码可执行性
- 本文生成代码语法有效率高于之前的 diffusion 方法
消融实验
-
本文提出各模块的消融,不用 code generation 或者 CPD 都会严重掉点
- grounding(在最后的去噪步骤选择最接近的词汇表标记)或 clamping(在每个去噪步骤选择最接近的词汇表标记)
- grounding(在最后的去噪步骤选择最接近的词汇表标记)或 clamping(在每个去噪步骤选择最接近的词汇表标记)
-
diffusion timestep 消融,最终使用 1200 步。随着扩散时间步的增加,编辑距离减小了。对于 CF 规则来说,下降速度要快得多,因为这些比生成完整程序和 bash 命令更简单。
代码生成样例
- 在 Python 基准测试样本上的连续去噪阶段
Thoughts
- 本文出发点很好,从自回归生成代码的本质问题出发,提出了一种基于 diffusion 的解决思路
- 在更多主流数据集上迭代到超过现有自回归方案,应该能有更大影响力