Plug-and-Play Diffusion Distillation

To reduce the expensive computational cost in inference time of diffusion models, previous papers have attempted to improve the sampling speed of diffusion models.

One of the straightforward ways is designing accurate ODE samplers [11, 10, 6]. For example, Denoising Diffusion Implicit Models (DDIM) [5] uses first-order Euler’s method which enables to reduce the inference timesteps.

On the other hand, there are previous works that attempt to incorporate distillation techniques to improve model inference efficiency. Distillation in deep learning refers to a process where a larger and more complex model (referred to as the “teacher” model) is employed to train a smaller and simpler model (the “student” model). The goal of distillation is to transfer the knowledge and information captured by the teacher model to the student model, enabling the student model to achieve similar performance with reduced complexity and computational requirements. Golnari et al. [3] proposed optimizing specific denoising steps by restricting noise computation to conditional noise and eliminating unconditional noise computation, thus reducing the complexity of target iterations. Salimans and Ho [23] and Meng et al. [14] distilled the model to achieve fewer sampling steps. However, these methods only focus on reducing inference timesteps and require progressive model distillation, which may require more time and computing resources to train these models.

Recently, LCM-LoRA[12] proposed a plug-and-play distillation approach utilizing LoRA, which has garnered significant attention. However, their method has the drawback of requiring double computation due to the use of classifier-free guidance. Our work, completed contemporaneously and accepted by CVPR, addresses similar challenges. We acknowledge the impact and relevance of their contribution.

CoDi[13], a concurrent work presented at CVPR, excels in producing high-quality images with very few steps (e.g., 1-4) across multiple tasks, including super-resolution, text-guided image editing, and depth-to-image generation. We acknowledge their valuable contribution.

Many researchers in the field of diffusion models have demonstrated the ability to control models using methods beyond just text input. Notably, the use of external models to inject features into diffusion models has yielded impressive results [7, 29, 8, 28]. For instance, ControlNet [29] proposed an external model that uses images, skeletons, edge maps, etc., as conditions to generate corresponding images. GLIGEN [8] successfully created desired objects within specific bounding boxes. IP-adaptor [28] introduced a method for generating images similar to a given image condition. These approaches all successfully manipulated images by injecting values into features through external models. However, these methods have focused on conditional image generation or editing, with no instances of applying them to distillation.

Under the continuous time setting, where $t\sim Uniform[0,1]$, the goal of a denoising diffusion model is to train a model $\epsilon_{\phi}$ that approximates noise given the diffused noisy real data $x\sim p_{\text{data}}$:

where $\omega(\lambda_{t})=\omega(\log\left(\alpha_{t}^{2}/\sigma_{t}^{2}\right)$) is a pre-defined weighted function that takes into the signal-to-noise ratio $\lambda_{t}$, which decreases monotonically with $t$. $x_{t}$ is a latent variable that satisfied $x\sim q(x_{t}|x)=\mathcal{N}(x_{t};\alpha_{t}x,\sigma_{t}^{2}I)$.

After training the model $\epsilon_{\phi}$, during the sampling stage, $x_{t}$ can be obtained by applying the SDE / ODE solver. For example, using DDIM:

where $N$ is the total number of sampling steps and $x_{1}\sim\mathcal{N}(0,I)$

Classifier-free guidance [4] proves to be a highly effective strategy for significantly enhancing the quality of samples in class-conditioned diffusion models. It adopted an unconditioned class identifier $\varnothing$ as a substitute for a separate classifier that is traditionally required to create a Gaussian distribution tailored to a specific class.[2] This approach finds widespread application in extensive diffusion models, including notable examples like DALL·E2 [18], GLIDE [15], and Stable Diffusion [19]. In particular, Stable Diffusion designs the diffused forward and reverse process in the VAE latent space, $z=E(x),x=D(z)$ where $E$ and $D$ denote the VAE encoder and decoder. In the process of generating a sample, classifier-free guidance carries out evaluations on both conditional score estimates and unconditional score estimates. Specifically, the computation of the noise sample $\tilde{\epsilon_{\phi}}(z_{t},c)$ follows the formulation

where $\epsilon_{\phi}$ is the score estimate function that is a parameterized neural network (U-Net). $\epsilon_{\phi}(z_{t},c)$ represents the text-conditioned term, while $\epsilon_{\phi,}(z_{t},\varnothing)$ corresponds to the unconditional term (null text). The parameter $g$ stands for the guidance value that scales the perturbation. In this paper, we use the Stable Diffusion’s VAE latent and omit the notation of Encoder and Decoder of VAE for brief.

Inspired by ControlNet [29], we design the external guide network for CFG distillation by using the guidance number as the input condition. After the first stage of the distillation (i.e. CFG distillation) has been accomplished, we follow prior distillation techniques to reduce the sampling steps. This is accomplished by enabling the model to progressively learn how to halve the sampling steps [23]. The specifics of the whole process will be elucidated in the following sections.

The overview of our CFG distillation method is illustrated in Figure 3. We would like to learn a model $\epsilon_{\theta}^{\prime}$ to achieve

where $g$ is the guidance number, $G$ is our student guided model, $\epsilon_{\phi}(z_{t},\varnothing)$ is the unconditioned U-Net forward pass, and $\epsilon_{\phi}(z_{t},c)$ is the conditioned U-Net forward pass. Precisely, $G$ takes the guidance as the input hint, along with time and text embedding and $z_{t}$, then injects its output feature maps to the decoder part of the original U-Net. The feature map injection can be viewed as the “guidance strength” that helps the U-Net to trade-off between sample quality and diversity. The pseudo algorithm is listed in Algorithm 1.

Typically, distillation involves initializing an entirely new model that has the same structure as the teacher model and trying to make it learn the teacher’s output and update the parameters of the entire student network. Instead, we use a small guide model on top of the teacher model, which leads to reduced computational overhead during training because the number of parameters in the guide models is relatively small compared to the whole U-Net. Also, this approach does not discard the teacher model after distillation training, but uses the trained guide model along with the teacher U-Net for faster inference without CFG. This feature makes it applicable to ”plug-and-play” to different types of fine-tuned diffusion models directly without retraining the guide model $G$.

In this section, we introduce two types of external guide model, full guide model and tiny guide model.

After training the guide model, $G$, we progressively distill it with fewer sampling steps required by incorporating with existing sampling-step-based distillation methods [23]. To elaborate, under the discrete time-step scenario, let $N$ stand for the original number of sampling steps, we trained a student model to the output of two-step DDIM sampling of the teacher in one step. Precisely, the initial sampler $f(z;\eta)$ maps a random noise $\epsilon$ to samples $x$ requires $N$ steps, is distilled into a new sampler $f(z;\theta)$ that requires $N/2$ steps. $f(z;\theta)$ will become the new teacher so that we can learn another sampler that requires $N/4$ steps. This procedure will be repeated several times until the ideal sampling steps needed will be achieved. In this section, again, we only learn the parameters from the guide model $G$ and fix the base model (U-Net) throughout the distillation progress. The small size of the guide model enables the parameters to be learned quickly.

We trained our model with LAION ($512\times 512$) dataset [24] with the Stable Diffusion v1.5 as our score-estimation model. In the training stage, a randomly sampled guidance number $g\in[2,9]$ is broadcast into the shape of (C, H, W), which becomes the input of the guide model. For the tiny architecture, the input is a 1d-array with the length of $C$ that passes through the zero modules along with timesteps and text embedding. We apply the $\varepsilon$-prediction model through the whole experiment. Since Stable Diffusion v1.5 is trained on $1000$ steps, we sample images with $1000$ steps as our Teacher output for our guide model to learn. We evaluate our methods with the COCO dataset [9]. We compare our method with DDIM [30] sampling and PLMS [10] sampling. The FLOPs and number of parameters for our models compared to the teacher classifier-gree guidance are listed in Table 1. We see that our full guide model only needs $\sim$ 0.67 of FLOPs of the teacher while our tiny model only computes   0.51 of FLOPs of the teacher model.

Figure 6 illustrates the qualitative comparison between student and teacher models on various text prompts with the same initial noise. We see the quality of images generated by our guide model is close to the generated images using classifier-free-guidance while our approach can generate images and nearly half the FLOP counts. A user study associated with images from teacher and student models can be found in the Appendix.

Furthermore, we generate images with fewer timesteps on Figure 5. We do not observe obvious quality degradation when decreasing our model steps to 16 and 8 with full guide model given a certain level of guidance ($g$ = 8). On the other hand, since tiny guide model has less capacity, it’s challenging for the student to fully mimic the teacher’s output given a continuous guidance input during the sampling steps distillation process (i.e. progressive distillation). We observe that the tiny guide model can achieve almost the same image quality as full guide model when the sampling steps are around 50, but when the number of steps are reduced to 8, then the performance of the tiny model will degrade drastically. This can be partially addressed by training the tiny model with fixed guidance. Furthermore, both of the models can achieve comparable results compared to Classifier-Free Guidance (DDIM $N\times 2$ steps). The qualitative result is shown in Table 2.

In this subsection, we show that plug in our guide model to different types of fine-tuned stable diffusion models from Dreambooth [20] without any additional training. Our objective is to demonstrate that our guide model can acquire a general latent representation of guidance, which possesses significant adaptability to various types of fine-tuned models without training. We focus on three different types of fine-tuned stable diffusion v1.5 models: watercolor style, realistic style, and 3D cartoon style. Then, we directly plug in our pretrained guide model $G$  to these fine-tuned models to modify their outputs. We run the models without classifier-free guidance and pass the guidance value to our guide module. For other distillation approaches [14], it may be necessary to distill a new model for each domain, which can be costly in terms of training and computation. Our approach removes this burden and make it easy to make different models finetuned for different domains nearly two times efficiently with no additional cost. We validate our approach by measuring FID and CLIP scores in generated images in different domains. Table 3 shows the FID scores and CLIP scores of the teacher CFG on these fine-tuned models versus the results of our tiny guide model injection approach. The generated pictures are sampled with 50 steps with guidance 8. We see that FID and CLIP scores of our tiny model which runs two times faster are comparable with CFG. In addition, qualitative results are shown in Figure 2 for the full guide model plug-ins and Figure 8 for tiny guide model plug-ins. The results indicate the great generalizability of our approach without needing to train the model for different domains.

In this experiment, our objective is to elucidate the latent representations within the feature map injections of our guide model $G$. To this end, we visualize the feature maps at various stages of the iteration process and under different guidance values. To the best of our knowledge, we are the first to visualize how classifier-free guidance emphasizes different patches of image generation in different timesteps. We are able to study this due to the architecture choice of our model that freezes the original model and adds the guide module as an additional component. By looking into feature maps of our guide module, we are able to get a better understanding of how classifier-free guidance impacts image generation.

For each layer of feature map injection, we computed the mean across various channels for each pixel and applied normalization. The number of DDIM steps used for sampling was 50.

Figure 7 displays the feature map injections throughout the sampling process. The values indicate that the initial stages of the sampling process are more critical with respect to Classifier-Free Guidance (CFG), as this is when the primary structure of the image is formed. In the middle stage, the main subjects of the image (e.g., a panda, bamboo) are more important. Thus CFG continues to play a role in these areas, while the background becomes less significant. Finally, in the last stage of the sampling, the feature map injections mainly focus on detail refinement on the edges with low strengths (i.e. values).

Additionally, an examination of the feature maps with varying guidance values, as shown in Figure 7, reveals a clear trend: with lower guidance, the feature map injections are less pronounced, whereas higher guidance results in more robust injections that more effectively steer the original diffusion model. Visualizations of other layers of feature maps can be found in the Appendix.