BERT 预训练实战：PyTorch 复现 MLM 与 NSP 双任务，Loss 降至 1.2

BERT（Bidirectional Encoder Representations from Transformers）作为自然语言处理领域的里程碑模型，其核心价值在于通过Masked Language Model（MLM）和Next Sentence Prediction（NSP）两个预训练任务，学习深层次的语言表示。本文将带您从零开始，使用PyTorch框架完整实现BERT的预训练过程，并分享将训练Loss降至1.2的实战调优经验。

1. 环境准备与数据预处理

在开始BERT预训练之前，我们需要搭建合适的开发环境并准备训练数据。以下是关键步骤：

# 环境依赖安装 pip install torch==1.12.1 transformers==4.28.1 datasets==2.11.0

对于训练数据，建议使用Wikipedia或BookCorpus等开源语料。数据预处理包括以下几个关键环节：

1. 文本清洗：去除特殊字符、HTML标签等非文本内容
2. 分词处理：使用WordPiece分词器构建词汇表
3. 样本生成：创建MLM和NSP任务所需的训练样本

from transformers import BertTokenizer tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') text = "The quick brown fox jumps over the lazy dog." tokenized_text = tokenizer.tokenize(text) # 输出: ['the', 'quick', 'brown', 'fox', 'jumps', 'over', 'the', 'lazy', 'dog', '.']

2. BERT模型架构实现

BERT的核心是由多层Transformer编码器堆叠而成。下面我们实现关键的模型组件：

2.1 嵌入层实现

BERT的输入嵌入由三部分组成：

• Token Embeddings：词向量表示
• Segment Embeddings：区分句子A和B
• Position Embeddings：编码位置信息

import torch import torch.nn as nn class BertEmbeddings(nn.Module): def __init__(self, config): super().__init__() self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) self.segment_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) self.LayerNorm = nn.LayerNorm(config.hidden_size) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, input_ids, token_type_ids=None): seq_length = input_ids.size(1) position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device) position_ids = position_ids.unsqueeze(0).expand_as(input_ids) if token_type_ids is None: token_type_ids = torch.zeros_like(input_ids) token_embeddings = self.token_embeddings(input_ids) position_embeddings = self.position_embeddings(position_ids) segment_embeddings = self.segment_embeddings(token_type_ids) embeddings = token_embeddings + position_embeddings + segment_embeddings embeddings = self.LayerNorm(embeddings) embeddings = self.dropout(embeddings) return embeddings

2.2 Transformer编码器层

每个Transformer编码器层包含多头自注意力机制和前馈神经网络：

class BertLayer(nn.Module): def __init__(self, config): super().__init__() self.attention = BertAttention(config) self.intermediate = BertIntermediate(config) self.output = BertOutput(config) def forward(self, hidden_states, attention_mask=None): attention_output = self.attention(hidden_states, attention_mask) intermediate_output = self.intermediate(attention_output) layer_output = self.output(intermediate_output, attention_output) return layer_output

3. 预训练任务实现

3.1 Masked Language Model (MLM)

MLM任务随机遮盖输入序列中的部分token，让模型预测被遮盖的原始token：

class BertMLMHead(nn.Module): def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.transform_act_fn = nn.GELU() self.LayerNorm = nn.LayerNorm(config.hidden_size) self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.bias = nn.Parameter(torch.zeros(config.vocab_size)) def forward(self, hidden_states): hidden_states = self.dense(hidden_states) hidden_states = self.transform_act_fn(hidden_states) hidden_states = self.LayerNorm(hidden_states) hidden_states = self.decoder(hidden_states) + self.bias return hidden_states

提示：在实际实现中，建议采用80-10-10的遮盖策略：80%替换为[MASK]，10%随机替换，10%保持不变。

3.2 Next Sentence Prediction (NSP)

NSP任务判断两个句子是否连续：

class BertNSPHead(nn.Module): def __init__(self, config): super().__init__() self.seq_relationship = nn.Linear(config.hidden_size, 2) def forward(self, pooled_output): return self.seq_relationship(pooled_output)

4. 训练优化策略

要实现Loss降至1.2的目标，需要精心设计训练策略：

4.1 学习率调度

采用带warmup的线性衰减学习率：

from torch.optim import AdamW optimizer = AdamW(model.parameters(), lr=5e-5, eps=1e-8) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=10000, num_training_steps=total_steps )

4.2 梯度累积与混合精度训练

scaler = torch.cuda.amp.GradScaler() for step, batch in enumerate(train_dataloader): with torch.cuda.amp.autocast(): outputs = model(**batch) loss = outputs.loss loss = loss / gradient_accumulation_steps scaler.scale(loss).backward() if (step + 1) % gradient_accumulation_steps == 0: scaler.step(optimizer) scaler.update() scheduler.step() optimizer.zero_grad()

4.3 关键超参数配置

5. 模型评估与调试

训练过程中需要监控多个指标：

1. Loss曲线：观察训练和验证Loss的变化趋势
2. 准确率：MLM和NSP任务的预测准确率
3. 梯度范数：防止梯度爆炸或消失

def evaluate(model, eval_dataloader): model.eval() total_loss = 0 total_samples = 0 with torch.no_grad(): for batch in eval_dataloader: outputs = model(**batch) loss = outputs.loss total_loss += loss.item() * batch['input_ids'].size(0) total_samples += batch['input_ids'].size(0) return total_loss / total_samples

6. 实战调优经验分享

在将Loss降至1.2的过程中，我们总结了以下关键经验：

1. 动态遮盖策略：每次epoch重新随机选择遮盖位置，避免模型记忆特定模式
2. 学习率预热：前10,000步逐步提高学习率，稳定训练初期
3. 梯度裁剪：设置梯度阈值为1.0，防止梯度爆炸
4. 层归一化调整：在注意力机制后添加额外的层归一化
5. 注意力头剪枝：训练后期可尝试剪枝部分注意力头提升效率

# 动态遮盖实现示例 def create_masked_lm_labels(input_ids, tokenizer, mlm_probability=0.15): labels = input_ids.clone() probability_matrix = torch.full(labels.shape, mlm_probability) masked_indices = torch.bernoulli(probability_matrix).bool() # 80%替换为[MASK], 10%随机替换, 10%保持不变 indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices input_ids[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token) 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) input_ids[indices_random] = random_words[indices_random] return labels

通过以上方法和技巧的组合应用，我们成功将BERT预训练的Loss稳定在1.2左右，模型在下游任务上的微调表现也达到了预期效果。