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

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:

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:

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

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

4.5 Drop-count sweep + per-block FFN placement

The recovery curves for every surrogate/drop config:

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:

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

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