Pixel-Space Post-Training of Latent-Diffusion Models

## Abstract

Abstract Latent diffusion models (LDMs) have made significant advancements in the field of image generation in recent years. One major advantage of LDMs is their ability to operate in a compressed latent space, allowing for more efficient training and deployment. However, despite these advantages, challenges with LDMs still remain. For example, it has been observed that LDMs often generate high-frequency details and complex compositions imperfectly.
We hypothesize that one reason for these flaws is due to the fact that all pre- and post-training of LDMs are done in latent space, which is typically 8 × 8 8 8 8\times 8 lower spatial-resolution than the output images. To address this issue, we propose adding pixel-space supervision in the post-training process to better preserve high-frequency details. Experimentally, we show that adding a pixel-space objective significantly improves both supervised quality fine-tuning and preference-based post-training by a large margin on a state-of-the-art DiT transformer and U-Net diffusion models
in both visual quality and visual flaw metrics, while maintaining the same text alignment quality.

## 1 Introduction

Diffusion models learn to sequentially denoise from random Gaussian noise to sharp images and have revolutionized the field of media generation and editing in recent years.
Latent diffusion models represent the most popular type of diffusion model because of their efficiency and simplicity.
State-of-the-art LDMs are typically pretrained on webscale data, resulting in “foundation models” (Rombach et al., ; Podell et al., ; Esser et al., ; Saharia et al., ; Imagen 3 Team, ; Dai et al., ; Ramesh et al., ; ; Betker et al., ).

These foundation models are then post-trained on a smaller, carefully curated dataset to improve quality through either supervised quality fine-tuning (SFT) (Dai et al., ) or human-in-the-loop preference modeling (Rafailov et al., ; Wallace et al., ; Meng et al., ).
Post-training of image foundation models is also utilized to create new models for a variety of applications, including controllable generation (Zhang et al., ), editing (Sheynin et al., ), 3D generation (Poole et al., ), video generation (Singer et al., ; Girdhar et al., ), and many others.

To achieve efficiency and simplicity, LDMs use a pretrained variational autoencoder (VAE) to compress images into latent representations. For example, in the original LDM paper (Rombach et al., ), the authors used a conv-based VAE to compress a $512\times 512\times 3$ image to $64\times 64\times 4$. This representation significantly speeds up training and reduces computational cost as the denoising diffusion model now operates in the $64\times 64\times 4$ space instead of the original $512\times 512\times 3$ space ($48\times$ compression). However, this comes at the cost of lossy compression, which can result in inaccuracies in or loss of high-frequency details.

The research community has invested considerable effort in improving fine details, including scaling up the model, carefully curating fine-tuning datasets, increasing the latent channel dimension (Dai et al., ), and designing more powerful decoders (Betker et al., ).

In this paper, we take a step back and hypothesize that the artifacts in LDMs are partially caused by the fact that all pretraining, post-training, and inference steps are done on a lower-resolution latent space. With this assumption, we propose adding an additional supervision term in the original pixel space during post-training by decoding the latent representation back and combining it with the original latent loss term. This approach aims to improve the quality of generated images by providing additional guidance in the pixel space, which helps to mitigate the loss of high-frequency details and artifacts introduced by the compression of the latent space. Figure [1](#S1.F1) demonstrates that our method can generate more stunning details when post-trained on the same dataset.

Figure: Figure 1: Enhancing LDMs with pixel-space objectives. We hypothesize that losses of details and artifacts in high-frequency details are partially caused by training on the lower-resolution latent space. We propose adding a pixel-space objective during LDM post-training. Our experiments show significant improvements in both DiT-based and UNet-based LDMs for reward-based and supervised fine-tuning.
Refer to caption: /html/2409.17565/assets/figures/teaser_jw.png

Through extensive experiments and independent human evaluation from annotators who have no knowledge of this project, we have found that the proposed method has the following advantages:

- 1.
Simplicity: The proposed method does not modify the architecture of the diffusion denoising model and can be seamlessly integrated into any LDM-based model without introducing new parameters, making it flexible and efficient.
- 2.
Effectiveness: Despite its simplicity, we found that the proposed method is surprisingly effective, resulting in a $18.2\%$ and $23.5\%$ improvements on visual appeal and visual flaws with supervised fine-tuning, and $17.8\%$ and $11.3\%$ improvements on preference-based fine-tuning on a DiT model on head-to-head A/B comparisons with the latent-space baseline.
- 3.
General applicability to models: The proposed method performs remarkably well in both DiT and U-Net based LDMs (Rombach et al., ; Dai et al., ).
- 4.
General applicability to post-training methods: The proposed method works well on both supervised fine-tuning and reward-based fine-tuning, and can be easily added to the future post-training methods researchers develop.

A secondary contribution of this paper is that we are the first paper that extends the recently proposed SimPO (Meng et al., ) preference optimization post-training technique to the image generation task and shows its effectiveness on the diffusion-based image generation domain.

## 2 Related Work

A comprehensive review of diffusion models is out of the scope of this section. Interested readers are referred to Fuest et al. () and Chan (). Here we highlight a few works that are closest-related to us.

### 2.1 Text-to-image diffusion model

Researchers have explored a variety of representations to train text-to-image diffusion models, including pixel-diffusion models (Ramesh et al., ; Saharia et al., ; Balaji et al., ), latent diffusion models (Rombach et al., ; Dai et al., ), and token-based generative transformers (Chang et al., ; Sun et al., ; Li et al., ). Pixel-diffusion models directly generate images in the pixel space, but due to computational constraints, they typically first generate images at a lower resolution (e.g., $64\times 64$) and then upsample them (sometimes multiple times) to achieve the target resolution in a cascade fashion (Saharia et al., ).

Latent Diffusion Models (LDMs), on the other hand, employ a pretrained autoencoder (Rombach et al., ) to compress the spatial dimensions of the image to be generated, typically by a factor of $8\times 8$, while moderately increasing the channel dimension from 3 (RGB) to 4. This approach significantly enhances training efficiency compared to pixel diffusion models, thereby facilitating various applications such as high-resolution (Chen et al., ) and real-time image generation (Kohler et al., ; Wimbauer et al., ). Early LDM models use convolutional U-Nets as the backbone diffusion model, such as LDM1.5 (Rombach et al., ) and Emu (Dai et al., ). Recently, the field has been dominated by diffusion transformers (DiTs), such as SD3 (Esser et al., ) and PixArt-$\alpha$ (Chen et al., ). PixArt-$\alpha$ incorporates cross-attention modules into DiT and trained the model on high-aesthetic data in its final training stage. However, all of these LDM methods are still trained in latent space, which might suffer from loss of details and artifacts due to low spatial resolution.

In this paper, we propose a novel approach to refining image quality in diffusion models by deploying a pixel-space objective function in the post-training stage. Our method does not depend on a particular diffusion model type and works equally well for both U-Nets and DiTs.

### 2.2 Supervised quality fine-tuning (SFT)

Supervised fine-tuning is crucial to the success of modern LLMs (Zhou et al., ; Touvron et al., ; Achiam et al., ). In image and vision, Dai et al. () proposes using a small set of extremely high-quality images to fine-tune a pretrained LDM model, resulting in significant improvements to the visual quality of generated images without sacrificing text-image alignment. Betker et al. () and Segalis et al. () propose using captions rewritten by vision language models to facilitate better learning, including during SFT.
However, none of these proposed methods explored the representation space in which the model was fine-tuned. In this paper, we propose supplementing the regular supervised fine-tuning loss with a pixel-space objective function. We experimented with two different models: a replacement-trained U-Net LDM-1.5 (Rombach et al., ) and a DiT model. Our results show that when fine-tuning on a small high-quality dataset, our proposed method can significantly improve generation quality and visual flaws.

### 2.3 Human Preference Based Post Training

Reinforcement learning represents another popular type of post-training technique. The seminal work of Schulman et al. () makes the policy gradient method practical.  Rafailov et al. () proposes doing direct preference optimization (DPO) with a reference model to improve the model quality on language models. Wallace et al. () and Black et al. () extend DPO to diffusion models. DPO optimizes diffusion models on paired human preference data by implicitly estimating a reward model.
Liang et al. () proposes doing step-aware DPO. Meng et al. (), on the other hand, remove the reference model to make reinforcement learning more direct and effective.
In this paper, we show that our proposed pixel-space objective also works well for reward-based post-training.

## 3 Method

Given an image $x\in\mathbb{R}^{H\times W\times 3}$ in RGB space, LDMs use an autoencoder $\mathcal{E}$ that encodes $x$ into a latent representation $z=\mathcal{E}(x)$. The decoder $\mathcal{D}$ then reconstructs the image from the latent, giving $\tilde{x}=\mathcal{D}(z)=\mathcal{D}(\mathcal{E}(x))$, where $z\in\mathbb{R}^{h\times w\times c}$.

### 3.1 Supervised Pixel-space Fine-tuning

Figure: Figure 2: Supervised fine-tuning with pixel-space loss. During fine-tuning, we decode the latent representation back to the pixel space and add a supervision in the output image resolution.
Refer to caption: /html/2409.17565/assets/figures/training_diagram.png

By denoising a normally distributed variable step-by-step, LDMs learn a data distribution $p_{\theta}(x)$. Therefore, they can be understood as a series of denoising autoencoders $\epsilon_{\theta}(z_{t},t)$; $t=1\dots T$ that are trained to predict the denoised variant of their input $z_{t}$ where $z_{t}$ is the noisy version of latent input $z$ at time $t$, $\epsilon$ is the original noise added to get $z_{t}$, and $\epsilon_{\theta}(z_{t},t)$ is the predicted noise. Furthermore, the noise added to $z_{t-1}$ to get $z_{t}$ is Gaussian with variance $\beta_{t}$. The standard objective function for LDMs is:

$$ $L_{\textit{SFT}}^{\textit{latent}}:=\mathbb{E}_{\mathcal{E}(x),\epsilon\sim\mathcal{N}(0,1),t}\left[\|\epsilon-\epsilon_{\theta}(z_{t},t)\|_{2}^{2}\right].$ (1) $$

Instead of working only in the latent space $\mathbb{R}^{h\times w\times c}$, we propose a loss function that incorporates the pixel space $\mathbb{R}^{H\times W\times 3}$ in the objective function. This is achieved by adding the noise $\epsilon$ to the latent image $z$ through the forward diffusion process $x_{t}=\sqrt{\bar{\alpha}_{t}}x_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon$, where $\alpha_{t}=1-\beta_{t}$ and $\bar{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{i}$, and decoding it back to the pixel space. The objective function then becomes

$$ $L_{\textit{SFT}}^{\textit{pixel}}:=\mathbb{E}_{\mathcal{E}(x),\epsilon\sim\mathcal{N}(0,1),t}\left[\|\mathcal{D}(\sqrt{\bar{\alpha}_{t}}z+\sqrt{1-\bar{\alpha}_{t}}\epsilon)-\mathcal{D}(\sqrt{\bar{\alpha}_{t}}z+\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}(z_{t},t))\|_{2}^{2}\right].$ (2) $$

We combine the latent objective, Equation [1](#S3.E1), with the pixel objective, Equation [2](#S3.E2), to obtain an objective that uses both the latent and pixel space, weighted by hyper-parameter $\lambda$.

$$ $L_{\textit{SFT}}:=L_{\textit{SFT}}^{\textit{latent}}+\lambda L_{\textit{SFT}}^{\textit{pixel}}.$ (3) $$

### 3.2 Pixel-space Fine-tuning using Reward Modeling

Define $x^{w}$ and $x^{l}$ to be the “winning” and “losing” samples from human annotations, then $z^{w}=\mathcal{E}(x^{w})$ and $z^{l}=\mathcal{E}(x^{l})$ represent the “winning” and “losing” samples in the latent space. Unlike regular supervised fine-tuning, fine-tuning with DPO utilizes a reference distribution $p_{\text{ref}}(x)$ and hyperparameter $\beta$ for regularization. Fine-tuning now aims to learn $p_{\theta}$, which is aligned to human preferences, while still remembering $p_{\text{ref}}$. The reward modeling objective for fine-tuning takes the form:

$$ $\displaystyle L_{\textit{dpo}}^{\textit{latent}}:=$ $\displaystyle-\mathbb{E}_{\mathcal{E}(x),\epsilon\sim\mathcal{N}(0,1),t}\log\sigma(-\beta(\|\epsilon^{w}-\epsilon_{\theta}(z_{t}^{w},t)\|_{2}^{2}-$ (4) $\displaystyle\|\epsilon^{w}-\epsilon_{\text{ref}}(z_{t}^{w},t)\|_{2}^{2}-(\|\epsilon^{l}-\epsilon_{\theta}(z_{t}^{l},t)\|_{2}^{2}-\|\epsilon^{l}-\epsilon_{\text{ref}}(z_{t}^{l},t)\|_{2}^{2}))).$ $$

Inspired by Meng et al. (), we remove the reference model and simplify the objective to

$$ $\displaystyle L_{\textit{simpo}}^{\textit{latent}}:=-\mathbb{E}_{\mathcal{E}(x),\epsilon\sim\mathcal{N}(0,1),t}\log\sigma\left(-\beta(\|\epsilon^{w}-\epsilon_{\theta}(z_{t}^{w},t)\|_{2}^{2}-(\|\epsilon^{l}-\epsilon_{\theta}(z_{t}^{l},t)\|_{2}^{2}))\right).$ (5) $$

Similar to supervised fine-tuning, we also incorporate calculations in the pixel space into our reward modeling. and define the pixel-space objective as

$$ $\displaystyle L_{\textit{simpo}}^{\textit{pixel}}:=$ $\displaystyle-\mathbb{E}_{\mathcal{E}(x),\epsilon\sim\mathcal{N}(0,1),t}\log\sigma($ (6) $\displaystyle-\beta($ $\displaystyle(\|\mathcal{D}(\sqrt{\bar{\alpha}_{t}}z^{w}+\sqrt{1-\bar{\alpha}_{t}}\epsilon^{w})-\mathcal{D}(\sqrt{\bar{\alpha}_{t}}z^{w}+\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}(z_{t}^{w},t))\|_{2}^{2}$ $\displaystyle-$ $\displaystyle(\|\mathcal{D}(\sqrt{\bar{\alpha}_{t}}z^{l}+\sqrt{1-\bar{\alpha}_{t}}\epsilon^{l})-\mathcal{D}(\sqrt{\bar{\alpha}_{t}}z^{l}+\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}(z_{t}^{l},t))\|_{2}^{2}))).$ $$

Combining latent and pixel terms and weighted by a constant $\mu$, we get

$$ $L_{\textit{reward}}=L_{\textit{simpo}}^{\textit{latent}}+\mu L_{\textit{simpo}}^{\textit{pixel}}.$ (7) $$

## 4 Experiment

We conduct a comprehensive qualitative and quantitative analysis, as well as ablation studies to show that our proposed loss function outperforms the regular latent space loss in both supervised fine-tuning and preference-based post-training.

### 4.1 Human evaluation

Like many recent studies, we found that a rigorous and independent human evaluation process is the most reliable way to evaluate different models. Commonly used metrics such as the FID score do not correlate well with human preference (Dai et al., ; Podell et al., ; Kirstain et al., ).

We contracted a team of paid and independent annotators who do not have contexts of our project to evaluate the generated images. We conducted A/B comparisons on visual flaws and visual appeal, as well as standalone evaluations on text alignment. We use a 600-prompt list in the GenAI MAGIC challenge for the evaluation (Tsai et al., ), where each example is annotated by at least 5 people and the majority decision is used.

Visual flaws. The annotators were presented with two images side-by-side, generated by two different models, without the prompt. The annotators were trained to identify major flaws (e.g., displaced body parts) and minor flaws (distorted eyes), and are asked to choose from “left wins”, “tie” and “right wins”.

Visual appeal. Similar to the visual flaws task, but the annotators are asked to compare which image is more aesthetically pleasing. Annotators were instructed to reject any examples where one image was photo-realistic and the other was stylized (e.g., a cartoon).

Text alignment. We predefined a list of binary questions for each prompt and asked annotators to answer yes or no. We calculated the text alignment rate by aggregating the results across all questions. For example, for the prompt “a cat and a dog”, the annotators are asked “is there one dog”, “is there one cat”, and “are there other animals present”.

### 4.2 experimental setup

Baseline. We tested our model on three models: 1) A 0.6B parameter DiT model with standard transformer and cross attention blocks and trained with high quality data in an “annealing” stage after pretraining to generate high quality images, 2) A 0.86B parameter U-Net with LDM1.5 architecture (Rombach et al., ), replacement-trained on 300M Shutterstock data without quality tuning, which thus generates lower-quality images without prompt engineering and 3) A larger U-Net based Emu model (Dai et al., ) that has been quality fine-tuned, and generates the highest quality images among the three.

Supervised fine-tuning. We curated a small, high-quality dataset of 1816 images for fine-tuning, following the practice of Dai et al. (). Since Emu is already quality fine-tuned, we focused on replacement-trained LDM1.5 and DiT. For supervised fine-tuning with a small dataset, style consistency is crucial. We found that using a hand-picked set of generated images from a high-quality model is sufficient. Examples of our curated fine-tune data are in Appendix Figure [6](#A1.F6).

Preference-based fine-tuning. We conducted experiments on the higher-quality U-Net Emu and DiT models. For each model, we generate 5 images per prompt and ask annotators to select a positive and negative pair. In instances where visual flaws and visual quality conflicted, we prioritize the image with fewest visual flaws as the positive example.

Implementation details. We run all experiments at $512\times 512$ resolution for LDM1.5 and DiT, and $768\times 768$ for Emu, using Adam optimizer with weight decay of $5e-6$. For preference-based fine-tuning, we set $\lambda=\mu=8.0$ for DiT and $2.0$ for Emu to balance the pixel loss magnitude and latent loss magnitude, running 100 epochs. For supervised fine-tuning, we empirically use 140 epochs. We ablate the hyper-parameters such as learning rate and batch size for each model and choose the best ones. Each experiment took 1-8 hours on 8 H100 GPUs to fine-tune one model. During inference, we use the standard DDIM solver with 50 steps with classifier-free guidance.

### 4.3 Experimental Results

Supervised Fine-tuning.
After supervised fine-tuning, our proposed loss improved visual flaws win rate from $17.7\%$ to $64.2\%$, and visual quality win rate from $47.9\%$ to $64.7\%$, compared to regular fine-tuning against the DiT baseline. When directly comparing the model fine-tuned with our loss to the model fine-tuned with the regular latent space loss, ours showed a $32.8\%$ vs $9.3\%$ win rate on visual flaws and $34.8\%$ vs $16.6\%$ win rate on visual appeal. We also found that supervised fine-tuning did not affect text alignment, with correct alignment rates of $74.0\%$, $74.3\%$ and $74.0\%$ for baseline, regular latent fine-tuning, and our fine-tuning respectively, (Table [1](#S4.T1)), all within the margin of error for the annotations. Qualitative examples in Figure [3](#S4.F3) demonstrate that our method generates fewer artifacts and much better fine details.

Figure: Figure 3: SFT with pixel-space loss: DiT. Fine-tuning with our method improves visual quality and reduces flaws. Zoom in to see details.
Refer to caption: /html/2409.17565/assets/x1.png

**Table 1: Pixel space loss improves supervised fine-tuning: DiT.**
|  |  | Visual Flaws |  | Visual Appeal |  | Text alignment |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model A | Model B | A Wins | Tie | B Wins |  | A Wins | Tie | B Wins |  | Model A | Model B |
| Regular SFT | Baseline | $17.7\%$ | $74.0\%$ | $8.3\%$ |  | $47.9\%$ | $19.3\%$ | $32.8\%$ |  | $74.3\%$ | $74.0\%$ |
| Ours | Baseline | $64.2\%$ | $26.5\%$ | $9.3\%$ |  | $64.7\%$ | $20.6\%$ | $14.8\%$ |  | $74.0\%$ | $74.0\%$ |
| Ours | Regular SFT | $\mathbf{32.8\%}$ | $57.8\%$ | $9.3\%$ |  | $\mathbf{34.8\%}$ | $48.6\%$ | $16.6\%$ |  | $74.0\%$ | $74.3\%$ |

**Table 2: Pixel space loss improves fine-tuning: Replacement trained LDM1.5.**
|  |  | Visual Flaws |  | Visual Appeal |  | Text alignment |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model A | Model B | A Wins | Tie | B Wins |  | A Wins | Tie | B Wins |  | Model A | Model B |
| Regular SFT | Baseline | $75.0\%$ | $13.7\%$ | $11.3\%$ |  | $79.6\%$ | $6.9\%$ | $13.5\%$ |  | $72.3\%$ | $61.4\%$ |
| Our SFT | Baseline | $74.2\%$ | $16.7\%$ | $9.1\%$ |  | $75.2\%$ | $10.8\%$ | $14.0\%$ |  | $72.1\%$ | $61.4\%$ |
| Our SFT | Regular SFT | $\mathbf{29.3\%}$ | $46.7\%$ | $24.0\%$ |  | $\mathbf{41.5\%}$ | $24.0\%$ | $34.4\%$ |  | $72.1\%$ | $72.3\%$ |

Table [2](#S4.T2) show the results for LDM1.5 (replacement trained). Our proposed SFT with pixel loss still improves over regular latent SFT in head-to-head comparisons (last row).
Comparing with the DiT experiments, the smaller difference between pixel and regular SFT when compared to the baseline is due to LDM1.5’s lower image generation quality, as it was only replacement trained with Shutterstock data. Therefore, both regularly SFT-ed and pixel SFT-ed models significantly outperform the baseline in terms of visual quality and flawlessness, and supervised fine-tuning in this case, focuses on learning the overall style and aesthetics of the fine-tune images instead of fine details.
See Figures [1](#S1.F1) for qualitative examples. Notably, text alignment also improved with SFT, consistent with findings from Dai et al. ().

**Table 3: Our proposed pixel-space objective also significantly improves DPO on the DiT model.**
|  |  | Visual Flaws |  | Visual Appeal |  | Text alignment |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model A | Model B | A Wins | Tie | B Wins |  | A Wins | Tie | B Wins |  | Model A | Model B |
| Regular DPO | Baseline | $27.5\%$ | $63.7\%$ | $8.8\%$ |  | $48.3\%$ | $39.7\%$ | $12.0\%$ |  | $72.8\%$ | $74.0\%$ |
| Ours | Baseline | $43.3\%$ | $43.7\%$ | $13.0\%$ |  | $61.2\%$ | $24.8\%$ | $14.0\%$ |  | $75.7\%$ | $74.0\%$ |
| Ours | Regular DPO | $\mathbf{31.0\%}$ | $49.3\%$ | $19.7\%$ |  | $\mathbf{43.7\%}$ | $30.4\%$ | $25.9\%$ |  | $75.7\%$ | $72.8\%$ |

Figure: Figure 4: Preference-based post-training with pixel-space loss: DiT. The model trained with our proposed objective function generates more stunning fine details and fewer artifacts.
Refer to caption: /html/2409.17565/assets/x2.png

Preference-based fine-tuning.
Our method also demonstrates exceptional performance in reward-based fine-tuning, generating significantly more impressive details than the baselines. The results are best demonstrated qualitatively in Figure [4](#S4.F4) and [5](#S4.F5). Quantitatively, as shown in Table [3](#S4.T3), compared to regular DPO, our proposed pixel objective function improves the win rate from $27.5\%$ to $43.3\%$ for visual flaws and $48.3\%$ to $61.2\%$ for visual appeal when evaluated against the baseline DiT.
When doing head-to-head comparison between our method and regular DPO, we achieve win rates of $31.0\%$ vs $19.7\%$ on visual flaws and $43.7\%$ vs $25.9\%$ on visual appeal.
Although unintended, our method also improves text alignment by $2.9\%$.

With U-Net based Emu (Figure [5](#S4.F5)), the baseline model already generates higher-quality flawless images in most cases. Therefore, the flaw comparison will result in ties in majority of the cases. Despite this, our proposed method still manages to improve visual flaws win rate from $2.7\%$ to $16.0\%$ and visual appeal from $36.3\%$ to $42.2\%$ as shown in Table [4](#S4.T4).

Figure: Figure 5: Preference-based post-training with pixel-space loss: Emu. Similar to DiT, the Emu model fine-tuned using our loss generates even better details, despite how the baseline Emu already generates good quality images with rich details. Zoom in to see the improvements.
Refer to caption: /html/2409.17565/assets/x3.png

**Table 4: Pixel-space objective also improves DPO on Emu.**
|  |  | Visual Flaws |  | Visual Appeal |  | Text alignment |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model A | Model B | A Wins | Tie | B Wins |  | A Wins | Tie | B Wins |  | Model A | Model B |
| Regular DPO | Baseline | $2.7\%$ | $95.3\%$ | $2.0\%$ |  | $36.3\%$ | $34.9\%$ | $28.8\%$ |  | $89.5\%$ | $89.2\%$ |
| Ours | Baseline | $16.0\%$ | $75.5\%$ | $8.5\%$ |  | $42.2\%$ | $31.4\%$ | $26.5\%$ |  | $89.3\%$ | $89.2\%$ |
| Ours | Regular DPO | $\mathbf{18.3\%}$ | $70.2\%$ | $11.5\%$ |  | $\mathbf{13.7\%}$ | $80.0\%$ | $6.3\%$ |  | $89.3\%$ | $89.5\%$ |

Additional qualitative examples. We provide additional qualitative examples for each experiment above in the Appendix.

### 4.4 Ablation studies

Latent vs Pixel vs Pixel+Latent Loss.
When only using the pixel space loss during supervised fine-tuning, we noticed that the resulting images had very clear details in the main focus of the image, but the background tends to be overly blurred as if they are photographs taken with an extremely narrow depth of field. As shown in Table [5](#S4.T5), using pixel space alone also significantly improves visual flaws, helping the images look more realistic and crisp. However by combining it with latent loss, we are able to significantly improve the visual appeal as well, resulting in images with more stunning details, especially in the background.

**Table 5: Ablation study: Supervised fine-tuning method: Latent, Pixel, vs Latent + Pixel (Ours).**
|  | Visual Flaws | Visual Appeal |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Model A | Model B | A Wins | Tie | B Wins | A Wins | Tie | B Wins |
| Latent only | Baseline DiT | 17.7% | 74.0% | 8.3% | 46.1% | 20.7% | 33.2% |
| Pixel only | Baseline DiT | 43.8% | 42.7% | 13.5% | 44.8% | 21.7% | 33.5% |
| Ours (Latent and Pixel) | Baseline DiT | 64.2% | 26.5% | 9.3% | 63.8% | 21.5% | 14.7% |
| Ours (Latent and Pixel) | Latent only | $\mathbf{32.8\%}$ | $57.8\%$ | $9.3\%$ | $\mathbf{34.8\%}$ | $48.6\%$ | $16.6\%$ |
| Ours (Latent and Pixel) | Pixel only | 22.7% | 60.5% | 16.8% | 47.5% | 38.5% | 14.0% |

Decoding Methodology.
Intuitively, to obtain as much image quality and details as possible, one may consider regressing to an objective function that compares the output to the original starting image in the pixel space. This involves two steps: first, transforming the predicted noise in the latent space back to $t=0$, then decoding into the pixel space using equation

$$ $x_{0}=\mathcal{D}(z_{0})=\mathcal{D}\biggl{(}\frac{1}{\sqrt{\bar{\alpha}_{t}}}(z_{t}-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}(z_{t},t))\biggr{)}.$ (8) $$

Although this method seems like it would be optimal in generating the detail and high quality desired in the final image,
the transformed $x_{0}$ has a greater variance for larger timesteps, causing the generated images to be blurry and fuzzy (Appendix Figure [7](#A1.F7)) since the fine-tuning process was trying to correct for the estimation error of $z_{0}$.

Based on this finding, we propose comparing the output directly with the sample in the pixel space at the current timestep to eliminate the unwanted variations in the previous method, using Equation [2](#S3.E2).
Empirically, we have found that even though noisy images are out-of-distribution for the autoencoder, it still does surprisingly well in reconstructing them.

Reference Model in Reward Modeling. Traditionally, DPO (Rafailov et al., ; Wallace et al., ) utilizes a reference model, but recently Meng et al. () proposed removing the reference model (SimPO) in LLMs and showed strong performance. Here, we tested out different combinations of using latent-space loss and pixel-space loss with and without the reference model, as shown in Table [6](#S4.T6).
We show that simply adding our proposed pixel loss term can already significantly improve the visual flaws metric using the standard DPO method ($53.0\%$ vs $27.5\%$ win rate compared to baseline DiT model). However, by incorporating SimPO, we achieve both significant improvement on visual flaws and visual appeal (improving from $27.5\%$ to $43.3\%$ in visual flaws and $47.2\%$ to $61.2\%$ in visual appeal).
To the best of our knowledge, this is also the first paper that demonstrates that the recent success of SimPO can also be extended to diffusion models.

**Table 6: Ablation Study: Reward modeling variations vs Baseline DiT**
|  | Visual Flaws | Visual Appeal |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Model | Win | Tie | Lose | Win | Tie | Lose |
| DPO latent (baseline) | 27.5% | 63.7% | 8.8% | 47.2% | 40.7% | 12.2% |
| DPO latent + DPO pixel | 53.0% | 31.5% | 15.5% | 49.0% | 21.0% | 30.0% |
| DPO latent + SimPO pixel | 50.7% | 35.7% | 13.7% | 41.3% | 22.8% | 35.8% |
| SimPO latent + SimPO pixel (Proposed) | 43.3% | 43.7% | 13.0% | $\mathbf{61.2\%}$ | $24.8\%$ | $14.0\%$ |

## 5 Limitations

Limitations of Baseline Models. Fine-tuning improvements are dependent on the quality of the original baseline model, and thus fine-tuning using pixel space loss may not always produce significant improvements. For example, if the original baseline model already has minimal flaws and high visual appeal, fine-tuning may not achieve many improvements.
In contrast, if the original baseline foundation model generates significant structural flaws that require the global understanding of the image composition, fine-tuning with our loss alone may not help eliminate them.

Limitations of Fine-tuning Dataset. Fine-tuning is also dependent on the quality and composition of the dataset used.
Our dataset was hand-curated and consisted of images that followed our definition of high quality and style. Changing the composition of this dataset would lead to different results.

Limitations of Human Evaluation. The images generated by the different models were evaluated by independent annotators. Although the annotators were trained on standards for visual flaws, visual appeal, and text alignment, these results may not fully reflect the real-world use of the models. Human evaluation is also inherently subjective and noisy in terms of aesthetics.

## 6 Conclusions

In this paper, we proposed a novel post-training objective function for latent diffusion models by incorporating a pixel-space loss with the commonly used latent-space fine-tuning loss. The resulting model shows noticeable improvement in visual flaws and visual appeal metrics in both supervised fine-tuning and preference-based post-training through rigorous human evaluations. The proposed objective function is simple and can be easily plugged into existing models such as DiT and U-Net.

#### Acknowledgments

We thank Yuval Kirstain, Xiaoliang Dai, and Zecheng He for their useful discussions.

## Appendix A Appendix

### A.1 Fine-tune dataset

Here we show some examples of our supervised fine-tune dataset, which are selected images generated by Emu (Dai et al., ).

Figure: Figure 6: Fine-tune dataset. Selected images in our supervised fine-tune dataset.
Refer to caption: /html/2409.17565/assets/x4.png

### A.2 Decoding Methodology

An alternative decoding methodology would be to transform the latent space back to $t=0$, and then decode it into the pixel space to obtain $x_{0}$, as discussed in Section [4.4](#S4.SS4). Figure [7](#A1.F7) shows that if one follows Equation [8](#S4.E8) to decode back to $t=0$, it leads to blurrier images for larger timesteps. Therefore, we chose to decode back directly at the present timestep, as discussed in Section [4.4](#S4.SS4) using Equation  [2](#S3.E2).

Figure: Figure 7: Decoding methodology. Transforming images back to $t_{0}$ causes the decoded images to be blurrier for larger timesteps. Therefore, we choose to decode at the present timestep.
Refer to caption: /html/2409.17565/assets/figures/decoding_method.png

### A.3 More examples

Figure: Figure 8: More examples: DiT with preference-based fine-tuning. We provide more examples here to demonstrate the improvements in high-frequency details and visual flaws.
Refer to caption: /html/2409.17565/assets/x5.png

Figure: Figure 9: More examples: U-Net Emu with preference-based fine-tuning. Similar to the DiT model, we also observe improvements for Emu, despite how the baseline model already generates quite high quality images in most cases.
Refer to caption: /html/2409.17565/assets/x6.png

Figure: Figure 10: More examples: DiT with supervised fine-tuning. The model fine-tuned with our objective function often generates sharper images, which is best appreciated when zooming in. The model fine-tuned with regular latent loss, on the other hand, is typically blurrier, such as the animals’ furs. Also as shown in many examples here, our model tends to generate fewer flaws.
Refer to caption: /html/2409.17565/assets/x7.png

Figure: Figure 11: More examples: U-Net LDM1.5 with supervised fine-tuning. The baseline replacement-trained LDM1.5 often generates images with bad composition and noticeable artifacts. Supervised fine-tuning alone helps improve the image quality significantly. Our loss further improves visual quality and flaws. Again, readers can better appreciate the improvements when zooming in to notice the improvements in the overall sharpness and fine details.
Refer to caption: /html/2409.17565/assets/x8.png

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: What is the difference in generated image quality when the loss is computed as in the paper (on actual images) versus on denoised latent? (where the denoising is done on the same schedule as for images)

What evidence is there that the high frequency details of themselves have changed, as opposed to the images just having superior layout/lighting/overall appeal?
soundness: 2
strengths: [S1] The pixel-space SFT and pixel-space DPO consistently have superior image quality win rate to baselines and default SFT/DPO, along with similar or better text-alignment.

[S2] The paper communicates its key ideas thoroughly and clearly, and provides good examples and metrics that are sufficient to convince of the efficacy of the proposed approach.

[S3] The method is flexible and can easily be applied to any model.
summary: The paper proposes to adjust the supervised-finetuning originally proposed by Emu, by computing the loss in pixel space instead of latent space. Furthermore, the loss is on the reverse-noised sample (noise prediction removed from image), rather than on the noise prediction itself. They demonstrate the efficacy of this change via evaluation by third party human annotators.
weaknesses: [W1] The efficacy of the method is not very well-connected to the motivation. Pixel space is not necessary for such stunning or complex images (although it may increase the likelihood the model generates such images). Without zoom-in figures, it is unclear how/why the higher resolution loss is helpful, since it's not obvious that the fine-grained details are any better, which should be the main outcome. Simply put, I see no evidence that "adding pixel-space supervision in the post-training process to better preserve high-frequency details." (abstract)

[W2] Related to W1, there seems to be a significant confounding factor: instead of the only difference between SFT and yours being the latent vs. pixel space, Equations 1 and 2 indicate that SFT operates on noise predictions only, whereas your loss operates on the input with the noise prediction used for reverse diffusion. Thus, it is difficult to gauge the difference that should be attributed to computing loss in pixel space, compared to computing loss after removing the noise prediction.

[W3] The proposed method is a very incremental adjustment on SFT (the key insight from Emu). SFT is simply performed with 2 additional steps: subtracting the noise from the input, and decoding to pixel space for MSE in pixel space. Without other compelling insights, this is only very slightly novel.

[W4] While human measurement is undoubtedly the gold standard, it would have been helpful to provide some automatic metrics to contextualize the results.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: All my questions are in weaknesses above.
soundness: 2
strengths: 1. The hypothesis that losses of details and artifacts in high-frequency details in Latent diffusion models (LDMs) are partially caused by training on the lower-resolution latent space is sound.
2. The proposed pixel-space objective is straightforward.
summary: This paper hypothesizes that losses of details and artifacts in high-frequency details in Latent diffusion models (LDMs) are partially caused by training on the lower-resolution latent space, and proposes adding a pixel-space objective during LDMs post-training. The experiments show improvements in both DiT-based and UNet-based LDMs for reward-based and supervised fne-tuning.
weaknesses: 1. The proposed method, with additional supervision in the pixel space, is too simplistic and lacks technical innovation.
2. The proposed method lacks comprehensive experimental validation; current experiments rely solely on human evaluation without additional quantifiable metrics.
3. The proposed method has only been validated under fine-tuning settings, without verification under pre-training settings.
4. The proposed method is relatively costly, as transforming from the latent space to the pixel space involves passing through a VAE, which can be computationally intensive.
5. The proposed method lacks theoretical and experimental explanations for why it is effective.
6. The quality of the writing is suboptimal, exhibiting problems with both logical coherence and linguistic fluency, particularly evident in the introduction section.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Please refer to the weaknesses.
soundness: 2
strengths: 1. Addressing the enhancement of detail realism in latent diffusion generation is an important problem.
2. Experimental results demonstrate that the proposed method achieves performance improvements over existing fine-tuning approaches.
summary: This paper addresses limitations in Latent Diffusion Models (LDMs) related to generating high-frequency details and complex compositions, which the authors attribute to conducting pre- and post-training exclusively within a low-resolution (8×8) latent space. To overcome these issues, the paper proposes integrating pixel-space supervision in the post-training phase, aiming to better capture high-frequency details.
weaknesses: 1. The proposed technical approach—simply adding a pixel-level loss—lacks novelty and may appear overly simplistic. The use of reinforcement learning within a diffusion model is also not groundbreaking, and the extension of SIMPo does not constitute a notable contribution. It would be better if the authors could clarify why the pixel-level loss is novel.

2. Technical details and motivations are not well-articulated: (1) Why does Equation 2 use a decoded ground truth image rather than the original image? If decoding already introduces detail loss, it is unclear how this approach can effectively enhance detail realism as intended. (2) In applying the pixel loss, is a random timestep selected? If so, when the noise level is high, how is one-step prediction accuracy ensured? It would be better if the authors could clearly clarify these questions.


3. The results presented in Figure 4 reveal that the method still struggles with preserving finer details, particularly in high-frequency regions, where the generated images display notable artifacts or blurriness. This suggests that while the proposed pixel-level supervision may have improved some aspects of detail generation, it falls short in producing consistently realistic textures across different parts of the image. Additionally, certain areas of the images, which typically require precise detailing—such as edges, textures, or small intricate patterns—do not appear as natural as expected, raising questions about the method's overall effectiveness in enhancing detail realism. 

4. The enhancement of natural image details is unconvincing. I suggest testing this approach in domains where detail accuracy is more critical, such as talking face generation. It would be beneficial to see if it truly improves fine details like eyes and teeth. The authors may consider using Diffusion-based foundation models, such as https://github.com/fudan-generative-vision/hallo.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Which DiT model does the authors use?

2. Is this method very easy to implment, like change a minor architecture, and add a loss term in the original loss?

3. Why the authors conduct experiments in SD1.5 and an unknown DiT? Why they do not use more recent UNet like SDXL, and use DiT in SD3?
soundness: 3
strengths: 1. This method have a relatively good motivation. 

2. The writing is clean and easy to follow.

3. The overall idea is general, so it can be implemented in both SFT and reward based finetuning. 

4. Although the method looks simple, extensive experiments show its effectiveness.
summary: This paper think the flaws occured in latent diffusion models is resulted from all the training process are conducted in the low-resolution latent space. Thus this manuscript proposes to add pixel-space supervision in the post-training process. This idea is simply and general, so it can have both implementation in supervised finetuning and reward-based finetuning. In practice, the authors modify SFT and DPO algorithm by mapping the latent back to the pixel space, and compute the pixel-space SFT/DPO as an additional loss term. The experiments is abundant, including lots of human evaluation, and verifies its effectiveness.
weaknesses: 1. The used DiT model detail is not explained.

2. The method is quite simple. The analysis of why this method can improve performance is missed.

3.  The authors do not conduct their experiments on more recent model. It is suggested to verity their method in more recent UNet and DiT models.

Minor: The text in equation should be put inside \text{...}