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

要旨

大規模なテキストから画像への拡散モデルは強力である一方で、計算コストが非常に高くなるという課題があります。既存のワンショットネットワークプルーニング手法は、拡散モデルの反復的なノイズ除去の性質上、直接適用することが困難です。このギャップを埋めるため、本論文ではOBS-Diffという新しいワンショットプルーニングフレームワークを提案します。これは、大規模なテキストから画像への拡散モデルを正確かつトレーニング不要で圧縮することを可能にします。具体的には、(i) OBS-Diffは古典的なOptimal Brain Surgeon (OBS)を再活用し、現代の拡散モデルの複雑なアーキテクチャに適応させ、非構造化、N:M半構造化、および構造化（MHAヘッドとFFNニューロン）スパース性を含む多様なプルーニング粒度をサポートします。(ii) プルーニング基準を拡散プロセスの反復的なダイナミクスに合わせるため、エラー蓄積の観点から問題を検討し、対数減少重み付けスキームを組み込んだ新しいタイムステップ対応Hessian構築を提案します。これにより、潜在的なエラー蓄積を軽減するため、初期のタイムステップにより大きな重要性を割り当てます。(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.