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 的基礎原理有更深入的了解。

審核編輯 ：李倩