ELECTRA

简介

ELECTRA 是一种自监督的语言表示学习方法。它可以在相对较少的计算资源下对 Transformer 网络进行预训练。ELECTRA 模型通过区分“真实”输入标记与由另一个神经网络生成的“虚假”输入标记来进行训练，这类似于 GAN 中的判别器。在小规模情况下，ELECTRA 即使只使用单个 GPU 训练也能取得优异的效果。而在大规模情况下，ELECTRA 在 SQuAD 2.0 数据集上达到了最先进的水平。

有关详细说明和实验结果，请参阅我们在 ICLR 2020 上发表的论文 ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators。

本仓库包含用于预训练 ELECTRA 的代码，其中包括在单个 GPU 上运行的小型 ELECTRA 模型。它还支持在下游任务上对 ELECTRA 进行微调，这些任务包括分类任务（例如 GLUE）、问答任务（例如 SQuAD）以及序列标注任务（例如 文本分块）。

此外，本仓库还包含 Electric 的代码，它是受 基于能量的模型 启发的 ELECTRA 版本。Electric 将 ELECTRA 更加严谨地视为一种“负采样”的 完形填空模型。它还可以高效地生成文本的 伪似然分数，这些分数可用于对语音识别或机器翻译系统的输出进行重新排序。有关 Electric 的详细信息，请参阅我们在 EMNLP 2020 上发表的论文 Pre-Training Transformers as Energy-Based Cloze Models。

已发布的模型

我们最初发布了三个预训练模型：

模型	层数	隐藏层大小	参数量	GLUE 分数（测试集）	下载链接
ELECTRA-Small	12	256	14M	77.4	链接
ELECTRA-Base	12	768	110M	82.7	链接
ELECTRA-Large	24	1024	335M	85.2	链接

这些模型是在不区分大小写的英文文本上训练的。它们分别对应于我们论文中的 ELECTRA-Small++、ELECTRA-Base++ 和 ELECTRA-1.75M。我们希望在未来发布其他模型，例如多语言模型。

在 GLUE 基准上，ELECTRA-Large 的得分略高于 ALBERT/XLNET，ELECTRA-Base 的得分优于 BERT-Large，而 ELECTRA-Small 的得分则略低于 TinyBERT（但未使用蒸馏技术）。详细的性能指标请参见下方的预期结果部分。

要求

• Python 3
• TensorFlow 1.15（尽管我们希望在未来支持 TensorFlow 2.0）
• NumPy
• scikit-learn 和 SciPy（用于计算部分评估指标）。

预训练

使用 build_pretraining_dataset.py 可以从原始文本转储中创建预训练数据集。该脚本具有以下参数：

• --corpus-dir: 包含要转换为 ELECTRA 示例的原始文本文件的目录。一个文本文件可以包含多个文档，各文档之间用空行分隔。
• --vocab-file: 定义 WordPiece 词汇表的文件。
• --output-dir: 存储 ELECTRA 示例的输出目录。
• --max-seq-length: 每个示例的最大标记数（默认为 128）。
• --num-processes: 如果大于 1，则会并行化多个进程（默认为 1）。
• --blanks-separate-docs: 是否将空行视为文档边界（默认为 True）。
• --do-lower-case/--no-lower-case: 是否将输入文本转换为小写（默认为 True）。

使用 run_pretraining.py 可以对 ELECTRA 模型进行预训练。该脚本具有以下参数：

• --data-dir: 存储预训练数据、模型权重等的目录。默认情况下，训练会从 <data-dir>/pretrain_tfrecords 加载示例，并从 <data-dir>/vocab.txt 加载词汇表。
• --model-name: 正在训练的模型名称。模型权重默认会保存到 <data-dir>/models/<model-name>。
• --hparams（可选）：包含模型超参数、数据路径等的 JSON 字典或 JSON 文件路径。支持的超参数请参阅 configure_pretraining.py。

如果训练被中断，再次运行带有相同参数的 run_pretraining.py 将会从中断处继续训练。

您可以通过以下步骤从已发布的 ELECTRA 检查点继续预训练：

1. 将 model-name 设置为指向已下载的模型（例如，如果您将权重下载到 $DATA_DIR/electra_small，则设置为 --model-name electra_small）。
2. 设置 num_train_steps（例如，在 --hparams 中添加 "num_train_steps": 4010000）。这将使小型模型再训练 10000 步（它已经训练了 400 万步）。
3. 提高学习率以适应线性学习率衰减。例如，若要以 2e-4 的初始学习率开始，应将 learning_rate 超参数设置为 2e-4 * (4e6 + 10000) / 10000。
4. 对于 ELECTRA-Small，还需要在 hparams 中指定 "generator_hidden_size": 1.0，因为我们并未为此模型使用小型生成器。

快速入门：预训练一个小型 ELECTRA 模型。

这些说明将预训练一个小型 ELECTRA 模型（12 层，隐藏层大小为 256）。遗憾的是，我们在论文中使用的数据并未公开，因此我们将改用由 Aaron Gokaslan 和 Vanya Cohen 发布的 OpenWebTextCorpus 数据集。完全训练好的模型（在 v100 GPU 上大约需要 4 天）在 GLUE 任务上的表现大致介于 GPT 和 BERT-Base 之间。默认情况下，该模型是在长度为 128 的序列上进行训练的，因此不适合用于问答任务。有关模型性能的更多详细信息，请参阅下方的“预期结果”部分。

设置

1. 将词汇表文件放置在 $DATA_DIR/vocab.txt 中。我们的 ELECTRA 模型均使用与英语小写 BERT 完全相同的词汇表，您可以从这里下载。
2. 下载 OpenWebText 数据集（12GB），并解压它（即运行 tar xf openwebtext.tar.xz）。将其放置在 $DATA_DIR/openwebtext 目录下。
3. 运行 python3 build_openwebtext_pretraining_dataset.py --data-dir $DATA_DIR --num-processes 5。此脚本会预处理/分词数据，并将示例以 tfrecord 文件的形式输出到 $DATA_DIR/pretrain_tfrecords 目录下。这些 tfrecord 文件大约需要 30GB 的磁盘空间。

预训练模型。

运行 python3 run_pretraining.py --data-dir $DATA_DIR --model-name electra_small_owt，即可在该数据集上对小型 ELECTRA 模型进行 100 万步的预训练。这在 Tesla V100 GPU 上大约需要略多于 4 天的时间。不过，该模型在经过 20 万步（约 10 小时的 v100 GPU 训练）后便能达到不错的效果。

若要自定义训练过程，可在运行命令中添加 --hparams '{"hparam1": value1, "hparam2": value2, ...}'。--hparams 也可以是一个包含超参数的 .json 文件路径。一些特别有用的选项包括：

• "debug": true：训练一个微型 ELECTRA 模型，仅进行几步。
• "model_size": one of "small", "base", or "large"：决定模型的大小。
• "electra_objective": false：训练一个采用掩码语言建模而非替换标记检测的模型（本质上是带有动态掩码、无下一句预测的 BERT）。
• "num_train_steps": n：控制预训练的持续时间。
• "pretrain_tfrecords": <paths>：指定预训练数据的位置。请注意，您需要指定具体的文件名，而不仅仅是目录（例如 <data-dir>/pretrain_tf_records/pretrain_data.tfrecord*）。
• "vocab_file": <path> 和 "vocab_size": n 可用于设置自定义的 WordPiece 词汇表。
• "learning_rate": lr, "train_batch_size": n 等参数可用于调整训练超参数。
• "model_hparam_overrides": {"hidden_size": n, "num_hidden_layers": m} 等参数可用于修改底层 Transformer 的超参数（"model_size" 标志会设置默认值）。

完整的支持超参数列表请参阅 configure_pretraining.py。

评估预训练模型。

若要在下游任务上评估模型，请参阅下方的微调说明。若要在 OpenWebText 数据上评估生成器和判别器，则运行 python3 run_pretraining.py --data-dir $DATA_DIR --model-name electra_small_owt --hparams '{"do_train": false, "do_eval": true}'。此命令将打印出生成器和判别器的准确率等评估指标，并将这些指标保存到 data-dir/model-name/results 文件中。

微调

使用 run_finetuning.py 脚本可以在下游 NLP 任务上对 ELECTRA 模型进行微调和评估。该脚本需要三个参数：

• --data-dir：存储数据、模型权重等的目录。默认情况下，脚本会从 <data-dir>/finetuning_data/<task-name> 加载微调数据，并从 <data-dir>/vocab.txt 加载词汇表。
• --model-name：预训练模型的名称；预训练权重应位于 data-dir/models/model-name 目录下。
• --hparams：包含模型超参数、数据路径等信息的 JSON 字典（例如 --hparams '{"task_names": ["rte"], "model_size": "base", "learning_rate": 1e-4, ...}'）。支持的超参数请参阅 configure_pretraining.py。除了字典之外，此参数也可以是一个包含超参数的 .json 文件路径。您必须指定 "task_names" 和 "model_size"（见下方示例）。

评估指标默认会保存到 data-dir/model-name/results，模型权重则默认保存到 data-dir/model-name/finetuning_models。默认情况下，评估是在开发集上进行的。若要自定义训练过程，可在运行命令中添加 --hparams '{"hparam1": value1, "hparam2": value2, ...}'。一些特别有用的选项包括：

• "debug": true：对微型 ELECTRA 模型进行几步骤的微调。
• "task_names": ["task_name"]：指定要训练的任务。此处使用列表是因为代码库理论上支持多任务学习，但请注意，这一功能尚未经过充分测试。
• "model_size": one of "small", "base", or "large"：决定模型的大小；您必须将其设置为与预训练模型相同的尺寸。
• "do_train" and "do_eval"：训练和/或评估模型（默认均为真）。若要使用 "do_eval": true 且 "do_train": false"，则需要指定 init_checkpoint，例如 python3 run_finetuning.py --data-dir $DATA_DIR --model-name electra_base --hparams '{"model_size": "base", "task_names": ["mnli"], "do_train": false, "do_eval": true, "init_checkpoint": "<data-dir>/models/electra_base/finetuning_models/mnli_model_1"}'。
• "num_trials": n：如果大于 1，则会使用不同的随机种子进行多次微调和评估。
• "learning_rate": lr, "train_batch_size": n 等参数可用于调整训练超参数。
• "model_hparam_overrides": {"hidden_size": n, "num_hidden_layers": m} 等参数可用于修改底层 Transformer 的超参数（"model_size" 标志会设置默认值）。

设置

获取一个预训练的 ELECTRA 模型，方法有两种：一是自行训练（参见上述预训练说明），二是下载官方发布的 ELECTRA 权重，并将其解压到 $DATA_DIR/models 目录下（例如，如果您使用大型模型，则应有一个 $DATA_DIR/models/electra_large 目录）。

在 GLUE 任务上微调 ELECTRA

通过运行此脚本下载 GLUE 数据。然后执行以下命令来整理数据：mv CoLA cola && mv MNLI mnli && mv MRPC mrpc && mv QNLI qnli && mv QQP qqp && mv RTE rte && mv SST-2 sst && mv STS-B sts && mv diagnostic/diagnostic.tsv mnli && mkdir -p $DATA_DIR/finetuning_data && mv * $DATA_DIR/finetuning_data。

接着运行 run_finetuning.py。例如，要在 MNLI 上微调 ELECTRA-Base 模型：

python3 run_finetuning.py --data-dir $DATA_DIR --model-name electra_base --hparams '{"model_size": "base", "task_names": ["mnli"]}'

或者使用上述说明预训练的小模型在 CoLA 上进行微调：

python3 run_finetuning.py --data-dir $DATA_DIR --model-name electra_small_owt --hparams '{"model_size": "small", "task_names": ["cola"]}'

在问答任务上微调 ELECTRA

该代码支持 SQuAD 1.1 和 2.0，以及 2019 年 MRQA 共享任务中的数据集。

• SQuAD 1.1：下载 train 和 dev 数据集，并将其移动到 $DATA_DIR/finetuning_data/squadv1/(train|dev).json。
• SQuAD 2.0：从 SQuAD 官网下载数据集，并将其移动到 $DATA_DIR/finetuning_data/squad/(train|dev).json。
• MRQA 任务：从 这里下载数据。将数据移至 $DATA_DIR/finetuning_data/(newsqa|naturalqs|triviaqa|searchqa)/(train|dev).jsonl。

然后运行（例如）：

python3 run_finetuning.py --data-dir $DATA_DIR --model-name electra_base --hparams '{"model_size": "base", "task_names": ["squad"]}'

此仓库使用由 SQuAD 作者和 MRQA 共享任务发布的官方评估代码来计算指标。

在序列标注任务上微调 ELECTRA

从 这里下载 CoNLL-2000 文本分块数据集，并将其放置在 $DATA_DIR/finetuning_data/chunk/(train|dev).txt 下。然后运行：

python3 run_finetuning.py --data-dir $DATA_DIR --model-name electra_base --hparams '{"model_size": "base", "task_names": ["chunk"]}'

添加新任务

在新任务上运行的最简单方法是实现一个新的 finetune.task.Task，将其添加到 finetune.task_builder.py 中，然后像往常一样使用 run_finetuning.py。对于分类、问答和序列标注任务，您可以继承 finetune.classification.classification_tasks.ClassificationTask、finetune.qa.qa_tasks.QATask 或 finetune.tagging.tagging_tasks.TaggingTask。

在数据预处理方面，我们使用与 BERT 相同的分词器。

预期结果

以下是 ELECTRA 在各项任务上的预期结果（分块任务使用测试集，其他任务使用开发集）。请注意，微调过程中的方差可能相当大，因此在多次从同一检查点进行微调时，某些任务的得分可能会出现较大波动。以下分数显示了大量随机种子下的中位数性能。ELECTRA-Small/Base/Large 是我们已发布的模型。ELECTRA-Small-OWT 是上述基于 OpenWebText 训练的模型（由于训练时间较短且数据集较小，其性能略逊于 ELECTRA-Small）。

	CoLA	SST	MRPC	STS	QQP	MNLI	QNLI	RTE	SQuAD 1.1	SQuAD 2.0	分块
指标	MCC	Acc	Acc	Spearman	Acc	Acc	Acc	Acc	EM	EM	F1
ELECTRA-Large	69.1	96.9	90.8	92.6	92.4	90.9	95.0	88.0	89.7	88.1	97.2
ELECTRA-Base	67.7	95.1	89.5	91.2	91.5	88.8	93.2	82.7	86.8	80.5	97.1
ELECTRA-Small	57.0	91.2	88.0	87.5	89.0	81.3	88.4	66.7	75.8	70.1	96.5
ELECTRA-Small-OWT	56.8	88.3	87.4	86.8	88.3	78.9	87.9	68.5	--	--	--

有关预训练期间各模型的损失和训练曲线，请参阅此处。

Electric

要训练 Electric，请使用与 ELECTRA 相同的预训练脚本和命令。在超参数中传递 "electra_objective": false 和 "electric_objective": true。我们计划很快发布预训练的 Electric 模型！

引用

如果您在论文中使用此代码，请引用原始论文：

@inproceedings{clark2020electra,
  title = {{ELECTRA}: Pre-training Text Encoders as Discriminators Rather Than Generators},
  author = {Kevin Clark and Minh-Thang Luong and Quoc V. Le and Christopher D. Manning},
  booktitle = {ICLR},
  year = {2020},
  url = {https://openreview.net/pdf?id=r1xMH1BtvB}
}

如果您使用此代码训练 Electric，请引用 Electric 论文：

@inproceedings{clark2020electric,
  title = {Pre-Training Transformers as Energy-Based Cloze Models},
  author = {Kevin Clark and Minh-Thang Luong and Quoc V. Le and Christopher D. Manning},
  booktitle = {EMNLP},
  year = {2020},
  url = {https://www.aclweb.org/anthology/2020.emnlp-main.20.pdf}
}

联系方式

如需帮助或遇到 ELECTRA 使用问题，请提交 GitHub 问题。

如需与 ELECTRA 相关的个人沟通，请联系 Kevin Clark（kevclark@cs.stanford.edu）。

ELECTRA 快速上手指南

ELECTRA 是一种高效的自监督语言表示学习方法。它通过训练模型区分“真实”输入 token 和由另一个神经网络生成的“伪造”token（类似 GAN 的判别器），从而在较少计算资源下实现强大的预训练效果。本指南将帮助你快速搭建环境并运行一个小规模的 ELECTRA 模型。

环境准备

在开始之前，请确保你的系统满足以下要求：

• 操作系统: Linux 或 macOS (Windows 需配合 WSL 使用)
• Python: 版本 3.x
• 深度学习框架: TensorFlow 1.15 (目前官方代码主要基于 TF 1.x，暂不直接支持 TF 2.x)
• 核心依赖:
  • NumPy
  • scikit-learn
  • SciPy

国内加速建议： 由于 TensorFlow 1.15 较旧且官方源下载可能较慢，建议使用清华或阿里镜像源安装依赖：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow==1.15.0 numpy scikit-learn scipy

安装步骤

1. 克隆代码仓库 首先获取 ELECTRA 的源代码：

   git clone https://github.com/google-research/electra.git
   cd electra
   

2. 安装 Python 依赖 确保已安装上述要求的 TensorFlow 1.15 及其他库。如果未安装，执行：

   pip install -r requirements.txt
   # 如果没有 requirements.txt，请手动安装：
   # pip install tensorflow==1.15.0 numpy scikit-learn scipy
   

3. 准备数据与词汇表 创建一个数据目录（例如 $DATA_DIR），并下载必要的文件：

   • 下载词汇表 (与 BERT uncased 相同)：

     mkdir -p $DATA_DIR
     wget -O $DATA_DIR/vocab.txt https://storage.googleapis.com/electra-data/vocab.txt
     

   • 下载预训练语料 (以 OpenWebText 为例)：

     cd $DATA_DIR
     # 下载并解压 (约 12GB)
     wget https://skylion007.github.io/OpenWebTextCorpus/openwebtext.tar.xz
     tar xf openwebtext.tar.xz
     mv openwebtext openwebtext_corpus
     

4. 构建预训练数据集 运行脚本将原始文本转换为 ELECTRA 所需的 TFRecord 格式：

   python3 build_openwebtext_pretraining_dataset.py --data-dir $DATA_DIR --num-processes 5
   

   注：生成的 tfrecord 文件将存储在 $DATA_DIR/pretrain_tfrecords，约占 30GB 磁盘空间。

基本使用

以下示例展示如何从头预训练一个小型的 ELECTRA 模型（12 层，隐藏层大小 256）。

1. 启动预训练

在单张 GPU（如 Tesla V100）上运行以下命令。完全训练约需 4 天，但训练 20 万步（约 10 小时）即可获得不错的效果。

python3 run_pretraining.py --data-dir $DATA_DIR --model-name electra_small_owt

自定义训练参数： 你可以通过 --hparams 参数调整模型大小、步数或开启调试模式。例如，仅进行少量步骤的调试训练：

python3 run_pretraining.py --data-dir $DATA_DIR --model-name electra_debug --hparams '{"debug": true}'

常用超参数说明：

• "model_size": 可选 "small", "base", "large"。
• "num_train_steps": 控制预训练总步数。
• "learning_rate": 学习率。

2. 评估预训练模型

训练完成后，可以在 OpenWebText 数据上评估生成器（Generator）和判别器（Discriminator）的性能：

python3 run_pretraining.py --data-dir $DATA_DIR --model-name electra_small_owt --hparams '{"do_train": false, "do_eval": true}'

评估结果将保存在 $DATA_DIR/electra_small_owt/results。

3. 下游任务微调 (Fine-tuning)

使用预训练好的模型进行下游任务（如 GLUE 基准测试中的 RTE 任务）微调：

python3 run_finetuning.py --data-dir $DATA_DIR --model-name electra_small_owt --hparams '{"task_names": ["rte"], "model_size": "small", "learning_rate": 1e-4}'

• task_names: 指定要微调的任务名称列表。
• model_size: 必须与预训练模型的大小保持一致。
• 微调后的模型权重默认保存在 $DATA_DIR/electra_small_owt/finetuning_models。