[論文レビュー] Structural Pruning for Diffusion Models

Gongfan Fang, Xinyin Ma+1 more|arXiv (Cornell University)|May 18, 2023

Music and Audio Processing被引用数 17

ひとこと要約

Diff-Pruningは、Taylor展開ベースの構造プリuning手法で、タイムステップとウェイトを剪定して事前学習済みの拡散モデルを圧縮し、元のトレーニングコストの約10–20%で約50%のFLOPs削減を実現し、生成挙動を保つ。

ABSTRACT

Generative modeling has recently undergone remarkable advancements, primarily propelled by the transformative implications of Diffusion Probabilistic Models (DPMs). The impressive capability of these models, however, often entails significant computational overhead during both training and inference. To tackle this challenge, we present Diff-Pruning, an efficient compression method tailored for learning lightweight diffusion models from pre-existing ones, without the need for extensive re-training. The essence of Diff-Pruning is encapsulated in a Taylor expansion over pruned timesteps, a process that disregards non-contributory diffusion steps and ensembles informative gradients to identify important weights. Our empirical assessment, undertaken across several datasets highlights two primary benefits of our proposed method: 1) Efficiency: it enables approximately a 50\% reduction in FLOPs at a mere 10\% to 20\% of the original training expenditure; 2) Consistency: the pruned diffusion models inherently preserve generative behavior congruent with their pre-trained models. Code is available at \url{https://github.com/VainF/Diff-Pruning}.

研究の動機と目的

• 拡散確率モデル（DPMs）を圧縮してトレーニングと推論のオーバーヘッドを削減する必要性を動機づける。
• 拡散モデルに特化した専用のプリuning法（Diff-Pruning）を提案する。
• 重要なウェイトと剪定すべきタイムステップを識別するためのTaylor展開ベースの基準を開発する。
• プリuningが生成品質と一貫性を、さまざまなデータセットに渡って保存または改善できることを示す。

提案手法

• ウェイトのサブ構造全体を削除して疎なパラメータ行列を得るとしてのモデル剪定。
• 各タイムステップの損失L_tのTaylor展開を用いてパラメータの重要性を推定し、タイムステップ全体での影響を集約する（Equation 7の派生版）。
• relative loss L_t/L_max のしきい値機構を用いて剪定すべきタイムステップを選択するタイムステップ認識剪定を導入する（Equation 9/10）。
• 部分的なタイムステップにわたって勾配を蓄積して各パラメータの頑健な重要度スコアを計算する（Equation 10）。
• 事前学習済み拡散モデルに対してワンショット剪定を適用し、対象データセットでファインチューニングを行う。
• 複数データセット（CIFAR-10、CelebA-HQ、LSUN、ImageNet-1K）にわたって、効率（パラメータ、MACs）、品質（FID）、一貫性（SSIM）を評価する。

Figure 1 : Diff-Pruning leverages Taylor expansion at pruned timesteps to estimate the importance of weights, where early steps focus on local details like edges and color and later ones pay more attention to contents such as object and shape. We propose a simple thresholding method to trade off the

実験結果

リサーチクエスチョン

• RQ1 構造的剪定は、広範な再訓練を要さずに拡散モデルの冗長な構成要素を正確に識別して削除できるか？
• RQ2 タイムステップの剪定とウェイトの剪定は、拡散モデルの内容生成とディテール生成にどのように影響するか？
• RQ3 剪定比率、回復努力、生成サンプル品質のトレードオフは、データセットとモデルタイプ（DDPMs、LDMs）全体でどうなるか？

主な発見

• Diff-Pruningは、元のトレーニングコストの約10%–20%を使用しつつ、約50%のFLOPs削減という substantial compression を達成。
• 剪定モデルは、事前学習済みモデルと同等の、あるいは向上する生成挙動とサンプルの一貫性を維持する（例：LSUN Churchでは training steps 0.5M 対 4.4M）。
• コンテンツに寄与するタイムステップは拡散の終端だけではなく、剪定にはタイムステップの重要度で重み付けを行い、内容とディテールのバランスを取る必要がある。
• 全タイムステップにわたる完全なTaylor展開はノイズの多い勾配を蓄積し得るため、閾値付きの部分的なTaylor展開を用いると剪定精度が向上する。
• LSUN Church/Bedroom および ImageNet-1K-LDMs では、剪定モデルが、ベースラインと比較してはるかに少ないパラメータとMACs で競争力のFID/SSIMを達成。
• Diff-Pruningは、CIFAR-10 および CelebA-HQ において、ランダム、振幅、素朴なTaylor剪定を一貫して上回る。

Figure 2 : Generated images of the pre-trained models [ 18 ] (left) and the pruned models (right) on LSUN Church and LSUN Bedroom. SSIM measures the similarity between generated images.