OBS-Diff: Точное прореживание диффузионных моделей за один шаг

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

October 8, 2025

Авторы: Junhan Zhu, Hesong Wang, Mingluo Su, Zefang Wang, Huan Wang

cs.AI

Аннотация

Крупномасштабные диффузионные модели для генерации изображений по тексту, несмотря на свою мощь, страдают от чрезмерно высоких вычислительных затрат. Существующие методы одноэтапного сжатия нейронных сетей (pruning) практически неприменимы к ним из-за итеративной природы процесса удаления шума в диффузионных моделях. Чтобы устранить этот разрыв, в данной статье представлен OBS-Diff — новый одноэтапный фреймворк для сжатия, который позволяет точно и без дополнительного обучения сжимать крупномасштабные диффузионные модели. В частности: (i) OBS-Diff возрождает классический метод Optimal Brain Surgeon (OBS), адаптируя его к сложным архитектурам современных диффузионных моделей и поддерживая различные уровни сжатия, включая неструктурированное, полуструктурированное N:M и структурированное (головы MHA и нейроны FFN) разрежение; (ii) Чтобы согласовать критерии сжатия с итеративной динамикой диффузионного процесса, рассматривая проблему с точки зрения накопления ошибок, мы предлагаем новую конструкцию гессиана, учитывающую временные шаги, которая включает схему весового убывания по логарифмическому закону, придавая большее значение ранним шагам для минимизации потенциального накопления ошибок; (iii) Кроме того, предлагается вычислительно эффективная стратегия последовательного группового сжатия, чтобы распределить затраты на дорогостоящий процесс калибровки. Эксперименты показывают, что OBS-Diff достигает наилучших результатов в одноэтапном сжатии диффузионных моделей, обеспечивая ускорение вывода с минимальной деградацией визуального качества.

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.