Research map / Image generation / Text-to-image diffusion
KOALA: Empirical Lessons Toward Memory-Efficient and Fast Diffusion Models for Text-to-Image Synthesis
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Distills SDXL into compact latent U-Nets (KOALA-1B and KOALA-700M) using self-attention feature distillation, high-resolution long-caption data (LAION-POP), and step-distilled teachers, yielding fast 1024px generation on 8GB GPUs.
Problem
Text-to-image diffusion models like SDXL have large U-Nets that demand expensive GPUs with large memory and are slow at inference, and prior block-removal distillation (BK-SDM) yields limited compression on SDXL while feature-level distillation remains underexplored.
Approach
Compress the SDXL U-Net via layer-wise removal of transformer layers at the lowest resolution plus block removal, and distill self-attention features (rather than the last-layer feature used by BK-SDM), combined with high-resolution long-caption data and step-distilled teachers (SDXL-Turbo/Lightning).
Key results
- KOALA-Lightning-700M achieves about 4x faster latency than SDXL-Base-1.0 (0.655s vs 3.229s at 1024px on RTX 4090) with HPSv2 31.50 vs 30.82 and CompBench 0.4505 vs 0.4445.
- Self-attention feature distillation gives the largest HPSv2 gain among feature types (26.74 vs 25.53 baseline, vs 26.63 for last feature).
- KOALA models generate 1024px images on an 8GB GPU (3060Ti) where SDXL-Base cannot run.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Distill self-attention features instead of relying on task loss alone (ablation) | KOALA-1B · Feature KD ablation | HPSv2: 25.53 -> 26.74 | Helped |
| Distill cross-attention, residual, FFN, or last features instead of task loss alone (ablation) | KOALA-1B · Feature KD ablation | HPSv2: 25.53 -> 26.11; 25.53 -> 26.27; 25.53 -> 26.48; 25.53 -> 26.63 | Helped |
| Apply self-attention KD at decoder stage UP-1 only rather than task loss alone (ablation) | KOALA-1B · Feature KD stage ablation | HPSv2: 25.53 -> 26.52 | Helped |
| Apply self-attention KD at encoder stages DW-2 or DW-3 only rather than task loss alone (ablation) | KOALA-1B · Feature KD stage ablation | HPSv2: 25.53 -> 25.32; 25.53 -> 25.57 | Mixed |
| Use SA-bottom rather than SA-interleave or SA-up feature locations (ablation) | KOALA-1B · SA location ablation | HPSv2: 26.74 vs 26.58 and 26.48 | Helped |
| Add last-feature KD at DW-1 and UP-3 to SA KD (ablation) | KOALA-1B · Feature KD combination ablation | HPSv2: 26.74 -> 26.98 | Helped |
| Add residual-feature KD at DW-1 and UP-3 to SA KD (ablation) | KOALA-1B · Feature KD combination ablation | HPSv2: 26.74 -> 26.94 | Helped |
| Add LF or Res feature KD at all stages rather than only at SA-absent stages (ablation) | KOALA-1B · Feature KD combination ablation | HPSv2: 26.98 -> 26.83; 26.94 -> 26.80 | Hurt |
| Distill SA, CA, Res, FFN, and LF at all stages (ablation) | KOALA-1B · Feature KD combination ablation | HPSv2: 26.74 -> 26.39 | Hurt |
| Add synthetic detailed captions to LAION-A-6+ (ablation) | KOALA-700M · Training data comparison | CompBench: 0.3791 -> 0.4168 | Helped |
| Use high-resolution, longer-caption LAION-POP instead of LAION-A-6+ with synthetic captions (ablation) | KOALA-700M · Training data comparison | CompBench: 0.4168 -> 0.4290 | Helped |
| Use SDXL-Turbo instead of SDXL-Base as the teacher (ablation) | KOALA-700M · Teacher model comparison | CompBench: 0.4290 -> 0.4470 | Helped |
Surprising findings
- Using SDXL-Turbo or SDXL-Lightning as teachers gives comparable or better student quality than SDXL-Base with 25 steps, despite fewer teacher denoising steps.
- Fewer high-resolution images with long captions (LAION-POP, 491K) outperform 8M low-resolution images with short captions on HPSv2 and CompBench.
Limitations
- Struggles with rendering legible long text, human hands, and complex multi-attribute prompts.
- Trained only on LAION-POP (relatively small, 491K samples), which may explain failure modes.
- Direct comparison with SSD models is confounded by different training datasets.
Experimental setup
Datasets: LAION-POP LAION-Aesthetics V2 6+ LAION-Aesthetics V2 6.5+ synCap (LLaVA-1.5 synthetic captions)
Benchmarks: HPSv2 T2I-CompBench
Metrics: HPSv2 score CompBench average score latency (s) memory (GB) U-Net parameters
Baselines: SDXL-Base-1.0 SDXL-Turbo SDXL-Lightning SSD-1B SSD-Vega Pixart-alpha Pixart-sigma SDM-v2.0 BK-SDM
Builds on: SDXL BK-SDM SDXL-Turbo SDXL-Lightning PCM
Closest papers by meaning
- PixArt-$\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis (ICLR 2024 · similarity 0.90)
- SnapGen: Taming High-Resolution Text-to-Image Models for Mobile Devices with Efficient Architectures and Training (CVPR 2025 · similarity 0.90)
- Würstchen: An Efficient Architecture for Large-Scale Text-to-Image Diffusion Models (ICLR 2024 · similarity 0.90)
- Any-Size-Diffusion: Toward Efficient Text-Driven Synthesis for Any-Size HD Images (AAAI 2024 · similarity 0.90)
- SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis (ICLR 2024 · similarity 0.89)
- One-Way Ticket: Time-Independent Unified Encoder for Distilling Text-to-Image Diffusion Models (CVPR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.