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

1. BERT模型的核心预训练任务解析

BERT（Bidirectional Encoder Representations from Transformers）作为自然语言处理领域的里程碑模型，其核心创新在于通过Masked Language Model（MLM）和Next Sentence Prediction（NSP）两个预训练任务，让模型学习到深层次的上下文语义表示。这两个任务的设计理念非常巧妙：MLM让模型学会理解单词在上下文中的含义，而NSP则让模型掌握句子间的逻辑关系。

在实际项目中，我发现很多开发者对这两个任务的具体实现存在困惑。比如MLM任务中15%的mask比例如何选择？NSP任务中正负样本如何构建？这些细节直接影响模型最终效果。下面我将结合PyTorch代码，带大家从零实现这两个关键任务。

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

2.1 安装必要依赖

首先需要安装PyTorch和相关工具库。建议使用Python 3.8+环境：

pip install torch==1.12.1 numpy==1.21.6

2.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%替换为随机token，10%保持不变）。这种设计让模型不能简单地依赖[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_index=0) # 忽略padding位置 optimizer = optim.Adam(model.parameters(), lr=0.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_model=512，n_layers=4通常是不错的起点。

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, dtype=torch.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

这里有几个工程实现要点：

1. 使用残差连接(residual connection)缓解梯度消失
2. Layer Normalization加速收敛
3. 注意力掩码(attn_mask)处理变长输入

7. 常见问题与解决方案

7.1 训练不收敛问题

如果遇到训练loss波动大或不收敛，可以尝试：

1. 减小学习率（如从0.001降到0.0001）
2. 增加warmup步数（前1000步线性增大学习率）
3. 检查梯度裁剪（gradient clipping）是否合理

7.2 过拟合处理

在小数据集上训练BERT容易过拟合，解决方法包括：

1. 增加Dropout（推荐0.1-0.3）
2. 使用更小的模型尺寸
3. 提前停止（early stopping）
4. 数据增强（如同义词替换）

7.3 显存不足问题

当遇到CUDA out of memory错误时：

1. 减小batch_size
2. 使用梯度累积（gradient accumulation）
3. 尝试混合精度训练（AMP）

我在实际部署中发现，使用梯度累积配合AMP能在几乎不损失精度的情况下，将显存占用降低40%以上。