OBS-Diff: Poda Precisas para Modelos de Difusão em Uma Única Etapa

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

October 8, 2025

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

cs.AI

Resumo

Modelos de difusão de texto para imagem em grande escala, embora poderosos, sofrem com custos computacionais proibitivos. Os métodos existentes de poda de rede em uma única etapa dificilmente podem ser aplicados diretamente a eles devido à natureza iterativa de remoção de ruído dos modelos de difusão. Para preencher essa lacuna, este artigo apresenta o OBS-Diff, uma nova estrutura de poda em uma única etapa que permite a compressão precisa e sem treinamento de modelos de difusão de texto para imagem em grande escala. Especificamente, (i) o OBS-Diff revitaliza o clássico Optimal Brain Surgeon (OBS), adaptando-o às arquiteturas complexas dos modelos de difusão modernos e suportando diversas granularidades de poda, incluindo esparsidade não estruturada, semi-estruturada N:M e estruturada (cabeças MHA e neurônios FFN); (ii) Para alinhar os critérios de poda com a dinâmica iterativa do processo de difusão, ao examinar o problema sob a perspectiva de acumulação de erros, propomos uma nova construção de Hessiana consciente do passo de tempo que incorpora um esquema de ponderação de diminuição logarítmica, atribuindo maior importância aos passos de tempo iniciais para mitigar a potencial acumulação de erros; (iii) Além disso, uma estratégia de poda sequencial por grupos computacionalmente eficiente é proposta para amortizar o caro processo de calibração. Experimentos extensivos mostram que o OBS-Diff alcança a poda em uma única etapa de ponta para modelos de difusão, proporcionando aceleração de inferência com degradação mínima na qualidade visual.

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.