File size: 7,286 Bytes

ec0a9aa

# Ctrl-World Token-Pruning Analysis (SiTo / ITM)

Empirical study of inference-time token pruning (SiTo, Importance Token Merge)
applied to the Ctrl-World action-conditioned video world model. Records the
profiling methodology, measurements, and conclusions for later write-up.

All numbers below are on `single_arm` multiview, `makovian` subset,
`episode_000255`, `num_inference_steps=25`, bf16, one H100-class GPU. Absolute
wall times were collected while the GPU was shared with other jobs, so treat
them as *relative* comparisons (dense baseline re-measured in every batch).

---

## 1. Model / workload

- Ctrl-World = SVD-based (Stable Video Diffusion) UNet, action-conditioned,
  autoregressive replay. 3-view multiview, 320x192-per-view stacked (height×3).
- Per episode: `interact_num` autoregressive segments × `num_inference_steps`
  denoise steps × UNet forward. `guidance_scale=1` by default (NO CFG).
- Attention layers: 64 diffusers `AttnProcessor` slots. Spatial self-attention
  sequence lengths observed: **2880 (72×40), 720 (36×20), 180 (18×10)**; temporal
  attention sequence length **11** (num frames), with batch = H·W.

## 2. Profiling methodology

- Wrapped `unet.forward`, `vae.encode/decode` with CUDA-synced timers.
- Registered forward pre/post hooks on `BasicTransformerBlock`,
  `TemporalBasicTransformerBlock`, `ResnetBlock2D`/`TemporalResnetBlock` to get
  per-module-type time.
- Instrumented `SVDSiToAttnProcessor._should_use_sito` to count fired seq_lens.
- Standalone micro-benchmarks: isolated `to_q/k/v + SDPA + to_out`, full
  `BasicTransformerBlock`, and full `TransformerSpatioTemporalModel` forwards,
  dense vs pruned, with plan `prepare()` timed separately.

## 3. Where the time goes

Episode wall-time split:

| Component | Share |
|---|---|
| UNet forward | ~64% |
| VAE decode | ~8% |
| VAE encode | ~1% |
| Python/loop/mediapy | ~27% |

UNet-internal split (per-module-type hooks):

| UNet component | Share of UNet |
|---|---|
| Spatial `BasicTransformerBlock` | ~16% |
| `TemporalBasicTransformerBlock` | ~22% |
| ResNet (2D + temporal) | ~31% |
| proj_in/out, time embed, other | ~31% |

Attention-only split: temporal attn (seq_len=11) ≈57% of attention time; spatial
attn (2880/720/180) the remaining ≈43%. **SiTo/ITM only touch spatial
self-attention ≈ 7% of total wall time.**

## 4. Per-layer prune profitability (micro-bench, dense vs pruned)

Isolated self-attention (`to_q/k/v` + SDPA + `to_out`), plan precomputed:

| seq_len | C | prune | speedup (SDPA-only) |
|---|---|---|---|
| 2880 | 320 | 0.45 | **1.71x** |
| 720 | 640 | 0.35 | 0.72x (slower) |
| 180 | 1280 | 0.20 | 0.76x (slower) |
| 11 (temporal) | 320 | 0.30 | 0.75x (slower) |

Only the largest stage (2880) benefits. Smaller stages have few tokens / high
channels; the index_select + scatter bookkeeping outweighs the attention FLOPs.

## 5. The two overhead killers

1. **`prepare()` cost > attention it prunes.**
   - SiTo `prepare()` ≈ 0.67 ms/call (normalize, mean, score matmul, argmax,
     patch layout, similarity recover), vs a dense 2880 attention ≈ 0.42 ms.
   - ITM `prepare()` ≈ **11.7 ms/block** (randperm, full argsort, src×dst
     similarity matmul, scatter). Catastrophic.
   - Mitigation implemented: **plan caching** across denoise steps
     (`plan_recompute_every`; the plan is spatial-structure-driven and stable),
     plus a **single `index_select` recover** (precomputed `gather_from_kept`)
     replacing two scatters + a gather. Cache-once (per episode) lifts the 2880
     layer from 0.42x → **1.26x**; 720/180 stay <1x even cached.

2. **ITM forces `guidance_scale → 2`.** ITM's importance = |cond − uncond| heat
   map needs classifier-free guidance, which doubles the UNet batch. Dense runs
   at guidance_scale=1. This alone makes vanilla ITM ~1.5–2x *slower* than dense
   regardless of pruning. Mitigation: block-hold `self_importance` mode derives
   importance from hidden-state feature L2 magnitude → runs at CFG=1.

## 6. Approaches tried & results

Dense baseline (25 steps): ~156–267 s depending on GPU sharing; PSNR ≈ 21.3.

| Approach | Speedup | PSNR | Note |
|---|---|---|---|
| SiTo per-attn, cache=5 | 0.94x | 19.8 | prepare still too frequent |
| SiTo per-attn, cache-once, ratio1, patch4 | ~1.0x | **21.4** | quality kept, no speedup |
| ITM per-attn (CFG=2) | 0.44x | 20.5 | CFG doubling dominates |
| ITM block-hold, self-importance, CFG=1 | 0.9x | 16–20.7 | attn+FF only, temporal untouched |
| **ST-hold keep0.9, stage2880 only** | ~1.0x | **18.9** | quality-leaning default |
| ST-hold keep0.8, all stages | 1.04x | 16.4 | |
| ST-hold keep0.5, all stages | **1.16x** | 13.7 | fastest, lossy |

### Spatio-temporal prune-once-hold (the only >1.1x path)

Key structural fact: the temporal block reshapes to `(batch·H·W, T, C)`, i.e. its
effective batch = number of spatial tokens. Pruning spatial tokens ONCE at the
`TransformerSpatioTemporalModel` input and holding them compressed through BOTH
the spatial and temporal blocks (unmerge only before `proj_out`) shrinks ~38% of
UNet instead of ~16%. Implemented in
`methods/prunning/SiTo/spatiotemporal_hold.py` (`--sito_spatiotemporal_hold`).

- Real speedup achieved: up to **1.16x** at keep_ratio 0.5.
- **Quality collapses**: pruning before the temporal block corrupts temporal
  dynamics; even keep_ratio 0.9 (drop 10%) gives PSNR ~18.9 and no real speed.
- `similarity_recover=True` (recover pruned token from most-similar kept token,
  cosine) adds ≈ +3.7 PSNR at equal speed vs nearest-index recover — keep it on.
- No config is simultaneously >1.1x and PSNR>18. Clear speed↔quality Pareto,
  no free lunch.

## 7. Ceiling argument (why 16% caps it)

If spatial blocks are sped k× and are fraction f of UNet, UNet speedup =
1/((1−f) + f/k). With f≈0.16 and an optimistic per-block k≈1.2, UNet≈1.03x; with
k≈1.95 (aggressive prune), UNet≈1.08x → episode ≈1.05x (UNet is 64% of wall).
Extending to spatial+temporal (f≈0.38) is what unlocks >1.1x, but at the quality
cost documented above.

## 8. Recommendation

- Token pruning (SiTo/ITM) on Ctrl-World: **quality-preserving ⇒ ~break-even
  speed; real speedup ⇒ large quality loss.** Fundamental to the small SVD UNet.
- For speedup that preserves quality, use **whole-UNet step caching**
  (DiCache/WorldCache) which skips entire UNet forwards (attacks the 64%).
  DiCache @25 steps kept PSNR ≈ 21.2 (≈ dense) in preliminary tests.
- Keep SiTo/ITM as measured benchmark data points (per project policy: run every
  method even when its speedup is weak/negative).

## 9. Reproduce

```bash
# dense baseline (25 steps)
python scripts/infer_single_arm_multiview_ctrlworld.py ... --num_inference_steps 25

# SiTo per-attn, quality-preserving cache-once (break-even speed)
... --use_sito --sito_max_downsample_ratio 1 --sito_plan_recompute_every 999999 \
    --sito_patch_h 4 --sito_patch_w 4 --sito_prune_ratio 0.7

# Spatio-temporal hold (fast, lossy)
... --use_sito --sito_spatiotemporal_hold --sito_st_keep_ratio 0.5 --sito_max_downsample_ratio 4
```

Run scripts: `scripts/run_single_arm_multiview_ctrlworld_{sito,itm}.sh`.
Related rules: `.cursor/rules/ctrlworld-pruning-insights.mdc`,
`.cursor/rules/acceleration-benchmark.mdc`.