model.md

Model

1. 共同目标

HY-WorldPlay 使用四步去噪。Predictor 的目标是在运行时替代一个或多个完整 54-layer DiT step：

Full step:       54 个 Teacher blocks + final layer
Predictor step:  2 个 Predictor blocks + frozen Teacher final layer

所有版本都使用目标 step 的 65 通道 I2V 输入、timestep/action frame condition、相机参数和 RoPE； 最终均预测 final-layer 前的 hidden，并复用冻结的 Teacher final layer 得到 velocity。residual head 零初始化，使初始状态满足 pred_hidden = anchor_hidden，等价于 Direct Reuse。

2. 版本总览

版本	image fusion	文本/Context	previous chunk	Predictor blocks	参数量
v1	current + same-chunk anchor	dense current/cached text，双流预测	无	完整 double-stream	339,893,248
v2	current + same-chunk anchor	Text K/V + exact history Vision K/V	无	AR vision-only	176,241,152
v3	current + anchor + previous chunk	Text K/V + exact history Vision K/V	有	AR vision-only	197,218,816
v4	current + anchor + previous chunk	Text K/V + BF16 pre-feature 重建 Vision K/V	有	AR vision-only	197,218,816
predictor-disca	current + same-chunk anchor	仅 Text K/V	无	AR vision-only	176,241,152

参数量为可训练参数。v4 另含 16,785,664 个冻结 Context projector 参数。

3. Predictor v1

实现：models/predictor.py。

• img_in、final layer 从 Teacher 加载并冻结；
• current image tokens 与 same-chunk Full-DiT anchor hidden 经过 image fusion；
• current text 与 text-only prefill 后的 cached text 经过独立 text fusion；
• 两个完整 image/text double-stream blocks 默认从 Teacher [1,52] 初始化；
• 不读取 Text K/V 或 history Vision K/V；
• 输出 anchor_hidden + residual_out(predictor_hidden)。

v1 是早期可行性验证代码路径。旧机器上的实验结果不再纳入当前项目报告；当前可核验的结果统一见 experiment.md。

4. Predictor v2

实现：models/predictor_v2.py。

v2 将 v1 的 dense text double-stream 改成 WorldPlay 原生 AR vision inference：

• current image tokens 与 same-chunk anchor hidden 做双输入 fusion；
• Predictor 内不执行 text stream；
• 每个 Predictor block 读取该 Teacher 层的 Text K/V 和 history Vision K/V；
• Vision K/V 包含 regular RoPE 与 ProPE 两个分支；
• 支持 [0,53] 或 [1,52] 两种 Teacher 初始化；
• img_in、final layer和 block text 参数冻结。

v2 数据直接保存每层 history Vision K/V，I/O 和磁盘开销较大。代码仍保留，但当前后续实验统一使用 v4 pre-feature schema。

5. Predictor v3

实现：models/predictor_v3.py。

v3 在 v2 上增加 previous-chunk same-timestep hidden：

current image tokens
same-chunk anchor final hidden
previous-chunk target-step final hidden
        -> TripleFeatureFusion -> two AR vision blocks

TripleFeatureFusion 分别 LayerNorm 三路输入，拼接后经过 6144→4096→2048 MLP。因为需要前一 chunk，chunk 0 不进入训练。其 Context 仍是 v2 的 direct Vision K/V。

6. Predictor v4

实现：models/predictor_v4.py。

v4 保留 v3 的三路 image fusion，但将磁盘中的 direct Vision K/V 替换为 BF16 Context pre-feature：

img_modulated
  -> frozen Teacher K/V projections + K norm
  -> Teacher-compatible RoPE/ProPE
  -> Context Vision K/V

训练时从 pre-feature 在线重建 K/V；推理 pipeline 已经持有精确 Teacher Context K/V，因此直接传入， 不重复运行 projector。动态 Context mask 会同时作用于 regular/ProPE 分支。

v4 的实验性 predictor_fusion_type 另支持：

• concat：默认且与旧 checkpoint 严格兼容的三路 LN + concat MLP；
• projection_add：三路 identity-init 独立投影、无 affine source norm、相加后 FusionNorm；
• token_gate_prev：在 projection_add 上，用 current/previous/Teacher time+action condition 预测 [B,N,1] 的 previous-chunk gate。第一版只支持已验证最好的 same_step previous condition， 不 gate same-chunk anchor。

正式分辨率下三路 fusion tensor 均为 [B,6240,2048]，其中 6240=4×30×52。

当前已完成两套正式权重：

• blocks [0,53]：checkpoints/predictor_v4_prefeature_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab；
• blocks [1,52]：checkpoints/predictor_v4_prefeature_blocks1-52_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab。

v4 能训练所有相邻转换 0→1、1→2、2→3。推理时可任意指定一个或多个替代 step；连续替代 时，前一个 Predictor 输出成为下一个 Predictor 的 same-chunk anchor。

7. Predictor-DisCa

实现：models/predictor_disca.py。

这是当前项目中与 v1、v4 并列的 predictor-disca 版本，不是新 HY-WorldPlay 主模型：

• current image tokens 与 same-chunk anchor hidden 做双输入 fusion；
• 不注入 previous-chunk hidden；
• 不注入 Context Vision K/V；
• 不加载或运行 frozen Context projector；
• 保留静态 Text K/V，因为 WorldPlay AR vision block 仍需要文本条件；
• 使用两个 Teacher-initialized AR vision blocks；
• 输出仍为 anchor_hidden + zero-init residual。

当前正式权重为 blocks [0,53]：

checkpoints/predictor_disca_blocks0-53_bf16_8gpu_flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab

它复现了 DisCa 的“两块轻量 Predictor 替代 Full DiT”思想，但没有复现原始 DisCa 的 dense text cache/fusion、固定 full-step cache 语义或对抗训练。因此准确定位是 DisCa-inspired WorldPlay ablation， 不是严格的 WorldPlay DisCa 复现。

8. 当前推荐对比

当前主要实验比较：

Full-DiT:          F-F-F-F
Direct Reuse:      F-R-R-F       # Full steps 0、3；steps 1、2 reuse
Predictor v4:      F-P-P-F
Predictor-DisCa:   F-P-P-F

v4 与 predictor-disca 使用相同 Teacher、数据、监督 pair、训练步数和评测集。两者的核心差异只在 previous-chunk hidden 与 Context Vision K/V，可用于分析 WorldPlay AR Context 的贡献。

ATC Transport 对齐 Causal-Forcing（2026-09-06）

models/predictor_atc.py 的 State Transport 现在默认使用与 Causal-Forcing-a predictor_training/atc_fusion.py 相同的实现（GlobalStateTransport）：当前 chunk 的 token 对上一 chunk 做全局 cross-attention， query/key 分别套用 Teacher 的 3-D RoPE（当前 chunk 帧位置 / 上一 chunk 帧位置），无可学习偏置， 残差加在上一 chunk 特征上；atc_previous_scope=last_frame 时 key 只取上一 chunk 最后一帧。 Transport 头数默认等于 Teacher 的 heads_num（head_dim 128 = sum(rope_dim_list)），其余维度 （hidden 2048、MLP 4096、gate 512）不变；anchor_only / raw_prev / no_gate / no_distance 消融开关保留。 原 3×3 局部窗口实现保留为 --atc_transport_mode local，2026-09-06 之前训练的 ATC checkpoint 都属于 该模式，续训或评估时需显式指定。诊断量新增 transport_attention_entropy_normalized 与 transport_argmax_displacement。实测 B=2、6240 token：global/chunk 25 ms、global/last_frame 15 ms、 local 40 ms（fwd+bwd，SDPA 后端，本环境无 flash_attn）。