Distribution Backtracking Builds A Faster Convergence Trajectory for Diffusion Distillation

Authors: Shengyuan Zhang, Ling Yang, Zejian Li, An Zhao, Chenye Meng, Changyuan Yang, Guang Yang, Zhiyuan Yang, Lingyun Sun
Primary area: generative models
Keywords: Diffusion Model, Diffusion Distillation, One-step Generation

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 two score functions on the samples generated by the student model.
However, there is a score mismatch issue in the early stage of the score distillation process, since 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 the teacher model and propose $\textbf{Dis}$tribution $\textbf{Back}$tracking Distillation ($\textbf{DisBack}$). DisBask is composed of two stages: $\textit{Degradation Recording}$ and $\textit{Distribution Backtracking}$. 
$\textit{Degradation Recording}$ is designed to obtain the convergence trajectory by recording the degradation path from the pre-trained teacher model to the untrained student generator.
The degradation path implicitly represents the intermediate distributions between the teacher and the student, and its reverse can be viewed as the convergence trajectory from the student generator to the teacher model.
Then $\textit{Distribution Backtracking}$ trains the student generator to backtrack the intermediate distributions along the path to approximate the convergence trajectory of the teacher model.
Extensive experiments show that DisBack achieves faster and better convergence than the existing distillation method and achieves comparable or better generation performance, with an FID score of 1.38 on the ImageNet 64$\times$64 dataset.
DisBack is easy to implement and can be generalized to existing distillation methods to boost performance.

Decision: accept (poster)
Reviewer scores: 8, 6, 3, 5

Reviews and discussion:
<thread>
creation date:
1730561449800

thread_type:
official_review

<thread_object>
creator: reviewer_hutp

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649553259

thread_type:
official_review

<thread_object>
creator: reviewer_jlt9

<content>
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.

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.

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.

questions:
My major worry is the reliability of performance and the laborious algorithm. Please respond to my worries.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730670587221

thread_type:
official_review

<thread_object>
creator: reviewer_vtdx

<content>
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.

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.

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.

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730758530312

thread_type:
official_review

<thread_object>
creator: reviewer_yzv7

<content>
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.

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.

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.

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).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732860702953

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We sincerely appreciate all the reviewers for their thorough examination and valuable suggestions. We are glad to hear that the proposed idea is intuitive, effective and novel (Reviewers vTDX, JLT9 and yzv7), the proposed method is easy to understand (Reviewers yzv7 and JLT9) and versatile (Reviewers HUTP and vTDX). We have revised the manuscript according to the comments of the reviewers (**highlighted in red**).

Here, we summarize and highlight our responses to the reviewers:
* We added a preliminary experiment to provide evidence for the score mismatch issue present in existing score distillation methods (Reviewer yzv7) and further conducted a detailed analysis of the motivation and effectiveness of the proposed DisBack method intuitively (Reviewers yzv7 and HUTP).
* We included an ablation study on the number of checkpoints and different degradation schemes to further explain and validate the effectiveness of DisBack (Reviewers yzv7, JLT9 and HUTP).
* We added a discussion on different types of mismatch issues (Reviewer vTDX) and a comparison with existing related work (Reviewer JLT9).
* We revised Fig. 1 to compare DisBack with Diff-instruct while accounting for the computational cost of the first stage (Reviewer vTDX, JLT9 and HUTP).
* We revised the unclear statements in the paper to avoid confusion (Reviewers vTDX, yzv7, and JLT9).

We reply to each reviewer's questions in detail below their reviews. Please kindly check out them. Thank you and please feel free to ask any further questions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.