David and Goliath: Small One-step Model Beats Large Diffusion with Score Post-training

Dear reviewer, we are delighted that you like the novelty of DIstar for the one-step diffusion model post-training. We appreciate your valuable suggestions. In the following paragraphs, we will address your concerns one by one.

Q1. Clarifications on image-data-free property and post-training setups.

A1. We will polish the setup as image-data free and post-training of one-step diffusion models.

Q2. (1) Broader evaluations, such as evaluations on challenging prompts like Parti and HPSv-benchs. (2) How does DIstar perform on models with the ability to take complex prompts, such as Pixart-alpha models? (3) Explorations on other rewards.

A2. Thanks for your useful feedback. In Table 1 and Table 2 in our paper, we have compared DIstar with other models on Parti prompts and HPSv2.1 benchmark. In the rebuttal period, we conduct two new experiments: (1) DIstar with SDXL-DPO (a pre-aligned SDXL model) as the reference diffusion; (2) DIstars with Pixel-art-alpha model that uses ImageReward as explicit reward models. We present more quantitative results in Table 2 (in rebuttal to Reviewer NM9A), and Table 1.

Table 1. Quantitative comparisons of one-step SD1.5 model and one-step Pixel-art-alpha (PAA) models and others in Preference Score on COCO-2017-validation prompt dataset.

The Pixel-art-alpha (PAA) diffusion model uses the T5 text encoder, which has a much longer context length. Besides, we also explore the use of ImageReward to see if DIstar is consistent across different reward models. So, we compare DIstar on the PAA model with other open-sourced models. As Table 1 shows, the PAA-DIstar model achieves a leading ImageReward, Aesthetic score, and HPSv2.0 score. It goes on par with the best PickScore as PAA-diffusion-25step shows.

Q3. Discussion on the mismatch setups introduced in MaPO.

A3. Thanks for the valuable comment. After a careful reading of the MaPO paper, we find it technically very interesting. MaPO introduced a novel approach to post-train diffusion models in cases of preference mismatch. We acknowledge that in our paper, we did not consider cases of preference mismatches, which is potentially important to study in our future. However, the DIstar and MaPO have essentially different targets in

(1) MaPO targets in diffusion models, while DIstar targets in post-training of one-step models;

(2) MaPO assumes access to the preference dataset, while DIstar assumes we have high-quality and diverse reward models;

However, we really appreciate the solid contributions of MaPO. We will add an additional discussion on MaPO in our revision. We are glad to explore DIstar in cases of preference mismatches introduced by MaPO.

Q4. Discuss DIstar with other related approaches.

A4. Please see A1 in rebuttal to Reviewer KEDK.

Q5. Are there any criteria the assistant model needs to satisfy?

A5. In general, arbitrary diffusion models can serve as an assistant model. But in practice, we initialize the assistant diffusion model with the preference diffusion model and find it works well.

Q6. Clarify more on three losses.

A6. Please see A3 in the rebuttal to Reviewer NM9A.

Q7. What are the memory requirements of this method? Is DIstar compatible with LoRA?

A7. We acknowledge that DIstar needs an additional assistant diffusion model, which brings more memory costs. However, in practice, we find that the additional cost is acceptable if we use Pytorch techniques like BF16 training, Distributed data-parallel (DDP), and gradient accumulations. We did not use LoRA in order to get the best performances.

Q8. Criteria for the assistant model. DIstar for few-step models?

A8. In general, arbitrary diffusion models can serve as an assistant model. DIstar for a few steps, possibly in our future work.

Q9. Other comments on paper writing.

A9. We will remove the blue colors. $p_\theta$ represents the distribution of the one-step model. $d'(y)$ in equation 3.4 represents the derivative of distance $d(y)$. We appreciate your valuable comments on the paper writing, and we will revise the paper writing in the revision. We will add experiment details in our revision.

We appreciate your valuable review. We sincerely hope our response has resolved your concerns. If you still have any concerns, please let us know.

Dear reviewer, we are glad that you like our novelty of regularizing the one-step diffusion model during post-training. In the following paragraphs, we will answer your questions one by one.

Q1. Experimental or theoretical evidence that shows mode-collapse issues of one-step diffusion distilling using KL-divergences.

A1. We appreciate your good question! We will answer your question from both theoretical and empirical perspectives.

Theoretical perspective: If we use $p_ r$ to represent reference distribution and $p_{g}$ the one-step generator distribution. The KL divergence between $p_ g$ and $p_ r$ is defined as: $D_ {KL}(p_ g||p_ r) = \mathbb{E}_ {x\sim p_g}\log \frac{p_ g (x)}{p_ r (x)}$ while the general score-based divergence with a form $\mathcal{D}_ {SD}(p_ g,p_ d) = \mathbb{E}_ {x\sim \pi} d(\nabla_ x \log p_ g (x) - \nabla_ x \log p_ d (x))$. KL divergence is long believed to be bothered by mode-collapse issues because the likelihood ratio $\frac{p_ g (x)}{p_ r (x)}$ will be ill-defined if the $p_ g$ and $p_ d$ have misaligned density support. However, the score-based divergence does not have such a "ratio" that will be ill-posed and, therefore, is safe in the case when $p_ g$ and $p_ r$ have misaligned support. Besides, the distance function $d()$ of general score-based divergence also brings flexibility in defining divergences to properly measure the difference between two distributions in a robust magnitude. For instance, our used Pseudo-Huber distance (line 259) has a self-normalization form, which also helps stabilize the gradient of score-divergence to a unit scale, potentially further helping address the mode-collapse issue.

Empirical perspective: We conduct a new experiment on one-step diffusion distillation on the CIFAR10 unconditional. The FID and Recall of KL and score-divergence model are (3.70,0.52) and (2.01,0.60). This shows that distilling using KL divergence leads to both worse FID and Recall.

Q2. Is aligning a distilled one-step diffusion model more effective than distilling from a pre-aligned multi-step diffusion model? If so or not so, why or why not?

A2. This is a very good question! We will answer via three perspectives:

• First of all, distillation takes far more than compute than post-training: roughly 8:1 GPU hours. Therefore, in this paper, we focus on post-training that improves existing models with minimal costs.

• Second, we appreciate your good intuition and find that if we replace the naive reference diffusion with an aligned diffusion (like SDXL-DPO) in DI-star, we can even accelerate the post-training process with better final results. In the rebuttal phase, we conduct a new experiment (in Table 2) that compares DI-star w/o pre-aligned reference diffusion, as well as other DMD2-like models.

Experiment Setup. We replace SDXL with SDXL-DPO as the reference diffusion model. Then, we compare models after post-training using DI-star, DMD, and DI++.

As Table 2 shows, in DI-star (Score-based PPO), both reward models and pre-aligned reference diffusion contribute to the performance improvements of one-step models in post-training. Besides, we find that DI++(KL-PPO) will harm the contribution of the reward model and pre-aligned diffusion.

Table 2. Quantitative comparisons of Preference Scores after post-training using SDXL-DPO on Parti prompts (left part) and HPSv2.1 score (right part).

Q3. More clarifications of the trade-offs between the three loss signals

A3.

• We find that CFG-implicit reward easily causes over-saturated colors. On the contrary, CFG-enhanced reference diffusion will not cause issues like over-saturated colors.

• ImageReward is conflict to CFG-rewards and CFG-enhanced divergences. However, PickScore is not conflicted with CFG-enhanced divergence loss, which makes both losses improve the one-step model during post-training.

We appreciate your valuable review. We sincerely hope our response has resolved your concerns. If you still have any concerns, please let us know.

Dear reviewer, we are delighted that you like our score-based reinforcement learning post-training of the one-step diffusion model. We appreciate your valuable suggestions. In the following paragraphs, we will address your concerns one by one.

Q1. Clarify the differences between Diff-Instruct-star, DMD, and other approaches.

A1. We are sorry for the confusion. In one-step diffusion distillation/post-training literature, there are typically two types of training criteria.

• The first type is trajectory-based methods, such as Progressive distillation, consistency distillation, and other variants.

• The second type is distribution matching based methods, which aim to minimize some probability divergence between one-step model distribution and teacher diffusion model distribution. Among them, Diff-Instruct first proposes one-step diffusion distillation by minimizing the integral KL divergence. DMD and DMD2 generalize the KL divergence minimization by incorporating a regression loss and a GAN loss. SiD[4] studies the minimize the one-step distillation by minimizing Fisher divergence instead of KL. Later, SIM unifies the one-step diffusion distillation by minimizing general score-based divergences (including Fisher divergence) with detailed theories that prove the gradient equivalence between SIM loss and the underlying divergence.

• Post-training: The reward-guided LCM first combines a reward model with consistency loss to train few-step diffusion models. DI++ then studied the post-training of the one-step model by introducing a KL divergence-PPO framework and achieved a very strong human preference quality. Our DIstar is inspired by DI++, by differs in (1) theoretical and empirical study the post-training via score-based-PPO instead of KL-based PPO; (2) introducing novel techniques that successfully scales the post-training of one-step diffusion model at 1024x1024 resolution, and outperforms current SoTA diffusion model (FLUX-dev) with only 1-step (1.8%) inference costs. This surprisingly strong performance also gives us the inspiration for our title Small one-step model beats large diffusion with score-based post-training.

In conclusion, DIstar follows the Occam's Razor principle: achieving better human-preference quality than open-sourced SoTA FLUX-dev with minimal losses (reward and general score-divergence).

Q2. Discussions on why the score-based divergence is potentially better in diversity-preserving than KL.

A2. We appreciate your good question! We will answer your question from both theoretical and empirical perspectives.

Theoretical perspective: If we use $p_ r$ to represent reference distribution and $p_{g}$ the one-step generator distribution. The KL divergence between $p_ g$ and $p_ r$ is defined as: $D_ {KL}(p_ g||p_ r) = \mathbb{E}_ {x\sim p_g}\log \frac{p_ g (x)}{p_ r (x)}$ while the general score-based divergence with a form $\mathcal{D}_ {SD}(p_ g,p_ d) = \mathbb{E}_ {x\sim \pi} d(\nabla_ x \log p_ g (x) - \nabla_ x \log p_ d (x))$. KL divergence is long believed to be bothered by mode-collapse issues because the likelihood ratio $\frac{p_ g (x)}{p_ r (x)}$ will be ill-defined if the $p_ g$ and $p_ d$ have misaligned density support. However, the score-based divergence does not have such a "ratio" that will be ill-posed and, therefore, is safe in the case when $p_ g$ and $p_ r$ have misaligned support. Besides, the distance function $d()$ of general score-based divergence also brings flexibility in defining divergences to properly measure the difference between two distributions in a robust magnitude. For instance, our used Pseudo-Huber distance (line 259) has a self-normalization form, which also helps stabilize the gradient of score-divergence to a unit scale, potentially further helping address the mode-collapse issue.

Empirical perspective: We conduct a new experiment on one-step diffusion distillation on the CIFAR10 unconditional. The FID and Recall of KL and score-divergence model are (3.70,0.52) and (2.01,0.60). This shows that distilling using KL divergence leads to both worse FID and Recall.

Q3. It would be good if the authors could present more qualitative examples comparing to KL-based models, such as DI++-SDXL, to verify the motivation of the method.

A3. In Fig. 4 in the appendix of the paper, we qualitatively compare the DIstar model with DI++, SIM, Diff-Instruct, DMD2, SDXL, and SDXL-DPO models. We clearly find that images generated by the DIstar model show better aesthetic quality, with gentle lights and colors. We also put more qualitative comparisons between DIstar-SDXL-DPO and DMD-SDXL-DPO in this anonymous link: https://anonymous.4open.science/r/distar_anonymous-E007/distar_dmd.png.

We appreciate your valuable review. We sincerely hope our response has resolved your concerns. If you still have any concerns, please let us know.