[論文レビュー] Successfully Applying Lottery Ticket Hypothesis to Diffusion Model

Chao Jiang, Bo Hui+2 more|arXiv (Cornell University)|Oct 28, 2023

Advanced Neuroimaging Techniques and Applications被引用数 4

ひとこと要約

本論文は、拡散モデルに初めてランダムチケット仮説（LTH）を適用し、最大99％のスパarsityを持つサブネットワークが、完全なモデルと同等の性能を達成することを示した。モジュール類似度（CKAを用いて測定）に基づいて層ごとに異なるスパarsityを導入することで、性能に損なわれることなく、よりスパースで効率的なウィンニングチケットを同定した。これにより、FLOPsとメモリ使用量が削減された。

ABSTRACT

Despite the success of diffusion models, the training and inference of diffusion models are notoriously expensive due to the long chain of the reverse process. In parallel, the Lottery Ticket Hypothesis (LTH) claims that there exists winning tickets (i.e., aproperly pruned sub-network together with original weight initialization) that can achieve performance competitive to the original dense neural network when trained in isolation. In this work, we for the first time apply LTH to diffusion models. We empirically find subnetworks at sparsity 90%-99% without compromising performance for denoising diffusion probabilistic models on benchmarks (CIFAR-10, CIFAR-100, MNIST). Moreover, existing LTH works identify the subnetworks with a unified sparsity along different layers. We observe that the similarity between two winning tickets of a model varies from block to block. Specifically, the upstream layers from two winning tickets for a model tend to be more similar than the downstream layers. Therefore, we propose to find the winning ticket with varying sparsity along different layers in the model. Experimental results demonstrate that our method can find sparser sub-models that require less memory for storage and reduce the necessary number of FLOPs. Codes are available at https://github.com/osier0524/Lottery-Ticket-to-DDPM.

研究の動機と目的

• ランダムチケット仮説が、特にノイズ除去拡散確率的モデル（DDPM）に適用可能かどうかを調査すること。
• 構造的プルーニングを用いて、拡散モデルの訓練および推論における高い計算コストとメモリ使用量を低減すること。
• 異なる層でプルーニング比を変化させることで、均一なスパarsityよりもスパースでより効率的なウィンニングチケットが得られるかどうかを調査すること。
• マグニチュードベースのプルーニングと元の重み初期化を用いて同定されたサブネットワークが、完全なモデルと同等の性能を達成することを実証的に検証すること。

提案手法

• CIFAR-10、CIFAR-100、MNISTデータセットにおけるU-NetベースのDDPMに、マグニチュードベースの反復的プルーニングを適用した。
• 1回のプルーニング戦略を採用し、反復的再訓練と重みの再初期化を用いてウィンニングチケットを同定した。
• 複数のウィンニングチケット反復において、中心化カーネル整合性（CKA）を用いてモジュール類似度を分析することで、層ごとの可変スパarsityを導入した。
• 観察された類似度トレンドに基づき、上流層では低いプルーニング比（高いスパarsity）を、下流層では高いプルーニング比を設定した。
• 標準的なDDPMハイパーパrameter（学習率、エポック数、時間ステップ）を用い、8台のA100 GPUで訓練した。
• FIDおよびISスコアを用いて性能を検証し、プルーニングされたサブネットワークと完全なモデルを比較した。

実験結果

リサーチクエスチョン

• RQ1 ランダムチケット仮説は、特にDDPMに成功裏に適用可能であり、スパースで高性能なサブネットワークを同定できるか？
• RQ2 LTHベースのプルーニングにおいて、すべての層に均一なスパarsityを適用すると、拡散モデルでは最適なウィンニングチケットが得られるか、それとも劣悪な結果が得られるか？
• RQ3 拡散モデルの異なる層において、ウィンニングチケット間の類似度はどのように変化するか？
• RQ4 層ごとの可変スパarsity戦略により、性能を損なわせることなく、ウィンニングチケットのスパarsityを向上させることができるか？
• RQ5 プルーニングは、プルーニングされた拡散モデルにおけるFLOPsとメモリ使用量にどのような影響を与えるか？

主な発見

• CIFAR-10、CIFAR-100、MNISTにおけるDDPMにおいて、性能に劣化を来すことなく最大99％のスパarsityを持つウィンニングチケットが成功裏に同定された。
• 中心化カーネル整合性（CKA）で測定したところ、ウィンニングチケット間の類似度は、上流層で下流層よりも一貫して高かった。
• 上流層では低いプルーニング比を、下流層では高いプルーニング比を設定する層ごとの可変スパarsity戦略により、均一なスパarsityよりもスパースで効率的なウィンニングチケットが得られた。
• 提案手法により、FLOPsが最大90％削減され、メモリ使用量も削減され、より効率的な訓練および推論が可能になった。
• 結果として、モジュール類似度に基づく構造的かつ非均一なプルーニングが、拡散モデルにおけるより効率的なサブネットワークを生み出せることを確認した。