CODEFUSION: A Pre-trained Diffusion Model for Code Generation

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=s1,s2,\dots ,sk$，输出是预测的代码片段 $y^{=}{y}^1,y^2,\dots ,y^s$。输入和输出都被填充到固定的维度 n。
• CODEFUSION有三个主要的基于 transformer 的组件（一个编码器E，一个去噪器N，一个解码器D）和一个分类头H。
  • 编码器（E，预训练的 CodeT5）：将分词话语 s 转换为向量表示 $Es=E(s)=e1,e2,\dots ,en$。
  • 去噪器 （N，10层 transformer 模块）：在编码的话语 Es 和时间 t 的条件下，去噪器（N）预测并移除噪声 εt，​ 从噪声程序嵌入 xt​ 获得预测的去噪程序嵌入 $x^0=N(xt,t,Es)$。
    • N 是一个带有 xt​ 和 Es 之间的交叉注意力以及 xt​ 上的完全自注意力的 transformer 模块。
  • 解码器（D，6 层 transformer 模块）：应用交叉注意力到 $x^0$​ 和 Es，以及在 $x^0$ 上的完全自注意力，来计算最终的隐藏表示 $Ds=d1,d2,\dots ,dn=D(x^0,Es)$。与以前的文本扩散方法不同，在那里标记是独立生成的，完全自注意力允许每个隐藏维度（$di$​）在关于其他维度的完整信息中生成。
  • 分类头 （H，单个全连接层）：Ds 被投影到实际的代码 token 上，通过一个分类头 H 计算代码 token 的分布 $p(y\mid di)$。我们不执行对这些 token 的搜索，并为每个 i 选择 $y^i=argmax{}_{y}p(y\mid di)$。

训练

• 两阶段训练：

  • 在代码片段上无监督预训练去噪器和解码器
  • 在（自然语言，代码片段）上有段度微调编码器、去噪器和解码器

• 根据之前在文本扩散上的工作，我们使用一个可训练的嵌入层 L 来将代码片段 y 嵌入到一个连续的空间，在时间步 t 我们可以添加（并移除）噪声 εt​。我们从之前在文本扩散上的工作中获得灵感，并通过结合来自解码器的隐藏表示 Ds 将来自 GENIE 的损失适应到 CODEFUSION。在时间步 t，损失计算为
  

  • 我们通过最小化预测噪声 ${\hat{\epsilon }}_t$​ 与实际噪声 ${\epsilon }_t$​ 之间的误差来训练去噪器 N。
  • 我们通过最小化隐藏表示 Ds 与嵌入代码之间的误差来训练解码器 D 和嵌入层 L。
  • 我们在分类头的输出上应用标准交叉熵损失，该头根据 Ds 产生预测的代码标记，并与真实代码片段 y 进行比较。损失函数允许我们使用扩散目标训练模型的三个主要组件（去噪器、解码器和分类头）。

• 预训练阶段的两个任务

  • 无监督的代码生成
  • 对代码的连续段落去噪（CPD）：这个特定于代码的CPD任务只掩盖与目标语言的标识符或内置关键词相关的标记

• 预训练和微调任务都使用 Lt 损失。因为在预训练中没有自然语言话语，所以去噪器N没有输入Es。在无监督代码生成任务中，Es被每个去噪时间步骤采样的高斯噪声所替代。在CPD任务中，Es通过将被掩盖的代码y通过编码器E计算得出。

推理

• 在推理过程中，用高斯噪声初始化 xt​ 并在 T 时间步骤上迭代移除（由调度器确定的）比例的噪声，以获得 ${\hat{x}}_0$​。在这个迭代去噪过程中，不使用解码器。在这个迭代过程之后，解码器产生最终预测的代码 $\hat{y}$​。我们对 $\hat{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（在每个去噪步骤选择最接近的词汇表标记）
    

• diffusion timestep 消融，最终使用 1200 步。随着扩散时间步的增加，编辑距离减小了。对于 CF 规则来说，下降速度要快得多，因为这些比生成完整程序和 bash 命令更简单。
  

代码生成样例

• 在 Python 基准测试样本上的连续去噪阶段

Thoughts

• 本文出发点很好，从自回归生成代码的本质问题出发，提出了一种基于 diffusion 的解决思路
• 在更多主流数据集上迭代到超过现有自回归方案，应该能有更大影响力