Hugging Face Transformers 中的 ERNIE 系列模型：架构原理、中文任务应用与源码级剖析

ERNIE（Enhanced Representation through kNowledge IntEgration）是百度提出的一系列面向自然语言理解任务的预训练模型，尤其在中文字词级任务上表现突出。本文以 ERNIE 官方模型文档 为主体，结合本仓库 transformers 的 ERNIE 模型源码 与 测试用例，完整讲解 ERNIE 的模型变体、配置参数、中文 Masked LM 实战用法，以及各下游任务头的底层实现细节，帮助你直接用 Pipeline、AutoModel 或各类 ErnieForXxx 模型完成中文文本理解任务。

ERNIE 模型家族概览与设计动机

ERNIE 一词由百度提出，核心思想是通过知识增强的掩码策略学习语言表示。与仅基于单字/子词随机掩码的预训练范式不同，ERNIE 引入了两种更高层的掩码单元：

• 实体级掩码（entity-level masking）：将句子中的命名实体（人名、地名、机构名等）整体掩码，迫使模型学习实体层面的语义知识；
• 短语级掩码（phrase-level masking）：将完整的短语作为一个掩码单元，帮助模型捕捉短语级的语义与搭配关系。

该模型系列先后经历了 ERNIE 1.0、ERNIE 2.0、ERNIE 3.0、ERNIE-Gram 与 ERNIE-health 等多代演进，属于百度提出的系列化预训练模型，在中文字词类任务中具有代表性。文档声明其由社区贡献者 nghuyong 移植到本仓库，官方代码（PaddlePaddle 框架）位于百度 PaddleNLP 项目，而本文所介绍的是 PyTorch 版本在 🤗 Transformers 中的实现。

说明：百度后来发布的 ERNIE 4.5 系列（含 MoE 变体）属于独立的模型家族与架构，不在本文范围，相关文档可参考 Ernie 4.5 与 Ernie 4.5 MoE。

可用模型变体一览

文档给出的模型检查点覆盖中文、英文两种语言以及多种尺寸规格，其中中文变体覆盖了从 nano 到 base 的多档规模，便于在不同算力预算下选择：

模型名称	语言	结构描述（层数 / 注意力头 / 隐层维度）
ernie-1.0-base-zh	中文	Layer: 12, Heads: 12, Hidden: 768
ernie-2.0-base-en	英文	Layer: 12, Heads: 12, Hidden: 768
ernie-2.0-large-en	英文	Layer: 24, Heads: 16, Hidden: 1024
ernie-3.0-base-zh	中文	Layer: 12, Heads: 12, Hidden: 768
ernie-3.0-medium-zh	中文	Layer: 6, Heads: 12, Hidden: 768
ernie-3.0-mini-zh	中文	Layer: 6, Heads: 12, Hidden: 384
ernie-3.0-micro-zh	中文	Layer: 4, Heads: 12, Hidden: 384
ernie-3.0-nano-zh	中文	Layer: 4, Heads: 12, Hidden: 312
ernie-health-zh	中文	Layer: 12, Heads: 12, Hidden: 768
ernie-gram-zh	中文	Layer: 12, Heads: 12, Hidden: 768

这些检查点由社区作者通过 Model Hub 分发（账号 nghuyong），模型细节以百度 PaddleNLP 官方文档为准。本仓库中 ERNIE 的整体训练目标仍然是 BERT 式的 Masked Language Modeling + Next Sentence Prediction（后者对应 ERNIE 的句子级预训练任务），同时保留 ERNIE 2.0/3.0 中可选的 task embedding 输入（详见下文"任务嵌入"部分）。

快速上手：用 [MASK] 补全实战中文预测

ERNIE 最常见的入门用例是预测被掩码的中文 token，文档给出了 Pipeline 与 AutoModel 两种等价的调用方式，此外任何 ErnieForMaskedLM 类模型都会自动注册进 fill-mask pipeline。

方式一：通过 Pipeline 一行完成

from transformers import pipeline


pipeline = pipeline(
    task="fill-mask",
    model="nghuyong/ernie-3.0-xbase-zh"
)

pipeline("巴黎是[MASK]国的首都。")

Pipeline 内部会自动完成分词、模型前向与 [MASK] 位置的概率解码，直接返回候选 token 及置信度，适合快速验证模型效果。

方式二：通过 AutoModel + AutoTokenizer 手动控制

如果需要精确掌控前向过程（例如批量推理、接自定义后处理或提取 logits），可以使用 AutoModelForMaskedLM：

import torch

from transformers import AutoModelForMaskedLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained(
    "nghuyong/ernie-3.0-xbase-zh",
)
model = AutoModelForMaskedLM.from_pretrained(
    "nghuyong/ernie-3.0-xbase-zh",
    device_map="auto"
)
inputs = tokenizer("巴黎是[MASK]国的首都。", return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(**inputs)
    predictions = outputs.logits

masked_index = torch.where(inputs['input_ids'] == tokenizer.mask_token_id)[1]
predicted_token_id = predictions[0, masked_index].argmax(dim=-1)
predicted_token = tokenizer.decode(predicted_token_id)

print(f"The predicted token is: {predicted_token}")

代码要点：AutoTokenizer 与 AutoModelForMaskedLM 会根据 model_type = "ernie" 自动解析到 ERNIE 对应的 tokenizer 与模型类；device_map="auto" 让模型自动落到可用设备；torch.where 定位序列中 [MASK] token 的位置，再从词表维度取 argmax 即可还原预测词（对示例句子应输出"法"）。也可用下述方式本地初始化一个随机权重的 ERNIE 模型用于实验（配置类 docstring 中的官方示例）：

from transformers import ErnieConfig, ErnieModel

# 初始化为 nghuyong/ernie-3.0-base-zh 风格配置
configuration = ErnieConfig()

# 以随机权重初始化模型
model = ErnieModel(configuration)

# 读取模型配置
configuration = model.config

源码架构剖析：基于 BERT 模块化的 ERNIE 实现

查看本仓库实现目录 src/transformers/models/ernie，共有四个文件：__init__.py（惰性加载导出）、configuration_ernie.py（配置类）、modular_ernie.py（模块化源文件）与由它生成的 modeling_ernie.py（完整建模实现）。其中 modular_ernie.py 清晰揭示了实现策略——ERNIE 大量复用 BERT 的现成组件：

• ErnieSelfAttention、ErnieCrossAttention、ErnieLayer、ErniePooler、ErnieLMPredictionHead、ErnieEncoder 分别直接继承自 BERT 的对应类；
• ErnieModel(BertModel)、ErnieForPreTraining(BertForPreTraining) 等头部类亦复用了 BERT 的建模逻辑。

即 ERNIE 在编码器骨干上与 BERT 同构（12/24 层 Transformer、GELU 激活等），差异集中在输入嵌入层与任务能力上。

差异点一：ERNIE 专属的 task embedding（任务嵌入）

ERNIE 2.0/3.0 的持续多任务学习会给不同预训练任务分配一个任务类型 id，让同一模型学习多类任务特性。ErnieEmbeddings（继承自 BertEmbeddings）新增了两项配置驱动的逻辑，见 modeling_ernie.py：

self.use_task_id = config.use_task_id
if config.use_task_id:
    self.task_type_embeddings = nn.Embedding(config.task_type_vocab_size, config.hidden_size)

在 forward 中，词嵌入、token type 嵌入、位置嵌入相加后，若开启 use_task_id，会额外将 task_type_ids 对应的嵌入加入求和结果，再统一经过 LayerNorm 与 dropout。task_type_ids 缺省时自动补零。由于任务数有限（默认 task_type_vocab_size = 3），该新增参数量可以忽略不计。

正是这一差异，导致 ERNIE 所有模型（从 ErnieModel 到各头部类）的 forward 签名都携带一个额外的 task_type_ids 参数。其官方语义为：任务类型嵌入用于刻画不同任务的特征，例如词感知预训练任务、结构感知预训练任务与语义感知预训练任务；每个任务被赋予一个 task_type_id，取值范围为 [0, config.task_type_vocab_size - 1]。

差异点二：从头模型到各任务头均传递 task 信息

ErnieModel.forward 会将其透传给嵌入层与编码器（见 modeling_ernie.py），并支持 use_cache 时的 DynamicCache/EncoderDecoderCache 缓存路径；ErnieForPreTraining、ErnieForCausalLM、ErnieForMaskedLM、ErnieForMultipleChoice、ErnieForNextSentencePrediction、ErnieForSequenceClassification、ErnieForTokenClassification、ErnieForQuestionAnswering 这 8 个任务头则全部把 task_type_ids 原样转发给底层 ernie 主干。

实现能力标注（从源码确认）

ErniePreTrainedModel 在 modeling_ernie.py 中声明了如下关键能力：

• config_class = ErnieConfig，base_model_prefix = "ernie"，预训练权重文件名与主干权重命名保持一致；
• supports_gradient_checkpointing = True：支持梯度检查点以省显存；
• _supports_flash_attn = True、_supports_sdpa = True、_supports_flex_attn = True、_supports_attention_backend = True：底层注意力层支持多种加速后端，可按硬件与场景选择 eager / SDPA / Flash Attention / Flex Attention；
• 权重初始化时对 LMPredictionHead 偏置、position_ids 缓冲与 token_type_ids 缓冲做了针对性初始化。

由于 ERNIE 与 BERT 共用结构，ErnieForMaskedLM、ErnieForPreTraining 的 _tied_weights_keys 将 LM 解码头与 ernie.embeddings.word_embeddings 的权重绑定（tie），这也解释了为何微调后权重文件体积与 BERT 保持一致。

ErnieConfig 配置参数详解

ErnieConfig 定义于 configuration_ernie.py，model_type 固定为 "ernie"。除继承 BERT 系通用参数外，它还新增了两个专属字段。下表给出所有字段的默认值与含义（源码直接可见）：

参数	默认值	说明
vocab_size	30522	词表大小（中文检查点实际词表可能不同，加载权重时以 config.json 为准）
hidden_size	768	隐层维度
num_hidden_layers	12	Transformer 层数
num_attention_heads	12	注意力头数
intermediate_size	3072	FFN 中间层维度
hidden_act	"gelu"	隐藏层激活函数
hidden_dropout_prob	0.1	隐藏层 dropout
attention_probs_dropout_prob	0.1	注意力概率 dropout
max_position_embeddings	512	最大位置编码长度
type_vocab_size	2	segment（句子 A/B）类型词表
task_type_vocab_size	3	ERNIE 2.0/3.0 特有：task_type_ids 的词表大小，即支持的任务类型总数
use_task_id	False	ERNIE 2.0/3.0 特有：是否启用 task embedding；开启后模型才接受 task_type_ids 输入
initializer_range	0.02	参数初始化范围
layer_norm_eps	1e-12	LayerNorm epsilon
pad_token_id	0	padding token id
use_cache	True	自回归解码时是否返回 KV 缓存
classifier_dropout	None	分类头 dropout，为 None 时回退到 hidden_dropout_prob
is_decoder	False	是否作为解码器（影响 use_cache 与注意力掩码）
add_cross_attention	False	是否加入交叉注意力层
bos_token_id / eos_token_id	None	起止符 id
tie_word_embeddings	True	是否绑定输入/输出词嵌入权重

需要特别强调的是：use_task_id 默认关闭，因此日常推理/微调无需关心 task 参数；只有当你想精确复现 ERNIE 2.0/3.0 论文中的多任务训练设定时，才需打开它并为每个训练任务分配 task_type_id。

下游任务模型逐一解析

文档为 8 个任务头提供了官方 API 索引。它们的输出结构、标签语义与损失函数可从 modeling_ernie.py 的源码逐一定位：

ErnieModel（主干模型）

返回 BaseModelOutputWithPoolingAndCrossAttentions，包含 last_hidden_state、pooler_output（[CLS] 池化结果）以及可选的 past_key_values。文档要求其输出需包含各层 hidden_states、attentions、cross_attentions 字段以支撑统一 API。

ErnieForPreTraining（预训练头）

联合 Masked LM 与 NSP 两个目标，输出 ErnieForPreTrainingOutput（即 BertForPreTrainingOutput 的别名，含 loss、prediction_logits、seq_relationship_logits）。源码中 loss = masked_lm_loss + next_sentence_loss，其中：

• labels（形状 (batch, seq_len)）：MLM 标签，-100 位置被忽略；
• next_sentence_label（形状 (batch,)）：取 0 表示 B 句是 A 句的后续，取 1 表示随机句。

文档附带的示例同时验证了该头可输出两组 logits，见 modeling_ernie.py。

ErnieForMaskedLM（掩码语言模型）

输出 MaskedLMOutput。cls 头作用于主干最后一层输出，标签同样采用 -100 忽略填充位的 CrossEntropy 约定（见 modeling_ernie.py）。中文 [MASK] 补全、语法纠错等任务常用该头。

ErnieForNextSentencePrediction（下句预测）

输入为句对（两段文本经 tokenizer 拼为单序列），输出二分类 logits（形状 (batch, 2)）。docstring 示例展示：当第二句为随机句时应得到 logits[0, 0] < logits[0, 1]。预训练语料判断与句间关系任务可选用。

ErnieForCausalLM（因果语言模型/解码器）

支持将 ERNIE 主干配置为 is_decoder=True 的自回归解码器，输出 CausalLMOutputWithCrossAttentions。源码新增 logits_to_keep 参数：提供标签计算 loss 时，可只计算末尾若干 token 的 logits 以节省显存；past_key_values 与 use_cache 提供增量解码。由于主干仍为双向 Transformer，ErnieForCausalLM 需结合解码配置（is_decoder）与注意力掩码使用。

ErnieForSequenceClassification（序列分类）

从 pooler_output 接 dropout 与分类头。源码会按 config.problem_type 自动选择损失：num_labels == 1 用 MSE（回归）、整数标签走单标签 CrossEntropy、否则走 BCEWithLogitsLoss（多标签），与 Transformers 其余模型行为一致。适用于情感分析、文本蕴含等句子级任务。

ErnieForMultipleChoice（多项选择）

输入形状为 (batch_size, num_choices, seq_len)，forward 内部先展平为 (batch*num_choices, seq_len) 过主干，再将 pooler_output 分类 logits reshape 回 (batch, num_choices) 计算 CrossEntropy。适合阅读理解选择、选项式问答。

ErnieForTokenClassification（Token 级分类）

对每层的每个 token 接线性分类头，输出 (batch, seq_len, num_labels) 的 logits，用逐位置 CrossEntropy 计算 loss。适用于命名实体识别（NER）、词性标注等序列标注任务。

ErnieForQuestionAnswering（抽取式问答）

对主干输出接 qa_outputs，输出 start_logits 与 end_logits 各为 (batch, seq_len)；start_positions/end_positions 越界位置会被 clamp 并以 ignore_index 忽略，最终 loss 为起止两个 CrossEntropy 的平均。可直接用于抽取式中文阅读理解（SQuAD 风格数据集）。

从测试文件 tests/models/ernie/test_modeling_ernie.py 可见，这 9 个模型类（ErnieModel + 8 个任务头）均被纳入通用一致性测试，且 pipeline_model_mapping 注册了 feature-extraction、fill-mask、text-classification、text-generation、token-classification、zero-shot 等 pipeline 映射，说明它们与 Transformers 的既有生态（pipeline、generate、设备放置等）完全兼容。

用测试用例验证行为边界

tests/models/ernie/test_modeling_ernie.py 中的 ErnieModelTester 使用微型配置（vocab_size=99, hidden_size=32, num_hidden_layers=2, num_attention_heads=4）做前向形状断言，值得关注的测试覆盖点包括：

• test_model、test_for_masked_lm、test_for_pretraining：验证 last_hidden_state/pooler_output/prediction_logits 形状；
• test_model_as_decoder、test_for_causal_lm_decoder、test_decoder_model_past_with_large_inputs：把 config.is_decoder=True、config.add_cross_attention=True 后验证 decoder 模式与 past_key_values 缓存一致性（带缓存与不带缓存输出的数值 allclose，容差 1e-3）；
• 各任务头的标签/位置参数（labels、next_sentence_label、start_positions、end_positions）形状校验；
• test_config：对 ErnieConfig 做通用配置一致性检查；
• @slow 的 test_model_from_pretrained：真实加载 nghuyong/ernie-1.0-base-zh 权重，验证转换后的 checkpoint 可正常从 Hub 恢复。

如果你本地已经安装好仓库源码环境与依赖，可在 tests/models/ernie 目录下运行对应测试来复现这些行为：

python -m pytest tests/models/ernie/test_modeling_ernie.py -k "test_model or test_for_masked_lm or test_for_pretraining"

（@slow 标记的联网测试需要先能访问模型 Hub。）

总结与实践建议

在 🤗 Transformers 中使用 ERNIE，可以归纳为三点：

1. 选择检查点：中文任务优先 ernie-3.0-* 系列，按显存从 nano/micro/mini/medium/base 中选取；英文任务使用 ernie-2.0-*；医疗领域中文可用 ernie-health-zh；
2. 按任务选头：掩码补全用 ErnieForMaskedLM/fill-mask pipeline；分类/序列标注/问答等直接用对应 ErnieForXxx 类；续写或生成类实验可把 ERNIE 当作 decoder 用 ErnieForCausalLM；
3. 关注 config：ERNIE 与 BERT 的结构差异集中在 task_type_vocab_size 与 use_task_id 两个字段，日常使用默认值即可。

如需进一步了解百度 ERNIE 的新一代架构（非本模型文档范围），可继续阅读仓库中的 Ernie 4.5 与 Ernie 4.5 MoE 文档，或浏览 nghuyong 账号下的全部 ERNIE 检查点列表以获取更多语言任务应用示例。