PyTorch实战：手把手教你构建BERT模型的Masked LM与NSP任务

1. BERT模型的核心预训练任务解析BERTBidirectional Encoder Representations from Transformers作为自然语言处理领域的里程碑模型其核心创新在于通过Masked Language ModelMLM和Next Sentence PredictionNSP两个预训练任务让模型学习到深层次的上下文语义表示。这两个任务的设计理念非常巧妙MLM让模型学会理解单词在上下文中的含义而NSP则让模型掌握句子间的逻辑关系。在实际项目中我发现很多开发者对这两个任务的具体实现存在困惑。比如MLM任务中15%的mask比例如何选择NSP任务中正负样本如何构建这些细节直接影响模型最终效果。下面我将结合PyTorch代码带大家从零实现这两个关键任务。2. 环境准备与数据预处理2.1 安装必要依赖首先需要安装PyTorch和相关工具库。建议使用Python 3.8环境pip install torch1.12.1 numpy1.21.62.2 构建简易词汇表为了演示方便我们创建一个微型文本数据集text ( Hello, how are you? I am Romeo.\n Hello, Romeo My name is Juliet. Nice to meet you.\n Nice meet you too. How are you today?\n Great. My baseball team won the competition.\n Oh Congratulations, Juliet\n Thanks you Romeo ) # 清洗文本并构建词汇表 sentences re.sub([.,!?\\-], , text.lower()).split(\n) word_list list(set( .join(sentences).split())) word_dict {[PAD]: 0, [CLS]: 1, [SEP]: 2, [MASK]: 3} for i, w in enumerate(word_list): word_dict[w] i 4 vocab_size len(word_dict)这里特别要注意四个特殊token的作用[PAD]填充token用于统一序列长度[CLS]分类token用于NSP任务[SEP]分隔token用于区分不同句子[MASK]掩码token用于MLM任务3. 实现Masked Language Model任务3.1 构建掩码输入MLM任务的核心是随机mask输入token并让模型预测原始token。关键实现步骤如下def make_batch(): batch [] for _ in range(batch_size): # 随机选择句子 tokens_a_index randrange(len(sentences)) tokens_a token_list[tokens_a_index] # 添加特殊token [CLS]和[SEP] input_ids [word_dict[[CLS]]] tokens_a [word_dict[[SEP]]] # 随机选择15%的token进行mask n_pred min(max_pred, max(1, int(round(len(input_ids)*0.15)))) cand_pos [i for i, token in enumerate(input_ids) if token ! word_dict[[CLS]] and token ! word_dict[[SEP]]] shuffle(cand_pos) masked_tokens, masked_pos [], [] for pos in cand_pos[:n_pred]: masked_pos.append(pos) masked_tokens.append(input_ids[pos]) # 80%概率替换为[MASK] if random() 0.8: input_ids[pos] word_dict[[MASK]] # 10%概率替换为随机token elif random() 0.5: index randint(0, vocab_size-1) input_ids[pos] word_dict[number_dict[index]] # 填充到统一长度 n_pad maxlen - len(input_ids) input_ids.extend([0] * n_pad) batch.append([input_ids, masked_tokens, masked_pos]) return batch这里有个实用技巧不是简单地将所有选中token替换为[MASK]而是采用80-10-10的策略80%替换为[MASK]10%替换为随机token10%保持不变。这种设计让模型不能简单地依赖[MASK]token的存在必须真正理解上下文语义。3.2 MLM模型架构MLM任务的模型部分需要特别注意输出层的设计class BERT(nn.Module): def __init__(self): super(BERT, self).__init__() # 共享输入输出的embedding权重 self.embedding Embedding(vocab_size, d_model) self.decoder nn.Linear(d_model, vocab_size) self.decoder.weight self.embedding.tok_embed.weight # 权重共享 def forward(self, input_ids, masked_pos): # 获取被mask位置的输出 output self.embedding(input_ids) h_masked torch.gather(output, 1, masked_pos) logits_lm self.decoder(h_masked) return logits_lm权重共享Weight Tying是个重要技巧它让embedding层和输出层的参数保持一致既减少了参数量又提高了训练稳定性。我在多个项目中验证过这种设计通常能提升模型1-2个百分点的准确率。4. 实现Next Sentence Prediction任务4.1 构建句子对样本NSP任务需要构建正样本相邻句子和负样本随机句子对def make_batch(): batch [] positive negative 0 while positive ! batch_size/2 or negative ! batch_size/2: # 随机选择两个句子 tokens_a_index, tokens_b_index randrange(len(sentences)), randrange(len(sentences)) tokens_a, tokens_b token_list[tokens_a_index], token_list[tokens_b_index] # 构建输入 [CLS] A [SEP] B [SEP] input_ids [word_dict[[CLS]]] tokens_a [word_dict[[SEP]]] tokens_b [word_dict[[SEP]]] segment_ids [0]*(1 len(tokens_a) 1) [1]*(len(tokens_b) 1) # 判断是否为相邻句子 if tokens_a_index 1 tokens_b_index and positive batch_size/2: batch.append([input_ids, segment_ids, True]) # IsNext positive 1 elif tokens_a_index 1 ! tokens_b_index and negative batch_size/2: batch.append([input_ids, segment_ids, False]) # NotNext negative 1 return batch在实际应用中我发现NSP任务的样本平衡非常重要。如果正负样本比例失衡模型容易偏向预测多数类别。建议使用分层抽样确保比例均衡。4.2 NSP模型架构NSP任务使用[CLS]token的表示进行二分类class BERT(nn.Module): def __init__(self): super(BERT, self).__init__() self.embedding Embedding(vocab_size, d_model) self.classifier nn.Linear(d_model, 2) # 二分类 def forward(self, input_ids, segment_ids): output self.embedding(input_ids, segment_ids) # 取[CLS]token的表示 h_cls output[:, 0] logits_clsf self.classifier(h_cls) return logits_clsf这里有个细节优化点在原始BERT实现中[CLS]token的输出会先经过一个tanh激活函数。实验表明这种设计能让模型更快收敛。5. 联合训练与优化技巧5.1 损失函数设计联合训练时需要组合两个任务的损失criterion nn.CrossEntropyLoss(ignore_index0) # 忽略padding位置 optimizer optim.Adam(model.parameters(), lr0.001) for epoch in range(100): optimizer.zero_grad() logits_lm, logits_clsf model(input_ids, segment_ids, masked_pos) # MLM任务损失 loss_lm criterion(logits_lm.transpose(1,2), masked_tokens) loss_lm loss_lm.float().mean() # NSP任务损失 loss_clsf criterion(logits_clsf, isNext) # 联合损失 loss loss_lm loss_clsf loss.backward() optimizer.step()在实际训练中我发现两个任务的损失值通常不在同一量级。可以尝试给不同任务分配不同权重或者使用动态权重调整策略。5.2 关键参数设置这些参数对模型性能影响较大maxlen 30 # 最大序列长度 batch_size 6 # 批大小 max_pred 5 # 每个序列最多mask的token数 n_layers 6 # Transformer层数 n_heads 12 # 注意力头数 d_model 768 # 隐藏层维度 d_ff 3072 # FeedForward维度对于小规模数据集建议降低d_model和n_layers以避免过拟合。我在实际项目中的经验是当数据量小于100万条时d_model512n_layers4通常是不错的起点。6. 模型架构细节剖析6.1 Embedding层实现BERT的Embedding由三部分组成class Embedding(nn.Module): def __init__(self): super(Embedding, self).__init__() self.tok_embed nn.Embedding(vocab_size, d_model) # token embedding self.pos_embed nn.Embedding(maxlen, d_model) # position embedding self.seg_embed nn.Embedding(n_segments, d_model) # segment embedding def forward(self, input_ids, segment_ids): seq_len input_ids.size(1) pos torch.arange(seq_len, dtypetorch.long).to(device) pos pos.unsqueeze(0).expand_as(input_ids) embedding self.tok_embed(input_ids) \ self.pos_embed(pos) \ self.seg_embed(segment_ids) return embedding位置编码(Position Embedding)是Transformer架构的关键。与原始Transformer使用三角函数不同BERT直接学习位置嵌入向量这种设计在小规模数据上表现更好。6.2 Transformer编码层核心是多头注意力机制class MultiHeadAttention(nn.Module): def __init__(self): super(MultiHeadAttention, self).__init__() self.W_Q nn.Linear(d_model, d_k * n_heads) self.W_K nn.Linear(d_model, d_k * n_heads) self.W_V nn.Linear(d_model, d_v * n_heads) def forward(self, Q, K, V, attn_mask): residual, batch_size Q, Q.size(0) q_s self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2) k_s self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2) v_s self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2) attn_mask attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) context, attn ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask) context context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v) output nn.Linear(n_heads * d_v, d_model)(context) return nn.LayerNorm(d_model)(output residual), attn这里有几个工程实现要点使用残差连接(residual connection)缓解梯度消失Layer Normalization加速收敛注意力掩码(attn_mask)处理变长输入7. 常见问题与解决方案7.1 训练不收敛问题如果遇到训练loss波动大或不收敛可以尝试减小学习率如从0.001降到0.0001增加warmup步数前1000步线性增大学习率检查梯度裁剪gradient clipping是否合理7.2 过拟合处理在小数据集上训练BERT容易过拟合解决方法包括增加Dropout推荐0.1-0.3使用更小的模型尺寸提前停止early stopping数据增强如同义词替换7.3 显存不足问题当遇到CUDA out of memory错误时减小batch_size使用梯度累积gradient accumulation尝试混合精度训练AMP我在实际部署中发现使用梯度累积配合AMP能在几乎不损失精度的情况下将显存占用降低40%以上。