Unified Latents (UL): How to train your latents

Jonathan Heek; Emiel Hoogeboom; Thomas Mensink; Tim Salimans

cs.LGFeb 19, 2026

Unified Latents (UL): How to train your latents

Authors: Jonathan Heek, Emiel Hoogeboom, Thomas Mensink, Tim Salimans

Organizations: Google DeepMind Amsterdam

Abstract

We present Unified Latents (UL), a framework for learning latent representations that are jointly regularized by a diffusion prior and decoded by a diffusion model. By linking the encoder's output noise to the prior's minimum noise level, we obtain a simple training objective that provides a tight upper bound on the latent bitrate. On ImageNet-512, our approach achieves competitive FID of 1.4, with high reconstruction quality (PSNR) while requiring fewer training FLOPs than models trained on Stable Diffusion latents. On Kinetics-600, we set a new state-of-the-art FVD of 1.3.

Figures & tables

Figure 1: Schematic overview of our model, include the Encoder ( $E_{\theta}$ ), the prior latent diffusion model ( $P_{\theta}$ ), and the diffusion decoder model ( $D_{\theta}$ ).

Figure 2: Unified Latents overview. An image

{\bm{x}}

is encoded to

{\bm{z}}_{\mathrm{clean}}

. A diffusion prior models the path from pure noise

{\bm{z}}_{1}

to a slightly noisy latent

{\bm{z}}_{0}

. This

{\bm{z}}_{0}

is then used by a diffusion decoder to reconstruct the image. The prior thus measures and regularizes the information content of

{\bm{z}}_{0}

Figure 3: Decoder weighting on

{\bm{{\epsilon}}}

-mse,

w_{\bm{{\epsilon}}}(\lambda_{t})=c_{\mathrm{lf}}\cdot\operatorname{sigmoid}(b-\lambda_{t})

, showing which noise levels are penalized (via a loss factor

c_{\mathrm{lf}}=1.6

in this case) and which noise levels are discounted. In theory, for weightings above

1

the latent model is preferred and for weightings below

1

the decoder is preferred. In practise, the decoder will model information even if the weighting is slightly above

1

Panel of Figure 4 — Figure 4: FID vs. training cost on ImageNet-512. UL outperforms all other approaches on base training compute versus generation equality We assume that one training iteration is three times as expensive as evaluating the model (i.e., forward pass, backprop to inputs, backprop to weights). Note that auto-encoder training cost is not included.

latents	gFID@30K	clip
UL (LF=1.5)	4.1	27.1
Pixel (no latents)	5.0	27.0
StableDiffusion	6.8	27.0

Table 1: Generation quality and text alignment for text-to-image models trained with Unified Latents, pixel diffusion (no latents), and StableDiffusion latents.

Figure 6: Image generation quality (left) and text alignment (right) against AutoEncoder Loss Factor for various base model sizes.

LF	bits/pixel	rFID@50k	PSNR	gFID (small)	gFID (medium)
1.3	0.035	0.79	25.7	1.42	1.37
1.5	0.059	0.47	27.6	1.54	1.31
1.7	0.083	0.36	28.9	1.77	1.38
1.9	0.101	0.31	29.6	2.02	1.45
2.1	0.116	0.27	30.1	2.38	1.58

Table 2: Increasing the loss factor leads to improved reconstruction metrics (rFID, PSNR) at the cost of increased bitrate in the latent encoding. For small models, the loss factor (and bits in the latent) matter a lot. For larger base models the loss factor is less sensitive.

Figure 7: Reconstruction quality vs loss factor. Fine details like small text are lost for low bitrate latents.

Figure 8: Image quality for various latent bitrates (FID vs bits/pixel) for a small model variant. Left: generation (gFID). Middle: reconstruction (rFID). Right: reconstruction PSNR. We sweep over sigmoid shift and loss factor. For each shift, we sweep over loss factors [1.5, 1.75, 2., 3., 4.].

# chan	rFID	gFID@50K
4	7.19	-
8	1.53	-
16	0.54	1.76
32	0.42	1.60
64	0.48	1.77

Table 3: FID metrics on ImageNet512 are insensitive to latent channel count. The AE is unable to obtain a good reconstruction quality (rFID) with to few channels (

\leq 8

latent shape ( $h\times w\times c$ )	rFID@50K,	gFID@50K
$64\times 64\times 32$	0.40	2.12
$32\times 32\times 32$	0.41	1.63
$16\times 16\times 32$	1.41	1.74

Table 4: FID metrics on ImageNet512 for AutoEncoders with spatial downsampling factors between 8x and 32x.

Prior	Reconstruction loss	Latent bpd	Latent bpd	rFID@50K	gFID@50K
		with prior	with base model
Diffusion	Diffusion	0.079	0.079	0.86	1.4
Diffusion	MSE	0.072	0.072	1.1	2.4
Normal	Diffusion	0.39	0.26	0.83	2.5

Table 5: Ablations on the auto-encoder training.

	bits/pixel	rFID@50k	gFID@50k
UL baseline (LF=1.5)	0.059	0.47	1.54
A. prior model	0.121	1.81	7.80
B. noisy latents	0.008	28.27	-
C. ImageNet data	0.034	1.37	1.63
D. learned variance	0.060	0.69	1.81

Table 6: Ablations study on Unified Latents components.

Panel of Figure 9 — Figure 9: FVD vs. training cost on Kinetics-600. Plotted until convergence.

Diffusion Autoencoder

See all 11 papers

End-to-End Training for Unified Tokenization and Latent Denoising
Mar 23, 2026Shivam Duggal, Xingjian Bai, Zongze Wu +5Latent Diffusion Models Image Generation
Massachusetts Institute of Technology · Adobe
UltraDiff: Transferring High-Fidelity Priors to Compressed Latent Spaces for High-Resolution Image Generation
Oct 15, 2026Jingjing Ren, Haitian Zheng, Haoyu Chen +5Latent Diffusion Models Diffusion Model Distillation
The Hong Kong University of Science and Technology (Guangzhou) · Adobe Research
How to Train Your Latent Diffusion Language Model Jointly With the Latent Space
May 8, 2026Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4Non-Autoregressive Text Generation Latent Diffusion Models
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2

Mar 23, 2026J/K move · Enter open · S save

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

Shivam Duggal, Xingjian Bai, Zongze Wu, Richard Zhang+4

Massachusetts Institute of Technology · Adobe

Latent Diffusion Models Image Generation Autoencoders

Open paper AskPDF arXiv

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.

Abstract

Diffusion Autoencoder

Explore similar work

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

UltraDiff: Transferring High-Fidelity Priors to Compressed Latent Spaces for High-Resolution Image Generation

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space