尧图网页设计 
TF Model Garden 中的 Roformer旋转位置编码 Transformer 的源码解析与预训练实战指南【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/modelsRoformerRotary Position Embedding Former是 TF Model Gardenofficial/目录下 TensorFlow 官方模型库中的 NLP 实验项目用 Keras 函数式 API 重新实现了“旋转位置编码”这一 Transformer 位置表示方案并完整打通了预训练与 GLUE 微调链路。本文基于 official/projects/roformer/README.md 给出的训练命令展开结合official/projects/roformer/下的模型源码、实验配置与测试用例讲清 Roformer 编码器结构、旋转位置编码的数学实现、预训练配置参数与数据格式读完你可以直接在该仓库内复现 Roformer 的 TPU 预训练并理解其与 BERT 的关键差异。一、项目结构与定位official/projects/roformer/是一个自包含的模型项目文件职责如下文件职责roformer.pyRoformerEncoderConfig与get_encoder工厂函数把配置绑定为RoformerEncoderroformer_encoder.pyRoformerEncoder编码器网络Keras Functional API 构建roformer_encoder_block.pyRoformerEncoderBlock单层 Transformer 块注意力 FFN LayerNorm 残差roformer_attention.pyRoformerAttention多头注意力层核心是旋转位置向量_build_trig_vectorroformer_experiments.py注册roformer/pretraining与roformer/glue两个实验配置train.py训练入口脚本experiments/roformer_base.yamlRoformer base 规模的完整实验配置roformer_attention_test.py、roformer_encoder_test.py、roformer_encoder_block_test.py注意力、编码器、编码器块的单元测试模型实现遵循 TF Model Garden 的标准范式通过official.core.exp_factory注册实验名如roformer/pretraining由统一的 train.py 驱动训练模型本体全部用tf_kerasKeras 3构建并注册为可序列化对象tf_keras.utils.register_keras_serializable(packageText)。二、README 训练命令逐行拆解README.md 给出的运行方式是DATA_PATH??? OUTPUT_DIR??? python3 train.py \ --experimentroformer/pretraining \ --config_fileexperiments/roformer_base.yaml \ --params_overridetask.validation_data.input_path${DATA_PATH},runtime.distribution_strategytpu \ --tpulocal \ --model_dir${OUTPUT_DIR} \ --modetrain_and_eval各参数的含义可对照 train.py 与 official/common/flags.py 中define_flags()的定义--experimentroformer/pretraining实验名必须与 roformer_experiments.py 中exp_factory.register_config_factory(roformer/pretraining)注册的工厂名一致train_utils.parse_configuration(FLAGS)会据此加载默认实验配置--config_fileexperiments/roformer_base.yaml在实验默认配置之上叠加的 YAML 配置相对train.py所在目录即roformer_base规模--params_overridetask.validation_data.input_path${DATA_PATH},runtime.distribution_strategytpu以点号.参数名值的形式做最终覆盖。这里把验证集数据路径指向DATA_PATH并显式将分布式策略设为 TPU与--tpulocal配合使用本地 TPU--tpulocal指定本地 TPU 地址由 official/common/distribute_utils.py 的get_distribution_strategy据此创建TPUStrategy--model_dir${OUTPUT_DIR}checkpoint 与 summary 输出目录。train.py 中可见只有mode包含train时才会serialize_config写出 YAML 配置文件避免纯 eval 任务与训练任务并发写同一文件产生竞争--modetrain_and_eval边训练边验证最终交给train_lib.run_experiment执行训练结束后train_utils.save_gin_config还会把 gin 配置落盘存档。另外train.py 支持runtime.mixed_precision_dtype开启混合精度mixed_float16/mixed_bfloat16TPU 场景通常配合 bfloat16 使用。三、预训练实验配置roformer_base.yaml 全参数解读experiments/roformer_base.yaml 对应“base”规格12 层 × 768 维与 BERT-base 同量级。关键参数分组如下模型部分task.model参数值说明encoder.typeany编码器类型选择器any下由roformer.RoformerEncoderConfig构建RoformerEncoderhidden_size768Transformer 隐层维度必须是num_attention_heads的整数倍见 roformer_encoder_block.py 的校验num_layers12Transformer 层数num_attention_heads12每层注意力头数单头维度 768/1264Roformer 要求该值为偶数见 roformer_attention.py 的assert self._key_dim % 2 0intermediate_size3072FFN 第一层维度且必须为偶数roformer_encoder_block.py 会直接抛错代码中通过 kwargs 别名映射到inner_dimhidden_activationgeluFFN 激活函数max_position_embeddings512决定旋转位置向量的预构建长度见下文vocab_size/type_vocab_size30522 / 2词表与 segment 类型数embedding_size768词嵌入宽度与 hidden_size 相等时省去投影层attention_dropout_rate/dropout_rate0.1注意力 dropout 与层间 dropoutinitializer_range0.02TruncatedNormal初始化标准差roformer.pycls_headsnext_sentence2 分类tanh下一句预测头数据部分task.train_data/task.validation_dataseq_length512、max_predictions_per_seq76、use_next_sentence_labeltrue、global_batch_size256、drop_remainder训练 true / 验证 false。注意 Roformer 实验在 roformer_experiments.py 中显式设置了use_v2_feature_namesTrue其含义见第五节。训练器部分trainer参数值说明optimizer_config.optimizeradamw默认weight_decay_rate0.01LayerNorm/layer_norm/bias不衰减roformer_experiments.py学习率polynomial初值 1e-4decay_steps1000000end_learning_rate0power1.0多项式线性衰减warmuppolynomialwarmup_steps100001 万步线性预热train_steps1000000总训练步数checkpoint_interval/validation_interval20000 / 1000checkpoint 与验证频率validation_steps64steps_per_loop/summary_interval50 / 50分布式循环步长与 summary 频率max_to_keep5最多保留 checkpoint 数roformer_experiments.py中还注册了第二个实验roformer/glue源码任务类型为sentence_prediction.SentencePredictionConfig学习率降至 3e-5、同样采用 AdamW 多项式 warmup用于在预训练权重上做 GLUE 句级分类微调。README 只写了 pretraining 命令但将--experiment换成roformer/glue并覆盖task.train_data.input_path等字段即可走微调流程。四、编码器结构与 BERT 的关键差异RoformerEncoder 用 Keras Functional API 构建输入为三个张量input_word_ids、input_mask、input_type_ids输出为sequence_output末层全部 token 表示、pooled_output[CLS] 位经pooler_transform的 tanh Dense 层以及每一层的encoder_outputs。与 BERT 编码器相比源码中有两处值得注意的“减法”没有绝对位置嵌入。roformer_encoder.py 中明确注释 “Roformer does not need a position embedding layer / does not have absolute position embedding”嵌入只是词嵌入 类型嵌入的相加再 LayerNorm Dropout_position_embedding_layer被显式置为NoneL238。位置信息完全由注意力层的旋转位置编码注入位置信息的载体是预构建的正弦/余弦向量每层RoformerEncoderBlock在构建RoformerAttention时传入q_max_sequence_length与kv_max_sequence_length均取自配置里的max_position_embeddings512三角向量在层初始化时一次性算好前向传播不再需要额外的位置嵌入查表。编码器其余部分保持经典 Transformer 结构若embedding_width ! hidden_size会插入一个embedding_projection的EinsumDense投影12 层RoformerEncoderBlock依次堆叠最后一层支持output_range切片以截取部分 token 输出。五、核心机制旋转位置编码Rotary Position EmbeddingRoformer 的技术核心在 roformer_attention.py。RoformerAttention继承自tf_keras.layers.MultiHeadAttention重写了两处逻辑。1. 三角向量构建_build_trig_vectorL23-L38def _build_trig_vector(length, key_dim): Builds the trig vector. tf_dtype tf_keras.mixed_precision.global_policy().compute_dtype position_ids tf.cast(tf.range(length), dtypetf_dtype) position_ids tf.expand_dims(position_ids, axis0) steps key_dim // 2 # 2 (i - 1) / key_dim (i - 1) / steps: (-1 achieved with zero-indexing) wavenumber_exponent -tf.cast(tf.range(steps), dtypetf_dtype) / steps wavenumbers tf.pow( tf.constant(10000.0, dtypetf_dtype), wavenumber_exponent ) vec tf.einsum(bl,d-bld, position_ids, wavenumbers) sin_vec tf.repeat(tf.sin(vec), repeats2, axis-1) cos_vec tf.repeat(tf.cos(vec), repeats2, axis-1) sin_vec, cos_vec tf.expand_dims(sin_vec, 2), tf.expand_dims(cos_vec, 2) return sin_vec, cos_vec其数学含义对位置 m、频率分量 i角度为m * 10000^(-2i/key_dim)即标准 sinusoidal 频率谱基频 1最高分量频率 1/10000与 BERT 位置嵌入使用同一频率公式但 Roformer 不在嵌入上加它而是作用于注意力投影之后的 Q/K 向量上。sin/cos_vec在最后一维做repeat(repeats2)是把每个频率分量“铺”到相邻的 (2i, 2i1) 两个维度上为后续的成对旋转做准备。2. Q/K 的旋转roformer_recompute_qkvL77-L91def roformer_recompute_qkv(self, q, k, v): ... q2 tf.stack([-q[..., 1::2], q[..., ::2]], axis4) q2 tf.reshape(q2, q_shape) k2 tf.stack([-k[..., 1::2], k[..., ::2]], axis4) k2 tf.reshape(k2, k_shape) ret_q q * self.q_cos_vec[:, 0:q_len, ...] q2 * self.q_sin_vec[:, 0:q_len, ...] ret_w k * self.k_cos_vec[:, 0:k_len, ...] k2 * self.k_sin_vec[:, 0:k_len, ...] return ret_q, ret_w, v实现上把 Q或 K的奇偶维度拆成两半q[..., ::2]、q[..., 1::2]构造旋转算子q2 [-q_odd, q_even]然后执行q·cos(mθ) q2·sin(mθ)——这正是二维平面旋转矩阵(cosθ, -sinθ; sinθ, cosθ)的成对向量形式每相邻两个维度构成一个旋转平面位置 m 决定旋转角mθ_i。V 向量不旋转直接返回。在call中L93-L117流程是线性投影得到 Q/K/V →roformer_recompute_qkv对 Q、K 施加位置旋转 → 复用父类MultiHeadAttention._compute_attention做标准缩放点积注意力含 mask→ 输出投影。由于 Q、K 都带位置旋转q_i·k_j的内积只依赖相对位置i-j这是旋转位置编码“相对性”的由来。测试印证roformer_attention_test.py 的test_trig_vector对length[8,50]、key_dim[64,128]的组合用tf.assert_equal逐位核对sin_emb/cos_emb与10000^(-i/half_d)公式的直接计算结果并校验奇偶维度重复一致test_attention_scores则验证带/不带 mask 时输出形状为[batch, seq_len, key_dim]。六、编码器块RoformerEncoderBlockRoformerEncoderBlock 是标准“注意力 FFN”块几个实现细节偶数维度硬约束__init__中要求inner_dim为偶数L87-L89build中要求hidden_size能被头数整除并计算单头维度attention_head_sizeL132-L136支持 pre-norm / post-norm 两种结构norm_first参数callL243-L310中 pre-norm 路径是“先 LayerNorm 再残差相加”post-norm 路径是“残差相加后再 LayerNorm”数值稳定性处理两个 LayerNorm 固定用dtypetf.float32L157-L164 注释说明为数值稳定混合精度下也未验证 fp16 安全FFN 激活层在mixed_bfloat16全局策略下回退到 float32 计算L172-L179输入形态灵活支持(input)、(input, attention_mask)、(query, key_value, attention_mask)三种输入后者为交叉注意力场景保留output_range切片设置后只对前output_range个 token 输出编码器中仅最后一层启用对应RoformerAttention构造时对q_sin_vec/q_cos_vec的同步切片roformer_attention.py。七、预训练数据格式与准备预训练任务由official.nlp.tasks.masked_lm.MaskedLMConfig驱动数据加载器是 official/nlp/data/pretrain_dataloader.py 的BertPretrainDataLoaderfile_type为tfrecord。Roformer 实验设置了use_v2_feature_namesTrue这决定了 TFRecord 中tf.Example的字段名use_v2_feature_namesTrue时pretrain_dataloader.py字段为input_word_ids、input_type_ids、input_mask、masked_lm_positions、masked_lm_ids、masked_lm_weights、next_sentence_labels为False时则使用历史 BERT 命名input_ids、segment_ids解码时所有 int64 特征会被 cast 成 int32TPU 不支持 int64见 pretrain_dataloader.py 注释因此 TFRecord 按int64写出即可。数据准备可参考仓库内的标准脚本 official/nlp/data/create_pretraining_data.py其命令行 flaguse_v2_feature_namesL80以随机掩码 15% 的方式生成 masked LM 样本。README 命令中的DATA_PATH即指向按上述格式v2 字段名、seq_length512、max_predictions_per_seq76切分好的 TFRecord 目录通过--params_override覆盖task.validation_data.input_path训练集路径同理可通过task.train_data.input_path覆盖。数据规格细节另见 official/nlp/docs/data_processing.md 与 official/nlp/data/README.md。八、从源码结构看可推断的适用边界从源码结构看Roformer 实现面向 BERT 式双目标预训练masked LM next sentence predictionnext_sentence分类头在配置中是固定项若不需要 NSP需在配置层去掉cls_heads并关闭use_next_sentence_label旋转向量在注意力层构建时按max_sequence_length预分配前向时按实际q_len/k_len切片索引self.q_cos_vec[:, 0:q_len, ...]因此推理序列不能超过 512且 Q 与 KV 可以不等长分别由q_max_sequence_length/kv_max_sequence_length控制roformer_attention_test.py、roformer_encoder_block_test.py、roformer_encoder_test.py 三个测试文件覆盖了三角向量正确性、注意力形状、块与整体编码器行为可作为修改实现时的回归基线运行前提依赖仓库根目录 requirements.txt 声明的 TensorFlow 环境--tpulocal需要可用的本地 TPU在纯 GPU/CPU 机器上应通过--params_override将runtime.distribution_strategy改为对应策略如mirrored再运行。九、小结official/projects/roformer用不到千行代码给出了一个结构清晰的旋转位置编码 Transformer 参考实现README.md 中的单条train.py命令背后是exp_factory注册的预训练实验、base 规模 YAML 配置、Keras 序列化的编码器三件套encoder / block / attention以及标准 v2 字段名的 TFRecord 预训练数据链路。对研究者而言最值得精读的是 roformer_attention.py 中_build_trig_vector与roformer_recompute_qkv两段代码——它们把“位置旋转”落到了奇偶维度拆分的两次张量操作上配合roformer_attention_test.py的公式级断言可以完整核对 Roformer 位置编码的每一处数学细节。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 TF Model Garden 中的 Roformer：旋转位置编码 Transformer 的源码解析与预训练实战指南
TF Model Garden 中的 Roformer旋转位置编码 Transformer 的源码解析与预训练实战指南【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/modelsRoformerRotary Position Embedding Former是 TF Model Gardenofficial/目录下 TensorFlow 官方模型库中的 NLP 实验项目用 Keras 函数式 API 重新实现了“旋转位置编码”这一 Transformer 位置表示方案并完整打通了预训练与 GLUE 微调链路。本文基于 official/projects/roformer/README.md 给出的训练命令展开结合official/projects/roformer/下的模型源码、实验配置与测试用例讲清 Roformer 编码器结构、旋转位置编码的数学实现、预训练配置参数与数据格式读完你可以直接在该仓库内复现 Roformer 的 TPU 预训练并理解其与 BERT 的关键差异。一、项目结构与定位official/projects/roformer/是一个自包含的模型项目文件职责如下文件职责roformer.pyRoformerEncoderConfig与get_encoder工厂函数把配置绑定为RoformerEncoderroformer_encoder.pyRoformerEncoder编码器网络Keras Functional API 构建roformer_encoder_block.pyRoformerEncoderBlock单层 Transformer 块注意力 FFN LayerNorm 残差roformer_attention.pyRoformerAttention多头注意力层核心是旋转位置向量_build_trig_vectorroformer_experiments.py注册roformer/pretraining与roformer/glue两个实验配置train.py训练入口脚本experiments/roformer_base.yamlRoformer base 规模的完整实验配置roformer_attention_test.py、roformer_encoder_test.py、roformer_encoder_block_test.py注意力、编码器、编码器块的单元测试模型实现遵循 TF Model Garden 的标准范式通过official.core.exp_factory注册实验名如roformer/pretraining由统一的 train.py 驱动训练模型本体全部用tf_kerasKeras 3构建并注册为可序列化对象tf_keras.utils.register_keras_serializable(packageText)。二、README 训练命令逐行拆解README.md 给出的运行方式是DATA_PATH??? OUTPUT_DIR??? python3 train.py \ --experimentroformer/pretraining \ --config_fileexperiments/roformer_base.yaml \ --params_overridetask.validation_data.input_path${DATA_PATH},runtime.distribution_strategytpu \ --tpulocal \ --model_dir${OUTPUT_DIR} \ --modetrain_and_eval各参数的含义可对照 train.py 与 official/common/flags.py 中define_flags()的定义--experimentroformer/pretraining实验名必须与 roformer_experiments.py 中exp_factory.register_config_factory(roformer/pretraining)注册的工厂名一致train_utils.parse_configuration(FLAGS)会据此加载默认实验配置--config_fileexperiments/roformer_base.yaml在实验默认配置之上叠加的 YAML 配置相对train.py所在目录即roformer_base规模--params_overridetask.validation_data.input_path${DATA_PATH},runtime.distribution_strategytpu以点号.参数名值的形式做最终覆盖。这里把验证集数据路径指向DATA_PATH并显式将分布式策略设为 TPU与--tpulocal配合使用本地 TPU--tpulocal指定本地 TPU 地址由 official/common/distribute_utils.py 的get_distribution_strategy据此创建TPUStrategy--model_dir${OUTPUT_DIR}checkpoint 与 summary 输出目录。train.py 中可见只有mode包含train时才会serialize_config写出 YAML 配置文件避免纯 eval 任务与训练任务并发写同一文件产生竞争--modetrain_and_eval边训练边验证最终交给train_lib.run_experiment执行训练结束后train_utils.save_gin_config还会把 gin 配置落盘存档。另外train.py 支持runtime.mixed_precision_dtype开启混合精度mixed_float16/mixed_bfloat16TPU 场景通常配合 bfloat16 使用。三、预训练实验配置roformer_base.yaml 全参数解读experiments/roformer_base.yaml 对应“base”规格12 层 × 768 维与 BERT-base 同量级。关键参数分组如下模型部分task.model参数值说明encoder.typeany编码器类型选择器any下由roformer.RoformerEncoderConfig构建RoformerEncoderhidden_size768Transformer 隐层维度必须是num_attention_heads的整数倍见 roformer_encoder_block.py 的校验num_layers12Transformer 层数num_attention_heads12每层注意力头数单头维度 768/1264Roformer 要求该值为偶数见 roformer_attention.py 的assert self._key_dim % 2 0intermediate_size3072FFN 第一层维度且必须为偶数roformer_encoder_block.py 会直接抛错代码中通过 kwargs 别名映射到inner_dimhidden_activationgeluFFN 激活函数max_position_embeddings512决定旋转位置向量的预构建长度见下文vocab_size/type_vocab_size30522 / 2词表与 segment 类型数embedding_size768词嵌入宽度与 hidden_size 相等时省去投影层attention_dropout_rate/dropout_rate0.1注意力 dropout 与层间 dropoutinitializer_range0.02TruncatedNormal初始化标准差roformer.pycls_headsnext_sentence2 分类tanh下一句预测头数据部分task.train_data/task.validation_dataseq_length512、max_predictions_per_seq76、use_next_sentence_labeltrue、global_batch_size256、drop_remainder训练 true / 验证 false。注意 Roformer 实验在 roformer_experiments.py 中显式设置了use_v2_feature_namesTrue其含义见第五节。训练器部分trainer参数值说明optimizer_config.optimizeradamw默认weight_decay_rate0.01LayerNorm/layer_norm/bias不衰减roformer_experiments.py学习率polynomial初值 1e-4decay_steps1000000end_learning_rate0power1.0多项式线性衰减warmuppolynomialwarmup_steps100001 万步线性预热train_steps1000000总训练步数checkpoint_interval/validation_interval20000 / 1000checkpoint 与验证频率validation_steps64steps_per_loop/summary_interval50 / 50分布式循环步长与 summary 频率max_to_keep5最多保留 checkpoint 数roformer_experiments.py中还注册了第二个实验roformer/glue源码任务类型为sentence_prediction.SentencePredictionConfig学习率降至 3e-5、同样采用 AdamW 多项式 warmup用于在预训练权重上做 GLUE 句级分类微调。README 只写了 pretraining 命令但将--experiment换成roformer/glue并覆盖task.train_data.input_path等字段即可走微调流程。四、编码器结构与 BERT 的关键差异RoformerEncoder 用 Keras Functional API 构建输入为三个张量input_word_ids、input_mask、input_type_ids输出为sequence_output末层全部 token 表示、pooled_output[CLS] 位经pooler_transform的 tanh Dense 层以及每一层的encoder_outputs。与 BERT 编码器相比源码中有两处值得注意的“减法”没有绝对位置嵌入。roformer_encoder.py 中明确注释 “Roformer does not need a position embedding layer / does not have absolute position embedding”嵌入只是词嵌入 类型嵌入的相加再 LayerNorm Dropout_position_embedding_layer被显式置为NoneL238。位置信息完全由注意力层的旋转位置编码注入位置信息的载体是预构建的正弦/余弦向量每层RoformerEncoderBlock在构建RoformerAttention时传入q_max_sequence_length与kv_max_sequence_length均取自配置里的max_position_embeddings512三角向量在层初始化时一次性算好前向传播不再需要额外的位置嵌入查表。编码器其余部分保持经典 Transformer 结构若embedding_width ! hidden_size会插入一个embedding_projection的EinsumDense投影12 层RoformerEncoderBlock依次堆叠最后一层支持output_range切片以截取部分 token 输出。五、核心机制旋转位置编码Rotary Position EmbeddingRoformer 的技术核心在 roformer_attention.py。RoformerAttention继承自tf_keras.layers.MultiHeadAttention重写了两处逻辑。1. 三角向量构建_build_trig_vectorL23-L38def _build_trig_vector(length, key_dim): Builds the trig vector. tf_dtype tf_keras.mixed_precision.global_policy().compute_dtype position_ids tf.cast(tf.range(length), dtypetf_dtype) position_ids tf.expand_dims(position_ids, axis0) steps key_dim // 2 # 2 (i - 1) / key_dim (i - 1) / steps: (-1 achieved with zero-indexing) wavenumber_exponent -tf.cast(tf.range(steps), dtypetf_dtype) / steps wavenumbers tf.pow( tf.constant(10000.0, dtypetf_dtype), wavenumber_exponent ) vec tf.einsum(bl,d-bld, position_ids, wavenumbers) sin_vec tf.repeat(tf.sin(vec), repeats2, axis-1) cos_vec tf.repeat(tf.cos(vec), repeats2, axis-1) sin_vec, cos_vec tf.expand_dims(sin_vec, 2), tf.expand_dims(cos_vec, 2) return sin_vec, cos_vec其数学含义对位置 m、频率分量 i角度为m * 10000^(-2i/key_dim)即标准 sinusoidal 频率谱基频 1最高分量频率 1/10000与 BERT 位置嵌入使用同一频率公式但 Roformer 不在嵌入上加它而是作用于注意力投影之后的 Q/K 向量上。sin/cos_vec在最后一维做repeat(repeats2)是把每个频率分量“铺”到相邻的 (2i, 2i1) 两个维度上为后续的成对旋转做准备。2. Q/K 的旋转roformer_recompute_qkvL77-L91def roformer_recompute_qkv(self, q, k, v): ... q2 tf.stack([-q[..., 1::2], q[..., ::2]], axis4) q2 tf.reshape(q2, q_shape) k2 tf.stack([-k[..., 1::2], k[..., ::2]], axis4) k2 tf.reshape(k2, k_shape) ret_q q * self.q_cos_vec[:, 0:q_len, ...] q2 * self.q_sin_vec[:, 0:q_len, ...] ret_w k * self.k_cos_vec[:, 0:k_len, ...] k2 * self.k_sin_vec[:, 0:k_len, ...] return ret_q, ret_w, v实现上把 Q或 K的奇偶维度拆成两半q[..., ::2]、q[..., 1::2]构造旋转算子q2 [-q_odd, q_even]然后执行q·cos(mθ) q2·sin(mθ)——这正是二维平面旋转矩阵(cosθ, -sinθ; sinθ, cosθ)的成对向量形式每相邻两个维度构成一个旋转平面位置 m 决定旋转角mθ_i。V 向量不旋转直接返回。在call中L93-L117流程是线性投影得到 Q/K/V →roformer_recompute_qkv对 Q、K 施加位置旋转 → 复用父类MultiHeadAttention._compute_attention做标准缩放点积注意力含 mask→ 输出投影。由于 Q、K 都带位置旋转q_i·k_j的内积只依赖相对位置i-j这是旋转位置编码“相对性”的由来。测试印证roformer_attention_test.py 的test_trig_vector对length[8,50]、key_dim[64,128]的组合用tf.assert_equal逐位核对sin_emb/cos_emb与10000^(-i/half_d)公式的直接计算结果并校验奇偶维度重复一致test_attention_scores则验证带/不带 mask 时输出形状为[batch, seq_len, key_dim]。六、编码器块RoformerEncoderBlockRoformerEncoderBlock 是标准“注意力 FFN”块几个实现细节偶数维度硬约束__init__中要求inner_dim为偶数L87-L89build中要求hidden_size能被头数整除并计算单头维度attention_head_sizeL132-L136支持 pre-norm / post-norm 两种结构norm_first参数callL243-L310中 pre-norm 路径是“先 LayerNorm 再残差相加”post-norm 路径是“残差相加后再 LayerNorm”数值稳定性处理两个 LayerNorm 固定用dtypetf.float32L157-L164 注释说明为数值稳定混合精度下也未验证 fp16 安全FFN 激活层在mixed_bfloat16全局策略下回退到 float32 计算L172-L179输入形态灵活支持(input)、(input, attention_mask)、(query, key_value, attention_mask)三种输入后者为交叉注意力场景保留output_range切片设置后只对前output_range个 token 输出编码器中仅最后一层启用对应RoformerAttention构造时对q_sin_vec/q_cos_vec的同步切片roformer_attention.py。七、预训练数据格式与准备预训练任务由official.nlp.tasks.masked_lm.MaskedLMConfig驱动数据加载器是 official/nlp/data/pretrain_dataloader.py 的BertPretrainDataLoaderfile_type为tfrecord。Roformer 实验设置了use_v2_feature_namesTrue这决定了 TFRecord 中tf.Example的字段名use_v2_feature_namesTrue时pretrain_dataloader.py字段为input_word_ids、input_type_ids、input_mask、masked_lm_positions、masked_lm_ids、masked_lm_weights、next_sentence_labels为False时则使用历史 BERT 命名input_ids、segment_ids解码时所有 int64 特征会被 cast 成 int32TPU 不支持 int64见 pretrain_dataloader.py 注释因此 TFRecord 按int64写出即可。数据准备可参考仓库内的标准脚本 official/nlp/data/create_pretraining_data.py其命令行 flaguse_v2_feature_namesL80以随机掩码 15% 的方式生成 masked LM 样本。README 命令中的DATA_PATH即指向按上述格式v2 字段名、seq_length512、max_predictions_per_seq76切分好的 TFRecord 目录通过--params_override覆盖task.validation_data.input_path训练集路径同理可通过task.train_data.input_path覆盖。数据规格细节另见 official/nlp/docs/data_processing.md 与 official/nlp/data/README.md。八、从源码结构看可推断的适用边界从源码结构看Roformer 实现面向 BERT 式双目标预训练masked LM next sentence predictionnext_sentence分类头在配置中是固定项若不需要 NSP需在配置层去掉cls_heads并关闭use_next_sentence_label旋转向量在注意力层构建时按max_sequence_length预分配前向时按实际q_len/k_len切片索引self.q_cos_vec[:, 0:q_len, ...]因此推理序列不能超过 512且 Q 与 KV 可以不等长分别由q_max_sequence_length/kv_max_sequence_length控制roformer_attention_test.py、roformer_encoder_block_test.py、roformer_encoder_test.py 三个测试文件覆盖了三角向量正确性、注意力形状、块与整体编码器行为可作为修改实现时的回归基线运行前提依赖仓库根目录 requirements.txt 声明的 TensorFlow 环境--tpulocal需要可用的本地 TPU在纯 GPU/CPU 机器上应通过--params_override将runtime.distribution_strategy改为对应策略如mirrored再运行。九、小结official/projects/roformer用不到千行代码给出了一个结构清晰的旋转位置编码 Transformer 参考实现README.md 中的单条train.py命令背后是exp_factory注册的预训练实验、base 规模 YAML 配置、Keras 序列化的编码器三件套encoder / block / attention以及标准 v2 字段名的 TFRecord 预训练数据链路。对研究者而言最值得精读的是 roformer_attention.py 中_build_trig_vector与roformer_recompute_qkv两段代码——它们把“位置旋转”落到了奇偶维度拆分的两次张量操作上配合roformer_attention_test.py的公式级断言可以完整核对 Roformer 位置编码的每一处数学细节。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考