report/REPORT.md

% FLUX.2 [klein] 4B → Compressed Student — Distillation Deep-Dive Report
% Single-day ablation study · 1× A100-80GB
% 2026-05-31

---

# Executive Summary

We set out to compress **FLUX.2 [klein] distilled 4B** (a 4-step, CFG-free rectified-flow MM-DiT)
into a smaller, faster text-to-image student via **depth-pruning + warm-started surrogate blocks +
short distillation recovery**, on a single A100-80GB. Over one day we ran a full ablation sweep:
block-selection strategies, four generations of surrogate architecture, recovery-recipe fixes
(including a divergence we diagnosed and corrected), an optimizer A/B, a drop-count sweep, and a
measured compute breakdown that re-pointed the strategy.

**Headline result — the speed↔quality frontier:**

| Config | Params | Smaller | Wall-clock speedup | Transformer-FLOP | Eval-loss | Verdict |
|---|---|---|---|---|---|---|
| teacher 4B | 3.876B | — | 1.00× | 1.00× | — | reference |
| v1 per-token, drop-12 | 2.441B | 37% | ~1.45× | ~1.64× | — | **collapsed** |
| per-token, drop-6 | 3.158B | 19% | 1.19× | 1.24× | 0.308 | ok, soft |
| linattn drop-6 (simple) | 3.177B | 18% | 1.15× | 1.23× | 0.253 | ok |
| **linattn drop-6 (RoPE+conv+warmstart)** | 3.177B | 18% | 1.15× | 1.23× | **0.231** | **best quality** |
| linattn drop-8 (+focused+FFN) | 2.995B | 23% | 1.20× | 1.28× | 0.269 | best perceptual |
| linattn drop-10 (mixed FFN) | 2.737B | 29% | 1.26× | 1.42× | ~0.322 | aggressive |

*Eval-loss = held-out velocity-matching loss vs the teacher (lower = closer to teacher).*

**Three load-bearing conclusions:**

1. **A per-token surrogate cannot replace an attention block** (it can't mix tokens) → aggressive
   pruning collapses. **Linear-attention surrogates** (token-mixing) recover far better and are the
   key unlock.
2. There is a **steep speed↔quality frontier**: best quality (0.231) sits at only 1.15× wall-clock;
   pushing to 1.26× costs ~40% more loss. The frontier is set by **how many** single blocks you
   remove, not surrogate cleverness.
3. The **5 double-stream blocks are only ~21% of compute** (measured) and do the fragile cross-modal
   binding — a bad target. **Single blocks are 78% of compute** and are the correct lever, but
   they're now near their useful limit. The real path to ~2× is **step reduction (4→2)**, a separate
   project.

---

# 0. Block-Indexing Convention (read this first)

Throughout this report, to avoid the ambiguity that bit us a few times:

> **Single-stream blocks are numbered `S0 … S19` (20 blocks), 0-indexed, counting from the START of
> the single-stream stack** — i.e. **`S0` = first / shallowest** (immediately after the 5 double
> blocks), **`S19` = last / deepest** (immediately before the output projection).
>
> **Double-stream blocks are `D0 … D4` (5 blocks)**, same convention (`D0` = first).
>
> Data flows: `input → D0…D4 (double) → S0…S19 (single) → norm_out → proj_out → output`.
> **"Deepest N"** always means the **highest-numbered** blocks (e.g. deepest-4 of the single stack =
> `S16,S17,S18,S19`). **Lower importance = safer to drop.**

Every experiment below lists its dropped / FFN block sets explicitly by these IDs.

---

# 1. Objective & Setup

- **Teacher:** `black-forest-labs/FLUX.2-klein-4B` — the *distilled* checkpoint (4-step, guidance-
  distilled → CFG-free, `is_distilled=true`), Apache-2.0. We deliberately use the distilled (not the
  50-step base) teacher: pure velocity matching cannot compress step count, so a 4-step student needs
  a 4-step teacher.
- **Hardware:** 1× A100-80GB. **No FlashAttention** (Blackwell-only path) → PyTorch SDPA. This is a
  prototyping rig; the design is meant to lift onto a B200 for the full run.
- **Goal:** a smaller / faster student that stays *recoverable*, plus a clean map of the
  size/speed/quality tradeoffs.
- **Eval metric:** a fixed held-out batch (16 cached latents) on which we measure the **velocity-
  matching loss** between student and teacher. It is a clean, apples-to-apples proxy; it tracks
  perceptual quality well but not perfectly (see §6.6).

---

# 2. Architecture (verified from the real config + measured)

`Flux2Transformer2DModel`: **d = 3072** (24 heads × 128), **5 double-stream + 20 single-stream**
blocks, `mlp_ratio=3`, `joint_attention_dim=7680` (Qwen3-4B conditioning), `in_channels=128`
(packed VAE latent), `guidance_embeds=false`. Transformer = **3.876B params**. VAE
`AutoencoderKLFlux2` (32 latent channels, batch-norm-normalized) decodes **once** at the end; the
**transformer runs 4×** (once per step); the **Qwen3-4B text encoder runs once**.

## 2.1 Double vs single blocks — internals & measured compute

**Double block (`Flux2TransformerBlock`, 245.4M):** an MM-DiT **dual-stream** block — it keeps image
and text as **separate streams**, joined only in attention:

```
attn (75.5M)     joint cross-modal attention: image to_q/k/v/out (4×d²) +
                 text add_q/k/v_proj/to_add_out (4×d²); both streams' Q/K/V concatenated
                 → every image token attends to every text token (and vice-versa)
ff      (84.9M)  image-stream FFN (its own)
ff_context(84.9M)text-stream FFN (its own, separate)
```
This is **where cross-modal binding happens** ("a *red* cube on a *blue* sphere").

**Single block (`Flux2ParallelSelfAttention`, 122.7M):** text+image **merged into one stream**, a
fused parallel attention+MLP with shared params over all 1536 tokens. These mostly **refine**.

**Measured compute split (one transformer forward, batch-2):**

![Compute breakdown + warm-start quality](/root/workspace/outputs/report_assets/breakdown.png)

```
full transformer forward:  136 ms
20 SINGLE blocks:          106 ms  → 78% of compute
 5 DOUBLE blocks:           28 ms  → 21% of compute
per-block: double 5.71 ms vs single 5.30 ms → only 1.08× (NOT 2× despite 2× params)
```

**Why params ≠ compute:** a double block's 245M params are *split across two shorter streams* (image
FFN on ~1024 tokens, text FFN on ~512), so it does roughly the same work as a single block's 122.7M
running on all 1536 tokens. **Takeaway: the 20 single blocks dominate compute → they are the correct
surgery target.** The 5 double blocks are a small, fragile slice (see §6.5).

---

# 3. Methodology

- **Surgery:** select least-useful single blocks, replace each with a cheaper **surrogate** that
  preserves the residual shortcut (`out = x + surrogate(x)`); keep the rest full.
- **Block selection:** v1 used **SVD-energy** of each block's residual (picked a bad contiguous set);
  v2+ used **leave-one-out importance ablation** — skip each block, measure the relative change in the
  final latent, drop the lowest-impact blocks. Importance is much better (§4.2).
- **Surrogate types (evolved across the day):** per-token low-rank → linear-attention → +RoPE +
  depthwise-conv + warm-start → +focused feature map + FFN (§4.4).
- **Recovery:** **freeze the entire pretrained network, train only the surrogates** (~0.6–3% of
  params), AdamW @1e-4 with cosine-to-floor LR + grad-clip + fp32 master on the trained params, bf16
  autocast compute. Primary loss = velocity matching to the frozen teacher (+ a light real-data
  flow-matching grounding term).

---

# 4. Experiments & Ablations

## 4.1 Surgery v1 — per-token low-rank, drop-12 → COLLAPSE

| | |
|---|---|
| Surrogate | per-token low-rank `x + B·σ(A·x)`, rank 512 |
| **Dropped (SVD-energy)** | **S0–S11** (first 12) |
| Kept full | S12–S19 |
| Params | 2.441B (−37%) |

A 7-agent visual-review workflow scored the warm-start student **5/5 severity (destroyed) in every
category** — flat color fields, no structure or text:

![v1 collapse: teacher vs per-token drop-12](/root/workspace/outputs/report_assets/v1_collapse.png)

**Root cause (the key architectural finding):** a per-token low-rank map is applied **independently
per token** — it *cannot move information between token positions*. But a dropped block was an
**attention** block, whose whole job is token-mixing. So the surrogate can only mimic the MLP-ish
per-token part (warm-start residual rel-err stuck at ~0.9 ≈ "predict nothing"). Removing 12 of 20
token-mixing layers severs the chain → collapse. **This single finding redirected the entire project
toward token-mixing surrogates.**

## 4.2 Surgery v2 — importance-selected drop-6 → FUNCTIONAL

| | |
|---|---|
| Surrogate | per-token low-rank (same as v1) |
| **Dropped (importance)** | **S12–S17** (6 least-important, middle-deep) |
| Kept full | S0–S11, S18, S19 |
| Params | 3.158B (−19%) |

Switching to importance ablation + dropping only 6 gave a **functional** (if soft) student even with
near-identity surrogates. Importance also revealed **low redundancy**: *every* single block shifts the
final latent by ≥43% when skipped — there are no "free" blocks.

## 4.3 Recovery recipe — divergence → fix → optimizer A/B

**Failure first.** Our initial recovery trained **all** weights with Muon at lr 0.02. It **diverged**:

![Divergence: training all weights](/root/workspace/outputs/report_assets/divergence.png)

Loss was stable ~0.83 to step ~120, then blew up to **8.08**; samples went to cyan checkerboard noise.
**Diagnosis:** Muon's 0.02 is a *bulk-pretraining* LR; applied to the well-pretrained **kept** blocks
it corrupts them faster than the surrogates can learn. *Never train the frozen-worthy weights.*

**The fix.** Freeze the whole pretrained net; train **only the surrogates** (≈0.6% of params); AdamW
@1e-4 (the diffusion/adapter LR regime), cosine LR with a floor, grad-clip 1.0, fp32 master on the
trained params. Result: **0.476 → 0.308 (−35%)**, monotonic, zero divergence, 28.7 GB (vs 53 GB
training all).

**Optimizer A/B (controlled — identical recipe, only the optimizer swapped):**

![Muon vs AdamW](/root/workspace/outputs/report_assets/muon_vs_adamw.png)

AdamW@1e-4 → **0.3081**, Muon@2e-3 → **0.3056** — a **statistical tie**. The earlier blow-up was *not*
Muon; it was lr-on-the-wrong-params. For a handful of adapter modules, AdamW is the simpler right tool;
Muon is reserved for a later full-network recovery.

## 4.4 Surrogate evolution — the heart of the study

We rebuilt the surrogate four times, each time at the same dropped positions (**S12–S17**) for a
controlled comparison, then swept drop-count.

**(a) Linear-attention surrogate** — O(N) linear attention (`elu+1` feature map), multi-head. Unlike
per-token, it **moves information between tokens**. Drop-6 → **0.253**, and it *crossed the per-token
surrogate's converged floor by step 50*. This validated token-mixing as the unlock.

**(b) + RoPE + depthwise-conv + warm-start** — RoPE (head_dim 128, the model's rotary) for spatial
awareness; a depthwise-conv local branch for high-frequency detail; and a **warm-start** that fits
each surrogate to mimic its teacher block. Drop-6 → **0.231 (best quality)**. Warm-start cut residual
rel-err **1.0 → ~0.65** (the per-token surrogate never got below ~0.9 — see breakdown graph, §2.1).

**(c) + focused feature map + FFN** — a **focused** (FLatten) feature map with a *learnable* sharpening
exponent makes linear attention's weights peaky/softmax-like instead of blurry; an **FFN** makes the
surrogate a real linear-attn *block*. Warm-start improved further to **0.47–0.59**. This version gave
the **best perceptual colors / local detail** (your eye preferred it even though its loss number is
higher — the metric doesn't fully capture perceptual sharpness).

**Per-surrogate-generation comparison (drop-6, final samples):**

![Per-token vs linear-attn](/root/workspace/outputs/surrogate_ab.png)

![Simple vs RoPE+conv+warmstart linear-attn](/root/workspace/outputs/linattn_upgrade_ab.png)

## 4.5 Drop-count sweep + per-block FFN placement

The recovery curves for every surrogate/drop config:

![Recovery curves](/root/workspace/outputs/report_assets/recovery_curves.png)

**Exact block tracking per run** (single-block IDs, `S0`=first … `S19`=last):

| Run | Surrogate | Dropped block IDs | FFN on (IDs) | Kept full | Params | Speedup | Loss |
|---|---|---|---|---|---|---|---|
| v1 | per-token | **S0–S11** | none | S12–S19 | 2.441B | ~1.45× | collapsed |
| v2 | per-token | **S12–S17** | none | S0–S11,S18,S19 | 3.158B | 1.19× | 0.308 |
| linattn drop-6 simple | linattn (elu+1) | **S12–S17** | none | S0–S11,S18,S19 | 3.177B | 1.15× | 0.253 |
| linattn drop-6 upgraded | linattn+RoPE+conv+ws | **S12–S17** | none | S0–S11,S18,S19 | 3.177B | 1.15× | **0.231** |
| linattn drop-8 | linattn+focused+FFN | **S10–S17** | **S10–S17 (all 8)** | S0–S9,S18,S19 | 2.995B | 1.20× | 0.269 |
| linattn drop-10 mixed | linattn+focused | **S8–S17** | **S14,S15,S16,S17 (deepest 4)** | S0–S7,S18,S19 | 2.737B | 1.26× | ~0.322 |

*Note the importance ablation consistently picks the **middle-deep** single blocks (S8–S17) as least
important; the very last two (S18, S19) and the early ones (S0–S7) are always kept.*

**Drop-8 vs drop-10 final samples:**

![drop-6 upgraded vs drop-8](/root/workspace/outputs/drop6_vs_drop8.png)

**On FFN placement (drop-10):** we put the FFN on the **deepest 4 dropped blocks (S14–S17)** and left
S8–S13 light. Rationale: deeper blocks have more direct influence on the output, and the *last* dropped
block should always carry an FFN. Warm-start confirmed it — the FFN blocks (S14–S17) reached rel-err
0.54–0.56 vs 0.61–0.70 for the light ones.

---

# 5. Results & The Frontier

![The speed-quality frontier](/root/workspace/outputs/report_assets/frontier.png)

- **Best quality:** drop-6 + RoPE+conv+warmstart → 0.231, 3.18B, 1.15× wall (~1.23× FLOP).
- **Best perceptual / balanced:** drop-8 +focused+FFN → 0.269, 3.00B, 1.20×.
- **Most compression/speed:** drop-10 mixed → ~0.322, 2.74B, 1.26× wall (1.42× FLOP).
- **Failed:** v1 per-token drop-12 (collapsed) — same aggressive drop a token-mixing surrogate might
  survive, but per-token cannot.

**Wall-clock understates the real compression.** At batch-1/512px, the Qwen3 text-encode (once) + VAE
decode (once) are fixed overhead (~22% of wall-clock) that no denoiser surgery touches. The
**transformer-FLOP column** is what the target hardware (B200, larger batch, cached text embeddings)
would approach — e.g. drop-10's 1.26× wall ≈ **1.42× on the real run**.

---

# 6. Key Findings & Lessons

1. **Function class is everything.** A per-token surrogate can't do attention's token-mixing → it
   caps recovery and collapses under aggressive pruning. Linear attention restores it (§4.1, §4.4).
2. **Steep speed↔quality frontier set by drop-count.** 0.231@drop-6 → 0.322@drop-10. The surrogate
   architecture shifts the curve a little; *how many blocks you remove* sets where you are on it.
3. **Recovery recipe:** freezing the base is **mandatory** (training all weights diverged to noise);
   AdamW ≈ Muon for adapter-style recovery; LR must be the adapter regime (~1e-4), not Muon's 0.02.
4. **LR schedule:** decaying cosine to **exactly 0** wastes the tail and breaks resume; we floor it at
   **30% of base** (1e-4 → 3e-5) so it decays visibly and settles (a constant LR jittered at the floor).
5. **Single blocks (78% of compute) are the target; the 5 double blocks (21%) are not.** Per-block,
   double ≈ single in speed (1.08×); the double blocks do fragile cross-modal binding (the first thing
   to degrade) — a bad place to cut for little gain (§6.5).
6. **The eval metric ≠ perception.** drop-8 had a *higher* velocity-loss (0.269) but the *best* colors
   / local detail by eye — the focused feature map + FFN improved perceptual quality the metric
   under-credits.
7. **Warm-start is limited by sequential mismatch.** We fit each surrogate on the *teacher's* block
   inputs, but at inference downstream surrogates see *student* inputs — so the local rel-err gains
   (→0.5) don't fully transfer end-to-end. A **progressive** warm-start would fix it.
8. **FFN buys quality, not speed.** It improved warm-start and perception but is heavy — at drop-8 it
   cut the speedup back to ~drop-6 levels. Capability (FFN) and speed (more drops) pull against each
   other.

## 6.5 Why we did NOT touch the double blocks (measured)
Shrinking the FFN on 2 of the 5 double blocks by 50% would touch `2/5 × 21% × ~60%(FFN) × 50%` ≈ **2%
of compute → ~1.02× speedup.** Even all five → ~1.05×. The leverage isn't there, and the cross-modal
attention is the most fragile thing in the model. Rejected on the math.

---

# 7. Speculations & Ideas (explored + future)

**Surrogate-side (to push the frontier):**
- **Gating (GLA-style)** data-dependent forget gates; **width** (heads × head_dim) as a capacity dial;
  **2D depthwise conv** on the image-token grid (vs current 1D over the sequence); **Hedgehog**
  learnable feature map trained to mimic softmax.
- **Progressive/sequential warm-start** (fit S-i, re-capture, fit S-(i+1)…) to fix the §6.7 mismatch.

**Training-signal upgrades (explained, not yet run):**
- **Trajectory velocity matching on the 4 schedule σ's** (vs current σ~U(0,1)): instead of matching at
  random noise levels off real-image interpolants, roll out the teacher's *actual* 4-step sampler and
  match at the 4 on-trajectory points the student will really visit — kills few-step exposure bias.
- **Feature matching on retained blocks (masked KD):** also match *intermediate* hidden states of the
  kept blocks (denser supervision; no projector since widths match), masking the few extreme "massive
  activations" so they don't dominate the loss.

**The real 2× levers (different axes):**
- **Step reduction (4→2)** — halves *all* transformer passes at once; needs step-distillation (DMD /
  consistency). **This is the biggest lever and the recommended next project.**
- **Text-encoder caching / batching** — removes the fixed ~22% wall-clock overhead.
- **B200 + FlashAttention + fp8 + larger batch + 300k-image data + offline latent shards** (the
  original plan's scale-up).

---

# 8. Recommendations / Where Things Stand

This was not "it'll never work" — we **mapped a real tradeoff** and built a working surgery + recovery
+ eval pipeline.

- **If you want one balanced single-block model:** lock **drop-7 with 3 FFN** (FFN on the deepest 3
  dropped incl. the last) — projected ~3.08B, ~1.17× wall, loss ~0.25 (the midpoint of drop-6/drop-8).
- **If quality is paramount:** drop-6 + RoPE+conv+warmstart (0.231, 1.15×).
- **Single-block surgery is near its useful limit** (~1.15–1.26× on this setup); squeezing further is
  diminishing returns.
- **For a genuine ~2×:** the next project is **step reduction (4→2 via step-distillation)**, plus
  text-encoder caching and the B200 scale-up — *not* more block surgery.

---

# Appendix A — Reproduction

```
scripts/01_inspect_model.py        # verify architecture
scripts/02_teacher_smoke.py        # teacher 4-step generation
scripts/03_build_student.py        # v1 surgery (SVD-energy, per-token, drop-12)
scripts/05_build_student_v2.py 6   # v2 surgery (importance, per-token, drop-6)
scripts/09_build_linattn.py 10 4   # linattn build: drop_k=10, FFN on deepest 4
scripts/08_train_recover.py 300 adamw 1e-4   # surrogate-only recovery (env: STUDENT_DIR, SCHED, MB)
scripts/10_bench.py                # inference speedup
scripts/make_report_assets.py      # graphs + montages
```
Data: 400 monet images cached @512 (`scripts/06_cache_data.py 400`). Held-out eval batch = first 16.

# Appendix B — Surrogate config (final linear-attn block)
`heads=4 × head_dim=128 (=512 inner, RoPE-compatible)`, `elu+1`→**focused** with learnable exponent,
`depthwise-conv kernel=5` local branch, optional **FFN** (`hidden=1024`) placed per-block via `ffn_idx`.
Output zero-init = identity start; warm-started to the teacher block. Per-block params: ~6.3M light,
~12.6M with FFN.