BERT模型MLM训练策略解析：15%掩码token的处理机制与实现细节

1次阅读

共计 1810 个字符，预计需要花费 5 分钟才能阅读完成。

BERT（Bidirectional Encoder Representations from Transformers）的核心预训练任务之一是掩码语言模型（Masked Language Model, MLM）。其核心思想是：随机遮盖输入句子中的部分 token（通常为 15%），让模型基于上下文预测被遮盖的原始单词。这种自监督训练方式使模型能够学习深层的双向语言表征。

对于被选中的 15%token 中的 80%，会直接替换为特殊符号[MASK]。这是最典型的情境，模型需要根据上下文推断被遮盖的词。例如：

原始句子："the cat sat on the mat"
处理后："the [MASK] sat on the mat"

10% 的概率会保留原始单词不变。这种策略迫使模型不仅要学会预测缺失词，还需要理解当前词是否合理。例如：

原始句子："the cat sat on the mat"
处理后："the cat sat on the mat"（无变化）

剩余 10% 的概率会替换为词表中的随机单词。这种噪声注入机制增强了模型的纠错能力。例如：

原始句子："the cat sat on the mat"
处理后："the dog sat on the mat"（"cat"→随机 "dog"）

数学表达为：

$$
\text{替换策略} =
\begin{cases}
[MASK] & \text{概率} = 0.8 \
\text{原词} & \text{概率} = 0.1 \
\text{随机词} & \text{概率} = 0.1
\end{cases}
$$

import torch
import random

def mask_tokens(inputs, tokenizer, mask_prob=0.15):
    """
    inputs: 输入 token_id 张量 [batch_size, seq_len]
    mask_prob: 掩码总概率（默认 15%）"""
    labels = inputs.clone()
    # 生成掩码位置矩阵（15% 概率）prob_matrix = torch.full(labels.shape, mask_prob)
    masked_indices = torch.bernoulli(prob_matrix).bool()

    with torch.no_grad():
        # 80% 概率替换为[MASK]
        mask_token = tokenizer.mask_token_id
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
        inputs[indices_replaced] = mask_token

        # 10% 概率替换为随机词
        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
        random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long)
        inputs[indices_random] = random_words[indices_random]

        # 剩余 10% 保持原词（labels 仍记录原始 id）return inputs, labels