tf-models-official 中的 FFF-NER：面向少样本命名实体识别的预训练式微调实现全解析

本文基于 tf-models-official 仓库中 official/projects/fffner 模块的 README 与配套源码，系统讲解 FFF-NER（Formulating Few-shot Fine-tuning for NER）这一少样本命名实体识别训练任务的完整落地方式：如何将预训练 BERT 转换为项目可用的 encoder 检查点、如何把原始语料转换为 tf_record、如何通过 official.nlp.train 入口在 Cloud TPU 上启动训练与评估，以及双头分类模型、负采样数据构造与 overall_f1 评估指标的底层实现细节。

1. 项目定位：把 NER 重新表述为"掩码跨度"的二分类+类型分类

项目 README 说明，FFF-NER 是论文 Formulating Few-shot Fine-tuning Towards Language Model Pre-training: A Pilot Study on Named Entity Recognition（arXiv:2205.11799，Zihan Wang 等，2022）的官方 TensorFlow 实现，README 同时标注该实现仍处于开发中（"This implementation is still under development"）。它被组织为 tf-models-official 的一个 project，复用 official/nlp 的训练框架、official/modeling 的 Transformer 组件与配置/注册体系。

与逐词打 BIO 标签的传统 NER 不同，该任务的建模思路是：把待判定的文本跨度（span）用 [MASK] 探针插入句子中，让模型回答两个问题——"这个位置的跨度是不是实体"（is_entity，2 类）与"实体属于哪种类型"（entity_type）。这一表述方式让 NER 任务在结构上更接近语言模型预训练时"根据上下文填空"的判别形式，从而更利于少样本场景下的微调。

2. 目录结构与入口：一个可注册的实验任务

模块文件组织如下（official/projects/fffner）：

文件	职责
train.py	训练入口，定义 flags 并调用 official.nlp.train.main
fffner_experiments.py	注册名为 fffner/ner 的实验配置工厂
fffner.py	从配置构建 FFFNerEncoder
fffner_encoder.py	核心 encoder：BERT 风格 Transformer + 位置探针
fffner_classifier.py	双分类头包装模型
fffner_prediction.py	Task：损失、指标、校验、评估日志聚合
fffner_dataloader.py	tf_record 数据加载器与数据配置
experiments/base_conll2003.yaml	CoNLL-2003 数据集基线实验配置
experiments/base_restaurants.yaml	餐厅点评数据集基线实验配置
utils/create_data.py	原始语料 → tf_record 的转换脚本
utils/convert_checkpoint_tensorflow.py	TF Hub BERT → FFFNerEncoder 检查点转换

fffner_experiments.py 中通过 @exp_factory.register_config_factory('fffner/ner') 把任务配置注册为 fffner/ner 实验：task 使用 FFFNerPredictionConfig，encoder 类型标记为 any（即加载自定义 encoder 配置）；trainer 默认 AdamW（weight_decay_rate 0.01，LayerNorm/layer_norm/bias 不施加权重衰减）+ 多项式学习率衰减（初始 2e-5，衰减至 0）+ 多项式 warmup，并附带 task.train_data.is_training != None 等约束断言。这正是 README 训练命令中 --experiment=fffner/ner 能够生效的注册机制。

3. 模型结构：FFFNerEncoder 的双探针机制

fffner_encoder.py 中的 FFFNerEncoder 在标准 BERT 编码器（默认 hidden_size=768、12 层、12 头、inner_dim=3072、近似 GELU 激活、TruncatedNormal(stddev=0.02) 初始化）的基础上，有两处关键扩展：

1. 额外的两个标量输入：在 self.inputs 字典中，除了常规的 input_word_ids、input_mask、input_type_ids，还包含 is_entity_token_pos 与 entity_type_token_pos（源码 L215-L220），分别记录"实体判别探针"与"类型判别探针"在序列中的位置下标。
2. 按位置取出的双池化头：call() 在最后一层输出上，用 tf.gather(last_encoder_output, indices=..., axis=1, batch_dims=1) 分别抽取两个探针位置的向量，再各自经过一个 tanh 激活的 Dense 池化层 pooler_transform_is_entity / pooler_layer_entity_type，拼接为 pooled_output（源码 L269-L286）。

fffner_classifier.py 中的 FFFNerClassifier 则把 encoder 输出再包装为最终模型：取 pooled_output 的第 0、1 两个向量位置，经 Dropout 后分别送入两个 ClassificationHead（num_classes_is_entity、num_classes_entity_type 可配置），模型输出为 [is_entity_logits, entity_type_logits] 列表。该类还实现了 checkpoint_items 属性，把 encoder 与两个分类头分键暴露，供检查点恢复使用。

任务侧的损失在 fffner_prediction.py 中定义：两个头各用一个 sparse_categorical_crossentropy（from_logits=True）交叉熵，二者相加并加均值；build_metrics 提供 cls_accuracy_is_entity 与 cls_accuracy_entity_type 两个稀疏分类准确率指标；initialize() 支持通过 init_checkpoint 以 partial 方式恢复预训练权重。

4. 数据构造：BIO 标签到"探针实例"的转换与负采样

README 指出原始数据集为 .words（分词文本）与 .ner（BIO 标签）成对文件，例如 /data/fffner_datasets/conll2003/few_shot_5_0.words 和 few_shot_5_0.ner。utils/create_data.py 的 NERDataset 负责把它们转成训练/评估用的 tf_record，核心逻辑包括：

• BIO → 跨度：bio_labels_to_spans() 将 B-/I- 前缀标签折叠为 (span_start, span_end, type) 三元组，并只保留在 entity_map.json 中出现过的实体类型（L164-L185）；
• 探针式序列化：process_word_list_and_spans_to_inputs() 对每个候选跨度，构造 [CLS] + 前文 + "[" + [MASK] + "]" + "[" + 跨度内词 + "]" + "[" + [MASK] + "]" + 后文 + [SEP] 的 token 序列，其中第一个 [MASK] 位置即 is_entity_token_pos，第二个即 entity_type_token_pos（L114-L162）。脚本还提供 ablation_not_mask、ablation_no_brackets、ablation_span_type_together 三个消融开关，对应论文中的消融实验；
• 负采样：训练集上，除真实实体跨度外，还会对"非实体"（O）跨度按概率采样负样本——prepare(negative_multiplier=3.) 中负样本数量约为 (len(tokens) + num_entities * 10) * 3，采样概率与跨度同实体的交叠程度相关（e^intersection_size），上限为所有可能跨度数 n(n+1)/2（L195-L248）。这保证了二分类头见到足够的负例；
• 测试集枚举全跨度：评估时（is_train=False）枚举句中所有可能的跨度，保证评估覆盖完整；
• tf_record 字段：每条样本写出 input_ids、input_mask、segment_ids、is_entity_token_pos、entity_type_token_pos、is_entity_label、entity_type_label、example_id、sentence_id、span_start、span_end 等 Int64 特征（L303-L347），输出为 {dataset_name}_{fold}.tf_record，例如 conll2003_few_shot_5_0.tf_record 与 conll2003_test.tf_record。

注意脚本使用 HuggingFace transformers.AutoTokenizer.from_pretrained("bert-base-uncased") 分词，默认序列上限 max_len = 128（L68），超长样本会被跳过。

5. 检查点转换：把 TF Hub 的 BERT 权重装进 FFFNerEncoder

训练前必须先把预训练语言模型转换为项目的 encoder 格式。utils/convert_checkpoint_tensorflow.py 默认加载 TF Hub 的 bert_en_uncased_L-12_H-768_A-12/3（即 base 尺寸 BERT uncased），流程为：

1. _get_tensorflow_bert_model_and_config() 从 Hub 模型权重名反推出 num_attention_heads、hidden_size、intermediate_size、层数、vocab_size 等配置（L27-L50）；
2. _create_fffner_model() 依配置构建 FFFNerEncoder；
3. convert() 逐层把 Hub 模型的 word_embeddings、layer_norm、position_embedding、Q/K/V/output 投影、intermediate/output FFN 权重拷贝进新 encoder（L73-L138）；
4. 最终 tf.train.Checkpoint(encoder=encoder).write(output_path)，输出路径默认为 tf-bert-uncased（L141-L175）——这正是两个实验 YAML 中 init_checkpoint: 'tf-bert-uncased' 所指的路径。

README 中给出的默认命令即：

python3 utils/convert_checkpoint_tensorflow.py

仓库另附 utils/convert_checkpoint_huggingface.py，用于从 HuggingFace 格式权重转换。

6. 数据加载：FFFNerDataLoader 的特征解析

训练/验证数据由 fffner_dataloader.py 加载。FFFNerDataConfig 的关键字段与默认值：

• input_path：tf_record 路径（YAML 中为 TODO 占位，训练时用 params_override 覆盖）；
• global_batch_size：默认 32（训练）/ 1024（验证）；
• seq_length：默认 128，与造数脚本一致；
• label_field_is_entity / label_field_entity_type：标签字段名，默认 is_entity_label / entity_type_label；
• file_type：默认 tfrecord；
• is_training、drop_remainder、include_example_id：验证集默认 is_training=False、drop_remainder=False、include_example_id=True（见 fffner_experiments.py L39-L41）。

_decode() 解析 tf.Example 后会把所有 int64 张量 cast 为 int32（注释说明原因是"tf.Example 只支持 int64，而 TPU 只支持 int32"）；_parse() 做键名映射，如 input_ids → input_word_ids、segment_ids → input_type_ids（L99-L120），最终字典恰好匹配 encoder 的输入签名。

7. 训练与评估：命令与实验配置逐项解读

README 给出的完整训练命令（在 official/projects/fffner 目录下执行）为：

PATH_TO_TRAINING_RECORD=conll2003_few_shot_5_0.tf_record # path to the training record
PATH_TO_TESTING_RECORD=conll2003_test.tf_record # path to the evaluation record
TPU_NAME="<tpu-name>"  # The name assigned while creating a Cloud TPU
MODEL_DIR=/tmp/conll2003_ew_shot_5_0 # directory to store the experiment
# Now launch the experiment.
python3 -m official.projects.mosaic.train \
  --experiment=fffner/ner \
  --config_file=experiments/base_conll2003.yaml \
  --params_override="task.train_data.input_path=${PATH_TO_TRAINING_RECORD},task.validation_data.input_path=${PATH_TO_TESTING_RECORD},runtime.distribution_strategy=tpu"
  --mode=train_and_eval \
  --tpu=$TPU_NAME \
  --model_dir=$MODEL_DIR

（该命令来自 README 原文，其中 --params_override 行末缺少反斜杠、MODEL_DIR 拼写 ew_shot 为原文如此，实际执行时请自行修正为 /tmp/conll2003_few_shot_5_0 并补上行续符。）造数命令则为：

export PATH_TO_DATA_FOLDER=/data/fffner_datasets/
export DATASET_NAME=conll2003
export TRAINING_FOLD=few_shot_5_0
python3 utils/create_data.py $PATH_TO_DATA_FOLDER $DATASET_NAME $TRAINING_FOLD

experiments/base_conll2003.yaml 的关键配置（CoNLL-2003，4 种实体类型）：

配置项	值	说明
task.init_checkpoint	tf-bert-uncased	第 5 节转换出的 encoder 检查点
task.model.num_classes_is_entity	2	"是否实体"二分类头
task.model.num_classes_entity_type	4	实体类型头（CoNLL-2003 对应 4 类）
task.train_data.global_batch_size / seq_length	32 / 128	与造数脚本的 128 上限对齐
task.validation_data.global_batch_size	1024	评估批大小
trainer.learning_rate	多项式：2.0e-05 → 0，decay_steps=2189	学习率线性衰减
trainer.train_steps	2189	注释 # 2335 * 30 / 32：2335 个样本跑 30 个 epoch，除以 batch size 32
trainer.validation_interval / checkpoint_interval	500 / 1000	评估/存盘频率
trainer.best_checkpoint_eval_metric	overall_f1（higher）	按整体 F1 选出最佳检查点，导出到 best_ckpts 子目录

experiments/base_restaurants.yaml 结构相同，差异在于 num_classes_entity_type: 8（餐厅数据集实体类型更多）以及 train_steps: 2518（注释 # 2686 * 30 / 32），说明换数据集时需要按"样本数 × epoch 数 / batch size"重算训练步数与 decay_steps。

训练入口 train.py 本身极薄：tfm_flags.define_flags() 定义通用 flags 后直接 app.run(train.main)，复用 official/nlp/train.py 的整套训练/评估/导出流程，因此 --mode=train_and_eval、--tpu、--model_dir、--params_override 等行为均与 tf-models-official 其他 NLP 项目一致。

8. 评估指标：从"逐跨度预测"到 raw_f1 / resolved_f1 / overall_f1

fffner_prediction.py 的 reduce_aggregated_logs() 实现了评估闭环：

1. 验证阶段每个样本都会把预测与标签连同 example_id、sentence_id、span_start、span_end 一并记录（validation_step 中的 logs）；
2. reduce_aggregated_logs 按 sentence_id 分组，把每个跨度的 is_entity 预测 softmax 后取最大概率，若其"是实体"概率大于"否"，则产生一条预测实体 (sentence_id, span_start, span_end, argmax(entity_type))；
3. resolve() 对同句内相互重叠的预测跨度按置信度降序去重，保留先命中的跨度并标记占用区间（L258-L275）；
4. 最终输出 6 个指标：raw_f1/raw_precision/raw_recall（未做重叠消解的预测）、resolved_f1/resolved_precision/resolved_recall（消解后），以及 overall_f1 = raw_f1 + resolved_f1——这与 YAML 中 best_checkpoint_eval_metric: "overall_f1" 的选型逻辑直接对应。

9. 运行前提与注意事项

• 依赖：README 的 Requirements 以 tf-models-official 的 PyPI 包与 TensorFlow 为准（README L21-L24）；造数脚本额外依赖 transformers 与 tqdm。
• 硬件：README 说明训练"可以运行在 Google Cloud Platform 的 Cloud TPU 上"，命令通过 --tpu=<tpu-name> 指定实例；distribution_strategy=tpu 由 params_override 显式给出。
• 数据格式：每个数据集目录需同时包含训练折（few_shot_X_Y.words/.ner）、test.words/test.ner 以及 entity_map.json（实体类型到类别的映射，其键顺序决定类别编号）。README 指向作者仓库中的示例数据与格式说明。
• 开发状态：README 明确标注实现仍在开发中，FFFNerTask.build_model 中对 XLNet encoder 分支直接 assert False, 'Not supported yet'（fffner_prediction.py L89-L90），即当前仅支持 BERT 风格的任意 encoder 配置。
• 可验证的模型测试：encoder 的构造行为有对应单元测试 fffner_encoder_test.py 可作参考。

10. 引用信息

按 README 建议，引用该实现时应引用原始论文：

@article{wang2022formulating,
  title={Formulating Few-shot Fine-tuning Towards Language Model Pre-training: A Pilot Study on Named Entity Recognition},
  author={Wang, Zihan and Zhao, Kewen and Wang, Zilong and Shang, Jingbo},
  journal={arXiv preprint arXiv:2205.11799},
  year={2022}
}

该模块以 Apache License 2.0 许可发布，维护者为 Zihan Wang。整体来看，official/projects/fffner 是一个"麻雀虽小五脏俱全"的 tf-models-official 项目范本：从 exp_factory 实验注册、task_factory 任务注册、data_loader_factory 数据加载，到检查点恢复与最佳模型导出，完整展示了如何把一篇少样本 NER 论文落地为可在 TPU 上跑通的配置驱动训练流水线。