OBS-Diff : Élagage précis pour les modèles de diffusion en une seule étape

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

October 8, 2025

papers.authors: Junhan Zhu, Hesong Wang, Mingluo Su, Zefang Wang, Huan Wang

cs.AI

papers.abstract

Les modèles de diffusion à grande échelle pour la génération d'images à partir de texte, bien que puissants, souffrent de coûts de calcul prohibitifs. Les méthodes existantes d'élagage de réseau en une seule étape ne peuvent guère leur être directement appliquées en raison de la nature itérative du processus de débruitage des modèles de diffusion. Pour combler cette lacune, cet article présente OBS-Diff, un nouveau cadre d'élagage en une seule étape qui permet une compression précise et sans entraînement des modèles de diffusion à grande échelle pour la génération d'images à partir de texte. Plus précisément, (i) OBS-Diff revitalise l'Optimal Brain Surgeon (OBS) classique, en l'adaptant aux architectures complexes des modèles de diffusion modernes et en prenant en charge diverses granularités d'élagage, y compris la sparsité non structurée, semi-structurée N:M et structurée (têtes MHA et neurones FFN) ; (ii) Pour aligner les critères d'élagage avec la dynamique itérative du processus de diffusion, en examinant le problème sous l'angle de l'accumulation d'erreurs, nous proposons une nouvelle construction de Hessian prenant en compte les étapes temporelles, intégrant un schéma de pondération à décroissance logarithmique, attribuant une importance accrue aux étapes initiales afin d'atténuer l'accumulation potentielle d'erreurs ; (iii) En outre, une stratégie d'élagage séquentiel par groupes, efficace sur le plan computationnel, est proposée pour amortir le processus coûteux de calibration. Des expériences approfondies montrent qu'OBS-Diff atteint un état de l'art en matière d'élagage en une seule étape pour les modèles de diffusion, offrant une accélération de l'inférence avec une dégradation minimale de la qualité visuelle.

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.