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

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Duggal, Shivam, Bai, Xingjian, Wu, Zongze, Zhang, Richard, Shechtman, Eli, Torralba, Antonio, Isola, Phillip, Freeman, William T.
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866918404298899456
author Duggal, Shivam
Bai, Xingjian
Wu, Zongze
Zhang, Richard
Shechtman, Eli
Torralba, Antonio
Isola, Phillip
Freeman, William T.
author_facet Duggal, Shivam
Bai, Xingjian
Wu, Zongze
Zhang, Richard
Shechtman, Eli
Torralba, Antonio
Isola, Phillip
Freeman, William T.
contents Latent diffusion models (LDMs) enable high-fidelity synthesis by operating in learned latent spaces. However, training state-of-the-art LDMs requires complex staging: a tokenizer must be trained first, before the diffusion model can be trained in the frozen latent space. We propose UNITE - an autoencoder architecture for unified tokenization and latent diffusion. UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator via weight sharing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes: tokenization infers latents from fully observed images, whereas generation infers them from noise together with text or class conditioning. Motivated by this, we introduce a single-stage training procedure that jointly optimizes both tasks via two forward passes through the same Generative Encoder. The shared parameters enable gradients to jointly shape the latent space, encouraging a "common latent language". Across image and molecule modalities, UNITE achieves near state of the art performance without adversarial losses or pretrained encoders (e.g., DINO), reaching FID 2.12 and 1.73 for Base and Large models on ImageNet 256 x 256. We further analyze the Generative Encoder through the lenses of representation alignment and compression. These results show that single stage joint training of tokenization & generation from scratch is feasible.
format Preprint
id arxiv_https___arxiv_org_abs_2603_22283
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle End-to-End Training for Unified Tokenization and Latent Denoising
Duggal, Shivam
Bai, Xingjian
Wu, Zongze
Zhang, Richard
Shechtman, Eli
Torralba, Antonio
Isola, Phillip
Freeman, William T.
Computer Vision and Pattern Recognition
Artificial Intelligence
Graphics
Machine Learning
Latent diffusion models (LDMs) enable high-fidelity synthesis by operating in learned latent spaces. However, training state-of-the-art LDMs requires complex staging: a tokenizer must be trained first, before the diffusion model can be trained in the frozen latent space. We propose UNITE - an autoencoder architecture for unified tokenization and latent diffusion. UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator via weight sharing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes: tokenization infers latents from fully observed images, whereas generation infers them from noise together with text or class conditioning. Motivated by this, we introduce a single-stage training procedure that jointly optimizes both tasks via two forward passes through the same Generative Encoder. The shared parameters enable gradients to jointly shape the latent space, encouraging a "common latent language". Across image and molecule modalities, UNITE achieves near state of the art performance without adversarial losses or pretrained encoders (e.g., DINO), reaching FID 2.12 and 1.73 for Base and Large models on ImageNet 256 x 256. We further analyze the Generative Encoder through the lenses of representation alignment and compression. These results show that single stage joint training of tokenization & generation from scratch is feasible.
title End-to-End Training for Unified Tokenization and Latent Denoising
topic Computer Vision and Pattern Recognition
Artificial Intelligence
Graphics
Machine Learning
url https://arxiv.org/abs/2603.22283