[論文レビュー] DiP-GO: A Diffusion Pruner via Few-step Gradient Optimization

Haowei Zhu, Dehua Tang+13 more|arXiv (Cornell University)|Oct 22, 2024

Optical Systems and Laser Technology被引用数 4

ひとこと要約

DiP-GO は、モデルの微調整を必要としない、差分可能なサブネット探索問題として定式化された、拡散モデルのための新しいプリーニング手法を提案する。特徴量の類似度に基づいて、非除去ステップ間でバックアップ接続を持つスーパーネットを構築し、最適なサブネットを少数ステップの勾配最適化によって特定する。最適化された損失関数を備えたプラグインプリーナーを導入し、Stable Diffusion 1.5 において、精度を損なわずに 4.4 倍の推論高速化を達成する。

ABSTRACT

Diffusion models have achieved remarkable progress in the field of image generation due to their outstanding capabilities. However, these models require substantial computing resources because of the multi-step denoising process during inference. While traditional pruning methods have been employed to optimize these models, the retraining process necessitates large-scale training datasets and extensive computational costs to maintain generalization ability, making it neither convenient nor efficient. Recent studies attempt to utilize the similarity of features across adjacent denoising stages to reduce computational costs through simple and static strategies. However, these strategies cannot fully harness the potential of the similar feature patterns across adjacent timesteps. In this work, we propose a novel pruning method that derives an efficient diffusion model via a more intelligent and differentiable pruner. At the core of our approach is casting the model pruning process into a SubNet search process. Specifically, we first introduce a SuperNet based on standard diffusion via adding some backup connections built upon the similar features. We then construct a plugin pruner network and design optimization losses to identify redundant computation. Finally, our method can identify an optimal SubNet through few-step gradient optimization and a simple post-processing procedure. We conduct extensive experiments on various diffusion models including Stable Diffusion series and DiTs. Our DiP-GO approach achieves 4.4 x speedup for SD-1.5 without any loss of accuracy, significantly outperforming the previous state-of-the-art methods.

研究の動機と目的

• 拡散モデル推論の高い計算コスト、特に多段階のノイズ除去に起因する問題に対処すること。
• 時系列にわたる動的特徴量の類似性を活用できない静的かつ非適応的なプリーニング戦略の限界を克服すること。
• 事前学習済みの拡散モデルの再訓練を必要とせず、効率的かつ高精度なモデルプリーニングを可能にすること。
• スーパーネットから最適なサブネット（サブネット）を特定する、差分可能で勾配最適化されたプリーニングフレームワークの開発。
• 多様な拡散アーキテクチャにおいて、顕著な高速化を実現しながら画像生成品質を維持すること。

提案手法

• 隣接するノイズ除去ステップ間の特徴量類似度に基づいてバックアップ接続を追加することで、標準的な拡散モデルにスーパーネットを構築する。
• プリーニングプロセスを、スーパーネットから最適なサブネットを導出するサブネット探索問題として定式化する。
• タスク固有の最適化損失関数を用いて、不要な計算パスを識別することができるプラグインプリーナーを設計する。
• 全再訓練と比較して計算コストを低減するため、少数ステップの勾配最適化を適用して最適なサブネットを効率的に探索する。
• ユーザー指定のプリーニング比および構造的制約を満たすように、最終的なサブネットを保証する軽量なポストプロセッシングステップを統合する。
• 拡張されたノイズ除去ステップにおける逆伝播処理でのメモリ問題を軽減するため、勾配チェックポイント技術を活用する。

Figure 1: Overview of the SuperNet and SubNet. Standard diffusion models execute the full inference path step by step. In our framework, we propose a SuperNet based on the original flow and integrate backup connections to facilitate block removal. This allows the partial inference SubNet to efficien

実験結果

リサーチクエスチョン

• RQ1 再訓練を回避するために、拡散モデルのプリーニングを差分可能なサブネット探索問題として定式化できるか？
• RQ2 隣接するノイズ除去ステップ間の特徴量類似度を効果的に活用することで、冗長な計算を削減できるか？
• RQ3 少数ステップの勾配最適化戦略は、従来のプリーニング手法と比較して、より効率的に高性能なサブネットを特定できるか？
• RQ4 生成品質を損なわず、どの程度まで推論を高速化できるか？
• RQ5 提案手法は、Stable Diffusion や DiT などの多様な拡散アーキテクチャに一般化できるか？

主な発見

• DiP-GO は、再訓練を要せず、Stable Diffusion 1.5 において画像品質に劣化を来さずに 4.4 倍の推論時間高速化を達成した。
• 本手法は、時系列にわたる動的特徴パターンを効果的に活用することで、従来の最先端のプリーニング技術を上回った。
• 本アプローチは、Stable Diffusion 1.5、2.1、XL および DiT など複数の拡散モデルに一般化可能であり、広範な適用性を示した。
• 最適化された損失関数を備えたプラグインプリーナーは、不要なコンponents を的確に同定し、効率的なサブネット発見を可能にした。
• 少数ステップの勾配最適化により、高パフォーマンスなサブネットへの収束が高速化され、探索コストが削減された。
• ポストプロセッシングステップにより、最終的なサブネットがプリーニング比制約を満たしつつ、モデルの能力を維持した。

Figure 2: Overview of our diffusion pruner. a) DiP-GO employs a pruner network to learn the importance scores of blocks in the diffusion sampling process. It takes $N\times T$ queries as input and passes them through stacked self-attention (SA) and fully connected (FC) layers to capture the structur