[论文解读] 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

一句话总结

该论文首次将彩票理论（Lottery Ticket Hypothesis, LTH）应用于扩散模型，证明了在高达99%稀疏度下，子网络仍可达到完整模型的性能表现。通过基于模块相似性（以CKA度量）的逐层可变稀疏度方法，该方法在不损失性能的前提下识别出更稀疏、更高效的彩票票券，显著降低了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）。
• 通过结构化剪枝降低扩散模型在训练与推理过程中的高计算与内存开销。
• 探究在不同层中采用可变剪枝率是否能获得比均匀稀疏度更稀疏、更高效的彩票票券。
• 通过实证验证：基于原始权重初始化的幅度剪枝所识别出的子网络，可与完整模型性能相当。

提出的方法

• 在CIFAR-10、CIFAR-100与MNIST数据集上，对基于U-Net的DDPM应用基于幅度的迭代剪枝。
• 采用一次性剪枝策略，结合迭代微调与权重重初始化，以识别彩票票券。
• 通过在多个彩票票券迭代过程中分析模块相似性（使用中心化核对齐，CKA），引入逐层可变稀疏度。
• 根据观察到的相似性趋势，配置上游层更低的剪枝率（更高稀疏度），下游层更高的剪枝率。
• 采用标准DDPM超参数（学习率、训练轮数、时间步长），并在8张A100 GPU上进行训练。
• 使用FID与IS指标验证性能，将剪枝后的子网络与完整模型进行对比。

实验结果

研究问题

• RQ1 彩票理论能否成功应用于扩散模型，特别是DDPM，以识别出稀疏且高性能的子网络？
• RQ2 在基于LTH的剪枝中，若所有层采用均匀稀疏度，是否会导致彩票票券性能最优或次优？
• RQ3 在扩散模型中，不同层之间的彩票票券相似性如何变化？
• RQ4 是否可通过逐层可变稀疏度策略，在不降低性能的前提下提升彩票票券的稀疏度？
• RQ5 剪枝对剪枝后扩散模型的FLOPs与内存占用有何影响？

主要发现

• 在CIFAR-10、CIFAR-100与MNIST上的DDPM中，成功识别出稀疏度高达99%的彩票票券，且性能未出现下降。
• 通过中心化核对齐（CKA）度量，彩票票券在上游层的相似性始终高于下游层。
• 采用逐层可变稀疏度策略（上游层剪枝率更低）所生成的彩票票券，其稀疏度高于均匀稀疏度策略，同时保持了性能表现。
• 所提方法实现了FLOPs最高达90%的减少，并降低了内存占用，从而实现了更高效的训练与推理。
• 结果证实，基于模块相似性的结构化、非均匀剪枝，可在扩散模型中生成更高效的子网络。