Representation Learning with Diffusion Models

TL;DR

This paper introduces LRDM, a framework that enables diffusion models to learn semantically meaningful representations, improving image synthesis and interpretability without additional training, by jointly training a representation encoder with the diffusion model.

Contribution

The paper proposes a novel joint training framework for diffusion models and a representation encoder to learn meaningful latent representations for image synthesis.

Findings

LRDM achieves competitive image generation results.

LRDM learns semantically meaningful representations.

Enables faithful image reconstructions and semantic interpolations.

Abstract

Diffusion models (DMs) have achieved state-of-the-art results for image synthesis tasks as well as density estimation. Applied in the latent space of a powerful pretrained autoencoder (LDM), their immense computational requirements can be significantly reduced without sacrificing sampling quality. However, DMs and LDMs lack a semantically meaningful representation space as the diffusion process gradually destroys information in the latent variables. We introduce a framework for learning such representations with diffusion models (LRDM). To that end, a LDM is conditioned on the representation extracted from the clean image by a separate encoder. In particular, the DM and the representation encoder are trained jointly in order to learn rich representations specific to the generative denoising process. By introducing a tractable representation prior, we can efficiently sample from the representation distribution for unconditional image synthesis without training of any additional model. We demonstrate that i) competitive image generation results can be achieved with image-parameterized LDMs, ii) LRDMs are capable of learning semantically meaningful representations, allowing for faithful image reconstructions and semantic interpolations. Our implementation is available at https://github.com/jeremiastraub/diffusion.