PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Bingliang Zhang, Wenda Chu, Yizhuo Li, Linjie Yang, Yisong Yue, Katherine L. Bouman, Yang Song, Qiushan Guo

arXiv:2610.36419v1cs.CV

TL;DR

End-to-end optimization with direct diffusion-training feedback has received less attention, and existing approaches lack effective mechanisms to support it. SpeeDiff jointly trains the VAE and diffusion model using Tweedie pixel reconstruction feedback and aligns both latents and intermediate diffusion features with a pretrained VFM. SpeeDiff with REPA++ achieves FID 1.50; on ImageNet 512×512, SpeeDiff-XL with a 32× compressed VAE achieves FID 1.53 without guidance.

  • Problem

    End-to-end optimization with direct diffusion-training feedback has received less attention, and existing approaches lack effective mechanisms to support it.

    Source: §1 · p. 2

  • Method

    SpeeDiff jointly trains the VAE and diffusion model using Tweedie pixel reconstruction feedback and aligns both latents and intermediate diffusion features with a pretrained VFM.

    Source: §1 · p. 2, §3.4 · p. 5

  • Results

    SpeeDiff with REPA++ achieves FID 1.50; on ImageNet 512×512, SpeeDiff-XL with a 32× compressed VAE achieves FID 1.53 without guidance.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that SpeeDiff sets new state-of-the-art results for both non-alignment and VFM-aligned latent diffusion models.

    Source: §5 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

Latent diffusion typically relies on a separately pretrained, frozen VAE, while prior VFM-based improvements do not address end-to-end feedback from diffusion training. SpeeDiff jointly trains the VAE and diffusion model, using pixel reconstruction feedback and representation alignment to support this approach.

1. Introduction

  • Prior VFM-based VAE improvements enhance efficiency and generative quality, but end-to-end diffusion feedback remains less explored and lacks effective supporting mechanisms.Vanilla joint training has been found to degrade overall performance.Source: §1 · p. 2
  • Naive diffusion feedback can collapse the latent space; Tweedie Pixel Reconstruction decodes a predicted clean latent for pixel-level feedback, alleviating collapse.The authors report that this feedback anchors the VAE and supports robust, well-converged end-to-end training.Source: §1 · p. 2
  • REPA++ aligns both latent representations and intermediate diffusion features with VFM representations, and is optimized end-to-end with the VAE and diffusion model.The architecture refinements also replace the conventional convolutional VAE with a ViT-VAE and modernize the DiT backbone.Source: §1 · p. 2
  • Adding REPA++ lowers FID from 2.42 without alignment to 1.50; separately, SpeeDiff-XL achieves FID 1.53 on ImageNet 512×512 without guidance.The 512×512 result uses an efficient 32× compressed VAE.Source: §1 · p. 2
  • SpeeDiff replaces the two-stage pipeline with single-stage training of the VAE and diffusion model from scratch, without stop-gradient operations.The authors report faster diffusion convergence and improved scaling behavior.Source: §1 · pp. 1–2

2. Background

The background describes VAEs as image tokenizers for compact latent diffusion and reviews VFM-based representation alignment and direct representation use. SpeeDiff adopts alignment in its end-to-end training pipeline.

2. Background

  • VAE tokenizers encode images into continuous latents and decode them for reconstruction; their objectives include reconstruction and variational KL losses, often complemented by perceptual or adversarial losses.The encoder defines a latent distribution conditioned on the image, and the decoder reconstructs the image from its latent.Source: §2 · p. 2
  • Latent diffusion operates in compact tokenizer-produced spaces to reduce the computational demand of pixel-space diffusion, and this work uses the stochastic-interpolant formulation.The interpolant blends a latent embedding with Gaussian noise, while flow matching trains a model to approximate the velocity field.Source: §2 · p. 2
  • SpeeDiff’s pipeline combines VAE reconstruction, diffusion flow matching, Tweedie pixel reconstruction, and REPA++ alignment branches.The figure describes jointly training the VAE and diffusion model from scratch without stop-gradient operations.Source: Figure 2 · p. 3
  • VFM integration includes representation alignment of latents or diffusion features and direct use of pretrained VFM representations; SpeeDiff adopts alignment as a flexible plug-and-play strategy.Prior alignment methods include cosine-similarity objectives and backpropagation of alignment loss through the diffusion model to the VAE encoder.Source: §2 · p. 3

3. Our Approach

SpeeDiff jointly trains the VAE and diffusion model from scratch, adding pixel-level reconstruction feedback to prevent latent collapse. Refined transformer architectures and REPA++ alignment further improve efficiency and generation quality.

3.1. Vanilla End-to-end Training

  • Vanilla end-to-end training reaches FID 33.95 on ImageNet 256×256 after 80 epochs, while its latent representation collapses and loses information needed for accurate pixel reconstruction.The latent error is nearly zero despite substantial pixel error; the latent distribution also develops sharp peaks, large channel biases, and suppressed variances.Source: §3.1 · pp. 3–4, Figure 3 · p. 4

3.2. Tweedie Pixel Reconstruction Loss

  • Adding TPR loss reduces FID on ImageNet 256×256 from 33.95 to 5.79 after 80 epochs and prevents latent collapse.TPR decodes the Tweedie estimate of the clean latent and penalizes pixel-space discrepancy from the clean image.Source: §3.2 · p. 4

3.3. Scaling with Refined Architecture

  • After 80 epochs, SpeeDiff without REPA++ reduces training cost from 436.29 to 334.98 GFLOPs and achieves FID 3.66 on ImageNet 256×256.This variant uses ViT-VAE and Refined-DiT with the TPR loss.Source: §3.3 · p. 4

3.4. Representation Alignment via REPA++

  • REPA++ combines alignment of clean latents and intermediate diffusion features to the same pretrained VFM representation, providing semantic supervision to both parts of the end-to-end model.The latent and intermediate features are each mapped through an MLP and aligned using patch-wise cosine similarity.Source: §3.4 · p. 5
  • FID reaches 1.69 on ImageNet 256×256 after 80 epochs with SpeeDiff using REPA++.REPA++ aligns latent and intermediate diffusion features with pretrained VFM representations.Source: §3.4 · p. 5

4. Experiments

The experiments compare SpeeDiff variants and examine how end-to-end training affects diffusion convergence and VAE latents. They also assess representation-alignment components and latent-channel choices.

4.3. Impact of E2E Training for Diffusion

  • SpeeDiff achieves state-of-the-art unguided performance in both representation-alignment settings on ImageNet 256², with gFID 1.69 after 80 epochs when using REPA++.At 10 epochs, the REPA++ model reaches gFID 7.36, surpassing vanilla SiT trained for 1,400 epochs.Source: Table 2 · p. 6, §4.3 · p. 7

4.1. Setup

  • The experiments use ImageNet generation benchmarks, compare unguided results, and include a 512² benchmark with a more efficient 32×-downsampled VAE.Training uses three model scales and 32 latent channels; sampling defaults to 250 NFEs.Source: §4.1 · p. 5, Table 2 · p. 6, Table 3 · p. 6

4.2. Component-wise Comparison

  • With TPR loss preventing latent collapse, end-to-end optimization consistently outperforms both a two-stage LDM baseline and a jointly trained stop-gradient baseline.The comparison isolates diffusion-gradient feedback to the VAE encoder.Source: §4.2 · p. 7

4.3. Impact of E2E Training for Diffusion

  • SpeeDiff shows a clearer, more monotonic scaling trend than prior approaches, and jointly scales its VAE and diffusion model rather than keeping the VAE fixed.The authors report that suitable parameter allocation across the encoder, decoder, and diffusion model improves performance.Source: §4.3 · pp. 7–8

4.4. Impact of E2E Training for VAE Latent

  • End-to-end diffusion-gradient feedback improves latent linear-probing accuracy on ImageNet by 9.45% and strengthens semantic organization.A frozen SpeeDiff-trained VAE also supports a new diffusion model that converges at a similar speed, suggesting the latent space is generation-friendly.Source: §4.4 · p. 8

4.5. Ablations

  • REPA++ combines latent and intermediate-feature alignment, while 32 latent channels provide the best overall performance across the tested settings.Latent-REPA improves VAE latent alignment but transfers weakly to the diffusion model; Diff-REPA improves intermediate-feature alignment, with effects diminishing in deeper layers.Source: §4.5 · p. 8, Table 6 · p. 8

5. Conclusion

The conclusion identifies latent collapse as the primary cause of vanilla end-to-end training failure and presents TPR as a stable route for diffusion-gradient supervision of the VAE. Architectural refinements and REPA++ further improve convergence and generation quality.

5. Conclusion

  • The authors identify latent-space collapse as the primary cause of vanilla end-to-end training failure.They report that TPR establishes a stable optimization pathway for direct diffusion-gradient supervision of the VAE.Source: §5 · p. 8
  • Architectural refinements and REPA++ further improve convergence speed and generative quality.Source: §5 · p. 8
  • SpeeDiff reports state-of-the-art results for both non-alignment and VFM-aligned latent diffusion models.Source: §5 · p. 8

A. Detailed Implementation of SpeeDiff

The appendix specifies SpeeDiff’s transformer VAE and diffusion backbone, training losses, and sampling configuration. Its pipeline combines reconstruction, flow matching, TPR, and optional REPA++ alignment.

A.1.1. ViT-VAE

  • The ViT-VAE replaces the convolutional VAE; its encoder produces 32-dimensional latent representations, and its decoder reconstructs pixels with a final identity-initialized convolution.Patch size is 16 for 256² generation and 32 for 512² generation.Source: §A.1.1 · p. 12

A.1.2. Refined-DiT

  • Refined-DiT uses a shared global modulation layer with learnable per-block biases to reduce parameters and computation while preserving generative performance.The appendix reports that the ViT-VAE is more scalable and memory-friendly than the baseline architecture.Source: §A.1.2 · p. 12, Table 7 · p. 12

A.2.5. REPA++

  • REPA++ aligns both the VAE latent and intermediate diffusion features with VFM representations using two MLP adapters.Source: §A.2 · p. 13, §A.2.5 · p. 14

A.2.1. Pipeline

  • SpeeDiff’s training pipeline combines VAE reconstruction, diffusion flow matching, TPR, and optional REPA++ alignment losses.The training pseudocode includes the data, encoder, decoder, diffusion model, and optional MLP adapters.Source: §A.2.1 · p. 13, §A.2 · p. 13

A.2.2. VAE Reconstruction

  • VAE reconstruction uses MSE, LPIPS, discriminator, and KL losses, with the KL weight chosen at the turning point that gives the best generation performance.Increasing KL weight reduces reconstruction performance while generation quality first improves and then degrades.Source: §A.2.2 · pp. 13–14

A.2.3. Diffusion Flow Matching

  • The diffusion model uses linear flow matching, interpolating the clean latent with Gaussian noise and training velocity prediction with a squared-error loss.Log-normal time sampling slightly accelerates Base-model training, but the speedup is negligible for Large and XL models.Source: §A.2.3 · p. 14

A.2.4. Tweedie Pixel Reconstruction (TPR) Loss

  • TPR estimates a clean latent from the noisy state and predicted velocity, decodes it, and applies pixel-space MSE and LPIPS losses against the input image.This adds a pixel-space anchor for the diffusion model.Source: §A.2.4 · p. 14

A.2.6. Loss Coefficients; A.3.1. AutoGuidance

  • Most loss weights default to 1.0, while discriminator and KL weights are 0.1 and 1 × 10^-5; sampling uses a separately trained smaller bad model for AutoGuidance.The bad model has half as many transformer blocks and attention heads, is paired with a frozen VAE, and is trained for 20 epochs.Source: §A.2.6 · p. 14, §A.3.1 · p. 15

A.3.2. Analysis of Unconditional Generation and CFG

The analysis attributes limited CFG gains with REPA++ to strong unconditional samples and reports larger gains from AutoGuidance. It also examines how joint VAE–diffusion scaling and layer allocation affect performance.

A.3.2. Analysis of Unconditional Generation and CFG

  • With REPA++, CFG yields limited gains over unguided generation, while AutoGuidance provides a substantial gain.The authors attribute the difference to strong null-class samples: CFG’s negative direction adds limited information, whereas AutoGuidance uses a degraded model’s conditional prediction.Source: §A.3.2 · p. 15, Table 9 · p. 16
  • Uncurated unconditional samples are shown for SpeeDiff-XL with REPA++ at 512² after 200 epochs.The samples use the null class.Source: Figure 6 · p. 15

B.1. Joint Scaling of VAE and Diffusion

  • SpeeDiff benefits from jointly scaling the VAE and diffusion model; a larger VAE provides a more expressive latent space and stronger feedback during end-to-end training.Source: §B.1 · p. 16
  • A 1:1:3 allocation of encoder, decoder, and diffusion layers gives a stable scaling trend, with roughly 40% of total layers assigned to the VAE.Encoder and decoder depths perform best when approximately balanced.Source: §B.1 · p. 16

B.2. SpeeDiff with Conventional Architecture

The refined transformer architectures improve SpeeDiff performance, while joint layer-allocation experiments identify configurations that work well across alignment settings.

B.2. SpeeDiff with Conventional Architecture

  • ViT-VAE and Refined-DiT consistently improve performance over conventional CNN-VAE and DiT variants across model scales.Source: §B.2 · p. 16
  • 1.69 gFID is the best overall ImageNet 256×256 result after 80 epochs with REPA++ and jointly scaled VAE; Refined-DiT-XL alone reaches 2.15.Without alignment, Refined-DiT improves gFID from 5.89 to 3.71 with CNN-VAE, and ViT-VAE reaches 3.66 at XL scale.Source: §B.2 · p. 16
  • A 1:1:3 encoder:decoder:diffusion layer allocation works well with and without REPA++, with about 40% of total layers assigned to the VAE.The encoder and decoder perform best with a balanced 1:1 ratio.Source: Figure 7 · p. 17

B.3. Analysis the Effect of KL Loss

The KL weight strongly shapes the latent space and downstream generation, while SpeeDiff's performance remains stable across a broad range of TPR loss weights.

B.3. Analysis the Effect of KL Loss

  • The KL weight balances VAE reconstruction and diffusion generation: zero weakens generation, while excessive weight degrades both reconstruction and generation.The average predicted standard deviation increases monotonically as KL weight rises; the authors select the turning point with the best generative performance.Source: §B.3 · p. 17

B.4. Sensitivity to TPR Loss Weight

  • A 4× increase in TPR loss weight, from 1.0 to 4.0, causes negligible generation-quality degradation, with gFID 6.33 vs. 6.20.Performance remains comparable across weights 0.5–4.0, and reconstruction metrics improve slightly at weight 4.0.Source: §B.4 · p. 17

B.5. Analysis of Practical Cost

Despite higher per-step compute and peak memory, SpeeDiff reaches the baseline's FID much faster, yielding a substantial time-level speedup.

B.5. Analysis of Practical Cost

  • Each end-to-end step is approximately 2× slower, with one VAE encoder call and two decoder calls in addition to diffusion.The comparison baseline uses a pretrained VAE and preprocessed latents, requiring one diffusion forward pass per step.Source: §B.5 · p. 18
  • 12.14× time-level speedup: SpeeDiff reaches the baseline's 7.93 FID in 0.17M steps, versus 4M steps for the two-stage baseline.This corresponds to a 23.52× step-level speedup; SpeeDiff also trains the VAE jointly rather than requiring prior VAE pretraining or finetuning.Source: §B.5 · p. 18, Table 13 · p. 18
  • SpeeDiff-XL uses 57.34 GB peak memory versus 26.97 GB for the two-stage baseline under the same training configuration.Source: §B.5 · p. 18

C. Additional Qualitative Samples

The appendix presents uncurated, no-guidance SpeeDiff-XL samples from ImageNet generation at both 256×256 and 512×512 resolutions.

C. Additional Qualitative Samples

Related papers

Loading 2610.36419v1…