BERT掩码语言模型(MLM)目标函数优化实战：从理论到高效实现

1次阅读

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

掩码语言模型（Masked Language Model, MLM）是 BERT 预训练的核心任务之一，其目标是通过预测被随机掩码的 token 来学习上下文相关的词表示。然而，在实际应用中，MLM 目标函数的计算往往成为性能瓶颈，主要体现在以下方面：

计算复杂度高：对于每个掩码位置，需要计算整个词表的 softmax，当词表规模较大（如 BERT-base 的 30,522 词）时，这部分计算开销巨大。
内存占用大：存储所有位置的 logits 和梯度需要大量显存，尤其在处理长序列时更为明显。
罕见词处理困难：传统交叉熵损失对罕见词（低频词）的优化效果不佳，影响模型整体表现。

MLM 目标函数本质上是带掩码的交叉熵损失，其数学表示为：

$$
\mathcal{L}{MLM} = -\frac{1}{N} \sum)
$$}^N \sum_{j=1}^V y_{ij} \log(p_{ij

其中：
– $N$ 是掩码位置数量
– $V$ 是词表大小
– $y_{ij}$ 是 one-hot 标签
– $p_{ij}$ 是模型预测的 softmax 概率

在 BERT 中，只有 15% 的 token 会被随机掩码，其中：

80% 替换为[MASK]
10% 替换为随机词
10% 保持不变

原始实现中，对每个掩码位置独立计算 softmax 会导致大量冗余计算。通过矩阵运算并行化可以显著提升效率：

# 原始实现（低效）loss = 0
for i in masked_positions:
    logits = model(input_ids, attention_mask)[i]
    loss += F.cross_entropy(logits, labels[i])

# 优化实现（高效）all_logits = model(input_ids, attention_mask)  # [seq_len, vocab_size]
masked_logits = all_logits[masked_positions]   # [num_masked, vocab_size]
loss = F.cross_entropy(masked_logits, labels[masked_positions])

传统 softmax 需要计算整个词表的概率分布，而实际上大部分词的概率接近零。通过动态负采样可以显著减少计算量：

对每个 batch，统计词频分布
对高频词按概率采样，确保覆盖重要上下文
对低频词随机采样，保持模型泛化能力

实现代码片段：

def dynamic_negative_sampling(logits, labels, k=5000):
    # logits: [batch_size, seq_len, vocab_size]
    # labels: [batch_size, seq_len]
    batch_probs = torch.softmax(logits.detach(), dim=-1)
    neg_probs = 1 - batch_probs.gather(-1, labels.unsqueeze(-1)).squeeze()
    neg_samples = torch.multinomial(neg_probs, k, replacement=True)
    return logits.gather(-1, neg_samples), labels.gather(-1, neg_samples)

利用 PyTorch 的 AMP（Automatic Mixed Precision）模块可以大幅减少显存占用并提升计算速度：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    logits = model(input_ids, attention_mask)
    loss = F.cross_entropy(logits[masked_positions], labels[masked_positions])

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

import torch
import torch.nn.functional as F
from transformers import BertForMaskedLM

class OptimizedBertMLM(torch.nn.Module):
    def __init__(self, model_name='bert-base-uncased'):
        super().__init__()
        self.bert = BertForMaskedLM.from_pretrained(model_name)

    def forward(self, input_ids, attention_mask, labels, use_amp=True):
        with torch.cuda.amp.autocast(enabled=use_amp):
            outputs = self.bert(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels,
                output_hidden_states=True
            )

            # 获取掩码位置的 logits
            masked_pos = (input_ids == self.bert.config.mask_token_id).nonzero(as_tuple=True)
            masked_logits = outputs.logits[masked_pos]
            masked_labels = labels[masked_pos]

            # 动态负采样
            if self.training:
                sampled_logits, sampled_labels = dynamic_negative_sampling(masked_logits.unsqueeze(0), 
                    masked_labels.unsqueeze(0)
                )
                loss = F.cross_entropy(sampled_logits.squeeze(0), sampled_labels.squeeze(0))
            else:
                loss = outputs.loss

        return {"loss": loss, "logits": outputs.logits}

使用梯度检查点（Gradient Checkpointing）：

model = BertForMaskedLM.from_pretrained(
    'bert-base-uncased', 
    gradient_checkpointing=True
)

序列分块处理：对于长文本，将序列分割为多个子序列分别处理

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化进程组
dist.init_process_group("nccl")
model = DDP(model.to(device), device_ids=[local_rank])

# 使用 ZeRO 优化器（需安装 deepspeed）# 在配置文件中添加："optimizer": {
    "type": "AdamW",
    "params": {
        "lr": 5e-5,
        "weight_decay": 0.01
    }
},
"zero_optimization": {
    "stage": 2,
    "offload_optimizer": {"device": "cpu"}
}

子词 tokenizer 优化：调整 BPE 算法的合并次数
损失函数加权：

class WeightedCrossEntropy(torch.nn.Module):
    def __init__(self, vocab_weights):
        super().__init__()
        self.weights = torch.tensor(vocab_weights).cuda()

    def forward(self, logits, labels):
        log_probs = F.log_softmax(logits, dim=-1)
        nll_loss = -log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
        weighted_loss = nll_loss * self.weights[labels]
        return weighted_loss.mean()

在 Wikipedia 数据集上的对比实验（batch_size=32，序列长度 =512）：