Diffusion Autoencoders: Toward a Meaningful and Decodable Representation

TL;DR

This work introduces diffusion autoencoders, which learn a decodable, meaningful latent by splitting the representation into a semantic vector $\mathbf{z}_{\text{sem}}$ and a stochastic code $\mathbf{x}_T$. A learnable semantic encoder and a conditional DDIM decoder enable near-exact image reconstruction while supporting semantically driven edits and interpolation; a latent DDIM is trained for sampling the semantic codes, allowing both conditional and unconditional generation. The approach yields faster denoising, effective few-shot conditional sampling, and competitive unconditional generation, addressing a key limitation of diffusion models as latent representations. Overall, diffusion autoencoders demonstrate that diffusion-based models can produce compact, interpretable representations suitable for real-image manipulation and downstream tasks beyond conventional GAN inversion.

Abstract

Diffusion probabilistic models (DPMs) have achieved remarkable quality in image generation that rivals GANs'. But unlike GANs, DPMs use a set of latent variables that lack semantic meaning and cannot serve as a useful representation for other tasks. This paper explores the possibility of using DPMs for representation learning and seeks to extract a meaningful and decodable representation of an input image via autoencoding. Our key idea is to use a learnable encoder for discovering the high-level semantics, and a DPM as the decoder for modeling the remaining stochastic variations. Our method can encode any image into a two-part latent code, where the first part is semantically meaningful and linear, and the second part captures stochastic details, allowing near-exact reconstruction. This capability enables challenging applications that currently foil GAN-based methods, such as attribute manipulation on real images. We also show that this two-level encoding improves denoising efficiency and naturally facilitates various downstream tasks including few-shot conditional sampling. Please visit our project page: https://Diff-AE.github.io/

Figures (22)

• Figure 1: Attribute manipulation and interpolation on real images. Diffusion autoencoders can encode any image into a two-part latent code that captures both semantics and stochastic variations and allows near-exact reconstruction. This latent code can be interpolated or modified by a simple linear operation and decoded back to a highly realistic output for various downstream tasks.
• Figure 2: Overview of our diffusion autoencoder. The autoencoder consists of a "semantic" encoder that maps the input image to the semantic subcode $(\mathbf{x}_0 \rightarrow \mathbf{z}_\text{sem})$, and a conditional DDIM that acts both as a "stochastic" encoder $(\mathbf{x}_0 \rightarrow \mathbf{x}_T)$ and a decoder $((\mathbf{z}_\text{sem}, \mathbf{x}_T) \rightarrow \mathbf{x}_0)$. Here, $\mathbf{z}_\text{sem}$ captures the high-level semantics while $\mathbf{x}_T$ captures low-level stochastic variations, and together they can be decoded back to the original image with high fidelity. To sample from this autoencoder, we fit a latent DDIM to the distribution of $\mathbf{z}_\text{sem}$ and sample $(\mathbf{z}_\text{sem}, \mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}))$ for decoding.
• Figure 3: Reconstruction results and the variations induced by changing the stochastic subcode $\mathbf{x}_T$. Each row corresponds to a different $\mathbf{z}_\text{sem}$, which completely changes the person, whereas changing the stochastic subcode $\mathbf{x}_T$ only affects minor details.
• Figure 4: Interpolation between two real images. In contrast to StyleGAN2 and DDIM, our method produces smooth and consistent results with well-preserved original details from both images.
• Figure 5: Real-image attribute manipulation results on two global attributes (gender, age) and two local attributes (smile, wavy hair) by moving $\mathbf{z}_\text{sem}$ along the positive or negative direction found by linear classifiers. The top two are from FFHQ karras_style-based_2019 and the bottom two are from CelebA-HQ karras_progressive_2018. Our method synthesizes highly-plausible and realistic results that preserve an unprecedented level of detail.