使用确定性掩码替代随机掩码微调BERT的实现方案

核心思路 #

既然已经提前处理好带[MASK]的文本和对应的完整标签文本，我们只需要对齐分词后的token位置，把非掩码位置的标签设为-100（PyTorch会自动忽略该值的损失计算），再用自定义的数据整理器替代随机掩码的默认工具即可。

1. 数据集预处理 #

先对text（带掩码）和label（完整文本）分别做分词，确保两者的token序列长度一致，然后生成模型需要的input_ids、attention_mask和labels：

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

def preprocess_function(examples):
    # 对带掩码的文本分词
    tokenized_text = tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=128,
        return_tensors="pt"
    )
    # 对完整标签文本分词
    tokenized_label = tokenizer(
        examples["label"],
        truncation=True,
        padding="max_length",
        max_length=128,
        return_tensors="pt"
    )
    
    # 生成labels：只有[MASK]对应的位置保留真实token，其余设为-100
    labels = tokenized_label["input_ids"].clone()
    mask_positions = (tokenized_text["input_ids"] == tokenizer.mask_token_id)
    labels[~mask_positions] = -100
    
    return {
        "input_ids": tokenized_text["input_ids"],
        "attention_mask": tokenized_text["attention_mask"],
        "labels": labels
    }

# 假设你的数据集是Hugging Face Dataset格式
processed_dataset = raw_dataset.map(preprocess_function, batched=True)

2. 自定义数据整理器 #

不需要用默认的DataCollatorForLanguageModeling，直接用基础的padding整理器即可，因为预处理已经完成所有掩码和标签设置：

from transformers import DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

3. 启动训练 #

用Trainer组件快速启动训练，也可以自己写PyTorch训练循环，这里以Trainer为例：

from transformers import BertForMaskedLM, TrainingArguments, Trainer

model = BertForMaskedLM.from_pretrained("bert-base-uncased")

training_args = TrainingArguments(
    output_dir="./custom_mask_bert",
    per_device_train_batch_size=8,
    num_train_epochs=3,
    logging_dir="./logs",
    logging_steps=10,
    save_steps=100,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=processed_dataset["train"],
    data_collator=data_collator,
)

trainer.train()

关键说明 #

• 预处理时必须保证text和label分词后的序列长度一致，否则掩码位置会错位，你已经提前处理过多token分词的情况，这一步没问题。
• 设置labels为-100的目的是让模型只计算掩码位置的预测损失，非掩码位置不参与损失计算，和默认MLM的逻辑一致，只是掩码位置是固定的。

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