背景
题目:ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators
机构:斯坦福、Google Brain
作者:Kevin Clark, Minh-Thang Luong, Quoc V. Le, Christopher D. Manning
论文地址:https://www.aminer.cn/pub/5e5e18d593d709897ce33b3c/electra-pre-training-text-encoders-as-discriminators-rather-than-generators
摘要
MLM(Masked language modeling)方式的预训练语言模型如BERT是在输入上用[MASK]遮蔽掉部分tokens,再训练一个模型以重建出原始的tokens。这种方式迁移到下游NLP任务时能够得到较好的结果,已然成为NLP任务中的标配。但是这种预训练方法往往需要大量的算力。
为此,本文提出一种样本效率更为高效的预训练任务:替换token检测(replaced token detection, RTD)。RTD不对输入进行遮蔽,而是从生成器中采样得到可信的tokens,再替换掉原始输入上的tokens。本文不再训练模型预测输入中的原始tokens,而是训练一个判别器以辨别输入中每个token是否是由生成器生成。
通过实验表明这种新的预训练任务比MLM更高效,该任务定义于全部的输入tokens,而非仅仅被遮蔽掉的那一部分小小的输入子集。实验结果表明本文方案所学习到的上下文表征大大优于相同模型大小、相同数据量和相同算力下的BERT。这种收益在小模型上更加显著,比如在GLUE自然语言理解基准上,单块GPU训练4天的结果优于GPT(算力比其高30倍+)。
本文的方法在大模型上也很有成效,只使用RoBERTa和XLNet四分之一算力即可取得与之比肩的结果。
方法替换Token检测(RTD)任务如Figure 2所示。本文方法需要训练两个神经网络分别是生成器G和判别器D。两者都包含了一个encoder,比如Transformer。encoder将输入token序列映射为一序列上下文表征向量。
给定位置t,生成器用softmax层输出x_t概率分布:其中e表示token的嵌入表征。给定位置t,判别器判断token x_t是否由生成器生成,其输出层是sigmoid:训练的生成器是为了执行MLM。给定输入列,MLM先选择一些随机的位置进行遮蔽,遮蔽集合记为m。这些被选中的tokens被用[MASK] token替换。生成器学习预测出被遮蔽掉的原始token。
至于训练的判别器则是为了能够鉴别出token是否是来自生成器。通过生成器生成的token替换原始字符,再训练判别器鉴别出哪些token和原始输入相匹配。总结下,模型根据以下式子重建输入:其中REPLACE是定义的替换操作。第一个是用[MASK]替换,第二个是用生成器生成的字符替换。
对应的损失函数分别如下:尽管看似GAN的训练目标函数,但是存在以下几点关键不同点:第一,当生成器碰巧生成正确的token,则该token被视为真实的,而非假冒的。实验发现这种方式可以一定程度上地改善下游任务。第二,生成器的训练是基于最大似然而非对抗训练。后者试图在训练过程中骗过判别器。对抗训练在这里很难实现,究其原因在于无法通过从生成器采样来进行反向传播。
尽管本文也尝试使用强化学习的方式训练生成器从而绕开这个问题(附件F),但是效果仍不及最大似然的训练方式。第三,不像GAN那样在生成器的输入中添加噪声。
最终最小化联合损失:实验过程中用单样本来近似损失的期望。同时不把判别器的损失反向传播到生成器,采样的步骤也做不到。在预训练之后,丢弃生成器,并在下游任务中对判别器进行微调。
实验
实验设置
为验证文本方法的有效性,在GLUE和SQuAD上进行实验。预训练数据集分两种情况:第一种是和BERT对比,则采用Wikipedia和BooksCorpus数据集,合计大概3.3 Billion tokens;第二种是和XLNet对比,则与其相同的数据集进行预训练,即ClueWeb、CommonCrawl和Gigaword,合计33 Billion tokens。
本文的模型架构和超参数总体上和BERT相同。在微调阶段,对于GLUE则在ELECTRA的顶部增加一个简单的线性分类器;对于SQuAD则借用XLNet中的问答模块,将其加到ELECTRA的顶部。之所以这部分借用XLNet,是由于该模块比BERT更精细,起始位置和终点位置是联合预测的,而不像BERT那样两者各自独立,另外对于SQuAD2.0该模块还有能否回答的分类器。
需要注意的是,一些评估数据集很小,这意味着微调模型的准确性可能会随着随机种子的不同而发生很大的变化。因此,最终报告结果是从相同的预先训练的检查点运行10次微调的中位数。
模型拓展
为了改善模型本文在模型中进行了以下几点拓展。除非特别说明,否则实验对比中模型大小和训练数据集均与BERT-Base相同。
第一,权重共享。生成器和判别器权重参数共享可以提高预训练的效率。当生成器和判别器大小相同,则二者的所有transformer权重都可以密切相连。但是,实验中发现较小的生成器更为高效,这使得只能共享生成器和判别器的嵌入参数(包括token嵌入和位置嵌入)。这种情况下,使用的嵌入大小是根据判别器隐状态的嵌入尺寸。在生成器中添加一个线性变换层以将嵌入投射为生成器隐状态尺寸的表征。
实验对比了生成器和判别器相同大小的情况。在GLUE中,权重完全无关时的得分为83.6,共享token嵌入时的得分为84.3,共享所有权重时的得分为84.4。可以猜想本文的模型得益于共享token嵌入的方案,可能是MLM在学习以下表征上更为有效:仅当输入中有生成器生成的token时判别器才会更新。生成器的softmax在整个词典上密集更新所有的token嵌入。
完全关联生成器和判别器的权重几乎没有什么提升,却需要生成器和判别器大小完全相同,这得不偿失。
基于上述这些发现,最终选用共享嵌入参数。
第二,较小的生成器。如果生成器和判别器大小相同,则训练ELECTRA每步需要的计算量大约是只使用MLM训练的两倍。为此,本文使用较小的生成器以降低该因素的影响。具体是降低层的大小,同时保持其他超参数恒定不变。
此外,本文还尝试了一个极简的unigram生成器,该生成器根据训练语料中的频率采样生成假的tokens。不同大小的生成器和判别器在GLUE上的得分如Figure 3中左图所示。从中发现生成器最佳的尺寸是判别器的1/4-1/2。这可能是由于当生成器过于强大,则对应的任务对判别器来说越难。
第三,训练方法。本文还进一步探索了ELECTRA的其他训练方法。尽管这些探索性工作并没有最终改善结果。探索的训练方法使用以下两阶段的训练流程:1. 只训练生成器,训练n步。2. 用生成器的权重初始化判别器。再训练判别器,训练n步,同时固定住生成器的权重。
上述提到的权重初始化要求生成器和判别器的大小相同。实验中发现,如果判别器没有这个权重初始化操作,那么有时会甚至在大多数类之外学习失败。
这可能是由于此时的生成器已远远领先于判别器。联合训练的方式另一方面天然地为判别器提供了一个总纲,在这个总纲中,生成器一开始很弱,但在整个训练过程中变得越来越好。此外本文还探索了如同GAN那般使用对抗方式训练生成器,具体使用强化学习来适应离散操作的采样生成器。具体参考附件F。
各种训练方法的对比如Figure 3右图所示。
从中可以看出,对于两阶段的训练方法,当将生成目标转为判别目标后,下游任务的性能得到显著提升,但是最终没有超过联合训练方案。至于对抗训练方案,尽管超过了BERT,但是仍然逊色于最大似然训练方法。这之间的差距可能源于以下2个因素:第一,对抗方式训练生成器在MLM上表现更差,只能取得58%的准确率,而最大似然的训练方案可以取得65%的准确率。
准确率较差的主要原因是在生成文本的庞大动作空间中,强化学习的样本效率较低。第二,对抗方式训练的生成器产生一个低熵输出分布,其中大部分概率集中在单个token上,这意味着生成器生成的样本中没有太多的多样性。
Small版模型先在Small版模型上对比本文预训练方案的有效性。
Small版模型与BERT-Base的超参数基本相同,只是将序列长度由原来的512缩短为128,batch size由原来的256降低为128,隐层维度由原来的768缩减为128。为公平起见,用上述超参数也训练一个BERT-Small。各个模型的Small版在GLUE dev数据集上的对比详情如Table 1所示。可以看出,在模型大小相同时ELECTRA-Small显著优于其他模型。
比如ELECTRA-Small高出BERT-Small大概5分,甚至高于更大的GPT模型。ELECTRA-Small的训练多数是收敛的,即使模型的训练时间更短(只有6个小时),仍然可以获得合理的性能结果。实验结果还表明ELECTRA在中度大小的模型上依然有效,ELECTRA-Base显著优于BERT-Base,甚至超越BERT-Large(GLUE得分为84)。
Large版模型
此外,训练ELECTRA-Large以验证RTD预训练任务在大规模先进的Transformers模型中的有效性。ELECTRA-Large和BERT-Large大小一致,但是后者训练耗时更久。ELECTRA-400k模型仅训练400k个steps,大概是RoBERTa的1/4;ELECTRA-1.75M模型则是训练了1.75M个steps,所耗算力与RoBERTa相当。
使用2048的batch size,数据集是来自于XLNet的预训练数据集。Large版模型在GLUE dev数据集上的实验结果如Table 2所示。可以看出ELECTRA-400K能够取得与RoBERTa、XLNet相比肩的性能结果,但所需的算力仅仅是RoBERTa和XLNet的1/4。这表明ELECTRA的采样效率增益在大模型上保持不变。
训练得越久的ELECTRA,比如ELECTRA-1.75M在大部分的GLUE任务上都优于其他模型且所需的预训练算力更少。令人惊讶的是,基线BERT模型得分明显低于RoBERTa-100k,这表明我们的模型可能受益于更多的超参数调优或RoBERTa训练数据。ELECTRA在GLUE test数据集上仍能维持收益,具体见Table 3。各个模型在SQuAD数据集上的表现如Table 4所示。
实验结果表明在相同算力资源下,ELECTRA的得分高于基于MLM的模型。比如ELECTRA-400K优于RoBERTa-100k和BERT基线模型(均使用相近的预训练算力)。ELECTRA训练得越久,结果能够进一步提升:在SQuAD数据集上ELECTRA-1.75M超出此前的其他模型。
ELECTRA-Base不仅大大优于相同模型大小的BERT-Base和XLNet-Base,甚至在多数指标上超越BERT-Large。ELECTRA在SQuAD 2.0上的表现优于SQuAD 1.1。这可能得益于RTD任务的引入,在RTD任务中模型能够辨别token的真伪,这对于有是否可回答类型问题的SQuAD 2.0数据集来说更具迁移能力。
有效性分析
为了进一步理解ELECTRA模型收益来源,对比了以下几种预训练目标:
ELECTRA 15%:这与ELECTRA大致相同,仅有的不同的在判别器的损失仅来自于输入中被遮蔽的15% tokens。换句话说,$
ext{L}_{ ext{Disc}}$中$i$的不再是1-n,而是$i ext{in} m$。
Replace MLM:该目标与MLM大致相同,所唯一不同的是不使用[MASK]替换tokens,而是用生成器生成的token替换。这一目标测试了ELECTRA的收益在多大程度上来自于预训练期间解决的将模型暴露于[MASK] tokens的不一致问题,而不是微调阶段。
All-TokensMLM:与Replace MLM类似,该目标的替换token也是来自生成器。
所不同的是,模型预测输入中的所有tokens,而不再仅仅是被遮蔽掉的那些。实验发现当模型有显性的拷贝机制时,效果更好。这种拷贝机制采用一个sigmoid层控制每个输入token的输出拷贝概率为D。模型的输出分布由输入token的D权值,再加上1-D乘以MLM softmax的输出。该模型本质上是BERT和ELECTRA的结合体。
需要注意的是,如果没有生成器生成替换字符,模型将很容易地学会从[MASK]标记的词汇表中进行预测,并从输入中直接拷贝其他token。
上述的对比结果如Table 5所示:从实验结果可以看出:第一,ELECTRA很大程度上受益于基于所有输入token定义的损失。采用15%子集定义的损失逊色于ELECTRA。
第二,从Replace MLM略胜BERT,可以看出BERT预训练时引入的[MASK] token导致微调的不一致,确实略微有损性能指标。第三,使用All-Tokens MLM方案的模型可以很大程度上缩减BERT和ELECTRA之间的差距。总之:ELECTRA模型的提升大部分可以归因于从所有tokens中学习,而小部分源于提升得益于减轻预训练-微调的不一致。
另外还对比了BERT和ELECTRA在不同模型尺寸下的表现,如Figure 4所示:可以看出,小模型中ELECTRA的收益更大,对于完全训练的小模型,ELECTRA也是显著优于BERT。这可以看出ELECTRA训练参数更为高效,估计是由于ELECTRA不必对每个位置可能的token的完整分布进行建模。但是仍然需要进行更多的分析来完全解释ELECTRA的参数效率问题。
总结
本文提出一个新的自监督任务:替换Token检测(RTD)。其核心是训练一个文本encoder以在输入tokens中鉴别出来自于小型生成器所生成的高质量负样本。对比于MLM,本文的预训练目标计算更高效,在下游任务中能取得更优的结果。本文的方法即使在算力相对较少的情况下仍然有效,这对于算力贫困人士来说真是大大的福音。