cs.LGFeb 19, 2026
SaveUnified Latents (UL): How to train your latents
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 ( ), the prior latent diffusion model ( ), and the diffusion decoder model ( ).
Figure 2: Unified Latents overview. An image is encoded to . A diffusion prior models the path from pure noise to a slightly noisy latent . This is then used by a diffusion decoder to reconstruct the image. The prior thus measures and regularizes the information content of .
Figure 3: Decoder weighting on -mse, , showing which noise levels are penalized (via a loss factor in this case) and which noise levels are discounted. In theory, for weightings above the latent model is preferred and for weightings below the decoder is preferred. In practise, the decoder will model information even if the weighting is slightly above .
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 ( ).
| latent shape ( ) | rFID@50K, | gFID@50K |
|---|---|---|
| 0.40 | 2.12 | |
| 0.41 | 1.63 | |
| 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.
Figure 9: FVD vs. training cost on Kinetics-600. Plotted until convergence.
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.
UltraDiff: Transferring High-Fidelity Priors to Compressed Latent Spaces for High-Resolution Image Generation
High-resolution synthesis with Latent Diffusion Models (LDMs) faces a fundamental trade-off between computational cost and generative fidelity. This trade off is largely governed by the latent compression factor f (e.g. f8 denotes f=8 and f32 f=32). While aggressive compression (f32) significantly improves efficiency, it suffers from severe information loss and ambiguity compared to standard f8 models. We propose UltraDiff, a framework that reformulates this challenge as joint latent diffusion learning followed by marginal score distillation. We introduce MixTeacher, a dual branch DiT that models the joint f8-f32 distribution, utilizing shared attention to resolve the ambiguity of the compressed space with high fidelity f8 priors. Crucially, we propose vector field distillation in the native f32 space. Theoretically, this compels the student to learn the conditional expectation of the teacher trajectory, effectively filtering out compression artifacts. Across the evaluated 1K to 4K settings, UltraDiff approaches f8 level quality while preserving f32 efficiency. With 20 sampling steps on a single A100 GPU, it generates a 2K image in 6.2 seconds over 40 times faster than our f8 baseline.
How to Train Your Latent Diffusion Language Model Jointly With the Latent Space
Latent diffusion models offer an attractive alternative to discrete diffusion for non-autoregressive text generation by operating on continuous text representations and denoising entire sequences in parallel. The major challenge in latent diffusion modeling is constructing a suitable latent space. In this work, we present the Latent Diffusion Language Model (LDLM), in which the latent encoder, diffusion model, and decoder are trained jointly. LDLM builds its latent space by reshaping the representations of a pre-trained language model with a trainable encoder, yielding latents that are easy to both denoise and decode into tokens. We show that naive joint training produces a low-quality diffusion model, and propose a simple training recipe consisting of an MSE decoder loss, diffusion-to-encoder warmup, adaptive timestep sampling, and decoder-input noise. Ablations show that each component substantially impacts generation performance. On OpenWebText and LM1B, LDLM achieves better generation performance than existing discrete and continuous diffusion language models while being faster, indicating that jointly learning the latent space is a key step toward making latent diffusion competitive for text generation.