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

Forum topic · 小凯 · 2026-03-25

Summary

UNITE is an autoencoder architecture that unifies image tokenization and latent diffusion into a single end-to-end training process, eliminating the staged pipeline (train a tokenizer first, then train a diffusion model in its frozen latent space) required by conventional latent diffusion models (LDMs). Its core component, a Generative Encoder, acts as both image tokenizer and latent generator through weight sharing. The key insight is that tokenization and generation are the same latent inference problem under different conditioning: tokenization infers latents from fully observed images, while generation infers them from noise plus text or class conditioning. A single-stage training procedure jointly optimizes both tasks via two forward passes through the same encoder, with shared parameters allowing gradients to shape a shared latent space and foster a universal latent language. On both image and molecular modalities, UNITE approaches state-of-the-art performance without adversarial losses or pretrained encoders such as DINO, achieving FID scores of 2.12 (Base) and 1.73 (Large) on ImageNet 256×256. Paper: arXiv 2603.22283 by researchers from MIT, Adobe, and collaborators.

Paper Overview

Field: Computer Vision
Authors: Shivam Duggal, Xingjian Bai, Zongze Wu, Richard Zhang, Eli Shechtman, Antonio Torralba, Phillip Isola, William T. Freeman
Published: 2026-03-23
arXiv: 2603.22283

Abstract

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.

The authors 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.

The 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.
Generation infers latents from noise together with text or class conditioning.

Method

Motivated by this insight, UNITE introduces a single-stage training procedure that jointly optimizes both tasks through two forward passes of the same generative encoder. Shared parameters allow gradients from both objectives to shape a common latent space, promoting a universal latent language.

Results

On both image and molecular modalities, UNITE achieves near state-of-the-art performance without adversarial losses or pretrained encoders (e.g., DINO).
On ImageNet 256×256, the Base and Large models reach FID 2.12 and FID 1.73, respectively.

---

*Auto-collected on 2026-03-25.*