[Paper Review] Successfully Applying Lottery Ticket Hypothesis to Diffusion Model

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

Advanced Neuroimaging Techniques and Applications4 citations

TL;DR

This paper首次 applies the Lottery Ticket Hypothesis (LTH) to diffusion models, demonstrating that subnetworks with up to 99% sparsity can match the performance of full models. By introducing layer-wise varying sparsity based on module similarity (measured via CKA), the method identifies sparser, more efficient winning tickets without performance loss, reducing FLOPs and memory usage.

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.

Motivation & Objective

• To investigate whether the Lottery Ticket Hypothesis applies to diffusion models, particularly denoising diffusion probabilistic models (DDPM).
• To reduce the high computational and memory costs of training and inference in diffusion models through structured pruning.
• To explore whether varying pruning ratios across different layers can yield sparser, more efficient winning tickets than uniform sparsity.
• To empirically validate that subnetworks identified via magnitude-based pruning with original weight initialization can match full model performance.

Proposed method

• Applies magnitude-based iterative pruning to U-Net-based DDPMs on CIFAR-10, CIFAR-100, and MNIST datasets.
• Uses a one-shot pruning strategy with iterative retraining and weight reinitialization to identify winning tickets.
• Introduces layer-wise varying sparsity by analyzing module similarity via Centered Kernel Alignment (CKA) across multiple winning ticket iterations.
• Configures lower pruning ratios (higher sparsity) in upstream layers and higher in downstream layers based on observed similarity trends.
• Employs standard DDPM hyperparameters (learning rate, epochs, time steps) and trains on 8 A100 GPUs.
• Validates performance using FID and IS metrics, comparing pruned subnetworks to full models.

Experimental results

Research questions

• RQ1 Can the Lottery Ticket Hypothesis be successfully applied to diffusion models, particularly DDPMs, to identify sparse, high-performing subnetworks?
• RQ2 Does uniform sparsity across all layers in LTH-based pruning lead to optimal or suboptimal winning tickets in diffusion models?
• RQ3 How does the similarity between winning tickets vary across different layers in a diffusion model?
• RQ4 Can layer-wise varying sparsity improve the sparsity of winning tickets without degrading performance?
• RQ5 What is the impact of pruning on FLOPs and memory usage in pruned diffusion models?

Key findings

• Winning tickets with up to 99% sparsity were successfully identified in DDPMs on CIFAR-10, CIFAR-100, and MNIST without performance degradation.
• The similarity between winning tickets was consistently higher in upstream layers than in downstream layers, as measured by Centered Kernel Alignment (CKA).
• A layer-wise varying sparsity strategy, with lower pruning in upstream layers, produced sparser winning tickets than uniform sparsity while maintaining performance.
• The proposed method achieved up to 90% reduction in FLOPs and reduced memory usage, enabling more efficient training and inference.
• The results confirm that structured, non-uniform pruning based on module similarity can yield more efficient subnetworks in diffusion models.