BERT 推出这一年来,除了 XLNet,其他的改进都没带来太多惊喜,无非是越堆越大的模型和数据,以及动辄 1024 块 TPU,让工程师们不知道如何落地。
ELECTRA 通过类似 GAN 的结构和新的预训练任务,在更少的参数量和数据下,不仅吊打 BERT,而且仅用 1/4 的算力就达到了当时 SOTA 模型 RoBERTa 的效果。
简介
ELECTRA 的全称是 Efficiently Learning an Encoder that Classifies Token Replacements Accurately,先来直观感受一下 ELECTRA 的效果:
右边的图是左边的放大版,纵轴是 GLUE 分数,横轴是 FLOPs (floating point operations),Tensorflow 中提供的浮点数计算量统计。从上图可以看到,同等量级的 ELECTRA 是一直碾压 BERT 的,而且在训练更长的步数之后,达到了当时的 SOTA 模型——RoBERTa 的效果。从左图曲线上也可以看到,ELECTRA 效果还有继续上升的空间。
模型结构
NLP 式的 Generator-Discriminator
ELECTRA 最主要的贡献是提出了新的预训练任务和框架,把生成式的 Masked language model(MLM) 预训练任务改成了判别式的 Replaced token detection(RTD) 任务,判断当前 token 是否被语言模型替换过。那么问题来了,我随机替换一些输入中的字词,再让 BERT 去预测是否替换过可以吗?可以的,因为我就这么做过,但效果并不好,因为随机替换太简单了。
那怎样使任务复杂化呢?咦,咱们不是有预训练一个 MLM 模型吗?
于是作者就干脆使用一个 MLM 的 G-BERT 来对输入句子进行更改,然后丢给 D-BERT 去判断哪个字被改过,如下:
于是,我们 NLPer 终于成功地把 CV 的 GAN 拿过来了!
Replaced Token Detection
但上述结构有个问题,输入句子经过生成器,输出改写过的句子,因为句子的字词是离散的,所以梯度在这里就断了,判别器的梯度无法传给生成器,于是生成器的训练目标还是 MLM(作者在后文也验证了这种方法更好),判别器的目标是序列标注(判断每个 token 是真是假),两者同时训练,但判别器的梯度不会传给生成器,目标函数如下:
因为判别器的任务相对来说容易些,RTD loss 相对 MLM loss 会很小,因此加上一个系数,作者训练时使用了 50。
另外要注意的一点是,在优化判别器时计算了所有 token 上的 loss,而以往计算 BERT 的 MLM loss 时会忽略没被 mask 的 token。作者在后来的实验中也验证了在所有 token 上进行 loss 计算会提升效率和效果。
事实上,ELECTRA 使用的 Generator-Discriminator 架构与 GAN 还是有不少差别,作者列出了如下几点:
实验及结论
创新总是不易的,有了上述思想之后,可以看到作者进行了大量的实验,来验证模型结构、参数、训练方式的效果。
Weight Sharing
生成器和判别器的权重共享是否可以提升效果呢?作者设置了相同大小的生成器和判别器,在不共享权重下的效果是 83.6,只共享 token embedding 层的效果是 84.3,共享所有权重的效果是 84.4。作者认为生成器对 embedding 有更好的学习能力,因为在计算 MLM 时,softmax 是建立在所有 vocab 上的,之后反向传播时会更新所有 embedding,而判别器只会更新输入的 token embedding。最后作者只使用了 embedding sharing。
Smaller Generators
从权重共享的实验中看到,生成器和判别器只需要共享 embedding 的权重就足矣了,那这样的话是否可以缩小生成器的尺寸进行训练效率提升呢?作者在保持原有 hidden size 的设置下减少了层数,得到了下图所示的关系图:
可以看到,生成器的大小在判别器的 1/4 到 1/2 之间效果是最好的。作者认为原因是过强的生成器会增大判别器的难度(判别器:小一点吧,我太难了)。
Training Algorithms
实际上除了 MLM loss,作者也尝试了另外两种训练策略:
Adversarial Contrastive Estimation:ELECTRA 因为上述一些问题无法使用 GAN,但也可以以一种对抗学习的思想来训练。作者将生成器的目标函数由最小化 MLM loss 换成了最大化判别器在被替换 token 上的 RTD loss。但还有一个问题,就是新的生成器 loss 无法用梯度下降更新生成器,于是作者用强化学习 Policy Gradient 的思想,将被替换 token 的交叉熵作为生成器的 reward,然后进行梯度下降。强化方法优化下来生成器在 MLM 任务上可以达到 54% 的准确率,而之前 MLE 优化下可以达到 65%。
Two-stage training:即先训练生成器,然后 freeze 掉,用生成器的权重初始化判别器,再接着训练相同步数的判别器。
对比三种训练策略,得到下图:
可见「隔离式」的训练策略效果还是最好的,而两段式的训练虽然弱一些,作者猜测是生成器太强了导致判别任务难度增大,但最终效果也比 BERT 本身要强,进一步证明了判别式预训练的效果。
Small model? Big model?
这两节真是吊打之前的模型,作者重申了他的主要目的是提升预训练效率,于是做了 GPU 单卡就可以愉快训练的 ELECTRA-Small 和 BERT-Small,接着和尺寸不变的 ELMo、GPT 等进行对比,结果如下:
数据简直优秀,仅用 14M 参数量,以前 13% 的体积,在提升了训练速度的同时还提升了效果,这里我疯狂点赞。
小 ELECTRA 的本事我们见过了,那大 ELECTRA 行吗?直接上图:
上面是各个模型在 GLUE dev/text 上的表现,可以看到 ELECTRA 仅用了 1/4 的计算量就达到了 RoBERTa 的效果。而且作者使用的是 XLNet 的语料,大约是 126G,但 RoBERTa 用了 160G。由于时间和精力问题,作者们没有把 ELECTRA 训练更久(应该会有提升),也没有使用各种榜单 Trick,所以真正的 GLUE test 上表现一般(现在的 T5 是 89.7,RoBERTa 是 88.5,没看到 ELECTRA)。
Efficiency Analysis
前文中提到了,BERT 的 loss 只计算被替换的 15% 个 token,而 ELECTRA 是全部都计算的,所以作者又做了几个实验,探究哪种方式更好一些:
1. ELECTRA 15%:让判别器只计算 15% token 上的损失
2. Replace MLM:训练 BERT MLM,输入不用 [MASK] 进行替换,而是其他生成器。这样可以消除这种 pretrain-finetune 直接的 diff。
3. All-Tokens MLM:接着用 Replace MLM,只不过 BERT 的目标函数变为预测所有的 token,比较接近 ELECTRA。
三种实验结果如下:
可以看到:
对比 ELECTRA 和 ELECTRA 15%:在所有 token 上计算 loss 确实能提升效果
对比 Replace MLM 和 BERT:[MASK] 标志确实会对 BERT 产生影响,而且 BERT 目前还有一个 trick,就是被替换的 10% 情况下使用原 token 或其他 token,如果没有这个 trick 估计效果会差一些。
对比 All-Tokens MLM 和 BERT:如果 BERT 预测所有 token 的话,效果会接近 ELECTRA
另外,作者还发现,ELECTRA 体积越小,相比于 BERT 就提升的越明显,说明 fully trained 的 ELECTRA 效果会更好。另外作者推断,由于 ELECTRA 是判别式任务,不用对整个数据分布建模,所以更 parameter-efficient。
总结
无意中发现了这篇还在 ICLR 盲审的 ELECTRA,读完摘要就觉得发现了新大陆,主要是自己也试过 Replaced Token Detection 这个任务,因为平时任务效果的分析和不久前看的一篇文章,让我深刻感受到了 BERT 虽然对上下文有很强的编码能力,却缺乏细粒度语义的表示,我用一张图表示大家就明白了:
这是把 token 编码降维后的效果,可以看到 sky 和 sea 明明是天与海的区别,却因为上下文一样而得到了极为相似的编码。细粒度表示能力的缺失会对真实任务造成很大影响,如果被针对性攻击的话更是无力,所以当时就想办法加上更细粒度的任务让 BERT 去区分每个 token,不过同句内随机替换的效果并不好,弱鸡的我也没有再往前想一步,不然就也 ICLR 了。相信这个任务很多人都想到过,不过都没有探索这么深入,这也告诫我们,idea 遍地都是,往下挖才能有 SOTA。
ELECTRA 是 BERT 推出这一年来我见过最赞的 idea,它不仅提出了能打败 MLM 的预训练任务,更推出了一种十分适用于 NLP 的类 GAN 框架。毕竟 GAN 太牛逼了,看到 deepfake 的时候我就想,什么时候我们也能 deepcheat,但听说 GAN 在 NLP 上的效果一直不太好(只懂皮毛,要学起来了,轻拍),这次 ELECTRA 虽然只用了判别器,但个人认为也在一定程度上打开了潘多拉魔盒。
另外,整篇文章都干货满满,不再像之前的 BERT+模型一样可以用「more data+params+steps+GPU+MONEY」简单概括。推荐大家去通读正文+附录,里面还有一些失败尝试我没有讲。
如果 ELECTRA 去直播,我一定给它刷一辆游艇。
论文:《ELECTRA: PRE-TRAINING TEXT ENCODERS AS DISCRIMINATORS RATHER THAN GENERATORS》
论文链接:https://openreview.net/pdf?id=r1xMH1BtvB