model.md

# Model

## 1. 共同目标

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

```text
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：

```text
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：

```text
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]`：

```text
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. 当前推荐对比

当前主要实验比较：

```text
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）。