OBS-Diff: Poda precisa para modelos de difusión en un solo paso

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

Resumen

Los modelos de difusión de texto a imagen a gran escala, aunque potentes, sufren de un costo computacional prohibitivo. Los métodos existentes de poda de redes en una sola pasada difícilmente pueden aplicarse directamente a ellos debido a la naturaleza iterativa de eliminación de ruido de los modelos de difusión. Para cerrar esta brecha, este artículo presenta OBS-Diff, un novedoso marco de poda en una sola pasada que permite la compresión precisa y sin entrenamiento de modelos de difusión de texto a imagen a gran escala. Específicamente, (i) OBS-Diff revitaliza el clásico Optimal Brain Surgeon (OBS), adaptándolo a las arquitecturas complejas de los modelos de difusión modernos y apoyando diversos niveles de granularidad en la poda, incluyendo la esparcidad no estructurada, semi-estructurada N:M y estructurada (cabezas de MHA y neuronas de FFN); (ii) Para alinear los criterios de poda con la dinámica iterativa del proceso de difusión, al examinar el problema desde una perspectiva de acumulación de errores, proponemos una construcción novedosa de Hessian consciente del paso de tiempo que incorpora un esquema de ponderación de disminución logarítmica, asignando mayor importancia a los pasos de tiempo iniciales para mitigar la posible acumulación de errores; (iii) Además, se propone una estrategia de poda secuencial por grupos computacionalmente eficiente para amortizar el costoso proceso de calibración. Experimentos extensivos muestran que OBS-Diff logra la poda en una sola pasada más avanzada para modelos de difusión, ofreciendo aceleración en la inferencia con una degradación mínima en la calidad 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.