如何进行MLM训练

1. 关于MLM

1.1 背景

作为 Bert 预训练的两大任务之一，MLM 和 NSP 大家应该并不陌生。其中，NSP 任务在后续的一些预训练任务中经常被嫌弃，例如 Roberta 中将 NSP 任务直接放弃，Albert 中将 NSP 替换成了句子顺序预测。

这主要是因为 NSP 作为一个分类任务过于简单，对模型的学习并没有太大的帮助，而 MLM 则被多数预训练模型保留下来。由 Roberta的实验结果也可以证明，Bert 的主要能力应该是来自于 MLM 任务的训练。

Bert为代表的预训练语言模型是在大规模语料的基础上训练以获得的基础的学习能力，而实际应用时，我们所面临的语料或许具有某些特殊性，这就使得重新进行 MLM 训练具有了必要性。

1.2 如何进行MLM训练

1.2.1 什么是MLM

MLM 的训练，在不同的预训练模型中其实是有所不同的。今天介绍的内容以最基础的 Bert 为例。

Bert的MLM是静态mask，而在后续的其他预训练模型中，这一策略通常被替换成了动态mask。除此之外还有 whole word mask 的模型，这些都不在今天的讨论范围内。

所谓 mask language model 的任务，通俗来讲，就是将句子中的一部分token替换掉，然后根据句子的剩余部分，试图去还原这部分被mask的token。

1.2.2 如何Mask

mask 的比例一般是15%，这一比例也被后续的多数模型所继承，而在最初BERT 的论文中，没有对这一比例的界定给出具体的说明。在我的印象中，似乎是知道后来同样是Google提出的 T5 模型的论文中，对此进行了解释，对 mask 的比例进行了实验，最终得出结论，15%的比例是最合理的（如果我记错了，还请指正）。

将15%的token选出之后，并不是所有的都替换成[mask]标记符。实际操作是：

• 从这15%选出的部分中，将其中的80%替换成[mask]；
• 10%替换成一个随机的token；
• 剩下的10%保留原来的token。

这样做可以提高模型的鲁棒性。这个比例也可以自己控制。

到这里可能有同学要问了，既然有10%保留不变的话，为什么不干脆只选择15%*90% = 13.5%的token呢？如果看完后面的代码，就会很清楚地理解这个问题了。

先说结论：因为 MLM 的任务是将选出的这15%的token全部进行预测，不管这个token是否被替换成了[mask]，也就是说，即使它被保留了原样，也还是需要被预测的。

2. 代码部分

2.1 背景

介绍完了基础内容之后，接下来的内容，我将基于 transformers 模块，介绍如何进行 mask language model 的训练。

其实 transformers 模块中，本身是提供了 MLM 训练任务的，模型都写好了，只需要调用它内置的 trainer 和datasets模块即可。感兴趣的同学可以去 huggingface 的官网搜索相关教程。

然而我觉得 datasets 每次调用的时候都要去写数据集的py文件，对arrow的数据格式不熟悉的话还很容易出错，而且 trainer 我觉得也不是很好用，任何一点小小的修改都挺费劲（就是它以为它写的很完备，考虑了用户的所有需求，但是实际上有一些冗余的部分）。

所以我就参考它的实现方式，把它的代码拆解，又按照自己的方式重新组织了一下。

2.2 准备工作

首先在写核心代码之前，先做好准备工作。
import 所有需要的模块：

import os
import json
import copy
from tqdm.notebook import tqdm

import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from transformers import BertForMaskedLM, BertTokenizerFast

然后写一个config类，将所有参数集中起来：

class Config:
    def __init__(self):
        pass
    
    def mlm_config(
        self, 
        mlm_probability=0.15, 
        special_tokens_mask=None,
        prob_replace_mask=0.8,
        prob_replace_rand=0.1,
        prob_keep_ori=0.1,
    ):
        """
        :param mlm_probability: 被mask的token总数
        :param special_token_mask: 特殊token
        :param prob_replace_mask: 被替换成[MASK]的token比率
        :param prob_replace_rand: 被随机替换成其他token比率
        :param prob_keep_ori: 保留原token的比率
        """
        assert sum([prob_replace_mask, prob_replace_rand, prob_keep_ori]) == 1,                 ValueError("Sum of the probs must equal to 1.")
        self.mlm_probability = mlm_probability
        self.special_tokens_mask = special_tokens_mask
        self.prob_replace_mask = prob_replace_mask
        self.prob_replace_rand = prob_replace_rand
        self.prob_keep_ori = prob_keep_ori
        
    def training_config(
        self,
        batch_size,
        epochs,
        learning_rate,
        weight_decay,
        device,
    ):
        self.batch_size = batch_size
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.device = device
        
    def io_config(
        self,
        from_path,
        save_path,
    ):
        self.from_path = from_path
        self.save_path = save_path

接着就是设置各种配置：

config = Config()
config.mlm_config()
config.training_config(batch_size=4, epochs=10, learning_rate=1e-5, weight_decay=0, device='cuda:0')
config.io_config(from_path='/data/BERTmodels/huggingface/chinese_wwm/', 
                 save_path='./finetune_embedding_model/mlm/')

最后创建BERT模型。注意，这里的 tokenizer 就是一个普通的 tokenizer，而BERT模型则是带了下游任务的 BertForMaskedLM，它是 transformers 中写好的一个类，

bert_tokenizer = BertTokenizerFast.from_pretrained(config.from_path)
bert_mlm_model = BertForMaskedLM.from_pretrained(config.from_path)

2.3 数据集

因为舍弃了datasets这个包，所以我们现在需要自己实现数据的输入了。方案就是使用 torch 的 Dataset 类。这个类一般在构建 DataLoader 的时候，会与一个聚合函数一起使用，以实现对batch的组织。而我这里偷个懒，就没有写聚合函数，batch的组织方法放在dataset中进行。

在这个类中，有一个 mask tokens 的方法，作用是从数据中选择出所有需要mask 的token，并且采用三种mask方式中的一个。这个方法是从transformers 中拿出来的，将其从类方法转为静态方法测试之后，再将其放在自己的这个类中为我们所用。仔细阅读这一段代码，也就可以回答1.2.2 中提出的那个问题了。

取batch的原理很简单，一开始我们将原始数据deepcopy备份一下，然后每次从中截取一个batch的大小，这个时候的当前数据就少了一个batch，我们定义这个类的长度为当前长度除以batch size向下取整，所以当类的长度变为0的时候，就说明这一个epoch的所有step都已经执行结束，要进行下一个epoch的训练，此时，再将当前数据变为原始数据，就可以实现对epoch的循环了。

class TrainDataset(Dataset):
    """
    注意：由于没有使用data_collator，batch放在dataset里边做，
    因而在dataloader出来的结果会多套一层batch维度，传入模型时注意squeeze掉
    """
    def __init__(self, input_texts, tokenizer, config):
        self.input_texts = input_texts
        self.tokenizer = tokenizer
        self.config = config
        self.ori_inputs = copy.deepcopy(input_texts)
        
    def __len__(self):
        return len(self.input_texts) // self.config.batch_size
    
    def __getitem__(self, idx):
        batch_text = self.input_texts[: self.config.batch_size]
        features = self.tokenizer(batch_text, max_length=512, truncation=True, padding=True, return_tensors='pt')
        inputs, labels = self.mask_tokens(features['input_ids'])
        batch = {"inputs": inputs, "labels": labels}
        self.input_texts = self.input_texts[self.config.batch_size: ]
        if not len(self):
            self.input_texts = self.ori_inputs
        
        return batch
        
    def mask_tokens(self, inputs):
        """
        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
        """
        labels = inputs.clone()
        # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
        probability_matrix = torch.full(labels.shape, self.config.mlm_probability)
        if self.config.special_tokens_mask is None:
            special_tokens_mask = [
                self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
            ]
            special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
        else:
            special_tokens_mask = self.config.special_tokens_mask.bool()

        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100  # We only compute loss on masked tokens

        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
        indices_replaced = torch.bernoulli(torch.full(labels.shape, self.config.prob_replace_mask)).bool() & masked_indices
        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)

        # 10% of the time, we replace masked input tokens with random word
        current_prob = self.config.prob_replace_rand / (1 - self.config.prob_replace_mask)
        indices_random = torch.bernoulli(torch.full(labels.shape, current_prob)).bool() & masked_indices & ~indices_replaced
        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)
        inputs[indices_random] = random_words[indices_random]

        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
        return inputs, labels

然后取一些用于训练的语料，格式很简单，就是把所有文本放在一个list里边，注意长度不要超过512个token，不然多出来的部分就浪费掉了。可以做适当的预处理。

[
"这是一条文本",
"这是另一条文本",
...,
]

然后构建dataloader：

train_dataset = TrainDataset(training_texts, bert_tokenizer, config)
train_dataloader = DataLoader(train_dataset)

2.4 训练

构建一个训练方法，输入参数分别是我们实例化好的待训练模型，数据集，还有config：

def train(model, train_dataloader, config):
    """
    训练
    :param model: nn.Module
    :param train_dataloader: DataLoader
    :param config: Config
    ---------------
    ver: 2021-11-08
    by: changhongyu
    """
    assert config.device.startswith('cuda') or config.device == 'cpu', ValueError("Invalid device.")
    device = torch.device(config.device)
    
    model.to(device)
    
    if not len(train_dataloader):
        raise EOFError("Empty train_dataloader.")
        
    param_optimizer = list(model.named_parameters())
    no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {"params": [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], "weight_decay": 0.01},
        {"params": [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], "weight_decay": 0.0}]
    
    optimizer = AdamW(params=optimizer_grouped_parameters, lr=config.learning_rate, weight_decay=config.weight_decay)
    
    for cur_epc in tqdm(range(int(config.epochs)), desc="Epoch"):
        training_loss = 0
        print("Epoch: {}".format(cur_epc+1))
        model.train()
        for step, batch in enumerate(tqdm(train_dataloader, desc='Step')):
            input_ids = batch['inputs'].squeeze(0).to(device)
            labels = batch['labels'].squeeze(0).to(device)
            loss = model(input_ids=input_ids, labels=labels).loss
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            model.zero_grad()
            training_loss += loss.item()
        print("Training loss: ", training_loss)

调用它训练几轮：

train(model=bert_mlm_model, train_dataloader=train_dataloader, config=config)

2.5 保存和加载

使用过预训练模型的同学应该都了解，普通的bert有两项输出，分别是:

• 每一个token对应的768维编码结果;
• 以及用于表征整个句子的句子特征。

其中，这个句子特征是由模型中的一个 Pooler 模块对原句池化得来的。可是这个Pooler的训练，并不是由 MLM 任务来的，而是由 NSP任务中来的。

由于没有 NSP 任务，所以无法对 Pooler 进行训练，故而没有必要在模型中加入 Pooler。所以在保存的时候需要分别保存 embedding和 encoder， 加载的时候也需要分别读取 embedding 和 encoder，这样训练出来的模型拿不到 CLS 层的句子表征。如果需要的话，可以手动pooling 。

torch.save(bert_mlm_model.bert.embeddings.state_dict(), os.path.join(config.save_path, 'bert_mlm_ep_{}_eb.bin'.format(config.epochs)))
torch.save(bert_mlm_model.bert.encoder.state_dict(), os.path.join(config.save_path, 'bert_mlm_ep_{}_ec.bin'.format(config.epochs)))

加载的话，也是实例化完bert模型之后，用bert的 embedding 组件和 encoder 组件分别读取这两个权重文件即可。

到这里，本期内容就全部结束了，希望看完这篇博客的同学，能够对 Bert 的基础原理有更深入的了解。

审核编辑 ：李倩