block_surgery_todo.md

# TODO — surrogate strengthening & roadmap

Backlog of knobs/experiments for the FLUX.2 klein 4B→~2B distillation. Ordered roughly by
ROI for image distillation. Status: [ ] todo · [~] in progress · [x] done.

## Surrogate capability knobs (linear-attention)
- [x] **RoPE positional encoding** on q,k (head_dim=128) — DONE. Helped text/structure legibility.
- [x] **Depthwise-conv local branch** on v — DONE (FLatten-style). In the upgraded run.
- [x] **Warm-start = mimic the teacher block** — DONE (RoPE-aware seq fit; residual rel-err 1.0→~0.65).
      NOTE: one-shot fit on *teacher* inputs → limited end-to-end benefit (sequential mismatch).
      Upgraded run (RoPE+conv+warmstart, drop6) floored at **0.231 vs simple linattn 0.253 (9% lower)**.
      Confound: same run also doubled data (400) + batch (16) + cosine-30% LR — not surrogate-only.
- [ ] **Sequential/progressive warm-start** (fit blk12, re-capture, fit blk13, …) to fix the mismatch.
- [x] **Feature map φ**: `elu+1` → **focused with learnable exponent** (FLatten) — DONE.
- [x] **Make it a real block**: added an FFN (ffn_hidden=1024) — DONE. Improves warm-start
      (rel-err → 0.47–0.59) BUT is heavy: at drop-8 it cut the speedup to 1.20× (~= drop-6).
      **Lesson: FFN buys quality at fixed drop, but kills the speed-via-more-drops strategy.**
- [ ] **Gating (GLA-style)** data-dependent forget gates — selectivity. Small cost.
- [ ] **Width** (heads × head_dim) — raw capacity dial. Cost linear in width.
- [ ] **Stack depth** (2 linattn ops) — alternative to FFN for capacity.
- [ ] **2D depthwise conv on the image-token grid** (vs current 1D over the sequence) — better
      spatial locality; needs text/image token split + 32×32 reshape.

## Surgery / compression frontier
MEASURED (eval velocity-loss floor, 512/4-step, 300 steps):
  teacher 4B            3.876B  1.00x   —
  per-token  drop6      3.158B  1.19x   0.308
  linattn    drop6      3.177B  1.15x   0.253  (RoPE+conv+warmstart: 0.231)
  linattn+FFN drop8     2.995B  1.20x   0.269
- [ ] **Speed needs LIGHT surrogates + more drops, NOT FFN.** drop-8+FFN proved adding capability
      while dropping more nets ~same speed (1.20×) + worse quality. For ~1.45×: drop 10–12 with a
      **light** surrogate (linattn+RoPE+conv+warmstart, NO FFN / tiny FFN).
- [ ] Re-run **importance ablation** per drop_k; never drop the un-prunable 5 double blocks (the
      ~31% FLOP floor) — revisit only if a much bigger speedup is needed.

## Training recipe
- [x] Freeze base, train surrogates only; AdamW≈Muon (tie); fp32 master; grad-clip.
- [x] **Cosine LR floor = 30% of base** (was 0/15%) so the LR visibly decays and settles — DONE.
- [ ] Size `STEPS` to the data (avoid dead cosine tail).
- [ ] Trajectory velocity matching on the 4 schedule sigmas (vs current σ~U(0,1)).
- [ ] Feature matching on retained blocks (masked KD).

## Scale-up (B200 target)
- [ ] FlashAttention-4 / FlexAttention, torch.compile, larger batch, 300k data + offline latent shards.

## Session outcome 2026-05-31 — full deep-dive report at `report/REPORT.pdf`

Block IDs: single-stream **S0..S19** (S0 = first/shallowest after the 5 double blocks; S19 = last/
deepest before output). "Deepest N" = highest IDs. Lower importance = safer to drop.

Frontier (measured, eval velocity-loss floor, 512/4-step):
```
  per-token  drop6  (drop S12-17)                     3.158B  1.19x  0.308
  linattn    drop6  (drop S12-17) simple              3.177B  1.15x  0.253
  linattn    drop6  (drop S12-17) +RoPE+conv+warmstart 3.177B 1.15x  0.231  <- best quality
  linattn    drop8  (drop S10-17) +focused+FFN(all 8) 2.995B  1.20x  0.269  <- best perceptual
  linattn    drop10 (drop S8-17)  FFN on S14-17       2.737B  1.26x  0.322  <- most compression
  v1 per-token drop12 (drop S0-11)                    2.441B ~1.45x  COLLAPSED
```
- [x] Measured compute split: **single blocks 78%, double 21%** (per-block 1.08x). Double-block FFN
      shrink ≈ **2% speedup → REJECTED**; double blocks = fragile cross-modal core, don't cut.
- [ ] **Recommended next single-block deliverable:** drop-7, FFN on the deepest 3 (incl. last, S17-ish).
      Projected ~3.08B, ~1.17x wall, loss ~0.25 (the drop-6/drop-8 midpoint).
- [ ] **The real ~2x lever (separate project):** step reduction 4→2 via step-distillation (DMD/consistency).
- [ ] Free one-time wins later: cache/batch the Qwen3 text-encode; B200 + FlashAttention + fp8 + 300k data.