515. Google AI4Code – Understand Code in Python Notebooks | AI4Code
该解决方案基于双编码器或类似多编码器的方法,即每个单元格单独由 CodeBERT 处理,代码/Markdown 交互通过一对并行的 TransformerDecoderLayer 进行建模:
我要么对代码和 Markdown 单元格使用共享的基于 CodeBERT 的编码器,要么对代码使用 CodeBERT,对 Markdown 使用多语言 BERT。最终的激活值在 token 上取平均。
代码单元的激活值通过零填充的 1D 卷积传递,将 N 个代码单元位置转换为单元之间的 N+1 个位置(用于放置 MD 单元格的 bin),每一点结合了前后单元的激活值。
在将位置编码添加到代码激活值之后(绝对编码和相对编码的结合)。
下一阶段是几层(2 或 6 层)Transformer 解码器,代码单元使用自注意力并将 MD 单元作为内存进行注意力计算,Markdown 也是如此:自注意力以及对代码单元输出的内存查询:
for step in range(self.num_decoder_layers):
x_code = self.code_decoders[step](x_code, x_md)
x_md = self.md_decoders[step](x_md, x_code)
代码和 Markdown 解码器的输出通过矩阵乘法结合,类似于 Transformer 中的注意力机制,产生 3 个输出:
形状为 (nb_code+1) * nb_markdown,表示 Markdown 单元格 i 是否在代码 bin j 之前
形状为 (nb_code+1) * nb_markdown,表示 Markdown 单元格 i 是否位于代码 bin j 处
形状为 nb_markdown * nb_markdown,表示 Markdown 单元格 i 是否在另一个 Markdown 单元格 j 之前
前两个输出用于将 MD 单元格放入代码单元之间的正确 bin 中,最后一个输出用于对单个 bin 内的 MD 单元格进行排序。训练期间使用了 BCE 损失。
利用第一个预测(MD 单元格 i 在代码 bin j 之前的概率),将 MD 单元格放入代码单元之间的正确 bin 中。选择位置的标准是最小化所有 bin 中错误放置概率的总和:
md_cell_bin_num = np.argmin([
-1 * md_after_code[md_idx, :i+1].sum() + md_after_code[md_idx, i+1:].sum()
for i in range(nb_bins)
])
这种选择 Markdown 单元格位置的方法比简单的“Markdown 单元格 i 是否位于代码 bin j 处”输出 + softmax/argmax 效果好得多。这种损失和排序方法背后的直觉是将其作为比赛指标中使用的预期单元格交换次数的代理。
模型是端到端训练的,批次等于单个 notebook。我想由于 notebook 大小可变,这使得优化任务更加困难。梯度裁剪起到了作用。
损失加权和 notebook 采样也对分数有显著影响。较大的 notebook 对最终分数影响较大(因为错误预测会导致大量位置反转),但错误预测误差是在大量的 nb_md * nb_code 预测上取平均的。为了解决这个问题,我更频繁地采样较大的 notebook(最有用),并且对于某些模型添加了 avg_loss + a * max_loss + b * sum_loss(a 非常小,为 0.02,添加 sum_loss 在初始实验中有帮助,但后来被禁用了)。
我尝试解决多语言 notebook 问题,我要么对代码和 MD 单元格使用相同的共享 CodeBERT 模型,要么对代码使用 CodeBERT,对 MD 使用 sentence-transformers/paraphrase-multilingual-mpnet-base-v2。
不同方法的性能差异示例(来自本地验证,其中一个折):