BERT掩码语言模型(MLM)核心公式解析与工程实践优化

1次阅读

没有评论

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

BERT 的输入由三部分组成：

$$\text{Input} = \text{Token Embeddings} + \text{Segment Embeddings} + \text{Position Embeddings}$$

Token Embeddings：通过 WordPiece 分词器将输入文本转换为子词 token
Segment Embeddings：区分句子 A 和句子 B（对于单句输入全为 0）
Position Embeddings：使用固定位置编码，最大支持 512 个 token

原始 MLM 任务的掩码概率为 15%，其中：

$$P(\text{mask}) = \begin{cases}
80\% & \text{替换为[MASK]} \
10\% & \text{随机替换} \
10\% & \text{保持原词}
\end{cases}$$

MLM 使用交叉熵损失：

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

其中 $M$ 是被掩码的 token 位置集合，$w_{\backslash i}$ 表示上下文信息。

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class OptimizedBertMLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.bert = BertModel(config)  # 原始 BERT 模型
        self.cls = nn.Linear(config.hidden_size, config.vocab_size)
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, input_ids, attention_mask, labels=None):
        # 使用梯度检查点减少显存占用
        outputs = checkpoint(
            self.bert,
            input_ids,
            attention_mask,
            use_reentrant=False
        )

        logits = self.cls(outputs.last_hidden_state)

        if labels is not None:
            # 只计算 mask 位置的 loss
            loss_mask = (labels != -100)
            logits = logits[loss_mask].view(-1, self.config.vocab_size)
            labels = labels[loss_mask].view(-1)
            loss = self.loss_fn(logits, labels)
            return loss
        return logits

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    loss = model(input_ids, attention_mask, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()