BERT掩码语言模型(MLM)预训练任务示意图解析与实现优化

1次阅读

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

BERT 模型的预训练阶段包含两个核心任务：掩码语言模型（MLM）和下一句预测（NSP）。其中 MLM 任务通过随机遮盖输入文本中的部分词汇，要求模型基于上下文预测被遮盖的原始词汇，从而使模型学习深层次的语义表示。这一机制使得 BERT 在多种 NLP 任务中展现出强大性能，包括文本分类、命名实体识别、问答系统等。

MLM 任务的核心价值在于：
– 通过双向上下文建模克服传统语言模型的单向性限制
– 学习词汇在不同语境下的多义表示
– 为下游任务提供通用的语义编码基础

[输入序列] The quick brown fox jumps over the lazy dog
[掩码后] The [MASK] brown fox [MASK] over the lazy [MASK]

BERT 采用三种掩码策略组合使用：

全词掩码（Whole Word Masking）
对完整词汇进行遮盖，例如将 ”jumping” 整体替换为[MASK]
需配合 WordPiece 分词器使用
缓解子词掩码带来的语义碎片化问题
子词掩码（Subword Masking）
对 WordPiece 分词后的子词单元进行遮盖
例如将 ”jumping” 分为 ”jump” 和 ”##ing” 后随机遮盖部分片段
增强模型对词缀和罕见词的处理能力
字符级掩码（Character-level Masking）
对单个字符进行随机遮盖
主要用于拼音文字语言处理
需配合字符级编码器使用

MLM 任务的损失函数采用交叉熵损失：
$$
\mathcal{L}{MLM} = -\sum)
$$
其中 $M$ 表示被掩码的词汇集合，$w_{\backslash M}$ 表示未被掩码的上下文。} \log P(w_i|w_{\backslash M

实际预训练过程中主要面临三大挑战：

计算资源消耗
标准 BERT-large 模型需 16-64 块 GPU 训练数天
显存占用随序列长度平方级增长
梯度同步通信开销大
收敛速度问题
早期训练阶段损失下降缓慢
高频词与低频词学习速度不均衡
固定掩码比例导致效率低下
OOV 处理困境
罕见词因采样不足导致表示质量差
专业领域术语覆盖不足
多语言场景下的字符集冲突

def dynamic_masking(
    input_ids: torch.Tensor,
    mask_prob: float = 0.15,
    vocab_size: int = 30522
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    动态生成掩码位置的实现
    Args:
        input_ids: 输入 token id 张量 [batch_size, seq_len]
        mask_prob: 掩码比例 (default: 0.15)
        vocab_size: 词表大小 (default: BERT-base 30522)
    Returns:
        masked_input: 掩码后的输入张量
        mask_labels: 被掩码位置的原始 token id
    """
    # 初始化标签为 -100（忽略 loss 计算）labels = input_ids.clone()
    probability_matrix = torch.full(labels.shape, mask_prob)

    # 特殊 token 不参与掩码
    special_tokens_mask = [token in [0, 101, 102] for token in input_ids.tolist()]
    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), 
        value=0.0
    )

    # 生成随机掩码
    masked_indices = torch.bernoulli(probability_matrix).bool()
    labels[~masked_indices] = -100  # 只计算被掩码位置的 loss

    # 80% 概率替换为[MASK]
    mask_token = 103  # [MASK]的 token id
    replace_mask = torch.bernoulli(torch.full(labels.shape, 0.8)).bool()
    masked_input = torch.where(masked_indices & replace_mask, mask_token, input_ids)

    # 10% 概率随机替换其他词
    random_words = torch.randint(vocab_size, labels.shape, dtype=torch.long)
    random_replace = torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
    masked_input = torch.where(
        masked_indices & ~replace_mask & random_replace, 
        random_words, 
        masked_input
    )

    return masked_input, labels

采用分层衰减学习率策略：
– 底层 embedding 层：1e-5
– 中间隐藏层：3e-5
– 顶层输出层：5e-5
– 每 1000 步线性衰减 10%

optimizer.zero_grad()
for i, (batch) in enumerate(dataloader):
    loss = model(batch).loss
    loss.backward()

    if (i+1) % 4 == 0:  # 每 4 个 batch 更新一次
        optimizer.step()
        optimizer.zero_grad()

在 GLUE 基准测试上的效果对比（BERT-base）：