OFA-Diffusion Compression: Compressing Diffusion Model in One-Shot Manner

Haoyang Jiang; Zekun Wang; Mingyang Yi; Xiuyu Li; Lanqing Hu; Junxian Cai; Qingbin Liu; Xi Chen

cs.CVApr 14, 2026

OFA-Diffusion Compression: Compressing Diffusion Model in One-Shot Manner

Authors: Haoyang Jiang, Zekun Wang, Mingyang Yi, Xiuyu Li, Lanqing Hu, Junxian Cai, Qingbin Liu, Xi Chen,

Organizations: Renmin University of China Beijing, China · Alibaba Inc. Hangzhou, China · Independent researcher · Tencent Inc. Shenzhen, China

Abstract

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 $\mathcal{P}_{i}$ , only the top- $\mathcal{P}_{i}$ 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 $\mathcal{P}_{i}<\mathcal{P}_{j}$ , in each layer, the corresponded module of $\mathcal{P}_{i}$ is a subset of the module corresponded to $\mathcal{P}_{j}$ .

Figure 2. The practical retention rates of parameters

\hat{\mathcal{P}}_{i}

over the ideal ones

\mathcal{P}_{i}

, 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	$\mathcal{P}$	$0.25\times$	$0.375\times$	$0.5\times$	$0.625\times$	$0.75\times$	$0.875\times$	$1.0\times$	Avg.	Training Overhead
		MACs	6.06G	9.42G	11.44G	13.41G	15.98G	18.47G	21.15G	Avg.	Training Overhead
			FID	FID	FID	FID	FID	FID	FID	FID	(K iterations)
CIFAR10 $32\times 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 (

\downarrow

) on CIFAR10

32\times 32

, FFHQ

64\times 64

, AFHQv2

64\times 64

and ImageNet

64\times 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.

{\dagger}

are the performance of pre-trained models. We mark the best results under each

\mathcal{P}_{i}

in bold.

Figure 4. Averaged latency of model on GPU or CPU evaluated on (a) U-Net CIFAR-10

32\times 32

with 35-step Heun’s sampler and (b) U-ViT CelebA

64\times 64

with 50-step DPM-Solver. The GPU and CPU latency are evaluated with a batch size of 64 and 1, respectively.

Datasets	Method	$\mathcal{P}$	$0.25\times$	$0.375\times$	$0.5\times$	$0.625\times$	$0.75\times$	$0.875\times$	$1.0\times$	Avg.	Training Overhead
		MACs	3.09G	4.28G	5.71G	7.12G	8.46G	9.64G	11.39G	Avg.	Training Overhead
			FID	FID	FID	FID	FID	FID	FID	FID	(K iterations)
CIFAR10 $32\times 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 (

\downarrow

) on CIFAR10

32\times 32

and CelebA

64\times 64

with a U-ViT backbone, and on MS-COCO2014

512\times 512

with a Stable Diffusion v1.5 backbone.

Datasets	Method	$0.25\times$	$0.375\times$	$0.5\times$	$0.625\times$	$0.75\times$	$0.875\times$	$1.0\times$	Avg.
CIFAR10 $32\times 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\times 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

\mathcal{P}_{i}

are bold marked.

Figure 5. Comparison of FID versus MACs with different reweighting ratios on CIFAR10 under U-Net.

m\to 1.0

means

w_{\mathcal{P}_{1}}=mw_{\mathcal{P}_{N}}

. 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 ( $Conv_{1}$ , $Conv_{2}$ and Self-Attention) in our compression framework. We omit GroupNorm and time embedding for simplicity.Given a retention rate $\mathcal{P}_{i}$ , only the top- $\mathcal{P}_{i}$ 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 $\mathcal{P}_{i}<\mathcal{P}_{j}$ , in each layer, the corresponded module of $\mathcal{P}_{i}$ is a subset of the module corresponded to $\mathcal{P}_{j}$ .

	EDM (UNet)				UViT		SD
Hyperparameters	CIFAR10 $32\times 32$	FFHQ $64\times 64$	AFHQv2 $64\times 64$	ImageNet $64\times 64$	CIFAR10 $32\times 32$	CelebA $64\times 64$	COCO $512\times 512$
Target retention rates $\|\mathcal{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

\mathcal{P}

on dataset CIFAR10

32\times 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\times 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\times$	$0.5\times$	$0.75\times$	$1.0\times$
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.

Panel of Figure 9 — Figure 9. Samples from subnetworks trained by different approaches with the same random seed on CIFAR10 $32\times 32$ (EDM unconditional sampling, U-Net backbone).

Panel of Figure 10 — Figure 10. Samples from subnetworks trained by different approaches with the same random seed on FFHQ $64\times 64$ (EDM unconditional sampling, U-Net backbone).

Panel of Figure 11 — Figure 11. Samples from subnetworks trained by different approaches with the same random seed on AFHQv2 $64\times 64$ (EDM unconditional sampling).

Panel of Figure 12 — Figure 12. Samples from subnetworks trained by different approaches with the same random seed on ImageNet $64\times 64$ (EDM conditional sampling, U-Net backbone).

Panel of Figure 13 — Figure 13. Samples from subnetworks trained by different approaches with the same random seed on cifar10 $32\times 32$ (U-ViT unconditional sampling, Transformer backbone).

Panel of Figure 14 — Figure 14. Samples from subnetworks trained by different approaches with the same random seed on CelebA $64\times 64$ (U-ViT unconditional sampling, Transformer backbone).

Panel of Figure 15 — Figure 15. Samples from subnetworks trained by different approaches with the same random seed on MS-COCO $512\times 512$ (Stable Diffusion text-to-image generation).

Diffusion Model Compression

See all 27 papers

OneComp: One-Line Revolution for Generative AI Model Compression
Mar 30, 2026Yuma Ichikawa, Keiji Kimura, Akihiro Yoshida +11Hardware-Aware Quantization Model Compression
Fujitsu Limited · RIKEN Center for AIP · Institute of Science Tokyo +2
Importance-Aware OBS Pruning for Diffusion Models
Jul 22, 2026Ba-Thinh Lam, Srijan Das, Hieu LeDiffusion Model Pruning Neural Network Compression
Department of Computer Science UNC Charlotte NC, USA
CoD-Lite: Real-Time Diffusion-Based Generative Image Compression
Apr 14, 2026Zhaoyang Jia, Naifu Xue, Zihan Zheng +7Generative Image Compression Learned Image Compression
University of Science and Technology of China · Communication University of China · Microsoft Research Asia

Mar 30, 2026J/K move · Enter open · S save

OneComp: One-Line Revolution for Generative AI Model Compression

Yuma Ichikawa, Keiji Kimura, Akihiro Yoshida, Yudai Fujimoto+10

Fujitsu Limited · RIKEN Center for AIP · Institute of Science Tokyo+2

Hardware-Aware Quantization Model Compression Neural Network Compression Post-Training Quantization Neural Network Quantization Mixed-Precision Quantization

Open paper AskPDF arXiv

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.

Abstract

Diffusion Model Compression

Explore similar work

OneComp: One-Line Revolution for Generative AI Model Compression

Importance-Aware OBS Pruning for Diffusion Models

CoD-Lite: Real-Time Diffusion-Based Generative Image Compression