PaperlayerView exact revision on arXiv

Source-linked AI summary

BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion

Bo-Kyeong Kim, Hyoung-Kyu Song, Thibault Castells, Shinkook Choi

arXiv:2305.15798v4cs.LG

TL;DR

Stable Diffusion text-to-image generation is computationally demanding, and prior efficiency work largely retained the original architecture. This paper removes U-Net blocks and retrains compact models with knowledge distillation; the resulting models achieve competitive zero-shot results against large-scale baselines while reducing model size and latency.

  • Problem

    Stable Diffusion text-to-image generation has high computational demands, while architectural reduction of large diffusion models remains less explored.

    Source: §1 · p. 2

  • Method

    The approach removes residual and attention blocks from the Stable Diffusion U-Net and retrains the compact model using feature-level knowledge distillation from the original.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Up to 51% reduction in model size and 43% improvement in latency on CPU and GPU accompany zero-shot results competitive with multi-billion-parameter models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The compact backbones are applicable to personalized generation and image-to-image translation, with inference on Jetson AGX Orin and iPhone 14 taking less than 4 seconds.

    Source: §1 · p. 3, §5.7 · p. 14

Abstract

from arXiv · show

Text-to-image (T2I) generation with Stable Diffusion models (SDMs) involves high computing demands due to billion-scale parameters. To enhance efficiency, recent studies have reduced sampling steps and applied network quantization while retaining the original architectures. The lack of architectural reduction attempts may stem from worries over expensive retraining for such massive models. In this work, we uncover the surprising potential of block pruning and feature distillation for low-cost general-purpose T2I. By removing several residual and attention blocks from the U-Net of SDMs, we achieve 30%~50% reduction in model size, MACs, and latency. We show that distillation retraining is effective even under limited resources: using only 13 A100 days and a tiny dataset, our compact models can imitate the original SDMs (v1.4 and v2.1-base with over 6,000 A100 days). Benefiting from the transferred knowledge, our BK-SDMs deliver competitive results on zero-shot MS-COCO against larger multi-billion parameter models. We further demonstrate the applicability of our lightweight backbones in personalized generation and image-to-image translation. Deployment of our models on edge devices attains 4-second inference. Code and models can be found at: https://github.com/Nota-NetsPresso/BK-SDM

1 Introduction

Stable Diffusion models offer capable text-to-image generation, but their computational demands can hinder use, while architectural reduction has received less attention than step reduction and quantization. This work removes U-Net blocks and uses feature distillation to build compact models, reporting lower resource demands and competitive generation with limited retraining.

Contributions

  • The block-removed models achieve compelling zero-shot results on par with multi-billion-parameter models.The experiments compress SDM-v1.4 and SDM-v2.1-base.Source: §1 · p. 2
  • Up to 51% reduction in model size and 43% improvement in latency on CPU and GPU are achieved by removing U-Net blocks.The paper identifies structural pruning of billion-scale foundation models as an underexplored direction, partly because retraining them is expensive.Source: §1 · p. 2
  • Feature distillation trains compact models with 13 A100 days and 0.22M LAION pairs, compared with over 6,000 A100 days and 2,000M pairs to train SDMs from scratch.The paper presents this as evidence that compact general-purpose diffusion models can be built with substantially fewer resources.Source: §1 · p. 3
  • The lightweight backbones are applicable to personalized generation and image-to-image translation, with inference on Jetson AGX Orin and iPhone 14 taking less than 4 seconds.The paper also releases its approach, model weights, and source code.Source: §1 · p. 3

2 Related Work

Related work addresses Stable Diffusion’s efficiency through fewer sampling steps, faster solvers, quantization, and implementation optimizations. The paper positions architectural compression and distillation-based compression of diffusion models as less-explored complements.

Large T2I diffusion models

  • Stable Diffusion models perform diffusion in a low-dimensional latent space and are used as the paper’s baseline because of open access and growing use in downstream tasks.Other large text-to-image systems use different pipelines, including super-resolution modules or a prior network followed by a diffusion decoder [59].Source: §2 · p. 3

Efficient diffusion models

  • Fewer-step distillation and fast high-order solvers [42] speed sampling, while network compression reduces per-step computation and can be combined with fewer steps.Quantization and implementation optimizations can also be combined with compact models.Source: §2 · p. 3

Distillation-based compression

  • Classical output- and feature-level knowledge distillation is established for smaller models, but its use for structurally compressed diffusion models had not been explored.Concurrent efficient U-Net approaches require much larger training resources, while Diff-Pruning uses width pruning and a comparable Taylor criterion that yields insufficient MAC reduction in the reported comparison.Source: §2 · pp. 3–4

3 Compression Method

The method reduces computation by removing blocks from the Stable Diffusion U-Net and retraining the resulting student to mimic the original. It offers several block-removal configurations and combines the denoising task loss with output- and feature-level distillation.

3 Compression Method

  • The method compresses the compute-heavy U-Net, reducing per-step computation to produce Block-removed Knowledge-distilled SDMs (BK-SDMs).The U-Net processes latent representations conditioned on text and timestep embeddings, predicting a noise residual at each denoising step.Source: §3 · p. 4

3.1 Compact U-Net Architecture

  • In down stages, the design retains the first residual-attention pairs and removes second pairs; in up stages, it retains the third pairs to use down-stage features and skip connections.Retaining the first down-stage pairs preserves the dimensions needed to initialize from corresponding pretrained weights.Source: (1) Fewer Blocks in the Down and Up… · pp. 4–5
  • Removing the entire mid-stage does not noticeably degrade generation quality without retraining and reduces parameters by 11%.Combining mid-stage removal with fewer blocks further reduces compute at the cost of a slight performance decline, so the paper offers it as an efficiency-versus-quality option.Source: (2) Removal of the Entire Mid-Stage · pp. 5–6
  • Pruning the innermost down and up stages produces BK-SDM-Tiny, the lightest model, while the architecture analysis identifies these stages and second down-stage pairs as relatively minor.The analysis measures generation-score changes when removing blocks; blocks with incompatible input and output channel dimensions are replaced with interpolation modules.Source: (2) Removal of the Entire Mid-Stage · p. 6, Alignment with Pruning Sensitivity… · p. 6, (3) Further Removal of the Innermost… · p. 6
  • The CLIP Score pruning-sensitivity criterion excessively prunes attention blocks and yields inferior performance to the proposed architectures.The paper’s architectures are designed using human knowledge and prior studies, with sensitivity results supporting some removal choices.Source: Alignment with Pruning Sensitivity… · p. 6

3.2 Distillation-based Retraining

  • For general-purpose text-to-image generation, the block-removed U-Net student learns to mimic the original U-Net teacher using denoising task loss and output-level knowledge distillation.The task loss trains the student to estimate sampled noise from a noisy latent code, paired text embedding, and timestep; output-level distillation matches student and teacher noise predictions.Source: §3.2 · pp. 6–7
  • Feature-level distillation matches teacher and student features at selected layers, with stage-end matching avoiding extra channel-matching regressors.The stage-end feature maps already have matching dimensions.Source: §3.2 · p. 7
  • The final training objective sums task, output-distillation, and feature-distillation losses, with both distillation weights set to 1.The authors report that this setting works without loss-weight tuning.Source: §3.2 · p. 7

4 Experimental Setup

The experiments evaluate distillation training, zero-shot text-to-image generation, and downstream applications using specified datasets, prompts, and metrics. Main retraining uses one A100 GPU, and inference generally uses 25 denoising steps.

Distillation Retraining

  • Distillation retraining primarily uses 0.22M image-text pairs from LAION-Aesthetics V2 6.5+, compared with over 2,000M pairs used to train the original SDMs.For dataset-size experiments, smaller datasets are sampled from L-Aes 6.5+ and larger datasets from L-Aes 6.25+.Source: §4 · p. 7

Zero-shot T2I Evaluation

  • Zero-shot evaluation uses 30K MS-COCO validation prompts, compares generated 256×256 images with the full validation set, and measures FID, IS, and CLIP score.Images are generated at 512×512 and downsampled to 256×256; CLIP-ViT-g/14 measures text-image correspondence.Source: §4 · p. 7

Downstream Tasks

  • Personalized-generation evaluation uses 30 DreamBooth subjects with 25 prompts and 4–6 images per subject, while image-to-image inputs come from Meng et al..Per-subject finetuning is evaluated with DINO and CLIP-I/CLIP-T scores.Source: §4 · pp. 7–8

Implementation Details

  • Main retraining uses a single NVIDIA A100 80G GPU, per-subject finetuning uses an RTX 3090, and inference generally uses 25 denoising steps with guidance scale 7.5.The default latent resolution is 64×64, yielding 512×512 images.Source: §4 · p. 8

5 Results

Across zero-shot MS-COCO evaluation, the compressed models perform competitively with large models, while block removal reduces computation and distillation supports retraining with limited resources. The paper also examines pruning criteria, training-resource trade-offs, and applications beyond general-purpose generation.

5.1 Comparison with Existing Works

  • Despite using only 0.22M samples and fewer than 1B parameters, the compressed models achieve zero-shot MS-COCO performance on par with existing large models.Visual comparisons also report more photorealistic outputs than AR- and GAN-based baselines.Source: §5.1 · p. 9, Figure 7 · p. 9

5.2 Computational Gain

  • Inference time is reduced by more than 30%, and BK-SDM-Tiny has 50% fewer parameters than the original SDM.Table 2 reports whole-model computation and latency for generating a single 512×512 image with 25 denoising steps.Source: Table 2 · p. 9, §5.2 · p. 10

5.3 Benefit of Distillation Retraining

  • Distillation improves generation quality across model types; output-level and feature-level distillation each add gains, and teacher-weight initialization is highly beneficial.The study evaluates total KD and its individual components on MS-COCO.Source: §5.3 · p. 10, Table 3 · p. 10, Table 4 · p. 10
  • With distillation, training is more stable and faster, and small- and tiny-size models outperform the bigger base-size model trained without KD.FID, IS, and CLIP score trends differ over training, so the authors caution against judging models using FID alone.Source: §5.3 · p. 11, Figure 9 · p. 11

5.4 Comparison with Different Pruning Criteria

  • The pruning approach achieves superior or comparable results to benchmark methods without an extra pruning-criteria calculation phase, while retaining a favorable performance–speed balance.Magnitude pruning removes critical outer blocks; Taylor pruning removes inner blocks, reduces MACs and speed less effectively, and requires intensive gradient calculations.Source: §5.4 · pp. 11–12, Table 5 · p. 12

5.6 Impact of Training Resources on Performance

  • Larger batches, more data, and longer training enhance performance, but require greater training and storage resources.The authors therefore primarily report models trained with fewer resources despite stronger-performing alternatives.Source: Table 7 · p. 12, Figure 11 · p. 12, §5.6 · p. 13

5.7 Application

  • BK-SDMs retain 95%–99% of original-SDM scores in personalized generation while reducing fine-tuning costs, and support text-guided stylization and edge-device inference.The reported edge-device setup uses 20 denoising steps on Jetson AGX Orin and 10 on iPhone 14; generated 512×512 images take within 4 seconds.Source: Table 8 · p. 13, Table 9 · p. 13, §5.7 · p. 14

6 Conclusion and Discussion

The paper concludes that block-removed Stable Diffusion models can perform competitively in zero-shot generation, with distillation enabling retraining under constrained resources. It also identifies inaccurate full-body human appearance as a shortcoming.

6 Conclusion and Discussion

  • Block-removed lightweight models achieve competitive zero-shot generation against large-scale baselines, and distillation supports retraining under very constrained resources.The authors describe architectural compression as orthogonal to methods that reduce sampling steps and say the approaches can be combined.Source: §6 · p. 14
  • The compact models inherit the source model’s high-fidelity image-generation capability but can generate full-body human appearances inaccurately.Source: §6 · p. 14

A U-Net Architecture and Distillation Retraining

The appendix describes the shared U-Net block configuration and the student’s feature- and output-distillation process. It also reports that removing the entire mid-stage reduces parameters without noticeably degrading generation quality.

A U-Net Architecture and Distillation Retraining

  • The approach applies directly to SDM versions v1.1–v1.5 and v2.0–v2.1, which share the same U-Net block configuration.The appendix also describes differences in attention-head counts and text-embedding dimensions between SDM-v1 and v2.Source: A U-Net Architecture and Distillation… · pp. 20–21
  • The compact student U-Net is formed by removing residual and attention blocks from the teacher and trained through feature and output distillation.For SDM-v1 and v2-base, the default latent resolution is 64×64, yielding 512×512 generated images.Source: Figure 16 · p. 20, Figure 15 · p. 20

B Impact of Mid-stage Removal

  • Removing the entire mid-stage from the original U-Net reduces the number of parameters while not noticeably degrading generation quality for many text prompts.The reported MS-COCO evaluation is at 256×256 with 30K samples; retraining is not performed.Source: B Impact of Mid-stage Removal · p. 22, Table 10 · p. 22

C Block-level Pruning Sensitivity Analysis

The appendix describes block-removal sensitivity analysis and presents comparisons with existing studies, personalized generation, and text-guided image-to-image translation.

C Block-level Pruning Sensitivity Analysis

  • Block-removal sensitivity is assessed by removing blocks where dimensions permit and using channel interpolation where input and output dimensions differ.The described choices include removing innermost stages and the second residual-attention pairs in down stages.Source: C Block-level Pruning Sensitivity… · p. 23, Figure 19 · p. 23

D Comparison with Existing Studies

  • The appendix presents zero-shot general-purpose text-to-image results alongside prior studies, with no CLIP-based reranking applied to SDM or the authors’ models.Source: Figure 20 · p. 24

E Personalized Generation

  • The personalized-generation examples mark each subject with an identifier and class noun; the caption states that compact models preserve subject appearance across backgrounds.Source: Figure 21 · p. 25

F Text-guided Image-to-Image Translation

  • The text-guided image-to-image examples present stylization of input images by the small models.Source: Figure 22 · p. 26

G Deployment on Edge Devices

BK-SDM models were evaluated on Jetson AGX Orin and iPhone 14, where they generated 512×512 images faster than the compared Stable Diffusion models.

G Deployment on Edge Devices

H Impact of Training Data Volume

The additional experiments compare BK-SDM outputs across training-data quantities and report improved visual outcomes as data increases.

H Impact of Training Data Volume

  • Fig. 28 presents Base, Small, and Tiny models trained on 0.22M pairs alongside corresponding models trained on 2.3M pairs.The larger-data models are Base-2M, Small-2M, and Tiny-2M.Source: H Impact of Training Data Volume · p. 30, Figure 28 · p. 31
  • Increasing training data improves visual outcomes, including image-text matching and differentiation between objects.This trend is illustrated for BK-SDM-Small.Source: Figure 27 · p. 30
  • The data-quantity experiments include visual comparisons of BK-SDM-Small trained with different amounts of data.The reported comparison discusses image-text matching and object differentiation as visual outcomes.Source: Figure 27 · p. 30, H Impact of Training Data Volume · p. 30

I Additional Experiments

Additional experiments examine structural variants, learning rates, skip-connection pruning, and implementation settings for retraining and inference.

I Additional Experiments

  • Removing the mid-stage yields a model between the original and base sizes, while CLIP-Score-based structural variants have suboptimal performance.The mid-stage option further reduces compute at a slight performance cost; the tables report the corresponding structural comparisons.Source: I Additional Experiments · p. 32, Table 11 · p. 32, Table 12 · p. 32
  • Learning rates of 5e-5 and 2.5e-5 yield good results, while extremely high or low values are detrimental.The comparison concerns BK-SDM-v2-Small.Source: I Additional Experiments · p. 32, Table 13 · p. 32
  • Inner skip concatenations are robust to zeroing and prunable, with second concatenations more removable than others; first residual blocks are often unprunable.The analysis corrupts skip features by assigning zero values and reports the comparison in Fig. 29.Source: I Additional Experiments · p. 33, Figure 29 · p. 33

J Implementation

  • Distillation retraining uses one A100 for 50K iterations and about 300 hours; personalized finetuning uses one RTX 3090 for 800 iterations.General-purpose retraining uses batch size 256 and 53GB GPU memory; personalized finetuning uses batch size 1, with BK-SDMs requiring 13–19GB versus 23GB for the original SDM.Source: J Implementation · p. 34
Loading 2305.15798v4…