【AI4Code】《CoSQA: 20,000+ Web Queries for Code Search and Question Answering》 ACL 2021

《CoSQA: 20,000+ Web Queries for Code Search and Question Answering》 ACL 2021

类似CLIP 做出了一个 NL-PL的query-key二分类数据集，然后类似CLIP一样双模态对齐训练，在此基础上添加了对比学习，并设计了两类数据扩增方法。双塔的encoder都是CodeBERT。

CoSQA数据集

文章想要实现的事情是我们可以像在网上搜图片一样，根据需求输入query，返回符合需求的代码实现（现在一般是返回博客）。文章花了很大力气构造一个这样的数据集，大概长这个样子。

还有大量的数据构造细节，比如部分满足query的需求，完全满足query的需求，满足少于50%的需求，仅和query有关等。

模型

模型的输入形式各为一条序列：[CLS] xxxxxxxx [SEP]。模型用的孪生网络，query和code都用同一个CodeBERT编码。模型的输出是 [CLS] 的表征。
$${\mathbf{q}}_i=\mathbf{Code}\mathbf{ERT}\left(q_i\right),{\mathbf{c}}_i=\mathbf{CodeB}\mathbf{ERT}\left(c_i\right)$$
模型不是简单的用q和c的内积计算相似度，而是再用一个MLP的计算二者的匹配关系。MLP的输出是向量，不是相似度分数
$${\mathbf{r}}^{(i,i)}=\tanh \left({\mathbf{W}}_1\cdot \left[{\mathbf{q}}_i,{\mathbf{c}}_i,{\mathbf{q}}_i-{\mathbf{c}}_i,{\mathbf{q}}_i⨀{\mathbf{c}}_i\right]\right)$$
再用一个单层NN计算二者的相似度：
$$s^{(i,i)}=\mathrm{sigmoid}\left({\mathbf{W}}_2\cdot {\mathbf{r}}^{(i,i)}\right)$$
然后用BCE loss训练：
$${\mathcal{L}}_b=-\left[y_i\cdot \log s^{(i,i)}+\left(1-y_i\right)\log \left(1-s^{(i,i)}\right)\right]$$

对比学习

除了BCE loss外，In-Batch Augmentation (IBA)和Query-Rewritten Augmentation (QRA)

IBA Loss

对于每一个query，同时选择当前batch内其他的code作为负样本，相当于对一个query扩增了多个code负样本
$${\mathcal{L}}_{ib}=-\frac{1}{n-1}\sum _{\begin{array}{c}j=1\\ j\ne i\end{array}}^n\log \left(1-s^{(i,j)}\right)$$

QRA Loss

由于Web query通常都很短，并且不保证语法正确，因此对于一对标签为1的 query-code pair，对query做一些重写修改，包括：随机删一个单词、随机调换两个单词位置、随机复制一个单词。

这样相当于对一个code扩增了多个Query正样本。在QRA的基础上还会应用IBA loss：
$${\mathcal{L}}_{qr}={\mathcal{L}}_b^{\prime }+{\mathcal{L}}_{ib}^{\prime }$$

实验

本文提出code contrastive learning method (CoCLR) 的是一种学习方法，实验是在预训练好的CodeBERT的基础上继续训练。

两个task一个是Code Question Answering 直接从训练集里划分出测试集，一个是代码搜索。

CodeBERT+CoSQA就是在BERT的基础上显式地对齐两种语言，有一定提升，但是效果最好的还是对比学习的数据扩增。其中最有用的是batch内负样本。query增强中交换单词顺序的提升较大，这也比较符合直觉，因为一般换交两个字往往不影响读阅理解。