LDM-is-AE: Latent Diffusion Model is an Auto-Encoder for End-to-End Image Generation
Zhengqiang Zhang, Lingchen Sun, Rongyuan Wu, Qiaosi Yi, Xiangtao Kong, Chaodong Xiao, Lei Zhang
cs.CV
2026-09-29
LDM-is-AE trains a tokenizer-free LDM by aligning DiT mid-features to images. ImageNet class-conditional FID is 1.80 at 256 and 1.90 at 512.
Standard latent diffusion is a two-stage pipeline. An auto-encoder is pretrained, often on OpenImages, to squash a 256×256 image into a compact latent such as 32×32. That tokenizer is then frozen and a denoiser (DiT, SiT, Stable Diffusion, FLUX) is trained inside it.
The cost is obvious: a separate tokenizer, and in strong systems a vision foundation model on top. The latent is also a poor fit. It is fitted for reconstruction and stays fixed while the generator learns. EQ-VAE and VAVAE already showed that a reconstruction-friendly space need not be generation-friendly. REPA-E, UNITE, and DSD jointly train tokenization and generation, but they still keep a separate encoder, decoder, or diffusion module.
LDM-is-AE takes a narrower cut. A DiT step already maps latent to feature and back to latent, an implicit decode-encode. Align the middle feature with image space and the standalone tokenizer can go.
A DiT expands a noisy latent zt into a high-dimensional feature F and projects it back. LDM-is-AE names the two halves:
A small MLP on each side matches F to a pixel-unshuffled image xu, with spatial resolution folded into channels. Image-space loss is pixel ℓ2 plus LPIPS, weighted by w(t)=(1-t)^2 so late, clean steps dominate and the scale tracks v-loss.
Forcing all of F into image-shaped channels starves denoising. DiT-D therefore splits its output: image-aligned F and an auxiliary F′. DiT-E reads γ(t)F+(1-γ(t))F′ with γ(t)=t^k. Near t=1, F owns the path so the clean step is a valid auto-encoder; at noisy t, F′ keeps capacity. k=3.0 wins on the small-model sweep.
DiT-E is residual: a normalized prediction added to a skip interpolated from the image-aligned feature, so it only learns a correction. Total loss is Lldm plus weighted image alignment. Clean z1 is encoded from xu by DiT-E at t=1 with gradients detached, so there is no extra reconstruction objective on that path. At t=1 the network runs z1 → F≈xu → ẑ1, which reversed is image → latent → image. With depth fixed, 2 blocks for DiT-E and 10 for DiT-D work best (FID 17.69, IS 119 on the ablation backbone). Decoding needs more capacity than encoding.
Class-conditional ImageNet, 50K images, 50-step Heun, CFG 2.2. The main model is 961M parameters, 300 epochs, 7.02×10^19 generator-only forward FLOPs. Separate AE/VFM pretraining is not in that budget.
| Method | Setup | FID↓ | IS↑ |
| DiT-XL/2 | two-stage, no VFM | 2.27 | 278 |
| SiT-XL/2 | two-stage, no VFM | 2.06 | 270 |
| REPA-E (scratch) | one-stage+DINOv2 | 1.67 | 310 |
| UNITE-XL | one-stage+separate decoder | 1.75 | 255 |
| JiT-H/16 | pixel diffusion | 1.86 | 303 |
| LDM-is-AE | one-stage generator only | 1.80 | 314 |
IS 314 is the best in the table at 256. Among methods without a VFM, FID 1.80 beats JiT-H/16 at 1.86 and PixelFlow at 1.98, and also beats PixNerd at 2.15 despite PixNerd using DINOv2. REPA-E trained from scratch reaches 1.67, but that run still has a separate AE and DINOv2. Generator FLOPs are about half of JiT-H/16 (7.02 vs 14.0) and below UNITE-XL at 12.0, which runs the generator twice per step. Two-stage systems with a VFM still win on FID: REPA-E tuned 1.12, RAE-DiT 1.13, DDT-XL/2 1.26.
At 512, LDM-is-AE records FID 1.90 and IS 320, ahead of JiT-H/32 at 1.94/309 and REPA-SiT-XL/2 at 2.08/275.
Reconstruction on ImageNet val-50K shows what the space is actually for. PSNR is 27.57, above VAVAE 26.59 and SDVAE 25.94. rFID is 0.7879, worse than VAVAE 0.2650. gFID is 1.821, better than VAVAE 2.566 and REPA-E 2.745. The reconstructions sit closer to the natural-image distribution than to the input distribution. Generation FID 1.80, reconstruction gFID 1.82, and FID of a random 50K training subset 1.74 land on top of each other. PSNR saturates around epoch 20; latent drift versus 10 epochs earlier plateaus around epoch 260. After reconstruction freezes, the space keeps moving with the denoising objective.
On JiT-B/16 without CFG for 200 epochs, vanilla FID is 30.27, LPIPS alone 27.57, full LDM-is-AE 25.95; IS goes 54.80 → 61.09 → 67.45. The same recipe on SiT-B/2 drops FID from 45.19 to 41.68. 128 latent channels work best.
For anyone trying to drop a standalone VAE, this is a reproducible one-stage recipe: no OpenImages tokenizer, no DINOv2. 961M, 300 epochs, roughly half of JiT's generator FLOPs, ImageNet FID 1.80 at 256 and 1.90 at 512. The pipeline is cheaper. It is not a new FID champion. VFM-backed two-stage models still sit lower. The paper floats the image-aligned mid-feature for controllable generation and editing; the main text has no such experiment. The rFID-up, gFID-down pattern is the actual scientific claim, a diffusion-native latent. Text-to-image is not measured in the main paper.
The authors flag residual latent prediction as a possible cap on latent-space exploration, which matches early PSNR saturation. Encoding is run with detached gradients, so reconstruction error does not flow into DiT-E on that path; joint learning is one-way there. Reported FLOPs count only generator forwards and omit AE/VFM pretraining, so cheaper is a chosen accounting. Main numbers use CFG 2.2. The 2/10 split, k, and channel width were swept on a small backbone with 10K images; whether XL wants the same split is untested. Text-to-image is deferred to an appendix with no figures in the main text. ImageNet class-conditional generation does not transfer to open-set T2I by assertion.