ELECTRA:用Bert欺骗Bert

18年Bert横空出世,刷新了各大榜单之后,对齐的改进层出不穷,Ernie, Albert大多数改进都没有新意,无非就是大力出奇迹的堆叠参数。ICLR2020 斯坦福和Google为我们提供了一种新思路,用Bert来"欺骗"Bert。今天我们来为大家介绍ELECTRA的思路。

背景

目前以Bert为代表的state of art的预训练模型都是基于MLM(Masked Language modeling)来进行预训练的,这些模型将输入的句子中15%的Mask掉,然后使用模型去预测被mask掉的原始内容。当然这些模型也面临一个问题就是模型因为参数过多,每次训练只能学习到训练数据的15%的内容,从而导致计算量过大的问题。

因此文章中提出了一种新的训练方法:随机替换句子中的token使用模型去判断这个token是否被替换过。ELECTRA的效果有多凶残呢,我们看下图,左图是右图的方法版,横轴是预训练的FLOPs(floating point operations),TF中的浮点数计算统计量,纵轴是GLUE的分数,在同等计算量的情况下,ELECTRA一直碾压Bert,在训练到一定程度之后可以到达RoBERTa的效果。

ELECTRA:用Bert欺骗Bert

训练方法

ELECTRA的主要贡献是在预训练中将MLM(Masked Language Model)替换为RTD任务(Replace Token Detection)预测Token是否被替换。这个任务由两个模型来实现,Generator和Discriminator。Generator负责生成被替换的Token,用Discriminator去判断每个Token是否是被替换过的。

ELECTRA:用Bert欺骗Bert

那么如何替换Token以及判断呢?下面我们来分别讲一下Generator和Discriminator.

1. Generator

有博主尝试过随机替换Token的方法,但是效果并不好,因为随机替换太简单了。那么文中是怎么做的呢?在MLM任务中我们会随机Mask掉一部分位置的Token并训练模型去预测这一部分,文中也借鉴了这个思想,使用Generator训练了一小的MLM任务,然后Discriminator去判断这些Token是否被Generator替换过。

对于特定位置t,我们假设该位置被mask掉了,那么该位置被预测为x_{t}的概率为:

ELECTRA:用Bert欺骗Bert

其中x=[x_{1},x_{2},...,x_{n}] 是输入的Token序列,h(x)=[h_{1},h_{2},...,h_{n}] 是经过MLM之后的输出,其中e(x_{t})^{T} 是token的embedding。

对于Generator的loss如何计算呢?可以看到文中是这么计算的,对我们mask的Token的概率进行log然后求期望。

ELECTRA:用Bert欺骗Bert

现在这么说或许还是有点模糊,具体来看一下ELECTRA是如何实现的吧。首先对输入的sequence按照一定概率进行mask,输入模型的config,预训练的input,不能被mask的位置和已经被mask的位置,然后我们对预训练的输入按照一定的概率产生返回结果我们对输入的数据进行mask,最后对input进行处理转换成一个dict,里面存储了input_id, masked_lmposition等内容。

def mask(config: configure_pretraining.PretrainingConfig,
inputs: pretrain_data.Inputs, mask_prob, proposal_distribution=1.0,
disallow_from_mask=None, already_masked=None):
# Get the batch size, sequence length, and max masked-out tokens
N = config.max_predictions_per_seq
B, L = modeling.get_shape_list(inputs.input_ids)
# Find indices where masking out a token is allowed
vocab = tokenization.FullTokenizer(
config.vocab_file, do_lower_case=config.do_lower_case).vocab
candidates_mask = _get_candidates_mask(inputs, vocab, disallow_from_mask)
# Set the number of tokens to mask out per example
num_tokens = tf.cast(tf.reduce_sum(inputs.input_mask, -1), tf.float32)
num_to_predict = tf.maximum(1, tf.minimum(
N, tf.cast(tf.round(num_tokens * mask_prob), tf.int32)))
masked_lm_weights = tf.cast(tf.sequence_mask(num_to_predict, N), tf.float32)
if already_masked is not None:
masked_lm_weights *= (1 - already_masked)
# Get a probability of masking each position in the sequence
candidate_mask_float = tf.cast(candidates_mask, tf.float32)
sample_prob = (proposal_distribution * candidate_mask_float)
sample_prob /= tf.reduce_sum(sample_prob, axis=-1, keepdims=True)
# Sample the positions to mask out
sample_prob = tf.stop_gradient(sample_prob)
sample_logits = tf.log(sample_prob)
masked_lm_positions = tf.random.categorical(
sample_logits, N, dtype=tf.int32)
masked_lm_positions *= tf.cast(masked_lm_weights, tf.int32)
# Get the ids of the masked-out tokens
shift = tf.expand_dims(L * tf.range(B), -1)
flat_positions = tf.reshape(masked_lm_positions + shift, [-1, 1])
masked_lm_ids = tf.gather_nd(tf.reshape(inputs.input_ids, [-1]),
flat_positions)
masked_lm_ids = tf.reshape(masked_lm_ids, [B, -1])
masked_lm_ids *= tf.cast(masked_lm_weights, tf.int32)
# Update the input ids
replace_with_mask_positions = masked_lm_positions * tf.cast(
tf.less(tf.random.uniform([B, N]), 0.85), tf.int32)
inputs_ids, _ = scatter_update(
inputs.input_ids, tf.fill([B, N], vocab["[MASK]"]),
replace_with_mask_positions)
return pretrain_data.get_updated_inputs(
inputs,
input_ids=tf.stop_gradient(inputs_ids),
masked_lm_positions=masked_lm_positions,
masked_lm_ids=masked_lm_ids,
masked_lm_weights=masked_lm_weights
)

文中的Generator采用了Bert模型,也正如文中所说,文中使用Bert对Mask的Token进行预测,对每一个位置的Mask计算loss最后求和,loss的计算过程如下。输入是我们刚才处理过的maskedinputs和Generator。为了将Bert计算的logits转换为预测的Label,代码在Generator之后加了一层全连接层和sofmax,然后将预测的label转为one_hot编码,然后采用上述公式计算Mask部分的loss。

def _get_masked_lm_output(self, inputs: pretrain_data.Inputs, model):
"""Masked language modeling softmax layer."""
masked_lm_weights = inputs.masked_lm_weights
with tf.variable_scope("generator_predictions"):
if self._config.uniform_generator:
logits = tf.zeros(self._bert_config.vocab_size)
logits_tiled = tf.zeros(
modeling.get_shape_list(inputs.masked_lm_ids) +
[self._bert_config.vocab_size])
logits_tiled += tf.reshape(logits, [1, 1, self._bert_config.vocab_size])
logits = logits_tiled
else:
relevant_hidden = pretrain_helpers.gather_positions(
model.get_sequence_output(), inputs.masked_lm_positions)
hidden = tf.layers.dense(
relevant_hidden,
units=modeling.get_shape_list(model.get_embedding_table())[-1],
activation=modeling.get_activation(self._bert_config.hidden_act),
kernel_initializer=modeling.create_initializer(
self._bert_config.initializer_range))
hidden = modeling.layer_norm(hidden)
output_bias = tf.get_variable(
"output_bias",
shape=[self._bert_config.vocab_size],
initializer=tf.zeros_initializer())
logits = tf.matmul(hidden, model.get_embedding_table(),
transpose_b=True)
logits = tf.nn.bias_add(logits, output_bias)
oh_labels = tf.one_hot(
inputs.masked_lm_ids, depth=self._bert_config.vocab_size,
dtype=tf.float32)
probs = tf.nn.softmax(logits)
log_probs = tf.nn.log_softmax(logits)
label_log_probs = -tf.reduce_sum(log_probs * oh_labels, axis=-1)
numerator = tf.reduce_sum(inputs.masked_lm_weights * label_log_probs)
denominator = tf.reduce_sum(masked_lm_weights) + 1e-6
loss = numerator / denominator
preds = tf.argmax(log_probs, axis=-1, output_type=tf.int32)
MLMOutput = collections.namedtuple(
"MLMOutput", ["logits", "probs", "loss", "per_example_loss", "preds"])
return MLMOutput(
logits=logits, probs=probs, per_example_loss=label_log_probs,
loss=loss, preds=preds)

当然文中的思路是很清晰,但是我有一点疑惑,就是在预训练过程中的loss是Generator和Discriminator的loss求和,当然为了保证效果loss肯定是希望变小的。不过为了保证随机生成的效果,这里应该是预测的和原文本出入比较大比较好,那么如何去平衡loss变小的问题呢?或许可以借鉴NSP(Next Sentence Prediction)任务中队token进行跨领域的替换?或者在最终的loss中按照权重,增大Discriminator的权重?这里还是比较让人迷惑的,而且文中说的天花乱坠的思路看来只不过是Bert的二次利用,感觉有一丢丢受骗的感觉。

2. Discriminator

Disctiminator根据generator生成的输入去判断是否是生成的Token,这样就将问题转化为一个二分类问题,对每一个Token我们将其分成是生成Token或者不是,那么如此使用交叉熵来表示就是个非常好的选择了。事实上,论文中对Discriminator的loss采用的也是交叉熵。

ELECTRA:用Bert欺骗Bert

下面我们来看一下代码中如何实现Discriminator,如论文中所述,Discriminator和Generator都采用Bert,不同于Generator,Discriminator的输入是经过Generator生成之后的fake_input,label表示Token是否是fake。

def _get_discriminator_output(self, inputs, discriminator, labels):
"""Discriminator binary classifier."""
with tf.variable_scope("discriminator_predictions"):
hidden = tf.layers.dense(
discriminator.get_sequence_output(),
units=self._bert_config.hidden_size,
activation=modeling.get_activation(self._bert_config.hidden_act),
kernel_initializer=modeling.create_initializer(
self._bert_config.initializer_range))
logits = tf.squeeze(tf.layers.dense(hidden, units=1), -1)
weights = tf.cast(inputs.input_mask, tf.float32)
labelsf = tf.cast(labels, tf.float32)
losses = tf.nn.sigmoid_cross_entropy_with_logits(
logits=logits, labels=labelsf) * weights
per_example_loss = (tf.reduce_sum(losses, axis=-1) /
(1e-6 + tf.reduce_sum(weights, axis=-1)))
loss = tf.reduce_sum(losses) / (1e-6 + tf.reduce_sum(weights))
probs = tf.nn.sigmoid(logits)
preds = tf.cast(tf.round((tf.sign(logits) + 1) / 2), tf.int32)
DiscOutput = collections.namedtuple(
"DiscOutput", ["loss", "per_example_loss", "probs", "preds",
"labels"])
return DiscOutput(
loss=loss, per_example_loss=per_example_loss, probs=probs,
preds=preds, labels=labels,
)

在上一部分讨论Generator的时候,我谈到了loss的问题,果然在这一部分就被打脸,因为我们目标就是使整体的loss最小化,而最小的loss果然给DIscriminator赋予了权值。而实际上在实现过程中Generator也有相应的权值:self.total_loss 的计算首先是Generator的权重乘以loss然后加上Discriminator的权重乘以loss。如此严谨有理有据也让我表示自己没有想多哈哈,内心还是有点小窃喜呢,嘻嘻嘻。

ELECTRA:用Bert欺骗Bert
def __init__(self, config: configure_pretraining.PretrainingConfig,
features, is_training):
# Set up model config
self._config = config
self._bert_config = training_utils.get_bert_config(config)
if config.debug:
self._bert_config.num_hidden_layers = 3
self._bert_config.hidden_size = 144
self._bert_config.intermediate_size = 144 * 4
self._bert_config.num_attention_heads = 4
# Mask the input
masked_inputs = pretrain_helpers.mask(
config, pretrain_data.features_to_inputs(features), config.mask_prob)
# Generator
embedding_size = (
self._bert_config.hidden_size if config.embedding_size is None else
config.embedding_size)
if config.uniform_generator:
mlm_output = self._get_masked_lm_output(masked_inputs, None)
elif config.electra_objective and config.untied_generator:
generator = self._build_transformer(
masked_inputs, is_training,
bert_config=get_generator_config(config, self._bert_config),
embedding_size=(None if config.untied_generator_embeddings
else embedding_size),
untied_embeddings=config.untied_generator_embeddings,
name="generator")
mlm_output = self._get_masked_lm_output(masked_inputs, generator)
else:
generator = self._build_transformer(
masked_inputs, is_training, embedding_size=embedding_size)
mlm_output = self._get_masked_lm_output(masked_inputs, generator)
fake_data = self._get_fake_data(masked_inputs, mlm_output.logits)
self.mlm_output = mlm_output
self.total_loss = config.gen_weight * mlm_output.loss

3. GAN

相信看到这里的小伙伴们和我心理都有个疑问,这个和GAN的区别是什么呢?文中对此也做了解释:

ELECTRA:用Bert欺骗Bert

这里我们已经大概了解了ELECTRA的设计思想了,文章认为MLM任务比较简单,而且Mask的Token位置比较少,导致模型学习到的内容有限。而基于此观点,文章设计了比较精巧的生成器-判别器模式,而为了保证模型可以学习到语料的全部内容,生成器也避免简单地随机替换。这种类似于GAN又和GAN有所区别的思路,初看让人激动不已。然而看了源码,我的激动也逐渐理性,相比于文章的花哨,源码看上去更像是对Bert的一次组装,而生成器和判别器一起训练,Loss一起计算也让我有一丢疑惑。但是整体来说也是一次让人激动地尝试。那么模型的效果如何呢?我们将在下一篇文章中进行解释,请大家保持期待。

实验结果

1. 共享参数(Weight Sharing)

Generator和Discriminator应该共享参数吗?文中尝试了3中方法,不共享参数,共享embedding和共享所有参数,从效果来看共享所有参数的效果是最优的,但是这也意味着生成器和分辨器要一样大,这真的有必要吗?

生成器的工作是预测Mask掉的Token,至于对不对并不重要,而且从某种角度上来说,预测的越离谱,可能越适合分辨器学习。而Discriminator面对的是一整句话,要逐Token的判断该Token是原生的还是非原生的,他要学习的东西相比于Generator不仅庞大而且复杂,让Discriminator和Generator一样大对于任重而道远的Discriminator过于残忍,而如果让Generator和Discriminator一样大,又过于浪费。

begin{array}[b] {|c|} hline 实验方法 & GLUE score\ hline 不共享参数 & 83.6 \ hline 共享embedding & 84.3 \ hline 共享所有参数 & 84.4 \ hline end{array}\

因此文中所采用的是共享embedding的方法。

2. Small generators

作者在保持hidden size的情况下降低层数,从而降低Generator的大小,Discriminator尝试了hidden size为256,512和768,这里我们发现秉承着大力出奇迹的原则,同等情况下,Discriminator的hidden size为768的效果最好。与此同时我们还发现在Discriminator保持不变的情况下,Generator的大小并不是越大越好,当Generator的大小是Discriminator的1/4~1/2的时候实验效果最好。

ELECTRA:用Bert欺骗Bert

这是为什么呢?当Generator变得复杂,可能会有两种情况发生,Generator对Token的预测都非常有效,没有起到欺骗的作用,或者Generator过拟合导致任务对Discriminator过于复杂而降低了学习效率。(Discriminator:救救孩子吧)

3. Training Algorithms

  1. Two-stage训练:训练完Generator之后使用Generator的权重初始化Discriminator,然后训练Discriminator。
  2. Adversarial Contrastive Estimation: 上一篇文章中我们介绍了ELECTRA和GAN的区别,Discriminator的梯度无法传递到Generator,文中尝试了用强化学习的方法来训练模型。作者将Generator最小化MLM Loss替换为最大化Discriminator被替换Token的RTD LOSS。这时作者又面临了问题,新的loss无法使用梯度上升去最大化loss,于是作者采用了policy gradient reinforcement learning来寻找最优的分辨器。

对于Two-stage的方法中我们需要注意用Generator的参数初始化Discriminator,那么两个大小应该是一样。既然如此,那么应该可以预测这样做的效果并不会很好,为什么呀?因为上一部分我们不是提到了Generator是Discriminator的1/4~1/2时效果最好吗。除此以外,作者还尝试了另一种训练方法。

对于Adversarial Contrastive Estimation的看法大家可以在回复中留言讨论。

最后的实验结果也显示原始的训练方法效果最好。

ELECTRA:用Bert欺骗Bert

4. Small model? Big model?

文中训练了Small Electra和Big Electra,模型规模大家可以通过下标判断。相比于Bert Base, Small Electra的参数都进行了缩小,Big Electra和Bert large的超参数保持一致,同时训练的时间要更长一点。

begin{array}[b] {|c|} hline & Small Electra & Big Electra & Bert Base\ hline Sequence Length & 128 & 512 & 512\ hline Batch size & 128 & 2048 & 256\ hline Hidden size & 256 & 1024 & 768\ hline Embedding size & 128 & 1024 & 768\ hline end{array}\

那么Electra的效果如何呢?我们先看Small Electra,效果可以说是一骑绝尘。同等规模的情况下效果远超其他模型,Small Electra的Glue score要比Bert small高4.8个点,但是当规模变大,Electra Base的效果比Bert Base高2.9个点。没有Electra small那么亮眼,但也足够激动人心。毕竟Bert缩小版Albert也只是在xxLarge的情况下打败了Bert,但是Electra却赢得毫不费力。

ELECTRA:用Bert欺骗Bert

下面我们来看Electra Large的效果,是否如Small一样激动人心呢?结果依然不负众望,在SST, MRPC任务中Electra Large略逊于RoBERTa,其他任务都取得了很不错的成绩。需要注意的是在Electra之用了1/4的计算量就打败了RoBERTa。

ELECTRA:用Bert欺骗Bert

4. Efficiency Analysis

前文中提到Bert只计算了被替换的Token的loss,Electra使用了全部的Token,作者进一步做了一些实验来探讨那种方法更好。作者进行了一下3个实验

  • ELECTRA 15%:使用Electra计算15%的loss
  • Replace MLM: 使用Bert训练在预训练的时候输入不用MASK而是用其他生成器的输出替换
  • All-TokensMLM:结合了Bert和Electra,Bert的预测变成了预测所有Token

实验结果如下:

ELECTRA:用Bert欺骗Bert

可以看到Electra 15%的效果和Bert相似,因此Bert之前只学习15%的Token的做法对于输入是有很大的信息损失的,而Electra的做法也弥补了这一损失。这也证明了之前作者的看法:Bert只学习15%的Token是不够的。Replace MLM的效果和Electra 15%的效果相差不大,这说明MASK的内容其实并不重要,重要的是要学习全部的输入序列。All-Tokens MLM的效果也解释了这一点。

通过这些实验结果,我想我们也可以理解为什么作者认为分辨器的任务比生成器复杂,分辨器的规模也要比生成器大。因为Input is all.(这句话是我加的)

一些想法

前文中提到我个人对Electra的loss不是很认同的看法,这部分我们看到作者对Loss也做了一些新的尝试,例如最大化分辨器被替换的Token的RTD Loss,虽然如此但我个人还是不是特别满意。因为我认为为了保证分辨器的效果,生成器要给他尽可能复杂的替换效果。但是这里的RTD Loss关注的是分辨器,生成器在整个过程中所受到的关注并不多。当然文中实验也提到过,生成器规模变大,分辨器的效果会相对变差。那么文中此时也没有提到生成器规模变大,是否会导致生成器的loss降低,预测效果提高?而如果证明这一点,那么我也可以确定我之前的猜测并没有错:分辨器还是需要足够的噪音以提高效

当然整体而言,这篇文章给了我们在Bert基础上进行尝试的工匠一个亮眼的提示,其他的论文大都想着对结构,超参等进行修改。而这篇论文回归数据本身,关注我们的输入,用Bert"欺骗'Bert,想想都让人兴奋呢。那么下一次让人惊喜的创意又会是什么呢?希望大家和我一起期待。