Towards Controllable Image Generation through Representation-Conditioned Diffusion Models

Technical Summary: Towards Controllable Image Generation through Representation-Conditioned Diffusion Models

Problem Statement
While diffusion models have revolutionized high-quality image synthesis and editing, controlling them to produce specific outputs remains a significant challenge. Unlike Generative Adversarial Networks (GANs), which offer well-structured and disentangled latent spaces facilitating precise, fine-grained modifications, diffusion models lack an equally effective latent space for direct control. Existing conditioning methods, such as text prompts or semantic maps, often rely on extensively annotated datasets or complex prompt engineering, which can limit precision and adaptability to unstructured data. Furthermore, while representation-conditioned diffusion models have been explored to improve unconditional generation quality, their potential for providing a controllable latent space has not been thoroughly tested.

Methodology
The authors propose a Representation-Conditioned Diffusion Model (RCDM) that leverages a pre-trained self-supervised model to guide the generation process. The architecture consists of two primary components (illustrated in Figure 1):

1. Representation Extraction: A pre-trained self-supervised encoder (specifically DINO) maps input images to a lower-dimensional representation distribution.
2. Conditional Generation: A denoising U-Net (based on the Latent Diffusion Model, LDM) is trained to map a noise distribution to the image distribution, conditioned on the representations extracted by the encoder.

The training pipeline involves two stages: first, mapping images to the representation space using the SSL encoder; second, training the diffusion model to denoise while conditioned on these representations. The authors adapt methods from prior work (using MoCo v3 and MAGE) but apply them within the LDM framework conditioned on DINO representations. Experiments were conducted on the LSUN Churches and Celeb-A datasets.

Key Contributions and Experiments
The paper presents initial experiments to validate the effectiveness of the representation space for controllable generation, focusing on three key properties:

• Perturbation Robustness: The authors tested the stability of the model by perturbing representation vectors with Gaussian noise ($\hat{C} = C + \lambda\epsilon$). Results indicate that while Diffusion Inversion suffers significant quality degradation at higher noise levels ($\lambda > 0.4$), the RCDM maintains image quality and consistency while modifying content.
• Interpolation Smoothness: Linear interpolation between two representation vectors ($\hat{C} = \alpha C_1 + (1-\alpha)C_2$) was performed. The RCDM demonstrated semantically smoother transitions between reference images compared to Diffusion Inversion (which became indistinguishable in the middle) and Stable Diffusion (which exhibited sudden transitions).
• Semantic Direction Discovery:
  • Supervised: By computing the mean representation of specific attributes (e.g., blonde hair in Celeb-A) and adding it to a reference vector, the model successfully modified features. However, the authors noted that dataset correlations (e.g., baldness correlating with male gender) could influence unintended attribute changes.
  • Unsupervised: Using Principal Component Analysis (PCA) on the representation space, the authors identified directions corresponding to semantic variations such as forehead size, hair length, and background color. While these directions did not achieve the same level of disentanglement as GANs, they demonstrated the feasibility of unsupervised direction discovery.

Results and Claims
The paper claims that representation-conditioned diffusion models exhibit promising properties in their conditioning space, specifically smoothness and a degree of disentanglement. The experiments suggest that using representations from a pre-trained self-supervised model not only improves the quality of unconditional image generation but also establishes a latent space suitable for controlled generation.

Significance
The authors position this work as a preliminary exploration into a promising direction for generative modeling. They argue that representation conditioning offers a compelling approach to enhancing both the quality and controllability of diffusion models without the heavy reliance on annotated datasets required by text-based conditioning. The paper concludes that future work could focus on adapting representations to further improve disentanglement and applying these findings to image editing and data augmentation in unconditional diffusion models.