Unified Latents (UL): How to train your latents

TL;DR

Unified Latents (UL) jointly trains an encoder, a diffusion prior, and a diffusion decoder so that the latent space used by a downstream latent-diffusion base model is regularized by a diffusion prior on its own structure rather than by a hand-tuned KL-to-Gaussian penalty. The encoder produces a deterministic latent that is then forward-noised to a fixed log-SNR, which links the encoder’s encoding precision to the prior’s minimum noise level and reduces the KL to a simple weighted MSE over noise levels. Two hyper-parameters — a loss factor on the decoder and a sigmoid bias — provide explicit, interpretable control over the bits per pixel in the latent, letting the team move the reconstruction-vs-modelability trade-off where they want it. Reported results: SOTA FVD on Kinetics-600 (1.3) and competitive ImageNet-512 FID at lower training FLOPs than equivalent models trained on Stable Diffusion latents.

Key claims

A deterministic encoder followed by forward-noising to a fixed final log-SNR makes the VAE’s KL term reduce to a weighted MSE over diffusion noise levels, giving a tight upper bound on the latent bitrate (bits/pixel) [§3.1, Eq. 3].
The encoder noise level and the prior’s minimum noise level should be tied (same log-SNR floor); this is the load-bearing design choice that absorbs the encoder distribution into the diffusion forward process and avoids the LSGM-style separate encoder-entropy term [§3.1].
Two hyper-parameters — decoder loss factor LF (typically 1.3–2.1) and sigmoid weighting bias — together control latent bitrate, and sweeping them traces out a clean rFID/PSNR vs gFID Pareto front [§3.2, Table 2].
Larger base models prefer higher-bitrate (more informative) latents; smaller base models prefer lower-bitrate latents because they can only fit simpler latents — i.e. optimal latent bitrate depends on base-model capacity [§5.3, §5.4, Table 2].
A two-stage recipe is needed for best generation: stage 1 trains encoder + prior + decoder with ELBO weighting (which is the only valid prior-training weighting), then stage 2 retrains a larger base model on the frozen encoder with sigmoid weighting [§3.3].
UL is largely insensitive to latent channel count (16–64 channels all viable at fixed downsampling) — once the channel count is above the encoder’s information cliff (≥16), the bitrate is governed by the loss factor, not the shape [§5.5, Table 3].
On ImageNet-512, UL beats both Stable Diffusion latents and a UNet pixel-space baseline at matched architecture and matched training FLOPs [§5.3, Fig. 4].
On Kinetics-600, the medium UL model reaches FVD 1.3, claimed SOTA, with a small variant already at FVD 1.7 [§5.7, Fig. 9].
On internal text-to-image, UL at LF=1.5 reaches gFID@30k of 4.1 vs 6.8 for Stable Diffusion latents and 5.0 for pixel diffusion at the same architecture family [Table 1].
Ablating the prior gradient (stop-grad on prior input, replace with hand-tuned KL) degrades gFID and forces fewer latent channels — the diffusion prior is doing real regularization work, not just bitrate accounting [§5.8, Table 6 row A].
A learned-variance encoder (classical VAE setup) is unstable in the noisy-latent regime and underperforms the fixed-variance simplification — the determinism + fixed noise is a useful constraint, not a limitation [§5.8, Table 6 row D].

Method

The encoder is a ResNet that produces a deterministic latent z̃, which is then forward-noised to a fixed log-SNR floor to produce the actual latent z. A separately trained diffusion prior model models the distribution of z from that fixed noise floor onward; because the encoder noise level is the same as the prior’s minimum noise level, the diffusion forward process absorbs the entire encoder noise distribution and the KL term collapses to a noise-weighted MSE (Eq. 3). The decoder is itself a diffusion model in pixel space conditioned on z, trained with sigmoid-reweighted ELBO and a “loss factor” LF that up-weights the decoder relative to the KL (controlling effective latent bitrate). To get good generation quality, a stage 2 then retrains a larger base diffusion model on the frozen encoder using sigmoid weighting (the prior from stage 1 is only used as a regularizer, not for sampling). Architectures: ResNet encoder, single-level ViT (8 blocks × 1024 ch) for the prior, two-stage ViT ([512,1024] ch × [6,16] blocks) for the stage-2 base model, UViT decoder. Latent shape sweep: 32 channels at 16×16 spatial works best for ImageNet-512.

Results

ImageNet-512: UL Pareto-dominates Stable Diffusion latents and a UNet pixel-space baseline on training-FLOPs vs gFID curves; same architecture trained on SD latents lags UL noticeably [§5.3, Fig. 4].
Kinetics-600: small model FVD 1.7, medium model FVD 1.3 — claimed SOTA, beating MAGVIT and W.A.L.T. on training-FLOPs vs FVD [§5.7, Fig. 9].
Text-to-image (internal, gFID@30k / CLIP): UL LF=1.5 → 4.1 / 27.1; pixel-only → 5.0 / 27.0; SD latents → 6.8 / 27.0 [Table 1].
Latent shape robustness: rFID@50k of 0.54 / 0.42 / 0.48 for 16/32/64 channels at fixed downsampling, gFID@50k 1.76 / 1.60 / 1.77 — flat from 16 channels onward [Table 3].
Bitrate sweep: at LF=1.5, 0.059 bits/pixel, PSNR 27.6, rFID 0.47, gFID@50k 1.31 on medium model; LF=2.1 doubles bitrate to 0.116, gains 2.5 PSNR, but loses 0.27 gFID on the medium model [Table 2].
Ablation (Table 6): removing the prior gradient → gFID 7.80; insufficient encoder noise → rFID collapses to 28.27; learned-variance encoder → gFID 1.81 vs UL’s 1.54 baseline.

Why it’s interesting

UL is the most disciplined account of the “what should a diffusion latent look like” question currently filed: instead of treating the VAE as a black-box compressor with a hand-tuned KL weight, it makes the encoder’s noise level a first-class hyperparameter equal to the prior’s noise floor, and turns the KL into an interpretable bitrate budget. That maps cleanly onto the team’s working hypothesis that latent design — not the diffusion model — is the binding constraint for training efficiency (see Diffusion training efficiency). It also adds a fourth distinct vector to the same Pareto-front comparison the team has been tracking: where PixelGen (PixelGen: Improving Pixel Diffusion with Perceptual Supervision) attacks the latent-vs-pixel question by replacing MSE with perceptual losses, RAEv2 (Improved Baselines with Representation Autoencoders) by swapping the encoder for frozen DINO features, and L2P (L2P: Unlocking Latent Potential for Pixel Generation) by transferring an LDM into pixel space, UL keeps the latent but redesigns the encoder/prior coupling so the latent is information-budgeted by construction. The Kinetics-600 SOTA at FVD 1.3 is the strongest video-side data point in the wiki for this whole family.