PaperlayerView exact revision on arXiv

Source-linked AI summary

OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot

Junhan Zhu, Hesong Wang, Mingluo Su, Zefang Wang, Huan Wang

arXiv:2510.06751v3cs.CV

TL;DR

Large-scale text-to-image diffusion models are computationally expensive, while existing one-shot pruning methods do not directly account for iterative denoising. OBS-Diff adapts OBS with timestep-aware Hessians and group-wise calibration, achieving state-of-the-art training-free pruning with inference acceleration and limited visual-quality degradation.

  • Problem

    Existing one-shot pruning methods are difficult to apply directly to large-scale text-to-image diffusion models because iterative denoising causes pruning errors to accumulate across timesteps.

  • Method

    OBS-Diff adapts Optimal Brain Surgeon with timestep-aware Hessian construction and group-wise sequential pruning for diverse diffusion-model architectures and sparsity patterns.

  • Results

    OBS-Diff achieves state-of-the-art training-free diffusion-model pruning, outperforming layer-wise baselines while accelerating inference and preserving visual quality across sparsity levels and patterns.

  • Takeaways & Limitations

    OBS-Diff provides a broadly applicable one-shot compression framework for large-scale text-to-image diffusion models without pruning-time training or fine-tuning.

  • Takeaways & Limitations

    FID may be unreliable for evaluating pruning methods because pruned models can occasionally outperform dense models on this metric without producing better visual results.

Abstract

from arXiv · show

Large-scale text-to-image diffusion models, while powerful, suffer from prohibitive computational cost. Existing one-shot network pruning methods can hardly be directly applied to them due to the iterative denoising nature of diffusion models. To bridge the gap, this paper presents OBS-Diff, a novel one-shot pruning framework that enables accurate and training-free compression of large-scale text-to-image diffusion models. Specifically, (i) OBS-Diff revitalizes the classic Optimal Brain Surgeon (OBS), adapting it to the complex architectures of modern diffusion models and supporting diverse pruning granularity, including unstructured, N:M semi-structured, and structured (MHA heads and FFN neurons) sparsity; (ii) To align the pruning criteria with the iterative dynamics of the diffusion process, by examining the problem from an error-accumulation perspective, we propose a novel timestep-aware Hessian construction that incorporates a logarithmic-decrease weighting scheme, assigning greater importance to earlier timesteps to mitigate potential error accumulation; (iii) Furthermore, a computationally efficient group-wise sequential pruning strategy is proposed to amortize the expensive calibration process. Extensive experiments show that OBS-Diff achieves state-of-the-art one-shot pruning for diffusion models, delivering inference acceleration with minimal degradation in visual quality.

1 INTRODUCTION

OBS-Diff addresses the computational burden and limited applicability of existing pruning methods for large-scale text-to-image diffusion models. It provides a one-shot, training-free framework supporting diverse architectures, pruning granularities, and efficient calibration.

  • OBS-Diff targets the prohibitive computational and memory demands of billion-parameter text-to-image diffusion models.
  • Existing pruning methods are often architecture-specific, computationally costly, and insufficiently explored for unstructured and semi-structured sparsity in large-scale text-to-image models.
  • OBS-Diff adapts Optimal Brain Surgeon to modern diffusion architectures and supports unstructured, N:M semi-structured, and structured pruning of attention heads or FFN neurons.
  • Timestep-aware Hessian construction prioritizes earlier denoising steps, while group-wise sequential pruning amortizes calibration costs across module packages.
  • OBS-Diff achieves state-of-the-art training-free diffusion pruning, accelerating inference while maintaining high visual quality across sparsity levels and patterns.

2 RELATED WORK

Prior diffusion-model pruning methods commonly depend on retraining, fine-tuning, or architecture-specific structured pruning. Related work also includes Hessian-based and LLM-oriented post-training methods, but their application to large-scale diffusion models remains limited.

  • Diff-pruning uses gradient-based structured pruning but was demonstrated on small non-text-to-image models and requires expensive retraining.
  • EcoDiff generalizes across text-to-image architectures but requires costly mask-learning training and extensive hyperparameter tuning.
  • Diffusion pruning research has largely emphasized architecture-specific structured methods, leaving large-scale unstructured and semi-structured pruning relatively unexplored.
  • Classical OBS and later layer-wise methods made Hessian-based pruning tractable through local approximations, while LLM methods extended efficient unstructured, semi-structured, and structured pruning.

3 PRELIMINARIES

Layer-wise post-training pruning minimizes output reconstruction error under target sparsity, and OBS uses second-order information to remove weights while compensating the remaining weights. Efficient implementations update inverse-Hessian information as pruning proceeds.

  • Layer-wise pruning finds a sparse weight matrix that minimizes output reconstruction error for given activations and target sparsity.
  • OBS-Diff extends this workflow by grouping modules into sequential packages, collecting calibration statistics jointly, and pruning package layers simultaneously.
  • OBS decouples the L2 reconstruction objective by weight-matrix row and uses a Hessian-based second-order approximation.
  • The OBS saliency score identifies the weight whose removal minimally increases error, while the closed-form update compensates the remaining weights in its row.
  • The saliency calculation uses the q-th diagonal element and q-th column of the inverse Hessian.
  • Pruning repeats until target sparsity is reached, with efficient inverse-Hessian updates avoiding repeated full matrix inversion.

4 METHODOLOGY

OBS-Diff adapts Optimal Brain Surgeon to iterative diffusion denoising with timestep-aware importance weighting and group-wise calibration. The framework also extends to semi-structured and structured pruning across diffusion-model components.

  • 4.1 TIMESTEP-AWARE HESSIAN CONSTRUCTION: Early denoising errors receive greater weight because they propagate through subsequent steps, motivating timestep-aware pruning criteria.The method assigns higher importance to earlier timesteps to account for error accumulation across the denoising trajectory.
  • 4.1 TIMESTEP-AWARE HESSIAN CONSTRUCTION: A logarithmically decreasing schedule makes α1 > α2 > ··· > αT, prioritizing the beginning of inference.The weights decrease smoothly while remaining positive across the T denoising steps.
  • 4.1 TIMESTEP-AWARE HESSIAN CONSTRUCTION: The Timestep-Aware Hessian aggregates second-order information across weighted denoising steps, making inverse-Hessian saliency sensitive to early-stage parameter importance.This adapts OBS to the temporal dynamics of diffusion generation.
  • 4.2 MODULE PACKAGES: A GROUP-WISE SEQUENTIAL PRUNING STRATEGY: Module Packages amortize calibration by collecting statistics for multiple layers during one complete denoising trajectory and pruning them simultaneously.Basic Units contain layers with mutually independent inputs, while packages group one or more Basic Units for collective processing.
  • 4.2 MODULE PACKAGES: A GROUP-WISE SEQUENTIAL PRUNING STRATEGY: Package processing updates the network sequentially between groups but keeps it static within each group, trading fewer calibration runs for higher memory use.The method reports low sensitivity to package granularity, allowing computational cost and memory to be balanced.
  • 4.3 EXTENSION TO SEMI-STRUCTURED AND STRUCTURED PRUNING: OBS-Diff supports 2:4 semi-structured pruning by removing the two lowest-saliency weights in each four-weight block.Its structured extensions aggregate weight saliency for FFN neurons and MHA heads, with modality-specific head rankings fused using reciprocal rank fusion.

5 EXPERIMENTS

OBS-Diff is evaluated across diverse diffusion architectures and pruning granularities, where it generally preserves semantic and visual quality better than competing methods while providing practical acceleration. Ablations further examine timestep weighting, calibration packaging, and calibration-set size.

  • Unstructured pruning: OBS-Diff consistently leads ImageReward across benchmarks and is best on CLIP score in most test cases, with only slight decreases from dense models.These results cover unstructured pruning comparisons on text-to-image diffusion models.
  • Unstructured pruning: At high sparsity, baseline quality collapses while OBS-Diff continues producing coherent images; pruning SD3-Medium completes in under 15 minutes on one RTX 4090.The reported collapse occurs at 60% sparsity on SD3.5-Large and 70% on Flux.1-dev.
  • Semi-structured pruning: Under 2:4 sparsity, OBS-Diff surpasses the strongest baseline in CLIP score (0.3129) and ImageReward (0.4493), despite Wanda’s slightly better FID of 32.08 versus 32.13.The comparison is reported for SD3.5-Large with pruning applied to MMDiT blocks 3 through 25.
  • Structured pruning: On SDXL at 30% sparsity, OBS-Diff achieves FID 29.75 versus EcoDiff’s 101.96, confirming stronger structured-pruning performance across U-Net and MMDiT architectures.The comparison is reported in Table 3 for the U-Net-based SDXL model and alongside Table 4 for SD3.5-Large.
  • Ablation study: A logarithmic decrease in timestep weights outperforms alternative weighting strategies, while more module packages reduce peak memory but increase runtime without a clear performance trend.Calibration-set ImageReward improves sharply through 100 prompts before plateauing.

6 CONCLUSION

OBS-Diff provides a one-shot, training-free pruning framework for large-scale text-to-image diffusion models. It combines timestep-aware Hessian construction, group-wise sequential pruning, and multiple sparsity structures to preserve quality at high sparsity.

  • 6 CONCLUSION: OBS-Diff combines timestep-aware Hessian construction with group-wise sequential pruning to address iterative denoising and calibration costs.The framework prioritizes critical early generation steps and balances memory overhead with computational efficiency.
  • 6 CONCLUSION: The framework supports unstructured, semi-structured, and structured pruning, including broad applicability across diffusion-model architectures.Structured pruning covers attention heads and FFN neurons.
  • 6 CONCLUSION: OBS-Diff establishes state-of-the-art training-free diffusion pruning while maintaining high generative quality, especially at high sparsity regimes.

B IMPLEMENTATION DETAILS

The implementation details define common experimental configurations and report the computational resources used for evaluation. These settings support controlled comparisons across large-scale diffusion models.

  • B IMPLEMENTATION DETAILS: Experiments use common configurations to enable controlled and fair comparisons across text-to-image generation tasks.The supplied implementation passage specifies shared settings for the evaluated methods.
  • B IMPLEMENTATION DETAILS: Text-to-image generation experiments use 512 × 512 output resolution and group parameters into four module packages for OBS-Diff and selected baselines.
  • B IMPLEMENTATION DETAILS: Large text-to-image models are pruned and evaluated training-free on NVIDIA RTX 4090 GPUs, while DDPM experiments use NVIDIA A100 GPUs.

C.1 MORE ANALYSIS FOR UNSTRUCTUREDLY PRUNED SD3-MEDIUM

Additional analyses show that OBS-Diff remains robust across unstructured and structured pruning settings, with stronger quality preservation than baselines at high sparsity. Its pruning overhead is slightly higher than DSnoT but remains close in runtime, while its pruned DDPM checkpoint also performs better after equal-budget fine-tuning.

  • C.1 MORE ANALYSIS FOR UNSTRUCTUREDLY PRUNED SD3-MEDIUM: At 60% sparsity, OBS-Diff retains a positive ImageReward while competing unstructured-pruning methods produce negative scores.The method outperforms all baselines across evaluated sparsity levels, with the largest advantage at higher sparsity.
  • C.1 MORE ANALYSIS FOR UNSTRUCTUREDLY PRUNED SD3-MEDIUM: 14.95 vs. 14.25 minutes: OBS-Diff takes slightly longer to prune than DSnoT on SD3-Medium at 50% sparsity.The reported overhead is characterized as marginal relative to the quality and robustness gains.
  • C.1 MORE ANALYSIS FOR UNSTRUCTUREDLY PRUNED SD3-MEDIUM: OBS-Diff maintains performance close to the dense model across structured-pruning sparsity levels, whereas L1-norm magnitude pruning degrades severely as sparsity increases.The comparison covers 10%, 15%, 20%, and 25% sparsity on Stable Diffusion 3-Medium.
  • C.1 MORE ANALYSIS FOR UNSTRUCTUREDLY PRUNED SD3-MEDIUM: Under an identical 100K-step fine-tuning budget, OBS-Diff surpasses Diff-Pruning on CIFAR-10 DDPMs by achieving a superior FID score.The result suggests the OBS-Diff-pruned model is a more effective checkpoint for subsequent fine-tuning.

D.1 ROBUSTNESS TO CFG SCALES

The pruned SD3-Medium model remains robust when inference conditions differ from calibration settings. It performs better at a higher CFG scale and benefits from additional sampling steps.

  • D.1 ROBUSTNESS TO CFG SCALES: 0.7044 vs. 0.6425: the pruned model achieves higher ImageReward at CFG 9.0 than at the calibration CFG 7.0.The evaluation uses the SD3-Medium model pruned at 50% unstructured sparsity.
  • D.1 ROBUSTNESS TO CFG SCALES: The fixed pruned model is evaluated beyond its calibration conditions across CFG scales and inference-step counts.The robustness tests use conditions different from CFG 7.0 and 25 steps.
  • D.1 ROBUSTNESS TO CFG SCALES: A model calibrated at 25 steps generates higher-quality images when evaluated with 50 inference steps.The result supports robustness across sampling-step counts.

D.3 ROBUSTNESS ACROSS SAMPLERS

OBS-Diff generalizes across diffusion samplers and calibration-data distributions. Its pruning preserves sampler rankings and maintains comparable evaluation performance under out-of-distribution calibration.

  • Sampler generalization: On SD3-Medium, calibration on Euler generalizes zero-shot to Heun with significant quality gains.This tests transfer from a first-order solver to a second-order solver.
  • Sampler generalization: Pruning preserves sampler rankings between dense and pruned SD v2.1 models, indicating solver-agnostic behavior.The comparison uses 40% unstructured pruning calibrated on PNDM; DPM++ remains the highest-performing sampler.
  • Calibration-data robustness: Calibrating on GCC3M yields performance nearly identical to MS-COCO calibration when both models are evaluated on MS-COCO validation data.The out-of-distribution GCC3M model performs on par with, and slightly better in FID than, the in-distribution model.

E ADDITIONAL QUALITATIVE RESULTS

Additional qualitative evaluations compare OBS-Diff with established unstructured and structured pruning baselines across multiple diffusion architectures and sparsity levels. The results support consistent visual quality across these settings and motivate extensions beyond text-to-image diffusion models.

  • Additional qualitative comparisons: OBS-Diff consistently maintains superior visual quality across diverse architectures and high-sparsity regimes.The comparisons cover unstructured pruning on SD3-Medium and Flux 1.dev, plus structured pruning on SD3.5-Large.
  • Unstructured pruning: Unstructured comparisons evaluate Magnitude, DSnoT, Wanda, and OBS-Diff on SD3-Medium at 20%, 30%, 40%, and 50% sparsity.The SD3-Medium images use the same prompt and negative prompt at 512 × 512 resolution.
  • Unstructured pruning: On Flux 1.dev, the qualitative unstructured comparison is conducted at 70% sparsity against Magnitude, DSnoT, and Wanda.The figure directly tests the methods in a high-sparsity setting.
  • Structured pruning: Structured comparisons on SD3.5-Large evaluate OBS-Diff against the L1-norm baseline at 15%, 20%, 25%, and 30% sparsity.The results are presented across multiple qualitative figures.
  • Future extensions: The paper identifies possible extensions to diffusion LLMs, video generation models, vision-language models, and combinations with other efficiency techniques.These directions are presented as future opportunities rather than evaluated results in this section.
Loading 2510.06751v3…