Latent diffusion model trained end to end for better image generation
LDM-is-AE: Latent Diffusion Model is an Auto-Encoder for End-to-End Image Generation
Computer Vision and Pattern Recognition
Summary
Usually, image generation uses two separate steps: one finds a simple way to represent images, and the other learns how to improve blurry images step by step. The authors found that the same model can do both steps together by treating part of the model as an internal auto-encoder, which means it learns to compress and reconstruct images internally while improving the details. This approach trains the model all at once so that the way it represents images is better suited for making high-quality images. Their experiments show this method produces very good images at different sizes.
What this means in practice
- •For graphics software developers: Develop end-to-end image generation tools that produce higher fidelity images without pretraining separate image encoders.
- •For creative ai service providers: Offer improved AI-based image generation services with better internal modeling for diverse image resolutions.$Commercial implications: This paper enables more integrated image generation models improving output quality and efficiency attractive for creative content platforms.
Authors
Zhengqiang Zhang, Lingchen Sun, Rongyuan Wu, Qiaosi Yi, Xiangtao Kong, Chaodong Xiao, Lei Zhang
Abstract
Latent Diffusion Models (LDMs) typically adopt a two-stage pipeline: an auto-encoder (AE) is first pre-trained to define a latent space, then a diffusion model is trained to perform denoising within it. Such a two-stage design introduces a representation mismatch, as the latent space is optimized for reconstruction rather than adapting the denoising dynamics. We reveal that the LDM itself is an AE, and consequently present LDM-is-AE, an end-to-end one-stage LDM training framework that eliminates the need for a separately trained tokenizer. Our key observation is that the LDM backbone actually performs a latent-to-feature-to-latent transformation at each denoising step, which can be interpreted as an internal decoding--encoding process. Leveraging this structure, we split the DiT backbone into two reciprocal components, DiT-E (i.e., DiT Encoding) and DiT-D (i.e., DiT Decoding), and impose image-space supervision on the intermediate features across all timesteps. Our model encourages the internal representation to align with the image domain throughout denoising, thereby establishing an explicit latent-to-image-to-latent path. At the zero-noise timestep, our model further performs an image-to-latent-to-image mapping, corresponding to an auto-encoding process. As a result, LDM-is-AE jointly learns latent representations and denoising dynamics in an end-to-end manner, yielding a diffusion-native latent space tailored to the generation process. Experiments demonstrate that LDM-is-AE exhibits highly competitive generation performance, achieving an FID of 1.80 and 1.90 on 256x256 and 512x512 class-conditional image generation, respectively.