Importance-Aware OBS Pruning for Diffusion Models

Not Every Pixel Matters Equally

OBS-style pruning such as OBS-Diff minimizes the average reconstruction error, so an error on the subject costs the same as an error on the background. At high sparsity, the subject is what breaks. We weight that error with a spatial importance map.

OBS-Diff · uniform objective

$\|\hat W X - W X\|_2^2$

Motorcycle loses its shape

CFG importance map — Importance map $A$ (CFG)

Ours · importance-aware objective

$\|A \odot (\hat W X - W X)\|_2^2$

Subject preserved (red boxes)

One-shot & training-free Same OBS pipeline & cost Any spatial signal

Structured pruning of PixArt-Σ. Prompt: "A black Honda motorcycle parked in front of a garage."

Abstract

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 the MS-COCO dataset, our approach consistently retains subject fidelity and structural correctness at high compression ratios where conventional pruning causes visible degradation. These results show that content-aware objectives are key to perceptually faithful compression of generative models.

Method: Importance-Aware OBS Pruning

OBS ranks each weight by its saliency $S_q = w_q^2 / 2[H^{-1}]_{qq}$, with the Hessian $H$ built from layer activations. This score is already an approximation, so we correct it with information about what the image should preserve: a spatial importance map re-weights the activations before the Hessian is accumulated.

1

Compute an importance map

Take the magnitude of the CFG response, averaged over channels and min-max normalized:

$$M_t = \big|\epsilon_\theta(x_t,t,y) - \epsilon_\theta(x_t,t,\varnothing)\big|$$

$$A_t=\lambda M_t$$

Canny edges or detector boxes can replace $M_t$.

2

Mask the reconstruction error

Errors in important regions count more:

$$\min_{\hat W}\sum_t \alpha_t\big\|A_t\odot(\hat W X_t - W X_t)\big\|_2^2$$

Linear layers act on channels, so $A_t$ commutes with $W$.

3

Swap in the new Hessian

With filtered activations $X'_t = A_t\odot X_t$:

$$H_{\text{imp}} = 2\sum_t \alpha_t\,\mathbb{E}\big[X'_t X_t'^{\top}\big]$$

The OBS closed-form weight update is unchanged.

Performance

Calibrated on 1,000 GCC3M prompts and evaluated on 1,000 MS-COCO 2017 prompts. The gap grows with sparsity.

SD3-Medium
PixArt-Σ
Structured
Other Signals
Targeted

Unstructured pruning on SD3-Medium

Sparsity	Method	CLIP (↑)	ImageReward (↑)	MUSIQ (↑)
0%	Dense	32.26	0.98	72.68
30%	OBS-Diff	32.25	0.96	72.53
30%	Ours (CFG)	32.24	0.98	72.41
40%	OBS-Diff	32.30	0.92	71.02
40%	Ours (CFG)	32.35	0.95	71.48
45%	OBS-Diff	32.25	0.84	68.76
45%	Ours (CFG)	32.30	0.87	69.86
50%	OBS-Diff	32.16	0.71	65.98
50%	Ours (CFG)	32.20	0.76	66.30

Unstructured pruning on PixArt-Σ

Sparsity	Method	CLIP (↑)	ImageReward (↑)	MUSIQ (↑)
0%	Dense	31.86	0.94	71.11
45%	OBS-Diff	31.63	0.78	70.12
45%	Ours (CFG)	31.73	0.85	70.77
50%	OBS-Diff	31.62	0.76	70.07
50%	Ours (CFG)	31.71	0.78	70.34
55%	OBS-Diff	31.65	0.64	68.25
55%	Ours (CFG)	31.69	0.70	69.25
60%	OBS-Diff	31.44	0.49	65.24
60%	Ours (CFG)	31.54	0.52	67.13

Structured pruning on PixArt-Σ

Sparsity	Method	CLIP (↑)	ImageReward (↑)	MUSIQ (↑)
0%	Dense	31.86	0.94	71.11
20%	OBS-Diff	31.65	0.74	69.57
20%	Ours (CFG)	31.68	0.79	71.37
30%	OBS-Diff	31.36	0.38	66.98
30%	Ours (CFG)	31.30	0.54	69.85
40%	OBS-Diff	30.84	−0.04	65.17
40%	Ours (CFG)	31.08	0.30	68.64

Alternative importance signals on PixArt-Σ (unstructured)

Sparsity	Method	CLIP (↑)	ImageReward (↑)	MUSIQ (↑)
0%	Dense	31.86	0.94	71.11
45%	OBS-Diff	31.63	0.78	70.12
	Ours (CFG)	31.73	0.85	70.77
	Ours (CFG + Det.)	31.73	0.82	70.87
	Ours (Canny)	31.73	0.82	70.81
50%	OBS-Diff	31.62	0.76	70.07
	Ours (CFG)	31.71	0.78	70.34
	Ours (CFG + Det.)	31.71	0.78	70.59
	Ours (Canny)	31.75	0.81	70.34
55%	OBS-Diff	31.65	0.64	68.25
	Ours (CFG)	31.69	0.70	69.25
	Ours (CFG + Det.)	31.79	0.75	69.62
	Ours (Canny)	31.83	0.76	69.12
60%	OBS-Diff	31.44	0.49	65.24
	Ours (CFG)	31.54	0.52	67.13
	Ours (CFG + Det.)	31.60	0.50	67.27
	Ours (Canny)	31.52	0.53	65.91

Category-targeted pruning on PixArt-Σ (60% unstructured)

General: calibrated on the category-diverse GCC3M set. Target: calibrated only on prompts that contain the category. All models are evaluated on the same category-specific test prompts.

Category	Method	CLIP (↑)	ImageReward (↑)	MUSIQ (↑)
Woman	OBS-Diff (general)	32.36	0.72	65.08
	OBS-Diff (target)	32.25	0.75	65.75
	Ours (general)	32.31	0.76	66.74
	Ours (target)	32.38	0.85	67.40
Cat	OBS-Diff (general)	32.58	0.28	64.35
	OBS-Diff (target)	32.45	0.29	69.81
	Ours (general)	32.66	0.41	69.52
	Ours (target)	32.28	0.34	71.64
Airplane	OBS-Diff (general)	29.35	0.45	67.81
	OBS-Diff (target)	28.62	0.43	69.90
	Ours (general)	29.22	0.62	68.19
	Ours (target)	28.76	0.51	70.81

Best value per sparsity level (or category) in green.

Application · Flexible importance guidance

Bias Pruning Toward What You Care About

The pruning formulation never changes. The only input you choose is the importance map $A_t$, which says where reconstruction errors should cost more. Changing the signal behind $A_t$ changes what the pruned model preserves: prompt content, whole objects, or fine structure.

Takeaway: the importance signal is a plug-in. Canny makes the method usable when CFG maps are unavailable or inconvenient to extract, and adding a detector steers pruning toward object-centric preservation. Every variant beats OBS-Diff on PixArt-Σ at all tested sparsities (Other Signals tab above).

Gallery

Outputs of our pruned models at their highest tested sparsity. Click any image to compare against OBS-Diff at every sparsity level. The last row shows category-targeted pruning, where calibration uses only prompts that contain a chosen category.

BibTeX

@inproceedings{lam2026importance,
  title     = {Importance-Aware OBS Pruning for Diffusion Models},
  author    = {Lam, Ba-Thinh and Das, Srijan and Le, Hieu},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026},
  eprint    = {2607.20048},
  archivePrefix = {arXiv}
}