Pós-Treinamento no Espaço de Pixels de Modelos de Difusão Latente

Pixel-Space Post-Training of Latent Diffusion Models

September 26, 2024

Autores: Christina Zhang, Simran Motwani, Matthew Yu, Ji Hou, Felix Juefei-Xu, Sam Tsai, Peter Vajda, Zijian He, Jialiang Wang

cs.AI

Resumo

Os modelos de difusão latente (LDMs) têm avançado significativamente no campo da geração de imagens nos últimos anos. Uma grande vantagem dos LDMs é a capacidade de operar em um espaço latente comprimido, permitindo um treinamento e implantação mais eficientes. No entanto, apesar dessas vantagens, desafios com os LDMs ainda persistem. Por exemplo, tem sido observado que os LDMs frequentemente geram detalhes de alta frequência e composições complexas de forma imperfeita. Nossa hipótese é que uma razão para essas falhas é devido ao fato de que todo o pré e pós-treinamento dos LDMs são feitos no espaço latente, que é tipicamente 8 vezes 8 de menor resolução espacial do que as imagens de saída. Para abordar esse problema, propomos adicionar supervisão no espaço de pixels no processo de pós-treinamento para preservar melhor os detalhes de alta frequência. Experimentalmente, mostramos que adicionar um objetivo no espaço de pixels melhora significativamente tanto a qualidade de ajuste fino supervisionado quanto o pós-treinamento baseado em preferência em uma grande margem em modelos de difusão DiT transformer e U-Net de última geração, tanto em qualidade visual quanto em métricas de falhas visuais, mantendo a mesma qualidade de alinhamento de texto.

English

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.