TensorFlow models 官方 NLP networks 模块解析:BertEncoder、ALBERT、XLNet 与 FNet 等 Keras 编码器架构全览

本篇基于官方 official/nlp/modeling/networks/README.md 及其对应源码,系统讲解 models 仓库中 NLP 建模的 "Networks" 模块:它如何用 tf.keras 把 BERT、ALBERT、MobileBERT、XLNet、FNet、Sparse Mixer 等编码器以及分类/跨度标注头等结构封装成带标准化配置的 Keras 模型。读完你将理解 Network 与 Layer/Task 的层级关系、各编码器关键参数的默认值与含义,以及如何在 Keras Functional 模型中组合编码器与任务头。

一、什么是 Networks:定位与设计原则

README 对 Network 的定义是:

Networks are combinations of tf.keras layers (and possibly other networks). They are tf.keras models that would not be trained alone. It encapsulates common network structures like a transformer encoder into an easily handled object with a standardized configuration.

即 Network 是由 Keras 层(以及其它 Network)组合而成的、不会单独训练的 Keras 模型,它把 Transformer 编码器这类通用网络结构封装成配置标准化的对象。这一点同样体现在包文档 official/nlp/modeling/networks/init.py 的 docstring 中,并统一从该包导出 BertEncoder、AlbertEncoder、MobileBERTEncoder、Classification、PackedSequenceEmbedding、SpanLabeling、XLNetBase、FNet、SparseMixer 等类。

从源码结构看,Networks 处在三层结构中:

• layers(位于 official/nlp/modeling/layers):可复用的 tf.keras.layers.Layer,如 TransformerEncoderBlock、OnDeviceEmbedding、PositionEmbedding、SelfAttentionMask;
• networks(本模块):把若干 layer 串成一个"半成品"模型,负责输入定义、嵌入组合、层堆叠和输出组织,但不含训练目标(没有损失函数与任务逻辑);
• tasks(位于 official/nlp/tasks):在 Network 之上接任务头、定义损失与指标,由 official/nlp/train.py 等训练入口驱动。

这一分层意味着:Network 的输出是中间张量(如 sequence_output、pooled_output),必须交给 Task 或外部损失函数才有意义。

模块清单

README 列出的 9 个核心网络如下(文件均位于 official/nlp/modeling/networks/):

网络类	文件	对应论文 / 说明
BertEncoder	bert_encoder.py	BERT:双向 Transformer 编码器,含嵌入查找、Transformer 层与池化层
AlbertEncoder	albert_encoder.py	ALBERT:嵌入参数分解为两个小矩阵,且跨层共享参数
MobileBERTEncoder	mobile_bert_encoder.py	MobileBERT:面向资源受限设备的紧凑 BERT
Classification	classification.py	单层隐藏层的分类/回归头(类数为 1 时为回归)
PackedSequenceEmbedding	packed_sequence_embedding.py	支持打包序列(packed sequences)与 position ids 的嵌入网络
SpanLabeling	span_labeling.py	单跨度标注头(每条样本预测一对 start/end 索引),用于 SQuAD 类任务
XLNetBase	xlnet_base.py	XLNet:相对位置编码、mask 计算、segment 矩阵、单/双流相对自注意力
FNet	fnet.py	用傅里叶混合层替代(全部或大部分)自注意力子层的编码器
SparseMixer	sparse_mixer.py	由异构编码器块组成:线性混合或注意力子层 + 稠密 MLP 或稀疏激活 MoE 子层

此外,从 official/nlp/modeling/networks/init.py 的导出还可以看到 BertEncoderV2、XLNetSpanLabeling、EncoderScaffold(encoder_scaffold.py)和 FunnelTransformerEncoder(funnel_transformer.py)等补充实现,说明该模块的导出面比 README 清单更宽。每个网络都配有同名 *_test.py 测试(如 bert_encoder_test.py),可作为行为验证依据。

二、BertEncoder:参数默认值、输入输出与调用链

BertEncoder 是模块中结构最完整、也最常作为模板参考的实现。README 指出它实现了 BERT 论文描述的双向 Transformer 编码器,"It includes the embedding lookups, transformer layers and pooling layer"(包含嵌入查找、Transformer 层与池化层)。bert_encoder.py 的 docstring 进一步说明:它不包含掩码语言模型或分类任务网络,且默认值取自 BERT-Base 实现,网络由 Keras Functional API 构建。

构造参数与默认值

源码 __init__ 签名(bert_encoder.py)给出的 BERT-Base 默认值如下:

参数	默认值	含义
vocab_size	必填	词表大小
hidden_size	768	Transformer 隐层维度,须能被注意力头数整除
num_layers	12	Transformer 层数
num_attention_heads	12	每个 Transformer 的注意力头数
max_sequence_length	512	编码器可消费的最大序列长度,决定位置嵌入的形状
type_vocab_size	16	type_ids 可取的类型数
inner_dim	3072	每个 Transformer 中两层前馈网络第一层 Dense 的输出维度
inner_activation	近似 GELU(gelu(approximate=True))	前馈网络第一层的激活
output_dropout / attention_dropout	0.1 / 0.1	注意力后与输出 dropout、注意力内部 dropout
initializer	TruncatedNormal(stddev=0.02)	全编码器权重初始化器
output_range	None	最后一层输出取 [0, output_range) 切片;None 表示全序列输出
embedding_width	None(即等于 hidden_size)	词嵌入宽度;不等于 hidden_size 时嵌入参数分解为 (vocab_size, embedding_width) 与 (embedding_width, hidden_size) 两个矩阵
embedding_layer	None	可传入自定义嵌入层;缺省时用 OnDeviceEmbedding 构建
norm_first	False	False 表示对注意力和中间稠密层的输出做 LayerNorm
dict_outputs	False	是否以字典形式组织模型输出
return_all_encoder_outputs	False	是否输出所有层编码结果
return_attention_scores	False	是否额外输出各层注意力分数,形状 [batch, heads, seq, seq]
return_word_embeddings	False	是否额外返回输入词嵌入序列

源码还做了向后兼容处理:旧参数名 intermediate_size、activation、dropout_rate、attention_dropout_rate 会被 pop 并映射到新参数(bert_encoder.py),sequence_length 则被标记为废弃。

内部调用链

BertEncoder.__init__ 的构图顺序(见 bert_encoder.py)是:

1. 声明三个 Functional 输入:input_word_ids、input_mask、input_type_ids(均为 int32);
2. OnDeviceEmbedding 查词嵌入,PositionEmbedding 生成位置嵌入,OnDeviceEmbedding(use_one_hot=True) 生成类型嵌入,三者经 tf.keras.layers.Add 相加;
3. 依次做 LayerNormalization(epsilon=1e-12、float32)与 Dropout;
4. 若 embedding_width != hidden_size,插入 EinsumDense('...x,xy->...y') 投影到 hidden_size;
5. SelfAttentionMask() 由 data 与 mask 生成注意力掩码;循环 num_layers 次调用 layers.TransformerEncoderBlock(命名为 transformer/layer_i),仅最后一层可应用 output_range 切片;
6. 对最后一层输出的首 token([:, 0, :],经 Keras 张量子脚本切片产生 SliceOpLambda 层)应用 pooler_transform(Dense + tanh)得到 pooled_output。

最终输出组织(bert_encoder.py):

• dict_outputs=True:输出为字典,含 sequence_output(末层输出)、pooled_output(CLS 池化)、encoder_outputs(逐层输出列表),可选 attention_scores、word_embeddings;
• dict_outputs=False:输出为张量列表 [sequence_output, cls_output](若开启 return_all_encoder_outputs 则首元素为逐层列表),可选追加 attention_scores。

这个"编码器 + 池化输出分离"的设计正是它与 Task 层衔接的接口:SQuAD 任务通常取 sequence_output,序列分类任务通常取 pooled_output。

同一文件中还有 BertEncoderV2(基于 tf.keras.layers.Layer 的子类实现,bert_encoder.py),与 Functional 版结构等价,额外支持 with_dense_inputs(在 call 时接受 dense_inputs/dense_mask/dense_type_ids 并在序列尾部拼接,见 bert_encoder.py)。二者均通过 get_config/from_config 支持 Keras 序列化,并在 from_config 中对"保存了共享嵌入层对象"的模型打印告警。

三、AlbertEncoder:嵌入分解与跨层参数共享

README 对 ALBERT 的概括是:与 BERT 相比,ALBERT 把嵌入参数分解为两个小矩阵,并在各层间共享参数。这两点在 albert_encoder.py 中都能找到直接对应:

嵌入分解。默认 embedding_width=128、hidden_size=768(albert_encoder.py)。词嵌入表宽度只有 128,加上位置/类型嵌入(同为 embedding_width 宽度)相加、LayerNorm、dropout 后,若 embedding_width != hidden_size,用 EinsumDense('...x,xy->...y', name='embedding_projection') 投影到 768(albert_encoder.py)。这与 BertEncoder 的 embedding_width 机制一致——ALBERT 只是把"分解"作为默认配置。

跨层参数共享。核心差异在这一段(albert_encoder.py):

shared_layer = layers.TransformerEncoderBlock(
    num_attention_heads=num_attention_heads,
    inner_dim=intermediate_size,
    inner_activation=activation,
    output_dropout=dropout_rate,
    attention_dropout=attention_dropout_rate,
    kernel_initializer=tf_utils.clone_initializer(initializer),
    name='transformer')
encoder_outputs = []
for _ in range(num_layers):
  data = shared_layer([data, attention_mask])
  encoder_outputs.append(data)

注意与 BertEncoder 对照:后者在循环内每层新建一个 TransformerEncoderBlock(命名为 transformer/layer_i),而 ALBERT 只构造一个名为 transformer 的块并在循环中重复调用同一实例——Functional 图中重复调用同一 Layer 即实现权重共享。这是 ALBERT 参数量大幅下降的直接原因。

其余默认值与 BERT-Base 对齐(num_layers=12、num_attention_heads=12、max_sequence_length=512、intermediate_size=3072、GELU、TruncatedNormal(stddev=0.02));输出同样是 [sequence_output, pooled_output] 或字典形式(含 encoder_outputs 逐层列表)。

四、MobileBERTEncoder:面向端侧的紧凑 BERT

README 将 MobileBERTEncoder 对应到 "MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices"。从 mobile_bert_encoder.py 的默认参数可以看出其"紧凑"体现在哪些维度:

参数	默认值	说明
word_vocab_size	30522	词表大小(BERT 风格词表)
word_embed_size	128	词嵌入宽度(远小于 hidden_size)
type_vocab_size	2	句子类型数
max_sequence_length	512	最大输入长度
num_blocks	24	Transformer 块数量
hidden_size	512	隐层维度
num_attention_heads	4	注意力头数
intermediate_size	512	前馈中间层维度
intermediate_act_fn	'relu'	前馈激活
hidden_dropout_prob / attention_probs_dropout_prob	0.1 / 0.1	dropout
intra_bottleneck_size	128	瓶颈维度
initializer_range	0.02	截断正态初始化标准差
use_bottleneck_attention	False	是否从瓶颈变换取注意力输入(为 True 时忽略下一项)
key_query_shared_bottleneck	True	是否共享 key/query 的线性变换
num_feedforward_networks	4	堆叠前馈网络数
normalization_type	'no_norm'	仅支持 no_norm 与 layer_norm
classifier_activation	False	是否对 [CLS] 表示做 tanh 池化激活
input_mask_dtype	'int32'	input_mask 张量类型

几个值得注意的实现细节:

1. 教师/学生范式在参数里体现:docstring 说明 normalization_type='no_norm' 代表论文中建议的学生模型的逐元素线性变换,'layer_norm' 用于教师模型(mobile_bert_encoder.py);
2. TFLite 量化友好设计:input_mask_dtype 参数说明,若要走不支持 Cast op 的 tf.lite 量化,可将其设为 float32 并直接喂 float32 的 mask,以避开计算图中的 tf.cast(mobile_bert_encoder.py);
3. gin 配置驱动:类上标注 @gin.configurable(mobile_bert_encoder.py),可直接用 gin 配置构造,这与仓库中其它网络的手工传参风格不同;
4. 构图上它使用 layers.MobileBertEmbedding 做嵌入(词嵌入 128 → 投影到 hidden_size),循环 num_blocks 个 layers.MobileBertTransformer 并始终取回 attention_score,最终输出字典包含 sequence_output、pooled_output、encoder_outputs、attention_scores(mobile_bert_encoder.py);classifier_activation=True 时池化层为 EinsumDense('ab,bc->ac', activation=tf.tanh)。

五、FNet 与 SparseMixer:用混合子层替代自注意力

README 将这两个编码器归为"更高效 BERT"方向:

• FNet:"FNet has the same structure as a Transformer encoder, except that all or most of the self-attention sublayers are replaced with Fourier sublayers"(结构与 Transformer 编码器相同,只是全部或大部分自注意力子层被傅里叶子层替换)。
• SparseMixer:"Sparse Mixer consists of layers of heterogeneous encoder blocks. Each encoder block contains a linear mixing or an attention sublayer together with a (dense) MLP or sparsely activated Mixture-of-Experts sublayer"(由异构编码器块组成,每块含线性混合或注意力子层,配稠密 MLP 或稀疏激活 MoE 子层)。

FNet 的关键配置

从 fnet.py 的默认参数看:

• mixing_mechanism 默认为 layers.MixingMechanism.FOURIER;docstring 明确实现默认对应经典 FNet Base 配置,同时支持更一般的混合模型(如 'Linear'、'HNet')以及注意力+混合混用的混合模型(如 'FNet-Hybrid');
• use_fft 默认为 False,即用 DFT 矩阵而非 FFT 计算傅里叶变换(仅对谱混合机制生效,具体取舍建议见 layers.FourierTransformLayer / layers.HartleyTransformLayer);
• attention_layers 默认为空元组,用于声明哪些层保留自注意力;docstring 给出的经验法则是若使用注意力层,建议放在最后几层;
• max_sequence_length=512 是唯一可消费的序列长度("The input length is fixed to 'max_sequence_length'"),因为它同时决定位置嵌入形状与混合矩阵尺寸;
• 其余参数(hidden_size=768、num_layers=12、num_attention_heads=12、inner_dim=3072、近似 GELU、0.1 dropout、TruncatedNormal(0.02))与 BERT-Base 默认值保持一致,便于公平对比。

SparseMixer 的关键配置

SparseMixer 的默认参数直接体现了"异构编码器块"思想:

hidden_size: int = 512,
num_layers: int = 14,
moe_layers: Sequence[int] = (5, 6, 7, 8),
attention_layers: Sequence[int] = (10, 11, 12, 13),
num_experts: int = 16,
train_capacity_factor: float = 1.,
eval_capacity_factor: float = 1.,
mixing_mechanism: layers.MixingMechanism = layers.MixingMechanism.LINEAR,

即默认 14 层中:第 5–8 层用稀疏激活 MoE(16 个专家,专家本身是与常规 MLP 相同结构的模块),第 10–13 层保留注意力,其余层为普通 MLP + 线性混合。两个值得注意的实现约定:

1. Fast Sparse Mixer 变体:docstring 指出将 *_capacity_factor 设为 0.5 即得到更稀疏、更快的 Fast Sparse Mixer,每个专家处理的 token 数约减少 50%;
2. 辅助损失传递:底层 MoeLayer 通过 Keras 的 add_loss() / add_metric() 传播 MoE 辅助损失与指标,使用它的上层模型需要自行收集这些损失(必要时通过 export_metrics 控制是否导出指标)。

同样地,max_sequence_length 是固定输入长度,输入必须与其一致。

六、XLNetBase:排列输入与多重注意力掩码

README 对 XLNetBase 的描述是:实现 "XLNet: Generalized Autoregressive Pretraining for Language Understanding" 中的基础网络,"It includes embedding lookups, relative position encodings, mask computations, segment matrix computations and Transformer XL layers using one or two stream relative self-attention"(含嵌入查找、相对位置编码、mask 计算、segment 矩阵计算,以及使用单流或双流相对自注意力的 Transformer XL 层)。

从 xlnet_base.py 的源码可以印证其复杂度主要在于掩码体系。_create_causal_attention_mask 的 docstring(xlnet_base.py)用一个 S=2、M=1 的示意矩阵解释了单向注意力下,如何在 [(batch dims), S, S+M] 的注意力分数矩阵上屏蔽 j > i 的位置;_combine_masks 支持以 and/or 逻辑组合两个掩码。而 _compute_attention_mask 的注释(xlnet_base.py)列出了 XLNet 中三种掩码:

• 因果注意力掩码:单向注意力时屏蔽当前位置之后的 token;
• 输入掩码:区分真实 token 与 padding(0 表示可注意,1 表示不可注意);
• 排列掩码:XLNet 预训练把输入序列分解为 factorization sequence z,在部分预测中于切分点 c 处切分,i <= c 的位置不可被注意、i > c 的位置可被注意。

该函数将它们广播合并,产出 query 与 content 两个注意力掩码,供 Transformer XL 的双流(one or two stream)相对自注意力使用。文件顶部还定义了 _SEG_ID_CLS = 2 常量(xlnet_base.py),用于 segment 矩阵中 [CLS] 的段类型编号。这些细节共同支撑了 README 中"relative position encodings, mask computations, segment matrix computations"的表述,测试可参见 xlnet_base_test.py。

七、任务头:Classification 与 SpanLabeling

README 把这两个头归入 Networks,因为它们同样是"不会单独训练"的 Keras 模型,直接挂在编码器输出之上。

Classification:分类与回归合一的单层头

Classification 的实现就是一个对 cls_output(shape (input_width,))做 Dense 变换的头:

• num_classes=1 时视为回归问题;
• output 取值 logits 或 predictions:后者在 logits 上追加 tf.nn.log_softmax 激活;源码特别处理了混合精度场景——当全局策略为 mixed_bfloat16 时,回退到 float32 做 log-softmax,因为 bf16 对 softmax 后交叉熵不稳定(classification.py);
• 初始化器默认 glorot_uniform;
• 需要留意:该类已被标记弃用,docstring 建议改用 layers.ClassificationHead(classification.py)。

SpanLabeling:面向 SQuAD 的单跨度标注头

SpanLabeling 的构造非常精炼:对 sequence_data(shape (batch, seq_len, input_width))施加一个 2 单元的 Dense 层预测 start 与 end 位置,再经 _split_output_tensor 拆成两路,分别做 tf.nn.log_softmax。output='logits' 输出 [start_logits, end_logits],output='predictions' 输出对数概率预测。

同文件还包含 XLNetSpanLabeling(XLNet 在 SQuAD 2.0 上用的复杂版,span_labeling.py):先稠密预测 start 位置,训练时用真实 start、推理时用束搜索(beam search)预测 end;默认 start_n_top=5、end_n_top=5、activation='tanh';call 的返回字典含 start_predictions、end_predictions、start_logits、end_logits,推理时追加 start_top_predictions、start_top_index、end_top_predictions、end_top_index,并额外输出 class_logits(是否可答的分类),其实现基于 CLS 表示与 start 概率加权的表示拼接(span_labeling.py)。其 docstring 也明确提示:compute_with_beam_search 不能用于 Functional API,所以它被实现为 tf.keras.layers.Layer 而非 Functional Model。

八、PackedSequenceEmbedding:支持打包序列的嵌入网络

README 的一句话是:"implements an embedding network that supports packed sequences and position ids"。结合 packed_sequence_embedding.py 的 docstring 与源码,其能力可拆解为:

1. 打包多条子序列(pack_multiple_sequences=True):多条序列可拼进一条长序列进行训练/推理且互不影响。实现上由 PackedSequenceMask 层完成——它假设父序列首 token 为 [CLS]、每条子序列以一个 [CLS] 开头且只含一个 [CLS],通过 cumsum 给不同子序列编号,再两两比较相等性,得到形状 [batch, seq, seq] 的子序列内注意力掩码;该掩码随后乘入 SelfAttentionMask 的输出(packed_sequence_embedding.py);
2. 可选显式 position ids(use_position_id=True):此时模型多一个 position_ids 输入;为 False 时位置 id 自动推断——不打包时为 0..seq_length-1,打包时每段子序列各自从 0 开始,由 PositionEmbeddingWithSubSeqMask 层配合子序列掩码完成(packed_sequence_embedding.py);
3. 嵌入组合流程与 BERT 一致:词嵌入 + 位置嵌入 + 类型嵌入相加 → LayerNorm(epsilon=1e-12)→ Dropout,embedding_width != hidden_size 时经 EinsumDense 投影(注意此处分量轴 bias_axes=None,不带偏置);
4. 输出为 [embeddings, attention_mask] 两个张量(packed_sequence_embedding.py)——即它只到"嵌入网络"为止,后续 Transformer 堆叠由外部完成,这正是它与 BertEncoder 的差异。

九、使用方式:组合、序列化与验证

综合以上源码,使用这些 Network 的标准姿势是:

1. 组合编码器与任务头。 以 Functional 风格为例,BertEncoder 与 Classification 的输出可以直接作为另一层模型输入:编码器输出 [sequence_output, pooled_output](dict_outputs=True 时为字典),任务头以 input_width 对齐 pooled_output 的最后维度(默认 hidden_size)。各网络的输入均为 (input_word_ids, input_mask, input_type_ids) 三件套,与仓库数据管线(见 official/nlp/data)产出的特征键一致。

2. 依赖标准化配置做保存/恢复。 所有网络都实现了 get_config() / from_config()(如 bert_encoder.py、span_labeling.py),并通过 @tf_keras.utils.register_keras_serializable(package='Text') 注册,因此可以在 Keras 序列化体系中按名称重建;BertEncoder.from_config 还会对含共享嵌入层对象的存档打印告警,提示重新训练时嵌入不再共享。

3. 用测试作为行为基准。 每个网络都有对应测试文件,例如 bert_encoder_test.py 验证 BERT 编码器的输出结构与数值,albert_encoder_test.py、fnet_test.py、sparse_mixer_test.py、span_labeling_test.py 等覆盖其余网络,修改或参照实现时可作为回归依据。

4. 选择合适的编码器。 选择逻辑可以直接映射到源码默认值:标准 BERT 用 BertEncoder(BERT-Base 默认,embedding_width 可选分解);参数敏感、追求小模型用 AlbertEncoder(128 宽嵌入 + 层共享);端侧部署用 MobileBERTEncoder(512 隐层、4 头、bottleneck 结构,注意 input_mask_dtype 与 TFLite 量化约束);长上下文/预训练用 XLNetBase;追求推理效率、可接受固定 512 输入长度时用 FNet(傅里叶混合)或 SparseMixer(MoE + 线性混合,注意收集 MoE 辅助损失)。

十、小结

official/nlp/modeling/networks 是 models 仓库 NLP 建模的"中间层":它把论文中的编码器结构固化为带 BERT-Base 级默认值的 Keras 模型,统一了输入(input_word_ids/input_mask/input_type_ids)与输出(sequence_output/pooled_output/encoder_outputs)约定,并通过 get_config/from_config 保证可序列化。对使用者,理解 BertEncoder 的参数表(第二节)是基础,其余编码器则是围绕"嵌入分解"(AlbertEncoder、MobileBERTEncoder)、"注意力替代"(FNet、SparseMixer)与"生成式预训练掩码"(XLNetBase)三条主线对它的变体;分类头与跨度标注头则演示了 Network 与 Task 的衔接方式。所有结论均可在 official/nlp/modeling/networks/ 目录下逐文件对照验证。