[Paper Review] Training VAEs Under Structured Residuals

Garoe Dorta, Sara Vicente+3 more|arXiv (Cornell University)|Apr 3, 2018

Industrial Vision Systems and Defect Detection28 references6 citations

TL;DR

This paper proposes a novel VAE architecture that models structured residual correlations using a learnable covariance matrix prediction network, enabling more realistic image generation and reconstruction. By incorporating structured uncertainty into the likelihood function with minimal parameter overhead, the method significantly improves likelihood and reconstruction quality over factorized Gaussian VAEs on CelebA and LSUN Churches datasets.

ABSTRACT

Variational auto-encoders (VAEs) are a popular and powerful deep generative model. Previous works on VAEs have assumed a factorized likelihood model, whereby the output uncertainty of each pixel is assumed to be independent. This approximation is clearly limited as demonstrated by observing a residual image from a VAE reconstruction, which often possess a high level of structure. This paper demonstrates a novel scheme to incorporate a structured Gaussian likelihood prediction network within the VAE that allows the residual correlations to be modeled. Our novel architecture, with minimal increase in complexity, incorporates the covariance matrix prediction within the VAE. We also propose a new mechanism for allowing structured uncertainty on color images. Furthermore, we provide a scheme for effectively training this model, and include some suggestions for improving performance in terms of efficiency or modeling longer range correlations.

Motivation & Objective

• To address the limitation of factorized Gaussian likelihoods in VAEs, which assume independent pixel-wise uncertainty and fail to model spatial structure in residuals.
• To develop a scalable method for incorporating structured covariance predictions into VAEs without drastically increasing model complexity.
• To enable effective training of VAEs with correlated likelihoods, avoiding poor local minima and improving performance on high-frequency image details.
• To extend the structured uncertainty modeling to color images using a YCbCr-based approach for better perceptual quality.

Proposed method

• Proposes a new VAE architecture that integrates a structured Gaussian likelihood by predicting the full covariance matrix from the latent code using a dedicated subnetwork.
• Uses a parameter-efficient design to limit the number of additional parameters, avoiding the $O(n_p^2)$ cost of full covariance matrices.
• Employs a YCbCr color space transformation to decouple luminance and chrominance, enabling structured uncertainty modeling on color images.
• Introduces a training scheme that stabilizes optimization by encouraging the VAE to reduce variance rather than modeling residual errors through covariance.
• Adapts and improves the covariance prediction mechanism from [7] using dilated convolutions and sparsity patterns to model long-range correlations efficiently.
• Applies a reparameterization trick to sample from the structured Gaussian likelihood during generation, ensuring differentiable training.

Experimental results

Research questions

• RQ1 Can structured residual correlations in VAE reconstructions be effectively modeled using a learnable covariance matrix prediction network?
• RQ2 Does incorporating structured uncertainty improve the likelihood and reconstruction quality of VAEs compared to factorized Gaussian likelihoods?
• RQ3 Can the proposed method maintain training stability and avoid poor local minima despite the increased complexity of covariance modeling?
• RQ4 How does the model perform on high-frequency image details such as hair and wrinkles compared to standard VAEs?
• RQ5 Can the method be extended to color images while preserving perceptual quality and computational efficiency?

Key findings

• The proposed VAE with structured residuals achieves a negative log likelihood (NLL) of $-6918 \pm 2423$ on the LSUN Churches dataset, significantly outperforming VAEs with diagonal and spherical likelihoods.
• On the CelebA dataset, the model produces sharper reconstructions and samples with high-frequency details like wrinkles and hair, which are blurred or missing in baseline VAEs.
• The model reduces reconstruction MSE to $614 \pm 286$ on LSUN, outperforming the diagonal VAE ($728 \pm 327$) and $\beta$-VAE ($800 \pm 354$), despite similar KL divergence.
• Samples from the proposed model show improved perceptual quality, with better-structured noise that captures fine textures, unlike the oversmoothed outputs of standard VAEs.
• The structured uncertainty branch successfully models complex, spatially correlated residuals without degrading training stability, even on complex scenes like church interiors.
• The method enables effective modeling of long-range correlations through dilated convolutions and sparsity patterns in the covariance prediction network.