Distribution Backtracking Builds A Faster Convergence Trajectory for Diffusion Distillation

## Abstract

Abstract Accelerating the sampling speed of diffusion models remains a significant challenge. Recent score distillation methods distill a heavy teacher model into a student generator to achieve one-step generation, which is optimized by calculating the difference between the two score functions on the samples generated by the student model.
However, there is a score mismatch issue in the early stage of the distillation process, because existing methods mainly focus on using the endpoint of pre-trained diffusion models as teacher models, overlooking the importance of the convergence trajectory between the student generator and the teacher model.
To address this issue, we extend the score distillation process by introducing the entire convergence trajectory of teacher models and propose Dis tribution Back tracking Distillation ( DisBack ). DisBask is composed of two stages: Degradation Recording and Distribution Backtracking . Degradation Recording is designed to obtain the convergence trajectory of the teacher model, which records the degradation path from the trained teacher model to the untrained initial student generator.
The degradation path implicitly represents the teacher model’s intermediate distributions, and its reverse can be viewed as the convergence trajectory from the student generator to the teacher model.
Then Distribution Backtracking trains a student generator to backtrack the intermediate distributions along the path to approximate the convergence trajectory of teacher models.
Extensive experiments show that DisBack achieves faster and better convergence than the existing distillation method and accomplishes comparable generation performance, with an FID score of 1.38 on the ImageNet 64x64 dataset.
Notably, DisBack is easy to implement and can be generalized to existing distillation methods to boost performance. Our code is publicly available on https://github.com/SYZhang0805/DisBack .

## 1 Introduction

Recently, a plethora of generative models have demonstrated remarkable performance across diverse domains such as images (Kou et al., ; Yin et al., ), audio (Evans et al., ; Xing et al., ), and videos (Wang et al., ; Chen et al., ). However, existing models still grapple with the “trilemma” problem, wherein they struggle to simultaneously achieve high generation quality, fast generation speed, and high sample diversity (Xiao et al., ). Generative Adversarial Networks (GANs) (Goodfellow et al., ) can rapidly produce high-quality samples but often face mode collapse issues. Variational Autoencoders (VAEs) (Kingma & Welling, ) offer stable training but tend to yield lower-quality samples. Recently, Diffusion models (DMs) have emerged as a competitive contender in the generative model landscape (Fan et al., ; Zhou et al., ; Xu et al., ). DMs can generate high-quality, diverse samples but still suffer from slow sampling speeds due to iterative network evaluations.

Figure: Figure 2: Several examples of 1024$\times$1024 images generated by our proposed one-step DisBack model distilled from SDXL (Podell et al., ).
Refer to caption: x4.png

To accelerate the sampling speed, the score distillation method tries to distill a harvy teacher model to a smaller student generator to reduce the sampling cost and achieve the one-step generation  (Bao et al., ; Luo et al., ; Yin et al., ), it optimizes the student model by calculating the difference between two score functions on the samples generated by the student. However, initially as the generated distribution is far from the training distribution, the generated sample lies outside the training data distribution. Thus, the predicted score of the generated sample from the teacher model does not match the sample’s real score in the training distribution.
This mismatch issue is reflected by unreliable network predictions of the teacher model, which prevents the student model from receiving accurate guidance and leads to a decline in final generative performance. We identified that this issue arises because existing score distillation models mainly focus on using the endpoint of the pre-trained diffusion model as the teacher model, overlooking the importance of the convergence trajectory between the student generator and the teacher model. Without the constraint of a convergence trajectory, the mismatch issue causes the student generator to deviate from a reasonable optimization path during training, leading to convergence to suboptimal solutions and a decline in final performance.

To address this problem, we extend the score distillation process by introducing the entire convergence trajectory of the teacher model and propose Distribution Backtracking Distillation (DisBack) for a faster and more efficient distillation. The construction of DisBack is based on the following insights. In practice, the convergence trajectory of most teacher models is inaccessible, particularly for large models like Stable Diffusion (Rombach et al., ). Because the trajectory of distribution changes is bidirectional, it is possible to construct a degradation path from the teacher model to the initial student generator, and the reverse of this path is viewed as the convergence trajectory of the teacher model. Compared with fitting the teacher model directly, fitting intermediate targets along the convergence trajectory can mitigate the mismatch issue. Thus, the DisBack incorporates degradation recording and distribution backtracking stages. In the degradation recording stage, the teacher model is tuned to fit the distribution of the initial student generator and obtains a distribution degradation path. The path includes a series of in-between diffusion models to represent the intermediate distributions of the teacher model implicitly. In the distribution backtracking stage, the degradation path is reversed and viewed as the convergence trajectory. Then the student generator is trained to backtrack the intermediate distributions along the trajectory to optimize toward the convergence trajectory of the teacher model. In practice, the degradation recording stage typically requires only a few hundred iterations. Therefore, the proposed method incurs trivial additional computational costs. Compared to the existing score distillation method, DisBack exhibits a significantly increased convergence speed (Fig. 1), and it also delivers superior generation performance (Fig. 2).

Our main contributions are summarized as follows. (1) We extend the score distillation process by introducing the entire convergence trajectory of the teacher models and propose Distribution Backtracking Distillation (DisBack), which achieves a faster and more efficient distillation (Sec. 4).
(2) Extensive experiments demonstrate that DisBack accelerates the convergence speed of the score distillation process while achieving comparable or better generation quality compared to existing methods (Sec. 5).
(3) The contribution of DisBack is orthogonal to those of other distillation methods. Researchers are encouraged to incorporate our DisBack training strategy into their distillation methods.

## 2 Related Works

#### Efficient diffusion models.

To improve the efficiency of the diffusion model, existing methods use the knowledge distillation method to distill a large teacher diffusion model to a small and efficient student diffusion model (Yang et al., ).
The progressive distillation model (Salimans & Ho, ) progressively distills the entire sampling process into a new diffusion model with half the number of steps iteratively.
Building on this, the classifier-guided distillation model (Sun et al., ) introduces a dataset-independent classifier to focus the student model on the crucial features to enhance the distillation process.
Guided-distillation model (Meng et al., ) proposes a classifier-free guiding framework to avoid the computational cost of additional classifiers and achieve high-quality sampling in only 2-4 steps.
Recently, Consistency Model (Song et al., ) uses the self-consistency of the ODE generation process to achieve one-step distillation, but this is at the expense of generation quality. To mitigate the surface of the sample quality caused by the acceleration, Consistency Trajectory Models (Kim et al., ) combines the adversarial training and denoising score matching loss to further improve the performance.
Latent Adversarial Diffusion Distillation (Sauer et al., ) leverages generative features from pre-trained latent diffusion models to achieve high-resolution, multi-aspect ratio, few-step image generation.

#### Score distillation for one-step generation.

Diff-Instruct (Luo et al., ) proposes a distillation method from the pre-trained diffusion model to the one-step generator that involves optimizing the generator by the gradient of the difference between two score functions. One score function represents the pre-trained diffusion distribution, while the other represents the generated distribution. Adversarial Score Distillation (Wei et al., ) further employs the paradigm of WGAN and retains an optimizable discriminator to improve performance. Additionally, Swiftbrush (Nguyen & Tran, ) leverages score distillation to distill a Stable Diffusion v2.1 into a one-step text-to-image generation model and achieve competitive results. DMD (Yin et al., ) suggests the inclusion of a regression loss between noisy images and corresponding outputs to alleviate instability in the distillation process in text-to-image generation tasks. DMD2 (Yin et al., ) introduces a two-time-scale update rule and an additional GAN loss to address the issue of generation quality being limited by the teacher model in DMD, achieving superior performance. Recently, HyperSD (Ren et al., ) integrate score distillation with trajectory segmented consistency distillation and the human feedback learning, which achieves SOTA performance from 1 to 8 inference steps.

## 3 Preliminary

In this part, we briefly introduce the score distillation approach. Let $q^{G}_{0}$ and $q^{G}_{t}$ be the distribution of the student generator $G_{\mathit{stu}}$ and its noisy distribution at timestep $t$. In addition, $q_{0}$ and $q_{t}$ are the training distribution and its noisy distribution at timestep $t$. By optimizing the KL divergence in Eq. (1), we can train a student generator to enable one-step generation (Wang et al., ).

$$ $\min_{\eta}D_{KL}\left(q_{0}^{G}\left(\bm{x}_{0}\right)\|q_{0}\left(\bm{x}_{0} \right)\right)$ (1) $$

Here $\bm{x}_{0}=G_{\mathit{stu}}(\bm{z};\eta)$ is the generated samples, and $\eta$ is the trainable parameter of $G_{\mathit{stu}}$.
However, due to the complexity of $q_{0}$ and its sparsity in high-density regions, directly solving Eq.(1) is challenging (Song & Ermon, ). Inspired by Variational Score Distillation (VSD) (Wang et al., ), Eq.(1) can be extended to optimization problems at different timesteps $t$ in Eq. (2). As $t$ increases, the diffusion distribution becomes closer to a Gaussian distribution.

$$ $\displaystyle\min_{\eta}\mathbb{E}_{t\sim\mathcal{U}(0,1),\epsilon\sim\mathcal {N}(0,I)}D_{KL}\left(q_{t}^{G}\left(\bm{x}_{t}\right)\|q_{t}\left(\bm{x}_{t} \right)\right)$ (2) $$

Here $\bm{x}_{t}$ is the noisy data and $p(\bm{x}_{t}\mid\bm{x}_{0})\sim\mathcal{N}(\bm{x}_{0},\sigma_{t}^{2}I)$. Theorem 1 proves that introducing the additional KL-divergence for $t>0$ does not affect the global optimum of the original optimization problem in Eq.(1).

###### Theorem 1 (The global optimum of training (Wang et al., 2023 ) )

Given $t>0$, we have,

$$ $D_{\mathrm{KL}}\left(q_{t}^{G}\left(\bm{x}_{t}\right)\|q_{t}\left(\bm{x}_{t} \right)\right)=0\Leftrightarrow D_{\mathrm{KL}}\left(q_{0}^{G}\left(\bm{x}_{0} \right)\|q_{0}\left(\bm{x}_{0}\right)\right)=0$ (3) $$

Therefore, by minimizing the KL divergence in Eq. (2), the student generator can be optimized through the following gradients:

$$ $\nabla_{\eta}D_{\mathrm{KL}}\left(q^{G}_{t}\left(\bm{x}_{t}\right)\|q_{t}\left (\bm{x}_{t}\right)\right)=\mathbb{E}_{t,\epsilon}\left[\left[\nabla_{\bm{x}_{t }}\log q^{G}_{t}\left(\bm{x}_{t}\right)-\nabla_{\bm{x}_{t}}\log q_{t}\left(\bm {x}_{t}\right)\right]\frac{\partial\bm{x}_{t}}{\partial\eta}\right]$ (4) $$

Here the score of perturbed training data $\nabla_{\bm{x}_{t}}\log q_{t}\left(\bm{x}_{t}\right)$ can be approximated by a pre-trained model $s_{\theta}$. The score of perturbed generated data $\nabla_{\bm{x}_{t}}\log q^{G}_{t}\left(\bm{x}_{t}\right)$ is estimated by another diffusion model $s_{\phi}$, which is optimized by the traditional diffusion training with generated data (Song et al., ):

$$ $\displaystyle\min_{\phi}\mathbb{E}_{t,\epsilon}\left\|s_{\phi}\left(\bm{x}_{t} ,t\right)-\frac{\bm{x}_{0}-\bm{x}_{t}}{\sigma_{t}^{2}}\right\|_{2}^{2}$ (5) $$

Thus, the gradient of student generator in Eq.(4) is estimated as

$$ $\nabla_{\eta}D_{\mathrm{KL}}\left(q^{G}_{t}\left(\bm{x}_{t}\right)\|q_{t}\left (\bm{x}_{t}\right)\right)\approx\mathbb{E}_{t,\epsilon}\left[s_{\phi}\left(\bm {x}_{t},t\right)-s_{\theta}\left(\bm{x}_{t},t\right)\right]\frac{\partial\bm{x }_{t}}{\partial\eta}$ (6) $$

The distribution of the student generator changes after its update. Therefore, $s_{\phi}$ also needs to be optimized based on the newly generated images to ensure the timely approximation of the generated distribution. Thus, the student generator and $s_{\phi}$ are optimized alternately.

In practice, $s_{\phi}$ has three initialization strategies: (1) $s_{\phi}$ is randomly initialized (Franceschi et al., ). (2) $s_{\phi}$ is initialized as $s_{\theta}$ or its LoRA (Hu et al., ; Wei et al., ). (3) $s_{\phi}$ is initialized by fitting the generated samples of student generator (Luo et al., ). Beyond unconditional image generation (Ye & Liu, ), this method has also been applied to tasks such as text-to-image and image-to-image generation across various structures (Yin et al., ; Hertz et al., ).

Figure: Figure 3: The overall framework of DisBack. Stage 1: An auxiliary diffusion model is initialized with the teacher model $s_{\theta}$ and then fits the distribution of the initial student generator $G_{\mathit{stu}}^{0}$. The intermediate checkpoints $\{s_{\theta_{i}}^{\prime}\mid i=0,\ldots,N\}$ are saved to form a degradation path. The degradation path is then reversed and viewed as the convergence trajectory. Stage 2: The intermediate node $s_{\theta_{i}}$ in the convergence trajectory is distilled to the generator sequentially until the generator converges to the distribution of the teacher model.
Refer to caption: x5.png

###### Theorem 1 (The global optimum of training (Wang et al., 2023 ) )

## 4 Method

### 4.1 Insight

In this section, we introduce the Distribution Backtracking Distillation (DisBack).
The key insight behind DisBack is the importance of the convergence trajectory. As mentioned in Sec.3, there are two score functions in score distillation, one representing the pretrained diffusion distribution and the other representing the generated distribution. The student model is optimized using the gradient of the difference between these two score functions. Existing methods (Luo et al., ; Yin et al., ; ) only consider the endpoint of the pre-trained diffusion model as the teacher model, overlooking the intermediate convergence trajectory between the student generator and the teacher model. The resulting score mismatch issue between the predicted scores of the generated sample from the teacher model and the real scores causes the student model to receive inaccurate guidance. It ultimately leads to a decline in final performance. Constraining the convergence trajectory between the student generator and the teacher model during the distillation process can mitigate the mismatch issue and help the student generator approximate the convergence trajectory of teacher models to achieve faster convergence. In practice, it is infeasible to obtain the convergence trajectory of most teacher models, especially for large models such as Stable Diffusion. Reversely, it is possible to obtain the degradation path from the teacher model to the initial student generator. The reverse of this degradation path can be viewed as the convergence trajectory of the teacher model. Based on the above insights, we structure the proposed DisBack method in two stages including the degradation recording stage and the distribution backtracking stage (Fig. 3).

### 4.2 Degradation Recording

This stage aims to obtain the degradation path from the teacher model to the initial student generator. The degradation path is then reversed and viewed as the convergence trajectory of the teacher model. The teacher model here is the pre-trained diffusion model $s_{\theta}$ and the student generator is represented by $G^{0}_{stu}$.

Let $s_{\theta}^{\prime}$ be a diffusion model initialized with the teacher model $s_{\theta}$, and it is trained on generated samples to fit the initial student generator’s distribution $q^{G}_{0}$ with Eq. (7). By saving the multiple intermediate checkpoints during the training, we can obtain a series of diffusion models $\{s_{\theta_{i}}^{\prime}\mid i=0,\ldots,N\}$, where $s_{\theta_{0}}^{\prime}=s_{\theta}\approx q_{0}$ and $s_{\theta_{N}}^{\prime}\approx q^{G}_{0}$. These diffusion models describe the scores of non-existent distributions on the path recording how the training distribution $q_{0}$ degrades to $q^{G}_{0}$. Algorithm 1 shows the process of obtaining the degradation path. Since distribution degradation is easily achievable, the degradation recording stage only needs trivial additional computational resources (200 iterations in most cases).

$$ $\displaystyle\min_{\theta}\mathbb{E}_{t,\epsilon}\left\|s_{\theta}^{\prime} \left(\bm{x}_{t},t\right)-\frac{\bm{x}_{0}-\bm{x}_{t}}{\sigma_{t}^{2}}\right\| _{2}^{2}$ (7) $$

Figure: Algorithm 1 Degradation Recording.

Figure: Algorithm 2 Distribution Backtracking.

### 4.3 Distribution Backtracking

Given the degradation path from the teacher model to the student generator, the reverse path is viewed as a representation of the convergence trajectory between the initial student generator $G^{0}_{\mathit{stu}}$ and the teacher model $s_{\theta}$. The key to the distribution backtracking is to sequentially distill checkpoints in the convergence trajectory into the student generator. The last node $s_{\theta_{N}}^{\prime}$ in the path is close to the initially generated distribution $q^{G}_{0}$. Therefore, during the distribution backtracking stage, we use $s_{\theta_{N-1}}^{\prime}$ as the first target to distill the student generator. When near convergence, we switch the target to $s_{\theta_{N-2}}^{\prime}$. The checkpoints $s_{\theta_{i}}^{\prime}$ is sequentially distilled to $G_{\mathit{stu}}$ until the final target $s_{\theta_{0}}^{\prime}$ is reached. During the distillation, the gradient of $G_{\mathit{stu}}$ is:

$$ $\text{Grad}(\eta)=\mathbb{E}_{t,\epsilon}\left[s_{{\phi}}\left(\bm{x}_{t},t \right)-s_{\theta_{i}}^{\prime}\left(\bm{x}_{t},t\right)\right]\frac{\partial \bm{x}_{t}}{\partial\eta}$ (8) $$

In this stage, $G_{\mathit{stu}}$ and $s_{\phi}$ are also optimized alternately and the optimization of $s_{\phi}$ is the same as in the original score distillation (Eq. 5). Compared to existing methods, the final target of DisBack is the same while constraining the convergence trajectory to achieve more efficient distillation of the student generator. Algorithm 2 summarizes the distribution backtracking stage.

**Table 1: The unconditional generation performance of DisBack. The FID ($\downarrow$) scores are shown. The best results are shown in bold.**
| Model | NFE ($\downarrow$) | FFHQ | AFHQv2 | LSUN-bedroom | LSUN-cat |
| --- | --- | --- | --- | --- | --- |
| DDPM (Ho et al., ) | 1000 | 3.52 | - | 4.89 | 17.10 |
| ADM (Dhariwal & Nichol, ) | 1000 | - | - | 1.90 | 5.57 |
| NSCN++ (Song et al., ) | 79 | 25.95 | 18.52 | - | - |
| DDPM++ (Song et al., ) | 79 | 3.39 | 2.58 | - | - |
| EDM (Karras et al., ) | 79 | 2.39 | 1.96 | 3.57 | 6.69 |
| EDM (Karras et al., ) | 15 | 15.81 | 13.67 | - | - |
| Diff-Instruct (Luo et al., ) | 1 | 19.93 | - | - | - |
| PD (Salimans & Ho, ) | 1 | - | - | 16.92 | 29.60 |
| CT (Song et al., ) | 1 | - | - | 16.00 | 20.70 |
| CD (Song et al., ) | 1 | 12.58 | 10.75 | 7.80 | 11.00 |
| DisBack | 1 | 10.88 | 9.97 | 6.99 | 10.30 |

## 5 Experiment

Experiments are conducted on different models across various datasets. We first compare the performance of DisBack with other multi-step diffusion models and distillation methods (Sec. 5.1). Secondly, we compare the convergence speed of DisBack with its variants without backtracking along the convergence trajectory (Sec. 5.2). Thirdly, further experiments are conducted to demonstrate DisBack’s effectiveness in mitigating the score mismatch issues (Sec. 5.3).
Then, we also conduct the ablation study to show the effectiveness of DisBack (Sec. 5.4). Finally, we show the results of DisBack on text-to-image generation tasks (Sec. 5.5).

**Table 2: The conditional generation performance of DisBack on ImageNet 64x64 dataset. The best results are shown in bold.**
| Model | NFE ($\downarrow$) | FID ($\downarrow$) |
| --- | --- | --- |
| DDPM (Ho et al., ) | 1000 | 3.77 |
| DDDM (Zhang et al., ) | 1000 | 2.11 |
| EDM (Karras et al., ) | 79 | 2.44 |
| EDM (Karras et al., ) | 15 | 10.46 |
| Moment Matching (Salimans et al., ) | 8 | 3.3 |
| SlimFlow (Zhu et al., ) | 1 | 12.34 |
| BOOT (Gu et al., ) | 1 | 12.30 |
| DDDM (Zhang et al., ) | 1 | 3.47 |
| CTM (Kim et al., ) | 1 | 2.06 |
| Sid (Zhou et al., ) | 1 | 1.52 |
| DMD2 (Yin et al., ) | 1 | 1.51 |
| Diff-Instruct (Luo et al., ) | 1 | 5.57 |
| PD (Salimans & Ho, ) | 1 | 8.95 |
| CT (Song et al., ) | 1 | 13.00 |
| CD (Song et al., ) | 1 | 6.20 |
| DisBack | 1 | 1.38 |

### 5.1 Quantitative evaluation

DisBack can achieve performance comparable to or even better than the existing diffusion models or distillation methods. Experiments are conducted on different datasets. (1) The unconditional generation on FFHQ 64x64, AFHQv2 64x64, LSUN-bedroom 256x256 and LSUN cat 256x256. (2) The conditional generation on ImageNet 64x64. The performance of DisBack is shown in Tab. 1 and Tab. 2. All the DisBack models are distilled from the pre-trained EDM model (Karras et al., ).

For unconditional generation, the one-step generator distilled by the DisBack achieves comparable performance across different datasets compared to multi-step generation diffusion models. Specifically, it also outperforms the original EDM model with 15 NFEs (10.88 of DisBack and 15.81 of EDM on FFHQ64). Compared to existing one-step generators and distillation models, DisBack achieves optimal performance. DisBack requires no training data and additional constraints during training. For conditional generation, the DisBack achieves the best performance compared to the existing models. In conclusion, DisBack can achieve competitive distillation performance compared to existing models.

**Table 3: Ablation study on constraining the convergence trajectory to the score distillation process. The FID ($\downarrow$) scores in each case are shown.**
| Model | FFHQ | AHFQv2 | ImageNet | LSUN-bedroom | LSUN-cat |
| --- | --- | --- | --- | --- | --- |
| DisBack | 10.88 | 9.97 | 1.38 | 6.99 | 10.30 |
| w/o Convergence Trajectory | 12.26 | 10.29 | 5.96 | 7.43 | 10.63 |

### 5.2 Convergence speed

We conducted a series of experiments to demonstrate the advantages of DisBack in accelerating the convergence speed of the score distillation process on unconditional FFHQ 64x64 (Karras et al., ), CIFAR10 (Krizhevsky, ), and conditional ImageNet 64x64 (Deng et al., ) datasets. Diff-Instruct Luo et al. () is the existing SOTA score distillation method, which can be regarded as a variation of DisBack not introducing the convergence trajectory. We compared the FID trends of DisBack and Diff-Instruct during distillation in the same situation.

The results are shown in Fig. 1. DisBack achieves a convergence speed 2.46 times faster than the variant without the constraint of convergence trajectory on the FFHQ 64x64 dataset and 13.09 times faster on the CIFAR dataset. For the conditional generation on the ImageNet 64x64 dataset, DisBack is 2.57 times faster. The fast convergence speed is because constraining the convergence trajectory of the generator provides a clear optimization direction, avoiding the generator falling into suboptimal solutions and enabling faster convergence to the target distribution.

### 5.3 Experiments on score mismatch issue

In this part, experiments are conducted to validate the positive impact of constraining the convergence trajectory and the effectiveness of mitigating the mismatch issues.
We propose a new metric called mismatch degree to assess whether the predicted score of the teacher model matches the distribution’s real score given a data distribution. This score is inspired by the score-matching loss.

$$ $d_{\mathit{mis}}=\mathbb{E}_{\bm{x}_{t}}\|s_{\theta}(\bm{x}_{t},t)-\nabla_{\bm {x}_{t}}\log p_{t}(\bm{x}_{t})\|_{2}$ (9) $$

Figure: Figure 4: The mismatch degree during the distillation process of Diff-Instruct and proposed DisBack. The standard deviation is visualized. DisBack effectively mitigates the mismatch degree during the distillation process.
Refer to caption: x6.png

Here $\bm{x}_{t}$ is the noisy data from the assessed distribution. Besides, $s_{\theta}(\bm{x}_{t},t)$ is the predicted score and $\nabla_{\bm{x}_{t}}\log p_{t}(\bm{x}_{t})$ is the real score. We use Stable Target Field (STF) (Xu et al., ) to approximate the real score $\nabla_{\bm{x}_{t}}\log p_{t}(\bm{x}_{t})$ of $x_{t}$ to the assessed distribution in experiments.

When the assessed distribution is close to the distribution represented by the teacher model $s_{\theta}$, the mismatch degree is small, and vice versa.
When calculated directly on the training data, the resulting mismatch degree represents the ideal lower bound. Thus, the mismatch degree can be used to assess the convergence degree of the generated distribution and visualize the speed of convergence under the constraining of the convergence trajectory.

We conduct experiments on FFHQ dataset with Diff-Instruct (Luo et al., ) as a baseline. We calculate the mismatch degree in the distribution of the student generator of both Diff-Instruct and DisBack. The pre-trained EDM model is chosen as the teacher model. In this scenario, the ideal lower bound is $0.037$. We visualized the mismatch degree in Fig. 4.
With the constraining of the convergence trajectory from Disback, the mismatch degree is lower during the distillation process, meaning the student generator converges faster and better and the DisBack achieves more effective distillation.

### 5.4 Ablation study

Ablation studies are conducted to compare the performance of DisBack with its variant without the constraint of the convergence trajectory. The results are shown in Tab. 3. Results show that the variant without the constraint of convergence trajectory suffers from a performance decay in different cases. This confirms the efficacy of constraining the convergence trajectory between the teacher model and the student generator can improve the final performance of the generation.

Figure: Figure 5: Generation samples by DisBack distilled from SDXL with 1024$\times$1024 resolution.
Refer to caption: x7.png

**Table 4: The results of text-to-image generation.**
| Model | FID ($\downarrow$) | NFE ($\downarrow$) | User Preference |
| --- | --- | --- | --- |
| SDXL | 19.36 | 100 | 38.7% |
| DisBack | 18.96 | 1 | 61.3% |

### 5.5 Text to Image Generation

Further experiments are conducted on text-to-image generation tasks. We use DisBack to distill the SDXL model (Podell et al., ) and evaluate the FID scores of the distilled SDXL and the original SDXL on the COCO 2014 (Radford et al., ). The user studies are conducted to further verify the effectiveness of DisBack. We randomly select 128 prompts from the LAION-Aesthetics (Schuhmann et al., ) to generate images and ask volunteer participants to choose the images they think are better. Detailed information about the user study is included in the Appendix B.3. The results of the FID evaluation and user study are presented in Tab 4. DisBack achieved better results in single-step generation compared to the original SDXL with the 100-step DDIM sampler (Song et al., ). The preference scores of DisBack over the original SDXL are $61.3\%$. Some generation samples are shown in Fig. 2 and Fig. 5.

We also conducted experiments on LCM-LoRA (Luo et al., ) on the COCO 2014 (Radford et al., ). The LCM-LoRA distilled using DisBack has the FID score of 36.37, while the FID of original LCM-LoRA is 78.26. Detailed experimental results are provided in the Appendix.

## 6 Conclusion

#### Summary.

This paper proposes Distribution Backtracking Distillation (DisBack) to introduce the entire convergence trajectory of the teacher model in the score distillation. The DisBack can also be used in the distillation of large-scale text-to-image models. DisBack performs a faster and more efficient distillation and achieves a comparable or better performance in one-step generation compared to existing multi-step generation diffusion models and one-step diffusion distillation models.

#### Limitation.

The performance of DisBack is inherently limited by the teacher model. The better the original performance of the teacher model has, the better the performance of DisBack also presents. Additionally, to achieve optimal performances in both accelerated distillation and generation quality, DisBack requires careful design of the distribution degradation path and the setting of various hyperparameters (such as how many epochs used to fit each intermediate node). While with no meticulous design it can also achieve better performance than models without backtracking, further exploration is required to enable the model to reach optimal performance.

## Appendix A More Experiment Results on Text-to-Image Generation

Figure: Figure 6: The FID scores of LCM-LoRA distilled from SD1.5 across training steps. DisBack achieves faster convergence and better performance.
Refer to caption: extracted/5877782/figures/FID_LCM_LoRA.png

We conducted the experiment on LCM-LoRA (Luo et al., ). LCM-LoRA is a Low-Rank Adaptation (LoRA) fine-tuned version of the Latent Consistency Model (LCM) Luo et al. (), applicable across fine-tuned Stable Diffusion models for high-quality, single-step or few-step generation. In this experiment, we use LCM-LoRA as the student generator and Stable Diffusion v1.5 as the teacher model. We observed that the score distillation underperforms when LCM-LoRA serves as the teacher model. This issue likely stems from the infeasibility of directly converting the outputs of LCM-LoRA into scores.

We distill the LCM-LoRA on the LAION-Aesthetic-6.5+ dataset (Schuhmann et al., ) and evaluate the FID scores on the COCO 2014 dataset (Radford et al., ) with the resolution of 512×512. 50,000 real images and 30,000 generated images were used to calculate FID scores. The 30,000 generated images were obtained by generating one image for each of the 30,000 distinct prompts. The results are shown in Tab. 5. The change in FID scores over training steps is illustrated in Fig. 6, showing that our DisBack method achieves a 1.5 times acceleration in convergence speed and yields superior generation performance within the same training period. Some generated samples of DisBack and original LCM-LoRA are shown in Fig. 7.

**Table 5: The distillation results of LCM-LoRA with NFE $=1$.**
| Model | FID ($\downarrow$) |
| --- | --- |
| LCM-LoRA | 78.26 |
| DisBack | 36.37 |

Figure: Figure 7: Text-to-images generation samples by original LCM-LoRA and DisBack model distilled from LCM-LoRA with 512$\times$512 resolution.
Refer to caption: x8.png

## Appendix B Implementation Details

### B.1 Dataset Setup

We experiment on the following datasets:

The FFHQ (Flickr-Faces-HQ) dataset (Karras et al., ) is a high-resolution dataset of human face images used for face generation tasks. It includes high-definition face images of various ages, genders, skin tones, and expressions from the Flickr platform. This dataset is commonly employed to train large-scale generative models. In this paper, we utilize a derivative dataset of the FFHQ called FFHQ64, which involves downsampling the images from the original FFHQ dataset to a resolution of 64$\times$64.

The AFHQv2 (Animal Faces-HQ) dataset (Choi et al., ) comprises 15,000 high-definition animal face images with a resolution of 512$\times$512, including 5,000 images each for cats, dogs, and wild animals. AFHQv2 is commonly employed in tasks such as image-to-image translation and image generation. Similar to the FFHQ dataset, we downscale the original AFHQv2 dataset to a resolution of 64$\times$64 for the experiment.

The CelebA (Celebrities Attributes) dataset (Liu et al., ) is a large-scale collection of facial images used for facial attribute recognition. With over 200,000 images of celebrities, each image is annotated with 40 attributes, including features like hairstyle, gender, and age. In addition to attribute annotations, each image also includes annotations for facial landmarks such as eyes, nose, and mouth. For our experiments, we utilize the 64$\times$64 version of the CelebA dataset.

The ImageNet dataset (Deng et al., ) was established as a large-scale image dataset to facilitate the development of computer vision technologies. This dataset comprises over 14,197,122 images spanning more than 20,000 categories, indexed by 21,841 Synsets. In this paper, we use the ImageNet64 dataset, a subsampled version of the ImageNet dataset. The Imagenet64 dataset consists of a vast collection of images with a resolution of 64$\times$64, containing 1,281,167 training samples, 50,000 testing samples, and 1,000 labels.

The LSUN (Large Scale Scene Understanding) dataset (Yu et al., ) is a large-scale dataset for scene understanding in visual tasks within deep learning. Encompassing numerous indoor scene images, it spans various scenes and perspectives. The original LSUN dataset comprises multiple sub-datasets; in this study, we experiment on the LSUN Church and Bedroom sub-datasets with a resolution of 256$\times$256.

### B.2 Experiment Setup

For experiments on FFHQ 64x64, AFHQv2 64x64, and ImageNet 64x63 datasets, the pre-trained models are provided by the official release of EDM Karras et al. () in [https://github.com/NVlabs/edm](https://github.com/NVlabs/edm). We use Adam optimizers to train the student generator $G$ and $s_{\phi}$, with both learning rates set to $1e^{-5}$. The training consisted of 500,000 epochs on four NVIDIA 3090 GPUs, and the batch size per GPU is set to 8. The training ratio between $s_{\phi}$ and $G$ remains at $1:1$. In the Degradation stage, we trained for 200 epochs total, saving a checkpoint every 50 epochs, resulting in a total of 5 intermediate nodes along the degradation path $\{s_{\theta_{i}}^{\prime}|i=0,1,2,3,4\}$. In the Distribution Backtracking stage, when $i\geq 3$, each checkpoint was trained for 1,000 steps. When $i<3$, each checkpoint was trained for 10,000 steps. The remaining steps were used to distill the original teacher model $s_{\theta_{0}}^{\prime}$.

For experiments on LSUN bedroom and LSUN cat datasets, the pre-trained models are provided by the official release of Consistency Model Song et al. () in [https://github.com/openai/consistency_models](https://github.com/openai/consistency_models). During the training, we set $\sigma_{max}$ to 80 and keep it constant during the single-step generation process. We use SGD and AdamW optimizers during training to train the generator $G$ and $s_{\phi}$, with learning rates set to $1e^{-3}$ and $1e^{-4}$, respectively. The training consisted of 10,000 epochs on one NVIDIA A100 GPU, and the batch size per GPU is set to 2. The training ratio between $s_{\phi}$ and $G$ remains at $4:1$. In the Degradation stage, we trained for 200 epochs total and saved the checkpoint every 50 epochs, resulting in a total of 5 intermediate nodes along the degradation path $\{s_{\theta_{i}}^{\prime}|i=0,1,2,3,4\}$. In the Distribution Backtracking stage, when $i\geq 3$, each checkpoint was trained for 500 steps. When $i<3$, each checkpoint was trained for 1000 steps. The remaining steps were used to distill the original teacher model $s_{\theta_{0}}^{\prime}$.

When distilling the SDXL model, the teacher model and the student generator are both initialed by the pre-trained SDXL model on the huggingface (model id is ‘stabilityai/stable-diffusion-xl-base-1.0’). We use Adam optimizers to train $G$ and $s_{\phi}$, with learning rates set to $1e^{-3}$ and $1e^{-2}$, respectively. The training consisted of 50,000 epochs on one NVIDIA A100 GPU, and the batch size per GPU is set to 1. The training ratio between $s_{\phi}$ and $G$ remains at $1:1$. The training prompts are obtained from LAION-Aesthetics. In the Degradation stage, we trained for 1,000 epochs total and saved the checkpoint every 100 epochs, resulting in a total of 10 intermediate nodes along the degradation path $\{s_{\theta_{i}}^{\prime}|i=0,1,...,9\}$. In the Distribution Backtracking stage, each checkpoint was trained for 1,000 steps. The remaining steps were used to distill the original teacher model $s_{\theta_{0}}^{\prime}$.

### B.3 User study setup

Firstly, we randomly selected 128 prompts from the LAION-Aesthetics (Schuhmann et al., ). Then we use the original SDXL model and the distilled SDXL model to generate 128 pairs of images. Subsequently, we randomly recruit 10 volunteers, instructing each to individually evaluate the fidelity, detail, and vividness of these pairwise images. 10 volunteers included 6 males and 4 females, aged between 24 and 29. 5 of them have artificial intelligence or related majors and the other 5 of them have other majors. They were given unlimited time for the experiment, and all of the volunteers completed the assessment with an average time of 30 minutes. Finally, we took the average of the evaluation results of 10 volunteers as the final user study result.

## Appendix C Theoretical Demonstration

### C.1 KL Divergence of DisBack

As the KL divergence follows

$$ $\displaystyle D_{\mathrm{KL}}(q\parallel p)=\mathbb{E}_{q}\left[\log\frac{q}{p }\right]$ (10) $$

The KL divergence of generated distribution and training distribution at timestep $t$ can be written as

$$ $\displaystyle D_{\mathrm{KL}}\left(q_{t}^{G}\left(\bm{x}_{t}\right)\parallel q _{t}\left(\bm{x}_{t}\right)\right)$ (11) $\displaystyle=\mathbb{E}_{x_{t}\sim q_{t}^{G}\left(\bm{x}_{t}\right)}\log\frac {q_{t}^{G}\left(\bm{x}_{t}\right)}{q_{t}\left(\bm{x}_{t}\right)}$ $\displaystyle=\mathbb{E}_{x_{0}\sim G(z;\eta)}\left[\log q_{t}^{G}\left(\bm{x} _{t}\right)-\log q_{t}\left(\bm{x}_{t}\right)\right]$ $\displaystyle=\mathbb{E}_{z}\left[\log q_{t}^{G}\left(\bm{x}_{t}\right)-\log q _{t}\left(\bm{x}_{t}\right)\right]$ $$

Thus, the gradient of KL divergence can be estimated as

$$ $\nabla_{\eta}D_{\mathrm{KL}}\left(q^{G}_{t}\left(\bm{x}_{t}\right)\|q_{t}\left (\bm{x}_{t}\right)\right)=\mathbb{E}_{t,\epsilon}\left[s_{\phi}\left(\bm{x}_{t },t\right)-s_{\theta}\left(\bm{x}_{t},t\right)\right]\frac{\delta\bm{x}_{t}}{ \delta\eta}$ (12) $$

### C.2 Stable Target Field

Given $\bm{x}_{0}\sim q_{0}$ is the training data, $\bm{x}_{t}\sim p(\bm{x}_{t}\mid\bm{x}_{0})$ is the disturbed data, Xu et al. (Xu et al., ) presents an estimation of the score as:

$$ $\nabla_{\bm{x}_{t}}\log p_{t}(\bm{x}_{t})=\frac{\nabla_{\bm{x}_{t}}p_{t}(\bm{x }_{t})}{p_{t}(\bm{x}_{t})}=\frac{\mathbb{E}_{\bm{x}_{0}}\nabla_{\bm{x}_{t}}p( \bm{x}_{t}\mid\bm{x}_{0})}{p_{t}(\bm{x}_{t})}$ (13) $$

The transition kernel $p(\bm{x}_{t}\mid\bm{x}_{0})$ follows the Gaussian distribution $p(\bm{x}_{t}\mid\bm{x}_{0})\sim\mathcal{N}(\mu_{t},\sigma_{t}^{2}I)$. Here $\mu_{t}=\bm{x}_{0}$ in Variance Exploding SDE (Song et al., ) but is defined differently in other diffusion models.

$$ $p(\bm{x}_{t}\mid\bm{x}_{0})=\frac{1}{\sqrt{(2\pi^{k})}\sigma_{t}}\exp(-\frac{( \bm{x}_{t}-\mu_{t})^{T}(\bm{x}_{t}-\mu_{t})}{2\sigma_{t}^{2}})$ (14) $$

$$ $\displaystyle\nabla_{\bm{x}_{t}}p(\bm{x}_{t}\mid\bm{x}_{0})$ (15) $\displaystyle=\nabla_{\bm{x}_{t}}\left[\frac{1}{\sqrt{(2\pi^{k})}\sigma_{t}} \exp(-\frac{(\bm{x}_{t}-\mu_{t})^{T}(\bm{x}_{t}-\mu_{t}))}{2\sigma_{t}^{2}}\right]$ $\displaystyle=p(\bm{x}_{t}\mid\bm{x}_{0})\nabla_{\bm{x}_{t}}(-\frac{(\bm{x}_{t }-\mu_{t})^{T}(\bm{x}_{t}-\mu_{t})}{2\sigma_{t}^{2}})$ $\displaystyle=p(\bm{x}_{t}\mid\bm{x}_{0})\frac{\mu_{t}-\bm{x}_{t}}{\sigma_{t}^ {2}}$ $$

Combine Eq. (13) to Eq. (15), we have

$$ $\nabla_{\bm{x}_{t}}\log p_{t}(\bm{x}_{t})=\mathbb{E}_{\bm{x}_{0}}\frac{p(\bm{x }_{t}\mid\bm{x}_{0})}{p_{t}(\bm{x}_{t})}\frac{\mu_{t}-\bm{x}_{t}}{\sigma_{t}^{ 2}}=\frac{1}{p_{t}(\bm{x}_{t})}\mathbb{E}_{\bm{x}_{0}}p(\bm{x}_{t}\mid\bm{x}_{ 0})\frac{\mu_{t}-\bm{x}_{t}}{\sigma_{t}^{2}}$ (16) $$

Let $B$ be a set of reference samples for Monte Carlo estimation, we have

$$ $p_{t}(\bm{x}_{t})=\mathbb{E}_{\bm{x}_{0}}p(\bm{x}_{t}\mid\bm{x}_{0})\approx \frac{1}{|B|}\sum_{x_{0}^{(i)}\in B}p(\bm{x}_{t}\mid\bm{x}_{0}^{(i)})$ (17) $$

Combine the Eq. (16) and Eq. (17), we can get

$$ $\nabla_{\bm{x}_{t}}\log p_{t}(x_{t})=\mathbb{E}_{\bm{x}_{0}}\frac{p(\bm{x}_{t} \mid\bm{x}_{0})}{p_{t}(\bm{x}_{t})}\frac{\mu_{t}-\bm{x}_{t}}{\sigma_{t}^{2}} \approx\frac{1}{p_{t}(\bm{x}_{t})}\frac{1}{|B|}\sum_{x_{0}^{(i)}\in B}p(\bm{x} _{t}\mid\bm{x}_{0}^{(i)})\frac{\mu_{t}-\bm{x}_{t}}{\sigma_{t}^{2}}$ (18) $$

Here the ”$\approx$” represents the Monte Carlo estimate.

Depending on the network prediction, the diffusion model can be divided into different types, including $\epsilon$ prediction (Karras et al., ) and $x_{0}$ prediction (Song et al., ; Ho et al., ; Nichol & Dhariwal, ). When the score $\nabla_{\bm{x}_{t}}\log p_{t}(x_{t})$ is estimated by Eq.(18), it can be converted to $\epsilon$, $x_{0}$ and $v$ by a series of transformations.

$$ $\hat{\epsilon}\approx-\sigma_{t}\nabla_{\bm{x}_{t}}\log p_{t}(x_{t})$ (19) $$

$$ $\hat{\bm{x}}_{0}\approx\nabla_{\bm{x}_{t}}\log p_{t}(x_{t})*\sigma_{t}^{2}+\bm {x}_{t}$ (20) $$

## Appendix D Discussion

### D.1 Training Efficiency of DisBack

While DisBack involves an iterative optimization process during training, the optimization objective of $s_{\phi}\left(\bm{x}_{t},t\right)$ aims to minimize the loss of the standard diffusion model based on Eq.(5), and the objective of student generator aims to minimize the KL divergence in Eq.(6). These two optimization processes do not entail an adversarial training as in GANs. Consequently, the optimization process tends to be more stable. A recent work Monoflow (Yi et al., ) also discusses in GANs training a vector field is obtained to guide the optimization of the generator, but the vector field derives from the discriminator and the instability is not mitigated.

$$ $\displaystyle\min_{\phi}\mathbb{E}_{t,\epsilon}\left\|s_{\phi}\left(\bm{x}_{t} ,t\right)-\frac{\bm{x}_{0}-\bm{x}_{t}}{\sigma_{t}^{2}}\right\|_{2}^{2}$ (21) $$

For the DisBack, training the student generator only requires two U-Nets to perform inference and subtraction. Training $s_{\phi}$ only involves training a single U-Net, and gradients do not need to be backpropagated to $G_{stu}$. Therefore, these models can be naturally deployed to different devices, making computational resource requirements more distributed. This ease of distribution allows for joint training on computational devices with limited capacity. In contrast, for GANs and VAEs, which require gradient propagation between models (discriminator to generator, decoder to encoder), computational requirements are more centralized, necessitating the use of a single device or tools like DeepSpeed to manage the workload.

### D.2 Vector Field

In our research, each of the estimated score functions $s_{\theta_{i}}$, for $i$ ranging from 0 to $N$, delineates a vector field $\mathbb{R}^{3\times W\times H}\mapsto\mathbb{R}^{3\times W\times H}$.
We make a strong assumption behind our proposed method that these score functions represent existing or non-existent distributions and that they altogether imply a transformation path between $s_{\theta}$ and the student generator $G_{0}$.
Nevertheless, a score fundamentally constitutes a gradient field, signifying the gradient of the inherent probability density.
A vector field is a gradient field when several conditions are satisfied, including path independence, continuous partial derivatives, and zero curls (Matthews, ).
The vector field, as characterized by the score functions, may not meet these conditions, and thus there is not a potential function or a probability density function.
Such deficiencies could potentially hinder the successful training of the student generator and introduce unforeseen difficulties in the distillation process.
Specifically, in instances where $s_{\theta}$ does not precisely represent a gradient field, a highly probable scenario considering $s_{\theta}$ is a neural network, the samples generated from $s_{\theta}$ could encompass failure cases.
Although our empirical studies exemplify the effectiveness of the proposed Generative Backtracking Model, the detrimental effects of the discussed issue remain unclear. We will further explore this issue in our future work.

## Appendix E Additional Details in Pre-experiments

### E.1 A Toy Experiment on Gaussian Mixture Distribution

To validate the feasibility of the proposed DisBack, we conduct experiments on two-dimensional Gaussian mixture data. First, we randomly select $10$ Gaussian distributions mixed as the training distribution $q_{0}$. Next, we construct a ResNet MLP as the two-dimensional diffusion model $s_{\theta}$ and train it using the created mixture Gaussian distribution. Similarly, we construct a simple MLP as the student generator $G_{stu}$ and train a model $s_{\phi}$ with the same architecture as $s_{\theta}$ using generated data. Therefore, we can use $s_{\theta}$ and $s_{\phi}$ to train the student generator $G_{stu}$. During the training process, we visualize the distribution of the student generator and training data to intuitively demonstrate the changes in the student generator distribution under the proposed training framework. The distribution of $G_{stu}$ during the training process is shown in Figure 8. As training progresses, the generated distribution $q^{G}$ initially expands outward and then gradually convergents towards the training distribution. The results show that the proposed method for training the student generator in the paper is effective.

Figure: (a) epoch=0
Refer to caption: extracted/5877782/SM/2D_1.png

### E.2 Gradient Orientation Verification of DisBack

As mentioned in Section 3.1 of the main paper, when updating $G_{stu}$ using Eq.(6), $s_{\theta}\left(\bm{x}_{t},t\right)$ provides a gradient towards the training distribution, while $s_{\phi}\left(\bm{x}_{t},t\right)$ provides a gradient toward the generated distribution.

To validate the correctness of these gradient directions, we experiment on two-dimensional data. We evenly sample $N$ data points within the range of $(x,y)\in[-6,6]$ as the noisy data $\bm{x}_{t}$. Subsequently, we depict the gradient directions of $\bm{x}_{t}$ based on $s_{\theta}$ and $s_{\phi}$ respectively. As shown in Figure 9, consistent with theoretical derivation, for any given $\bm{x}_{t}$, the gradient direction of $s_{\theta}\left(\bm{x}_{t},t\right)$ points toward the training distribution, and the magnitude of the gradient decreases as the distance to the training distribution decreases. Similarly, for any given $\bm{x}_{t}$, the gradient direction of $s_{\phi}\left(\bm{x}_{t},t\right)$ points toward the generated distribution.

Figure: (a) The gradient direction of $s_{\theta}\left(\bm{x}_{t},t\right)$
Refer to caption: extracted/5877782/SM/theta_gradient.png

## Appendix F Additional Samples from DisBack

### F.1 Samples of Image Generation

We provide additional samples from DisBack on FFHQ $64\times 64$ (Figure 10), AFHQv2 $64\times 64$ (Figure 11), ImageNet $64\times 64$ (Figure 12), LSUN Bedroom $256\times 256$ (Figure 13) and LSUN Cat $256\times 256$ (Figure 14).

### F.2 Samples Text-to-Image generation

## Appendix G Ethical Statement

### G.1 Ethical Impact

The potential ethical impact of our work is about fairness.
As “human face” is included as a kind of generated image, our method can be used in face generation tasks. Human-related datasets may have data bias related to fairness issues, such as the bias to gender or skin color. Such bias can be captured by the generative model in the training.

### G.2 Notification to Human Subjects

In our user study, we present the notification to subjects to inform the collection and use of data before the experiments.

> Dear volunteers, we would like to thank you for supporting our study. We propose the Generative Backtracking Model, which uses a backtracking strategy to achieve stable and fast distillation and high-quality single-step generation. All information about your participation in the study will appear in the study record. All information will be processed and stored according to the local law and policy on privacy. Your name will not appear in the final report. Only an individual number assigned to you is mentioned when referring to the data you provided. We respect your decision whether you want to be a volunteer for the study. If you decide to participate in the study, you can sign this informed consent form.

The Institutional Review Board approved the use of users’ data of the main authors’ affiliation.

Figure: Figure 10: Additional Samples form conditional FFHQ 64x64.
Refer to caption: extracted/5877782/figures/gbm_ffhq.png

Figure: Figure 11: Additional Samples form conditional AFHQv2 64x64.
Refer to caption: extracted/5877782/figures/gbm_afhqv2.png

Figure: Figure 12: Additional Samples form conditional ImageNet 64x64.
Refer to caption: extracted/5877782/figures/gbm_imagenet.png

Figure: Figure 13: Additional Samples form conditional LSUN bedroom.
Refer to caption: extracted/5877782/figures/gbm_bedroom.png

Figure: Figure 14: Additional Samples form conditional LSUN cat.
Refer to caption: extracted/5877782/figures/gbm_cat.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: The original motivation suggests that the teacher training trajectory could be used, but "the convergence trajectory of most teacher models is inaccessible". It's not clear to me that the training trajectory of the teacher would be useful for distillation, as it may not align well with the student distribution either. Did you explore DisBack using the training trajectory for trained diffusion models?

Did you explore distillation into smaller student models? Mismatched architectures could be a useful application of DisBack too.

Do you have any samples from the models along the intermediate teacher trajectory? Do these produce sensible samples at all?

Overall, I like the proposed DisBack algorithm and feel that it is sufficiently novel and performant to justify publication. I would give the paper a higher score if the authors provided some more experimental investigation into why their method is successful.


Minor:

Fig 2. includes the same sample twice (top middle, and middle).
soundness: 2
strengths: This paper proposes a novel technique that is intuitive and is shown to work well. The experiments show a clear improvement in the convergence speed of the DisBack method and the qualitative results show high-quality single-step samples. The authors apply their method to a variety of teacher models over multiple datasets and demonstrate success and failure cases (in the appendix).

The paper is easy to follow. Technical content is presented clearly.
summary: This paper introduces a novel approach to improve diffusion distillation. The key idea is to include a trajectory of teacher distributions for the student to match. This improves convergence and the final quality of the student model.
weaknesses: I felt that the paper writing in places suggested more than was provided. For example, the authors claim that they "identified this issue arises because existing score distillation methods focus on using the endpoint". However, the authors provide no such identification in their own work. They provided sufficient evidence that utilizing a trajectory of distributions improves the student but this has not been shown to be necessary. There may be alternative approaches that work well while using only the endpoint. This is present elsewhere in the work, e.g. "the fat convergence speed is because constraining the convergence trajectory of the generator provides a clear optimization direction", this is a strong technical claim that has not been adequately explored.

The authors could do more to explain why DisBack is successful. Some theoretical analysis could help to explain the improved convergence speed, or experiments designed to show the convergence more carefully. For instance, I'd be interested to see how quickly the student model converges to each point on the trajectory, and how closely it matches the distribution. Presumably, this behaviour must be better than linear for DisBack to succeed, which is interesting. Overall, I felt that the idea worked well and was intuitive, but I didn't understand why it worked so well after reading the paper.

The ablation study is minimal (and I would argue does not qualify as an ablation study). The authors only explore DisBack and the original method side-by-side. Instead, they could also investigate the effect of the number of degradation path checkpoints, different student initializations, different degradation schemes (like using the training trajectory), etc.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See Weaknesses.
soundness: 2
strengths: 1. The paper has a clear motivation: to address the initial score mismatch problem in diffusion distillation. Specifically, the authors first identify the suboptimal performance of existing diffusion distillation methods as being due to the score mismatch issue at the beginning of student model training, and then propose a novel approach to resolve it. 
2. The proposed approach is intuitive and effective. It makes sense to follow the degradation path from the teacher model to the student model in reverse during distillation, providing a progressive learning signal for the student and mitigating the initial score mismatch problem. In practice, by introducing this backtracking strategy, the distillation process is shown to be significantly faster than its variant without this technique. 
3. The proposed approach is versatile, as it is orthogonal to other diffusion distillation approaches.
summary: In this paper, a novel approach is proposed to further improve existing diffusion distillation methods. The proposed approach leverages the convergence trajectory from teacher model to the initial state of student model to guide the training of student model backwards, which mitigates the score mismatching problem at the beginning of the distillation. Empirical results have demonstrated the superior performance of the proposed approach.
weaknesses: 1. The state-of-the-art claim for the proposed method is misleading. According to the experimental setup in Appendix B.2, the proposed method trains for 500,000 **epochs** ($\approx500K \times \frac{|D|}{|B|}$ **iterations or steps**, where $|D|$ is the data size and $|B|$ is the batch size). This number is significantly higher than for the baselines. For example, Diff-Instruct only trains for ${\color{red}50K}$ **iterations** on ImageNet $64\times 64$, while DisBack (this paper) uses about ${\color{red}500K\times 40K=20G}$ **iterations** ($|D|=1,281,167$ and $|B|=32$), which is approximately $40,0000$ times larger. Even if "epochs" actually refers to steps (if it is a typo), it still represents 10 times the training length compared with the Diff-Instruct baseline. Additionally, the result (${\color{green}1.51}$) of DMD2 on ImageNet $64\times 64$ is achieved by training for ${\color{red}200K}$ **iterations**. With the extended training setup (${550K}$ **iterations** in total), DMD2 could achieve an FID of ${\color{green}1.28}$, which is lower than DisBack's ${\color{green}1.38}$. This raises concerns that the proposed strategy may not fully account for the state-of-the-art performance showcased. This is also supported by their ablation study and Table 3. The variant ("w/o convergence trajectory") is essentially Diff-Instruct, as noted in the main text on Lines 391-392. However, even this variant, when trained under the same setting, shows better performance on FFHQ (12.26) versus the original Diff-Instruct (19.93).

2. The speedup shown in Figure 1 is only plotted for epochs 0 to 2000, which covers only the early stage of the distillation. More epochs, either until convergence or until training budgets are exhausted, are needed to better understand how the backtracking strategy behaves throughout training.

3. Although the entire concept revolves around backtracking the degradation path, in actual training, each intermediate checkpoint is only trained for as few as $1000$ steps (for FFHQ, AFHQv2, and ImageNet at $64\times64$ resolution), while the remaining steps are trained with the original teacher model. This means that the proposed backtracking is used for only a small fraction of the student model's training, which makes it even harder to attribute the superior performance to the proposed strategy.

---

confidence: 5
contribution: 3
initial_rating: 3
presentation: 3
questions: My major worry is the reliability of performance and the laborious algorithm. Please respond to my worries.
soundness: 3
strengths: - The idea is simple but effective. It is makes sense that distilled from degraded teacher make the distillation faster.
- Degradation Recording algorithm looks reasonable. The degraded teacher finally converges at the initialized student distribution, which make the student easy to learn in the early stage.
- The result compared to Diff-instruct seems the algorithm is effective.
summary: This paper propose a DisBack, which is a new distillation method of diffusion models. On the top of the Diff-Instruct, this DisBack propose a better training algorithm. While Diff-Instruct only use the pre-trained diffusion teacher, DisBack makes a series of degraded teachers, and use that teachers iteratively. This makes the student models easy to learn a teacher distribution.
weaknesses: - My major worry is that I can not trust the performance. The paper distilled from EDM in ImageNet 64 which have 2.44 FID, but the distilled student has the performance of 1.38 FID. In my understanding, I think the upper bound performance of the student is teacher.

- I also can not believe user preference study compared to the SDXL teacher. How can it better than teacher? 

- The ablation on the number of degraded teacher N is missing. I want to see progressive performance boosting from N=1 (equivalent to the Diff-Instruct) to N= large.

- Is there any scheduling algorithm that changes the teacher in stage 2? It may requires lots of trials to find the schedule that determine when to change that target teacher from degraded to the original.

- Figure 1 is a little bit over-claimed. This algorithm should contains the training costs of stage 1.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - The score estimation (7) is not general for all noising methods. For example, the score estimation of ddpm has a mean scale $\alpha_t$. When do distillation, should the teacher and student noising methods keep consistent?
- Compared with Diff-Instruct which only training student models to fit one teacher model, Algorithm 2 needs to fit $N-1$ intermediate checkpoints, what about the training overhead of this part? In Fig.1, did the epochs for DisBack in x-axis record from the degradation recording stage or from which other points?
- Any experiments to show the influence of the number of degradation checkpoints and the number of degradation epochs? Will more checkpoints and epochs mitigate the mismatch better?
soundness: 2
strengths: 1. DisBack addresses the common “score mismatch” issue in score distillation by incorporating the entire convergence trajectory. DisBack enables the student generator to align more accurately with the teacher model, leading to faster convergence and better optimization paths.
2. DisBack is designed to be easily integrated into current distillation frameworks, providing a versatile tool to further boost performance in generative model distillation.
summary: This paper introduces Distribution Backtracking Distillation (DisBack), a method to accelerate sampling in diffusion models by addressing the “score mismatch” issue common in traditional score distillation approaches. Unlike existing methods that rely solely on the endpoint of a pre-trained teacher model, DisBack captures the full convergence path between the teacher and student models. It does this through two stages: Degradation Recording, which records a degradation path from the teacher to the untrained student model, and Distribution Backtracking, where the student generator retraces this path to improve alignment with the teacher model.
weaknesses: **Major**:
- While the authors claim DisBack is orthogonal to those of other distillation methods, there is no evidence to support this point. It would be valuable if the authors could provide further experiments to show it can be incorporated into other distillation methods, like consistency distillation or adversarial score distillation. 
- The paper aims to mitigate the score mismatch issue by employing degradation recording as convergence trajectory for distillation. The mismatch between the predicted score of generated samples and the model's prediction will be degraded but the mismatch between the model's prediction and the teacher's score prediction will be larger in degradation path.  This suggests a potential tradeoff between these two types of mismatches, which could impact the final model’s performance. Providing further analysis or empirical results on this point would strengthen the motivation and effectiveness of this approach.

**Minor**:
- In Eq.(6), $\partial x_t/\partial \eta$ should be included in expectation, same as (8). 
- Better to use bold $\epsilon$ for noise and show the relationship between $\epsilon$ and $x_t$. 
- In Algorithm 1&2, since the loss includes the expectation w.r.t. $t$ and $\epsilon$, the line to calculate $x_t = x_0 + \sigma_t \epsilon$ is unnecessary and misleading. 
- Labels in Fig.7 are wrong.

Meta-review: This paper targets enhancing diffusion distillation through a novel means of resolving the score mismatch issue by incorporating a trajectory of teacher distributions for student matching, which boosts convergence and student model quality. Reviewers generally find the method novel and captivating, with significant experimental outcomes. The experiments evidently enhance DisBack's convergence speed, and qualitative results present high-quality single-step samples. However, concerns pertain to justifying experimental results and analyzing the sensitivity of backtracking schedule choices. AC agrees this is a novel contribution to the community, hence recommending it as a poster.

Decision: Accept (Poster)