| --- |
| license: apache-2.0 |
| tags: [video-world-model, robotics, robocasa, wan2.2, context-pooling, ibot] |
| --- |
| # WAM_DIT4DIT — context pooling on the **finalized B-method** base (RoboCasa kitchen) |
| |
| Wan2.2-TI2V-5B video DiT + 48-joint action head, `training_mode=joint`. The base is suhyeok's finalized |
| B-method recipe: a teacher-forced (sigma=0.25) self-EMA teacher plus an iBOT prototype loss at L18<->L18, |
| gamma=0.01, two-view. On top of it the 3 PAST cond latent frames are pooled into one motion frame before a |
| chosen block. |
|
|
| **These are NOT the surrogate ctxpool runs.** The surrogate line (older base, pd8 x GA1) lives in |
| `hmkang/wam_ctxpool_xattn` and `hmkang/wam_ctxpool_avg`. Do not compare across the two sets. |
|
|
| | folder | pooling | plate | |
| |---|---|---| |
| | `xattn_ffn_L3/` | cross-attn + FFN refine before block L3 | 4 GPU x pd4 x GA2, GC=0 | |
| | `avg_st_L3/` | past pooled over time AND space -> 1 token, before L3 | same | |
| | `avg_L6/` | past 3 latent frames averaged, before block L6 | same | |
| | `avg_L3/` | past 3 latent frames averaged, before block L3 | same | |
| | `xattn_L3/` | cross-attn (current queries the past, rope-aware, zero-init `o`), before L3, no FFN | same | |
| | `lq_q96_L3/` | 96 LEARNABLE queries (no rope) cross-attend the past; current excluded from the kv | same | |
| | `lq_q48_L3/` | same design, 48 slots — the compression axis (6x instead of 3x) | same | |
| | `control_3latin/` | NO pooling, 3-latin geometry (cond 3 + future 2 = 480 tok/row) | same | |
| | `control_4latin/` | NO pooling, 4-latin geometry (cond 4 + future 2 = 576 tok/row) | same | |
|
|
| Geometry: 4-latin (`num_frames_in=25`, `num_frames_out=41`, fdf 2) = 4 cond + 2 future latent slots, |
| 96 tokens per latent frame, 576 tokens per row. Effective batch 16 clips x GA 2 x 2 views = 64 rows, |
| the same plate the base run uses via pd8 x GA1. |
|
|
| `WAM_DINO_GA_AWARE=1` makes the iBOT center and the teacher-head EMA fire once per optimizer step, which is |
| what makes GA=2 the same experiment as GA=1 at twice the per-device batch. deepspeed 0.19.6. |
|
|
| Every 20k steps. Weights + configs only (no optimizer state). Code: https://github.com/HEMMO0208/wam |
| (branch `ctxpool-bmethod`; `gr00t/model/wam_dit4dit/wan22_ctxpool.py`). |
|
|
