Research map / Image generation / Autoregressive and latent generation

End-to-End Training for Unified Tokenization and Latent Denoising

NeurIPS 2026 · Image generation › Autoregressive and latent generation · approach: Latent diffusion with learned tokenizers · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes UNITE, a single-stage end-to-end autoencoder where one weight-shared Generative Encoder serves as both image tokenizer and latent flow denoiser, trained from scratch without adversarial losses or pretrained encoders.

Problem

Latent diffusion models require a staged pipeline (train tokenizer, freeze, then train denoiser); joint end-to-end training of tokenizer and generator often yields degenerate solutions and typically relies on pretrained encoders like DINO for stability.

Approach

Tokenization and generation are treated as the same latent inference problem under different conditioning (full observation vs. noise plus conditions), so a single weight-shared ViT Generative Encoder performs both via two forward passes per step, with a stop-gradient on clean latents before noising.

Key results

• UNITE-B reaches FID 2.12 on ImageNet 256 and UNITE-L 1.73, with UNITE-XL at 1.75, without adversarial loss or pretrained encoders.
• UNITE-B achieves rFID 1.01 without adversarial loss, improving to 0.51 with a lightweight GAN decoder fine-tune, versus SD-VAE 0.62 and RAE 0.58.
• On QM9, UNITE-S reaches 99.37% reconstruction match and 99.71% uniqueness, surpassing ADiT tokenizer (97.20%) and ADiT uniqueness (97.76%).

Surprising findings

• Removing the stop-gradient improves the separate encoder-denoiser ablation (gFID 2.60 to 2.24 per the text) but UNITE performs best with stop-gradient, suggesting extra gradient interference under weight sharing.
• Generation metrics can keep improving while the denoising loss rises during joint training, resembling GAN-like dynamics.

Limitations

• Relies on LPIPS loss, which uses a pretrained VGG network, so it is not entirely free of pretrained components.
• Linear probing accuracy of the Generative Encoder is only around 30%, comparable to VAEs/VQGANs, so discriminative strength is limited.
• Adversarial fine-tuning was only applied to the decoder due to compute limits; joint encoder-decoder adversarial training was not explored.

Experimental setup

Datasets: ImageNet-1K QM9 MP20

Benchmarks: ImageNet 256x256 QM9 MP20

Metrics: FID IS rFID match rate RMSD validity uniqueness structural validity compositional validity

Baselines: JiT-B/16 JiT-L/16 JiT-H/16 JiT-G/16 DiT-XL/2 SiT-XL/2 REPA-SiT-XL/2 RAE ADiT EDM GeoLDM ViTok-B/16

Builds on: REPA REPA-E RAE Unified Latents JiT rectified flow

Closest papers by meaning

• Masked Autoencoders Are Effective Tokenizers for Diffusion Models (ICML 2025 · similarity 0.91)
• Boosting Generative Image Modeling via Joint Image-Feature Synthesis (NeurIPS 2025 · similarity 0.90)
• Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think (NeurIPS 2025 · similarity 0.90)
• Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models (CVPR 2025 · similarity 0.89)
• Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.89)
• SoftVQ-VAE: Efficient 1-Dimensional Continuous Tokenizer (CVPR 2025 · similarity 0.89)