深入解析BERT掩码语言模型(MLM)目标函数：从理论到PyTorch实现

1次阅读

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

BERT（Bidirectional Encoder Representations from Transformers）的核心创新之一就是采用了掩码语言模型（Masked Language Model, MLM）作为预训练目标。与传统的单向语言模型不同，MLM 通过随机掩盖输入序列中的部分 token（标记），让模型基于上下文双向预测被掩盖的内容。这种设计使 BERT 能够捕获更丰富的上下文信息。

MLM 的损失函数采用标准的交叉熵（Cross-Entropy）形式。对于被掩盖的 token，其损失计算如下：

$$
\mathcal{L}{MLM} = -\sum)
$$} \log P(x_i | x_{\backslash i

其中：
– $masked$ 表示被掩盖的 token 位置集合
– $x_i$ 是第 i 个位置的真实 token
– $x_{\backslash i}$ 表示除第 i 个位置外的所有输入 token

import torch
import random

def random_mask(input_ids, mask_token_id, vocab_size, mask_prob=0.15):
    """
    对输入序列进行随机掩码处理
    :param input_ids: 输入 token ID 序列 [batch_size, seq_len]
    :param mask_token_id: [MASK]标记的 ID
    :param vocab_size: 词表大小
    :param mask_prob: 掩码概率
    :return: 处理后的 input_ids, 掩码位置的 labels
    """
    labels = input_ids.clone()
    # 创建随机掩码矩阵
    probability_matrix = torch.full(labels.shape, mask_prob)
    # 特殊 token 不参与掩码（如 [CLS], [SEP] 等）special_tokens_mask = (input_ids == 0) | (input_ids == 1)  # 示例，实际需根据 tokenizer 调整
    probability_matrix.masked_fill_(special_tokens_mask, value=0.0)

    masked_indices = torch.bernoulli(probability_matrix).bool()
    labels[~masked_indices] = -100  # 只计算被掩盖位置的损失

    # 80% 的概率替换为[MASK]
    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
    input_ids[indices_replaced] = mask_token_id

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

    return input_ids, labels

import torch.nn as nn

def mlm_loss(logits, labels):
    """
    计算 MLM 损失
    :param logits: 模型输出 [batch_size, seq_len, vocab_size]
    :param labels: 标签 [batch_size, seq_len]，其中非掩码位置为 -100
    :return: 标量损失值
    """
    loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
    # 将 logits 和 labels 展平
    active_loss = labels.view(-1) != -100
    active_logits = logits.view(-1, logits.size(-1))[active_loss]
    active_labels = labels.view(-1)[active_loss]

    return loss_fct(active_logits, active_labels)