hmkang/wam_ctxpool_bmethod

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

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