Tuning Timestep-Distilled Diffusion Model Using Pairwise Sample Optimization

## Abstract

Abstract Recent advancements in timestep-distilled diffusion models have enabled high-quality image generation that rivals non-distilled multi-step models, but with significantly fewer inference steps. While such models are attractive for applications due to the low inference cost and latency, fine-tuning them
with a naive diffusion objective would result in degraded and blurry outputs.
An intuitive alternative is to repeat the diffusion distillation process with a fine-tuned teacher model, which produces good results but is cumbersome and computationally intensive: the distillation training usually requires magnitude higher of training compute compared to fine-tuning for specific image styles.
In this paper, we present an algorithm named pairwise sample optimization (PSO), which enables the direct fine-tuning of an arbitrary timestep-distilled diffusion model.
PSO introduces additional reference images sampled from the current time-step distilled model, and increases the relative likelihood margin between the training images and reference images. This enables the model to retain its few-step generation ability, while allowing for fine-tuning of its output distribution. We also demonstrate that PSO is a generalized formulation which can be flexibly extended to both offline-sampled and online-sampled pairwise data, covering various popular objectives for diffusion model preference optimization.
We evaluate PSO in both preference optimization and other fine-tuning tasks, including style transfer and concept customization. We show that PSO can directly adapt distilled models to human-preferred generation with both offline and online-generated pairwise preference image data.
PSO also demonstrates effectiveness in style transfer and concept customization by directly tuning timestep-distilled diffusion models.

## 1 Introduction

Figure: Figure 1: Illustration of images sampled from: (a) original timestep-distilled diffusion models, (b) Fine-tuned distilled model with diffusion objective (Ho et al., ), and (c) Fine-tuned distilled models with our PSO, and It can be seen that simply tuning distilled models with the vanilla diffusion loss leads to blurry, degraded generation, while our method can steer the distilled model toward better alignment with human preference & prompts, and style-transferred generation. Prompt from left to right: A Pirate in a Pirateship.// a woman with long hair next to a luminescent bird.// Photograph of a wall along a city street with a watercolor mural of foxes in a jazz band.// A stern-faced, brown-feathered owl Pokémon with a leaf-shaped crown and piercing red eyes.// A cute rabbit.
Refer to caption: /html/2410.03190/assets/figures/teaser_arrow.jpg

Diffusion models (Ho et al., ; Song et al., ; Nichol & Dhariwal, ; Ho & Salimans, ; Karras et al., ; Rombach et al., ) have shown strong capabilities in generating high-fidelity images, marking a significant advancement in the field of generative modeling. Despite their impressive performance, a notable challenge is the high inference cost due to its iterative denoising nature. To address this issue, various methods are proposed to accelerate the sampling process of diffusion models, including improving the efficiency of samplers (Karras et al., ; Lu et al., ; ; Liu et al., ; Zhao et al., ) and employing model distillation techniques (Song et al., ; Luo et al., ; Sauer et al., ; Xu et al., ; Lin et al., ; Sauer et al., ; Yin et al., ; ; Ren et al., ; Kohler et al., ) to reduce the number of inference steps. Recent advancements in trajectory distillation methods and distribution matching techniques, often enhanced by adversarial learning at scale, have shown considerable promise in generating high-fidelity images in extremely low steps such as one to four steps.

Despite significant advancements in timestep-distilled diffusion models, it remains unclear how to effectively fine-tune or customize such distilled models. Naively tuning the distilled model with diffusion loss will make the generation results blurry, as shown in Figure [1](#S1.F1) (b). An alternative approach is to fine-tune or customize the original diffusion model, and then repeat the diffusion distillation process to create a distilled model variant. However, the large computation cost of diffusion distillation, when compared with the customization training used for applications (*cf*., *3840* A100 GPU hours for SDXL-DMD2 (Yin et al., ) and *0.25* A100 GPU hours for Dreambooth concept customization (Ruiz et al., )), often makes such distilled model tuning approach less feasible.

In this paper, we present an effective algorithm named pairwise sample optimization (PSO), designed to directly tuning a timestep-distilled diffusion model for different user preferences or customizing it towards new image domains. PSO maximizes the likelihood ratio between a given pair of target and reference images, where the target image is sampled from the distribution we wish to tune the distilled model towards, and the reference image is sampled from the generative distribution of the tuned distilled models on-the-fly. By optimizing this relative likelihood objective, PSO effectively aligns the timestep-distilled models with target distributions while best preserving its original few-step generation ability learned through diffusion distillation. Our formulation of PSO can be flexibly adapted to both offline and online settings, and unifies various prior works (Wallace et al., ; Yang et al., ) as special cases to our framework.

We validate the effectiveness of the proposed method across a broad range of tasks, including human preference tuning, style transfer, and concept customization. For preference tuning and style transfer, we construct or directly utilize target-reference images with the given prompts, *e.g.*, win-lose pairs from the user preference data (Kirstain et al., ), and adopt our offline PSO objective to tune the distilled models. We also examine our online objective in preference tuning, where we sample pairs of images from distilled model and use a reward mode to assign the label of target or reference. For concept customization, we use the given concept images as target and sample reference images from the tuned model. In all tasks and settings, PSO enables the lightweight fine-tuning of arbitrary distilled models while preserving their few-step generation capability, as shown in Figure [1](#S1.F1) (b).

We evaluate our method on various datasets and benchmarks. For human preference tuning, we benchmark PSO on the standard Pick-a-Pic (Kirstain et al., ) and PartiPrompts (Yu et al., ) datasets. We also evaluate distilled model fine-tuning by experimenting the image style transfer task with the Pokemon dataset (Pinkney, ) and the concept customization task with the exeamples in Dreambooth (Ruiz et al., ). We observe that PSO shows comparable performance to the oracle results of fine-tuning a new teacher model and conduct the diffusion distillation training, while being much less computationally expensive. Our approach also significantly outperforms other designed baselines achieved via the use of LoRA parameters.

Our contributions are summarized as follows.

- •
We present pairwise sample optimization (PSO) that can directly tune timestep-distilled diffusion models toward a new output distribution for various customization applications.
- •
We show that PSO is a generalized formulation that covers both offline and online preference optimization scenarios, encompassing previous works, such as DiffusionDPO (Wallace et al., ) and D3PO (Yang et al., ) as special cases, which can be flexibly adapted to various settings and tasks.
- •
Extensive of experiments demonstrate the effectiveness of PSO in human preference tuning, as well as other domain transfer tasks such as style transfer and concept customization.

## 2 Method

In this section, we present our method of pair-wise sample optimization (PSO) for directly fine-tuning timestep-distilled diffusion models. We first show the PSO formulation with both reference and data trajectories, and the extend it to both online and offline settings.

Figure: Figure 2: Demonstration of the proposed pairwise sample optimization. To tune the generative distribution $p_{\theta}$ to $p_{data}$, we sample a pair of images together with their trajectories of the same prompt, where we adopt our Markov Decision Process (MDP) formulation for the timestep-distilled diffusion model to efficiently sample the backward denoising trajectories $\{x^{\rho}_{t_{n}}\}$, while sampling $\{x^{\tau}_{t_{n}}\}$ from data via the forward diffusion process. The sampled trajectories are then sent to the final objective to move the generation trajectory aligned with the forward process from $p_{data}$.
Refer to caption: /html/2410.03190/assets/figures/main_fig.png

### 2.1 Pairwise Sample Optimization for Timestep-Distilled Diffusion Models

Consider a timestep-distilled diffusion model $p_{\theta}(x_{0}|c)$ that samples clean images $x_{0}$ conditioned on the prompt $c$ from initial noise $x_{T}\sim N(0,I)$ with very few timesteps $N,N=1\sim 4$, we denote its sampling trajectory as $\{x_{t_{N}},x_{t_{N-1}},...,x_{t_{1}},x_{t_{0}}\},\text{where }t_{N}=T,t_{0}=0$.
Our target is to fine-tune the model towards the given target distribution $p_{data}(x^{\tau}_{0},c)$, i.e., maximizing the log-likelihood $\log p_{\theta}(x^{\tau}_{0}|c)$, where $x^{\tau}_{0}\sim p_{data}$. Directly adopting the DDPM (Ho et al., ) formulation and minimizing the diffusion objective $||\epsilon_{\theta}(x^{\tau}_{t_{n}},t)-\epsilon||^{2}$ leads to blurry, degraded generation. To mitigate this issue, we introduce the reference sample $x^{\rho}_{0}$ from the current model $p_{\theta}$ and recast the optimization as maximizing the relative likelihood between the target and reference samples. Intuitively, by maximizing $\log\frac{p_{\theta}(x^{\tau}_{0}|c)}{p_{\theta}(x^{\rho}_{0}|c)}$, we can steer the generative distribution $p_{\theta}$ towards the data distribution $p_{data}$. Meanwhile, we avoid directly minimizing the diffusion objective, which can potentially help preserve its ability on few-step generation.

Formally, given a pair of positive and negative images with the same prompt $c$, $x^{\tau}_{0}\sim p_{data}(x^{\tau}_{0}|c),x^{\rho}_{0}\sim p_{\theta}(x^{\rho}_{0}|c)$, we take inspiration from direct preference optimization (Rafailov et al., ) and maximize the margin between $\log p_{\theta}(x^{\tau}_{0}|c)$ and $\log p_{\theta}(x^{\rho}_{0}|c)$, regularized by the pre-trained timestep-distilled model $\theta_{pre}$,

$$ $\mathcal{L}=-\mathbb{E}_{(x^{\tau}_{0},x^{\rho}_{0},c)}\left[\log\sigma\left(\beta\log\frac{p_{\theta}(x^{\tau}_{0}|c)}{p_{\text{pre}}(x^{\tau}_{0}|c)}-\beta\log\frac{p_{\theta}(x^{\rho}_{0}|c)}{p_{\text{pre}}(x^{\rho}_{0}|c)}\right)\right],$ (1) $$

where $\sigma$ denotes the sigmoid function, $\beta$ is the reguarlization weight.

In diffusion models, obtaining $p_{\theta}(x_{0}|c)$ requires integrating over intermediate states which is costly. Instead, we consider the whole sampling trajectory $\{x_{t_{n}}|c\}_{n=0}^{N}$, and maximize the margin between their trajectory joint likelihoods. Specifically, as shown in Figure [2](#S2.F2), we sample the data trajectory $\{x^{\tau}_{t_{0}:t_{N}}|c\}$ with diffusion forward process, and sample the reference trajectory $\{x^{\rho}_{t_{0}:t_{N}}|c\}$ with the generative reverse process,

$$ $\displaystyle p_{data}(x^{\tau}_{t_{0}:t_{N}}|c)$ $\displaystyle=p_{data}(x^{\tau}_{t_{0}}|c)\prod_{n}q(x^{\tau}_{t_{n}}|x^{\tau}_{t_{n-1}})$ (2) $\displaystyle p_{\theta}(x^{\rho}_{t_{0}:t_{N}}|c)$ $\displaystyle=p_{\theta}(x^{\rho}_{t_{N}}|c)\prod_{n}p_{\theta}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c).$ $$

Substitute the trajectory formulation into Eq. [1](#S2.E1), and we obtain our PSO objective below,

$$ $\displaystyle\mathcal{L_{\text{PSO}}}$ $\displaystyle=-\mathbb{E}_{(x^{\tau}_{0:t_{N}},x^{\rho}_{0:t_{N}},c)}\left[\log\sigma\left(\beta\sum_{n}\left(\log\frac{p_{\theta}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)}{p_{\text{pre}}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)}-\log\frac{p_{\theta}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)}{p_{\text{pre}}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)}\right)\right)\right]$ (3) $\displaystyle=-\mathbb{E}_{(x^{\tau}_{0:t_{N}},x^{\rho}_{0:t_{N}},c)}\biggl{[}\log\sigma\biggl{(}\beta\sum_{n}\biggl{(}D_{\text{KL}}[q(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n},t_{0}})\parallel p_{\text{pre}}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)]$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}-D_{\text{KL}}[q(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n},t_{0}})\parallel p_{\theta}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)]-\log\frac{p_{\theta}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)}{p_{\text{pre}}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)}\biggr{)}\biggr{)}\biggr{]}.$ $$

The derivation is provided in Appendix [A](#A1).

To expand the joint likelihood terms, we need to define the transition kernel $p_{\theta}(x_{t_{n-1}}|x_{t_{n}},c)$. Inspired by previous works (Black et al., ; Yang et al., ), we formulate the denoising sampling process of timestep-distilled models as a Markov Decision Process (MDP).

#### MDP Formulation for Timestep-Distilled Diffusion Models.

Let $\epsilon_{\theta}(x_{t_{n}},t_{n},c)$ denote the timestep-distilled diffusion model, which predicts $x_{0}$ as $f_{\theta}(x_{t_{n}},t_{n},c)=f(\epsilon_{\theta}(x_{t_{n}},t_{n},c),x_{t_{n}},t_{n})$ given the corresponding noisy latent $x_{t_{n}}$ (Song et al., ; Yin et al., ; ). In the iterative few-step sampling process, the distilled model first predicts $x_{0}$ at $t_{n}$, and then add noise back to noise-level $t_{n-1}$,

$$ $x_{t_{n-1}}=\sqrt{\bar{\alpha}_{t_{n-1}}}f_{\theta}(x_{t_{n}},t_{n},c)+\sqrt{1-\bar{\alpha}_{t_{n-1}}}z,z\sim N(0,I),$ (4) $$

where $\bar{\alpha}$’s are the forward process coefficients. The Markov Decision Process for distilled models can be then formulated as,

$$ $\displaystyle s_{n}=(x_{t_{n}},t_{n}),~{}~{}a_{n}=x_{t_{n-1}},~{}~{}P(s_{n+1}|s_{n},a_{n})=\delta(x_{t_{n-1}},t_{n-1},c)$ $\displaystyle\pi_{\theta}(a_{n}|s_{n})=N(\sqrt{\bar{\alpha}_{t_{n-1}}}f_{\theta}(x_{t_{n}},t_{n},c),1-\bar{\alpha}_{t_{n-1}})=N(\mu_{\theta}(x_{t_{n}},t_{n},c),\sigma_{t_{n}}^{2}I),$ $$

where $s_{n},a_{n}$ denote the state and action, $P(s_{n+1}|s_{n},a_{n})$ denotes the transition kernel, $\delta$ is the Dirac function, and $\pi_{\theta}$ is the policy. For the last timestep, the distilled model directly predict the clean image without noise added. This makes the final timestep a deterministic transition, $p_{\theta}(a_{N}|s_{N})=\delta(f(x_{t_{1}},t_{1},c))$, so we remove the last time step from the training. With this formulation, we have $p_{\theta}(x_{t_{n-1}}|x_{t_{n}})=\pi_{\theta}(a_{n}|x_{n})$.

#### Final PSO Loss.

By plugging the MDP action-state conditional distribution, we obtain the objective for our pairwise sample optimization,

$$ $\begin{split}\mathcal{L_{\text{PSO}}}=-\mathbb{E}\Bigg{[}&\log\sigma\Bigg{(}-\beta\cdot\sum_{n=2}^{N}\Big{(}\left(||\epsilon^{\tau}-\epsilon_{\theta}(x^{\tau}_{t_{n}},t_{n},c)||^{2}-||\epsilon^{\tau}-\epsilon_{\text{pre}}(x^{\tau}_{t_{n}},t_{n},c)||^{2}\right)\\ &-\frac{1}{2\sigma_{t_{n}}^{2}}\left(||x^{\rho}_{t_{n-1}}-\mu_{\theta}(x^{\rho}_{n},t_{n},c)||^{2}-||x^{\rho}_{t_{n-1}}-\mu_{\text{pre}}(x^{\rho}_{n},t_{n},c)||^{2}\right)\Big{)}\Bigg{)}\Bigg{]}.\end{split}$ (5) $$

The derivation is provided in Appendix [A](#A1).

### 2.2 Offline and Online Extension

Our pairwise sample optimization can be extended to online and offline settings as follows,

#### Offline PSO.

In practice, we may have pre-sampled $\{x^{\rho}_{t_{0}}\}$ from models that have similar generation distribution with $p_{\theta}(\cdot|c)$, which removes the need for reverse sampling $\{x^{\rho}_{t_{n}}\}$ from $p_{\theta}$ again. In this case, we extend our PSO formulation by approximating $p_{\theta}(x^{\rho}_{t_{1}:t_{N}}|x^{\rho}_{0},c)$ with the forward trajectory $q(x^{\rho}_{t_{1}:t_{N}}|x^{\rho}_{0})$. Specifically, we approximate the per-step transition kernel $p_{\theta}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)$ with $\int p(x^{\rho}_{t_{0}}|c)q(x^{\rho}_{t_{n-1}}|x_{t_{n}},x_{t_{0}})$, and we obtain our offline PSO objective,

$$ $\begin{split}\mathcal{L_{\text{PSO-Offline}}}=-\mathbb{E}\Bigg{[}\log\sigma&\Bigg{(}-\beta\cdot\sum_{n=2}^{N}\Big{(}||\epsilon^{\tau}-\epsilon_{\theta}(x^{\tau}_{t_{n}},t_{n},c)||^{2}-||\epsilon^{\tau}-\epsilon_{\text{pre}}(x^{\tau}_{t_{n}},t_{n},c)||^{2}\\ &-||\epsilon^{\rho}-\epsilon_{\theta}(x^{\rho}_{t_{n}},t_{n},c)||^{2}\,+\,||\epsilon^{\rho}-\epsilon_{\text{pre}}(x^{\rho}_{t_{n}},t_{n},c)||^{2}\Big{)}\Bigg{)}\Bigg{]},\end{split}$ (6) $$

which can be seen as a variant of the offline Diffusion-DPO (Wallace et al., ).

#### Online PSO.

There are also cases where we decide the data and the reference distribution in an online manner, e.g., with a reward model. In this case, we can substitute the forward trajectory in Eq. [3](#S2.E3) with the generative trajectory,

$$ $\begin{split}\mathcal{L_{\text{PSO-Online}}}=-\mathbb{E}\Bigg{[}\log\sigma&\Bigg{(}-\frac{\beta}{2\sigma_{t_{n}}^{2}}\cdot\sum_{n=2}^{N}\Big{(}||x^{\tau}_{t_{n-1}}-\mu_{\theta}(x^{\tau}_{n},t_{n},c)||^{2}-||x^{\tau}_{t_{n-1}}-\mu_{\text{pre}}(x^{\tau}_{n},t_{n},c)||^{2}\\ &-||x^{\rho}_{t_{n-1}}-\mu_{\theta}(x^{\rho}_{n},t_{n},c)||^{2}+||x^{\rho}_{t_{n-1}}-\mu_{\text{pre}}(x^{\rho}_{n},t_{n},c)||^{2}\Big{)}\Bigg{)}\Bigg{]},\end{split}$ (7) $$

which can be seen as a variant of the online DPO for diffusion models (Yang et al., ).

## 3 Experiments

In this section, we validate our method with multiple tasks and experiments. First, we demonstrate the effectiveness of our proposed PSO in the human-preference tuning task. We then present results on other general fine-tuning tasks for timestep-distilled diffusion models with PSO, including style transfer and concept customization. Our experiments cover various main-stream timestep-distillation methods, including Distribution Matching (DMD) (Yin et al., ; ), Adversarial Distillation (ADD) (Sauer et al., ), and Latent Consistency Models (LCM) (Luo et al., ). We use the distilled versions of SDXL (Podell et al., ) with these methods as our base models, i.e., SDXL-DMD2 (4 step) (Yin et al., ), SDXL-Turbo (Sauer et al., ), and SDXL-LCM (Luo et al., ).

Figure: (a) SDXL-DMD2
Refer to caption: /html/2410.03190/assets/figures/dmd_dpo_0.jpg

### 3.1 PSO for Human-Preference Fine-tuning

#### Human Preference Tuning with Offline PSO.

We first examine the effectiveness of PSO in human-preference fine-tuning in the offline setting. Specifically, we consider pre-sampled reference image samples $\{x^{\rho}_{0}|c\}$ as discussed in Sec. [2.2](#S2.SS2), which are negative samples in the human preference dataset in this case, and the data samples $\{x^{\tau}_{0}|c\}$ are the preferred images in the dataset. We train all the base timestep-distilled diffusion models with the Pick-a-Picv2 (Kirstain et al., ) dataset. After removing the pairs with tied preference labels, we end up with 851,293 pairs, with 58,960 unique prompts. We adopt the PSO-Offline objective in Eq. [6](#S2.E6) and fine-tune all our base timestep-distilled models. The training details and hyperparameters are provided in Appendix [B.1](#A2.SS1).

#### Human Preference Tuning with Online PSO.

We also validate the proposed method with the online human-preference tuning task, where we decide sampled data belonged to $p_{data}$ or $p_{\theta}$ throughout the training based on the human preference model PickScore (Kirstain et al., ). We tune all the base distilled models with the PSO objective in Eq. [7](#S2.E7) with a subset of training prompts from Pick-a-Pic v2 (Kirstain et al., ). All the other training details are provided in Appendix [B.1](#A2.SS1).

**Table 1: Human preference tuning Results of SDXL-DMD2.**
| Dataset | Method | PickScore | CLIP Score | ImageReward | Aesthetic Score |
| --- | --- | --- | --- | --- | --- |
| Pickapic Test | SDXL | 22.30 | 0.3713 | 0.8556 | 6.060 |
| SDXL-DPO | 22.60 | 0.3787 | 1.0075 | 6.040 |  |
| SDXL-DMD2 | 22.35 | 0.3679 | 0.9363 | 5.937 |  |
| SDXL-DPO + DMD2-LoRA | 22.20 | 0.3673 | 0.9287 | 5.759 |  |
| Offline-DPO w/ SDXL-DMD2 | 22.46 | 0.3690 | 0.9381 | 5.994 |  |
| Online-DPO w/ SDXL-DMD2 | 22.73 | 0.3671 | 0.9773 | 6.077 |  |
| Parti-Prompts | SDXL | 22.77 | 0.3607 | 0.9142 | 5.750 |
| SDXL-DPO | 22.92 | 0.3674 | 1.1180 | 5.795 |  |
| SDXL-DMD2 | 22.99 | 0.3607 | 1.0713 | 5.671 |  |
| SDXL-DPO + DMD2-LoRA | 22.76 | 0.3644 | 1.0638 | 5.513 |  |
| Offline-DPO w/ SDXL-DMD2 | 23.07 | 0.3649 | 1.0964 | 5.715 |  |
| Online-DPO w/ SDXL-DMD2 | 23.29 | 0.3634 | 1.1702 | 5.836 |  |

**Table 2: Human preference tuning results of SDXL-Turbo.**
| Dataset | Method | PickScore | CLIP Score | ImageReward | Aesthetic Score |
| --- | --- | --- | --- | --- | --- |
| Pickapic Test | SDXL-Turbo-4step | 22.22 | 0.3610 | 0.9300 | 5.987 |
| Offline-DPO w/ SDXL-Turbo-4step | 22.40 | 0.3634 | 0.9695 | 6.029 |  |
| Online-DPO w/ SDXL-Turbo-4step | 22.71 | 0.3647 | 0.9882 | 6.157 |  |
| SDXL-Turbo-1step | 22.29 | 0.3642 | 0.8830 | 6.061 |  |
| Offline-DPO w/ SDXL-Turbo-1step | 22.40 | 0.3663 | 0.9073 | 6.072 |  |
| Online-DPO w/ SDXL-Turbo-1step | 22.62 | 0.3661 | 0.9113 | 6.137 |  |
| Parti-Prompts | SDXL-Turbo-4-step | 22.88 | 0.3594 | 1.0173 | 5.709 |
| Offline-DPO-SDXL-Turbo-4step | 22.96 | 0.3642 | 1.0509 | 5.746 |  |
| Online-DPO-SDXL-Turbo-4step | 23.23 | 0.3632 | 1.0893 | 5.837 |  |
| SDXL-Turbo-1-step | 22.78 | 0.3596 | 0.9246 | 5.706 |  |
| Offline-DPO-SDXL-Turbo-1step | 22.86 | 0.3631 | 0.9664 | 5.755 |  |
| Online-DPO-SDXL-Turbo-1step | 22.96 | 0.3632 | 0.9762 | 5.795 |  |

Figure: Figure 4: Experiments on tuning SDXL-Turbo for style transfer with PSO. The proposed method effectively tunes the distilled model to generate images following the targeted style. Prompts from top to bottom row: A stern-faced, brown-feathered owl Pokémon with a leaf-shaped crown and piercing red eyes stands ready for a battle.// robotic cat with wings.// A cute bunny rabbit.
Refer to caption: /html/2410.03190/assets/figures/style_trans_poke.jpg

#### Evaluation and Benchmarks.

We use the Pick-a-Pic test set and PartiPrompts (Yu et al., ) as the evaluation benchmarks, where we fix the random seed for all models and generate images for all prompts in the dataset. As for the evaluation metrics, we adopt PickScore (Kirstain et al., ), CLIP score (Radford et al., ), ImageReward (Xu et al., ), and Aesthetic score following Rafailov et al. (). We benchmark distilled models fine-tuned with PSO with various benchmarks and settings. For DMD (Yin et al., ) distillation method, We compare our fine-tuned SDXL-DMD2 with un-distilled multi-step(25 steps) SDXL models, SDXL with offline human preference tuning, SDXL-DPO (Rafailov et al., ). We also benchmark with timestep-distilled models, SDXL-DMD2, SDXL with DMD2 LoRA applied, and SDXL-DPO with DMD2 LoRA applied. We adopt the similar evaluation settings for SDXL-LCM, as detailed in Appendix [B.1](#A2.SS1). For SDXL-Turbo, we fine-tune the model with 4 steps ($N=4$), and we evaluate the results on both 1-step and 4-step settings to show the timestep generalization ability of the proposed method.

#### Results.

Table [1](#S3.T1) shows the experiment results with DMD2 (Yin et al., ) as the base model.
Directly applying the DMD2 LoRA to human-preference fine-tuned SDXL-DPO (row with SDXL-DPO + DMD2 LoRA) achieves slightly inferior results compared with the SDXL-DMD2 baseline. It suffers from a drop in PickScore (drops from 22.35 to 22.20 in PickaPic-test, and 22.99 to 22.77 in Parti-Prompts) and Aesthetic Score (drops from 5.937 to 5.759 in PickaPic-test, and 5.671 to 5.513 in Parti-Prompts). The performance drop indicates the need to redo the diffusion distillation from SDXL-DPO to obtain a distilled model better aligned with human preference, which is costly ($\sim$ 4000 A100 GPU hours (Yin et al., )) for each fine-tuning. In contrast, our method achieves superior results compared with the SDXL-DMD2 baseline. For both Pick-a-Pic test set and Parti-Prompts, our Offline-PSO fine-tuned model achieves higher results in all metrics compared with baselines, showing its dominance in visual quality and human preference. For online-PSO where both the data and reference sets are updated together with the model, it achieves the best results in PickScore, ImageReward, and Aesthetic Score among all distilled models. Furthermore, its performance can even surpass the performance of multi-step fine-tuned SDXL-DPO, which further highlights the effectiveness of our method. For instance, our method achieves PickScore of 22.73 and 23.29 in two datasets, while SDXL-DPO has scores of 22.60 and 22.92, Qualitatively, as shown in Figure [3](#S3.F3), both PSO-Offline and PSO-Online fine-tuned models demonstrate better improvement in image quality in terms of prompt following, visual appeal, and the intricacy of details within each image. Moreover, Online PSO fine-tuned SDXL-DMD2 shows notable improvement over other baselines in terms of aesthetic appeal and high-frequency details, underscores the effectiveness of the proposed method.

Table [2](#S3.T2) reports the results on tuning SDXL-Turbo (Sauer et al., ). We observe a similar trend as in Table [1](#S3.T1), where both Offline-PSO and Online-PSO achieves superior results of all metrics in 4-step generation compared with the base distilled SDXL-Turbo. Moreover, our fine-tuned model also achieves superior results in 1-step generation, showcasing the generalization ability of the proposed method. We provide qualitative results of SDXL-Turbo, along with the results on SDXL-LCM in Appendix [B.1](#A2.SS1), where the proposed consistently surpass baselines in terms of visual appeal and intricate detail generation.

Figure: Figure 5: Experiments of tuning SDXL-Turbo for concept customization with PSO. The proposed method can effectively tune SDXL-Turbo to generate images that contain the given objects.
Refer to caption: /html/2410.03190/assets/figures/dreambooth_dog.png

### 3.2 PSO Fine-tuning for Style Transfer

In this section, we demonstrate that our PSO can steer the timestep-distilled models towards a specific generation style. We select the Pokemon dataset (Pinkney, ) that contains Pokemon-style image-text pairs, and we utilize the proposed PSO to have the model generate images following the Pokemon style in the dataset. Specifically, we are interested in distilled models that has no equivalent distillation LoRA, so we select SDXL-Turbo (Sauer et al., ). To simplify the experimental costs, we adopt the offline-PSO setting that we sample the reference set using the training prompts in the Pokemon dataset. Other tarining details and hyperparameters are provided in Appendix [B.2](#A2.SS2). We show the generation results of the Offline-PSO fine-tuned SDXL-Turbo in Figure [4](#S3.F4). We observe that along the fine-tuning progress, the generated images gradually shift from the original generative distribution to match the style of the given Pokemon style.

### 3.3 PSO Fine-tuning for Concept Customization

In this section, we validate the PSO fine-tuning method with concept customization with SDXL-Turbo (Sauer et al., ). Following Dreambooth (Ruiz et al., ), we aim at tuning the model to associate the special token with the given object or concept represented in few (around 5) images, and making it generalize on unseen environments. Specifically, we adopt the PSO objective in Eq. [3](#S2.E3), and we sample the reverse generative trajectories in each fine-tuning step. Training details and hyperparameters are given in Appendix [B](#A2). We show the qualitative results in Figure [5](#S3.F5), where we can see the PSO fine-tuned SDXL-Turbo generate images that contain the exact the same object shown in the training input images. This concept customization results further demonstrate the effectiveness and its generalization among tasks of PSO. More results are provided in Appendix [B.3](#A2.SS3).

## 4 Related Works

### 4.1 Diffusion Model Distillation

Various timestep distillation methods have been proposed recently to accelerate the generation speed of the diffusion models. The ability of few-step generation ability is obtained by either distilling the PF-ODE from the pre-trained multi-step diffusion model (Song et al., ; Salimans & Ho, ; Kim et al., ), or matching the denoising distribution with an explicit analytical distribution matching loss (Yin et al., ; ) or by learning an implicit discriminator (Sauer et al., ; Xu et al., ; Sauer et al., ; Lin et al., ). There are also recent works that combine those approaches to further boost the generation quality (Ren et al., ; Chadebec et al., ). However, the timestep-distilled models generally lose the ability to predict the local score, indicating that they can not be directly fine-tuned with the original score-matching loss towards a target distribution, as shown in Figure [1](#S1.F1). In this work, we tackle this problem with the incorporation of the reference samples, and formulate the pairwise sample optimization objective to directly fine-tune timestep-distilled diffusion models.

### 4.2 Preference Optimization for Diffusion Models

To control the generation behavior of the diffusion models, researchers propose the markov decision process (MDP) formulation and utilize reinforcement learning (RL) to tune the multi-step diffusion modelsBlack et al. (); Fan et al. () to generate images that have higher aesthetic values, better prompt-following ability, alignment with human preference, and diversity (Kirstain et al., ; Xu et al., ; Miao et al., ; Liu et al., ). Inspired by Rafailov et al. (), researchers have further derived direct preference optimization (DPO) from the RL formulation for diffusion models with both online (Yang et al., ; ; Liang et al., ) and offline settings (Wallace et al., ). However, all those methods are designed for tuning base multi-step diffusion models, which are not applicable to distilled few-step models. In this work, we formulate the MDP for timestep-distilled models for direct fine-tuning.

### 4.3 Diffusion Models Tuning for Style Transfer and Customization

Given the base multi-step diffusion models, a large corpus of works has focused on fine-tuning the models toward a specific style (Huang et al., ; Wang et al., ; Sohn et al., ; Zhang et al., ), and customizing the models with a given concept (Ruiz et al., ; Kumari et al., ). However, those methods are specifically designed for the corresponding task, and are general limited for multi-step base diffusion models. In our work, we target at directly tuning timestep-distilled diffusion models with a unified objective that can be also applied to human-preference tuning.

## 5 Conclusion

In this work, we presented pairwise sample optimization (PSO) for directly fine-tuning timestep-distilled diffusion models. By constructing pairs of data and reference images and utilizing preference optimization to increase the likelihood margin between them, PSO enables manipulation of the model’s generative distribution while preserving its few-step generation capability. We re-formulated preference optimization for timestep-distilled models and demonstrated its effectiveness in aligning model outputs with human preferences using both offline and online-generated pairwise preference data. Furthermore, we extended PSO to style transfer and concept customization tasks by constructing pairwise source-target image pairs. Our approach provides an efficient solution for adapting timestep-distilled diffusion models to downstream applications without the need for costly re-distillation. Our work can be further strengthened with enhanced fine-tuning results in personalized generation, and including experiments in controllable image generation.

## Ethical Statement

Our research on fine-tuning timestep-distilled diffusion models with the proposed PSO generally does not has ethical concerns, except the usage of human preference data. We acknowledge the potential for introducing or amplifying biases present in human judgments, which could lead to unfair or discriminatory outputs. To address this, we commit to implementing diverse sampling strategies for both annotators and image datasets, and developing fairness metrics to assess and mitigate potential biases. We also recognize the importance of transparency in our fine-tuning process and the need to protect data privacy. Additionally, we are mindful of the potential for misuse of these models to generate misleading content and the environmental impact of computational resources required. We are committed to ongoing evaluation and improvement of our methods to ensure responsible development and deployment of this technology, including establishing guidelines for ethical use and exploring ways to enhance interpretability and minimize environmental impact.

## Appendix

## Appendix A Derivations

### A.1 PSO Derivation

In this section, we present the derivation for Eq. [3](#S2.E3) in Sec. [2.1](#S2.SS1). Given the forward sampled data trajectory $x^{\tau}_{t_{0}:t_{N}}|c$ and reverse sampled reference trajectory $x^{\rho}_{t_{0}:t_{N}}|c$, we represent our PSO objective as increasing the likelihood gap between two trajectories,

$$ $\mathcal{L}=-\mathbb{E}_{(x^{\tau}_{t_{0}:t_{N}},x^{\rho}_{t_{0}:t_{N}},c)}\left[\log\sigma\left(\beta\log\frac{p_{\theta}(x^{\tau}_{t_{0}:t_{N}}|c)}{p_{\text{pre}}(x^{\tau}_{t_{0}:t_{N}}|c)}-\beta\log\frac{p_{\theta}(x^{\rho}_{t_{0}:t_{N}}|c)}{p_{\text{pre}}(x^{\rho}_{t_{0}:t_{N}}|c)}\right)\right],$ (8) $$

With the formulation in Eq. [2](#S2.E2), we can write the joint likelihood in the above equation as,

$$ $\begin{split}\mathcal{L}=-&\mathbb{E}_{(x^{\tau}_{t_{0}:t_{N}}\sim p_{data}(x^{\tau}_{t_{0}:t_{N}}),x^{\rho}_{t_{0}:t_{N}}\sim p_{\theta}(x^{\rho}_{t_{0}:t_{N}}|c)}\bigg{[}\log\sigma\bigg{(}\\ &\beta\sum_{n=1}^{N}\Big{(}\log\frac{p_{\theta}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)}{p_{\text{pre}}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)}-\log\frac{p_{\theta}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)}{p_{\text{pre}}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)}\Big{)}\Bigg{)}\Bigg{]}.\end{split}$ (9) $$

For the first log-likelihood ratio in the bracket, it can be further derived as,

$$ $\displaystyle L_{1}$ $\displaystyle=-\mathbb{E}_{x^{\tau}_{t_{0}:t_{N}}\sim p_{data}(x^{\tau}_{t_{0}:t_{N}})}\Big{[}\log\sigma(\beta\sum_{n=1}^{N}\log p_{\theta}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)-\beta\sum_{n=1}^{N}\log p_{\text{pre}}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c))\Big{]}$ (10) $\displaystyle=-\mathbb{E}_{x^{\tau}_{t_{0}:t_{N}}\sim p_{data}(x^{\tau}_{t_{0}})q(x_{t_{1}:t_{N}}|x_{t_{0}})}$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}\Big{[}\log\sigma(\beta\sum_{n=1}^{N}\log p_{\theta}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)-\beta\sum_{n=1}^{N}\log p_{\text{pre}}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c))\Big{]}$ $$

Given that each sampling step in the forward process can be written as $x_{t_{n}}\sim q(x^{\tau}_{t_{n}}|x^{\tau}_{t_{n-1}},x^{\tau}_{0})$ and we know the marginal $q(x_{t_{n}}|x_{0})$, we can sample $x^{\tau}_{t_{n-1}}\sim q(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},x_{0})$ given $x^{\tau}_{t_{n}},x^{\tau}_{0}$ with the Bayes equation. In this way, the sub-objective above can be further derived as,

$$ $\displaystyle L_{1}$ $\displaystyle=-\mathbb{E}_{q}\Big{[}\log\sigma(\beta\sum_{n=1}^{N}\log\frac{p_{\theta}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)}{q(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},x^{\tau}_{0})}-\beta\sum_{n=1}^{N}\log\frac{p_{\text{pre}}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)}{q(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},x^{\tau}_{0})})\Big{]}$ (11) $\displaystyle=-\mathbb{E}_{q}\Big{[}\log\sigma\Big{(}\beta(D_{\text{KL}}[q(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n},t_{0}})\parallel p_{\text{pre}}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)]-D_{\text{KL}}[q(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n},t_{0}})\parallel p_{\theta}(x^{\tau}_{t_{n-1}}|x^{\tau}_{t_{n}},c)])\Big{)}\Bigg{]}$ $\displaystyle=-\mathbb{E}_{q}\Big{[}\log\sigma\Big{(}-\beta\omega_{t_{n}}\left(||\epsilon^{\tau}-\epsilon_{\theta}(x^{\tau}_{t_{n}},t_{n},c)||^{2}-||\epsilon^{\tau}-\epsilon_{pre}(x^{\tau}_{t_{n}},t_{n},c)||^{2}\right)\Big{)}\Bigg{]},$ $$

where the third equation adopts the reparameterization in the original diffusion objective derivation (Black et al., ), and $\omega_{t_{n}}$ is the reweighting which is usually set to constant in practice.

Meanwhile, the second log-likelihood ratio in Eq. [9](#A1.E9), can be further simplified with Eq. [2](#S2.E2) as,

$$ $\log\frac{p_{\theta}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},c)}{p_{\text{pre}}(x^{\rho}_{t_{n-1}}|x^{\rho}_{t_{n}},x^{\rho}_{0})}=-\frac{1}{2\sigma_{t_{n}}^{2}}(\lVert x^{\rho}_{t_{n-1}}-\mu_{\theta}(x^{\rho}_{t_{n}},t_{n},c)\rVert_{2}^{2}-\lVert x^{\rho}_{t_{n-1}}-\mu_{\text{pre}}(x^{\rho}_{t_{n}},t_{n},c)\rVert_{2}^{2})$ (12) $$

By substituting Eq. [11](#A1.E11) and Eq. [12](#A1.E12) into Eq. [9](#A1.E9) and removing the last sampling step $t_{1}$, we got the PSO objective as in Eq. [5](#S2.E5)

$$ $\begin{split}\mathcal{L_{\text{PSO}}}=-\mathbb{E}\Bigg{[}&\log\sigma\Bigg{(}-\beta\cdot\sum_{n=2}^{N}\Big{(}\left(||\epsilon^{\tau}-\epsilon_{\theta}(x^{\tau}_{t_{n}},t_{n},c)||^{2}-||\epsilon^{\tau}-\epsilon_{pre}(x^{\tau}_{t_{n}},t_{n},c)||^{2}\right)\\ &-\frac{1}{2\sigma_{t_{n}}^{2}}\left(||x^{\rho}_{t_{n-1}}-\mu_{\theta}(x^{\rho}_{n},t_{n},c)||^{2}-||x^{\rho}_{t_{n-1}}-\mu_{pre}(x^{\rho}_{n},t_{n},c)||^{2}\right)\Big{)}\Bigg{)}\Bigg{]}\end{split}$ (13) $$

### A.2 MDP formulation for SDXL-Turbo

The MDP formulation presented in Sec. [2.1](#S2.SS1) is applicable to DMD (Yin et al., ) and LCM (Luo et al., ). Here we present the MDP formulation for SDXL-Turbo (Sauer et al., ). The Euler ancestral sampler can be defined as,

$$ $x_{\sigma_{n-1}}=x_{\sigma_{n}}+s_{\theta}(x_{\sigma_{n}};\sigma_{n})\cdot(\bar{\sigma}_{n}-\sigma_{n})+\hat{\sigma}_{n}\cdot z,~{}~{}z\sim N(0,I),$ (14) $$

where $\bar{\sigma}_{n}=\sqrt{(\sigma_{n}^{2}-\sigma_{n-1}^{2})\frac{\sigma_{n-1}^{2}}{\sigma_{n}^{2}}},\hat{\sigma}_{n}=\sqrt{\sigma_{n-1}^{2}-\sigma_{n}^{2}},s_{\theta}(x,\sigma_{n})=(x_{\sigma_{n}}-f_{\theta}(x_{\sigma_{n}},\sigma_{n},c))/\sigma_{n}$, and $\sigma_{n}$ is the designated noise levels in Euler ancestral sampler used in SDXL-Turbo which has one-to-one mapping to $t_{n}$.

The MDP formulation can then be represented as,
The Markov Decision Process for distilled models can be then formulated as,

$$ $\displaystyle s_{n}=(x_{\sigma_{n}},\sigma_{n}),~{}~{}a_{n}=x_{\sigma_{n-1}},~{}~{}P(s_{n+1}|s_{n},a_{n})=\delta(x_{\sigma_{n-1}},\sigma_{n-1},c)$ $\displaystyle\pi_{\theta}(a_{n}|s_{n})=N(x_{\sigma_{n}}+s_{\theta}(x_{\sigma_{n}},\sigma_{n})(\bar{\sigma}_{n}-\sigma_{n}),\hat{\sigma}_{n})^{2}I),$ $$

## Appendix B Supplementary Experimental Settings and Results

### B.1 Human Preference Tuning

#### Datasets.

We adopt the Pick-a-Pic v2 (Kirstain et al., ) dataset for the offline human preference tuning task following Wallace et al. (). It consists of pairwise preferences for images generated by SDXL-beta and Dreamlike, a fine-tuned version of SD1.5 (Rombach et al., ). The prompts and preferences were collected from users of the Pick-a-Pic web application. After removing $\sim 12\%$ pairs with tied preference, we obtain $\sim 850K$ win-lose pairs, which are used as data-reference pairs within our Offline-PSO objective as in Eq. [6](#S2.E6). For the online human preference setting, we use a subset of 4K prompts from Pick-a-Pic training prompts as the training prompts, and we sample a pair of images given a prompt, assign the target & reference labels based on the PickScore (Kirstain et al., ) as a discriminator.

**Table 3: Results of SDXL-LCM**
| Dataset | Method | PickScore | CLIP Score | ImageReward | Aesthetic Score |
| --- | --- | --- | --- | --- | --- |
|  | SDXL-LCM | 21.9072 | 0.3572 | 0.6717 | 5.8642 |
|  | SDXL + LCM-LoRA | 21.8520 | 0.3553 | 0.6840 | 5.9447 |
|  | SDXL-DPO + LCM-LoRA | 22.3178 | 0.3633 | 0.7883 | 5.9286 |
|  | Offline-DPO w/ SDXL-LCM | 22.3953 | 0.3668 | 0.8129 | 5.9949 |
|  | Online-DPO w/ SDXL-LCM | 22.6353 | 0.3613 | 0.8343 | 6.0117 |
| Parti-Prompts | SDXL-LCM | 22.4761 | 0.3492 | 0.6738 | 5.5867 |
| SDXL + LCM-LoRA | 22.4332 | 0.3481 | 0.6571 | 5.5141 |  |
| SDXL-DPO + LCM-LoRA | 22.7854 | 0.3551 | 0.7778 | 5.7090 |  |
| Offline-DPO w/ SDXL-LCM | 22.8131 | 0.3592 | 0.8131 | 5.8367 |  |
| Online-DPO w/ SDXL-LCM | 22.9451 | 0.3581 | 0.8351 | 5.8572 |  |

#### Parameter Efficient Fine-tuning.

As we aim at designing a light-weight method for fine-tuning distilled diffusion models, it’s natural for us to adopt the parameter efficient fine-tuning methods which majorly are based on matrix low-rank decomposition. LoRA (Hu et al., ) and following works (Liu et al., ; Dettmers et al., ) decompose the weight matrix along the channel dimension into two low rank matrices, which induces much lower costs in fine-tuning compared with tuning the original full-rank weight. Meanwhile, (Chen et al., ) proposes the filter subspace decomposition for weight matrices. The filter subspace decompositon method (Qiu et al., ) has shown effectiveness in continual learning Miao et al. (); Chen et al. (), video representation learning Miao et al. (), domain adaptation (Wang et al., ), and generative tasks (Wang et al., ; ).
There are some other works on fine-tuning the SVD decomposition of weights(Han et al., ) or fine-tunig bias parameters(Xie et al., ).
Considering the implementation and the integration with other codebase, we choose LoRA (Hu et al., ) to fine-tune timestep-distilled diffusion models.

#### Training Details.

We use LoRA (Hu et al., ) to fine-tune all the distilled diffusion models efficiently, we set LoRA rank $r=16$ for SDXL-DMD2 and SDXL-Turbo, and $r=32$ for SDXL-LCM in both online and offline human preference tuning experiments. As for other training hyperparameters, we set the number of training distilled steps $N=4$, which is the same as the number of sampling steps of these distilled models, and we set the regularization weight $\beta=50$ for Offline preference tuning, and $\beta=5$ for Online preference tuning. We set the batch size to 64 and the learning rate to $1e-5$ for offline experiments, and train for 5k steps. For online preference tuning, we first sample 128 pairs of images with 128 training prompts, which are further labeled as reference and target with PickScore (Kirstain et al., ) with a batch size of 64, and then we conduct online PSO on the sampled pairs for 1 epoch with a batch size of 32 and learning rate $1e-5$ and train for 15k steps.

#### Evaluation.

For all baseline distilled models and our PSO fine-tuned models, we conduct testing on Pickapic (Kirstain et al., ) test and PartiPrompts (Yu et al., ). They contain 500 and 1632 prompts respectively. We sample for each prompt a testing image with 4 sampling steps (we also conduct a 1-step sampling test for SDXL-Turbo, which is trained under a 4-step schedule). Each sampled image is then evaluated by PickScore (Kirstain et al., ), CLIP Score (Radford et al., ), ImageReward (Xu et al., ), and Aesthetic Score (Schuhmann et al., ). We report the average scores in the tables.

Figure: Figure 6: Human Preference Tuning with SDXL-DMD2 addidional results, Prompts from left to right: Sign that says DON’T FART// Warrior with an axe on a exotic animal with long hair style loisel// Movie Still of The Joker wielding a red Lightsaber, Darth Joker a sinister evil clown prince of crime, HD Photograph// A swirling, multicolored portal emerges from the depths of an ocean of coffee, with waves of the rich liquid gently rippling outward. The portal engulfs a coffee cup, which serves as a gateway to a fantastical dimension. The surrounding digital art landscape reflects the colors of the portal, creating an alluring scene of endless possibilities.
Refer to caption: /html/2410.03190/assets/figures/dpo_dmd_app0.jpg

Figure: Figure 7: Human Preference Tuning with SDXL-DMD2 addidional results. Prompts from left to right: a red hair girl look at you, distant view.// A dog with a sign that reads ”Hello”.// anime pencil concept style.// A galaxy-colored figurine floating over the sea at sunset, photorealistic.
Refer to caption: /html/2410.03190/assets/figures/dpo_dmd_app1.jpg

#### Additional Results.

We provide more qualitative results for SDXL-DMD2 human preference tuning in Figure [6](#A2.F6) & [7](#A2.F7). As shown in those figures, images generated by our Online-PSO and Offline-PSO tuned models demonstrate better aesthetic values, prompt following, and detailed accuracy. For SDXL-Turbo, we provide qualitative illustrations in Figure [8(a)](#A2.F8.sf1) for 1-step evaluation and Figure [8(b)](#A2.F8.sf2) for 4-step evaluation. Again, we see that our PSO-tuned models generate more visual-appeal images in both 1-step and 4-step generations. Specifically, we observe better visual detail generation, for instance, human body, objects, and object consistency in images generated by our method. And finally, we provide quantitative results for SDXL-LCM (Luo et al., ) in Table [3](#A2.T3). As shown in the table, our method achieves higher results in all metrics compared with baselines. Specifically, our method achieves better results compared with SDXL-DPO + LCM-LoRA, which is the preference-tuned multi-step SDXL equipped with LCM LoRA, which further justify and validate our method.

Figure: (a) Preference Fine-tuning Results with PSO on 1step SDXL-Turbo.
Refer to caption: /html/2410.03190/assets/figures/dpo_turbo_app_1step.jpg

### B.2 PSO for Style Transfer

We adopt SDXL-Turbo (Sauer et al., ) for this experiment with the Pokemon dataset (Pinkney, ). We first construct a reference set by sampling images from SDXL-Turbo using the training prompts in the pokemon dataset. We add LoRA of rank $r=16$ to the attention layers and set $\beta=50$. Then, we adopt the objective of offline-PSO for fine-tuning. We use batch size of 8, learning rate of 1e-5, and train for 3k steps.

### B.3 PSO for Concept Customization

We also use SDXL-Turbo for this task with the Dreambooth dataset (Ruiz et al., ). We adopt the PSO objective in Eq. [3](#S2.E3) in the way that we use the given input images as the targeted set, and sample reference images from SDXL-Turbo with the prompt ’A photo of [V] concept’. We set lora rank $r=16$, $\beta=5$, learning rate of 5e-5 and train for 1000 steps with batch size of 4.

We provide additional results in Figure [9](#A2.F9). Our PSO can effectively customize the distilled model with the given animal or object.

Figure: Figure 9: Experiments of Tuning SDXL-Turbo for personalized generation with PSO. The proposed method can effectively tune SDXL-Turbo to generate images that contain the given objects.
Refer to caption: /html/2410.03190/assets/figures/dreambooth_cat.png

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 6
presentation: 2
questions: Please refer to Weakness
soundness: 3
strengths: 1. This paper addresses a very important and practical issue in the diffusion-model field.
2. The proposed pairwise sample optimization (PSO) algorithm is novel and has thereotical derivation to increase its rigorism。
3. This paper considers both offline and online setting, whichi can unifies various prior works.
4. Experimental studies are conducted on different acceleration methods and applied on several downstream finetune tasks, which make this PSO more persuasive.
summary: This paper proposes a new and important issue about diffusion models, namely direct finetuning a timestep-distilled diffusion model without apparent performance degradation. Since finetuning a timestep-distilled diffusion model with a naive diffusion loss will result in degraded generated images, this paper proposes a so-called pairwise sample optimization algorithm that increases a relative likelihood margin between the training images and reference images. The authors conduct several experiments on different acceleration methods, including DMD2, Turbo, and LCM on different finetune tasks, i.e., humen preference tuning, style transfer, and customized image generation. Experimental results can demonstrate the effectiveness of the PSO
weaknesses: 1. The analysis and description of the proposed PSO are slightly incomplete. For instance:
- This paper doesn't explain why "directly adopting the DDPM  formulation and minimizing the diffusion objective $||\epsilon\theta(x_{t_n}^\tau) − ε||^2$ leads to blurry."
- Why you choose to "recast the optimization as maximizing the relative likelihood between the target and reference samples?'' I think this is the most crucial insight of this paper.
- How to chooose the "pre-sampled reference image samples''?
2. The experimental results in Table 1 are a little confusing. Why do you choose "SDXL-DPO + DMD2-LoRA'' as a compared baseline? Use DMD2-LoRA based on a SDXL-DPO model? Why not compare with "SDXL-DMD2-LoRA  + DPO'' by using the distilled SDXL-DMD2 as base model?
3. Some typos: Figure 1(b) -> Figure 1(c) in Line 107; the extend it to -> then extend it to in Line 148; no explanation for $p_{pre}$ in Eq (1) (maybe $p_{\theta_{pre}}$)

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Why is DPO used in all the tables instead of your method, PSO?

In Table 1, why are your results bolded even when they perform worse than SDXL-DPO?

How should we interpret that the online update version shows improvements in almost all metrics but generally performs worse than the offline version on the CLIP score?
soundness: 2
strengths: Fine-tuning time-distilled diffusion models is very useful.

The proposed method generalizes well across different tasks and base models.
summary: This paper focuses on fine-tuning distilled diffusion models. The authors introduce a pairwise optimization approach, which combines preference optimization over positive-negative pairs and alignment with a frozen distilled few-step diffusion model. They present both offline and online training versions of this loss function. Experiments on tasks of human-preference tuning, style transfer, and concept customization show it outperforms the original DPO.
weaknesses: Although the experiments cover many tasks and base models, there aren’t many comparisons with other methods. The authors only compare their method to the original DPO in Table 1 and Table 3, and there are no comparisons in Table 2. Including comparisons with more fine-tuning methods would make the results more comprehensive and highlight the advantages of the proposed approach.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: This paper is good to me. I hope author can provide an algorithm for complete method for better readability.
soundness: 4
strengths: - This problem is both novel and relevant, as many time-distilled diffusion models incorporate multiple objectives, such as adversarial training, making efficient fine-tuning essential.

- The proposed method’s use of pair optimization is innovative, providing a critical solution to eliminate the need for the cumbersome teacher model distillation process.

- Extensive experiments underscore the effectiveness of this approach.
summary: The motivation behind this work stems from the observation that models distilled over multiple timesteps can suffer from blurriness and reduced quality if distilled further using a naive diffusion objective. A straightforward solution would be to use a teacher model and repeat the distillation process, but this approach is overly cumbersome.
To address this, the authors propose an optimization method that fine-tunes the distilled model without requiring a teacher model. PSO introduces additional self-generated images and optimizes the relative likelihood margin, preserving the model's ability to generate high-quality images in fewer steps while still allowing fine-tuning.
weaknesses: - The notation used in the paper lacks clarity and can be confusing for readers, particularly with symbols like ρ and p.

- Additionally, Figure 2 is somewhat abstract. Including an algorithm for the complete method would improve readability and help readers understand the approach more efficiently.

- it would be beneficial to discuss how this method compares to current DiT-like image generation models, such as Flux.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - Why was offline PSO used for style transfer and online PSO for concept customization? Could the reverse approach also work?

- In Tables 1 and 2, is there a specific reason why "Offline-DPO" and "Online-DPO" are listed instead of "Offline-PSO" and "Online-PSO"?

- What reward model was used for online PSO?
soundness: 3
strengths: - This paper addresses the important problem of finetuning timestep-distilled diffusion models.

- PSO, while similar to DPO, introduces key differences, notably the ability to support both offline and online training.

- Finetuning timestep-distilled models with PSO proves more effective than attaching timestep-distilled LoRA to finetuned diffusion models, and the experiments demonstrate the potential for style transfer and customization.
summary: This paper proposes PSO, an objective for finetuning timestep-distilled diffusion models like SDXL-DMD2, SDXL-Turbo, and SDXL-LCM. These models, trained with GAN loss or consistency distillation loss, can no longer be trained with traditional diffusion loss. The proposed PSO is a modification of DPO, which is widely used in human preference optimization for LLMs and diffusion models. Unlike DPO, PSO uses negative samples generated by timestep-distilled models and supports both offline and online generation. PSO shows effectiveness in human preference tuning, style transfer, and concept customization for these models.
weaknesses: - The paper primarily demonstrates the effectiveness of PSO in human preference tuning. However, direct finetuning of timestep-distilled models holds more significance in concept customization rather than human preference tuning. As the authors mention in the introduction, customization methods like Dreambooth require significantly less time compared to timestep distillation methods like DMD. On the other hand, human preference tuning, such as DPO, still requires lengthy training and may not offer much advantage over timestep distillation. To further highlight the necessity of PSO in human preference tuning, it would be beneficial to compare the computational cost of timestep distillation versus human preference tuning.

- For the reasons mentioned above, PSO should further showcase its capabilities in concept customization and style transfer on datasets from Dreambooth (Ruiz et al.) and StyleDrop (Sohn et al). Specifically, following these works, text-image alignment could be compared using CLIP, and image-reference similarity could be assessed using DINO. Additionally, results in Figures 5 and 9 appear blurry, so a quantitative comparison using CLIP and DINO is essential.

- PSO uses generated negative samples, which is similar to SPIN [A]. A discussion on the differences between PSO and SPIN would be beneficial. A discussion on the differences between PSO and SPIN could be added to the related works section.

[A] Chen et al., Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models, ICML 2024.

Meta-review: All reviewers agree to accept the paper. Reviewers appreciate the novel task, the innovative method, and extensive experiments. Please be sure to address the reviewers' comments in the final version.

Decision: Accept (Poster)