BERT预训练原理详解：从掩码语言模型(MLM)到下一句预测(NSP)的实战指南

1次阅读

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

BERT（Bidirectional Encoder Representations from Transformers）是一种基于 Transformer 架构的预训练语言模型。与传统的单向语言模型不同，BERT 通过以下两个关键任务进行预训练：

掩码语言模型（MLM）：让模型学会根据上下文预测被掩盖的单词
下一句预测（NSP）：让模型理解句子间的关系

这种双向训练方式使 BERT 能够捕获更丰富的语义信息，在各种 NLP 任务中表现出色。

BERT 的输入由三部分组成：

Token Embeddings：词嵌入表示
Segment Embeddings：区分不同句子的嵌入（用于 NSP）
Position Embeddings：位置编码

MLM 的核心策略是随机掩盖输入文本中的部分单词，然后让模型预测这些被掩盖的单词。具体实现：

随机选择 15% 的 token 进行掩码
其中 80% 替换为 [MASK] 标记
10% 替换为随机单词
10% 保持不变

MLM 使用交叉熵损失函数，只计算被掩盖位置的预测损失：

criterion = nn.CrossEntropyLoss(ignore_index=-100)  # 忽略未被掩盖的位置
loss = criterion(masked_lm_logits.view(-1, vocab_size), masked_lm_labels.view(-1))

这个比例是经过实验验证的平衡点：

太低：模型学习信号不足
太高：破坏句子结构，影响模型理解

NSP 任务需要构造句子对：

正样本：实际连续的句子（50%）
负样本：随机组合的句子（50%）

模型需要预测第二个句子是否是第一个句子的实际后续：

# [CLS] token 的表示用于 NSP 分类
seq_relationship_logits = nn.Linear(hidden_size, 2)(pooled_output)

NSP 特别有利于需要理解句子关系的任务，如：

问答系统
文本蕴涵
对话系统

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# 处理 MLM
def mask_tokens(inputs, tokenizer):
    labels = inputs.clone()
    # 创建掩码矩阵 (15% 概率)
    probability_matrix = torch.full(labels.shape, 0.15)
    masked_indices = torch.bernoulli(probability_matrix).bool()

    # 80% 替换为[MASK]
    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
    inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.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% 保持不变
    return inputs, labels

import torch.nn as nn
from transformers import BertModel

class BertForPretraining(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.bert = BertModel(config)
        self.cls = BertOnlyMLMHead(config)
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
        outputs = self.bert(input_ids, attention_mask, token_type_ids)
        sequence_output = outputs.last_hidden_state
        pooled_output = outputs.pooler_output

        # MLM 预测
        prediction_scores = self.cls(sequence_output)

        # NSP 预测
        seq_relationship_score = self.seq_relationship(pooled_output)

        return prediction_scores, seq_relationship_score

# 初始化模型和优化器
model = BertForPretraining(config)
optimizer = AdamW(model.parameters(), lr=5e-5)

for epoch in range(epochs):
    for batch in dataloader:
        input_ids, attention_mask, token_type_ids, masked_lm_labels, next_sentence_labels = batch

        # 前向传播
        masked_lm_logits, seq_relationship_logits = model(input_ids, attention_mask, token_type_ids)

        # 计算损失
        mlm_loss = criterion(masked_lm_logits.view(-1, vocab_size), masked_lm_labels.view(-1))
        nsp_loss = criterion(seq_relationship_logits.view(-1, 2), next_sentence_labels.view(-1))
        total_loss = mlm_loss + nsp_loss

        # 反向传播
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()