OFA-Diffusion Compression: Compressing Diffusion Model in One-Shot Manner
This work addresses the practical need for efficient deployment of diffusion models across devices with varying resource constraints by enabling one-shot compression, reducing the computational cost of model adaptation.
The paper proposes a once-for-all (OFA) compression framework for diffusion models that generates multiple subnetworks with different computational costs in a single training process, reducing the overhead of repeated compression for deployment on diverse devices. The method achieves satisfactory performance with significantly lower training overhead compared to existing compression techniques.
The Diffusion Probabilistic Model (DPM) achieves remarkable performance in image generation, while its increasing parameter size and computational overhead hinder its deployment in practical applications. To improve this, the existing literature focuses on obtaining a smaller model with a fixed architecture through model compression. However, in practice, DPMs usually need to be deployed on various devices with different resource constraints, which leads to multiple compression processes, incurring significant overhead for repeated training. To obviate this, we propose a once-for-all (OFA) compression framework for DPMs that yields different subnetworks with various computations in a one-shot training manner. The existing OFA framework typically involves massive subnetworks with different parameter sizes, while such a huge candidate space slows the optimization. Thus, we propose to restrict the candidate subnetworks with a certain set of parameter sizes, where each size corresponds to a specific subnetwork. Specifically, to construct each subnetwork with a given size, we gradually allocate the maintained channels by their importance. Furthermore, we propose a reweighting strategy to balance the optimization process of different subnetworks. Experimental results show that our approach can produce compressed DPMs for various sizes with significantly lower training overhead while achieving satisfactory performance.