BERT预训练实战：如何通过MLM和NSP任务构建高效语言模型

1次阅读

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

在自然语言处理领域，BERT（Bidirectional Encoder Representations from Transformers）已成为基石模型。其核心创新在于通过掩码语言模型（MLM）和下一句预测（NSP）两个预训练任务，使模型能学习深层次的上下文表征。然而在实际工程化过程中，开发者常面临预训练效率低下、任务融合困难等挑战。本文将深入解析这两大任务的实现细节，并提供可落地的优化方案。

MLM 任务的核心是随机掩盖输入序列中 15% 的 token（基于 BERT 论文实证效果），让模型预测被掩盖的原始词。这个比例经过严格验证：

低于 10% 会导致训练信号不足
高于 20% 可能破坏语义完整性

具体实现时需注意：

在构建 DataLoader 时，需要先对原始文本进行 WordPiece 分词
随机选择 15% 的 token 进行以下处理：
80% 概率替换为 [MASK]
10% 概率替换为随机词
10% 概率保持原词（增加模型纠错能力）

# PyTorch 实现示例
def create_masked_lm_predictions(tokens, mask_prob=0.15):
    cand_indices = [i for i, token in enumerate(tokens) if token != '[CLS]' and token != '[SEP]']
    num_to_mask = min(int(len(cand_indices) * mask_prob), 512-2)  # 防止溢出
    random.shuffle(cand_indices)

    masked_tokens = tokens.copy()
    labels = [-100] * len(tokens)  # PyTorch 的 ignore_index

    for index in cand_indices[:num_to_mask]:
        # 80-10-10 策略
        rand_prob = random.random()
        if rand_prob < 0.8:
            masked_tokens[index] = '[MASK]'
        elif rand_prob < 0.9:
            masked_tokens[index] = random.choice(vocab_list)
        labels[index] = vocab_dict[tokens[index]]

    return masked_tokens, labels

NSP 任务要求模型判断两个句子是否为连续文本。关键点在于采样策略：

同文档采样 ：从同一文档中取连续句子作为正例，随机组合作为负例
跨文档采样 ：从不同文档随机取句构建负例（噪声更大但数据多样性更强）

实验表明，同文档采样在领域特定任务表现更好，而跨文档采样对通用语料更有效。建议根据下游任务需求选择：

def get_next_sentence_sample(sent_a, sent_b, corpus):
    # 50% 概率选择真实下一句
    if random.random() < 0.5:
        is_next = True
        sent_b = get_actual_next_sentence(sent_a, corpus)
    else:
        is_next = False
        sent_b = get_random_sentence(corpus)
    return sent_a, sent_b, is_next

两种任务的损失需要加权合并，典型比例为 1:1（MLM loss + NSP loss）：

# 带梯度累积的训练步骤
def train_step(batch, model, optimizer, grad_accum_steps=4):
    input_ids, segment_ids, input_mask, lm_labels, is_next = batch

    with autocast():  # 混合精度训练
        mlm_loss, nsp_loss = model(
            input_ids=input_ids,
            token_type_ids=segment_ids,
            attention_mask=input_mask,
            masked_lm_labels=lm_labels,
            next_sentence_label=is_next
        )
        total_loss = mlm_loss + nsp_loss

    # 梯度累积
    total_loss = total_loss / grad_accum_steps
    scaler.scale(total_loss).backward()

    if (step + 1) % grad_accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

使用 AMP（Automatic Mixed Precision）可显著减少显存占用：

16 位浮点计算：节省约 50% 显存
32 位主权重：保持数值稳定性

实测在 V100 上训练 BERT-base 时：

FP32 模式：约需 16GB 显存
AMP 模式：仅需 9GB 显存

DataParallel 已逐渐被更高效的方案取代：

DistributedDataParallel：
每个 GPU 维护独立进程
支持多机训练
需配合 torch.distributed.launch 使用
Horovod：
支持 TensorFlow/PyTorch
环形梯度聚合优化通信

warmup 步数需与 batch size 关联：

warmup_steps = max(1000, total_steps * 0.1)  # 至少 1000 步
if batch_size > 256:  # 大 batch 需要更长 warmup
    warmup_steps = int(warmup_steps * (batch_size / 256) ** 0.5)

当验证集指标波动时：