PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Autoencoders: Toward a Meaningful and Decodable Representation

Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, Supasorn Suwajanakorn

arXiv:2111.15640v3cs.CVcs.LG

TL;DR

DPM latent variables lack semantic meaning for representation learning, while GAN inversion struggles to faithfully reconstruct real images. The paper combines a learnable semantic encoder with a diffusion decoder, yielding compact, meaningful codes, near-exact decoding, efficient denoising, and downstream sampling and editing capabilities.

  • Problem

    DPM latent variables lack high-level semantics and desirable representation properties, while GAN inversion struggles to faithfully reconstruct input images.

  • Method

    A learnable encoder extracts global semantics into a non-spatial zsem vector, while a DDIM decoder models stochastic image variations and a latent DDIM models zsem distributions.

  • Results

    The resulting representation supports near-exact decoding, competitive unconditional sampling, few-shot conditional generation, real-image attribute manipulation, and improved denoising efficiency.

  • Takeaways & Limitations

    Separating semantic and stochastic information makes diffusion representations useful for interpolation, real-image editing with preserved details, and sampling from compact semantic spaces.

  • Takeaways & Limitations

    Out-of-distribution images may reconstruct well but yield semantic and stochastic subcodes that cannot be reliably interpreted or interpolated.

Abstract

from arXiv · show

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/

1. Introduction

The paper asks whether diffusion models can learn image representations that retain high-level semantics while supporting near-exact reconstruction. It proposes combining a learnable semantic encoder with diffusion-based modeling of stochastic variations, enabling representation-driven applications.

  • Motivation: Diffusion models generate realistic images but their latent variables generally lack semantic meaning and useful representation properties.In contrast, GAN inversion yields semantic codes but often fails to reconstruct real inputs faithfully.
  • Motivation: The target representation must capture both high-level semantics and low-level stochastic variations to remain decodable.The paper addresses this requirement through a two-level latent representation.
  • Approach: A learnable encoder discovers compact semantic information, while a conditional DPM decodes images and models remaining stochastic details.This separates the semantic subcode from the stochastic subcode.
  • Contributions: The resulting representation is linear, semantically meaningful, and decodable, while encoding real images without optimization and preserving original details.These properties support real-image interpolation and attribute manipulation.
  • Contributions: The framework supports competitive unconditional generation and convincing few-shot conditional generation from only a handful of labeled examples.The few-shot setup does not use the additional contrastive learning required by prior work.

2. Background

Diffusion models learn denoising processes that transform Gaussian noise into images, but their naturally obtained latent variables mainly encode stochastic image degradation rather than high-level semantics. DDIM adds a deterministic process that permits accurate inversion and reconstruction, yet its image latent remains semantically weak.

  • Diffusion models: Diffusion and score-based models learn denoising processes across noise levels to map Gaussian noise into clean image samples.The process uses successive denoising passes from a Gaussian prior.
  • Diffusion models: A Gaussian diffusion process progressively adds noise to x0, with βt controlling noise levels and αt accumulating the schedule.The reverse process approximates p(xt−1|xt) using a parameterized Gaussian distribution.
  • Latent variables: DPM forward-process variables x1:T are stochastic sequences of Gaussian image degradation and therefore contain little semantic information.This motivates mechanisms that make diffusion latent variables more meaningful.
  • DDIM: DDIM provides a deterministic generative process while preserving the original DDPM marginal distribution.It shares the DDPM objective and solution but differs in sample generation.
  • DDIM: Running DDIM backward encodes an image into xT and can reconstruct it accurately, but xT still lacks high-level semantics.Thus, DDIM alone supplies fidelity without a meaningful representation.

3. Diffusion autoencoders

The diffusion autoencoder separates image information into a learned semantic vector and a DDIM-derived stochastic code. A conditional DDIM decoder uses both codes, allowing the semantic encoder to summarize global content while the stochastic encoder preserves residual detail.

  • Architecture: A semantic encoder maps x0 to zsem, while a conditional DDIM uses zsem and xT to decode the image.The conditional DDIM also acts as the stochastic encoder that obtains xT by reversing its deterministic process.
  • Semantic encoder: The semantic code is a non-spatial d = 512 vector designed to represent global semantics rather than spatially localized information.Its design resembles a StyleGAN style vector and targets smooth latent interpolation.
  • Diffusion decoder: The decoder models pθ(xt−1|xt, zsem) and receives the combined latent z = (zsem, xT).Its noise-prediction network is implemented with a modified UNet conditioned on the semantic code.
  • Semantic encoder: The semantic encoder summarizes the image with information that helps the decoder denoise and predict the output image.Conditioning on information-rich zsem supports a more efficient denoising process.
  • Stochastic encoder: The stochastic encoder captures information left out by the limited-capacity semantic code, enabling near-detail reconstruction alongside a high-level representation.The stochastic encoder is used for encoding but is not needed during training.

4. Sampling with diffusion autoencoders

Because conditioning removes the decoder’s unconditional modeling role, the framework fits a separate latent DDIM to the distribution of semantic codes. Sampling then combines a generated semantic code with Gaussian stochastic noise before decoding.

  • Latent sampling: A latent DDIM models the distribution of zsem so the diffusion autoencoder can generate samples after conditioning the image decoder.The latent DDIM is trained with Llatent.
  • Latent sampling: The latent DDIM uses deep MLPs with skip connections because the semantic code is non-spatial data.The authors empirically found L1 loss preferable to L2 for this latent model.
  • Sampling procedure: Unconditional sampling draws zsem from the latent DDIM, unnormalizes it, samples xT ∼N(0, I), and decodes the pair.The semantic encoder remains fixed while the latent DDIM is trained post-hoc.
  • Training strategy: Post-hoc latent-DDIM training enables quick experiments with different latent models and avoids constraining the semantic code with a prior loss.This preserves the expressiveness of zsem during autoencoder training.

5. Experiments

Experiments show that the two-part latent code separates global semantics from stochastic detail, enabling smooth real-image interpolation, attribute manipulation, and accurate reconstruction with compact semantic codes.

  • 5.1. Latent code captures both high-level semantics and low-level stochastic variations: The semantic subcode zsem controls global appearance, while stochastic subcode xT changes only minor details such as hair, skin, eyes, and mouth.Varying zsem produces different people, facial shapes, illumination, and overall structures.
  • 5.2. Semantically meaningful latent interpolation: Our method produces smooth interpolations between real images while preserving details from both endpoints, unlike the non-smooth transitions of DDIM.StyleGAN interpolations are smooth, but their endpoints do not resemble the input images as closely.
  • 5.3. Attribute manipulation on real images: Linear movement in zsem changes global attributes such as gender and age or local attributes such as smile and wavy hair while preserving unrelated image details.The autoencoder trained on FFHQ generalizes to CelebA-HQ without fine-tuning, and local edits keep most of the remaining image stationary.
  • 5.4. Autoencoding reconstruction quality: The diffusion autoencoder achieves reconstruction quality competitive with NVAE while providing useful high-level semantics in a compact 512D zsem.A 64D zsem is already on par with StyleGAN2 inversion in its 512D W space, suggesting effective compression.
  • 5.7. Unconditional sampling: The method achieves FID scores comparable to D2C for few-shot conditional generation and remains competitive with DDIM for unconditional generation.The few-shot comparison does not use the self-supervised contrastive learning used by D2C.
  • 5.5. Faster denoising process: Conditioning the DDIM decoder on zsem improves image quality at small timestep counts and enables more accurate x0 prediction in fewer denoising steps.The comparison uses a DDIM baseline with matching UNet hyperparameters but without zsem conditioning.

6. Related work

Diffusion autoencoders pursue useful image representations while retaining near-perfect decoding, sharing goals with established autoencoder families and addressing limitations associated with GAN inversion.

  • Unlike standard DPM latents, the proposed representations are intended to support downstream tasks while preserving reconstruction capability.
  • VAE-based autoencoders face posterior collapse and prior-hole problems, whereas the paper states that DPMs do not.
  • Diffusion autoencoders aim to provide useful image representations that can also be decoded back near perfectly.
  • GAN image representation requires optimization-based inversion or an external encoder, which can limit reconstruction fidelity or produce high-dimensional codes outside the learned manifold.

7. Limitations & Discussion

The framework combines semantic and stochastic codes to support near-exact decoding and downstream editing, while retaining limitations in generation speed, scale-specific control, and spatial precision.

  • Non-spatial latent codes suit global semantics, but precise local or spatial reasoning may benefit from incorporating 2D latent maps.
  • Generation remains slower than GANs despite reducing the DDIM timesteps needed for high-quality samples.
  • Unlike StyleGAN, the diffusion autoencoder lacks scale-specific generation control.
  • The method separately infers semantic and stochastic information, enabling near-exact decoding and compact semantics for downstream tasks.

A.1. Latent DDIM architectures

The latent DDIM uses an MLP with skip connections to model the semantic-code distribution, with implementation choices evaluated across schedulers and training settings.

  • The latent DDIM architecture uses an MLP with skip connections and was found to be reasonably effective and fast.
  • The architecture and hyperparameters are specified in Figure 8 and Table 6.
  • On LSUN Horse, the constant 0.008 beta scheduler achieved validation FID 10.50 versus 13.36 for the linear beta scheduler.
  • Similar latent-DDIM settings worked across datasets, with network depth and total training iterations tuned separately for each dataset.
  • Linear classifiers are trained on normalized zsem for attribute manipulation and class-conditional sampling.
  • Latent DDIM training uses only a fraction of the computational resources required by diffusion autoencoder training.

C. Does the latent DDIM memorize its input?

Experiments indicate that the latent DDIM does not memorize training images and that stochastic-code variation changes minor details while preserving overall structure.

  • Nearest-neighbor comparisons using LPIPS and image- and semantic-space MSE found substantially different generated images, suggesting no memorization problem.
  • With the same semantic code, varying xT produces realistic images whose changes are limited to minor details such as hair pattern while preserving overall structure.
  • The latent representation is evaluated through linear-classification AUROC on 40 CelebA-HQ attributes, comparing zsem with StyleGAN’s inverted W-space code.

F. Real-image interpolation results

The method produces smooth real-image interpolations across faces, bedrooms, and horses, while outperforming DDIM on interpolation smoothness.

  • The method interpolates real images from FFHQ, LSUN-Bedroom, and LSUN-Horse datasets.Examples include people with and without glasses, bedrooms with different styles and angles, and horses with different poses.
  • The method handles challenging morphs between different identities, bedroom views, and horse poses.
  • Perceptual Path Length measures perceptual change along interpolation paths using a small step ϵ = 10^-4.The metric averages perceptual differences over sampled latent-code pairs and interpolation positions.
  • The method significantly outperforms DDIM in interpolation smoothness.PPL is computed over 200 samples, corresponding to 400 FFHQ images.

G. Real-image attribute manipulation results

Real-image attribute manipulation changes smiling, hair, age, and gender while retaining original content, and compares favorably with StyleGAN-W on the supplied FID evaluation.

  • The method manipulates smiling, wavy hair, aging, and gender attributes on real FFHQ and CelebA-HQ images.The generated results are described as highly realistic and plausible.
  • Manipulation moves z_sem linearly along a target attribute direction and selects an image-specific step using classifier confidence.The target confidence is set to the median confidence of real positive images.
  • The manipulated images are closer to real positive images than StyleGAN-W for 4 out of 5 attributes and preserve original content for all 5 attributes.FID comparisons use real positive and real negative images as reference distributions.
  • Identity preservation is evaluated with cosine similarity between ArcFace embeddings of the input and manipulated images.The comparison uses four CelebA-HQ attribute classes and notes that lower StyleGAN scores are partly attributable to poor inversion.

I. Class-conditional samples

The paper demonstrates few-shot class-conditional sampling with sparse labels and reports realistic generation and manipulation capabilities alongside distributional and detection caveats.

  • Class-conditional samples: Class-conditional samples for Blond, Non-blond, Female, and Male use 100 labeled examples and 10k unlabeled examples.A linear classifier is trained separately for each attribute.
  • Unconditional samples: The diffusion autoencoder produces uncurated unconditional samples on FFHQ, LSUN-Bedroom, and LSUN-Horse.
  • Out-of-distribution images: Out-of-distribution images can still be reconstructed closely, but inferred semantic and stochastic subcodes may fall outside learned distributions.For face shifts and horse images encoded by a face-trained model, noise maps retain input details and are not normally distributed.
  • Detection: Although generated images look highly realistic, another neural network can detect certain artifacts, with diffusion-model fake detection reported as currently > 90% effective.Higher T produces samples that are harder to detect.
Loading 2111.15640v3…