Importance-Aware OBS Pruning for Diffusion Models

Conference: NeurIPS 2026 (acceptance information supplied in the task metadata)
arXiv: 2607.20048v3
Area: Model Compression
Keywords: diffusion models, parameter pruning, spatial importance, Optimal Brain Surgeon, training-free compression

TL;DR

The paper injects prompt-related spatial importance into the layer-wise OBS reconstruction objective to improve subject fidelity in highly sparse diffusion models without fine-tuning, but its global mask-coefficient ablation has an unexplained inconsistency with the stated exact equations.

Background & Motivation

Text-to-image diffusion models have large parameter counts and repeatedly execute their denoising networks. OBS-Diff already extends second-order parameter pruning with Optimal Brain Surgeon (OBS) to diffusion models: it collects layer activations across denoising timesteps, estimates the reconstruction cost of removing a weight, and analytically adjusts other weights to compensate. This approach requires no retraining, but average reconstruction error need not capture the aspects of generation that matter most to human observers.

Feature errors of the same magnitude can have different consequences when they affect background textures rather than faces or object boundaries. Standard layer-wise reconstruction does not explicitly distinguish these locations. At high sparsity, an image can remain broadly related to its prompt while suffering subject deformation or missing details. The authors also note that the original diffusion task loss need not be stationary on an individual calibration batch or timestep. However, this is separate from whether the teacher-based layer reconstruction objective is an exact quadratic, and the two issues should not be conflated.

Rather than redesigning the pruning solver, the paper changes the distribution of errors that the solver must protect: it identifies locations where the prompt has a stronger influence, then makes those locations contribute more to curvature statistics. Core Idea: filter calibration activations with spatial importance, replace uniform layer-wise reconstruction with content-aware reconstruction, and retain OBS ranking and analytical compensation.

Method

Overall Architecture

The inputs are a pretrained diffusion model, calibration prompts, and a target sparsity level; the output is a parameter-pruned model. Calibration executes conditional and unconditional denoising branches to extract a spatial CFG signal, applies the corresponding location weights to layer activations, and constructs a Hessian from the weighted activations before OBS ranking, removal, and compensation.

This is an offline calibration and parameter-modification procedure, not the training of a new network. Dense-layer outputs provide reconstruction references rather than manually labeled supervision. Once pruning is complete, ordinary generation uses the pruned model and does not need to recompute pruning importance maps for every generation.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    P["Dense model and calibration prompts"] --> A["CFG Spatial Importance"]
    A --> B["Importance-Weighted Hessian"]
    P -.->|Dense-layer reconstruction reference| B
    B --> C["OBS Ranking and Compensation"]
    C --> D["Parameter-pruned model"]
    Q["Inference prompt and noise"] --> E["Ordinary denoising generation"]
    D --> E

Key Designs

1. CFG Spatial Importance: use conditional responses to identify locations worth protecting

CFG stands for classifier-free guidance. Importance here is neither the conditional prediction itself nor text cross-attention: it is the magnitude of the difference between conditional and unconditional predictions at the same noise state. The authors first take the absolute channel-wise differences and then average across channels. Responses with opposite signs therefore do not cancel, and the result is a two-dimensional spatial map.

\[ \tilde{M}_{t}(i,j)=\frac{1}{C}\sum_{c=1}^{C}\left|\epsilon_{\theta}(x_{t},t,y)_{c,i,j}-\epsilon_{\theta}(x_{t},t,\varnothing)_{c,i,j}\right|. \]

The channel index and text condition are represented by different symbols. Spatial min–max normalization is then performed independently for every sample and denoising timestep. It is not batch-wise normalization and does not retain response magnitudes across timesteps on a common scale.

\[ M_{t}=\frac{\tilde{M}_{t}-\min(\tilde{M}_{t})}{\max(\tilde{M}_{t})-\min(\tilde{M}_{t})+\epsilon},\qquad A_{t}=\lambda M_{t}. \]

The small constant appears only in the source's normalization denominator. The paper explicitly uses \(A_{t}=\lambda M_{t}\) without an additive unit background term. Low-response locations can approach zero and contribute almost nothing to the objective. CFG measures the prompt's influence on a prediction; it is not validated human-saliency ground truth and does not guarantee that all important small objects are highlighted.

The same interface also accepts Canny edge maps, or CFG responses restricted to object regions detected by YOLOv8. The former emphasizes boundaries and textures, while the latter emphasizes detected objects. The detector provides spatial selection rather than additional training supervision. These signals change the preservation objective, not the OBS solver itself.

2. Importance-Weighted Hessian: change the error metric rather than multiply the final scores by a mask

The authors broadcast the importance map across output channels, multiply reconstruction residuals location-wise, and then sum squared errors. The timestep weight \(\alpha_{t}\) follows OBS-Diff's logarithmically decreasing setting to emphasize earlier denoising steps. It and the spatial map weight different dimensions of the objective.

\[ \mathcal{J}(\hat{W}_{l})=\sum_{t=1}^{T}\alpha_{t}\left\|A_{t}\odot\left(\hat{W}_{l}X_{l,t}-W_{l}X_{l,t}\right)\right\|_{F}^{2},\qquad \operatorname{Sparsity}(\hat{W}_{l})=S. \]

For a linear layer operating independently on spatial tokens along the channel dimension, the location weights can be moved onto its input activations. The filtered activations then enter the existing second-order statistics. The modification thus occurs during curvature estimation rather than directly assigning each weight a retention probability from an image map.

\[ X'_{l,t}=A_{t}\odot X_{l,t},\qquad H_{l,\mathrm{imp}}=2\sum_{t=1}^{T}\alpha_{t}\mathbb{E}\left[X'_{l,t}X_{l,t}^{\prime\top}\right]. \]

Squaring is an important detail: the residual is multiplied by the spatial weight before the squared norm is taken. Consequently, each location's contribution to the expanded Hessian is multiplied by \(A_{t}(i,j)^{2}\), not \(A_{t}(i,j)\). Using an input-token vector for a single output row gives the following equivalent expansion. This is the note's algebraic expansion of the source objective, not an additional equation proposed by the paper.

\[ H_{l,\mathrm{imp}}=2\sum_{t=1}^{T}\alpha_{t}\mathbb{E}\left[\sum_{i,j}A_{t}(i,j)^{2}x_{l,t,i,j}x_{l,t,i,j}^{\top}\right]. \]

The spatial map thereby changes the relative composition of the input covariance. Channel directions activated at important locations receive greater weight in the objective, but individual retention decisions still depend on weight magnitudes and correlations in the inverse Hessian. This cannot be reduced to “large foreground activations must be retained.” Nor can the commutation argument between spatial weighting and channel-wise linear mapping be extended without qualification to arbitrary spatial operators.

3. OBS Ranking and Compensation: retain the analytical solver and check what global scaling actually changes

For one output row, OBS constrains a selected weight to become zero while allowing the others to change jointly to minimize quadratic reconstruction error. The minimum objective increase is its saliency score. Lower-saliency parameters are removed first, while the analytical update exploits channel correlations to distribute the removed weight's reconstruction error across other weights.

\[ S_{l,q}=\frac{w_{q}^{2}}{2[H_{l,\mathrm{imp}}^{-1}]_{qq}},\qquad \delta w=-\frac{w_{q}}{[H_{l,\mathrm{imp}}^{-1}]_{qq}}H_{l,\mathrm{imp}}^{-1}\mathbf{e}_{q}. \]

These expressions correspond to the ranking formula in the main text and the compensation derivation in Appendix C, assuming an invertible Hessian. Pruning follows OBS-Diff's procedure and Cholesky-based inverse-Hessian updates. The structured variant uses the baseline's aggregation strategy; the paper does not introduce a separate structured-pruning solver.

Appendix C explicitly states that, with calibration activations fixed and the dense layer serving as teacher, both the residual and first-order gradient vanish at \(\hat{W}_{l}=W_{l}\). This layer-wise surrogate is an exact quadratic in the weights being pruned. Nonstationarity of the original diffusion task should not be interpreted as an inevitably neglected first-order term in this surrogate. The actual approximations concern the relationship between surrogate reconstruction and final generation quality, and the pruning procedure's approximation to joint optimization.

Another reproducibility check concerns scaling: if the mask shape is fixed and all layers and timesteps share one positive global coefficient, changing that coefficient does not alter relative spatial weights. The following scaling relations follow directly from the paper's equations. Here, \(H_{0}\) denotes the Hessian for the same importance map with a unit coefficient, not the unweighted OBS Hessian.

\[ H_{\lambda}=\lambda^{2}H_{0},\qquad H_{\lambda}^{-1}=\lambda^{-2}H_{0}^{-1},\qquad S_{\lambda,q}=\lambda^{2}S_{0,q},\qquad \delta w_{\lambda}=\delta w_{0}\quad(\lambda>0). \]

Under the exact equations, all scores scale uniformly, rankings remain unchanged, and the scale cancels in compensation updates. Nevertheless, the main text reports different quality for different positive coefficients and interprets this as stronger or weaker semantic guidance. The stated equations alone do not support that explanation. The full source gives no damping or stabilization definition that explains the difference, so an implementation mechanism should not be invented for the authors. Moreover, \(\lambda=0\) makes the weighted Hessian zero and renders the usual inverse formula unusable; it does not automatically recover unweighted OBS.

Loss & Training

The method performs no fine-tuning or gradient-based training. Calibration data and dense-layer outputs construct the local reconstruction objective; subsequent weight changes are analytical OBS compensation rather than additional training epochs.

Calibration uses 1,000 GCC3M samples, a batch size of 2, and 10 denoising steps; testing uses 25 steps. Calibration CFG scales are 7.0 for SD3-Medium and 4.5 for PixArt-Σ, while both use 7.0 at test time. The default mask coefficient is 1.0.

The full source does not completely specify the exact timestep-weight implementation, engineering details of spatial mapping across layers, or inverse-matrix numerical handling. Reproduction should verify these elements rather than assume that retaining the existing pipeline constitutes a complete implementation disclosure.

Key Experimental Results

Main Results

The following selection contains key unstructured-pruning results from Tables 1 and 2; higher is better for every metric. CLIP Score proxies global image-text alignment, ImageReward proxies learned human preferences, and MUSIQ measures image quality without directly relying on text.

Model	Sparsity	Method	CLIP Score	ImageReward	MUSIQ
SD3-Medium	Dense	Original model	32.26	0.98	72.68
SD3-Medium	30%	OBS-Diff	32.25	0.96	72.53
SD3-Medium	30%	Ours CFG	32.24	0.98	72.41
SD3-Medium	50%	OBS-Diff	32.16	0.71	65.98
SD3-Medium	50%	Ours CFG	32.20	0.76	66.30
PixArt-Σ	Dense	Original model	31.86	0.94	71.11
PixArt-Σ	60%	OBS-Diff	31.44	0.49	65.24
PixArt-Σ	60%	Ours CFG	31.54	0.52	67.13

The source contains a dataset-scope conflict: the experimental settings specify testing on 1,000 prompts from MS-COCO 2017 validation, but Table 1 labels its results as GCC3M, and the captions of Tables 2 and 5 also mention GCC3M samples. The original numerical values are retained here without reassigning all tables to either test dataset.

Ablation Study

The following results come from Table 6 for PixArt-Σ at 60% sparsity. They are the paper's reported observations, not evidence that the global-scale consistency issue above has been resolved.

Method	Mask coefficient	CLIP Score	ImageReward
OBS-Diff	Not applicable	31.44	0.49
Ours	0.05	31.54	0.52
Ours	0.1	31.60	0.54
Ours	0.3	31.58	0.55
Ours	0.5	31.57	0.53
Ours	1.0	31.54	0.52
Ours	2.0	31.56	0.51
Ours	3.0	31.49	0.50

The default is 1.0, whereas the best tabulated CLIP value occurs at 0.1 and the best ImageReward at 0.3. The default therefore should not be described as the ablation-optimal setting. More importantly, these results do not establish an explainable sensitivity of the exact objective's ranking to a positive global coefficient.

Key Findings

• Improvement is not universal across metrics and settings: at 30% sparsity on SD3-Medium, Ours has slightly lower CLIP and MUSIQ than OBS-Diff and improves only ImageReward. Advantages are clearer at high sparsity.
• Signal choice involves trade-offs: in Table 3's PixArt-Σ setting at 60%, CFG+Detector achieves MUSIQ of 67.27 versus CFG's 67.13, but ImageReward is 0.50 versus CFG's 0.52. Canny achieves ImageReward of 0.53 without the highest MUSIQ.
• Structured pruning also has supporting evidence: in Table 5's PixArt-Σ setting at 40%, ImageReward increases from -0.04 with OBS-Diff to 0.30 with Ours, and MUSIQ rises from 65.17 to 68.64. However, no corresponding real inference-latency table is provided.
• Category-targeted calibration does not improve every metric: in Table 4's Cat setting, the targeted variant achieves MUSIQ of 71.64 versus the general variant's 69.52, while CLIP decreases from 32.66 to 32.28 and ImageReward from 0.41 to 0.34.

Appendix A uses a five-level human image-quality rating. The table below corresponds to Table 7's PixArt-Σ evaluation at 60% sparsity on 20 images. It is a selected small-scale study, not an average effect established across all models, prompts, or random seeds.

Method	CLIP Score	ImageReward	MUSIQ	Human score
OBS-Diff	31.90	0.72	69.83	2.59
Ours	31.06	0.78	72.93	3.97

Human ratings strongly favor Ours even though its CLIP score is lower. Table 8 also reports per-sample preference agreement with human judgments of 25.00%, 50.00%, and 70.00% for CLIP, ImageReward, and MUSIQ, respectively, on this evaluation set. This illustrates limitations of the automatic metrics rather than a generally validated metric-calibration result.

Highlights & Insights

• Incorporating perceptual priorities into the spatial reconstruction metric is more structured than appending a heuristic to final rankings. It preserves analytical OBS compensation while changing the activation distribution that the weights must jointly protect.
• Content guidance is not equivalent to a higher global alignment score. The opposing human-rating and CLIP results suggest that compression evaluation should inspect object structure rather than rely exclusively on global embedding similarity.
• Calibration distribution and spatial maps provide distinct control interfaces: the former determines which concepts are frequently observed, and the latter determines which locations within those samples matter most. Both can support application customization, but content outside the preservation priority also needs evaluation.

Limitations & Future Work

• The authors acknowledge noisy or incomplete CFG maps for abstract prompts, small objects, and cluttered scenes. Detector-based guidance also inherits detection failures and category biases.
• The global-coefficient ablation is not reconciled with the exact equations' scale invariance. The small normalization-denominator constant does not explain why subsequent global scaling would change exact rankings. Implementation and experimental conditions need clarification rather than a guessed damping mechanism.
• Nonstationarity of the original task must be distinguished from the exact quadratic surrogate in the appendix. The paper does not prove that lower weighted layer error necessarily improves final generation quality.
• Calibration and test dataset descriptions conflict, and the small human study lacks sufficient cross-model and cross-seed evidence. The main text's reference to SD3 at 60% also exceeds the highest 50% shown in the corresponding main table and figures.
• Unstructured sparsity is not a measured speedup: practical latency gains depend on hardware or sparse-kernel support, and the paper does not demonstrate real inference acceleration. Fine-tuning recovery, broader generation tasks, and robustness on non-target categories remain open.

Related Work & Insights

• Compared with OBS-Diff: the method inherits timestep weighting, second-order ranking, and compensation. Its main change is curvature statistics computed from spatially filtered activations. The contribution is therefore a perceptually relevant objective modification, not a wholly new pruning optimizer.
• Compared with SparseGPT and layer-wise OBS: the methods share activation-covariance estimation and compensation through correlated weights. This paper emphasizes spatiotemporal content differences in diffusion calibration; exact local reconstruction should not be mistaken for an exact final-task guarantee.
• Compared with importance-based token merging: both can use CFG to identify important regions, but token merging changes computation on intermediate representations while this method modifies parameters. They could be combined in principle; the paper does not report measured gains for a joint method.

Rating

• Novelty: 3/5. Spatially aware reconstruction is a clear, targeted extension, while the solver largely follows OBS.
• Experimental Thoroughness: 3/5. Two backbones, multiple sparsity forms, and guidance signals are covered, but dataset-label conflicts, human-study size, and missing latency results constrain the conclusions.
• Writing Quality: 2/5. The core pipeline is understandable, but the global-coefficient explanation conflicts with the exact equations, and some experimental claims exceed the tabulated scope.
• Value: 4/5. The paper offers a reusable content-aware compression interface, with implementation reproducibility and mathematical consistency still requiring verification.