SD3.5-Flash: Distribution-Guided Distillation of Generative Flows
Hmrishav Bandyopadhyay, Rahim Entezari, Jim Scott, Reshinth Adithyan, Yi-Zhe Song, Varun Jampani
cs.CV, cs.AI
2025-09-26
Stability AI distills SD3.5 Medium to 2/4 steps via timestep-shared DMD. 8-bit without T5 uses 6.61GB; 4-step is ~18x faster than the teacher and wins a 124-person quality study.
Rectified-flow models such as SD3.5 draw a nearly straight path from noise to pixels. Quality is high; the compute is not. The paper's own bar is 25+ steps, 16GB+ VRAM, 30+ seconds per image. Phones and consumer GPUs cannot host that.
Few-step distillation is the usual way out. Distribution Matching Distillation (DMD) trains a cheap student by matching teacher and student distributions through the gradient of a KL divergence, without walking the full trajectory. That recipe works on ordinary diffusion. On flow models it goes noisy: the standard trick re-noises a finished sample x0 to a random t and then reads scores. Rectified flow pairs noise and image on a straight ODE. Adding fresh noise kicks the sample off that path, so velocity predictions and gradients rot. With only two or four steps left, later iterations cannot clean up the mess. Capacity is also tight, so aesthetics and prompt alignment fight over the same weights. SD3.5 still carries T5-XXL, which dominates peak memory.
The teacher is 2.5B SD3.5 Medium. Training images come from 8B SD3.5 Large, 32 steps, CFG 4.0. On 8 H100s the student is first pretrained for 2,000 iterations (about 17 hours).
A trajectory-guidance loss comes first. At the student's own timesteps (four points for the 4-step model), the student velocity is trained to match the teacher's integrated displacement along that segment. That parks the student near the teacher ODE before DMD starts.
DMD is then rewritten with timestep sharing. Instead of re-noising x0 to a random t, matching runs on the partially denoised xt already sitting on the student trajectory. A few-step student estimates x0 poorly at high noise; fake endpoints poison the score. Sharing timesteps means alignment only happens at t values the sampler will actually visit. Timestep diversity drops; ablations still prefer the textures and layouts.
Prompt alignment is patched with split-timestep fine-tuning. The checkpoint is cloned into two branches, one on t in (0, 500], one on (500, 1000], each with EMA decay 0.99. After convergence the branches are interpolated 3:7 (low-noise : high-noise), a ratio picked by GenEval. This is used only for the 4-step model; the 2-step run did not show the same jump.
The adversarial term follows NitroFusion: the proxy student is a feature extractor, with 8-layer MLP heads on MM-DiT layers 3, 4, 5, 6, 8, 10 and 11. Teacher synthetics count as real. Heads reset with probability 0.005 per iteration. The discriminator and proxy take 10 updates per generator update. The 2-step student continues from the 4-step checkpoint and adds an MSE on Gram matrices of teacher and student features. Appendix numbers: stage two of the 4-step run is 800 iterations (6 hours), 2-step is another 1,200 (9 hours), split-timestep another 400 (4 hours).
For deployment, T5-XXL is replaced by null embeddings (SD3.5 was pretrained with encoder dropout). CLIP-L and CLIP-G stay. MM-DiT is quantized from 16-bit to 8-bit, then to 6-bit via CoreML on Apple silicon, with RMSNorm rewritten so the Neural Engine keeps precision.
COCO-30K automatic metrics on an RTX 4090:
| Method | Steps | Latency | Peak VRAM | ImageReward | GenEval |
| SD3.5M teacher | 50 | 10.58s | 19.47GB | 0.91 | 0.64 |
| SWD-M | 4 | 0.66s | 17.88GB | 1.12 | 0.72 |
| Flash 16-bit+T5 | 4 | 0.58s | 17.58GB | 1.10 | 0.70 |
| Flash 8-bit, no T5 | 4 | 0.61s | 6.61GB | 1.08 | 0.68 |
| SANA-Sprint 1.6B | 2 | 0.24s | 10.17GB | 1.01 | 0.73 |
| Flash 16-bit+T5 | 2 | 0.39s | 17.58GB | 1.00 | 0.70 |
Against the 50-step teacher, 4-step 16-bit is about 18x faster (10.58s / 0.58s). ImageReward rises from 0.91 to 1.10, GenEval from 0.64 to 0.70, Aesthetic Score from 5.99 to 6.38. FID goes the other way: Flash 29.80, teacher 20.06, SDXL-DMD2 16.64. The paper blames the teacher family. SD3.5M already has a worse FID than SDXL, and sibling distillations SWD and Turbo sit high as well.
The same-family 4-step baseline SWD-M scores GenEval 0.72 and IR 1.12, both a notch above Flash. At 2 steps, SANA-Sprint 1.6B posts GenEval 0.73, IR 1.01, 0.24s; Flash 16-bit with T5 is 0.70 / 1.00 / 0.39s. Sprint wins the automatic table.
A user study used 124 annotators who already rank generated images, 507 prompts, 4 seeds, 3 votes per pair. Quality and prompt adherence were split: quality votes hid the prompt; adherence votes ignored aesthetics. Flash beat other few-step models and the 50-step teacher on quality. Adherence differed by less than ±1.6% across methods, a near tie. Flash variants sit at the top of the ELO chart at every consumer compute budget; the paper does not print the numeric ELO values.
On-device 4-step, 6-bit, no T5: iPhone A17 takes 8.32s at 768px and 3.25s at 512px; iPad M4 takes 6.44s and 2.62s. Full 16-bit 1024px on a 32GB M3 MacBook Pro takes 18.65s.
Ablations are qualitative. Drop the adversarial term and training falls apart. Skip pretraining and color plus layout go first. Random DMD timesteps hurt texture and composition. No discriminator refresh yields over-smooth images. There is no numeric ablation table.
For anyone who wants the SD3.5 look on a consumer GPU or a phone, this is a deployable path. Four steps, drop T5, 8-bit: peak VRAM falls from about 18GB to 6.61GB, and a 4090 writes a 1024px image in about 0.6s. Timestep sharing is a targeted fix for DMD on straight flows; later distillations of SD3 or Flux can reuse it. Split-timestep adds capacity only during training.
This is not a new few-step SOTA. SWD-M is slightly ahead on GenEval and IR. SANA-Sprint wins the 2-step automatic metrics. The product is a memory ladder and an on-device latency kit for the SD3.5 family, not a single leaderboard number.
The paper says the quiet part: distillation trades quality and diversity on hard prompts, and dropping T5 makes complex layouts worse. FID on COCO is clearly worse than SDXL-family distillations. Even if that is teacher DNA, the model is skewed toward aesthetic scores rather than the COCO photo distribution.
The user study split quality from adherence. Quality won; adherence tied. That sits awkwardly next to the qualitative claim of stronger composition. Ablations have no numbers. SD3.5 Large and Flux.1-dev are omitted because the teachers would not fit consumer hardware, so the few-step field still lacks that comparison. Training images come from 8B Large while the teacher is 2.5B Medium; some of the student's skill may be the Large data distribution, not a pure Medium distill.