The Diffusion Probabilistic Model (DPM) achieves remarkable performance in image generation, while its increasing parameter size and computational overhead hinder its deployment in practical applications. To improve this, the existing literature focuses on obtaining a smaller model with a fixed architecture through model compression. However, in practice, DPMs usually need to be deployed on various devices with different resource constraints, which leads to multiple compression processes, incurring significant overhead for repeated training. To obviate this, we propose a once-for-all (OFA) compression framework for DPMs that yields different subnetworks with various computations in a one-shot training manner. The existing OFA framework typically involves massive subnetworks with different parameter sizes, while such a huge candidate space slows the optimization. Thus, we propose to restrict the candidate subnetworks with a certain set of parameter sizes, where each size corresponds to a specific subnetwork. Specifically, to construct each subnetwork with a given size, we gradually allocate the maintained channels by their importance. Furthermore, we propose a reweighting strategy to balance the optimization process of different subnetworks. Experimental results show that our approach can produce compressed DPMs for various sizes with significantly lower training overhead while achieving satisfactory performance.
Figures & tables
Figure 1. Overview of our OFA-Diffusion Compression framework. On the left, we present the structure of U-ViT architecture. On the right, we show the details of layers in blocks (FFN and Self-Attention) in our compression framework. We omit GroupNorm and time embedding for simplicity. Given a retention rate Pi , only the top- Pi important parameters will be preserved in each block, then these budgets will be unevenly allocated to each layer in a greedy manner. By doing so, for two retention rates Pi<Pj , in each layer, the corresponded module of Pi is a subset of the module corresponded to Pj .
Figure 2. The practical retention rates of parameters P^i over the ideal ones Pi , under the structures of U-Net and U-ViT used in Section 6 .
Figure 3. Comparison between our proposed subnetwork construction approach and other ablations. For simplicity, “Arch” and “Imp” denote the “Architecture” and “Importance”, respectively. All OFA networks are under the same training and evaluation settings. The results are obtained under U-Net, all the FID scores of 28 subnetworks are plotted. We also report the averaged FID across all subnetworks for each method. We can observe that our method consistently achieves a better averaged FID on all datasets.
Datasets
Method
P
0.25×
0.375×
0.5×
0.625×
0.75×
0.875×
1.0×
Avg.
Training Overhead
MACs
6.06G
9.42G
11.44G
13.41G
15.98G
18.47G
21.15G
FID
FID
FID
FID
FID
FID
FID
FID
(K iterations)
CIFAR10 32×32
Separate Compression ( Fang et al., 2023 )
3.60
2.74
2.50
2.28
2.21
2.11
1.97 †
2.49
200N
Naive OFA ( Cai et al., 2020 )
5.31
3.87
3.16
3.00
2.75
2.60
2.66
3.34
200
Ours
4.42
2.97
2.60
2.34
2.30
2.32
2.30
2.75
200
+ Fine-tune 10K
3.45
2.64
2.36
2.28
2.15
2.05
1.95
2.41
200+10N
Table 1. FID scores ( ↓ ) on CIFAR10 32×32 , FFHQ 64×64 , AFHQv2 64×64 and ImageNet 64×64 . N is the number of subnetworks (28 in our experiments) while we show 7 subnetworks here for space limits. The training overhead of our OFA network remains constant as the number of compressed subnetworks increases. “+Fine-tune" denotes the specific subnetworks are further fine-tuned after grabbing weights from our OFA network. † are the performance of pre-trained models. We mark the best results under each Pi in bold.
Figure 4. Averaged latency of model on GPU or CPU evaluated on (a) U-Net CIFAR-10 32×32 with 35-step Heun’s sampler and (b) U-ViT CelebA 64×64 with 50-step DPM-Solver. The GPU and CPU latency are evaluated with a batch size of 64 and 1, respectively.
Datasets
Method
P
0.25×
0.375×
0.5×
0.625×
0.75×
0.875×
1.0×
Avg.
Training Overhead
MACs
3.09G
4.28G
5.71G
7.12G
8.46G
9.64G
11.39G
FID
FID
FID
FID
FID
FID
FID
FID
(K iterations)
CIFAR10 32×32
Separate Compression ( Fang et al., 2023 )
6.83
5.58
5.28
5.39
5.55
5.63
3.73 †
5.43
50N
Naive OFA ( Cai et al., 2020 )
20.96
11.41
8.52
8.56
8.18
7.82
7.38
10.40
50
Ours
12.58
9.46
7.89
7.35
6.96
6.66
6.49
8.20
50
+ Fine-tune 10K
8.99
7.60
7.03
6.63
6.34
6.03
5.74
6.91
50+10N
Table 2. FID scores ( ↓ ) on CIFAR10 32×32 and CelebA 64×64 with a U-ViT backbone, and on MS-COCO2014 512×512 with a Stable Diffusion v1.5 backbone.
Datasets
Method
0.25×
0.375×
0.5×
0.625×
0.75×
0.875×
1.0×
Avg.
CIFAR10 32×32
Uniform
6.38
3.62
3.01
2.55
2.48
2.35
2.34
3.25
Sandwich
4.71
3.30
3.21
2.68
2.70
2.29
2.60
3.07
Ours
4.42
2.97
2.60
2.34
2.30
2.32
2.30
2.75
FFHQ 64×64
Uniform
5.10
3.95
3.08
3.15
2.78
2.82
2.93
3.40
Sandwich
4.30
3.82
3.41
3.04
2.90
2.87
2.63
3.28
Ours
4.27
3.63
3.04
3.03
2.68
2.90
2.82
3.20
Table 3. Ablation of training strategy in OFA-Diffusion Compression framework evaluated on U-Net. The best results for each Pi are bold marked.
Figure 5. Comparison of FID versus MACs with different reweighting ratios on CIFAR10 under U-Net. m→1.0 means wP1=mwPN . We also report the averaged FID over subnetworks. For efficiency, we evaluate FID on 10K generated samples.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6. Overview of our OFA-Diffusion Compression framework. On the left, we present the structure of U-Net architecture. On the right, we show the details of layer in blocks and the compressed modules ( Conv1 , Conv2 and Self-Attention) in our compression framework. We omit GroupNorm and time embedding for simplicity.Given a retention rate Pi , only the top- Pi important parameters will be preserved in each block, then these budgets will be unevenly allocated to each layer in a greedy manner. By doing so, for two retention rates Pi<Pj , in each layer, the corresponded module of Pi is a subset of the module corresponded to Pj .
EDM (UNet)
UViT
SD
Hyperparameters
CIFAR10 32×32
FFHQ 64×64
AFHQv2 64×64
ImageNet 64×64
CIFAR10 32×32
CelebA 64×64
COCO 512×512
Target retention rates ∣P∣
28
28
28
22
19
19
28
Batch size
512
256
256
512
512
256
64
Learning rate
1e-3
2e-4
2e-4
1e-4
2e-4
1e-4
5e-5
Training iterations
200K
200K
200K
400K
50K
100K
50K
EMA half-life (Mimg)
0.5
0.5
0.5
50
3.5
1.8
0.4
Appendix
Table 4. Hyperparameters of our OFA diffusion network on each dataset.
Figure 7. The convergence of subnetworks with different parameter retention rates P on dataset CIFAR10 32×32 . The y-axis is the relative FID scores, to indicate the convergence rate under different retention rates.
Figure 8. Visualization of normed importance scores across layers in the block at 8×8 resolution. At the top of each bar, we also present the FID increases (denoting performance drop) if we remove the corresponding layer, which shows the sensitivity of layers.
Method
0.25×
0.5×
0.75×
1.0×
Ours (Algorithm 1 )
5.80
4.73
4.37
4.20
(a) Update Importance Scores
9.52
5.92
4.44
4.39
(b) Split Time
8.00
5.84
4.87
4.34
Appendix
Table 5. Comparison of varying importance scores of channels as discussed in Section E . For efficiency, we evaluate FID on 10K generated samples.
Figure 9. Samples from subnetworks trained by different approaches with the same random seed on CIFAR10 32×32 (EDM unconditional sampling, U-Net backbone).
Figure 10. Samples from subnetworks trained by different approaches with the same random seed on FFHQ 64×64 (EDM unconditional sampling, U-Net backbone).
Figure 11. Samples from subnetworks trained by different approaches with the same random seed on AFHQv2 64×64 (EDM unconditional sampling).
Figure 12. Samples from subnetworks trained by different approaches with the same random seed on ImageNet 64×64 (EDM conditional sampling, U-Net backbone).
Figure 13. Samples from subnetworks trained by different approaches with the same random seed on cifar10 32×32 (U-ViT unconditional sampling, Transformer backbone).
Figure 14. Samples from subnetworks trained by different approaches with the same random seed on CelebA 64×64 (U-ViT unconditional sampling, Transformer backbone).
Figure 15. Samples from subnetworks trained by different approaches with the same random seed on MS-COCO 512×512 (Stable Diffusion text-to-image generation).
Deploying foundation models is increasingly constrained by memory footprint, latency, and hardware costs. Post-training compression can mitigate these bottlenecks by reducing the precision of model parameters without significantly degrading performance; however, its practical implementation remains challenging as practitioners navigate a fragmented landscape of quantization algorithms, precision budgets, data-driven calibration strategies, and hardware-dependent execution regimes. We present OneComp, an open-source compression framework that transforms this expert workflow into a reproducible, resource-adaptive pipeline. Given a model identifier and available hardware, OneComp automatically inspects the model, plans mixed-precision assignments, and executes progressive quantization stages, ranging from layer-wise compression to block-wise refinement and global refinement. A key architectural choice is treating the first quantized checkpoint as a deployable pivot, ensuring that each subsequent stage improves the same model and that quality increases as more compute is invested. By converting state-of-the-art compression research into an extensible, open-source, hardware-aware pipeline, OneComp bridges the gap between algorithmic innovation and production-grade model deployment.
Yuma Ichikawa, Keiji Kimura, Akihiro Yoshida +11
Fujitsu Limited · RIKEN Center for AIP · Institute of Science Tokyo +2
We propose importance-aware pruning for diffusion models, a training-free framework that prioritizes preserving parameters critical to semantically salient image regions. To do so, we incorporate spatial importance maps -- derived from conditioning signals or model attention -- into the pruning objective. This produces parameter rankings aligned with perceptual relevance rather than uniform reconstruction error. On MS-COCO dataset, our proposed approach consistently retains subject fidelity and structural correctness at high compression ratios where conventional pruning causes visible degradation. These results demonstrate that content-aware objectives are key to perceptually faithful compression of generative models.
Ba-Thinh Lam, Srijan Das, Hieu Le
Department of Computer Science UNC Charlotte NC, USA
Recent advanced diffusion methods typically derive strong generative priors by scaling diffusion transformers. However, scaling fails to generalize when adapted for real-time compression scenarios that demand lightweight models. In this paper, we explore the design of real-time and lightweight diffusion codecs by addressing two pivotal questions. First, does diffusion pre-training benefit lightweight diffusion codecs? Through systematic analysis, we find that generation-oriented pre-training is less effective at small model scales whereas compression-oriented pre-training yields consistently better performance. Second, are transformers essential? We find that while global attention is crucial for standard generation, lightweight convolutions suffice for compression-oriented diffusion when paired with distillation. Guided by these findings, we establish a one-step lightweight convolution diffusion codec that achieves real-time 60~FPS encoding and 42~FPS decoding at 1080p. Further enhanced by distillation and adversarial learning, the proposed codec reduces bitrate by 85% at a comparable FID to MS-ILLM, bridging the gap between generative compression and practical real-time deployment. Codes are released at https://github.com/microsoft/GenCodec/tree/main/CoD_Lite
Zhaoyang Jia, Naifu Xue, Zihan Zheng +7
University of Science and Technology of China · Communication University of China · Microsoft Research Asia