导语
- 会议:ACL 2022
- 链接:arxiv.org/abs/2203.03…
1 简介
已有的三种类型的Code LLM(Encoder-only、Decoder-only、Encoder-decoder)各自擅长特定的任务,本文提出了UniXcoder,一个统一的跨模态预训练的编程语言模型,支持与代码相关的理解和生成任务。UniXcoder利用带有前缀适配器的掩码注意矩阵来控制每个token对上下文的访问。与目前统一的编码器解码器模型相比,UniXcoder可以更好地应用于自回归任务,例如代码补全更加高效。除了将代码作为唯一的输入,作者还考虑了多模态内容,如代码注释和抽象语法树(AST)来增强代码表示。
本文提出了一种一对一映射的方法,将AST转换为一个保留树中所有信息的序列结构,然后将该序列作为输入来增强代码表示。UniXcoder使用三种语言建模任务进行预训练:MLM(BERT采用的)、自回归式语言建模(GPT采用的)和去噪语言建模(T5、BART等采用的),这可以使模型支持各种类型的下游任务。此外,作者引入了两个预训练任务来学习可以表示代码片段语义的嵌入。一种是多模态对比学习,它利用AST来增强代码片段嵌入的语义,另一种是跨模态生成,它利用代码注释来对齐编程语言之间的嵌入。
在理解和生成任务上的实验结果表明,UniXcoder在大多数任务中都达到了SOTA。进一步的分析表明,AST和代码注释都可以增强UniXcoder以更好地捕获代码语义。
2 相关工作
略
3 UniXCoder
3.1 输入表示
图1给出了一个AST的示例,其中表示了代码中丰富的结构信息。然而,AST通常表示为树,不能直接用作Transformer的输入。为了与代码注释并行编码AST,作者提出了一个在算法1中描述的一对一映射函数F,将AST转换为保留所有结构信息的序列。
算法整体流程如下:首先将树的root放入函数,接着函数将递归式的将中间节点和叶子节点添加到序列中,以“parameters→(data)”为例,映射函数F将子树转换为“<parameters,left> (data) <parameters,right>”。作者证明,这个映射算法是一对一可逆的。
模型的最终输入是注释W和序列化后的AST树F(T(C))。
3.2 模型结构
图2展示了整个模型的架构,该模型将N个Transformer层应用于代码注释和带前缀的序列化AST以产生隐藏状态H,其中前缀p∈{[Enc], [Dec], [E2D]}
表示模型的行为。
如图所示,模型通过修改Attention中的Mask矩阵来实现不同的行为:
- Encoder:将Mask矩阵中所有值设为0;
- Decoder:将矩阵的上三角部分值设为负无穷;
- Encoder-Decoder:将所有输入的token对应的mask设置为0,输出则按照解码的顺序来设置。
作者也采用三种语言建模的目标函数来进行学习:
- Masked Language Modeling:即BERT采用的方式;
- Unidirectional Language Modeling:即GPT采用的方式;
- Denoising Objective:即T5、BART等采用的方式。
除此之外,作者还新提出了两个目标函数用于Code Fragment Representation Learning,即:
- Multi-modal contrastive learning:采用了一个对比学习损失,使用不同的dropout对相同的输入进行forward得到的表示作为正例,同一batch中的其他表示作为负例。
- Cross-modal generation:即在输入AST序列的情况下,生成对应的注释文本。
由于引入了AST,AST展开后的序列导致输入增加了大量额外的tokens(70% longer)。因此,在微调阶段UniXcoder仅使用AST的叶子节点,但这样会造成训练和验证数据形式不一致。为了缓解这个gap,在预训练阶段作者设置了0.5的概率随机丢弃输入序列中的非叶子节点。
4 实验
作者在许多任务上取得了SOTA的表现,
为了进一步评估代码片段表示的性能,作者还提出了一种称为Zero-shot Code-to-Code Search的新任务。给定一个源代码作为查询,该任务的目标是在Zero-shot设置下从候选集合中检索具有相同语义的代码。该任务可以通过检索具有相同语义的源代码,帮助用户从一个PL转换到另一个PL。
消融研究进一步表明,多模态数据和代码片段表示预训练任务都可以增强UniXcoder。
5 总结
为了同时支持与代码相关的理解和生成任务,本文提出了UniXcoder,这是一个统一的预训练模型,它集成了来自代码注释和AST的语义和语法信息。本文提出了一种一对一的映射方法来将AST转换为序列结构,并提出了两个新的预训练任务来学习代码片段表示。为了进一步研究代码表示的性能,本文提出了一个新的零镜头代码到代码搜索的下游任务,并为此任务创建了一个数据集。实验表明,UniXcoder在大多数任务上的性能明显优于以往的工作。进一步的消融研究也表明,AST和代码注释都可以增强UniXcoder,并揭示了提出的映射函数和预训练任务的有效性。