Pixel-Space Post-Training of Latent Diffusion Models

TL;DR

This work tackles the limitation of latent-diffusion models losing high-frequency details due to latent-space bottlenecks by introducing pixel-space supervision during post-training. It formulates a pixel-space objective that complements the standard latent-space loss, and demonstrates substantial gains in visual quality and reduction of artifacts for both DiT and U-Net LDMs, while preserving text alignment. The approach applies to both supervised quality fine-tuning and reward-based post-training, and further extends SimPO-style reward modeling to diffusion models. The results, validated by rigorous independent human evaluation, show that the simple, architecture-agnostic pixel-space loss is a practical, effective enhancement for diffusion-based image generation.

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

Figures (11)

• 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.
• 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.
• Figure 3: SFT with pixel-space loss: DiT. Fine-tuning with our method improves visual quality and reduces flaws. Zoom in to see details.
• 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.
• 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.