如何在 MLM 和 NSP 的新域上从头开始训练 BERT？

Question

我正在尝试使用我自己的数据集和 HuggingFace 库从头开始训练 BERT 模型。我想以具有原始 BERT 模型的确切架构的方式训练模型。

在原始论文中，它指出：“BERT 在两个任务上进行训练：预测随机掩码标记 (MLM) 和预测两个句子是否相互跟随 (NSP)。 SCIBERT 遵循与 BERT 相同的架构，但在科学文本上进行了预训练。”

我试图了解如何在上述两个任务上训练模型。目前，我将模型初始化如下：

from transformers import BertForMaskedLM
model = BertForMaskedLM(config=config)

但是，它只适用于 MLM，而不适用于 NSP。我如何也可以使用 NSP 初始化和训练模型，或者我的原始方法是否可以正常使用？

我的假设是

1. 使用 BertForPreTraining 初始化（对于 MLM 和 NSP），或者

2. 用BertForMaskedLM完成训练后， 初始化相同的模型并再次训练 BertForNextSentencePrediction（但这种方法的计算和 资源会花费两倍……）

我不确定哪一种方法是正确的。任何见解或建议将不胜感激。

Answer 1

我建议做以下事情：

1. 首次对 MLM 目标进行预训练 BERT。 HuggingFace 提供了一个脚本，专门用于根据您自己的数据在 MLM 目标上训练 BERT。你可以找到它here。正如您在run_mlm.py 脚本中所见，它们使用AutoModelForMaskedLM，您可以指定任何您想要的架构。

2. 其次，如果要训练下一句预测任务，您可以定义一个BertForPretraining 模型（其顶部同时具有 MLM 和 NSP 头），然后从您在步骤 1 中训练的模型中加载权重，然后在下一个句子预测任务上进一步对其进行预训练。

更新：显然下一句预测任务确实有助于提高 BERT 在某些 GLUE 任务上的性能。请参阅 BERT 的作者this talk。

Answer 2

您可以使用 BertForPretraining TextDatasetForNextSentencePrediction DataCollatorForLanguageModeling 和 Trainer 的组合轻松地从头开始训练 BERT 进行 MLM 和 NSP 任务。

我不建议您先训练模型 MLM，然后再训练 NSP，这可能会导致灾难性遗忘。它基本上是忘记了你从之前的培训中学到的东西。

1. 加载预训练的分词器。

from transformers import BertTokenizer
bert_cased_tokenizer = BertTokenizer.from_pretrained("/path/to/pre-trained/tokenizer/for/new/domain", do_lower_case=False)

1. 使用BertForPretraining 初始化您的模型

from transformers import BertConfig, BertForPreTraining
config = BertConfig()
model = BertForPreTraining(config)

1. 为 NSP 任务创建数据集。 TextDatasetForNextSentencePrediction 将标记并为句子创建标签。您的数据集应位于以下format：（或者您可以只修改现有代码）

(1) 每行一个句子。理想情况下，这些应该是实际的句子 (2) 文档之间的空行

Sentence-1 From Document-1
Sentence-2 From Document-1
Sentence-3 From Document-1
...

Sentence-1 From Document-2
Sentence-2 From Document-2
Sentence-3 From Document-2

from transformers import TextDatasetForNextSentencePrediction
dataset = TextDatasetForNextSentencePrediction(
    tokenizer=bert_cased_tokenizer,
    file_path="/path/to/your/dataset",
    block_size = 256
)

4. 使用DataCollatorForLanguageModeling 屏蔽和传递从TextDatasetForNextSentencePrediction 生成的标签。 DataCollatorForNextSentencePrediction 一直是 removed，因为它与 DataCollatorForLanguageModeling 做同样的事情

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=bert_cased_tokenizer, 
    mlm=True,
    mlm_probability= 0.15
)

5. 训练和保存

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir= "/path/to/output/dir/for/training/arguments"
    overwrite_output_dir=True,
    num_train_epochs=2,
    per_gpu_train_batch_size= 16,
    save_steps=10_000,
    save_total_limit=2,
    prediction_loss_only=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("path/to/your/model")