Distilling Stable Diffusion into smaller, faster replicas

What it does

Segmind Distilled Diffusion is an unofficial implementation of the BK-SDM paper that trains shrunken Stable Diffusion U-nets—dubbed sd_small and sd_tiny—using knowledge distillation. A full-sized teacher model (Realistic Vision V4.0) guides the student through MSE losses on final noise predictions and intermediate block outputs, producing checkpoints that plug directly into the Hugging Face diffusers library. The sd_tiny variant lands at roughly 323 million parameters versus the base U-net’s 860 million, while the project claims up to double the inference speed and 30% lower VRAM use on an NVIDIA A100.

The interesting bit

The project does not rewrite the diffusion pipeline; it swaps in a smaller-capacity U-net and retrains it to parrot the teacher at every block level. That makes the boring part—the loss function, a weighted sum of output and feature-level MSE—the actual engine of compression, effectively treating model slimming as a supervised regression task against the original.

Key highlights

• Two shrunken U-net variants: sd_small (~579M params) and sd_tiny (~323M params) versus the base ~860M.
• Training loss combines task MSE with output-level and per-block feature-level distillation losses against the teacher.
• Pretrained checkpoints (segmind/small-sd, segmind/tiny-sd, and a portrait-tuned tiny model) are ready for standard diffusers inference.
• Claims up to 2× faster inference and 30% lower VRAM footprint on an A100, plus faster LoRA and DreamBooth fine-tuning.
• Training script extends the standard diffusers text-to-image finetuning loop rather than reinventing it.

Caveats

• The authors warn the distilled models are early-phase and not yet production quality.
• Composability and multi-concept generation are currently weak spots; the models work best when further fine-tuned or LoRA-trained on specific styles.
• Checkpoints produced by the training script ship with a broken config.json that must be manually replaced before reloading.

Verdict

Worth a look if you need a drop-in, smaller Stable Diffusion backbone for style-specific LoRA work and can tolerate some generalization trade-offs. Skip it if you need a general-purpose, out-of-the-box foundation model that handles complex multi-concept prompts reliably.