Shaping latent spaces improves guided image generation quality
Learning to Be Guided: Shaping Latent Geometry for Classifier-Free Guidance
Machine LearningArtificial Intelligence
Summary
Guided image generation models create images based on input conditions, but their quality depends on how these conditions are arranged internally. The authors studied how different class conditions are organized in the model's latent space and found that making class differences more distinct helps improve image quality under guidance. They introduced a technique called Prototype Regularization that better organizes this latent space, leading to clearer and more accurate generated images. This method works well across different types of conditions and datasets.
What this means in practice
- •For image generation engineers: Improve the guided image generation quality by reorganizing latent space representations with prototype regularization to reduce artifacts and boost fidelity.
- •For video synthesis developers: Apply prototype regularization to latent video representations for better guided video generation with clearer condition distinctions.
Authors
Haoming Liu, Junhao Zhu, Shaoyi Zheng, Yuanhe Guo, Jinnuo Liu, Shenji Wan, Hongyi Wen
Abstract
Classifier-free guidance (CFG) is a standard component of conditional diffusion models, and guided generation performance has become the \textit{de facto} measure of state-of-the-art models. Recent work on latent diffusion models seeks to learn latent spaces that are easier to model with diffusion. However, under CFG, generation additionally depends on how conditions are arranged in the latent space. We study this arrangement under discrete class conditions, termed \emph{class geometry}, and investigate what class geometry allows a latent space to work well with CFG. Under a Gaussian class model, we show that shifting variance from intra- to inter-class at fixed total variance makes classes more distinguishable under noise, and derive a noisy posterior objective whose per-sample descent direction is proportional to the oracle CFG correction. Based on the analysis, we propose Prototype Regularization (PR), a lightweight surrogate for the intractable noisy posterior objective, which applies a cosine prototype classifier to spatially pooled latents. With a controlled training budget (20 epochs for VAE, 80 epochs for DiT), adding PR to an LDM baseline reduces guided gFID on ImageNet from 3.31 to 2.28, surpassing a DINOv3-aligned VA-VAE (2.44) with negligible overhead. We further demonstrate that PR generalizes to pseudo-labels derived from clustering, continuous text-embedding conditions, and video latents. Together, our findings highlight class geometry as a critical factor in guided generation and take a first step toward CFG-compatible representation learning. Our project page is available at: https://maps-research.github.io/cfg-geometry/.