超越 BERT 模型的 ELECTRA 代码解读

转载来源：https://zhuanlan.zhihu.com/p/139898040
作者：feelingzhou，腾讯 WXG 应用研究员。

1、概述

在 2019 年 11 月份，NLP 大神 Manning 联合谷歌做的 ELECTRA 一经发布，迅速火爆整个 NLP 圈，其中 ELECTRA-small 模型参数量仅为 BERT-base 模型的 1/10，性能却依然能与 BERT、RoBERTa 等模型相媲美，得益于 ELECTRA 模型的巧妙构思 LOSS, 在 2020 年 3 月份 Google 对代码做了开源，下面针对 Google 放出的 ELECTRA 做代码做解读，希望通过此文章大家能在自己文本数据、行为序列数据训练一个较好的预训练模型，在业务上提升价值。

2、ELECTRA 模型

2.1 总体框架

ELECTRA 模型 (BASE 版本) 本质是换一种方法来训练 BERT 模型的参数；BERT 模型主要是利用 MLM 的思想来训练参数，直接把需要预测的词给挖掉了，挖了 15% 的比例。由于每次训练是一段话中 15% 的 token，导致模型收敛更新较慢，需要的语料也比较庞大。同时为了兼顾处理阅读理解这样的任务，模型加入了 NSP，是个二分类任务，判断上下两句是不是互为上下句；而 ELECTRA 模型主要借助于图像领域 gan 的思想，利用生成器和判别器思想，如下图所；ELECTRA 的预训练可以分为两部分，生成器部分仍然是 MLM，结构与 BERT 类似，利用这个模型对挖掉的 15% 的词进行预测，并将其进行替换，若替换的词不是原词，则打上被替换的标签，语句的其他词则打上没有替换的标签，判别器部分是训练一个判别模型对所有位置的词进行替换识别，此时预测模型转换成了一个二分类模型。这个转换可以带来效率的提升，对所有位置的词进行预测，收敛速度会快的多，损失函数是利用生成器部分的损失和判别器的损失函数以一个比例数 (官方代码是 50) 相加。

2.2 代码框架

2020 年 3 月份 Google 开源了 ELECTRA 模型代码，见代码连接，其主要代码框架如下：

下面对该代码框架做一一说明：

finetune： 该文件夹下面的代码主要是对已经训练好的 ELECTRA 模型做微调的代码例子，如文本分类、NER 识别、阅读理解等任务，这个和 BERT 的任务一致，在这里不做过多累赘。

model/modeling.py： 该文件主要是 bert 模型的实现逻辑以及 bert 模型的配置读取代码，在 ELECTRA 模型的预训练阶段生成和判别阶段做调用，另外在做各种 finetuneing 任务会调用。

model/optimization.py： 该文件主要是对优化器的实现，主要是对 AdamWeightDecay 的实现，可以自己加 lamb 等优化方法的实现。

model/tokenization.py： 该文件主要是 WordPiece 分词器的实现，可以对英文、中文分词，在将文本转化为 tfrecord 的时候会用的到。

pretrain/pretrain_data.py: 该文件是主要作用是对 ELECTRA 模型在 pretraining 对 tfrecords 文件读取、collections.namedtuple 更新上的一些逻辑实现。

pretrain_helpers.py： 该文件是 pretraining 阶段核心功能实现，主要实现逻辑对序列做动态随机 mask 以及对已经 mask 的序列做 unmask。

util/training_utils.py： 该文件主要在训练阶段实现了一个 Hook，在训练阶段为了打印更多日志信息。

util/utils.py： 该文件主要是一些基础方法，如增删文件夹、序列化反序列化、读取配置文件等。

build_openwebtext_pretraining_dataset.py、build_pretraining_dataset.py： 这两个文件功能类似，但是用到的数据源不一样，主要是把文本文件转化为 tfrecord 文件，tfrecord 文件的 key 包括 input_ids、input_mask、segment_ids，生成的 tfrecord 文件不像 bert 预训练需要的文件那样，不需要再生成 masked_lm_positions,masked_lm_ids,masked_lm_weights ，这几个 key 会在模型 pretraining 阶段自动生成，与此同时 mask 也是随机动态的，类似于 RoBerta，不像 BERT 那样固定。里面写了个多线程加速；对于大型的文件，还是使用 spark 将文本转化较适宜。

configure_finetuning.py： finetuneing 阶段的一些超参数配置，google 这次放出的代码参数并没有使用 tf.flags。

configure_pretraining.py： pretraining 阶段的一些超参数配置，google 这次放出的代码参数并没有使用 tf.flags。

run_finetuning.py： 模型做 finetuning 逻辑，加载已经训练好的 ELECTRA 模型做微调。

run_pretraining.py： 模型 pretraining 逻辑，也是 ELECTRA 模型最核心的逻辑，下面会加以详细说明。

2.3 pretraining 阶段

ELECTRA 模型 pretraining 阶段是最核心的逻辑，代码是在 run_pretraining.py 里面，下面会加以详细说明，整体阶段理解绘制了一张图，逻辑见如下图：

2.3.1 主方法入口

在 run_pretraining.py 文件中主方法有三个必须的参数见：

def main():parser = argparse.ArgumentParser(description=__doc__)parser.add_argument("--data-dir", required=True,help="Location of data files (model weights, etc).")parser.add_argument("--model-name", required=True,help="The name of the model being fine-tuned.")parser.add_argument("--hparams", default="{}",help="JSON dict of model hyperparameters.")args = parser.parse_args()if args.hparams.endswith(".json"):hparams = utils.load_json(args.hparams)else:hparams = json.loads(args.hparams)tf.logging.set_verbosity(tf.logging.ERROR)train_or_eval(configure_pretraining.PretrainingConfig(args.model_name, args.data_dir, **hparams))

三个必须参数见如下： –data-dir： 表示 tfrecord 文件的地址，一般是以 pretrain_data.tfrecord-0-of * 这种格式，调用 build_pretraining_dataset.py 文件生成，默认生成 1000 个 tfrecord 文件，数目可以自己改，此外需要注意的是切词需要制定 vocab.txt，训练中文的模型词典指定 BERT 模型那个 vocab.txt 词典即可，同理用于英文的模型训练。

–model-name： 表示预训练模型的名字一般是 electar，可以自己设定。

–hparams：，一般是个 json 文件，可以传递自己的参数进去，比如你要训练的模型是 small、base、big 等模型，还有 vocab_size，一般中文是 21128，英文是 30522。还有模型训练是否是测试状态等参数，一般我训练中文模型 hparams 参数是的 config.json 是:

{"model_size": "base","vocab_size": 21128
}

详细的参数可以去看 configure_pretraining.py，一般你传进去的参数进去会更新里面的超参数。

程序入口训练模型：

train_or_eval(configure_pretraining.PretrainingConfig(args.model_name, args.data_dir, **hparams))

还有一个入口，只 see.run() 一次，用于测试，见如下：

train_one_step(configure_pretraining.PretrainingConfig(args.model_name, args.data_dir, **hparams))

2.3.2 数据 mask

训练模型主要是代码是 PretrainingModel 类的定义，在 PretrainingModel 里面 程序首先对输入的 tfrecord 文件做随机 mask，

# Mask the inputmasked_inputs = pretrain_helpers.mask(config, pretrain_data.features_to_inputs(features), config.mask_prob)

用于生成含有 masked_lm_positions,masked_lm_ids,masked_lm_weights 等 key 的 tfrecord 文件，随机 MASK 实现的主要逻辑是调用 pretrain_helpers.mask() 来实现，其中用到了随机生成多项分布的函数 tf.random.categorical，这个函数目的是随机获取 masked_lm_positions、masked_lm_weights，再根据 masked_lm_positions 调用 tf.gather_nd 做索引截取来获取 masked_lm_ids。

2.3.3 Generator BERT

数据获取之后往下一步走就是生成 Generator BERT 阶段的模型，调用方法见如下：

generator = self._build_transformer(masked_inputs, is_training,bert_config=get_generator_config(config, self._bert_config),embedding_size=(None if config.untied_generator_embeddingselse embedding_size),untied_embeddings=config.untied_generator_embeddings,)

这里主要用于 Generator 阶段 BERT 模型生成，同时生成 MLM loss 和 Fake data，其中 Fake data 非常核心。 MLM loss 生成见代码，和 BERT 的逻辑几乎一样 ：

def _get_masked_lm_output(self, inputs: pretrain_data.Inputs, model):"""Masked language modeling softmax layer."""masked_lm_weights = inputs.masked_lm_weightswith 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_tiledelse: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-6loss = numerator / denominatorpreds = 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)

Fake data 数据生成逻辑见下面代码，这里调用了 unmask 函数和上面提到的 mask 函数作用相反，把原来 input_ids 随机 mask 的函数还原回去生成一个 input_ids_new，再利用谷生成模型生成的 logit 取最大索引去还原原来被 mask 调的 input_ids, 生成一个 updated_input_ids，判断 input_ids_new 和 updated_input_ids 是否相等，生成 true label

def _get_fake_data(self, inputs, mlm_logits):"""Sample from the generator to create corrupted input."""inputs = pretrain_helpers.unmask(inputs)disallow = tf.one_hot(inputs.masked_lm_ids, depth=self._bert_config.vocab_size,dtype=tf.float32) if self._config.disallow_correct else Nonesampled_tokens = tf.stop_gradient(pretrain_helpers.sample_from_softmax(mlm_logits / self._config.temperature, disallow=disallow))sampled_tokids = tf.argmax(sampled_tokens, -1, output_type=tf.int32)updated_input_ids, masked = pretrain_helpers.scatter_update(inputs.input_ids, sampled_tokids, inputs.masked_lm_positions)labels = masked * (1 - tf.cast(tf.equal(updated_input_ids, inputs.input_ids), tf.int32))updated_inputs = pretrain_data.get_updated_inputs(inputs, input_ids=updated_input_ids)FakedData = collections.namedtuple("FakedData", ["inputs", "is_fake_tokens", "sampled_tokens"])return FakedData(inputs=updated_inputs, is_fake_tokens=labels,sampled_tokens=sampled_tokens)

2.3.4 Discrimina BERT

利用上一步生成的 Fake data，作为 Discrimina BERT 的输入，见代码：

if config.electra_objective:discriminator = self._build_transformer(fake_data.inputs, is_training, reuse=not config.untied_generator,embedding_size=embedding_size)disc_output = self._get_discriminator_output(fake_data.inputs, discriminator, fake_data.is_fake_tokens)

获取二分类的损失函数，代码见：

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) * weightsper_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"])

2.3.5 总的损失函数

上面一步骤求出了 disc_output.loss 也就是 sigmod 的 loss，代码见：

self.total_loss = config.gen_weight * mlm_output.lossself.total_loss += config.disc_weight * disc_output.loss

这里 config.gen_weight=1 以及 config.disc_weight=50，这里 sigmod 的损失函数设置为 50，作者也没给明确的答复。

2.3.6 模型优化以及 checkpoint

上面一步已经求出了总的损失函数，下一步则是做模型优化训练以及做 checkpoint，程序入口在 train_or_eval() 里面，代码见:

is_per_host = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V2tpu_cluster_resolver = Noneif config.use_tpu and config.tpu_name:tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(config.tpu_name, zone=config.tpu_zone, project=config.gcp_project)tpu_config = tf.estimator.tpu.TPUConfig(iterations_per_loop=config.iterations_per_loop,num_shards=(config.num_tpu_cores if config.do_train elseconfig.num_tpu_cores),tpu_job_name=config.tpu_job_name,per_host_input_for_training=is_per_host)run_config = tf.estimator.tpu.RunConfig(cluster=tpu_cluster_resolver,model_dir=config.model_dir,save_checkpoints_steps=config.save_checkpoints_steps,tpu_config=tpu_config)model_fn = model_fn_builder(config=config)estimator = tf.estimator.tpu.TPUEstimator(use_tpu=config.use_tpu,model_fn=model_fn,config=run_config,train_batch_size=config.train_batch_size,eval_batch_size=config.eval_batch_size)

可以看到 google 开源的代码主要是 TPU 的一些钩子，改成 GPU 也比较简单，在 BERT 里面就有 GPU 相关的钩子，下面就会讲到。

2.4 finetuning 阶段

ELECTRA finetuning 阶段给出了不少例子，也比较简单，在 finetuneing 文件下面，这里不做过多的说明，和 bert 类似，唯一要改的就是把 TPU 相关的设置改为 GPU、CPU 即可。

2.5 序列训练改进

上面代码主要存在两个问题，第一个是 TPU 设置的问题，并不是人人都是土豪，还是要适配 GPU 的训练，第二个就是假如我想训练一个 vocab 比较大的序列模型，上面模型是训练不动的，loss 方面改为负采样的形式。

2.5.1 TPU 改 GPU 训练

global_step = tf.train.get_or_create_global_step()optimizer = optimization.AdamWeightDecayOptimizer(learning_rate=learning_rate)train_op = optimizer.apply_gradients(zip(grads, tvars), global_step)update_global_step = tf.assign(global_step, global_step + 1, name='update_global_step')output_spec=tf.estimator.EstimatorSpec(mode=mode,predictions=probabilities,loss=total_loss,train_op=tf.group(train_op, update_global_step))run_config = tf.estimator.RunConfig(model_dir=FLAGS.modelpath,save_checkpoints_steps=500)bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)model_fn = model_fn_builder(bert_config=bert_config,num_labels=FLAGS.n_class,is_training=True,init_checkpoint=FLAGS.init_checkpoint,learning_rate=FLAGS.learning_rate,use_one_hot_embeddings=False)estimator = tf.estimator.Estimator(model_fn=model_fn,                                    config=run_config)total_files=glob.glob("/tf*")random.shuffle(total_files)eval_files=total_files.pop()input_fn_train=lambda:input_fn(total_files,FLAGS.batch_size,num_epochs=N)input_fn_eval=lambda:input_fn(eval_files,FLAGS.batch_size,is_training=False)train_spec = tf.estimator.TrainSpec(input_fn=input_fn_train, max_steps=10000)eval_spec = tf.estimator.EvalSpec(input_fn=input_fn_eval,steps=None, start_delay_secs=30, throttle_secs=30)tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

2.5.2 负采样改造

主要是对 mlm loss 做改造：

def get_masked_lm_output(bert_config, input_tensor, output_weights, positions,label_ids, label_weights):with tf.variable_scope("cls/predictions"):with tf.variable_scope("transform"):input_tensor = tf.layers.dense(input_tensor,units=bert_config.hidden_size,activation=modeling.get_activation(bert_config.hidden_act),kernel_initializer=modeling.create_initializer(bert_config.initializer_range))input_tensor = modeling.layer_norm(input_tensor) #batch*10*embeding_size# The output weights are the same as the input embeddings, but there is# an output-only bias for each token.output_bias = tf.get_variable("output_bias",shape=[bert_config.vocab_size],initializer=tf.zeros_initializer())label_ids = tf.reshape(label_ids, [-1,1])label_weights = tf.reshape(label_weights, [-1])per_example_loss = tf.nn.sampled_softmax_loss(weights=output_weights,biases=output_bias,labels=label_ids,inputs=input_tensor,num_sampled=N,num_classes=bert_config.vocab_size)numerator = tf.reduce_sum(label_weights*per_example_loss )denominator = tf.reduce_sum(label_weights) + 1e-5loss = numerator / denominatorreturn (loss, per_example_loss)

3、总结

前段时间挺忙，有比较多的新 idea 出来没来的及看，上周末花了一天时间看了下 electra 源码，并做记录，也看到不少团队做了一些中文 electra 预训练模型，虽然 electra 没有达到 state of the art，和 roberta 差距可以忽视，但是这种训练方式这是一个很棒的 idea，其收敛速度是其他以 bert 为基础为改造的模型不能比的，在序列建模就有非常重大的研究意义，欢迎一起交流。

4、参考文献

[1]https://openreview.net/forum?id=r1xMH1BtvB
[2]https://zhuanlan.zhihu.com/p/104956125