LAPTOP-Diff: Layer Pruning and Normalized Distillation for Compressing Diffusion Models

TL;DR

The paper tackles the memory and latency barriers of diffusion models for on-device use by introducing LAPTOP-Diff, which couples automated one-shot layer pruning with a normalized feature distillation during retraining. The core ideas are a combinatorial optimization-based pruning criterion with a proven additivity property and a reweighting scheme that mitigates imbalanced feature losses during distillation. Empirical results on SDXL and SDM-v1.5 show that around 50% pruning yields minimal performance drop (4.0% PickScore decline) while offering substantial speedups, outperforming handcrafted layer-removal approaches. The approach provides a scalable path to efficient diffusion models, enabling practical on-device generation with preserved visual quality and text-image alignment.

Abstract

In the era of AIGC, the demand for low-budget or even on-device applications of diffusion models emerged. In terms of compressing the Stable Diffusion models (SDMs), several approaches have been proposed, and most of them leveraged the handcrafted layer removal methods to obtain smaller U-Nets, along with knowledge distillation to recover the network performance. However, such a handcrafting manner of layer removal is inefficient and lacks scalability and generalization, and the feature distillation employed in the retraining phase faces an imbalance issue that a few numerically significant feature loss terms dominate over others throughout the retraining process. To this end, we proposed the layer pruning and normalized distillation for compressing diffusion models (LAPTOP-Diff). We, 1) introduced the layer pruning method to compress SDM's U-Net automatically and proposed an effective one-shot pruning criterion whose one-shot performance is guaranteed by its good additivity property, surpassing other layer pruning and handcrafted layer removal methods, 2) proposed the normalized feature distillation for retraining, alleviated the imbalance issue. Using the proposed LAPTOP-Diff, we compressed the U-Nets of SDXL and SDM-v1.5 for the most advanced performance, achieving a minimal 4.0% decline in PickScore at a pruning ratio of 50% while the comparative methods' minimal PickScore decline is 8.2%.

Figures (7)

• Figure 1: Architecture of SDM's U-Net and knowledge distillation framework, using SDXL as an example. Other SDMs' U-Net have slightly different architectures, e.g., SDM-v1.5 has 4 Dn and Up stages instead of 3 and has fewer transformer layers.
• Figure 2: L2-Norms of feature maps at the end of each stage of the SDXL teacher and the feature loss terms at the 15K-th iteration of retraining.
• Figure 3: Unnormalized feature loss terms, i.e., $\left|\left|f^i_T-f^i_S\right|\right|_2^2$, of each stage using vanilla feature distillation or our normalized feature distillation, at the 15K-th iteration of retraining.
• Figure 4: Visual comparison with SSD and Vega in DDIM 25 steps.
• Figure 5: Validations of the additivity properties of different pruning criteria on SDXL and SDM-v1.5. Each point represents an observed value pair of the approximated criterion and the real criterion.