导语
- 会议:ICLR 2023在投
- 链接:openreview.net/pdf?id=VPCi…
1 简介
尽管预训练语言模型在代码领域取得了不错的表现,但现有的模型仍具有两个主要的局限:
- 现有的模型大多是Encoder-only或Decoder-only的架构,仅仅适用于部分任务;具体而言,Encoder-only的模型更擅长理解类的任务,Decoder-only的模型更擅长生成类的任务。一些工作采用了Encoder-Decoder的架构来适应各种任务类型,然而这些模型都无法匹敌对应任务上的X-only模型的表现。而且有工作发现,Encoder-Decoder的模型在代码合成任务上无法与GPT样式的模型那样进行in-context learning;
- 现存的模型采用的自监督学习方法并不适合迁移到一些下游任务中。比如,基于T5的模型大多采用span denoising的预训练目标函数,并不适用于自回归形式的代码合成任务。此外,大多数模型没有特定的预训练任务来确保清晰的文本/代码表示学习,而这对于理解文本到代码检索等任务至关重要。
为了解决上述问题,本文提出“CodeT5Mix”,这是一种新的预训练语言框架,用于代码理解和生成(参见图1概述)。具体来说,CodeT5Mix通过不同的阶段训练不同的Transformer组件,并通过参数共享的方式尽量减少参数增加。实验结果表明,CodeT5Mix在几乎所有的任务上都取得了SOTA表现。
2 相关工作
这项工作提出了一种灵活的Transformer混合架构,可以在各种模式下运行:Encoder-only、Decoder-only和Encoder-Decoder。与该工作最类似的工作为yuque.antfin.com/ai-live/ypf…。
3 CodeT5Mix:一个Encoder-Decoder Transformer的混合模型
CodeT5Mix是在CodeT5模型基础上拓展而来的,使用Encoder-Decoder Transformer的灵活混合架构(第3.1节)。在这种混合架构中,作者用增强的学习目标预训练模型,在单模态和多模态数据预训练的两个主要阶段(第3.2节)使用一组不同的自我监督任务。最后,使用预训练的模型,可以传输和激活不同的组件,为不同的下游任务提供不同的功能(第3.3节)。预训练和微调过程的概述可以在图1中看到,更多的细节在图2和图3中。
3.1 模型架构
CodeT5Mix由多个基于transformer的编码器和解码器组成。它们中的每一个都可以操作特定的功能,结合起来,它们可以作为统一的多任务模型:
- Bimodal encoder:通过双向Self-attention将文本或代码片段编码为连续表示。与BERT类似,在输入前加上一个特殊的token [CLS],并将其在最终Transformer层的输出嵌入视为相应输入文本或代码的表示。在面对下游任务时,添加一个线性层,将输出经L2归一化后映射到256维向量。
- Bimodal matching decoder:它的目的是预测文本和代码片段是否具有相同的语义。给定一个代码示例,该解码器首先将其传递给嵌入层和因果自注意层。然后将自注意表示传递给交叉注意层,交叉注意层查询来自文本表示的相关信号(从双模态编码器接收)。一个特定于任务的[Match]令牌被添加到代码输入序列之前,以通知解码器文本代码匹配功能,并且一个[EOS]令牌被添加到代码输入的末尾。最后一个解码器层的[EOS]输出嵌入作为文本代码跨模态对齐表示;
- Unimodal generation decoders:可以用编程语言/自然语言的模态生成输出序列。这些解码器遵循与Bimodal matching decoder相同的设计,具有因果注意层和交叉注意层。当输入是一个文本样本时,使用一个代码生成解码器,并在输入序列的第一个标记前加上一个[CDec]标记。该解码器在代码生成功能中运行。当输入是一个代码示例时,使用一个文本生成解码器,并在输入序列前加上一个[TDec]令牌。该解码器在文本生成(即代码摘要)功能中运行。
注:综上所述,作者一共引入了一个Encoder组件,即Bimodal encoder;两种三个Decoder组件,即Bimodal matching decoder、Unimodal(Text/Code) generation decoders.通过进行组合(或单独使用),将能够应对所有的代码相关任务。
3.2 预训练目标函数
预训练共分为两个阶段,第一阶段只在代码的单模态数据上进行训练,第二阶段则进行双模态训练。
3.2.1 仅在Code单模态数据上预训练
第一阶段使用两种预训练的目标函数(图2),分别是:
- Span Denoising:类似于T5的训练方式,即使用[MASK] token来随机遮盖Code中的一些Span,让模型预测被遮住的部分。
- Causal Language Modeling:类似于GPT的训练方式,即给出前面部分,让模型预测后面的token。
3.2.2 在双模态数据上预训练
第二阶段在双模态数据上进行训练,这些并行的语料来自CodeSearchNet数据集。主要有三种不同的训练方式:
- Text-Code Contrastive Learning:使用对比学习的方式来训练将正样本的文本-编码对的表示靠近在一起并将负样本的文本-编码对的表示分开来,以此构造对齐的文本和代码编码器的特征空间。该任务已被证明对代码理解有效。此任务仅使用Bimodal encoder组件进行训练。
- Text-Code Matching:此任务使用Bimodal matching decoder,其目的是学习更好的双模态表示,以捕获文本和代码形式之间的细粒度对齐。作者在Decoder的输出嵌入之上构建了一个线性层,用于二进制分类任务,它预测文本代码对是正的(匹配)还是负的(不匹配)。
- Text-Code Dual Generation:本任务的重点是通过双重多模式转换实现文本和代码之间的跨模式生成目标:文本到代码生成(代码合成任务)和代码到文本生成(代码摘要任务)。每次转换分别激活相应的Unimodal(Text/Code) generation decoder。该任务是一个有效的学习目标,可以缩小基于生成的下游任务的预训练和微调阶段之间的差距。
3.2.3 使用高效的权重共享策略实现特定任务
虽然只使用一个编码器来处理双模态数据,但三个解码器之间采用了部分共享方案(图3,相同的颜色代表了参数共享,可以看到作者进区分了FFN层)。具体而言,共享自注意层和交叉注意层的参数。作者发现,从这些层学习到的上下文表示应该在文本代码匹配和文本代码双生成任务之间共享,从而实现上下文级别的跨任务泛化。
3.3 下游任务上的微调
通过前面的描述,作者一共设计了1个Encoder、3个Decoder,通过将他们组合拼装,就可以实现不同的功能。
- 理解类任务:只使用Encoder,适配于defect detection和clone detection任务;对于text-to-code retrieval,,则可以加上Bimodal matching decoder。
- Decoder-only任务:使用[CLM]作为Encoder的输入开始部分,类似于第一阶段那样;
- Seq2seq生成任务:使用Encoder和对应的单模态Decoder组件(输出要求是Text的用Text Decoder,是Code的用Code Decoder)。
4 实验
作者在广泛的任务上进行了实验,
同时,作者进行了消融实验来对比不同组件的作用。
最后,作者展示了所提出的模型可以进一步用于检索式增强的代码生成。
5 总结
本文提出了CodeT5Mix,这是一种新的基于代码的预训练模型,它混合了编码器-解码器Transformer模块,可以灵活组合用于广泛的代码理解和生成任务。为了训练CodeT5Mix,作者引入了一组不同的预训练任务,包括去噪、因果语言建模、对比学习和匹配,以从纯代码数据和双模态代码-文本数据中学习丰富的表示。作者设计了一种有效的与任务特定的FFN的权重共享策略,以减少任务间的干扰,同时保持模型参数的可承受性。在7个下游任务上的大量实验验证了该模型的优越性。同时,实验结果进一步展示了CodeT5Mix可以完全支持半参数检索增强代码生成方法,因为该模型可以有效地在代码检索和生成任务中运行。