PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffuseVAE: Efficient, Controllable and High-Fidelity Generation from Low-Dimensional Latents

Kushagra Pandey, Avideep Mukherjee, Piyush Rai, Abhishek Kumar

arXiv:2201.00308v3cs.LGcs.CV

TL;DR

DiffuseVAE addresses the mismatch between diffusion models’ strong synthesis but slow sampling and VAEs’ compact latents but weaker sample quality. It combines a VAE with a conditional diffusion model, yielding controllable low-dimensional generation, improved speed–quality tradeoffs, competitive benchmark quality, and generalization to conditioning noise. Its semantics remain dependent on the first-stage VAE, especially when its coarse reconstruction is uninformative.

  • Problem

    Diffusion models lack low-dimensional interpretable latents and are slow to sample, while standard VAEs commonly provide compact latents but produce lower-quality images.

  • Method

    DiffuseVAE combines a VAE with a conditional DDPM in a two-stage framework, using VAE reconstructions and latent representations to condition diffusion generation.

  • Results

    DiffuseVAE provides controllable low-dimensional synthesis, improves the speed–quality tradeoff, achieves comparable quality to recent state-of-the-art benchmarks, and generalizes to different conditioning noise types.

  • Takeaways & Limitations

    The framework equips diffusion models with a usable low-dimensional latent space while retaining high-quality generation and enabling faster sampling.

  • Takeaways & Limitations

    Final-sample semantics depend largely on the first-stage VAE, and uninformative blurry conditioning can produce unpredictable refinements.

Abstract

from arXiv · show

Diffusion probabilistic models have been shown to generate state-of-the-art results on several competitive image synthesis benchmarks but lack a low-dimensional, interpretable latent space, and are slow at generation. On the other hand, standard Variational Autoencoders (VAEs) typically have access to a low-dimensional latent space but exhibit poor sample quality. We present DiffuseVAE, a novel generative framework that integrates VAE within a diffusion model framework, and leverage this to design novel conditional parameterizations for diffusion models. We show that the resulting model equips diffusion models with a low-dimensional VAE inferred latent code which can be used for downstream tasks like controllable synthesis. The proposed method also improves upon the speed vs quality tradeoff exhibited in standard unconditional DDPM/DDIM models (for instance, FID of 16.47 vs 34.36 using a standard DDIM on the CelebA-HQ-128 benchmark using T=10 reverse process steps) without having explicitly trained for such an objective. Furthermore, the proposed model exhibits synthesis quality comparable to state-of-the-art models on standard image synthesis benchmarks like CIFAR-10 and CelebA-64 while outperforming most existing VAE-based methods. Lastly, we show that the proposed method exhibits inherent generalization to different types of noise in the conditioning signal. For reproducibility, our source code is publicly available at https://github.com/kpandey008/DiffuseVAE.

1 Introduction

DiffuseVAE combines VAE latent representations with diffusion models to address the quality–latency tradeoff and provide controllable low-dimensional generation. It reports strong synthesis quality, rapid sampling, controllability, and robustness to conditioning noise across image benchmarks.

  • Motivation: VAEs provide low-dimensional representations but often produce blurry samples, whereas diffusion models achieve strong image synthesis quality while lacking compact latents and requiring expensive iterative sampling.These complementary limitations motivate combining the two model families.
  • Method: DiffuseVAE uses a two-stage framework in which a VAE models a conditioning signal and a DDPM models training data conditioned on that signal and its low-dimensional latent representation.Under simplifying choices, this becomes a VAE generator followed by a conditional DDPM refiner.
  • Controllable synthesis: Major sample structure is controlled through the low-dimensional VAE latent, while diffusion noise controls minor stochastic details.This separates semantic or structural control from stochastic refinement.
  • Speed and quality: 16.47 FID versus 34.36 for the corresponding DDIM model at T=10 steps on CelebA-HQ-128 demonstrates DiffuseVAE’s reported speed–quality advantage.The method can generate plausible samples in as few as 10 reverse-process sampling steps when combined with DDIM sampling.
  • Benchmark results: DiffuseVAE achieves synthesis quality comparable to recent state-of-the-art methods on CIFAR-10, CelebA-64, and CelebA-HQ while retaining a low-dimensional latent code.The framework also generalizes to different noise types in the DDPM conditioning signal.

2 Background

VAEs learn latent representations and reconstructions through probabilistic encoder–decoder modeling, while DDPMs learn to reverse a fixed forward noising process. DiffuseVAE uses these components in a two-stage process where VAE reconstructions condition diffusion generation.

  • Variational Autoencoders: VAEs learn low-dimensional latent representations with an encoder–decoder formulation and optimize an evidence lower bound on data log-likelihood.The approximate posterior and likelihood can be parameterized by neural networks using the reparameterization trick.
  • Diffusion models: DDPMs use a fixed forward Markov chain that gradually destroys data structure and a learned reverse process that denoises toward the original sample.The forward transitions use Gaussian noise governed by a variance schedule.
  • Diffusion sampling: With sufficiently large T and a suitable variance schedule, the forward process approaches an isotropic Gaussian distribution that can be sampled before running the reverse process.Sampling starts from noise with the same dimensionality as a training data point.
  • DiffuseVAE process: DiffuseVAE’s generative process first encodes x0 with a VAE to produce a reconstruction, then uses that reconstruction to condition a second-stage DDPM.The figure presents this as a two-stage generator–refiner pipeline.

3 DiffuseVAE: VAEs meet Diffusion Models

DiffuseVAE combines a VAE with a conditional diffusion model through a two-stage framework, using VAE-derived representations to condition diffusion while retaining a low-dimensional latent space. For unconditional modeling, the framework reduces to training a VAE first and a DDPM conditioned on its reconstructions, with simplifying assumptions governing the forward and reverse processes.

  • DiffuseVAE framework: DiffuseVAE factorizes the joint distribution over data, conditioning signal, latent code, and diffusion states, then trains with a surrogate posterior and an ELBO.The VAE decoder and conditional diffusion reverse process have separate parameters, while the VAE recognition network supplies the latent posterior.
  • Design choices: For unconditional modeling, the conditioning signal y is set to x0, and the reverse diffusion process is conditioned on the VAE reconstruction rather than directly on latent code z.The reconstruction is a deterministic function of z, allowing the second-stage DDPM to use it as its conditioning signal.
  • Training procedure: DiffuseVAE training reduces to sequentially optimizing a VAE on x0, then optimizing the DDPM while freezing the VAE encoder and decoder.The resulting two-stage procedure trains the first-stage VAE before the second-stage conditional diffusion model.
  • Simplifying assumptions: The framework assumes forward transitions are conditionally independent of the VAE reconstruction and latent code, while reverse transitions depend only on the VAE reconstruction.These assumptions simplify the conditional diffusion formulation and its parameterization.
  • Generator-refiner formulation: During inference, Stage-2 DDPM refinement transforms VAE-generated samples into refined outputs, using T=1000 reverse-process steps in the illustrated framework.The generator-refiner view depicts VAE samples as inputs to the second-stage diffusion model.
  • Conditional diffusion: The forward process incorporates the VAE reconstruction, and at T with a suitable noise schedule the terminal distribution is approximately N(ˆx0, I).The reconstruction therefore serves as the base measure for the reverse process during inference.

4 Experiments

Experiments show that DiffuseVAE refines low-quality VAE generations, supports latent interpolation and controllable edits, and improves sampling speed–quality tradeoffs while retaining competitive benchmark quality. The model also generalizes to alternative noisy conditioning signals, although sampling artifacts and prior mismatch remain practical limitations.

  • 4.1 Generator-refiner framework: DiffuseVAE refines blurry VAE samples into higher-quality images, supporting its generator–refiner interpretation across both model formulations.On CelebA-HQ-128, the method improves the baseline VAE FID by about eight times.
  • 4.2.1 DiffuseVAE Interpolation: VAE-latent interpolation preserves major image structure, while DDPM stochasticity can make minor details vary non-smoothly.Shared DDPM stochasticity improves consistency and enables deterministic sampling across interpolated outputs.
  • 4.2.1 DiffuseVAE Interpolation: Interpolating xT changes minor features while preserving face orientation, gender, and facial expressions, indicating distinct roles for the VAE and DDPM representations.The VAE latent primarily controls structure and diversity, whereas xT carries minor stochastic information.
  • 4.2.2 From Interpolation to Controllable Generation: Vector arithmetic in the VAE latent space enables single-attribute and composite image manipulation on CelebA-HQ-256.The demonstrated attributes include gender, age, and hair texture, with shared DDPM latents used to maintain consistent styles.
  • 4.4 State-of-the-art comparisons: DiffuseVAE is competitive with state-of-the-art image generators and outperforms most VAE-based methods on CIFAR-10, CelebA-64, and CelebA-HQ-256.Evaluation uses FID and IS, but CelebA-HQ-256 FID is computed from 10k rather than the customary 30k samples due to compute limitations.
  • 4.5 Generalization to Noisy Conditioning: A pre-trained DiffuseVAE generalizes to different noisy conditioning signals, though generated refinements can exhibit artifacts and degraded sample quality.At T = 1000, the unconditional DDPM baseline outperforms DiffuseVAE, plausibly because of VAE prior mismatch; Ex-PDE reduces this gap but does not fully close it.

5 Related Work

DiffuseVAE is positioned among diffusion, latent-space, and autoencoding approaches by combining a low-dimensional VAE representation with diffusion-based refinement. Related methods differ in latent modeling, probabilistic sampling, and conditioning design.

  • Recent diffusion models achieve strong image-synthesis performance but require expensive iterative sampling and lack low-dimensional latent representations.
  • DiffuseVAE uses a VAE latent representation to condition diffusion-based generation, while exploring controllable synthesis and generalization to noisy conditioning signals.
  • The authors describe latent diffusion as complementary because the VAE prior in DiffuseVAE could itself be modeled with a diffusion model.
  • Unlike DiffAE, DiffuseVAE has a probabilistic formulation that enables straightforward sampling and can also use external density estimators on its latent space.

6 Limitations and Discussion

The discussion identifies limitations arising from generator-refiner dependence, uninformative conditioning, correlated latent attributes, separate training, and weak conditioning mechanisms. These constraints motivate stronger VAE and diffusion designs.

  • Final-sample semantics depend largely on the coarse VAE-generated sample, so semantically poor generator outputs can persist after diffusion refinement.
  • The VAE prior–aggregated-posterior mismatch persists despite Ex-PDE estimation and remains visible as a sample-quality gap relative to unconditional DDPM.
  • Uninformative, overly blurry VAE conditioning can produce unpredictable DDPM refinements, while correlated latent directions create artifacts in controllable synthesis.
  • The work does not explore how end-to-end diffusion training affects the VAE latent space, leaving potential benefits of joint training unresolved.
  • Stronger conditioning mechanisms could reduce reliance on the stochastic DDPM sub-code and improve generalization to downstream tasks such as super-resolution and denoising.

Broader Impact Statement

The supplied passages describe broader applications and risks alongside background on DDPM forward and reverse processes. They emphasize potential use beyond images while warning that generative systems may be misused.

  • DiffuseVAE may extend beyond image modeling to modalities such as speech and text, with possible relevance to bias and privacy concerns.
  • The same generative capabilities could be misused to produce fake or misleading information, motivating responsible research and deployment.
  • DDPMs combine a forward noising process with a learned reverse denoising process for generation from an isotropic Gaussian terminal distribution.
  • The reverse process is trained to predict noise added to data at each timestep using a simplified objective associated with strong sample quality.
  • DiffuseVAE initializes diffusion from a VAE reconstruction and is motivated by improving the low-time-step speed–quality tradeoff.
  • Compared with standard unconditional DDPM/DDIM models, DiffuseVAE adds a low-dimensional latent code for downstream uses including controllable synthesis.

D Detailed Proofs

The appendix formalizes DiffuseVAE as a joint VAE–diffusion generative process with a tractable surrogate posterior and ELBO-based training objective. It also reports conditioning comparisons and modified diffusion marginals.

  • DiffuseVAE factorizes generation over an image, conditioning signal, VAE latent representation, and diffusion trajectory.
  • Because the joint posterior is intractable, the derivation introduces a factorized surrogate posterior over diffusion states and the VAE latent.
  • The resulting ELBO combines the VAE prior, decoder likelihood, and conditional diffusion-model terms.
  • The appendix derives expectations and variances for the modified forward process, including E(x1) = √1 −β1x0 + ˆx0 and Var(xt) = (1 −¯αt)I.
  • Table 7 compares FID across DiffuseVAE conditioning schemes, while Table 8 compares two-stage and end-to-end training on CIFAR-10.

E Justification of the design choices in DiffuseVAE

DiffuseVAE’s design choices favor reconstruction-based conditioning and two-stage training to improve flexibility, performance, and computational practicality.

  • Conditioning choices: The choice y=x0 is motivated by refining blurry VAE samples with a DDPM.This design makes the training-data signal itself the conditioning target for the refinement process.
  • Conditioning choices: The framework conditions the second-stage DDPM on VAE reconstructions, enabling refinement from samples produced by another model or real images.Conditioning directly on reconstructions is more flexible than conditioning only on the lower-dimensional latent code.
  • Conditioning choices: Conditioning only on VAE reconstructions slightly outperformed the variant that also used the latent code and time-step embedding on CIFAR-10.The Formulation-1 variant was trained for around 1.1M steps before comparison.
  • Training strategy: Two-stage training performed better than end-to-end training during CIFAR-10 inference and is more suitable for limited computational resources.The comparison used models trained for 400k steps.
  • Training strategy: Hyperparameters for VAE and DDPM training are listed in Table 9, with shared model and training hyperparameters across both DiffuseVAE formulations.The passage identifies Table 9 as the source of the training setup details.

F Training and Hyperparameter details

The section documents preprocessing, architectures, training resources, evaluation protocols, controllable latent manipulation, and speed–quality benchmarking for DiffuseVAE.

  • Training and preprocessing: Stage-1 VAE inputs were normalized to [0.0, 1.0], while Stage-2 DDPM inputs were scaled to [-1.0, 1.0] with random horizontal flips.The scaling procedure also applied to unconditional baselines and both DiffuseVAE formulations.
  • Training and preprocessing: VAE latent codes used 1024 dimensions for LHQ-256 and CelebA-HQ, and 512 dimensions for CIFAR-10 and CelebA-64.The VAE used residual-block architectures, while DDPM networks followed prior diffusion-model architectures.
  • Training and preprocessing: Training used Nvidia 1080Ti GPUs and cloud TPUv2-8 or TPUv3-8 systems according to dataset resolution.CIFAR-10 and CelebA-64 used GPUs; higher-resolution CelebA-HQ and LHQ models used TPUs.
  • Evaluation: FID evaluation used 10k samples for CelebA-HQ-128, 50k for CIFAR-10 and CelebA-64 comparisons, and 30k for CelebA-HQ-256.FID and IS were computed with torch-fidelity.
  • Sample-quality evaluation: Table 10 compares first-stage VAE generator samples with final DiffuseVAE refiner samples across benchmarks.The table provides quantitative support for the qualitative generator–refiner results.
  • Controllable synthesis: Attribute directions are computed from averaged positive-minus-negative VAE latent differences and applied by adding the direction to a negative-image latent.The scalar λ controls edit coarseness, with larger values usually producing coarser generations.
  • Speed–quality evaluation: The study compares DDIM and DiffuseVAE using spaced reverse-process sampling, selecting the linear or quadratic schedule with lower FID.The quadratic schedule was used for benchmarking across datasets.
  • Speed–quality evaluation: DiffuseVAE provides a better speed–quality tradeoff than standard DDPM on CelebA-64, but trails DDPM on CIFAR-10 except at T=10.The comparison uses FID as the quality measure.

G.4 State-of-the-art Comparisons

State-of-the-art comparisons examine model size, runtime, temperature scaling, and speed–quality behavior across DiffuseVAE and competing diffusion-based systems.

  • Model size and runtime: LSGM uses an order of magnitude larger VAE backbones and denoising decoders than DiffuseVAE on CIFAR-10 and CelebA-HQ-256.The comparison is based on model parameter counts.
  • Model size and runtime: LSGM consumes 5.1GB versus around 2.00GB for DiffuseVAE when generating batches of 64 CIFAR-10 samples on one Nvidia 1080Ti GPU.The lower DiffuseVAE memory use is consistent with its smaller model size.
  • Model size and runtime: 67.03s with 140 NFEs for LSGM was faster than around 103.13s with 1000 NFEs for DiffuseVAE on the same batch-generation setup.The passage hypothesizes that latent-space diffusion drives LSGM’s runtime advantage, while noting its loss of a compact latent space.
  • Temperature sampling: Temperature scaling uses a Gaussian initial DDPM latent whose standard deviation is multiplied by λ.For CelebA-HQ-256, λ=0.8 alleviated color bleeding and was used for FID reporting; other benchmarks did not use scaling.

G.6 DiffuseVAE Training Dynamics and Stability

The training-dynamics analysis addresses stability and convergence, comparing DiffuseVAE’s losses and learning curves across CIFAR-10 and CelebA-HQ-128.

  • Training stability: Hierarchical VAEs can be unstable to train and require heuristics such as spectral normalization and gradient clipping, whereas standard VAEs generally avoid these issues.The passage states that stability is not guaranteed for hierarchical VAEs even with such heuristics.
  • Loss dynamics: Figure 11 tracks KL loss, reconstruction loss, and total loss during VAE training on CIFAR-10 and CelebA-HQ-256.Rows correspond to datasets, and columns correspond to the three loss quantities.
  • Training stability: DiffuseVAE is reported to be more stable to train than corresponding hierarchical VAE and GAN-based counterparts.The stated explanation links this stability to the training behavior of the standard VAE component.
  • Convergence: CIFAR-10 models began to slightly overfit after 2000 epochs, so that checkpoint was used for subsequent analysis.CelebA-HQ-128 training stopped after the maximum compute budget of 1000 epochs.
  • Convergence: Figure 12 compares FID-versus-epoch learning curves for DiffuseVAE formulations on CIFAR-10 and CelebA-HQ-128 with T=1000 inference steps.The curves were used to assess model convergence.

G.8 Additional Samples

Additional figures illustrate DiffuseVAE samples, controllable synthesis, generator-refiner comparisons, and robustness to conditioning noise across CelebA-HQ-128, CelebA-HQ-256, and LHQ-256.

  • Additional samples: DiffuseVAE generates qualitative samples on the LHQ-256 and CelebA-HQ-256 datasets under DDIM and temperature-scaled sampling settings.The figures include LHQ-256 samples with T=1000, CelebA-HQ-256 samples with T=50 using DDIM, and additional CelebA-HQ-256 samples with T=1000 and temperature scaling factor 0.8.
  • Controllable synthesis: Controllable synthesis examples use vector addition and subtraction operations on CelebA-HQ-128 samples.Green boxes denote vector addition, while red boxes denote vector subtraction.
  • Generator-refiner framework: Composite edits are shown before and after refinement, comparing first-stage VAE outputs with corresponding DiffuseVAE samples.The comparison isolates the visual refinement provided by the generator-refiner framework.
  • Noise generalization: As Gaussian conditioning noise increases, DiffuseVAE outputs become significantly worse on CelebA-HQ-128.The figure varies the noise standard deviation σ and uses T=100 with DDIM sampling for final samples.
  • Latent control: Sharing DDPM latents between samples makes generation effectively driven by the low-dimensional VAE latent space.These CelebA-HQ-256 samples use T=1000 and a temperature scaling factor of 0.8.
Loading 2201.00308v3…