Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation

The proposed method is mainly driven by the desire to end up with a fast, robust, and reliable approach that would be easily transposed to different use cases. Given a set of data $x_{0}\in\mathcal{X}$ such that $x_{0}\sim p(x_{0})$ and $z_{0}=\mathcal{E}(x_{0})$ the associated latent variables, the main idea of the proposed approach is quite similar to diffusion models. Given a noise schedule defined by $\alpha(t)$ and $\sigma(t)$, we propose to create a noisy latent sample $z_{t}$ with $t\sim\pi(t)$ as specified in Eq. (1) and train a function $f_{\theta}^{\mathrm{student}}$ to predict a denoised version $\tilde{z}_{0}$ of the original sample $z_{0}$. The main difference with a diffusion model is that instead of using $z_{0}$ as a target, we propose to leverage the knowledge of the teacher model and use a sample belonging to the data distribution learned by the teacher model $p_{\phi}^{\mathrm{teacher}}(z_{0})$. In other words, we use the teacher model and an ODE solver to generate a denoised latent sample $\tilde{z}_{0}^{\mathrm{teacher}}(z_{t})$ that belongs to the learned data distribution and use it as a target for the student model. The main distillation loss writes as follows:

where $\pi(t)$ denotes the distribution over the timesteps, and $\tilde{z}_{0}^{\mathrm{teacher}}(z_{t})$ is obtained by running several steps of an ODE solver $\Psi$ on a teacher model $\varepsilon_{\phi}^{\mathrm{teacher}}$ starting from $z_{t}=\alpha(t)\cdot z_{0}+\sigma(t)\cdot\varepsilon$. A similar idea was employed in [62] but the authors generate fully synthetic samples meaning that the samples $z_{t}$ are pure noise, $z_{t}\sim\mathcal{N}\left(0,\mathbf{I}\right)$. In contrast, in our approach, we hypothesize that allowing $z_{t}$ to retain some information from the ground-truth encoded sample $z_{0}$ could enhance the distillation process. As in [42], when distilling a conditional DM, we also perform Classifier-Free Guidance (CFG) [16] with the teacher to better enforce the model to respect the conditioning. This technique actually significantly improves the quality of the generated samples by the student as shown in Sec. 5. Additionally, it eliminates the need for conducting CFG during inference with the student, further decreasing the method’s computational cost by halving the NFEs for each step. The value of the guidance scale $\omega$ used during training is part of the ablations presented in Sec. 5.2 but in practice $\omega$ is uniformly sampled in $[\omega_{\min},\omega_{\max}]$ where $0\leq\omega_{\min}\leq\omega_{\max}$. As mentioned in Sec. 3.2, our approach bears resemblance to existing consistency models [68, 42]. However, rather than relying on a previous instance of the student model to estimate the origin of the PF-ODE, we directly employ the teacher model coupled with an ODE solver to generate the target. We observed that these ingredients enhance the stability of the training procedure.

The cornerstone of our approach hinges on the selection of the timestep probability density function, denoted as $\pi(t)$. According to the continuous modeling, exposed in [67], DMs are trained to remove noise from a latent sample $z_{t}$ for any given continuous time $t$. However, since we aim at achieving few steps data generation (typically 1-4 steps) at inference time, the learned function $\varepsilon_{\theta}$ will only be evaluated at a few discrete timesteps $\{t_{i}\}_{i=1}^{K}$.

In order to tackle this issue and enforce the distillation process to focus on the most relevant timesteps, we propose to select $K$ (typically 16, 32, or 64) uniformly spaced timesteps in the range $[0,1]$ and assign a probability to each of them according to a probability mass function $\pi(t)$. We choose $\pi(t)$ as a mixture of Gaussian distribution controlled by a series of weights $\{\beta_{i}\}_{i=1}^{K}$

where the mean of each Gaussian is controlled by $\{\mu_{i}=i/K\}_{i=1}^{K}$ and the variance is fixed to $\sigma=\sqrt{0.5/K^{2}}$. This approach is such that when distilling the teacher only a small number of $K$ discrete timesteps will be sampled instead of the continuous range $[0,1]$¹¹1In practice when training a DM, the range $[0,1]$ is actually discretized (typically into 1000 timesteps) for computational purposes.. Moreover, the distribution $\pi$ is defined such that out of the $K$ selected timesteps, the 4 timesteps used at inference for 1, 2 and 4 steps generation are over-sampled (typically we set $\beta_{i}>0$ if $i\in[\frac{K}{4},\frac{K}{2},\frac{3K}{4},K]$ and $\beta_{i}=0$ otherwise). Unlike other methods [61, 62] we do not only focus on those 4 timesteps since we noticed that it can lead to a reduction of diversity in the generated samples. This is in particular emphasized in the ablation study presented in Sec. 5.2. In practice, we notice that a warm-up phase is beneficial to the training process. Therefore, we decide to start by first imposing a higher probability to the timesteps corresponding to the least added amount of noise by setting $\beta_{K/4}=\beta_{K/2}=0.5$ and $\beta_{i}=0$ otherwise. We then progressively shift the probability mass towards full noise to favor single-step generation while still over-sampling the targeted 4 timesteps by setting a strictly positive value for $\beta_{i}$ where $i\equiv 0[K/4]$, and $\beta_{i}=0$ otherwise. An example for $\pi$ with $K=32$ is illustrated in Figure 3. As pictured in the figure, the $[0,1]$ interval is split into 32 timesteps. During the warm-up phase, the probability mass allocates a higher probability to timesteps $[0.25,0.5]$ to ease the distillation process. As the training progresses, the probability mass function is then shifted towards full noise to favor single-step generation while always allocating a higher probability to the 4 timesteps $[0.25,0.5,0.75,1]$. The impact of the timesteps distribution is further discussed in Sec. 5.2.

To further enhance the quality of the samples, and since it proved very efficient in several works proposed in the literature to achieve few steps image generation [76, 61, 33, 28, 62], we have also decided to incorporate an adversarial objective. The core idea is to train the student model to generate samples that are indistinguishable from the true data distribution $p(x_{0})$. To do so, we propose to train a discriminator $D_{\nu}$ to distinguish the generated samples $\tilde{x}_{0}$ from the real samples $x_{0}\sim p(x_{0})$. As proposed in [33, 62], we also apply the discriminator directly within the latent space. This approach circumvents the necessity of decoding the samples using the VAE, a process outlined in [61], that proves to be expensive and hampers the method’s scalability to high-resolution images.

Drawing inspiration from [33, 62], we propose an approach where both the one-step student prediction $\tilde{z}_{0}$ and the input latent sample $z_{0}$ are re-noised following the teacher noise schedule. This process uses a timestep $t^{\prime}$, which is uniformly chosen from the set $[0.01,0.25,0.5,0.75]$. The samples are first passed through the frozen teacher model, followed by the discriminator, to yield a real or fake prediction. When employing a UNet architecture [58] for the teacher model, our approach focuses on utilizing only the encoder portion of the UNet, generating an even more compressed latent representation and further reducing the parameter count for the discriminator. We carefully select specific timesteps to enable the discriminator to effectively differentiate between samples based on both high and low-frequency details, as discussed by [33]. Note that in our proposed method, the discriminator is the only component we train, while the teacher model remains frozen. The adversarial loss $\mathcal{L}_{\mathrm{adv}}$ and discriminator loss $\mathcal{L}_{\mathrm{discriminator}}$ write as follows:

where $\nu$ denotes the discriminator parameters. We opt for these particular losses due to their reliability and stability during training, as observed in our experiments. Additionally, we stress the impact of the chosen adversarial loss $\mathcal{L}_{\mathrm{adv}}$ through an ablation study detailed in Section 5.2. In practical terms, the discriminator’s architecture is designed as a straightforward Convolutional Neural Network (CNN) featuring a stride of 2, a kernel size of 4, SiLU activation [13, 53] and group normalization [75].

Inspired by the work of [77], we also propose to introduce a Distribution Matching Distillation (DMD) loss to ensure that the generated samples closely mirror the data distribution learned by the teacher [12, 9, 31]. Specifically, this involves minimizing the Kullback–Leibler (KL) divergence between the distribution of samples from the student model $p^{\mathrm{student}}_{\theta}$ and $p^{\mathrm{teacher}}_{\phi}$, the data distribution learned by the teacher model [73]:

Taking the gradient of the KL divergence with respect to the student model parameters $\theta$ leads to the following update rule:

where $s^{\mathrm{teacher}}$ and $s^{\mathrm{student}}$ are the score functions of the teacher and student distributions respectively. Inspired by [77], the one-step student prediction is re-noised using a uniformly sampled timestep $t^{\prime\prime}\sim\mathcal{U}([0,1])$ and the teacher noise schedule. The new noisy sample is passed through the frozen teacher model to get the score function for the teacher distribution $s^{\mathrm{teacher}}(f_{\theta}^{\mathrm{student}}(z_{t^{\prime\prime}},t^{\prime\prime}))=-(\varepsilon_{\phi}^{\mathrm{teacher}}(x_{t^{\prime\prime}},t^{\prime\prime})/\sigma(t^{\prime\prime}))$. In our approach, we utilize the student model for the score function of the student distribution, instead of a dedicated diffusion model as referenced in [77]. This choice significantly reduces the number of trainable parameters and computational costs. Moreover, the use of the distribution matching loss, examined in the ablation study section (Sec. 5.2), has proven to enhance the quality of generated samples and their prompt adherence.

While striving for robustness and versatility, we also aimed to design a model with a minimal number of trainable parameters, since it involves the loading of computationally intensive functions (teacher and student). To do so, we propose to rely on the parameter-efficient method LoRA [20] and apply it to our student model. This way, we drastically reduce the number of parameters and speed up the training process.

In a nutshell, our student model is trained to minimize a weighted combination of the distillation Eq. (6), the adversarial Eq. (8), and the distribution matching Eq. (9) losses:

The training process is detailed in Alg. 1 and illustrated in Figure 4. In more detail, we first pick a random sample $x_{0}\sim p(x_{0})$ belonging to the unknown data distribution. This sample is then encoded with an encoder $\mathcal{E}$ to get the corresponding latent sample $z_{0}$. A timestep $t$ is drawn according to the timesteps probability mass function $\pi$ detailed in Sec. 4.2 to create a noisy sample $z_{t}$ using Eq. (1). The teacher model $\varepsilon_{\phi}^{\mathrm{teacher}}$ and the ODE solver $\Psi$ are then used to solve the PF-ODE and so generate a synthetic sample $\tilde{z}_{0}^{\mathrm{teacher}}$ belonging to the distribution learned by the teacher model. At the same time, the student model $f_{\theta}^{\mathrm{student}}$ is used to generate a denoised sample $\tilde{z}_{0}^{\mathrm{student}}=f_{\theta}^{\mathrm{student}}(z_{t},t)$ in a single step. The distillation loss is then computed according to Eq. (6). Then, we re-noise the one-step student prediction $\tilde{z}_{0}^{\mathrm{student}}$ as well as the input latent sample $z_{0}$ and compute the adversarial loss as explained in Sec. 4.3. Finally, for distribution matching, we take again the one-step student prediction $\tilde{z}_{0}^{\mathrm{student}}$ and re-noise it using a uniformly sampled timestep $t\sim\mathcal{U}([0,1])$. The new noisy sample is passed through the teacher model to get the teacher score $s^{\mathrm{teacher}}$ function while we use the student model (and not a dedicated diffusion model as in [77]) to get the student score function $s^{\mathrm{student}}$. The distribution matching loss is then computed as explained in Sec. 4.

Overall, our proposed method relies on the training of only a few number of parameters. This is achieved through applying LoRA to the student model, utilizing a frozen teacher model for the adversarial approach, and employing the student denoiser directly rather than introducing a new diffusion model to calculate the fake scores for the distribution matching loss. This approach not only drastically cuts down on the number of parameters but also accelerates the training process.