Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, Liang Zheng
9 min
Latent Diffusion Models (LDMs) are powerhouse generative models that create images by diffusing noise in a compressed latent space rather than raw pixels. This space is produced by a Variational Autoencoder (VAE), acting as a tokenizer that squeezes high-res images into a compact representation. Traditionally, LDM training is a rigid two-stage affair: first train the VAE standalone, then freeze it and train the diffusion model (like a Scalable Image Transformer or SiT) on top. Why freeze? Naive end-to-end training—backpropagating diffusion loss through the VAE—causes latent space collapse, where the VAE simplifies representations too much, crippling generation quality.
REPA-E flips this script. It introduces a representation-alignment (REPA) loss that enables true end-to-end (E2E) joint tuning of VAE encoder and diffusion model. The result? Massive speedups (17× over prior REPA baselines, 45× over vanilla), better FID scores (e.g., 1.12 on ImageNet 256×256 without guidance), and surprisingly, improved VAE latent spaces that adaptively fix their own flaws.
In standard LDMs (e.g., Stable Diffusion), the VAE is fixed post-training (Fig. 1a). This creates a mismatch: the VAE's latents might be noisy (SD-VAE has high-frequency junk) or over-smoothed (IN-VAE lacks detail), suboptimal for diffusion denoising. Tuning VAEs separately is tricky—prior work like REPA aligned representations without E2E, but training stayed slow and staged.
Naive E2E (Fig. 1b) uses diffusion loss directly on VAE params, but it degrades latents further, as diffusion signals don't guide representation learning well. REPA-E (Fig. 1c) swaps in REPA loss + a clever BatchNorm layer with EMA stats for stable normalization, unlocking joint optimization without collapse.
REPA loss aligns diffusion model outputs with pretrained perceptual features (e.g., from DINO/CLIP), measured via CKNNA (Centered Kernel Nearest Neighbors Alignment) score. CKNNA quantifies how well latents capture semantic structure—higher scores predict better gFID. Unlike diffusion loss (noisy, sequence-dependent), REPA is a clean perceptual proxy that backprops effectively to VAE encoder.
Crucially, REPA prevents collapse by encouraging latents to match rich, pretrained representations. A stop-gradient on decoder outputs and BatchNorm (using dataset EMA, not mini-batch) dynamically normalizes latents cheaply, avoiding recompute.
PCA visualizations (Fig. 2a) are telling: project VAE latents to 3-channel RGB via top PCs. SD-VAE starts noisy; E2E + REPA smooths it. IN-VAE starts blurry; E2E adds detail. This automatic adaptation happens without VAE-specific losses—the diffusion task itself shapes better latents via REPA.
(Fig. 2b) shows gFID gains across VAEs after 400K steps. Tuned VAEs are drop-in upgrades for any LDM.
REPA-E trains SiT diffusion models to SOTA: FID 4.07 (vs. 5.9 baseline) at 17-45× speed on ImageNet. Training curves (Fig. 1d) converge faster, as E2E avoids VAE-diffusion mismatch.
Broader impact: E2E challenges 'freeze the tokenizer' dogma in diffusion. Tuned VAEs boost downstream tasks; REPA generalizes beyond images. Simple, effective—code at https://end2end-diffusion.github.io.
In this paper we tackle a fundamental question: "Can we train latent diffusion models together with the variational auto-encoder (VAE) tokenizer in an end-to-end manner?" Traditional deep-learning wisdom dictates that end-to-end training is often preferable when possible. However, for latent diffusion transformers, it is observed that end-to-end training both VAE and diffusion-model using standard diffusion-loss is ineffective, even causing a degradation in final performance. We show that while diffusion loss is ineffective, end-to-end training can be unlocked through the representation-alignment (REPA) loss -- allowing both VAE and diffusion model to be jointly tuned during the training process. Despite its simplicity, the proposed training recipe (REPA-E) shows remarkable performance; speeding up diffusion model training by over 17x and 45x over REPA and vanilla training recipes, respectively. Interestingly, we observe that end-to-end tuning with REPA-E also improves the VAE itself; leading to improved latent space structure and downstream generation performance. In terms of final performance, our approach sets a new state-of-the-art; achieving FID of 1.12 and 1.69 with and without classifier-free guidance on ImageNet 256 x 256. Code is available at https://end2end-diffusion.github.io.
Alex: Huh. So it's like cheating the training to make one step easy, but hurting the overall picture. How do they know the alignment loss avoids that?
Sam: They measure alignment by comparing how closely the noise-remover's internal views of images match those from a fixed feature extractor—a pre-trained system that pulls out key patterns like edges and colors from real photos, without changing itself. A score called CKNNA checks this match across image patches; higher scores mean the views line up better. The paper finds higher CKNNA during training links directly to sharper final images, as shown in their graphs—unlike the diffusion signal, alignment keeps summaries detailed.
Alex: Right, so alignment acts as a stand-in goal for good images. But if the compressor's summaries limit how high that score can go, how do they push it further?
Sam: In the standard setup, the alignment score plateaus because the compressor's features can't improve. By sending the alignment signal back through the compressor too, they raise the score notably higher—like tuning the source to better match the listener. They add a few practical tweaks: a normalization layer between parts, using running averages of stats to keep inputs steady without re-scanning the whole dataset each time; researchers call this a batch-norm layer with EMA. They also block the diffusion signal from reaching the compressor using a stop-gradient operation, so it only trains the noise-remover.
Alex: So the alignment signal tunes everything end-to-end without the simplifying hack, and these tweaks make it feasible.
Sam: Yes, and the tuned compressor even slots into other systems for better results overall. It's a meaningful refinement in how these generators train.
Alex: Okay, so the tweaks stabilize the process. How do they actually combine all these signals during training without one overpowering the others?
Sam: They blend three main signals into one overall goal. First, the usual noise-removal signal trains just the main generator part, without flowing back to the compressor. Second, the alignment signal adjusts the generator, compressor, and a small extra matching layer together. Third, they add a set of checks just for the compressor—like comparing its summaries to the original images in different ways, plus a balancing term to keep summaries probabilistic and varied. This mix lets everything learn at once.
Alex: Right, so the compressor's extra checks prevent it from drifting too far while alignment pulls everything in sync. Does that lead to noticeably better images right from early training?
Sam: Yes, the paper shows this through side-by-side image comparisons from checkpoints during training. When generating the same picture prompt at early steps, their approach produces clearer shapes and details compared to the prior method, which looks blurrier and less structured. By later steps, the gap widens: their images hold up better overall. This early strength suggests the balanced signals build a solid foundation faster.
Alex: Huh, so it's not just the end result, but steady progress from the start. And this holds up across different model sizes?
Sam: The experiments confirm it improves performance substantially regardless of model scale—from smaller ones around the size of a smartphone app to larger ones fitting on high-end computers. They measure image quality with a standard score that compares generated pictures to real ones; adding their tuning drops that score notably lower, meaning better matches to reality.
Alex: So the end-to-end balance accelerates reliable scaling without special tricks. One thing you touched on—the automatic improvement to the compressor's summary space. The paper shows visualizations of that. What do they reveal?
Sam: Pre-trained compressors often have flaws in their summary spaces. Some pack in too much random junk, like static overwhelming the signal in a weak radio reception; others smooth everything out too much, losing fine details like blurring a photo until textures vanish. The researchers plot these spaces in two dimensions to spot the issues—spreading points show variety and cleanliness, clusters of mess mean problems. With their end-to-end tuning, these flaws fix themselves across various compressor designs: noisy ones quiet down, smooth ones gain detail.
Alex: Huh. So it adapts no matter the starting design, pushing toward a better structure automatically.
Sam: Yes, the paper tests variations systematically. They try different feature matchers—pre-trained systems spotting patterns in images—and different points in the noise-remover to apply alignment. In each case, the approach improves image quality notably over the baseline. This shows robustness: it works regardless of those choices.
Alex: Right, so not tied to one setup. What about the tweaks like blocking the noise signal or the normalization layer—do ablations confirm their roles?
Sam: Ablations leave out one piece at a time. Without blocking the noise signal from reaching the compressor, quality crashes badly, as it reintroduces simplification. Skipping the normalization layer between parts worsens scores noticeably, since evolving summaries need steady scaling—like adjusting volume on a changing amp. Both confirm the components keep training stable.
Alex: Okay, that makes the balance clearer. And they even train the whole thing from scratch, without pre-trained compressors?
Sam: They do. Starting with random weights for both compressor and noise-remover, it still outperforms the baseline that needs a separate pre-training stage—reaching strong quality in far fewer steps. Pre-initializing helps a bit more, but the method stands alone.
Alex: So these tuned compressors aren't just for this setup...
Sam: Exactly. Once tuned, they drop into other generators as replacements, speeding up and sharpening results there too. It's a practical carryover benefit.
Alex: Makes sense for practical builds. Though, what are the main limits to keep in mind?
Sam: A few practical ones stand out. It relies on a pre-trained pattern-spotter for the alignment signal, so you need that ready first. The compressor also requires extra checks during tuning to stay stable, and testing focuses mainly on one dataset size.
Alex: Fair points—keeps expectations grounded. Thanks for breaking it down, Sam. That's our look at unlocking better training for image generators. Thanks for listening to ResearchPod.