文献阅读笔记 # CodeBERT: A Pre-Trained Model for Programming and Natural Languages

• 《CodeBERT: A Pre-Trained Model for Programming and Natural Languages》EMNLP 2020 (CCF-B)
• 作者主要是来自哈工大、中山大学的 MSRA 实习生和 MSRA、哈工大的研究员。
• 资源：code | pdf
• 相关资源：RoBERTa-base | CodeNN
• 词汇： bimodal: 双模态/双模；unimodal: 本文应该要理解为单模态；

Abstract

本文提出用于程序语言(PL)和自然语言(NL)的双模态预训练模型 CodeBERT。CodeBERT 学习支持下游任务(如: natural language code search, code documentation generation)的通用表示，然后基于混合的目标函数对其进行训练，该目标函数包含预训练任务和 replaced token detection(detect plausible alternatives sampled from generators) 任务。这使得我们能利用 NL-PL 对的 “bimodal” data 和 “unimodal” data。前者为模型训练提供 input tokens，后者用来帮助学习更好的 generators。评估：

• 在 natural language code search, code documentation generation 任务中取得了 2020 年 SOTA。
• 在 zero-shot setting 和 fix model params 的条件下，在本文构建的 NL-PL probing dataset 上测试取得 2020 年 SOTA。

【个人理解】这里的 RTD 任务类似 GAN，数据是 mask 的数据，然后通过 generator 生成可能的候选方案，然后再用 discriminator 判断候选方案是否正确。

1 Introduction

大规模预训练模型，如 ELMo、GPT、BERT、XLNet、RoBERTa 等模型，极大提高了各种 NLP 任务的 SOTA。这些预训练模型通过自监督目标在大规模无标签的文本中学习有效的上下文表示，如 masked language modeling 从人工屏蔽的输入序列中预测原始词。预训练模型的成功也推动了多模态预训练模型的发展，如 ViLBERT、VideoBERT 等。

本文提出的 NL-PL 双模态预训练模型 CodeBERT 在包含6种编程语言的数据中进行预训练，其中的双模态数据是指代码和其对应的函数级别的自然语言形式的文档。模型训练在类似多语言 BERT 的设置下进行，输入的时候不会特别说明程序语言的类型。

本文主要贡献：1）首个面向多编程语言的 NL-PL 预训练大模型；2）CodeBERT 在自然语言代码搜索和代码文档生成任务中都取得了当前(2020年) SOTA。3）创建了用于研究代码预训练模型的 probing 能力的数据集。

2 Background

2.1 Pre-Trained Models in NLP

• self-supervised：无需人工标注自动获得用于预训练的监督信息；

• 主要的学习目标是 language modeling and its variations；

  • GPT：给定前文 {w1,w2,...,wk−1}\{w_1,w_2,...,w_{k-1}\}{w1​,w2​,...,wk−1​} 预测下一个词 $w_k$。

• 由于预训练的最终目标不是训练一个好的语言模型，因此最好同时考虑前后文，以学习更好的通用上下文表示。这导致 BERT 采用了 masked language modeling objective。
  【说实话，论文这里的表述方式我觉得很怪，这里真的有因果关系？怎么去定义“好”的语言模型？为什么不是为了得到“好”的语言模型就要考虑前后文？好的语言模型是指能预测下文的能力才叫好吗？】

2.2 Multi-Modal Pre-Trained Models

多模态预训练模型：学习不同模态输入之间的隐式对齐。本文不但包含 NL-PL 的双模态数据，还包括大量的单模态数据（不带文档的代码）；

• ViLBERT ：
  • 数据： image caption data(language-image)；
  • 自监督任务：reconstructing categories of masked image region or masked words 同时预测标题是否描述了图像的内容；
• VideoBERT：
  • 数据：language-video；
  • 自监督任务：trained by video and text masked token prediction；

同期一项工作：采用 masked language modeling 和 next sentence prediction 作为训练 BERT 的目标，并在 Python 源码上训练模型。 这里的 sentence 指：Python 标准定义的 logical code line。在预训练过程上 CodeBERT 与上述工作的差异是：1）CodeBERT 以跨模态方式进行训练，同时利用双模态 NL-PL 数据和单模态 PL/NL 数据；2）CodeBERT 在6种编程语言上进行预训练；3）CodeBERT 基于一种新的 replaced token detection 目标进行训练。

3 CodeBERT

3.1 Model Architecture

CodeBERT: multi-layer bidirectional Transformer，结构与 RoBERTa-base 完全相同，因此不再赘述，模型参数共 125M。

3.2 Input/Output Representations

预训练阶段的输入：$[CLS],w_1,w_2,..,w_n,[SEP],c_1,c_2,...,c_m,[EOS]$；(NL Text + PL Code)

• $[CLS]$ 开头特殊标记，其 hidden representation 通常被认为是用于分类或排序的 aggregated sequence representation；
• 按照 Transformer 处理文本的方式，将 NL 视为单词序列，将其切分为 WordPiece；此处将一段代码看作 token sequence；

输出：1）对 NL/PL 输出每个 token 的 contextual vector representation；2）$[CLS]$ token 作为 aggregated sequence representation；

3.3 Pre-Training Data

分别用 bimodal/unimodal data 训练 CodeBERT；

• bimodal：parallel data of natural language-code pairs
• unimodal：codes without paired natural language texts and natural language without paired codes

数据：datapoints from Github Repos

• bimodal datapoint：带有对应注释的函数
• unimodal code：不带注释的函数

用于预训练的数据集的构造方法 爬取公开非 fork 的 github repo，根据下述规则过滤：1）至少被一个别的项目使用；2）each documentation is truncated to the first paragraph（看下面的图）；3）注释不能少于 3 tokens；4）函数不能小于三行；5）函数名不能带有 test；

3.4 Pre-Training CodeBERT

使用了两个预训练任务来预训练 CodeBERT：

• masked language modeling (MLM)：应用于 bimodal data of NL-PL pairs；
• replaced token detection (RTD)：使用大量 unimodal data，如无注释的代码；

Objective #1: Masked Language Modeling (MLM)
输入: datapoint of NL-PL pair x={w,c}x=\{w,c\}x={w,c}；$w$：NL words sequence；$c$：PL tokens sequence；
首先为 NL、PL 随机选择约 ~15% 的 token 替换为 $[MASK]$。

MLM 的目标是预测 masked position 处的 original token，下面的 $p^{D_1}$ 是判别器 discriminator，从大量的词汇表中预测出一个 token：

Objective #2: Replaced Token Detection (RTD)

RTD 目标最初是用于高效学习自然语言预训练模型。本文同时采用双模态和单模态数据用 RTD objective 进行训练。这里有两个 generators：NL generator $p^{G_w}$，PL generator $p^{G_c}$ 都用于对 masked positions 生成 plausible alternatives。

判别器 discriminator 被训练来判断一个词是否是原始词，这是一个二分类问题。对输入的每个位置，都应用 RTD objective，这与 GAN 的区别是：如果 generator 恰好生成了正确的 token，该 token 的 label 是 “real” 而不是 “fake”。【因为 GAN 里是通过 generator 和 discriminator 左右互搏，造 fake 和 鉴别 fake 来实现模型的提升，所以 generator 的目标是生成 fake 数据】

$\delta (i)$：指示函数；$p^{D_2}$：预测 i-th word 是原始词的概率的判别器。

有很多方式实现 generator，本文实现了2个具有 bidirectional contexts 的 n-gram 语言模型，分别用于 NL、PL，也分别从相应的数据点学习。这个方法也很容易推广到学习双模态生成器或以一种联合方式使用像基于 Transformer 这样的神经架构的复杂生成器，这部分作为未来可以拓展的工作。

当前的 PL 训练数据是单模态代码数据，NL 训练数据来自双模态数据中的代码文档。最终的 Loss 函数：

3.5 Fine-Tuning CodeBERT

将 CodeBERT 应用到下游的 NL-PL 任务：

• natural language code search：用与预训练相同的方式输入数据，用$[CLS]$的表示来度量代码和自然语言查询中的语义相关性。
• code-to-text generation：使用 encoder-decoder 结构，用 CodeBERT 初始化生成模型的编码器。

4 Experiment

4.1 Natural Language Code Search

Natural Language Code Search 任务：通过输入的自然语言查询代码集合中语义最相似的代码；

• 数据集：CodeSearchNet corpus
  • 训练数据集和验证数据集都平衡了正负样本。负样本也平衡了随机替换 NL $(c,\bar{w})$ 或者随机替换 PL $(\bar{c},w)$的两类样本。
• 评估方式：通过一组固定的 999 distractor codes 来为每对测试数据 $(c,w)$ 计算 Mean Reciprocal Rank (MRR) 指标。然后进一步计算测试数据中的所有编程语言的宏观平均的 MRR 来作为整体评估指标。
  注意：这个指标与原始论文中的 AVG 指标不同，原始论文中的 answer 从6种候选语言中获得。但本文对每种编程语言都微调了特定于语言的模型。使用二分类 loss 训练每个模型，其 softmax 层连接到 $[CLS]$ 的表示层。详细的微调超参数见附录 B.2。

Model Comparisons：

上面的模型都通过把代码视为 tokens sequence 来进行预训练。本文也以 MLM 任务只在 code 上训练 RoBERTa。CodeBERT(MLM) 从头开始学习的性能超越了 RoBERTa。用 RoBERTa 初始化 CodeBERT 可以提高性能。

4.2 NL-PL Probing

Task Formulation and Data Construction：在相关研究 NLP probing 实验之后，本文研究 NL-PL probing。由于目前没有面向该问题的工作，因此本文创造了数据集，并形式化该问题。

对给定的 NL-PL 对 $(c,w)$，NL-PL probing 的目标是测试模型在干扰下正确预测/恢复 masked token of interest 的能力(可能是$c_i$，也可能是$w_j$)。干扰有两种：

• 用于 MLM 目标的整体目标词典；
• 由专家对被测能力的理解而筛选或设计的少量候选样本（完型填空）；

本文采用第二种干扰，将 NL-PL probing 任务设计为有多个选项的 QA 问题。问题是完型填空，一个特定的 token 被 $[MASK]$ 替换，然后专家设计了一些候选答案。具体而言，本节分别在 NL 侧和 PL 侧进行评估。而为了减少数据构建的工作量，自动从 CodeSearchNet 数据集中的验证集和测试集中的 NL-PL pairs 获取数据（该数据在预训练阶段不可见）。

为了在 NL 方面评估，选择了 NL-PL pairs 中 NL docs 包含6个关键词 (max,maximize,min,minimize,less,greater) 之一的 pair，然后合并1/2关键词和3/4关键词的 pairs，将数据分为4组候选。任务是让预训练模型从4个选项中选择1个正确的选项。即，该设置的输入包括完整的代码和带有 mask 的文档。

对于 PL 方面，选择包含关键字 max 和 min 的代码，并将任务设计为二选一问题。由于代码补全是一个重要的场景，我们想测试模型仅仅基于前文的 PL context 预测正确 token 的能力。因此，此处为 PL 添加了一个额外配置，其输入是完整的 NL 文档和前面的代码。

Model Comparisons：这里采用 CodeBERT(MLM) 因为其输出层自然契合 probing。结果表明在 NL 和 PL probing 上 CodeBERT 都比基线方法更好，只有在 proceding context 设置下更差，这表明 code completion 是一个挑战性的任务，这项研究留给未来的工作。

本节进一步的对 NL-PL probing 进行 case study。分别 mask NL token、PL token，然后报告 RoBERTa 和 CodeBERT 的预测概率。我们可以看到 RoBERTa 在下图的两种情况下都失败了，而 CodeBERT 在 NL 和 PL 设置中都做出了正确的预测。

4.3 Code Documentation Generation

CodeBERT 的预训练目标不包括 generation-base objective，但本文想调查 CodeBERT 在生成任务上的扩展能力。本节研究 code-to-NL generation 并报告在 CodeSearchNet Corpus(6种编程语言) 上的文档生成任务的结果。因为生成的文档很短并且 higher order n-grams may not overlap，本节通过使用 smoothed BLEU score 解决了该问题。

Model Comparisons：baselines model 如带有 attention 机制的 RNN 模型、Transformer、RoBERTa 和仅在代码上预训练的模型。为了证明 CodeBERT 在 code-to-NL generation 任务上的有效性，采用了不同的预训练模型作为编码器并保持超参数一致（详细的超参数见附录B.3）。

• 可以发现在 PL 上预训练的模型效果好于 RoBERTa，这证明在程序语言上预训练的模型可以改善 code-to-NL generation。

4.4 Generalization to Programming Languages NOT in Pre-training

本节评估在预训练中未见过的编程语言上模型的泛化性，选择的任务：natural language summary of a C# code snippet。在 CodeNN 论文中的数据集上进行实验，包含了 66015 个来自 StackOverflow 的 QA 对。这个数据集由于数量级小，因此具有挑战性。采用 smoothed BLEU-4 score 评估模型。

Model Comparisons：但这里效果仅次于 Code2Seq，可能原因是因为 Code2Seq 使用在其 AST 上使用了 compositional paths。但 CodeBERT 只将原始代码作为输入。本文通过按照一定顺序遍历 AST 的树结构来训练 CodeBERT，但这样做并没有对生成任务带来改进。因此结合 AST 来改进 CodeBERT 也是未来研究的一个潜在方向。

5 Conclusion

• 本文提出首个面向 NL-PL 的双模态预训练模型 CodeBERT，并基于 bimodal 和 unimodal 数据共同训练，微调后的模型在自然语言代码搜索、代码文档生成等下游任务中取得了当前的 SOTA。
• 为了进一步研究预训练模型中蕴含的知识，设计了 NL-PL probing 任务，并创建了用于该任务的数据集。将 probing 任务看作完型填空问题(cloze-style answer selection problem)，并对 NL 和 PL 的部分设计了干扰项。结果显示，当模型参数固定不变时，CodeBERT 优于 RoBERTa 和只使用代码 continuously trained 的模型。

潜在研究方向

• 可以根据 bimodal evidence 学习更好的生成器或者更复杂的神经网络结构来改善 replaced token detection objective。
• CodeBERT 的 loss 函数主要针对 NL-PL understanding tasks，尽管在 code-to-documentation generation 任务上取得了很高的 BLEU scores，但还可以通过修改为与 generation-related learning 相关的目标函数来提升。
• 将 AST 纳入预训练也是一个值得探索的方向。
• 将 CodeBERT 应用到更多的 NL-PL 相关的任务中，扩展到更多编程语言，获得更好的泛化性： 探索灵活和强大的 domain/language adaptation 方法。

Appendix

A Data Statistic

B Train Details

B.1 Pre-training

• train CodeBERT on one NVIDIA DGX-2 machine using FP16；16 卡 NVIDIA Tesla V100；32GB内存；
• batchsize: 2048，learning rate：5e-4；Adam Optimizer；warmup steps：10K；max sequence length：512；max training step：100K；

B.2 CodeSearch

• fine-tuning 时 learning rate：1e-5, batch size: 64, max sequence length: 200, max fine-tuning epoch: 8;
• Adam Optimizer

B.3 Code Summarization on Six Programming Languages

• Transformer with 6 layers, 768 dimensional hidden states and 12 attention heads as our decoder in all settings；
• max length of input and inference as 256 and 64;
• Adam optimizer, learning rate: 5e-5,batch size: 64;
• perform early stopping on the development set

B.4 Code Summarization on C#

• 2-layer GRU with an attention mechanism as our decoder for a comparison
• fine-tune models using a grid search with the following set of hyper-parameters: batchsize is in {32, 64} and learning rate is in {2e-5, 5e-5}

C Learning Curve of CodeSearch

D Late Fusion

本节分析 CodeBERT 是否适合作为 unified encoder。CodeBERT 首先对 NL 和 PL 分别进行编码，然后通过内积计算相似度，这样代码搜索就相当于在共享向量空间中寻找最近的代码。这样做还有助于预先计算代码表示的线上系统中使用 CodeBERT。在运行时，系统只需要计算 NL 的表示和基于向量的点积。根据以下目标对 CodeBERT 进行微调：最大化 ground truth 的内积，同时让干扰项的内积最小化。

我们只在数据量少的语言上使用了这个设定。可以看到 CodeBERT 比 RoBERTa 和仅用代码预训练的模型表现得更好。late fusion 与 standard 方法效果相当，但更有效，因为这种设定可用于线上系统。

E Case Study