OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot
Authors: name: Junhan Zhu, name: Hesong Wang, name: Mingluo Su, name: Zefang Wang, name: Huan Wang
Organization: Westlake University

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.