OFA-Diffusion Compression: Compressing Diffusion Model in One-Shot Manner

TL;DR

This paper introduces a one-shot compression framework for diffusion probabilistic models, enabling efficient deployment across devices with different resource constraints by producing multiple subnetworks during a single training process.

Contribution

The proposed OFA-Diffusion Compression method allows for one-shot training of multiple subnetworks with different sizes, reducing overhead compared to traditional multiple compression processes.

Findings

Produces compressed DPMs of various sizes with lower training overhead.

Maintains satisfactory performance across different subnetworks.

Uses importance-based channel allocation and reweighting strategies.

Abstract

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.