如何从BERT下一句预测模块提取合理的句间跟随概率？

问题描述 #

使用BertForNextSentencePrediction模块获取句子B跟随句子A的概率时，输出数值总是接近1或0，呈现布尔逻辑特性，不符合对连续概率的需求。例如：

• 句子对“我今天可能去商店。”和“你喜欢吐司抹黄油吗？”输出0.9999562501907349（预期概率应较低）
• 句子对“我今天可能去商店。”和“那个男人咳嗽了。”输出4.5933375076856464e-05（预期概率虽低但不应趋近于0）

当前使用的代码如下：

from transformers import BertTokenizer, BertForNextSentencePrediction
from torch.nn import functional as F

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
nsp_model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')

def next_sentence_prediction(sentence_1, sentence_2):
  encoding = tokenizer.encode_plus(sentence_1, sentence_2, return_tensors='pt')
  outputs = nsp_model(**encoding)[0]
  softmax = F.softmax(outputs, dim = 1)
  print(softmax[0][0].item())

原因分析 #

预训练阶段的BERT NSP任务是二分类任务：正例是文档中真实连续的句子对，负例是随机抽取的无关句子对。这种极端的训练数据分布导致模型学会输出极端概率，而非连续的关联度概率——它的目标是区分“是/不是下一句”，而非量化“有多可能是下一句”。

解决方案 #

1. 微调模型（最优方案） #

用包含不同关联程度句子对的自定义数据集微调模型，让模型学习连续的概率分布：

• 准备数据集：收集句子对，并标注它们的真实跟随概率（比如从0到1的连续值），或者按关联度划分多个等级（如“强关联”“中等关联”“弱关联”）。
• 修改模型输出层：如果是回归任务，将原模型的二分类输出层替换为单神经元的线性层（输出连续值）；如果是多分类任务，调整输出层神经元数量对应等级数。
• 训练模型：使用MSE损失（回归）或交叉熵损失（多分类）进行训练，让模型学会输出符合预期的连续概率。

示例微调框架（回归任务）：

import torch
import torch.nn as nn
from transformers import BertTokenizer, BertForNextSentencePrediction

# 加载预训练模型并替换输出层
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
nsp_model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')
# 替换原二分类层为回归层
nsp_model.classifier = nn.Linear(nsp_model.classifier.in_features, 1)
nsp_model.config.num_labels = 1

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.AdamW(nsp_model.parameters(), lr=5e-5)

# 示例训练循环（需替换为自定义数据集）
for epoch in range(3):
    nsp_model.train()
    # 假设train_dataloader是自定义数据集的加载器，包含sentence1, sentence2, label（0-1的概率）
    for batch in train_dataloader:
        encoding = tokenizer(batch['sentence1'], batch['sentence2'], return_tensors='pt', padding=True, truncation=True)
        labels = batch['label'].unsqueeze(1).float()
        
        outputs = nsp_model(**encoding)[0]
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

2. 用句子相似度模型替代 #

如果无法微调，可使用Sentence-BERT计算句子对的语义相似度，以此作为跟随概率的近似——语义越相似的句子，越有可能是连续的下一句：

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer('all-MiniLM-L6-v2')

def get_next_sentence_prob(sentence1, sentence2):
    emb1 = model.encode(sentence1, convert_to_tensor=True)
    emb2 = model.encode(sentence2, convert_to_tensor=True)
    # 计算余弦相似度，归一化到0-1区间作为概率近似
    cos_sim = util.cos_sim(emb1, emb2).item()
    return (cos_sim + 1) / 2  # 从[-1,1]映射到[0,1]

3. 温度缩放（临时调整方案） #

如果必须使用原预训练NSP模型，可通过温度缩放软化概率分布，但这只是人为调整，并未改变模型的认知，效果有限：

from transformers import BertTokenizer, BertForNextSentencePrediction
from torch.nn import functional as F

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
nsp_model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')

def next_sentence_prediction(sentence_1, sentence_2, temperature=5):
  encoding = tokenizer.encode_plus(sentence_1, sentence_2, return_tensors='pt')
  outputs = nsp_model(**encoding)[0]
  # 除以温度系数后再做softmax，温度越高概率分布越平缓
  softmax = F.softmax(outputs / temperature, dim = 1)
  print(softmax[0][0].item())

内容的提问来源于stack exchange，提问作者TheGoldenPatrik1