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

cs.CVOct 26, 2026

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

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

Organizations: ByteDance Seed · California Institute of Technology · Meta

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.

Figures & tables

Figure 2: Detailed training pipeline of SpeeDiff. The forward pass consists of four major branches: (1) Reconstruction branch: $\mathcal{L}_{\text{VAE}}$ is computed over reconstructions from the VAE encoder and decoder. (2) Diffusion branch: follows the standard stochastic interpolant and leverages flow prediction as diffusion loss. (3) Tweedie Pixel Reconstruction branch: clean latents are estimated via Tweedie’s formula, decoded, and compared against ground-truth images. (4) REPA++ branch: both latents and intermediate diffusion features are aligned with VFM representations. Notably, SpeeDiff doesn’t requires stop-gradient operation and train both VAE and diffusion from scratch.

Figure 3: Impact of end-to-end training on latent space. (a) Comparing the denoised latent with its clean counterpart in both latent and pixel space reveals that vanilla end-to-end training collapses the latent representation, while TPR loss preserves meaningful reconstruction signals. (b) KDE visualizations show that vanilla training produces highly peaked, non-Gaussian latent distributions, whereas TPR regularizes the latent space and prevents degeneracy. (c) Channel-wise statistics further indicate that vanilla training induces large per-channel biases and severely suppressed variances, whereas TPR maintains a more balanced and normalized latent representation. Results are computed over 1000 validation images. For (c), channel bias is defined as the per-channel absolute mean, and channel variance is computed as the per-channel variance across samples and spatial locations.

Diffusion Loss for VAE	Configuration	gFID $\downarrow$
Detached	Two-stage Baseline	13.02
(w/o direct gradient feedback)	Single-stage Baseline	14.21
	+ TPR Loss	11.65
	+ Refined Architecture	7.52
	+ REPA++	3.47
End-to-end	Vanilla E2E	33.95

Table 1: From baseline to SpeeDiff. We evaluate FID without guidance on ImageNet 256

\times

256 generation after 80 training epochs.

Method	VAE		Diffusion		Epochs	Generation w/o guidance					Generation w/ guidance					rFID $\downarrow$
Method	Model	Params (M)	Model	Params (M)	Epochs	gFID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$	gFID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$	rFID $\downarrow$
Non-representation Alignment Methods
MaskDiT [ 64 ]	SD-VAE	84	DiT-XL/2	675	1600	5.69	10.34	177.9	0.74	0.60	2.28	5.67	276.6	0.80	0.61	0.61
DiT [ 41 ]					1400	9.62	6.85	121.5	0.67	0.67	2.27	4.60	278.2	0.83	0.57
SiT [ 37 ]					1400	8.61	6.32	131.7	0.68	0.67	2.06	4.50	270.3	0.82	0.59
FasterDiT [ 59 ]					400	7.91	5.45	131.3	0.67	0.69	2.03	4.63	264.0	0.81	0.60

Table 2: ImageNet 256

\times

256 generation benchmark. We report detailed results for both non–representation alignment and representation alignment methods. We additionally include results at 80 epochs for convergence speed comparison. SpeeDiff achieves state-of-the-art performance without guidance in both settings (

\downarrow

lower is better;

\uparrow

higher is better).

Method	VAE			Diffusion			Epochs	Generation w/o guidance			Generation w/ guidance
Method	Model	Params (M)	FLOPs (G)	Model	Params (M)	FLOPs (G)	Epochs	gFID $\downarrow$	sFID $\downarrow$	IS $\uparrow$	gFID $\downarrow$	sFID $\downarrow$	IS $\uparrow$
U-ViT [ 2 ]	SD-VAE	84	541 + 1240	U-ViT-H/4	501	133	400	-	-	-	4.05	6.44	263.79
DiT [ 41 ]				DiT-XL/2	675	524	600	12.03	7.12	105.25	3.04	5.02	240.82
SiT [ 37 ]				DiT-XL/2	675	524	600	-	-	-	2.62	4.18	252.21
MaskDiT [ 64 ]				DiT-XL/2	675	524	800	10.79	13.41	145.08	2.50	5.10	256.27
EDM2 [ 26 ]				EDM2-XS	125	46	1600	3.53	-	-	2.91	-	-

Table 3: ImageNet 512

\times

512 generation benchmark. We employ a more efficient 32

\times

downsampled VAE while keeping all other settings consistent with the 256

\times

256 setup. SpeeDiff offers flexible VAE configurations and enables efficient high-resolution generation without sacrificing performance (

\downarrow

lower is better;

\uparrow

higher is better).

Figure 4: Selected class-conditioned samples from SpeeDiff-XL at 512 $\times$ 512 (left) and 256 $\times$ 256 (right) resolutions without guidance.

Figure 5: Analysis of latent and intermediate diffusion features. (a) Visualization of latents using Robust PCA [ 4 ] on 1000 images. (b) Linear probing accuracy on the ImageNet validation set, which shows end-to-end training with diffusion gradient feedback automatically improve semantics. (c) CKNNA evaluation of latent and intermediate diffusion features for the REPA++ components.

Configuration	gFID $\downarrow$ (10 epochs)	gFID $\downarrow$ (80 epochs)
SpeeDiff-XL Pretrained VAE	6.94	1.73
Reference : SpeeDiff-XL	7.36	1.69

Table 4: Can SpeeDiff pretrained VAE be used independently? A diffusion model trained on top of the frozen SpeeDiff pretrained VAE converges at nearly the same rate as SpeeDiff itself. Results are reported on ImageNet 256

\times

256.

Configuration	gFID $\downarrow$ (80 epochs)
REPA++ (DINOv3-ViT-L/16 [ 50 ] )	1.69
REPA++ (DINOv2-ViT-L/14 [ 40 ] )	1.76
Latent-REPA (DINOv3-ViT-L/16 [ 50 ] )	2.28
Diff-REPA (DINOv3-ViT-L/16 [ 50 ] )	2.15

Table 5: Ablation on REPA++. We evaluate different variants of REPA++ components on ImageNet 256

\times

256. REPA++ with DINOv3-L yields the best overall performance.

Panel of Table 6 — Table 6: Ablations on latent channels. Results show that 32 channels offer the best trade-off between reconstruction and generation. The results are evaluate on ImageNet 256 $\times$ 256 after 80 epochs.

Latent Diffusion Models

Latent diffusion models denoise compact representations to generate or transform images, video, audio, and other complex data efficiently.

See all 207 papers

DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment
Mar 23, 2026Xin Cai, Zhiyuan You, Zhoutong Zhang +1Diffusion-Based Image Generation Diffusion Model Compression
Multimedia Laboratory, The Chinese University of Hong Kong · Shanghai AI Laboratory · Adobe NextCam +1
SSDD: Single-Step Diffusion Decoder for Efficient Image Tokenization
Oct 6, 2025Théophane Vallaeys, Jakob Verbeek, Matthieu CordImage Tokenization Image Reconstruction
Meta Fundamental AI Research · Sorbonne University
DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence
Sep 30, 2026Xu Huang, Ye Huang, Zijun Liao +6Autoencoders Diffusion-Based Image Generation
Peking University · Singapore University of Technology and Design

Mar 23, 2026J/K move · Enter open · S save

DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment

Xin Cai, Zhiyuan You, Zhoutong Zhang, Tianfan Xue

Multimedia Laboratory, The Chinese University of Hong Kong · Shanghai AI Laboratory · Adobe NextCam+1

Diffusion-Based Image Generation Diffusion Model Compression Diffusion Models Latent Diffusion Models Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Reducing token count is crucial for efficient training and inference of latent diffusion models, especially at high resolution. A common strategy is to build high-compression image tokenizers with more channels per token. However, when trained only for reconstruction, high-dimensional latent spaces often lose meaningful structure, making diffusion training harder. Existing methods address this with extra objectives such as semantic alignment or selective dropout, but usually require costly diffusion retraining. Pretrained diffusion models, however, already exhibit a structured, lower-dimensional latent space; thus, a simpler idea is to expand the latent dimensionality while preserving this structure. We therefore propose Detail-Aligned VAE, which increases the compression ratio of a pretrained VAE with only lightweight adaptation of the pretrained diffusion backbone. DA-VAE uses an explicit latent layout: the first

C

channels come directly from the pretrained VAE at a base resolution, while an additional

D

channels encode higher-resolution details. A simple detail-alignment mechanism encourages the expanded latent space to retain the structure of the original one. With a warm-start fine-tuning strategy, our method enables

1024 \times 1024

image generation with Stable Diffusion 3.5 using only

32 \times 32

tokens,

4\times

fewer than the original model, within 5 H100-days. It further unlocks

2048 \times 2048

generation with SD3.5, achieving a

6\times

speedup while preserving image quality. We also validate the method and its design choices quantitatively on ImageNet.

Abstract

Latent Diffusion Models

Explore similar work

DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment

SSDD: Single-Step Diffusion Decoder for Efficient Image Tokenization

DC-SAE: Deep Compression Semantic Autoencoder for Faster Diffusion Convergence