TensorFlow Models NLP Modeling Library：用 Keras 原语与 Scaffold 类组装 Transformer 架构的技术详解

本文围绕 TensorFlow Models 仓库中的 NLP Modeling Library 展开：它是一套由 tf.keras.Layer 与 tf.keras.Model 原语构成的 NLP 建模组件库，可自由组装成 Transformer 类模型，且同时兼容 TF1 与 TF2。读完后，你将掌握该库 layers/networks/models/losses 四层结构的职责划分，理解 TransformerScaffold 与 EncoderScaffold 两个脚手架类的参数注入机制，并能参照仓库中 BERT/ALBERT 的真实实现，用这套原语构建和改造自己的 Transformer 编码器。

一、四层组件结构：layers、networks、models 与 losses

库总览文档 将组件按抽象层级分为四类，每一层都有独立的 README 说明：

• layers（目录文档）：NLP 模型的基本构建块，均为可复用的 tf.keras 层，既可单独使用，也可拼装出新的层或模型。
• networks（目录文档）：层的组合体，是不单独训练的子网络。它把 Transformer 编码器这类常见网络结构封装成带有标准化配置的、易于操作的 tf.keras 对象。
• models（目录文档）：可以训练的组合体。仓库提供了若干预置模型用于训练编码器网络，它们既是便捷函数，也是规范示例。
• losses（目录文档）：NLP 任务中常用的损失计算，目前提供 weighted_sparse_categorical_crossentropy_loss，计算带权重的按 batch 稀疏分类交叉熵损失。

1.1 layers：原子能力清单

layers 目录 下共 70 余个源文件，覆盖注意力变体、位置编码、嵌入、掩码与文本处理等方向，核心成员包括：

组件	文件	作用
MultiHeadAttention	attention.py	可选掩码的 query/key/value 多头注意力；当 from_tensor 与 to_tensor 相同时即自注意力
CachedAttention	attention.py	带缓存的注意力层，用于自回归解码
BigBirdAttention	bigbird_attention.py	稀疏注意力，把二次复杂度降为线性
KernelAttention	kernel_attention.py	核特征图线性点积注意力，含 Performer 等方法
MatMulWithMargin	mat_mul_with_margin.py	带 margin 的矩阵乘法，用于检索/排序任务的训练
MultiChannelAttention	multi_channel_attention.py	多头注意力的多通道变体，可合并多个流用于交叉注意力
TalkingHeadsAttention	talking_heads_attention.py	Talking-Heads 注意力
Transformer / TransformerDecoderBlock	transformer.py	可选掩码的完整 Transformer 与解码器块（自注意力 + 交叉注意力 + FFN）
ReuseMultiHeadAttention / ReuseTransformer	reuse_attention.py	复用低层注意力分数，避免高层重复计算
ReZeroTransformer	rezero_transformer.py	采用 ReZero 的 Transformer，支持大深度快速收敛
OnDeviceEmbedding	on_device_embedding.py	面向 TPU 模型设计的高效嵌入查表
PositionalEmbedding	position_embedding.py	按 BERT 论文实现的位置嵌入
SelfAttentionMask	self_attention_mask.py	从二维掩码张量生成三维注意力掩码
MaskedSoftmax	masked_softmax.py	可选掩码的 softmax；掩码为 1 处放行、为 0 处置零
MaskedLM	masked_lm.py	掩码语言模型头，需要传入嵌入表变量
ClassificationHead	cls_head.py	序列嵌入上的池化头，常用于分类任务
GatedFeedforward	gated_feedforward.py	GLU 门控前馈层
MultiHeadRelativeAttention / TwoStreamRelativeAttention	relative_attention.py	支持相对位置编码的注意力变体（Transformer-XL / XLNet 风格）
TransformerXL	transformer_xl.py	含 TransformerXLBlock 的完整 Transformer-XL
MobileBertEmbedding / MobileBertTransformer	mobile_bert_layers.py	MobileBERT 论文的嵌入层与 Transformer 层
BertPackInputs / BertTokenizer / SentencepieceTokenizer	text_layers.py	将原始文本分词并打包为 BERT 输入
TransformerEncoderBlock	transformer_encoder_block.py	单层可选掩码的 Transformer 编码器块

此外还有 FactorizedEmbedding（factorized_embedding.py）、MixtureOfExperts（moe.py）、Mixing（mixing.py）等面向长序列与高效建模的层，说明该库的组件面远不止 BERT 一种结构。

1.2 networks：不单独训练的子网络

networks 目录 收录的组件是“结构级”封装：

• BertEncoder（bert_encoder.py）：双向 Transformer 编码器，含嵌入查表、Transformer 层与池化层；
• AlbertEncoder（albert_encoder.py）：将嵌入参数分解为两个小矩阵、并在层间共享参数；
• MobileBERTEncoder（mobile_bert_encoder.py）：面向资源受限设备的紧凑编码器；
• Classification（classification.py）：单隐层，可作分类头，类别数设为 1 时即回归头；
• PackedSequenceEmbedding（packed_sequence_embedding.py）：支持 packed 序列与位置 id 的嵌入网络；
• SpanLabeling（span_labeling.py）：基于单个 Dense 隐层的单跨度预测头（每 batch 样本预测一个起止索引），可用于 SQuAD 任务；
• XLNetBase（xlnet_base.py）：XLNet 基础网络，含嵌入查表、相对位置编码、掩码与分段矩阵计算、单/双流相对自注意力的 Transformer XL 层；
• FNet（fnet.py）：用傅里叶子层替代自注意力子层的编码器；
• SparseMixer（sparse_mixer.py）：由异构编码器块组成，每个块包含线性混合或注意力子层 + 稠密 MLP 或稀疏激活 MoE 子层。

1.3 models：可训练的完整模型

models 目录 中预置了“开箱可训练”的模型：

• BertClassifier（bert_classifier.py）：带单一 Classification 头的简单分类模型，类别数为 1 时兼作回归模型；
• BertTokenClassifier（bert_token_classifier.py）：对序列输出嵌入做逐 token 分类；
• BertSpanLabeler（bert_span_labeler.py）：单跨度起止索引预测器，适配 SQuAD 类任务；
• BertPretrainer（bert_pretrainer.py）：掩码语言模型 + 分类头的预训练模型；
• DualEncoder（dual_encoder.py）：适配检索任务的双编码器；
• Seq2SeqTransformer（seq2seq_transformer.py）：原始 Transformer 的 seq2seq 模型；
• T5Transformer（t5.py）：独立的 T5 模型，与已发布的 T5 架构及转换检查点兼容，模块以 tf.Module 实现，可用 Keras 自定义层包裹后使用。

1.4 losses：任务损失

losses 目录 目前提供 weighted_sparse_categorical_crossentropy_loss（实现文件），用于计算带权重的按 batch 稀疏分类交叉熵，是 tagging 等序列标注类预训练任务的标准损失。

二、Scaffold 设计：不 fork 整个 Transformer 即可改造注意力与编码器

总览文档（official/nlp/modeling/README.md）强调：除了预定义原语，该库还提供脚手架（scaffold）类，便于在新架构上做实验——例如想换一种注意力原语，无需 fork 整个 Transformer 对象。仓库中对应两个 scaffold：

2.1 TransformerScaffold：可替换注意力与前馈的 Transformer 层

TransformerScaffold 是一个 Keras 层，实现 "Attention Is All You Need" 的 Transformer 结构，同时开放注意力层与前馈层两个自定义点。从源码看，其构造函数签名（见 transformer_scaffold.py#L84-L103）为：

参数	默认值	说明
num_attention_heads	必填	注意力头数；输入最后一维必须可被其整除
inner_dim	768	两层前馈网络第一层 Dense 的输出维度（兼容旧参数字段 intermediate_size）
inner_activation	gelu	前馈第一层的激活函数
attention_cls	attention.MultiHeadAttention	注意力层的类或实例
attention_cfg	None	用于实例化 attention_cls 的配置字典
feedforward_cls / feedforward_cfg	None	前馈层的类/实例与配置；为 None 时使用论文标准前馈
dropout_rate / attention_dropout_rate	0.0	注意力后/输出处的 dropout；注意力层内部 dropout
norm_first	False	True 时先归一化再进注意力与中间 Dense（pre-norm），否则对二者输出做归一化（post-norm）
norm_epsilon	1e-12	LayerNorm 的 epsilon
kernel_initializer 等	glorot_uniform 等	Dense 层的初始化器、正则器与约束

自定义注入的解析逻辑在 build 方法中（transformer_scaffold.py#L155-L205）：内部 get_layer_instance 支持三种输入形态——直接传 Layer 实例则原样使用；传已序列化的 config dict 则反序列化后递归处理；传类时用 attention_cfg 实例化，若 attention_cfg 为 None 则回落到默认配置：

default_attention_cfg = {
    "num_heads": self._num_heads,
    "key_dim": self._attention_head_size,   # hidden_size // num_heads
    "dropout": self._attention_dropout_rate,
    "name": "self_attention",
}

call 方法（transformer_scaffold.py#L294-L359）支持三类输入：单张量、[input, attention_mask] 二元组、[input, key_value, attention_mask] 三元组——后两者分别对应自注意力与带显式 key/value 的注意力，且 key_value 缺省时自动取 input。前馈块的约定也值得注意：若未提供 feedforward_cls，内部用 EinsumDense("abc,cd->abd") 自建两层前馈；若提供了自定义块且 norm_first=False，则假定该块自身负责 LayerNorm 与残差连接，源码注释中点名 GatedFeedforward 作为具体示例。另外两个实现细节：LayerNorm 固定用 float32 以保证数值稳定性；mixed_bfloat16 策略下中间激活会回落 float32（注释说明 bfloat16 配合 LAMB 优化器时收敛不佳）。该层同时被 @gin.configurable 装饰，可直接在 gin 配置文件中按名实例化。

2.2 EncoderScaffold：BERT 风格编码器的三大注入点

EncoderScaffold 是一个 Keras Model，实现 BERT 论文中的双向 Transformer 编码器骨架，提供两个自定义注入点：

• embedding_cls / embedding_cfg / embedding_data：自定义嵌入子网络，替换标准嵌入逻辑。自定义网络需接受原始输入并输出两个值：形状 (batch_size, seq_length, hidden_size) 的嵌入张量和形状 (batch_size, seq_length, seq_length) 的注意力掩码。若做标准预训练，还需通过 embedding_data 传入词嵌入表引用（供 get_embedding_table() 使用）。未设置时，默认嵌入逻辑在构造函数内按 Keras Functional API 构建（encoder_scaffold.py#L143-L196）：三个输入 input_word_ids、input_mask、input_type_ids（均为 (seq_length,) 的 int32），经 OnDeviceEmbedding 词嵌入 + PositionEmbedding 位置嵌入 + one-hot OnDeviceEmbedding 类型嵌入相加，再做 LayerNorm 与 Dropout；embedding_cfg 需给出 vocab_size、type_vocab_size、hidden_size、max_seq_length、seq_length、initializer、dropout_rate 等键。
• hidden_cls / hidden_cfg：自定义隐层，替换编码器中的 Transformer 实例化。支持类、实例、以及类/实例混合的列表三种形态；列表形态时其长度必须等于 num_hidden_instances，否则抛出 RuntimeError（encoder_scaffold.py#L204-L209）。hidden_cfg 缺省需提供 num_attention_heads、intermediate_size、intermediate_activation、dropout_rate、attention_dropout_rate、kernel_initializer 等键；当 hidden_cls 未设置时默认使用 layers.Transformer。此外还有 mask_cls（默认 layers.SelfAttentionMask）与 mask_cfg 两个配套参数，负责由输入与二维 mask 生成传给隐层的注意力掩码。
• 输出形态控制：pooled_output_dim 决定池化维度，池化层是对最后一层输出取第 0 个 token（last_layer_output[:, 0, :]）后的 tanh Dense（encoder_scaffold.py#L246-L254）；return_all_layer_outputs 可返回所有层的序列输出，dict_outputs 则输出 sequence_output / pooled_output / encoder_outputs 三个键的字典；layer_norm_before_pooling 建议在 transformer 层设 norm_first=True 时同步开启。

其他值得留意的参数：feed_layer_idx 会把层索引写入 hidden_cfg['layer_idx']，layer_idx_as_attention_seed 会按层深拷贝 hidden_cfg 并为各层设置不同的 attention_cfg['seed']，recursive 模式则让隐层以 (data, state) 递归形式传递状态（初始为 None）。整个网络用 Functional API 构建后再调用 super().__init__，因而具备功能式模型的全部特性（encoder_scaffold.py#L267-L275）。

仓库附带了配套教程 customize_encoder.ipynb，演示如何借助 scaffold 类构建新架构；总览文档中链接的 NLP modeling library intro.ipynb 则演示如何用上述原语从零组装 Transformer 类 NLP 模型。

三、真实落地：BERT/ALBERT 模型如何建立在该库之上

总览文档最后指出：仓库中的 BERT 与 ALBERT 模型正是用这套库实现的，代码示例可在对应模型文件夹中找到。这一点在 BertEncoderV2 的源码中得到印证——其默认参数即 BERT-Base 的配置（bert_encoder.py#L86-L107）：

参数	默认值	含义
vocab_size	必填	词表大小
hidden_size	768	Transformer 隐层维度
num_layers	12	Transformer 层数
num_attention_heads	12	每层注意力头数，要求 hidden_size 可整除
max_sequence_length	512	决定位置嵌入变量形状
type_vocab_size	16	类型 id 取值范围
inner_dim	3072	前馈第一层输出维度
inner_activation	近似 gelu	前馈激活函数
output_dropout / attention_dropout	0.1	注意力后/输出处与注意力内部 dropout
initializer	TruncatedNormal(stddev=0.02)	全部权重的初始化器
embedding_width	None	若非空且异于 hidden_size，嵌入参数分解为 [vocab, embedding_width] 与 [embedding_width, hidden_size] 两个矩阵（ALBERT 风格）
norm_first	False	前归一化开关，与 Scaffold 同名参数语义一致
return_attention_scores	False	额外输出所有层注意力分数列表

从源码结构看，BertEncoderV2 通过消费 V1 风格的旧参数名（intermediate_size、activation、dropout_rate、attention_dropout_rate）并映射到新命名（bert_encoder.py#L108-L120），保证了 gin 配置与旧检查点脚本的向后兼容。ALBERT 实现可进一步参见 albert_encoder.py 与 models 目录中的预训练模型，它们分别消费了 embedding_width 分解嵌入与层间共享等机制——这正是 scaffold 注入点思想在最终产品模型上的体现。

四、使用建议与延伸阅读

结合文档与源码，一个典型的使用路径是：

1. 在 layers 中挑选注意力与前馈原语（如 MultiHeadAttention、GatedFeedforward），需要整块 Transformer 时直接用 Transformer 或 TransformerScaffold；
2. 需要完整 BERT 风格编码器时，用 EncoderScaffold 并通过 embedding_cls / hidden_cls 替换单个组件，而不是复制整段编码器代码；
3. 在 models 中选取可训练模型（如 BertPretrainer、BertSpanLabeler）作为“规范示例”，对照 networks 实现 理解其内部组装方式；
4. 参考 docs/nlp/index.ipynb（原语组装教程）与 docs/nlp/customize_encoder.ipynb（scaffold 定制教程）完成端到端实践。

需要说明的适用前提：该库面向 TF1/TF2 双兼容场景，源码中标注了部分混合精度与 bfloat16 下的数值处理策略（如 LayerNorm 固定 float32、mixed_bfloat16 下中间激活回落 float32），在 TPU 与大规模预训练环境中应按仓库内注释保持默认配置；对组件的详细行为边界，建议直接查阅各目录下与源码同名的 _test.py 测试文件（如 transformer_scaffold_test.py、encoder_scaffold_test.py）作为行为验证依据。