---
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`).