PaperlayerView exact revision on arXiv

Source-linked AI summary

Unified Latents (UL): How to train your latents

Jonathan Heek, Emiel Hoogeboom, Thomas Mensink, Tim Salimans

arXiv:2602.17270v1cs.LGcs.CV

TL;DR

Unified Latents addresses how to regularize latents that will later be modeled by diffusion, where conventional VAE-style regularization makes bitrate difficult to reason about. It jointly trains an encoder, diffusion prior, and diffusion decoder with aligned noise objectives, achieving reported gains in generation quality and training efficiency while making the reconstruction–modeling trade-off explicit. The approach remains computationally costly at sampling time without decoder distillation.

  • Problem

    Conventional VAE-style KL regularization makes latent information content difficult to control, while latent information density trades off against reconstruction quality and diffusion-modeling difficulty.

  • Method

    Unified Latents jointly trains an encoder, diffusion prior, and diffusion decoder, fixing encoder noise and aligning it with the prior’s minimum noise level.

  • Results

    Unified Latents outperforms existing methods in training efficiency and generation quality while providing stable, interpretable control over latent information through simple hyper-parameters.

  • Takeaways & Limitations

    Unified Latents makes the reconstruction–modeling trade-off explicit and offers a principled approach to latent design for scaling latent diffusion models.

  • Takeaways & Limitations

    Diffusion decoders are an order of magnitude more expensive to sample from than GAN decoders, making Unified Latents significantly costlier without decoder distillation.

Abstract

from arXiv · show

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.

1. Introduction

Unified Latents addresses the difficulty of regularizing latents for diffusion models by co-training a diffusion prior and decoder. It systematically navigates the information–reconstruction trade-off through fixed-noise encoding and aligned diffusion objectives.

  • Motivation: Diffusion latents improve scalability for high-resolution generation, but their information content remains difficult to control with conventional regularization.The original Latent Diffusion Model uses a VAE-style KL penalty whose weight must be set manually because the decoder lacks a likelihood-based loss.
  • Motivation: Lower-information latents are easier to model but can lose high-frequency details, producing worse PSNR or reconstruction artifacts.The paper frames latent design as a trade-off between information content, reconstruction quality, and generation performance.
  • Motivation: Latent channel count determines capacity: fewer channels ease modeling at reconstruction cost, while more channels improve reconstruction but demand greater modeling capacity.The paper proposes navigating this trade-off systematically.
  • Approach: Unified Latents co-trains a diffusion prior on the learned latents to regularize representations for their subsequent diffusion modeling.The approach jointly encodes, regularizes, and decodes the latent representation.
  • Approach: The method fixes encoder noise, aligns the prior with the minimum noise level, reduces the KL term to a weighted noise-level MSE, and reweights the decoder ELBO.These design choices connect latent regularization to the diffusion objectives used for prior modeling and decoding.

2. Background

The background develops variational latent modeling with diffusion-based priors and decoders. It emphasizes how diffusion losses decompose across noise levels and why latent priors require unweighted ELBOs to prevent information from concentrating in discounted regions.

  • Variational Autoencoders: Variational inference derives an ELBO for modeling images x with a latent variable z₀.The ELBO provides a principled objective for latent representation learning.
  • Variational Autoencoders: In this work, both the decoder pθ(x|z₀) and prior pθ(z₀) are learned with diffusion models.This places diffusion modeling on both the reconstruction and latent-prior terms.
  • Diffusion Models: Diffusion models learn to reverse gradual destruction processes, supporting compression, likelihood estimation, and sampling from continuous distributions.The destruction level is parameterized by a logSNR schedule.
  • Diffusion Models: Diffusion-model information requirements and losses are decomposed over noise levels, with schedule-dependent weighting used in standard image-quality-oriented training.The learned model predicts clean data from noisy inputs, while reweighted ELBOs can emphasize image quality.
  • Weighting diffusion ELBOs: A diffusion prior should use unweighted ELBO loss because weighted losses could let the encoder hide information at the most discounted noise levels.This issue arises because latent encoders can exploit noise-level weighting differently from ordinary data modeling.

3. Unified Latent Diffusion

Unified Latents trains deterministic, fixed-noise encodings with a diffusion prior and diffusion decoder, linking encoding noise to prior precision. Reweighted decoder losses and staged prior retraining balance latent information, reconstruction quality, and modeling complexity.

  • Training procedure: Unified Latents trains an encoder, diffusion prior, and diffusion decoder in stages, then retrains a base model on frozen latents.Stage one regularizes encoding and reconstruction; stage two trains a new model on the resulting latents.
  • Encoding and prior: The diffusion prior models the path from pure noise z1 to slightly noisy z0, which conditions the decoder and regularizes latent information.This construction makes the prior measure and regularize the information content of the decoder’s latent input.
  • Encoding and prior: The encoder produces a deterministic latent that is forward-noised to a fixed minimum noise level shared with the prior.The approach links encoder noise to the prior’s maximum precision instead of learning a flexible encoder distribution.
  • Decoding: The image-space diffusion decoder conditions on both noisy image xt and latent z0, using reweighting to favor decoder modeling at selected noise levels.Discounting low-noise levels encourages the decoder to model high-frequency features, while the loss factor controls the balance with the latent model.
  • Decoding: A decoder loss factor of 1.3 to 1.7 is sufficient in most experiments, while higher-information latents improve reconstruction but shift modeling complexity to the base model.The loss factor and sigmoid bias jointly control the amount of information retained in the latents.
  • Base model training: ELBO-trained priors do not produce good samples, so the prior is retrained because equal weighting of latent frequencies limits sampling performance.The base model uses the same fixed maximum log-SNR as the prior and samples the final prediction rather than z0.

4. Related Work

Prior work combines diffusion with latent decoding, priors, or both, but differs in how latents are regularized and optimized. Unified Latents jointly trains these components for generation efficiency while using deterministic fixed-noise encoding to simplify the objective.

  • Diffusion decoders: Diffusion-decoder methods use diffusion reconstruction, but several retain channel bottlenecks rather than a learned prior for latent regularization.DiffuseVAE, SWYCC, and ε-VAE represent distinct decoder-based approaches described in the paper.
  • Diffusion priors: LSGM jointly trains a diffusion prior but requires a separate encoder entropy term that introduces training instability.Unified Latents instead absorbs the encoder distribution into fixed noise in the diffusion forward process.
  • Diffusion decoders and priors: DiffAE uses diffusion for encoding and decoding, whereas Unified Latents jointly trains encoder, prior, and decoder for generation quality.The distinction is whether the latent encoder is pretrained for semantic meaning or optimized within the generative system.
  • Latent diffusion and efficient autoencoders: Latent Diffusion uses channel-bottlenecked latents and a small KL penalty but lacks a principled way to control latent information.Efficient autoencoders improve compression but do not address the interaction between autoencoder design and downstream diffusion modeling.
  • Self-supervised representations: Other recent approaches replace autoencoders with semi-supervised representations such as SigLip or Dino.These methods form a separate line of latent-representation work discussed in the related-work section.

5. Experiments

The experiments evaluate Unified Latents for pre-training efficiency and generation quality across ImageNet-512 and text-to-image settings. UL outperforms Stable Diffusion-latent and other approaches in reported perceptual quality and training-cost trade-offs, while reconstruction and text-alignment metrics expose bitrate and model-size effects.

  • Evaluation scope: UL is evaluated for the relationship between training compute and generation quality on ImageNet-512, Kinetics-600, and large-scale text-to-image and text-to-video datasets.The evaluation avoids fine-tuning stages and MS-COCO evaluations to reduce confounding factors.
  • Evaluation metrics: FID and FVD measure image and video quality, PSNR measures reconstruction fidelity, and estimated bits per dimension quantify latent information.FID is used for images, FVD for videos, and PSNR complements FID because perceptually similar reconstructions can differ from their originals.
  • ImageNet-512: UL outperforms other approaches in the ImageNet-512 training-cost versus generation-performance trade-off.The comparison counts training cost using a threefold approximation for gradient computation, while excluding autoencoder training cost.
  • Autoencoder transfer: Using an out-of-distribution autoencoder produces no significant training-efficiency difference, while in-distribution autoencoders seem slightly better for small models with low-information latents.The out-of-distribution autoencoder was trained on an internal text-to-image dataset.
  • Text-to-image: Lower AutoEncoder loss factors improve image quality by gFID, especially for smaller models, while very low factors slightly reduce text alignment.Text alignment can be improved by applying guidance.
  • Text-to-image: UL significantly outperforms pixel-diffusion and Stable Diffusion-latent baselines on perceptual quality, with slightly better text alignment.The comparison adds convolution blocks for higher resolution and does not compensate UL for additional FLOPs used by the other models.

5.4. Latent bitrate tuning

The experiments tune latent bitrate through the loss factor and decoder bias, balancing reconstruction quality against the modeling difficulty of more informative latents. This trade-off is more consequential for small base models than for larger ones.

  • Model-size dependence: For smaller base models, lower bitrates are typically optimal because more informative latents defer more modeling complexity to the base model.The combined autoencoder and base-model stack is optimized for generation FID rather than reconstruction FID alone.
  • Alternative tuning: Decoder bias and loss factor provide alternative settings with roughly equal performance-versus-latent-bitrate curves.Lower decoder biases typically require higher loss factors.
  • Loss-factor tuning: Increasing the loss factor improves reconstruction metrics rFID and PSNR but increases the bitrate of the latent encoding.The loss factor controls how much information flows through the latent.
  • Reconstruction quality: Fine details such as small text are lost for low-bitrate latents.This illustrates the reconstruction-quality boundary imposed by aggressive bitrate reduction.

5.5. Latent shape

Unified Latents remains largely insensitive to latent shape over the tested channel and spatial ranges, while reconstruction quality and generation modeling differ at constrained settings. Compared with a normal-prior VAE, its latent design achieves better reconstruction–generation trade-offs.

  • Latent channels: FID metrics are mostly insensitive to latent channel count, but reconstruction quality becomes inadequate at very low channel counts.The autoencoder cannot obtain good rFID with 8 or fewer channels.
  • Spatial downsampling: 32 latent channels work across the tested spatial downsampling factors, while 16x downsampling appears easier for the decoder than 8x at similar reconstruction quality.The reported rFID values are similar for 16x and 8x downsampling, but 16x produces lower gFID.
  • Training procedure: Training the encoder with a diffusion prior and then training the remaining components separately causes only a slight performance decrease.This provides an alternative to simultaneous training of two diffusion models.
  • Normal-prior comparison: A normal-prior VAE requires higher-bitrate latents for good reconstruction, producing latents that are harder to learn and worse gFID.The comparison links the higher bitrate to both modeling difficulty and generation quality.
  • Video generation: On Kinetics-600, the medium Unified Latents model achieves 1.3 FVD, while the small model achieves 1.7 FVD.The paper reports that the medium model's 1.3 FVD is state of the art.

5.8. Ablations

The ablations show that Unified Latents depends on jointly using a learned diffusion prior, noisy latents, and fixed encoder variance for stable, useful representations. Removing or weakening these design choices degrades bitrate estimation, reconstruction, stability, or generation quality.

  • Ablation setup: The ablation study evaluates Unified Latents components alongside a traditional VAE encoder that predicts mean and variance.The results are reported in Table 6.
  • Prior regularization: Stopping gradients from the prior forces stronger KL regularization and only achieves reasonable bitrate and gFID after reducing latent channels from 32 to 8.This ablation tests whether the prior improves and regularizes the latents.
  • Noisy latents: With λ_z(0) = 10 and σ ≈ 0.007, the prior fails to accurately model latent bitrate, while reconstructions are too low quality for a useful base model.The loss instead models most information in the decoder.
  • Training data: Training on a text-to-image dataset strongly affects rFID while generation continues to work well.The authors hypothesize that high-frequency statistical differences contribute to the rFID sensitivity.
  • Encoder variance: At λ_z(0) = 5, learned noise quickly drops to zero and training becomes unstable; at λ_z(0) = 10, the encoder injects additional noise but yields worse gFID.The authors conclude that fixed encoder variance improves stability and performance.

6. Discussion

The discussion presents Unified Latents as a principled way to jointly train latent representations, diffusion priors, and diffusion decoders while exposing reconstruction–modeling trade-offs. It also identifies scope boundaries involving scaling, decoder cost, comparisons, and decoder choice.

  • Scaling: Larger base models benefit from more informative latents, but optimal bitrate scaling laws remain future work and depend on dataset, metrics, and architecture.The paper proposes studying these laws at production scale.
  • Scope: The framework is primarily evaluated on images with some video extension, although the authors suggest discrete diffusion decoders could support latent compression for text.The text application is presented as a theoretical possibility.
  • Decoder choice: Diffusion decoders model distributions, whereas GAN decoders can produce better-looking images and rFID scores through mode-collapsing training.The paper identifies decoder choice as a possible explanation for differing reconstruction metrics.
  • Comparison boundaries: Comparisons between latent diffusion models are complicated by differences in autoencoder training data, including large external datasets used by some methods.The paper's experiments mostly use ImageNet, unlike the original Stable Diffusion autoencoder.
  • Sampling cost: Diffusion decoders cost an order of magnitude more to sample than GAN decoders, making Unified Latents significantly more expensive without decoder distillation.This is a deployment cost boundary rather than a training-quality result.
  • Conclusion: Unified Latents jointly trains an encoder, diffusion prior, and diffusion decoder, outperforming existing methods in training efficiency and generation quality.The approach provides stable, interpretable control over latent information through simple hyper-parameters.

A. Additional samples

Figure 10 shows uncherry-picked text-to-image generations produced with guidance set to 2 across varied prompts.

  • Text-to-image samples: The figure presents generations for six prompts spanning people, animals, landscapes, objects, and food.The examples include rain, a misty forest, a water figure, a crocodile, a dancing fox, and ice cream.

B. End-to-end latent training

The authors also explored training the encoder, decoder, and base diffusion model end-to-end in a single stage, but this approach did not achieve FID below 2 in the first attempt.

  • The authors tested end-to-end single-stage training of the encoder, decoder, and base diffusion model.
  • The first attempt shifted the decoder diffusion loss toward noisier data and combined it with the base model’s standard ELBO loss.
  • FID remained above 2 with this end-to-end approach, despite stable joint training.
Loading 2602.17270v1…