TensorFlow Models NLP 建模层库详解official/nlp/modeling/layers 中 Transformer 与注意力机制的源码级剖析【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models本文围绕 Layers 模块文档 展开系统讲解该仓库 NLP 建模层库中 20 余个核心 Keras 层——多头注意力、稀疏/线性注意力、带缓存的解码器、掩码与位置编码、任务头MaskedLM、分类头等——的设计意图与源码实现。读完后你将能够直接选用、组合这些层构建新的tf.kerasNLP 模型并理解BertEncoder、AlbertEncoder等上层网络是如何调用这些积木搭出完整 Transformer 的。1. LayersNLP 模型的基本构建块模块入口文档的定义非常直接见 READMELayers are the fundamental building blocks for NLP models. They can be used to assemble newtf.keraslayers or models.也就是说这一层级的库不做端到端模型而是把注意力机制、掩码、位置编码、任务头、分词打包输入等能力拆成可复用的tf_keras.layers.Layer供上层网络encoder/network自由组装。模块统一导出集中在 official/nlp/modeling/layers/init.py其中可以看到大量层使用tf_keras.utils.register_keras_serializable(packageText)注册使其可参与 Keras 模型序列化/保存导出列表比 README 覆盖的条目更宽还包括MoeLayer混合专家、FactorizedEmbedding、PackBertEmbeddings、TNTransformerExpandCondense、TransformerScaffold、PerDimScaleAttention、MultiQueryAttention等进阶实现均位于 layers 目录 下。下文的组织方式与 README 的条目一一对应并按注意力 → 块结构 → 掩码/Softmax → 位置编码 → 任务头 → 文本前处理展开。2. 核心注意力层2.1 MultiHeadAttention 与 CachedAttentionREADME 指出 MultiHeadAttention 实现了可选掩码的 query/key/value 注意力Attention Is All You Need当from_tensor与to_tensor相同时即自注意力。从源码看attention.py 中的实现是一行别名MultiHeadAttention tf_keras.layers.MultiHeadAttention即直接复用 Keras 内置多头注意力本仓库真正新增的是同文件中的CachedAttention——一个继承自tf_keras.layers.MultiHeadAttention的带缓存注意力层专门服务自回归解码普通路径decode_loop_step is None_update_cache用tf.concat把历史 key/value 与新 key/value 沿序列维拼接再写回cache[key]、cache[value]返回全长度张量参与后续 einsumTPU 特化路径传入decode_loop_step缓存张量形状固定无法动态拼接于是用tf.one_hot(decode_loop_step, key_seq_dim)构造单步索引把当前步的 key/value 以加权方式叠加进固定形状的缓存避免 shape 变化导致 TPU 重编译call(query, value, keyNone, attention_maskNone, cacheNone, decode_loop_stepNone, ...)返回(attention_output, cache)请求return_attention_scoresTrue时再附带attention_scores。这一cache decode_loop_step接口正是后文TransformerDecoderBlock逐 token 生成的底层支撑。2.2 TalkingHeadsAttention 与 MultiChannelAttentionTalkingHeadsAttentionTalking-Heads Attention 论文变体在 softmax 前/后对各头之间做线性交谈打破多头之间完全独立的假设MultiChannelAttention多头注意力的一种变体可把多路流合并来做交叉注意力。从 transformer.py 的TransformerDecoderBlock可以看到它的典型用法——构造参数multi_channel_cross_attentionTrue时encoder-decoder 交叉注意力默认替换为MultiChannelAttention且call的输入需要额外第 5 个张量doc-attention 概率见源码中对inputs[-1]的解包。2.3 BigBirdAttention把二次复杂度降为线性BigBirdAttention 实现 Big Bird 论文的稀疏注意力。源码中与注意力本体配合的掩码构造函数最能体现其局部窗口 全局 token 随机块的组合结构create_band_mask_from_inputs(from_blocked_mask, to_blocked_mask)由分块 2D 掩码生成局部窗口band的 3D 注意力掩码形状为[batch, 1, L/b - 4, b, 3*to_block_size]即每个 query 块只看邻近 3 个 key 块bigbird_block_rand_mask(...)按行生成随机块邻接表last_idx可限制随机块只能选在序列前缀内保证因果性create_rand_mask_from_inputs(...)把随机块索引展开成逐头 3D 掩码供注意力打分阶段屏蔽。文件顶部还定义了MAX_SEQ_LEN 4096即该实现面向的序列长度上限。模块同时导出BigBirdMasksinit.py 第 24 行用于封装上述三类掩码的生成。2.4 KernelAttention核特征图的线性注意力KernelAttention 把自注意力表示为核特征图的线性点积利用矩阵乘法结合律把复杂度从 O(n²) 降到 O(n)README 说明其涵盖了 Linear Attention、Performer、Random Feature Attention 三类方法。从 kernel_attention.py 源码看实现依赖一组窗口/分块工具函数pad_to_chunk_length/split_tensor_into_chunks把序列切成长度可整除的 chunk配合 TPU 分块计算rectangular_window_sum用前缀和差值在滑动矩形窗口上做求和O(n) 实现局部平滑weighted_window_sum用tf.nn.depthwise_conv2d实现加权滑动窗等价于对核特征做因果卷积式归一化配套的 KernelMask 把普通 2D 输入掩码[batch, seq]转成KernelAttention需要的掩码格式。2.5 ReuseMultiHeadAttention 与 ReuseTransformerReuse Transformer 一族的 TF 实现位于 reuse_attention.py 与 reuse_transformer.py。其思想是相邻层注意力分数高度冗余高层可复用小一层的分数省去部分点积计算。源码要点reuse_attention参数0不复用-1表示全部头复用源码中会归一化为num_heads中间值表示部分头复用且构造时会校验取值在[-1, num_heads]内当复用头数 总头数时_build_from_signature把 query/key/value/输出投影拆成value_reuse与value_new两组EinsumDense_compute_attention中先对新头正常计算new_scores可选叠加相对位置偏置再把传入的reuse_scores[:, :reuse_heads, :, :]与之沿头维tf.concat全部复用时reuse_heads num_heads则直接使用reuse_scores、完全跳过 Q/K 投影与点积相对位置偏置use_relative_peTrue且非全复用时额外创建形状[1, num_heads - reuse_heads, 2*pe_max_seq_length - 1]的relative_position_embeddings变量pe_max_seq_length默认 512_compute_relative_position通过索引矩阵查表得到逐头偏置dtype 跟随 Keras 混合精度全局策略mixed_bfloat16/mixed_float16/float32。2.6 ReZeroTransformerReZeroTransformer 在标准 Transformer 块的残差上引入逐元素可学习标量门控ReZero 论文用于缓解深层 Transformer 的收敛问题。它与TransformerEncoderBlock结构同构仅残差路径不同可对照阅读 rezero_transformer.py 及其测试 rezero_transformer_test.py。2.7 相对位置注意力与 TransformerXLMultiHeadRelativeAttentionTransformer-XL 的相对位置编码注意力变体且源码层面扩展支持了 XLNet 提出的基于 segment 的注意力偏置TwoStreamRelativeAttentionXLNet 的双流query 流 content 流相对自注意力TransformerXL包含TransformerXLBlock一个或双流相对自注意力 前馈网络与TransformerXL管理 attention bias 及堆叠多个 block两个类二者均从init.py 第 76-77 行导出。2.8 MobileBert 专用层MobileBertEmbedding 与 MobileBertTransformer 分别实现 MobileBERT 论文提出的轻量嵌入层与 Transformer 层。init.py 同时还导出了同文件中的MobileBertMaskedLM说明该目录对 MobileBERT 提供了嵌入 块 MLM 头的成套组件。3. Transformer 块Transformer、TransformerDecoderBlock 与 TransformerEncoderBlock3.1 Transformer已弃用的别名层Transformer 直接继承TransformerEncoderBlock并转发全部参数intermediate_size→inner_dim、dropout_rate→output_dropout等但构造时会发出明确的弃用告警TheTransformerlayer is deprecated. Please directly useTransformerEncoderBlock.transformer.py同文件还定义了CompiledTransformer在call上叠加tf_function_if_eager(experimental_compileTrue)来自 util.py用于 TF Function 自动编译加速的场景。阅读旧配置gin 配置里常见Transformer键名时需要知道它等价于TransformerEncoderBlock。3.2 TransformerEncoderBlock当前推荐的标准块TransformerEncoderBlock 是多头注意力 两层前馈的标准编码器块实现也是本仓库 encoder 网络的默认积木。构造参数源码签名transformer_encoder_block.py参数默认值说明num_attention_heads必填注意力头数inner_dim必填前馈中间层宽度inner_activation必填前馈激活函数output_rangeNone对输入序列取[0, output_range)切片None表示不切norm_firstFalseFalse为 Post-LN对块输出归一化True为 Pre-LN对输入归一化norm_epsilon1e-12LayerNorm/RMSNorm 的 epsilonuse_rms_normFalse用同文件定义的RMSNorm替代 LayerNormoutput_dropout/attention_dropout/inner_dropout0.0三处独立 dropoutnum_kv_headsNone指定 KV 头数Multi-Query/GQA 风格linformer_dimNone低秩线性注意力投影维度use_sigmoid_attn/sigmoid_attn_biasFalse/Nonesigmoid 注意力开关与偏置return_attention_scoresFalse是否额外返回注意力分数同文件的RMSNorm实现也很简洁对输入先转 float32计算平方均值后inputs * rsqrt(var epsilon) * scale再转回原 dtype且scale权重关闭 autocast 以保持精度。上层网络的真实调用可佐证其地位bert_encoder.py、bert_encoder.py、albert_encoder.pyALBERT 的参数共享层即复用同一个TransformerEncoderBlock实例、seq2seq_transformer.py 都在实例化该块。3.3 TransformerDecoderBlockTransformerDecoderBlock 是解码器单层由三个子层构成源码 docstring 与build一致自注意力默认类为attention.CachedAttention可经self_attention_cls替换因此天然带cache机制encoder-decoder 交叉注意力默认MultiHeadAttentionmulti_channel_cross_attentionTrue时换成MultiChannelAttention也可用cross_attention_cls显式指定位置前馈网络EinsumDense(abc,cd-abd)→ 激活 → dropout → 输出投影。call(inputs, cacheNone, decode_loop_stepNone)中inputs为四元组(input_tensor, memory, attention_mask, self_attention_mask)多通道交叉注意力时为五元组源码按norm_first分支在 Pre-LN/Post-LN 两种残差排布间切换最终返回(layer_output, cache)——把缓存回传给下一层从而支撑整段解码循环。build中还有一个实用约束输入必须为三维[batch, sequence, width]且width % num_attention_heads 0否则抛出ValueError。4. 掩码与 Softmax4.1 SelfAttentionMaskSelfAttentionMask 从 2D 掩码生成 3D 自注意力掩码。实现原理get_mask函数输入inputs形状[batch, from_seq_length, ...]to_mask形状[batch, to_seq_length]int321 表示有效、0 表示需屏蔽先把to_maskreshape 为[batch, 1, to_seq_length]再用tf.broadcast_to广播到[batch, from_seq_length, to_seq_length]。这避免了显式构造平方张量内存开销仅为广播视图。该文件同时提供独立的get_mask(inputs, to_mask, dtypeNone)函数供非 Layer 场景直接调用。4.2 MaskedSoftmaxMaskedSoftmax 实现带可选掩码的 softmaxREADME 中1 表示放行、0 表示屏蔽被屏蔽位置输出近似为 0的语义在源码中对应一个关键细节adder (1.0 - tf.cast(mask, scores.dtype)) * _large_compatible_negative(scores.dtype) scores adder其中_large_compatible_negative对 float32 返回-1e9而 float16 无法表示-1e9于是返回tf.float16.min——这是半精度训练下防止负得不够大导致泄漏的实现要点。另外mask_expansion_axes当掩码比分数张量维度少时循环tf.expand_dims到指定轴使[B, T, S]的掩码能适配[B, H, T, S]的分数normalization_axes默认(-1,)多轴归一化时改用exp(scores - reduce_logsumexp(...))的数值稳定写法而不是tf.nn.softmax后者只支持单轴。5. 位置编码PositionEmbedding 按 BERT 论文方式创建可学习的位置嵌入。构造参数max_length必填动态序列最大长度、initializerglorot_uniform、seq_axis1在哪个轴上加嵌入。其文档示例即完整可运行片段position_embedding PositionEmbedding(max_length100) inputs tf_keras.Input((100, 32), dtypetf.float32) outputs position_embedding(inputs)PositionEmbedding 的build按max_length创建嵌入表、按输入末维确定宽度同文件还导出RelativePositionBias与RelativePositionEmbedding见init.py分别服务相对位置偏置与相对位置嵌入两条路线。6. 任务头与不确定性建模6.1 MaskedLMMaskedLM 是 BERT 的掩码语言模型头README 强调它假设外部传入嵌入表变量。源码印证了这一约定构造参数embedding_table必须来自 encoder 的get_embedding_table()、activation、initializerglorot_uniform、outputlogits | predictions非法取值直接抛ValueErrorbuild中从嵌入表形状(vocab_size, hidden_size)反推维度依次创建Dense(hidden_size)transform/dense、LayerNormalization(epsilon1e-12)transform/LayerNorm与形状[vocab_size]的output_bias/biascall(sequence_data, masked_positions)先_gather_indexes取出被掩码位置再投影、归一化、与词表嵌入做点积输出 logits。文档给出的最小用法encoder modeling.networks.BertEncoder(...) lm_layer MaskedLM(embedding_tableencoder.get_embedding_table())6.2 ClassificationHead 与 GaussianProcessClassificationHeadClassificationHead 是在嵌入序列上做池化的分类头。源码参数inner_dim0/None 时只建输出投影、num_classes、cls_token_idx0在序列第几位做池化通常取 [CLS] 位、activationtanh、dropout_rate0.0内部结构为pooler_dense可选→ Dropout →logits输出层。call(features, only_projectFalse)支持只取池化向量不接分类投影。同文件的GaussianProcessClassificationHead是 SNGP光谱归一化神经网络高斯过程分类头依赖 gaussian_process.py 的 RandomFeatureGaussianProcess随机特征 GP见 Random Features for Large-Scale Kernel Machines与 spectral_normalization.py 的 SpectralNormalizationtf.Wrapper对内部层应用谱范数正则三者组合实现了距离感知的不确定性估计分类头。6.3 其他头/正则层MatMulWithMargin带 margin 的矩阵乘法层用于检索/排序类任务ADD 模型的双编码器加性 margin softmaxGatedFeedforwardGLU 变体门控前馈层GLU Variants Improve TransformerOnDeviceEmbedding为 TPU 模型设计的高效嵌入查找层适合词表极大的场景。7. 文本前处理层从原始文本到 BERT 输入text_layers.py 提供把原始文本 → 模型输入整条流水线 Keras 层化的能力README 列出的三个类及init.py 补充导出的FastWordpieceBertTokenizer均在其中BertTokenizerWordPiece 分词 特殊 token 处理SentencepieceTokenizerSentencePiece 分词配合 train_sentencepiece.py 训练出的模型文件BertPackInputs把分词结果按 BERT 的input_word_ids / input_mask / segment_ids约定打包、padding 成模型输入张量。这使得数据管道可以不落地 TFRecord而是在图中直接完成分词与打包。8. 组合使用从 README 条目到真实模型把 README 的条目按积木分类汇总如下全部文件路径均位于 official/nlp/modeling/layers/类别层实现文件基础注意力MultiHeadAttentionattention.py解码缓存注意力CachedAttentionattention.py头间交互TalkingHeadsAttention、MultiChannelAttentiontalking_heads_attention.py、multi_channel_attention.py长序列稀疏注意力BigBirdAttention、BigBirdMasksbigbird_attention.py线性注意力KernelAttention、KernelMaskkernel_attention.py注意力分数复用ReuseMultiHeadAttention、ReuseTransformerreuse_attention.py、reuse_transformer.py深层收敛ReZeroTransformerrezero_transformer.py相对位置/长上下文MultiHeadRelativeAttention、TwoStreamRelativeAttention、TransformerXL、TransformerXLBlockrelative_attention.py、transformer_xl.py轻量 BERTMobileBertEmbedding、MobileBertTransformer、MobileBertMaskedLMmobile_bert_layers.pyTransformer 块Transformer弃用、CompiledTransformer、TransformerDecoderBlock、TransformerEncoderBlocktransformer.py、transformer_encoder_block.py掩码/SoftmaxSelfAttentionMask、MaskedSoftmaxself_attention_mask.py、masked_softmax.py位置编码PositionEmbedding、RelativePositionBias、RelativePositionEmbeddingposition_embedding.py任务头MaskedLM、ClassificationHead、GaussianProcessClassificationHeadmasked_lm.py、cls_head.py不确定性/正则RandomFeatureGaussianProcess、SpectralNormalizationgaussian_process.py、spectral_normalization.py前馈/门控GatedFeedforwardgated_feedforward.py排序MatMulWithMarginmat_mul_with_margin.pyTPU 嵌入OnDeviceEmbeddingon_device_embedding.py文本前处理BertTokenizer、SentencepieceTokenizer、FastWordpieceBertTokenizer、BertPackInputstext_layers.py一个典型组合对应仓库内 bert_encoder.py 的结构BertTokenizer/BertPackInputs产输入 → 词嵌入 PositionEmbedding segment 嵌入 → 堆叠 N 个TransformerEncoderBlock(num_attention_heads..., inner_dim..., inner_activationgelu)→MaskedLM预训练或ClassificationHead微调。每个层都带get_config/from_config如 reuse_attention.py 中保存query/key/value形状以便from_config触发重建因此整条模型可以随 Keras 序列化保存与加载。9. 小结official/nlp/modeling/layers/ 是仓库 NLP 建模的最小抽象层注意力、块、掩码、位置、任务头各自成文件、成层测试一一对应如 attention_test.py、transformer_test.py便于单独验证与替换面向性能演进提供了成体系的替代件BigBirdAttention稀疏、KernelAttention线性、ReuseMultiHeadAttention跨层复用、OnDeviceEmbeddingTPU面向生产与序列化层均注册为 Keras 可序列化对象TransformerEncoderBlock是现行推荐块旧Transformer已弃用、CompiledTransformer提供自动编译加速。理解本文内容后你可以在 official/nlp/modeling/networks/ 与 official/nlp/modeling/models/ 中对照实际 encoder/模型实现按需替换注意力或头结构而不必重写整个训练管线。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
