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

Xin Cai; Zhiyuan You; Zhoutong Zhang; Tianfan Xue

cs.CVMar 23, 2026

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

Authors: Xin Cai, Zhiyuan You, Zhoutong Zhang, Tianfan Xue

Organizations: Multimedia Laboratory, The Chinese University of Hong Kong · Shanghai AI Laboratory · Adobe NextCam · CPII under InnoHK

Abstract

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.

Figures & tables

Figure 2: Overview of our method . Left: Illustration of our Detail-Aligned VAE (DA-VAE) , which encodes a high-resolution image using the same number of visual tokens as the base image. Right: Zero initialization of the linear layer for detail latent. At the beginning of training, the model keep pretrained diffusion model capability of generating images at the base resolution.

Figure 3: Effect of the proposed latent alignment loss on the learned detail feature $\mathbf{z}_{d}$ . In each pair, the right column shows training with only reconstruction loss (no alignment), and the left column shows training with our alignment loss. (a) DA-VAE on VA-VAE: the detail features (points colored by class) become more class-separable and well clustered under alignment, suggesting that $\mathbf{z}_{d}$ inherits the semantic structure of the original latent. (b) DA-VAE on SD3-VAE: alignment encourages the detail branch to capture fine-grained textures while preserving the global image layout, instead of collapsing into noisy residuals.

Figure 4: Comparison of zero initialization and random initialization . Benefiting from our zero initialization, the model starts from a well-behaved point and converges faster during training.

Method	Training Regime	Reconstruction		Token-nums (# of latent tokens)	Training Epoches	FID-50k $\downarrow$		Inception Score $\uparrow$
		AutoEncoders	rFID			w/o CFG	w/ CFG
DiT-XL † [ 30 ]	Scratch	SD-VAE (f8c4p2)	0.48	32 $\times$ 32	2400	12.04	3.04	255.3
REPA † [ 51 ]	Scratch				200	-	2.08	274.6
DiT-XL † [ 30 ]	Scratch	DC-AE (f32c32p1)	0.66	16 $\times$ 16	2400	9.56	2.84	117.5
DC-Gen-DiT-XL † [ 19 ]	Fine-tune				80	8.21	2.22	122.5
LightningDiT-XL ∗ [ 46 ]	Scratch	VA-VAE (f16c32p2)	0.50	16 $\times$ 16	80	21.79	3.98	229.7

Table 1: ImageNet

512\times 512

comparison in training regime, efficiency, and performance. Training Regime: Scratch trains the generator from random initialization for the target setting; Fine-tune starts from a pretrained generator (or a closely-related pretrained checkpoint) and adapts it to the target setting (e.g., resolution/tokenizer/architecture change). † indicates numbers are directly copied from the corresponding papers; ∗ follows the original paper’s from-scratch setting.

Figure 5: Qualitative samples from our model trained at $512\times 512$ resolution on ImageNet.

	Reconstruction				Generation
Autoencoder	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	FID-10k $\downarrow$
SD-VAE (f8c4p4)	0.48	29.22	0.13	0.79	58.17
DC-AE (f32c32p1)	0.66	27.78	0.16	0.74	35.97
VA-VAE (f16c32p2)	0.50	28.43	0.13	0.78	44.65
DA-VAE (f32c128p1)	0.47	28.53	0.12	0.78	31.51

Table 2: Performance comparison of different autoencoders . All generation models were trained from scratch.

Figure 6: Comparison between our method ( $1024\times 1024$ ) and Stable Diffusion 3.5 ( $1024\times 1024$ by $512\times 512$ upsample).

Figure 7: Comparison between our method and Stable Diffusion 3.5 on a $2048\times 2048$ resolution. Zoom in for a better view.

Method	AutoEncoders	Token-nums (# of latent tokens)	Params (B)	Throughput (img/s)	FID $\downarrow$	CLIP Score $\uparrow$	GenEval $\uparrow$
PixArt- $\Sigma$	—	64 $\times$ 64	0.6	0.40	6.15	28.26	0.54
Hunyuan-DiT	—	64 $\times$ 64	1.5	0.05	6.54	28.19	0.63
SANA-1.5	DC-AE(f32c32p1)	32 $\times$ 32	4.8	0.26	5.99	29.23	0.80
FLUX-dev	FLUX-VAE(f8c16p2)	64 $\times$ 64	12	0.04	10.15	27.47	0.67
SD3-medium	SD3-VAE (f8c16p2)	64 $\times$ 64	2.0	0.36	11.92	27.83	0.62
SD3.5-medium	SD3-VAE (f8c16p2)	64 $\times$ 64	2.5	0.25	10.31	29.74	0.63

Table 3: Comparison of our method with SOTA approaches in efficiency and performance. FID and CLIP Score are reported on MJHQ-30K (1024

\times

1024). Throughput is measured on a single A100 GPU (BF16, batch size 10). Data sources: the first five baselines (PixArt-

\Sigma

, Hunyuan-DiT, SANA-1.5, FLUX-dev, and SD3-medium) are copied from [ 44 ] under the same evaluation protocol.

Alignment	Reconstruction				Generation
loss weight $\boldsymbol{\lambda_{\text{align}}}$	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	FID-10k $\downarrow$
0.0	0.59	29.23	0.11	0.80	16.37
0.1	0.55	28.70	0.12	0.79	9.58
0.5	0.47	28.53	0.12	0.78	9.27
1.0	0.63	27.90	0.14	0.76	9.23

Table 4: Ablation on alignment-loss weight. Increasing

\lambda_{\text{align}}

slightly degrades reconstruction (higher rFID / LPIPS, lower PSNR / SSIM) but improves generation quality (lower gFID), with the best trade-off at a moderate weight.

Method	Alignment	Zero Init	Weight Scheduler	FID-10k $\downarrow$
Ours (full)	✓	✓	✓	9.27
w/o alignment	✗	✓	✓	16.37
w/o zero init	✓	✗	✓	29.73
w/o weight scheduler	✓	✓	✗	9.80

Table 5: Ablation on three components . Our full model enables all three (✓); each ablation disables exactly one component (✗).

Stage	Hyper-parameter	lightningDiT-XL [ 46 ] , Class-to-image	SD3.5-M [ 38 ] , Text-to-image
DA-VAE Training	learning rate	1e-4	1e-4
	batch size	128	16
	training steps	100K	10K
	optimizer	AdamW, betas=[0.5, 0.9]	AdamW, betas=[0.9, 0.999]
	loss weights $(\lambda_{L},\lambda_{1},\lambda_{\mathrm{adv}},\lambda_{KL},\lambda_{\mathrm{align}})$	(1.0, 1.0, 0.1, 1e-6, 0.5)	(1.0, 2.0, 0.1, 1e-7, 1.0)
DiT Fine-Tuning	learning rate	2e-4	1e-4

Table S1 : Training and sampling hyperparameters for lightningDiT-XL and SD3.5-M.

Figure S1 : DA-VAE architecture instantiated on SD3-VAE. We reuse the convolutional encoder

E_{f}

and decoder

D_{f}

blocks from SD3-VAE, but remove its original feature-to-latent and latent-to-feature heads. Instead, a lightweight downsampling module maps the shared

512\times H\times W

feature map to a more compressed detail latent

z_{d}

and a parallel base latent

z

of the same shape, while a symmetric upsampling module concatenates

(z,z_{d})

, upsamples them back to

512\times H\times W

, and feeds the result into the reused decoder backbone. This yields a higher-compression latent space with explicit base and detail channels, while keeping most of the VAE architecture intact.

Panel of Figure 14 — Figure S2 : Ablation on detail channels in the DA-VAE decoder on ImageNet. (a) Reconstruction metrics for different decoder variants. (b) Visual examples showing that randomizing or zeroing the detail latent either destroys the image or removes fine-grained details such as faces and text. Please zoom in for best view.

Figure S6 : Genearated examples by our DA-VAE enhanced SD3.5-M. Please zoom in for best view.

Explore similar work

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
DiT-IC: Aligned Diffusion Transformer for Efficient Image Compression
Mar 13, 2026Junqi Shi, Ming Lu, Xingchen Li +3Learned Image Compression Generative Image Compression
School of Electronic Science and Engineering, Nanjing University, Nanjing 210023, China
APEX: Asynchronous Prefix Denoising for Extreme Compression
Oct 9, 2026David Dinkevich, Nisan Chiprut, Yoav HaCohen +1Latent Diffusion Models Diffusion-Based Image Generation
Lightricks · Hebrew University of Jerusalem

Sep 30, 2026J/K move · Enter open · S save

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

Xu Huang, Ye Huang, Zijun Liao, Yuwei Niu+5

Peking University · Singapore University of Technology and Design

Autoencoders Diffusion-Based Image Generation Diffusion Model Acceleration T2I Diffusion Models Visual Token Compression

Open paper AskPDF arXiv

High-compression tokenizers are essential for scaling latent image generative models. However, aggressive compression creates a fundamental tradeoff between reconstruction fidelity and generation efficiency: high compression image encoder always increases the learning difficulty of diffusion training, resulting in slow model convergence. Recent representation autoencoders speed up the diffusion training by improving the latent feature's expressive capability by replacing VAE encoders with pretrained semantic encoders, yet they are typically limited to moderate compression and lose pixel-level details necessary for faithful reconstruction. To achieve both high compression and fast diffusion training, we propose DC-SAE, a Decoupled Compact Semantic Autoencoder designed for high-compression image generation with accelerated diffusion model convergence. DC-SAE consists of two key components: (1) a macro-level architecture design that leverages semantic encoders to enable higher compression ratios, and (2) a pixel-level encoder that preserves low-level details, ensuring high-fidelity image reconstruction. We empirically demonstrate that DC-SAE performs strongly on image generation tasks, achieving both compact latent representations and efficient training dynamics. Specifically, on the ImageNet dataset with

512 \times 512

resolution, DC-SAE achieves

32\times

spatial compression, with 29.79 PSNR and 3.37 gFID, substantially outperforming the previous state-of-the-art high-compression tokenizer baselines DC-AE by 13.5% and 54.9% on PSNR and gFID, respectively, maintaining comparable throughput and faster diffusion model training convergence. Beyond class-conditional generation, a

1.6

B-parameter DiT using DC-SAE achieves 0.84 on GenEval and 86.007 on DPG-Bench for text-to-image generation at

1024\times1024

resolution.

Abstract

Explore similar work

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

DiT-IC: Aligned Diffusion Transformer for Efficient Image Compression

APEX: Asynchronous Prefix Denoising for Extreme Compression