Transformers 掩码语言建模（MLM）实战：用 DistilRoBERTa 微调 ELI5 数据集并实现 Fill-Mask 推理

2026-09-04 17:10:34作者：薛曦旖Francesca

🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.

项目地址：https://gitcode.com/GitHub_Trending/tra/transformers

本篇技术指南基于 Transformers 官方文档掩码语言建模任务文档展开，完整覆盖掩码语言建模（Masked Language Modeling，MLM）的核心流程：加载并预处理 ELI5 数据集、使用 DataCollatorForLanguageModeling 实现动态掩码、通过 Trainer 微调 DistilRoBERTa、以困惑度（Perplexity）评估模型，并最终用 fill-mask pipeline 或原始 logits 完成推理。读完本文，你可以独立复现一篇从数据到上线推理的完整 MLM 微调链路，并理解 mlm_probability、动态填充（dynamic padding）、-100 标签掩码等关键参数在源码层面的真实行为。

什么是掩码语言建模

掩码语言建模让模型预测序列中某个被遮蔽的 token。与普通因果语言建模（从左到右单向预测）不同，MLM 模型可以对上下文做双向注意力——模型能够同时访问被遮蔽位置左侧和右侧的所有 token。这使它特别适合那些需要充分理解完整句/段上下文的任务，BERT 就是典型的掩码语言模型。

本指南将以 DistilRoBERTa（模型 ID distilbert/distilroberta-base）为底座，在 ELI5 数据集的 r/askscience 子集上完成微调，并演示两种推理方式：

使用 [~transformers.Trainer] 微调 DistilRoBERTa；
使用训练好的模型进行推理（fill-mask pipeline 与底层 logits 两种方式）。

开始之前，先安装必要依赖：

pip install transformers datasets evaluate

建议登录 Hugging Face 账号，以便加载并分享模型：

>>> from huggingface_hub import notebook_login

>>> notebook_login()

第一步：加载 ELI5 数据集

先用 Datasets 库加载 ELI5-Category 数据集的前 5000 条训练样本。只取子集是为了先验证整条链路跑通，再投入完整数据集：

>>> from datasets import load_dataset

>>> eli5 = load_dataset("eli5_category", split="train[:5000]")

接着把 train 集按 8:2 拆分为训练集与测试集：

>>> eli5 = eli5.train_test_split(test_size=0.2)

查看一条样本，可以看到其完整结构：

>>> eli5["train"][0]
{'q_id': '7h191n',
 'title': 'What does the tax bill that was passed today mean? How will it affect Americans in each tax bracket?',
 'selftext': '',
 'category': 'Economics',
 'subreddit': 'explainlikeimfive',
 'answers': {'a_id': ['dqnds8l', 'dqnd1jl', 'dqng3i1', 'dqnku5x'],
  'text': ["The tax bill is 500 pages long and there were a lot of changes still going on right to the end. It's not just an adjustment to the income tax brackets, it's a whole bunch of changes. As such there is no good answer to your question. The big take aways are: - Big reduction in corporate income tax rate will make large companies very happy. - Pass through rate change will make certain styles of business (law firms, hedge funds) extremely happy - Income tax changes are moderate, and are set to expire (though it's the kind of thing that might just always get re-applied without being made permanent) - People in high tax states (California, New York) lose out, and many of them will end up with their taxes raised.",
   'None yet. It has to be reconciled with a vastly different house bill and then passed again.',
   'Also: does this apply to 2017 taxes? Or does it start with 2018 taxes?',
   'This article explains both the House and senate bills, including the proposed changes to your income taxes based on your income level. URL_0'],
  'score': [21, 19, 5, 3],
  'text_urls': [[], [], [],
   ['https://www.investopedia.com/news/trumps-tax-reform-what-can-be-done/']],
 'title_urls': ['url'],
 'selftext_urls': ['url']}

字段虽然多，但真正有用的是文本内容。语言建模任务的一个特点是无需人工标注（属于无监督任务）——"下一个词"本身就是标签。

第二步：数据预处理（Preprocess）

加载 DistilRoBERTa 分词器

MLM 的预处理入口是加载 DistilRoBERTa 的分词器，用于处理目标文本字段：

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilroberta-base")

展平嵌套结构

从上面的样本可以看出，文本位于 answers 这一嵌套字典内部。需要用 Datasets 的 flatten() 把嵌套结构展平：

>>> eli5 = eli5.flatten()
>>> eli5["train"][0]
{'q_id': '7h191n',
 'title': 'What does the tax bill that was passed today mean? How will it affect Americans in each tax bracket?',
 'selftext': '',
 'category': 'Economics',
 'subreddit': 'explainlikeimfive',
 'answers.a_id': ['dqnds8l', 'dqnd1jl', 'dqng3i1', 'dqnku5x'],
 'answers.text': ["The tax bill is 500 pages long and there were a lot of changes still going on right to the end. ...",
  'None yet. It has to be reconciled with a vastly different house bill and then passed again.',
  'Also: does this apply to 2017 taxes? Or does it start with 2018 taxes?',
  'This article explains both the House and senate bills, including the proposed changes to your income taxes based on your income level. URL_0'],
 'answers.score': [21, 19, 5, 3],
 'answers.text_urls': [[], [], [],
  ['https://www.investopedia.com/news/trumps-tax-reform-what-can-be-done/']],
 'title_urls': ['url'],
 'selftext_urls': ['url']}

展平后，每个子字段都变成了带 answers 前缀的独立列，text 字段此时是一列列表（每个问题对应多条回答）。我们不逐句处理，而是把每条样本的回答列表拼成一个字符串再统一分词：

>>> def preprocess_function(examples):
...     return tokenizer([" ".join(x) for x in examples["answers.text"]])

用 [~datasets.Dataset.map] 批量应用。batched=True 让一次处理多条样本，num_proc 控制进程数，remove_columns 删掉所有不再需要的列：

>>> tokenized_eli5 = eli5.map(
...     preprocess_function,
...     batched=True,
...     num_proc=4,
...     remove_columns=eli5["train"].column_names,
... )

按 block_size 分块（group_texts）

此时数据集已是 token ID 序列，但部分序列会超过模型的最大输入长度。用第二个预处理函数完成两件事：

把整个 batch 内所有序列首尾拼接起来；
按 block_size 切成定长小段。block_size 必须小于模型最大输入长度，并且要能放进 GPU 显存。

>>> block_size = 128

>>> def group_texts(examples):
...     # 拼接所有文本。
...     concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
...     total_length = len(concatenated_examples[list(examples.keys())[0]])
...     # 丢弃尾部不足一个 block 的零头；若模型支持 padding，也可以改为补齐而不是丢弃，
...     # 这部分可按需自定义。
...     if total_length >= block_size:
...         total_length = (total_length // block_size) * block_size
...     # 切成 block_size 长度的块。
...     result = {
...         k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
...         for k, t in concatenated_examples.items()
...     }
...     return result

然后应用到全量数据：

>>> lm_dataset = tokenized_eli5.map(group_texts, batched=True, num_proc=4)

官方训练脚本 run_mlm.py 中的 group_texts 与此逻辑完全一致：拼接、total_length = (total_length // max_seq_length) * max_seq_length 截零、再按定长切片；且当 map 以 batched=True 运行时，每次处理 1000 条文本，每 1000 条会产生一次零头丢弃（见脚本注释）。

动态填充与随机掩码：DataCollatorForLanguageModeling

效率更高的做法是动态填充：组 batch 时才把每条序列填充到"本 batch 内最长句"，而不是把整个数据集填充到最大长度。为此把序列结束符设为 pad 符，并通过 mlm_probability 控制每个 epoch 随机遮蔽的比例：

>>> from transformers import DataCollatorForLanguageModeling

>>> tokenizer.pad_token = tokenizer.eos_token
>>> data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)

源码层面，DataCollatorForLanguageModeling 的完整参数与默认值如下，这里结合文档只暴露了 mlm_probability，其余参数同样值得关注：

参数	默认值	说明
`tokenizer`	必填	用于编码数据的分词器
`mlm`	`True`	是否启用掩码语言建模。为 `False` 时退化为因果 LM：标签与输入相同，pad 位置置 `-100`
`whole_word_mask`	`False`	是否整词遮蔽（依赖 fast tokenizer 的 offset mapping）
`mlm_probability`	`0.15`	每个 token 被随机掩码的概率，必须介于 0 和 1
`mask_replace_prob`	`0.8`	被选中掩码的 token 中，替换为 `[MASK]` 的比例（BERT 风格：80% 换 `[MASK]`、10% 换随机 token、10% 保持原样）
`random_replace_prob`	`0.1`	替换为词表中随机 token 的比例；两项之和不得超过 1
`pad_to_multiple_of`	`None`	若设置，则把序列填充到该值的整数倍
`return_tensors`	`"pt"`	返回张量类型，可选 `"pt"` / `"np"`
`seed`	`None`	掩码随机数种子；多进程 DataLoader 下会为每个 worker 派生独立种子

几个关键的实现细节（见 src/transformers/data/data_collator.py）：

__post_init__ 强校验：mlm=True 时若分词器没有 mask_token 会直接抛错，并提示"应传 mlm=False 改为因果语言建模"；mlm_probability 越界、概率之和超过 1 也会抛错。
torch_mask_tokens 生成标签：非掩码位置的标签统一置为 -100（loss 计算时忽略），只有被掩码位置保留待预测的 token ID，这就是 mlm_probability=0.15 在训练中的实际作用点。
special_tokens_mask 加速：若分词时传了 return_special_tokens_mask=True，collator 会直接复用该列，避免每次 batch 都重新计算特殊 token 掩码。官方脚本正是因此默认加上该参数。
整词遮蔽的约束：whole_word_mask=True 要求 fast tokenizer（依赖 offset mapping），且会强制 mask_replace_prob=1（不支持随机 token 替换）。
填充：collator 在 torch_call 中调用 tokenizer.pad 做 batch 内动态 padding，并支持 pad_to_multiple_of；官方脚本在 --line_by_line 且 fp16 时会把它设为 8（见 run_mlm.py）。

第三步：训练（Train）

如果对用 Trainer 微调模型的流程还不熟悉，可先参考官方 PyTorch Trainer 训练指南 docs/source/ar/training.md。

用 [~transformers.AutoModelForMaskedLM] 加载 DistilRoBERTa：

>>> from transformers import AutoModelForMaskedLM

>>> model = AutoModelForMaskedLM.from_pretrained("distilbert/distilroberta-base")

剩下三步：

在 [~transformers.TrainingArguments] 中配置训练参数。唯一必填项是 output_dir（模型保存位置）；设置 push_to_hub=True 可训练完成后直接推送到 Hub（前提：已登录 Hugging Face）；
将模型、数据集、data collator 一并传入 [~transformers.Trainer]；
调用 [~transformers.Trainer.train] 开始微调。

>>> training_args = TrainingArguments(
...     output_dir="my_awesome_eli5_mlm_model",
...     eval_strategy="epoch",
...     learning_rate=2e-5,
...     num_train_epochs=3,
...     weight_decay=0.01,
...     push_to_hub=True,
... )

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     train_dataset=lm_dataset["train"],
...     eval_dataset=lm_dataset["test"],
...     data_collator=data_collator,
...     tokenizer=tokenizer,
... )

>>> trainer.train()

参数要点：

output_dir="my_awesome_eli5_mlm_model"：检查点与最终模型的落盘目录；
eval_strategy="epoch"：每个 epoch 结束后在测试集上评估一次；
learning_rate=2e-5：微调（而非从头训练）常用的量级；
num_train_epochs=3：5000 条样本的轻量试验规模下 3 个 epoch 即可；
weight_decay=0.01：对权重施加 L2 正则；
push_to_hub=True：训练结束后自动推送模型。

评估与分享

训练结束后，用 [~transformers.Trainer.evaluate] 评估模型并计算困惑度（Perplexity）：

>>> import math

>>> eval_results = trainer.evaluate()
>>> print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}")
Perplexity: 8.76

困惑度即 exp(eval_loss)：eval_loss 是掩码位置上的平均对数损失，取指数后得到"模型在预测被遮蔽 token 时平均面对多少个候选词"。文档给出的 8.76 是其在 ELI5 子集上的示例输出，可用作复现时的参考量级（实际数值会随数据与随机掩码略有浮动）。

最后把模型推送到 Hub 供他人使用：

>>> trainer.push_to_hub()

推理：两种 fill-mask 用法

模型微调完成后即可用于推理。

方式一：fill-mask pipeline（最简）

准备一段含 <mask> 占位的文本，用 [~transformers.pipeline] 创建 fill-mask 管线，top_k 控制返回的候选数量：

>>> from transformers import pipeline

>>> mask_filler = pipeline("fill-mask", "username/my_awesome_eli5_mlm_model")
>>> mask_filler(text, top_k=3)
# text = "The Milky Way is a <mask> galaxy."
[{'score': 0.5150994658470154,
  'token': 21300,
  'token_str': ' spiral',
  'sequence': 'The Milky Way is a spiral galaxy.'},
 {'score': 0.07087188959121704,
  'token': 2232,
  'token_str': ' massive',
  'sequence': 'The Milky Way is a massive galaxy.'},
 {'score': 0.06434620916843414,
  'token': 650,
  'token_str': ' small',
  'sequence': 'The Milky Way is a small galaxy.'}]

从源码看，FillMaskPipeline 的后处理逻辑正是：取掩码位置概率的 topk(top_k)，逐个用候选 token 替换原始 input_ids 中 <mask> 的位置、过滤 pad token 后解码出 sequence，并输出 score（概率）、token（token ID）、token_str（解码文本）四个键。若分词器没有定义 mask_token，pipeline 会直接抛出 PipelineException。它还支持 targets 参数——只在一组候选词中比较概率（词表外的词会被分词取首 token 并告警）。该行为的测试覆盖见 tests/pipelines/test_pipelines_fill_mask.py。

方式二：手动分词 + logits（底层）

如果想拿到更细粒度的控制，可以手动完成整条链路。

第 1 步：分词并拿到 input_ids，同时定位 <mask> 的位置：

>>> from transformers import AutoTokenizer
import torch

>>> tokenizer = AutoTokenizer.from_pretrained("username/my_awesome_eli5_mlm_model")
>>> inputs = tokenizer(text, return_tensors="pt")
>>> mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]

第 2 步：把输入喂给模型，取出掩码位置的 logits：

>>> from transformers import AutoModelForMaskedLM

>>> model = AutoModelForMaskedLM.from_pretrained("username/my_awesome_eli5_mlm_model")
>>> logits = model(**inputs).logits
>>> mask_token_logits = logits[0, mask_token_index, :]

第 3 步：对掩码位置 logits 做 top-k，解码并替换原文本中的 <mask>：

>>> top_3_tokens = torch.topk(mask_token_logits, 3, dim=1).indices[0].tolist()

>>> for token in top_3_tokens:
...     print(text.replace(tokenizer.mask_token, tokenizer.decode([token])))
The Milky Way is a spiral galaxy.
The Milky Way is a massive galaxy.
The Milky Way is a small galaxy.

两种方式输出的候选完全一致——pipeline 只是把"分词 → 定位 mask → 取 logits → top-k → 替换解码"这条链路封装成了三行代码。

小结

MLM 是双向注意力的无监督预训练范式：标签来自"被遮蔽的下一个词"本身，无需人工标注；
数据侧的关键技巧是 flatten() → 拼接回答文本 → group_texts 按 block_size=128 分块，既消除了超长序列，又把零散短文本聚合为定长训练块；
DataCollatorForLanguageModeling 同时解决动态填充与随机掩码两件事：mlm_probability=0.15 决定掩码密度，-100 标签保证 loss 只算在掩码位置上，mask_replace_prob=0.8 / random_replace_prob=0.1 复现 BERT 的 80/10/10 替换策略；
训练侧只需在 TrainingArguments 里补上 output_dir，即可用 Trainer 完成微调、按 epoch 评估、以 exp(eval_loss) 换算困惑度，并一键 push_to_hub；
推理侧，fill-mask pipeline 是快速验证的首选，手动 logits 路径则适合需要嵌入更大系统或做多掩码/受限候选（targets）等定制场景。