OBS-Diff: Potatura Precisa per Modelli di Diffusione in Modalità One-Shot

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

October 8, 2025

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

cs.AI

Abstract

I modelli di diffusione testo-immagine su larga scala, sebbene potenti, soffrono di costi computazionali proibitivi. I metodi esistenti di pruning one-shot difficilmente possono essere applicati direttamente a essi a causa della natura iterativa di denoising dei modelli di diffusione. Per colmare questa lacuna, questo articolo presenta OBS-Diff, un nuovo framework di pruning one-shot che consente una compressione accurata e senza addestramento di modelli di diffusione testo-immagine su larga scala. Nello specifico, (i) OBS-Diff rivitalizza il classico Optimal Brain Surgeon (OBS), adattandolo alle complesse architetture dei moderni modelli di diffusione e supportando diverse granularità di pruning, inclusa la sparsità non strutturata, semi-strutturata N:M e strutturata (teste MHA e neuroni FFN); (ii) Per allineare i criteri di pruning con le dinamiche iterative del processo di diffusione, esaminando il problema da una prospettiva di accumulo di errori, proponiamo una nuova costruzione timestep-aware dell'Hessian che incorpora uno schema di ponderazione a decrescita logaritmica, assegnando maggiore importanza ai timestep iniziali per mitigare il potenziale accumulo di errori; (iii) Inoltre, viene proposta una strategia di pruning sequenziale efficiente dal punto di vista computazionale per ammortizzare il costoso processo di calibrazione. Esperimenti estesi dimostrano che OBS-Diff raggiunge lo stato dell'arte nel pruning one-shot per i modelli di diffusione, offrendo un'accelerazione dell'inferenza con una degradazione minima della qualità visiva.

English

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.