Tuning Timestep-Distilled Diffusion Model Using Pairwise Sample Optimization

Thanks to the reviewer for the constructive feedback, we provide our responses below and hope it can alleviate your concerns.

Q1: Compare the computational costs of distillation with human preference tuning using PSO.

We compare the computational cost of the distillation of SDXL-DMD2 with PSO tuning for human preference by measuring their training GPU hours. Note that a similar level of computational cost of human preference tuning is also reported in MaPO[1] (our cost is lower than MaPO as we remove the redundant data from the dataset). As shown in the table below, the computational cost of human preference tuning with Offline-PSO is significantly less than the diffusion model distillation. Specifically, the preference tuning cost of PSO is roughly 2.0%(78/3840) of the distillation cost of DMD2. Thus, our method demonstrates significant advantages in training efficiency compared with tuning followed by timestep distillation.

Q2: More quantitative results using CLIP and DINO scores.

Thanks to the reviewer for bringing this out. We conduct quantitative experiments with the proposed method on concept customization. Specifically, we adopt the experimental setup in Dreambooth by evaluating our method and the baseline with all the 30 concepts in the Dreambooth dataset. For evaluation, we sample images with 25 test prompts for each concept. As for the metrics, we follow the reviewer's suggestion and use the CLIP score for text-image alignment (CLIP-T), and the CLIP score with DINO score for fidelity (DINO and CLIP-I). We compare our method with customizing the multi-step diffusion model (Dreambooth w/ SDXL) and the distilled baseline model before customization (SDXL-Turbo).

As shown in the table above, the proposed PSO demonstrates effectiveness given the quantitative results. Our method shows superior results than SDXL-Turbo and comparable results with multi-step customization baseline. Specifically, compared with Dreambooth w/ SDXL, our method achieves a higher text-alignment CLIP-T score while using a magnitude fewer inference steps.

Q3: Comparison with SPIN.

Thanks to the reviewer for bringing this out. Our method is targeting at tuning few-step timestep-distilled diffusion models, whereas SPIN fine-tunes the multi-step regular diffusion models. We also show the effectiveness of PSO in general diffusion model tuning tasks including preference tuning, style transfer, and concept customization, while SPIN solely focuses on preference tuning. Moreover, we would like to state that SPIN consists of training iterations, where each iteration can be seen as Offline-PSO with offline sampled reference data, thus SPIN can be viewed as a multi-stage Offline-PSO. We add this discussion in the related works.

Q4: Why offline PSO for style transfer and online for customization.

We would like to clarify that we do not use Online-PSO for customization where both target and reference images are sampled from the distilled diffusion model. As stated in Line 459, we adopt the PSO objective for customization defined in Eq. (3), where the reference images are sampled throughout the training. We adopt the Offline-PSO setting for style transfer for efficiency concerns as it takes a longer training time. We conduct an ablation style transfer experiment with Eq. (3) and it works as well as Offline-PSO as in the paper.

Q5: Table typos.

Thanks to the reviewer for pointing out the typos. We have changed that in the revision.

Q6: Reward model for online PSO.

We adopt Pickscore in the Online-PSO for human preference tuning. We will clarify that in the revision.

Reference:

1. Hong, J., Paul, S., Lee, N., Rasul, K., Thorne, J., & Jeong, J. (2024). Margin-aware Preference Optimization for Aligning Diffusion Models without Reference. arXiv preprint arXiv:2406.06424.

Thanks to the reviewer for acknowledging the novelty and our contribution. We respond to your questions below.

Q1: Notation clarification.

Thanks to the reviewer for pointing this out. We will change the notation of the reference image $\rho$ in the revision.

Q2: Additional algorithm table.

Thanks to the reviewer for bringing this out. We provide an algorithm table in the revision.

Q3: DiT based diffusion models.

Thanks to the reviewer for this question. Our method is agnostic to model architectures. While we mainly validate our methods on U-Net based timestep-distilled diffusion models, i.e., SDXL distillation variants, we suppose the method can also be applied to DiT-based distilled models like FLUX-schnell. We will explore it in the future work.

Thanks to the reviewer for the constructive feedback. We provide our response below and hope it can alleviate your concerns.

Q1: More fine-tuning baselines for comparisons.

Thanks to the reviewer for bringing this up. We add more baselines including the supervised fine-tuning(SDXL-SFT), RPO[1], MaPO[2], and SPO[3] as additional human preference tuning baselines in Table 1 in the revision and show it below. Note that the proposed PSO targets directly tuning distilled diffusion models in a few-step(1-4 step) generation regime, while other fine-tuning baselines including DPO tune the multi-step(50 steps) diffusion models. To make it a fair comparison, we mainly focus on comparing with the few-step distilled models in the tables, indicated by the gray and black text color in Table 1. Nevertheless, as shown in the table below, the distilled few-step model (SDXL-DMD2) tuned with both Online-PSO consistently achieves the best results in Pickscore (as our major focus is human preference) and comparable results on other metrics compared with competitive multi-step diffusion models which are tuned with various human preference tuning baselines. This further highlights the effectiveness of the proposed method.

Q2: Table typos.

Thanks to the reviewer for pointing out the typo. We have changed them to the correct names in the revision.

Q3: Question on bold results.

As we stated above, our method tunes the few-step distilled diffusion models and we mainly focus on comparing with few-step diffusion model baselines. We provide an updated table that mainly compares the proposed method with few-step baselines below. As shown in the table, the proposed method achieves the best results and thus we make our results bold.

Thanks to the reviewer for acknowledging our contribution. We provide our response to the questions below.

Q1: More Analysis on PSO.

• Why directly minimizing the diffusion objective leads to blurry. The diffusion objective, $||\epsilon_\theta'(x_t, t)-\epsilon||^2$, trains the diffusion model $\epsilon_\theta'$ to predict the score as, $\nabla_x\log p(x_t)=- \epsilon_\theta'(x_t, t) / \sigma_t$, which be used to solve probability flow ODE(PF ODE) as the diffusion model sampling, $\frac{dx_t}{dt} = f(t)x_t + \frac{g^2(t)}{2\sigma_t}\epsilon_\theta'(x_t, t)$. Solving the reverse ODE, i.e., sampling from multi-step diffusion models, usually requires 50 ODE solver steps. That is, the model trained with the diffusion objective is used for 'local' prediction(small timestep range, e.g., $t=975$ to $t=950$) on PF ODE. On the other hand, the trained timestep-distlled model $\epsilon_\theta$ is trained to directly predict the endpoint $x_0$ of the PF ODE from arbitrary point $x_t$, which is a 'global' prediction, and obtain the few-step generation ability as a result. In this way, directly tuning the distilled model using the original diffusion objective will distort its 'global' prediction ability from $x_t→ x_0$, and resulting in blurry results, which is shown in Figure 1.

• Why recast the optimization as minimizing the relative likelihood. We choose to recast the optimization as maximizing relative likelihood between target and reference samples to avoid directly tuning the distilled model with 'local' score matching objective, as stated above, which leads to blurry few-step generation results. Inspired by some initial human preference tuning experiments with pairs of data, we find that by instead maximizing the relative likelihood ratio $\log\frac{p_θ(x^τ_0|c)}{p_θ(x^ρ_0|c)}$, we can preserve the few-step generation ability while effectively steering the generative distribution toward desired targets. Thus, we design PSO that maximizes the relative likelihood of pairwise data, which is further validated in other tuning experiments besides human preference tuning.

• How to choose pre-sampled images. We select the pre-sampled reference images from the model distribution, which contributes in PSO to preserve the few-step generation ability. Specifically, in human preference tuning, we directly adopt the negative samples in the pickapicv2 dataset as the reference images, as most of them are sampled with SDXL, which is close to the generative distribution of the distilled models. While in style tuning, we pre-sample the reference images from the distilled model using the prompts in the style transfer dataset.

Q2: Question on Baselines selected, why SDXL-DPO + DMD2 LoRA, rather SDXL-DMD2 + DPO?

Thanks to the reviewer for bringing this up. We would like to compare to the baseline where the multi-step diffusion model is tuned for human preference first (SDXL-DPO), and then distilled for few-step generation (DMD2). This baseline would be too costly to obtain (i.e., 3840 A100 GPU hours in Line 52), and we choose the DMD2 LoRA applied to SDXL-DPO (SDXL-DPO + DMD2 LoRA) as an efficient proxy. Meanwhile, as stated in Line 244, our Offline-PSO has a very similar formulation to Offline Diffusion-DPO. Thus, the required baseline SDXL-DMD2-LoRA + DPO can be seen as Offline-PSO results in the Tables (as discussed in Line 95), which shows effectiveness in tuning timestep-distilled diffusion models, and yet is not as effective as the proposed Online-PSO in human preference tuning.

Q3: Typos.

Thanks to the reviewer for pointing them out. We have revised that in the revision.