TensorFlow Model Garden 中的 Perceiver IO：从潜在向量编码到 GLUE 微调的完整配置与实践

Perceiver IO 是一种面向"结构化输入与输出"的通用 Transformer 变体架构，其核心思想是用一组可学习的**潜在向量（latent array）**替代传统自注意力对长序列的直接建模。本文基于 TensorFlow Model Garden 仓库中 official/projects/perceiver 目录下的官方实现文档与源码，完整讲解该实现的架构构成、Wiki+Books 预训练配置、GLUE 微调训练配方，以及实验结果与论文的差异。读完后，你将理解 Perceiver 编码器/解码器在 TF2 中的具体落地方式，并能依据仓库中的 YAML 配置与实验工厂（experiment factory）完整复现预训练与下游微调流程。

一、项目定位：TF2 版 Perceiver IO 实现

Perceiver IO 模块的定位直接写在其 README 中：

TF2 implementation of Perceiver (https://arxiv.org/abs/2107.14795).

即它是 Google 论文 Perceiver IO: A General Architecture for Structured Inputs & Outputs 的 TensorFlow 2 参考实现。README 同时说明，预训练（pretrain）、微调（finetune）与从头训练（train from scratch）所需的脚本入口统一组织在项目的实验目录下；在仓库中，对应的实际落点为 official/projects/perceiver/train.py 入口与 official/projects/perceiver/configs/experiments/ 下的 YAML 实验配置文件。

从目录结构看，该实现分为四层：

• configs/：实验注册与配置类定义，见 configs/perceiver.py 与 configs/encoders.py；
• modeling/layers/：底层注意力组件，包括 Encoder（自注意力 + 交叉注意力处理栈）、Decoder（交叉注意力解码器）与 utils；
• modeling/networks/：面向序列任务的封装，包括 SequenceEncoder 与 PositionalDecoder；
• modeling/models/ 与 tasks/：MLM 预训练模型 Pretrainer、分类模型 Classifier，以及对应的 pretrain、sentence_prediction 任务。

二、架构解析：编码器、潜在数组与解码器

2.1 序列编码器：词嵌入 + 位置编码 + 潜在查询

SequenceEncoder 是 Perceiver 处理文本的入口网络。它的职责（源码 docstring 明确说明）是：

Assumes positional learned encoding for latent inputs and embeddings. Creates an embedding table with vocab size. It uses the perceiver encode processor to encode the input and process the latent representation. It can be pretrained on masked LM and reused for fine-tuning.

其 call 方法（sequence_encoder.py）的处理流程是：

1. 通过 OnDeviceEmbedding 将 input_word_ids 映射为词嵌入，再叠加可学习输入位置编码 PositionEmbedding；
2. 构造形状为 [batch, z_index_dim, d_latents] 的潜在数组初始状态 z，其位置编码 z_pos_enc 由独立的 PositionEmbedding 层生成——这正是 Perceiver 的"可学习潜在查询"；
3. 将词嵌入序列与潜在查询一起送入 Perceiver Encoder，输出 dict(latent_output=z)，即压缩后的潜在表征。

默认超参数来自 SequenceEncoderConfig：d_model=768（输入/输出特征维度）、d_latents=1280（潜在向量宽度）、z_index_dim=256（潜在向量个数，即潜在数组长度）、max_seq_len=512、vocab_size=30522（WordPiece 词表）、embedding_width=768。

2.2 Perceiver 编码处理栈：先交叉注意力，再多次自注意力

核心组件 Encoder 的类注释说明它实现的是 Perceiver 论文的 "Encoder and Processor stack"，由 SelfAttention 与 CrossAttention 模块组合。其 build/call 方法（encoder.py）揭示了具体计算顺序：

1. 一次交叉注意力：以潜在数组 z 为 query、词嵌入序列为 key/value，构建 attention mask（make_cross_attention_mask）后，让 z 从长序列中提取信息；
2. 若干自注意力块：交叉注意力之后，对 z 串联执行 num_self_attends_per_block × num_blocks 个 TransformerEncoderBlock，块间完全共享输入形式（自注意力）。

Perceiver Base 在仓库配置 EncoderConfig 中的取值为：num_self_attends_per_block=26、num_blocks=1、self/cross_attention_num_heads=8、qk_last_dim=256（8 头 × 每头 32 维）、v_last_dim=1280（与 d_latents 一致）、dropout_prob=0.0。值得注意的是，这一"1 个块 × 26 次自注意力"的配置与常见 Transformer 的"8 层堆叠"形态不同——从源码结构看，它把计算深度集中在潜在数组的自注意力循环上，使每步复杂度只与 z_index_dim=256 而非序列长度相关，这正是 Perceiver 能处理长序列的机制所在。

2.3 解码器：用位置编码作为查询的通用读出

Perceiver IO 的"IO"体现在输出端：同一个潜在数组 z 可以接不同解码器。仓库中的实现是 PositionalDecoder：

• 它构造一个 PositionEmbedding 作为解码查询（decoder_pos_enc，最大长度即 output_index_dim），把潜在数组当作 key/value 送入基础 Decoder（单块交叉注意力 TransformerEncoderBlock，shape_for_attn="kv"），输出 sequence_output；
• 依据 output_index_dim 的不同，同一套代码派生出两类配置（perceiver.py）：
  • MaskedLMDecoderConfig：output_index_dim=512，与序列等长，用于 MLM 预训练时读出每个位置的预测；
  • ClassificationDecoderConfig：output_index_dim=1，用于下游分类任务读出一个表示。

两者共享 d_model=768、d_latents=1280、z_index_dim=256，且 DecoderConfig 固定 num_heads=8、v_last_dim=768、use_query_residual=False。

三、预训练：BERT Wiki Books 配方

README 中 "BERT Wiki Books Pretrain" 一节说明：模型配置可在 configs 与 experiment 目录中推导得出，除分词器与数据外基本遵循论文 Table 8 / Table 9 的设置。预训练结果与论文对比如下（引自 README）：

Model	Tokenizer	Pretrain Data	Batch Size	Steps	Val MLM Accuracy
Perceiver IO Base (paper)	SentencePiece	T5 + Wiki	512	500 k	N/A
Perceiver IO Base (ours)	WordPiece	Wiki + Books	512	500 k	68.69 %

3.1 预训练 YAML 配置逐项解读

预训练实验配置文件为 wiki_books_pretrain.yaml，其关键字段与源码配置的对应关系如下：

task:
  init_checkpoint: ''        # 冷启动时无预训练权重
  train_data:
    drop_remainder: true
    global_batch_size: 512   # 与 README 表格中 Batch Size=512 一致
    # 使用 glob 匹配所有分片；00141-of-00500-eval 保留作验证
    input_path: ''
    is_training: true
    max_predictions_per_seq: 76   # 每条序列 512 个 token 中约 15% 位置被掩码
    seq_length: 512
    use_next_sentence_label: false # 纯 MLM，不用 NSP（对齐 T5 风格词表任务）
    use_position_id: false
    use_v2_feature_names: true
  validation_data:
    drop_remainder: false
    global_batch_size: 512
    is_training: false
    # 其余字段与训练集相同
trainer:
  checkpoint_interval: 20000   # 每 2 万步存一次检查点
  max_to_keep: 5               # 最多保留 5 份
  steps_per_loop: 1000
  summary_interval: 1000
  train_steps: 500000          # 与 README 的 500 k 步一致
  validation_interval: 1000
  validation_steps: 64

其中 use_next_sentence_label: false 与 use_v2_feature_names: true 来自配置工厂 perceiver_wordpiece_pretrain：该工厂以 BertPretrainDataConfig 构建数据管道，global_batch_size=512，实验名注册为 perceiver/wordpiece_pretrain。

3.2 预训练的优化器与学习率

与 YAML 配套的是 configs/perceiver.py 中的 _MLM_WORDPIECE_TRAINER：

• 优化器：LAMB，weight_decay_rate=0.01，权重衰减排除 LayerNorm、layer_norm、bias；
• 学习率：余弦衰减，initial_learning_rate=1.25e-3、decay_steps=500000（与总步数对齐）；
• 预热：线性 warmup 1000 步，从 0 升到初始学习率。

MLM 头本身由 Pretrainer 承载。从源码看，Pretrainer 在编码器输出 latent_output 的基础上经过解码器得到 sequence_output，再用 layers.MaskedLM（pretrainer.py）按 masked_lm_positions 抽取被掩码位置并产生 mlm_logits——MLM 输出层共享词嵌入表（embedding_table=self.encoder.get_embedding_table()），mlm_activation 默认 'gelu'、初始化范围 mlm_initializer_range=0.02（见 PretrainerConfig）。

3.3 训练入口与实验注册

统一入口 train.py 只做三件事：注册 Perceiver 配置（configs.perceiver）、注册任务（tasks.pretrain、tasks.sentence_prediction）、调用 official.nlp.train.main 驱动训练：

from official.common import flags as tfm_flags
from official.nlp import train

# pylint: disable=unused-import
from official.projects.perceiver.configs import perceiver
from official.projects.perceiver.tasks import pretrain
from official.projects.perceiver.tasks import sentence_prediction

if __name__ == '__main__':
  tfm_flags.define_flags()
  app.run(train.main)

因此运行方式为"入口脚本 + 实验名 + YAML"的组合。以预训练为例，需将 YAML 填入 task.train_data.input_path 等数据路径，并以类似下面的方式启动（--exp_name 与 perceiver_wordpiece_pretrain 中注册的实验名对应，数据目录按自身分片情况填充）：

python official/projects/perceiver/train.py \
  --exp_name perceiver/wordpiece_pretrain \
  --config_file official/projects/perceiver/configs/experiments/wiki_books_pretrain.yaml \
  --data_dir /path/to/your/bert_pretrain_data \
  --model_dir /path/to/your/model_dir \
  --tpu_spec tpu

train.py 同时导入了 sentence_prediction 任务并注册了两个下游实验：perceiver/word_piece_sentence_prediction（TF 记录输入）与 perceiver/word_piece_raw_sentence_prediction（原始文本输入，对应 SentencePredictionTextDataConfig），均开启 XLA 运行时（enable_xla=True）。

四、GLUE 微调：训练配方与结果

4.1 微调优化器：LAMB + 多项式学习率衰减

下游微调与预训练使用不同的 Trainer 配置。configs/perceiver.py 中的 _SENTENCE_PREDICTION_TRAINER 为：

• 优化器：LAMB，权重衰减同样排除归一化层与 bias；
• 学习率：多项式衰减，initial_learning_rate=3.0e-5、end_learning_rate=0.0、power=1.0，decay_steps 按任务数据集的 10 个 epoch 总步数设定；
• 预热：线性 200 步。

以 glue_sst.yaml（SST-2）为例：

task:
  hub_module_url: ''
  model:
    num_classes: 2
  metric_type: 'accuracy'
  train_data:
    drop_remainder: true
    global_batch_size: 32      # GLUE 微调统一 batch size 32
    input_path: ''
    is_training: true
    seq_length: 128           # 下游任务序列长度 128
    label_type: 'int'
  validation_data:
    drop_remainder: false
    global_batch_size: 32
    is_training: false
    seq_length: 128
    label_type: 'int'
trainer:
  checkpoint_interval: 1000
  optimizer_config:
    learning_rate:
      polynomial:
        decay_steps: 21040    # 100% of train_steps
        end_learning_rate: 0.0
        initial_learning_rate: 3.0e-05
        power: 1.0
      type: polynomial
  steps_per_loop: 1000
  summary_interval: 1000
  # Training data size 67,349 examples, 10 epochs.
  train_steps: 21040
  validation_interval: 1052
  # Eval data size = 872 examples.
  validation_steps: 28
  best_checkpoint_export_subdir: 'best_ckpt'   # 按最佳指标导出最优检查点
  best_checkpoint_eval_metric: 'cls_accuracy'
  best_checkpoint_metric_comp: 'higher'

configs/experiments/ 下为 GLUE 全部任务提供了同构配置：glue_cola.yaml、glue_mnli_m.yaml、glue_mnli_mm.yaml、glue_mrpc.yaml、glue_qnli.yaml、glue_qqp.yaml、glue_rte.yaml、glue_sst.yaml、glue_stsb.yaml。以 glue_qqp.yaml 为例，其 train_steps=113700（363,849 条训练样本 × 10 epochs）、validation_steps=1264（40,430 条评测样本），同样按 cls_accuracy 导出 best_ckpt。

微调侧的网络装配由 Classifier 完成：它与 Pretrainer 一样支持"编码器直接输出 sequence_output"或"编码器输出 latent_output + 解码器读出"两种路径；分类头为 layers.ClassificationHead，默认初始化标准差取 1/sqrt(输出维度)。微调加载预训练权重通过任务配置 init_checkpoint 字段（SentencePredictionConfig 中定义）指定。

4.2 GLUE 单任务微调结果

README 给出了基于上述预训练权重的 GLUE 单任务微调成绩（配置对应论文 Table 10），并强调这些都是单任务（single-task）微调：

Model	Tokenizer	Pretrain Data	CoLA	MNLI-m/mm	MRPC	QNLI	QQP	RTE	SST-2	STS-B	Average
Perceiver IO Base (paper)	SentencePiece	T5 + Wiki	47.11 %	84.53/85.03 %	87.25 %	92.12 %	90.22 %	65.23 %	94.38 %	88.18 %	81.16 %
Perceiver IO Base (ours)	WordPiece	Wiki + Books	63.23 %	84.29/84.52 %	87.74 %	91.43 %	91.22 %	70.76 %	94.15 %	89.85 %	84.09 %

平均分按 8 个任务计算，其中 MNLI-matched 与 MNLI-mismatched 先取平均再计为一个任务：

Average = (63.23 + (84.29 + 84.52) / 2 + 87.74 + 91.43 + 91.22 + 70.76 + 94.15 + 89.85) / 8

README 同时明确记录了与论文的差异："~+2.93 average GLUE accuracy compared to paper results"，即仓库实现平均分 84.09% 相对论文报告的 81.16% 高出约 2.93 个百分点（主要差距来自 CoLA：63.23% 对 47.11%）。

五、与论文的差异、适用前提与引用

在把仓库结果用于复现或二次开发时，以下几点适用前提值得注意：

1. 分词与数据不同：本实现使用 WordPiece（BERT 风格词表，vocab_size=30522）与 Wiki+Books 数据，而论文使用 SentencePiece 与 T5+Wiki；因此结果不可与论文数字直接对等比较，仓库 README 已单列"Discrepancy with the paper"说明；
2. 配置可复现性：模型结构参数完全由 configs/perceiver.py 中的 dataclass 与 encoders.py 中的 build_encoder 构建函数确定，YAML 只覆盖数据与训练器字段——这保证了"配置即实验"；
3. 任务形态：该实现覆盖两类任务——MLM 预训练（perceiver/wordpiece_pretrain）与句子对预测/分类微调（perceiver/word_piece_sentence_prediction、perceiver/word_piece_raw_sentence_prediction），预训练与微调通过 init_checkpoint 衔接；
4. 测试佐证：各核心组件均配有测试，如 encoder_test.py、pretrainer_test.py、pretrain_test.py，可在改动配置后先行回归验证。

仓库也提供了官方引用信息（引自 README 的 Citing 一节）：

@misc{tensorflowmodelgarden2022,
  author = {Hongkun Yu and Chen Chen and Xianzhi Du and Yeqing Li and
            Abdullah Rashwan and Le Hou and Pengchong Jin and Fan Yang and
            Frederick Liu and Jaeyoun Kim and Jing Li},
  title = {{TensorFlow Model Garden}},
  howpublished = {\url{https://github.com/tensorflow/models}},
  year = {2020}
}

六、小结

TensorFlow Model Garden 的 Perceiver IO 实现把论文的"编码器-潜在数组-解码器"三段式结构完整落地为可配置的 TF2 工程：

• 编码器侧：SequenceEncoder（词嵌入 + 双套位置编码）+ Encoder（1 次交叉注意力 + 26 次潜在自注意力），复杂度由 z_index_dim=256 而非序列长度主导；
• 解码侧：PositionalDecoder + Decoder 以"位置编码即查询"的方式，用同一潜在数组分别支撑 512 长的 MLM 读出与 1 维的分类读出；
• 训练侧：预训练（LAMB + 余弦衰减，batch 512，500k 步）与 GLUE 微调（LAMB + 多项式衰减，batch 32，10 epochs，按最佳 cls_accuracy 导出检查点）各有独立、可直接读取的 YAML 配方；
• 结果侧：Wiki+Books 预训练达到 68.69% 的 Val MLM 准确率，GLUE 单任务微调平均 84.09%，并在 README 中如实标注了相对论文约 +2.93 的平均分差异及其分词/数据来源。

这套代码的价值在于：它既是一个可运行的 Perceiver IO 文本实现，也是一份展示了"如何用 Model Garden 的实验工厂 + dataclass 配置 + YAML 覆盖"这一标准范式来组织预训练/微调实验的样例，值得在复现长序列 Transformer 变体或构建自定义解码头时作为参照。