BERT模型MLM训练策略深度解析：15%掩码Token的处理机制与优化实践

1次阅读

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

BERT（Bidirectional Encoder Representations from Transformers）通过掩码语言模型（MLM）任务实现双向上下文建模。MLM 的核心思想是随机遮盖输入文本中的部分 token，让模型根据上下文预测被遮盖的原始内容。这种设计让模型能够学习深层次的语义表示，而不是简单地记忆词汇共现模式。

基础采样率：BERT 默认对每个 token 独立进行 15% 概率的采样，这意味着：
短文本可能实际被 mask 的 token 比例波动较大
长文本会严格趋近 15% 的比例
处理策略细分：被选中的 token 按以下比例分配：
80% 替换为 [MASK] 标记
10% 随机替换为其他词汇表 token
10% 保持原 token 不变

80% 掩码：强制模型进行真实预测任务
10% 随机替换：增强模型对错误输入的鲁棒性
10% 保持不变 ：平衡标签数据的分布，防止模型过度依赖[MASK] 标记

import torch
import random

def mlm_mask_tokens(inputs, tokenizer, mlm_prob=0.15):
    """
    inputs: 原始输入 token_ids (batch_size, seq_len)
    return: 处理后的 token_ids 和对应的 labels
    """
    labels = inputs.clone()
    # 创建概率矩阵
    prob_matrix = torch.full(labels.shape, mlm_prob)
    # 特殊 token 不参与 mask（CLS/SEP/PAD）special_tokens = [tokenizer.cls_token_id, tokenizer.sep_token_id, tokenizer.pad_token_id]
    prob_matrix[[t in special_tokens for t in labels]] = 0

    masked_indices = torch.bernoulli(prob_matrix).bool()
    labels[~masked_indices] = -100  # 忽略未 mask 位置的 loss 计算

    # 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% 随机替换
    vocab_size = len(tokenizer)
    indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
    random_words = torch.randint(0, vocab_size, labels.shape, dtype=torch.long)
    inputs[indices_random] = random_words[indices_random]

    # 剩余 10% 保持不变
    return inputs, labels

关键优化点：

使用矩阵运算替代循环，提升 GPU 利用率
通过布尔掩码实现高效的条件赋值
提前过滤特殊 token 避免无效计算

动态分块：当序列超过 512token 时：
按句子边界分割
保持重叠区域（约 10% 长度）
最后拼接各块预测结果
内存优化：
使用混合精度训练
梯度检查点技术

数据同步问题：各 GPU 应使用不同的随机种子生成 mask 模式
解决方案：
在 DataLoader 中设置 worker_init_fn
通过 distributed.get_rank()获取设备编号作为随机种子

如何设计动态调整的 mask 比例策略（如随训练轮次变化）？
对于专业领域文本，是否需要调整随机替换策略中的词汇采样分布？

《RoBERTa: A Robustly Optimized BERT Pretraining Approach》
DOI: 10.48550/arXiv.1907.11692
《ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators》
DOI: 10.48550/arXiv.2003.10555

经过实践发现，保持原 token 的 10% 比例对金融 / 医疗等专业领域文本尤为重要，因为这些领域的术语替换容易引入噪声。建议在实际应用中通过小规模实验找到最适合当前数据分布的参数组合。

正文完