There is No VAE: End-to-End Pixel-Space Generative Modeling via Self-Supervised Pre-training

TL;DR

The paper introduces EPG, a two-stage, end-to-end pixel-space generative framework that pre-trains an encoder to learn stable semantic representations along a diffusion trajectory and then fine-tunes a randomly initialized decoder for diffusion or consistency objectives. By combining contrastive and representation-consistency learning, the method achieves state-of-the-art results on ImageNet in pixel-space diffusion (FID as low as 1.58 on 256 and 2.35 on 512 with 75 NFEs) and strong one-step consistency model performance (FID 8.82 on ImageNet-256), while using substantially less training compute than VAE-based approaches. The work also demonstrates the first direct-pixel-space consistency-model training on ImageNet-256 without pre-trained VAEs or diffusion models, highlighting improved efficiency and stability from the proposed SSL-inspired pre-training. Overall, the approach narrows the performance and efficiency gap between pixel-space and latent-space generative methods and offers a scalable path for high-resolution image synthesis.

Abstract

Pixel-space generative models are often more difficult to train and generally underperform compared to their latent-space counterparts, leaving a persistent performance and efficiency gap. In this paper, we introduce a novel two-stage training framework that closes this gap for pixel-space diffusion and consistency models. In the first stage, we pre-train encoders to capture meaningful semantics from clean images while aligning them with points along the same deterministic sampling trajectory, which evolves points from the prior to the data distribution. In the second stage, we integrate the encoder with a randomly initialized decoder and fine-tune the complete model end-to-end for both diffusion and consistency models. Our framework achieves state-of-the-art (SOTA) performance on ImageNet. Specifically, our diffusion model reaches an FID of 1.58 on ImageNet-256 and 2.35 on ImageNet-512 with 75 number of function evaluations (NFE) surpassing prior pixel-space methods and VAE-based counterparts by a large margin in both generation quality and training efficiency. In a direct comparison, our model significantly outperforms DiT while using only around 30\% of its training compute.

Figures (19)

• Figure 1: (Left) Our diffusion model achieves SOTA performance while maintaining significant inference efficiency. The $x$-axis indicates the log-scaled GFLOPs for generating an image. The bubble size denotes number of model parameters. (Right) Images generated by our diffusion model.
• Figure 2: Overview of our framework. During pre-training, we train encoder by learning visual semantics (blue branch) and aligning them along ODE sampling trajectories (yellow branch). $\theta^-$ is the exponential moving average of $\theta$, and $sg$ is the stop gradient operation. After pre-training, we discard the projector and train $E_\theta$ along side a randomly initialized decoder $D_\theta$ in an end-to-end manner under specific generative training settings (e.g., denoising or consistency training). The encoder $E_\theta$ maps noisy images into features, which are then reconstructed into clean pixels by the decoder $D_\theta$. In our work, we choose vision transformer as our backbone.
• Figure 3: Class-conditional generation performance comparing with DiT and SiT.
• Figure 3: Performance in downstream tasks scales with pre-training compute budget. DM, CM: fine-tune under the settings of diffusion or consistency model.
• Figure 4: The network architecture of EPG utilized in downstream generative tasks.