SpeeDiff: Scalable Pixel-Anchored End-to-End Latent Diffusion Model

AI summary

Overview

Research area: Generative computer vision — latent diffusion models (LDMs), visual tokenizers (VAEs), and large-scale class-conditional image generation on ImageNet.

Technical level: Advanced. The paper assumes familiarity with diffusion/flow-matching objectives, stochastic interpolants, Tweedie's formula, VAE/ELBO training, and vision foundation model representation alignment.

Scope: The paper proposes SpeeDiff, a single-stage framework that trains a transformer-based VAE and a diffusion transformer jointly from scratch (no stop-gradient) using a new Tweedie Pixel Reconstruction loss to prevent latent collapse, plus a representation-alignment scheme called REPA++.

What This Paper Is About

Standard latent diffusion trains a VAE first, freezes it, and then trains a diffusion model inside that fixed latent space — so the tokenizer never receives feedback about whether its latents are actually easy to generate from. The authors show that simply letting the diffusion loss backpropagate into the VAE encoder (naive end-to-end training) fails badly, because the latent space collapses into a degenerate representation. SpeeDiff's goal is to make that end-to-end coupling work, so that generative supervision actively shapes a "generation-friendly" latent space and speeds up convergence dramatically.

Key Contributions

1. A diagnosis of why vanilla end-to-end LDM training fails. The authors show that unrestricted backpropagation of the diffusion loss into the VAE encoder produces latent space collapse: near-zero latent reconstruction error alongside large pixel reconstruction error, multi-peaked non-Gaussian latent distributions, large per-channel biases, and severely suppressed per-channel variances.
2. The Tweedie Pixel Reconstruction (TPR) loss. A simple pixel-space MSE that decodes the Tweedie-estimated clean latent (ẑ₀ = z_t − t·V_θ(z_t, t)) from an intermediate noisy state and compares it to the ground-truth image, anchoring the VAE and preventing degeneracy without any stop-gradient.
3. A refined, fully transformer-based architecture. A ViT-based VAE (ViT-VAE) replaces the convolutional VAE, and the DiT backbone is modernized (RMSNorm, SwiGLU, 2D RoPE, global modulation with per-block learnable biases, patch size 1), reducing total training cost from 436.29 to 334.98 GFLOPs.
4. REPA++ representation alignment. An end-to-end alignment scheme that combines Latent-REPA (aligning clean latents to VFM features via a two-layer MLP and patch-wise cosine similarity) with Diff-REPA (aligning intermediate diffusion features), both optimized jointly with the VAE and diffusion model.

Main Findings

• TPR rescues end-to-end training. Vanilla E2E reaches an FID of 33.95 on ImageNet 256×256 after 80 epochs; adding the TPR loss drops it to 5.79, and adding the refined architecture brings it to 3.66.
• End-to-end beats detached training at every stage. With a stop-gradient (no direct gradient feedback), the two-stage baseline scores 13.02 and the single-stage baseline 14.21; adding TPR gives 11.65, the refined architecture 7.52, and REPA++ 3.47 — all worse than the corresponding end-to-end variants (33.95 → 5.79 → 3.66 → 1.69).
• Final ImageNet 256×256 result. SpeeDiff-XL (w/ REPA++) achieves an FID of 1.50 without guidance after 200 epochs; the variant without REPA++ achieves 2.42, reported as a new state of the art among methods that do not use VFM alignment.
• Strong 80-epoch results. SpeeDiff-XL reaches 1.69 (w/ REPA++) and 3.66 (w/o REPA++) without guidance after only 80 epochs.
• Large convergence speedups. SpeeDiff-XL (w/ REPA++) reaches a gFID of 7.36 without guidance within 10 epochs, already surpassing vanilla SiT trained for 1400 epochs. The paper reports over 140× faster training than Vanilla SiT and 61× faster than REPA.
• High-resolution result. On ImageNet 512×512 with a 32× compressed VAE (f32d32), SpeeDiff-XL reaches an FID of 1.53 without guidance after 200 epochs.
• The learned VAE is independently useful. Freezing the converged SpeeDiff-XL VAE and training a fresh diffusion model on it converges at nearly the same rate (gFID 6.94 at 10 epochs and 1.73 at 80 epochs, versus 7.36 and 1.69 for SpeeDiff-XL itself), suggesting the method acts like latent pretraining.
• Semantic gains in the latent space. Latent linear probing accuracy on ImageNet improves by 9.45% with diffusion-gradient feedback, and Robust PCA visualizations show smoother, more semantically organized latents versus independently trained VAEs.
• REPA++ components complement each other. Ablations at 80 epochs: full REPA++ with DINOv3-ViT-L/16 gives 1.69; with DINOv2-ViT-L/14, 1.76; Latent-REPA alone, 2.28; Diff-REPA alone, 2.15. Latent-REPA most improves VAE latent alignment, Diff-REPA most improves intermediate diffusion features, and their addition covers both.
• 32 latent channels work best. The paper reports that 32 channels offer the best trade-off between reconstruction and generation across the settings tested.

Methodology in Plain English

The team starts from the observation that a frozen VAE is trained only to reconstruct pixels, so its latent space may not be convenient for a diffusion model to model. Their fix is to remove the freeze: the VAE encoder, VAE decoder, and diffusion transformer are all trained together from scratch, with the diffusion loss flowing all the way back into the encoder.

The obstacle is that this raw coupling destroys the latent space — the diffusion model finds a way to score well on its own latent objective while the latents stop carrying the information needed to reconstruct images. To block that shortcut, the authors add a second, pixel-level check. At each training step they take the noisy latent, use the diffusion model's own velocity prediction to estimate what the clean latent should be (Tweedie's formula), decode that estimate with the VAE decoder, and penalize the difference from the real image. Because this loss is measured in pixels and depends on the diffusion prediction, the encoder cannot drift into a degenerate latent space without being penalized.

On top of this, they swap the CNN-based VAE for a transformer-based VAE and update the diffusion transformer with modern components, which lowers compute. They also add REPA++, which pulls both the clean latents and the diffusion model's intermediate features toward features from a pretrained vision foundation model (DINOv3-ViT-L/16, 0.3B parameters) using patch-wise cosine similarity through small MLP adapters.

Training uses ImageNet at 256×256 (and 512×512 by changing only the first patch-embedding layer to 32× compression), 200 epochs / 1M iterations, batch size 256, Adam with a constant learning rate of 1e-4, and EMA with decay 0.999 for the ViT-VAE and 0.9999 for Refined-DiT. Sampling uses an Euler–Maruyama SDE solver with 250 NFEs and AutoGuidance rather than CFG. Evaluation uses 50,000 generated samples with the official OpenAI evaluation suite, reporting gFID, sFID, IS, precision, recall, plus rFID, PSNR, linear probing, and CKNNA.

Why This Matters

Impact on research. The paper reframes a known negative result — that end-to-end LDM training degrades quality — as a fixable optimization pathology rather than a fundamental limitation, and it offers a minimal, cheap remedy. It also challenges the default two-stage recipe and the assumption that better tokenizers must come from vision foundation models, showing that diffusion gradients alone improve latent semantics.

Real-world applications:

• Faster, cheaper training of text-to-image and video generation backbones, where VAE/diffusion co-design currently requires separate expensive stages.
• High-resolution synthesis pipelines, since the same setup works at 512×512 with a 32× compressed VAE while keeping other settings unchanged.
• Reusable tokenizers: a SpeeDiff-trained VAE can be frozen and reused for new diffusion models, useful for teams that want a shared visual tokenizer across many models.
• Strong guidance-free generation (FID 1.50 at 256×256), which reduces the inference cost of running multiple forward passes for classifier-free guidance.

Industry relevance. The reported compute reduction (436.29 → 334.98 GFLOPs) and the large convergence speedups target the dominant cost of training generative image models, which is directly relevant to organizations training diffusion models at scale.

Future Directions

• Extending beyond class-conditional ImageNet. The paper reports only ImageNet class-conditioned generation at 256×256 and 512×512; text-to-image, video, and larger-resolution regimes are not evaluated and remain open.
• Improving the TPR loss itself. The paper notes TPR is a pixel-space MSE that "can be further enhanced with perceptual losses such as LPIPS," which is not ablated in the provided content.
• Better balancing of Latent-REPA and Diff-REPA. REPA++ is a simple additive combination; the differing layer-wise effects observed in the CKNNA analysis suggest weighted or layer-aware schedules could help.
• Scaling and allocation studies. The appendix reportedly examines parameter allocation across the VAE encoder, VAE decoder, and diffusion model (Sec. B.1); the paper only studies B, L, and XL scales, so behavior at larger scales is not reported.
• Dependence on foundation models. The best results (1.50 at 256×256, 1.53 at 512×512) require a pretrained VFM (DINOv3-ViT-L/16), while the alignment-free variant reaches 2.42; closing that gap without external models is an open question.

Target Audience

Diffusion-model researchers and engineers working on latent diffusion architectures, visual tokenizers, or training-efficiency methods; practitioners at scale who need faster convergence and cheaper high-resolution generation; and graduate students with a background in generative modeling who want a clear case study in diagnosing and fixing latent-space collapse. Beginners will find the high-level motivation accessible, but the loss derivations and ablations assume intermediate-to-advanced familiarity.

Authors’ abstract

We present Scalable Pixel-anchored End-to-end Diffusion (SpeeDiff), a latent diffusion method that jointly trains the VAE and the diffusion model from scratch. In principle, joint training allows the diffusion loss gradient to directly guide the VAE encoder, encouraging the formation of a generation-friendly latent space and potentially yielding faster convergence than the conventional two-stage approach with a pretrained frozen VAE. However, a naive end-to-end implementation severely degrades performance, as unrestricted backpropagation of the diffusion loss leads to latent space collapse. Our main technical contribution is a simple yet effective Tweedie Pixel Reconstruction (TPR) loss, which provides additional pixel-level feedback by decoding a predicted clean latent from an intermediate noisy state using Tweedie's formula, thereby alleviating collapse. Furthermore, our method enables jointly scaling a fully transformer-based architecture and enhances representation alignment within the end-to-end framework. Our SpeeDiff-XL model achieves over 140x and 61x faster training compared to Vanilla SiT and REPA, respectively, while attaining an FID of 1.50 without guidance on ImageNet 256x256 generation. With a more efficient 32x compressed VAE, our model further reaches an FID of 1.53 without guidance on ImageNet 512x512 generation.