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

Forum topic · 小凯 · 2026-03-25

Summary

UNITE is an autoencoder architecture that unifies image tokenization and latent diffusion into a single-stage training process, proposed by researchers including Phillip Isola and William T. Freeman (arXiv 2603.22283). Traditional latent diffusion models require a complex staged pipeline: a tokenizer must be trained first before the diffusion model can be trained in the frozen latent space. UNITE replaces this with a Generative Encoder that serves as both image tokenizer and latent generator via 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. Joint optimization is achieved through two forward passes of the same encoder per training step, with shared parameters shaping a general latent language. UNITE reaches near state-of-the-art performance on image and molecule modalities without adversarial losses or pretrained encoders like DINO, achieving FID scores of 2.12 (Base) and 1.73 (Large) on ImageNet 256x256.

Paper Overview

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

Introduction

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.

Method: UNITE

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.

Key Insight

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.

Single-Stage Training

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

Results

On image and molecule modalities, UNITE achieves near state-of-the-art performance without adversarial losses or pretrained encoders (e.g., DINO):

ImageNet 256x256 FID 2.12 (Base model)
ImageNet 256x256 FID 1.73 (Large model)