Structured Uncertainty in the Observation Space of Variational Autoencoders
Authors: James Langley
Miguel Monteiro
Charles Jones
Nick Pawlowski
Ben Glocker
Venue: Accepted by TMLR ()

Abstract: Variational autoencoders (VAEs) are a popular class of deep generative models with many variants and a wide range of applications. Improvements upon the standard VAE mostly focus on the modelling of the posterior distribution over the latent space and the properties of the neural network decoder. In contrast, improving the model for the observational distribution is rarely considered and typically defaults to a pixel-wise independent categorical or normal distribution. In image synthesis, sampling from such distributions produces spatially-incoherent results with uncorrelated pixel noise, resulting in only the sample mean being somewhat useful as an output prediction. In this paper, we aim to stay true to VAE theory by improving the samples from the observational distribution. We propose SOS-VAE, an alternative model for the observation space, encoding spatial dependencies via a low-rank parameterisation. We demonstrate that this new observational distribution has the ability to capture relevant covariance between pixels, resulting in spatially-coherent samples. In contrast to pixel-wise independent distributions, our samples seem to contain semantically-meaningful variations from the mean allowing the prediction of multiple plausible outputs with a single forward pass.

Published in Transactions on Machine Learning Research (10/2022)
Structured Uncertainty in the Observation Space of Variational Autoencoders
James Langley
jublangley@gmail.cmm com com
Deportment of Computing, Imperial College London, UK
Miguel Montciro
Department of Computing, Imperial College London, UK
Charles Jones
chartes.jonnss//88mperiddd ack
Department of Computing, Imperial College London, UK
Nick Pawlowski
upeoloosski@niicrsoff.ccom
Microsoft Research, Cambridge, UK
Department of Compuling, Imperial College London, UK
Ben Glocker
b. buglockerr@kmmeriil.cccco
Department of Compuling, Imperial College London, UK
Reviewed on OpenReview: https: // openreview. net/ forum? id= cap 7h9g5c4
Abstract
Variational autoencoders (VAEs) are popular class of deep generative models with many variants and a wide range of applications. Improvements upon the standard VAE mostly focus on the modelling of the posterior distribution over the latent space and the proper- ties of the neural network decoder. In contrast, improving the model for the observational distribution is rarely considered and lypically defaults to a pixel-wise independent calegor- ical or normal distribution. In image synthesis, sampling from such distributions produces opatiall--incohernee results with uncorrelated pixel noise, resulting in only the sample mean being somewhat useful as an output prediction. In this paper, aim to stay true to VAE theory by improving the samples from the observational distribution. We propose SOS- VAF, an alternative model for the observation space, encoding spatial dependencies via low-rank paramotorisation. We demonstrate that this new obscrvational distribution has the ability to capture relevant covariance between pixels, resulting spatially-coherent samples. In contrast pixel-wise independent distributions, our samples seem contain semantically- meaningful variations from the mean allowing the prediction of multiple plausible outputs with single forward pass.
1 Introduction
Generative modelling one of the cornerstones of modern machine learning. One of the most used and widespread classes of generative models is the Variational Autoencoder (VAE) (Kingma & Welling, 2014; 2019). VAEs explicitly model the distribution of observations by assuming latent variable model with low- dimensional latent space and using a simple parametric distribution in observation space. Using a neural network, VAEs decode the latent space into arbitrarily complex observational distributions.
Despite many improvements on the VAE model, one oveen-overlooked aspect is the choice of observational distribution. As an explicit likelihood model, the VAE assumes a distribution obscrvation space using delta distribution would not allow gradient based optimisation. Most current implementations, however, employ only simple models, such pixel-wise independent normal distributions, which eases optimisation

---

Published in Transactions on Machine Learning Research (10/2022)
but limits expressivity. Else, the likelihood term is often replaced by reconstruction loss which, in the case of an L2 loss, implicitly assumes an independent normal distribution. Following this implicit assumption, samples are then generated by only predicting the mean, rather than sampling in observation space.
An application where this disconnect becomes apparent is image syn- thesis. The common choices for observational distributions are pixel- wise independent categorical or normal distributions. For pixel-wise independent distributions, regardless of other model choices, sam- pling from the joint distribution over pixels will result in spatially- incoherent samples due to independent pixel noise (cf. Figure 1). To address this problem, researchers use the predicted distributions to calculate the log- log-liielhhod in favour of the mean when generating samples or reconstructing imputs. However, this solution is akin to ignoring the issue rather than altempting to solve it.
the objective but then discard them
In this work, explore what happens when we strictly follow VAE theory and sample from the predicted observational distributions. We illustrate the problem of spatial incoherence that arises from using pixel-wise independent distributions. We propose SOS-VAE, using a spatially dependent joint distribution over the observation space and compare to the previous scenario. We further compare the samples to the mean of the predicted observational distribution, which is typically used when synthesising images. We note that, in this work, we are not focusing on absolute image quality. Instead, we aim to point to an issue often overlooked in VAE theory and application, which affocts most state-of-the-art methods. Thus we analyse the relative difference between using and not. using a joint pixel dependent observational distribution for basic VAE. Yet, our findings are of broad relevance and our proposed observational distribution can be used in more advanced VAE variants, such the NVAE (Vahdat & Kautz, 2020).
Figure 1: Top: Samples generated with a 3-VAE exhibiting pixel-wise in- dependent noise. Bottom: Samples with our SOS-VAE realistic and spatially coherent.
2 Related Work
Modern generative models can be divided into two classes, implicit likelihood models, such as Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) and diffusion models (Song & Durkan, 2021), and explicit likelihood models, such as VAËs (Kingın & Welling, 2014; Rezende et al., 2014; Kingma & Welling, 2019), flow models (Dinh et al., 2017; Kingma & Dhariwal, 2018; Dinh et al., 2015) and anto-regressive models (Van Den Oord et al., ba;;;;;; Salimans et al., 2017). Despile implicit likelihood models having achieved impressive results in terms of sample quality without the need for explicitly modelling the observation space (Karras et al., 2020; Brock et al., 2019), interest in explicit likelihood models have prevailed due to their appealing properties, such as ease of likelihood estimation.
One of the most popular and successful explicit likelihood models is the VAE (Kingma & Welling, 2014; Rezende et al., 2014; Kingma& & Welling, 2019). Since its introduction, there have been numerous extensions. For example, Van Den Oord et al. (2017); Razavi et al. (2019) quantize the latent space to achieve better image quality; Higgins et al. (2017); Chen et al. (2017) modify the latent posterior to obtain disentangled and interpretable representations; and Vahdat & Kautz (2020); Sonderby et al. (2016) use hierarchical architectures to improve sample quality. A prominent example is the NVAE (Vahdat & Kautz, 2020) whose output distribution produces samples that are sharp and of high quality. The distribution that is being used, however, is still pixel-wise independent (discretized logistic mixture), and cannot model spatially coherent samples in the observation space. The very nature of a pixel-wise independent distribution does not allow it to model spatial dependencies between pixels. Any variation to the mean remains uncorrelated, which can be defined as noise even if it is not visually apparent. The success of the NVAE's high quality samples largely stems from the hierarchical architectural design, using spatial latent variables and other design choices that result. in spatial coherence, but not the pixel-wise independent distribution, which remains to
2

---

Published in Transactions on Machine Learning Research (10/2022)
have disadvantages that our work addresses. SOS-VAE- are complementary to the contributions in NVAE, addressing a different problem that persists in these recent approaches. Thus, our contributions such as the ability to generate multiple plausible predictions and allowing efficient interactive editing with a single forward pass stand independently to recent architectural contributions.
Like most other explicit likelihood models, the VAE requires the choice of a parametric observational dis- tribution. This choice is often pixel-wise independent. As a result, practitioners use the distribution to calculate the likelihood but use its expected value when sampling, as the samples themselves, are noisy and of limited use in most applications. However, according to the theory and as previously pointed out by Stirn & Knowles (2020) and Detlefsen et al. (2019), latent space sample should entail distribution over obser- vations and not a single point. Despite attempts to enforce spatial dependencies in the decoder architecture (Miladinovic et al., 2021), without a pixel-dependent joint likelihood, the observation samples will remain noisy. Notable exceptions are auto-regressive models and auto-regressive VAE decoders (Van Den Oord et al., 2017; Razavi et al., 2019; Guirajani et al., 2017: Nash et al., 2021). Unlike other explicit likelihood models, aulo-regressive models jointly model the observational distribution by sequentially decoding pixels while conditioning on previously decoded values. While this sampling procedure results in spatially coherent samples, it is computationally expensive and uncertainty estimation is not trivial.
In the context of non-auto-reerresivee VAE decoders, work locusing on modelling a joint observational dis- tribution that accounts for pixel dependencies is limited. Monteiro et al. (2020) use a low-rank multivariate normal distribution to produce spatially consistent samples in a segmentation setting. However, they focus on discriminative models only. For generative models, a notable exception is a work by Dorta et al. (2018) which, similarly to our proposed method, employs non-diagonal multivariate normal distribution over ob- servation space. The key difference is the choice of parameterisation used for the covariance matrix. Dorta et al. (2018) predict the sholesky-decomposee precision matrix, which grows quadratically with the size of the image. To address this computational constraint, the authors use a sparse decomposition. The imple- mentation of sparse matrices incurs additional overhcad and complexity required to avoid a conversion to dense matrices. Results from our head-lo-head comparison show quantitative and qualitative improvements for our method compared to Dorta et al. (2018).
3 Methods
3.1
Variational Autoencoders
We start by briefly revising the theory of the standard VAE as proposed by Kingma & Welling (2014) and Rezende et al. (2014). Given a random variable x and its latent representation z whose posterior p(z)x) is intractable, the aim of the VAE to obtain efficient inference and learning for the generative model p(x 22) Since the latent posterior is intractable, the VAE uses variational inference to approximate the posterior p(zix) using a variational distribution q(z)x). The approximate latent posterior and the generative model are parameterised by neural networks. For dataset X = N the VAE objective for a single data-point is given by maximising the evidence lower bound with respect to the networks' parameters:
L(O, Φi x() -DxL (z ( [[p(z)] + E. (e
[Iogro (x [2)]]
(1)
where p(z) is the latent prior, go go(z) (0) is the latent variational posterior parameterised by neural network with parameters o, and ()(x 2z) is the posterior likelihood parameterised by parameters O. The latent prior and variational latent posterior are chosen to be normally distributed such that the KL divergence can be computed in closed form. The likelihood cannot be computed in closed form since the derivative of the lower bound w.r.t & is problematic due to the stochastic expectation operator. Consequently, Monte-Carlo estimation and the ce-paraametrristion trick are used to yield the Stochastic Gradient Variational Bayes (SGVB) estimator (Kingma & Welling, 2014).
3.2 The Problem with Images
The likelihood term in the ELBO (equation 1) requires the choice of a parametric observational distribution to calculate the likelihood that the data comes from the predicted distribution: po(x|z). z). Tt. follows that each
3

---

Published in Transactions on Machine Learning Research (10/2022)
latent sample entails distribution over observations and that the predicted distribution should be as close as possible to the observed distribution for its samples to look like real observations. The standard choice of the observational distribution for image data is pixel-wise independent Bernoulli or Gaussian distributions since these make the likelihood calculation straightforward. However, pixcl-wise independent distributions ignore the strong spatial structure the data, therefore, are poor choices for observational distributions in images. This is evidenced by practitioners ignoring sampling from the observational distribution in favour of its expected value to avoid pixel-wise independent noise from corrupting samples, as shown in figure 1. This workaround technique has become the de facto standard for generating images with VAEs, despite theory describing the modelling of an observed distribution, not just its expected value. Models that only use the expected value are deviating from theory to cover for an inadequate choice of observational distribution. By incorporaling spatial dependencies in the predicted distribution, we aim to overcome this limitation and generate more realistic samples under the observed distribution, without relying only on the expected value.
3.3 Structured Observation Space VAE
We propose replacing the joint pixcl-wisc independent distributions used in the observational postcrior with a distribution that explicitly models the dependencies between pixels. Specifically, following recent developments in discriminative models for segmentation (Monteiro et al., 2020), we use multivariate normal distribution with fully populated covariance matrix po (x|z))=N(H, E), which we compute efficiently using a low-rank parameterisation. This small modification can be applied to most existing VAE architectures.
Given an image x or sample y with width W, height H, and channels, we consider its flat representation with size W- H C and use a threo-headed VAE decoder which outputs:
The mean of the multivariate normal distribution ja with size S;
The elements d of a positive diagonal matrix D with S diagonal elements;
The covariance factor matrix P with size S x R where R is the rank hyperparameter chosen for the parameterisation.
The low-rank covariance matrix can be computed as Σ PPT + D which has S (S R) degrees of freedom as opposed to 2 for full-rank covariance matrix. The rank hyperparameter is a trade-off between expressivity and computational cost. As hyperparameter, the only unbiased way of selecting is through ablation. In practice, we found it sensible to select it based on the available computation resources. To calculate the likelihood of a multivariate normal distribution, we must invert the covariance matrix and compute its determinant. For a full-rank covariance matrix, the computational cost and memory footprint make this infeasible for all except very small images since these costs scale quadratically with image size. However, with the low-rank parametrisation, the full matrix never needs be used directly, The inverse and determinant of the covariance matrix can be efficiently computed using the Woodbury matrix identity Woodbury (1950) and the matrix determinant lemma, respectively. Since both the covariance factor P and the diagonal elements d scale linearly with the size of the image, the computational costs of computing the logelikel scales linearly as well, in contrast to a full-rank covariance matrix where costs increase quadratically.
The proposed modification on the likelihood distribution does not affect the theory of the VAE. Thus the SGVB estimator (equation 1) is applicable without modification. However, we found optimising the SGVB estimator with a free non-diagonal covariance to be unstable due to exploding variance. The instability comes from there being two routes that minimise the likelihood. Either to find appropriate mean and variance around it or to keep incrcasing the variance (uncertainty about the mcan). The second route is undesirable and results in implausible samples (with overly bright colours and high contrast; cf. Appendix A). This problem is not unique to our implementation; the stability of variance networks has been discussed before (Stimn & Knowles, 2020; Detlefsen et al., 2019). To address this issuc, we introduce an entropy penalty on the observational distribution. Constraining the entropy constrains the variance indirectly, thus giving preference to low-variance solutions. Therefore, we compute the entropy of the normal distribution in closed
4

---

Published in Transactions on Machine Learning Research (10/2022)
form and add it to the objective function. The new objective, now also including 3 penalty (Higgins et al., 2017) is given in by:
É(0 xxxx)))= =
t ( [og pe (x( H (x
(2)
where (z) is the entropy of the predicted observational distribution. In addition to the entropy con- straint, we used the following optimisation tricks to further incentivise the low variance solution: we freeze the variance heads of the decoder and train the decoder with only the mean head for a few epochs Mon- teiro et al. (2020); we fix the covariance diagonal to a small positive scalar in the low-rank parametrisation D ele
Multiple- Multiple-constrainn constraint optimisation can lead to bad convergence properties, as result, we opted to employ soft-constraints for the entropy and KT. divergence using the modified differential method of multipliers (Platt & Barr, 1988), for its agreeable convergence and stability properties, which results in the following Lagrangian formulation:
[(, xx0())x())) -E (a [Log [008p((x(0)1) [z]
(3)
- [DKL DKL ((00(2)(())lp((zz) EKL
-XH [1(x(012)= EH
where 3 and ε are the Lagrangian multiplier and the slack variable, respectively, for the 3-VAE constraint and λy and εH are the Lagrangian multiplier and the slack variable, respectively, for the entropy constraint. The slack variables are tunable parameters of the model. Since they have semantic values, they can be chosen empirically rather than through blind trial and error. The slack variables for the entropy and KL constraints can be selected by observing the respective values from an unconstrained model. These observations give the orders of magnitude for the slack variables before further tuning. For our SOS-VAE, training remains stable for slack variables in the correct order of magnitude, however we found the entropy constraint to be detrimental the training of a 3-VAF, which can be explained by the decreased expressivity of the 3-VAE compared to our SOS-VAE.
The prior work of Dorta et al. (2018) uses non-diagonal covariance matrix for the distribution over the observation space, however our choice of parameterisation differs, as discussed in section 2. Furthermore, our choice to use the modified differential method of multipliers differs and our introduction of an entropy constraint is novel, as best we know, for this use case.
4
Experiments and results
4.1 3-VAE vs. Structured Observation Space VAE
We start by comparing a B-VAE, with a pixel-wise independent normal observational distribution, to the proposed SOS-VAE method, with a low-rank multivariate normal observational distribution. We choose a 3-V (Higgins et al., 2017) baseline rather than standard VAE so that it can be trained using the same Lagrangian method with 3 being Lagrangian multiplier. This allows for fairer testing since the comparison is between models that differ only by our contributions. This is the only modification to the baseline over a standard VAE. We perform the comparison in two datasets: the CELEBA dataset (Liu et al., 2015) and the UK Biobank (UKBB) Brain Imaging dataset (Miller et al., 2016). For all models, we use a latent space of dimension 128. For the low-rank model, we use a rank of 25. For the CELEBA dataset we use a target KL loss ÊKL, of 45 for both models and εи 504750 for our model. For the UKBB datasct we use a target KL loss, EKT, of 15 for both models and Er 198906 for our model. Figures 2a & 2b and 2c & 2d show the qualitative results for the comparison.
In Figures/ 2a and 2c, we see that the samples of the 3-VAE exhibit uncorrelated pixel noise around the mean, resulting from the pixel-wise independent joint observational distribution. In contrast, in Figures 2b and 2d,
We found 10-5 yield good results.
5

---

Published in Transactions on Machine Learning Research (10/2022)
μ
/
y
y
y
μ
d
d
(1)
(a) 3-VAE on CELEBA
(b) SOS-VAE CELEBA
μ
y
y
μ
y 4
d
(i)
Σ (1)
Σ
(c) 3-VAE UKBB Brain Scans
(d) SOS- VAE UKBB
Figure 2: Qualitative results comparing 8-VAE (2a) and SOS- VAS (2b) the CELEBA dataset. Same comparison on the UKBB dataset (2c & 2d). The rows from top to bottom: the mean of the observational distribution, a sample from the observational distribution, the difference between the mean and the given sample, the pixcl-wise independent variance per pixel, a slice of the covarinnce matrix: positive covariance and negative covariance to the central pixel.

---

Published in Transactions on Machine Learning Research (10/2022)
Dorta et al.
Ours
Figure 3: A comparison of 100 64x64 grayscale samples from Structured Uncertainty Prediction Network VAE (Dorta et al., 2018) (left) and from SOS-VAE (right).
we see that the samples produced by our SOS- contain mumantically-mcannnggii variations around the mean and are spatially coherent, as illustrated in the difference (row 3) between the mean (row 1) and the sample (row 2). Looking the variance of the two mothods (row 4), we see a significant difference in the regions where each model uncertain, highlighting the difference in behaviour between the two predicted distributions. The predicted covariance (rows 5 and 6) for the low-rank model contains a structure that pertains to the image content. These figure rows represent positive and negative covariance to the central pixel indicating global covariance can is modelled. This structure results in spatially coherent samples as opposed to the noisy samples of the 3-VAE, which are a consequence of the diagonal covariance. Interestingly, we observe more variation in the means of the B-VAE, suggesting that as more variation can be modelled in the observation space, less ncods to be modelled in the latent space.
Quantitative evaluation of generative modelling is inherently difficult task due to its subjective nature. While measuring the log- lieeliood is the obvious choice, it is often not indicative of sample quality (Theis et al., 2016; Borji, 2019). The Fréchet Inception Distance (FID) (Heusel et al., 2017) is the current standard choice of metric due to its consistency with human perception (Borji, 2019). We nole this metric is not without criticism (Borji, 2019; Razavi et al., 2019). We use it to evaluate our generative models (Seitzer, 2020) and report the results in Table 1. As in (Heusel et al., 2017), we generate 50, 000 samples from each of our models and propagate them, in turn, through the Inception-v3 model sufficiently represent the set of synthesizable samples. This reduces sample bias to produce reliable results.. The results do not represent the absolute performance of SOS-OVVL VAE, but rather the relative difference when compared to 3-VAE with pixel- wise independent observational distribution while everything else is constant. We emphasise, the proposed method is compatible with generative models from other works. We observe that, for both datasets, the proposed SOS-VAE achieves a lower FID score than the 8-VAE. Notably, the samples from the model with a low-rank multivariate normal observational distribution outperform the means of the 8-VAE, indicating that sampling from the observational distribution, as theory entails, does not reduce image quality.
Table 1 also reports the results for the cropped and aligned CELEBA datasct, resized to 64x64 and converted to grayscale². This enables the comparison to Dorta et al. (2018)'s Structured Uncertainty Prediction Network trained with equivalent parameters. Note that the image size and number of colour channels is constrained by the code of Dorta et al. (2018). We used our own implementation of Dorta et al. (2018)'s model predominantly using their original code, the same dataset and learning rate, trained for the same
For this dataset, our methad En -1760.

---

Published in Transactions on Machine Learning Research (10/2022)
Table 1: FID metric results for a 3-VAB with pixel-wise independent observational distribution, a Struc-
tured Uncertainty Prediction Network (Dorta et al., 2018) and our SOS-VAE. Lower FID scores represent better performance.
Method
Dataset
FID +
3- 3-VAE (MEANS)
121.65
3-VAF (SAMPLES)
196.40
CELERA
SOS- VAE (MEANS)
132.93
SOS- VASE (SAMPLES)
104.62
3-VAE (MEANS)
211.24
V-V (SAMPLES)
332.89
UKBB
SOS- OOS- (MEANS)
141.87
SOS-VAE (SAMPLES)
79.712
3-VAE (MEANS)
90.97
3-VAE (SAMPLES)
257.11
64x64
DORTA ET AL. (MEANS)
73.44
GRAYSCALE
DORTA ET AL. (SAMPLES)
77.72
CELEBA
SOS-VAE (MEANS)
70.43
SOS-VAE (SAMPLES)
74.66
number of epochs. Using the FID metric, we find that our SOS-VAE outperforms Dorta et al. (2018) in both means and samples respectively. Qualitatively comparing samples from Dorta et al. (2018) and SOS-VAE (Figure 3), both exhibit spatially-correlatee variations from the mean. However, in the samples generated with the method by Dorta et al. (2018), the variations exhibit high-frequency distortions. In contrast, the samples obtained with our SOS-VAL show semantically-mennt meaningful, spatially-correlateed variations globally across the generated image.
4.2 Interpolation in the Observation Space
To explore the expressiveness of the representations captured in the observation space, we visualise a con- tinuous range of samples from the proposed method. We perform spherical linear interpolation over the observation space to capture range of plausible images between two initial samples, Ya and yg. This is shown in equation 4, where Wp E R and Wa E ((SxC) are both auxiliary noise variables, slerp is the spherical interpolation function (see appendix D) and 1 E [0, is the interpolation factor.
y, P Pwp. + Vews where wp. slerp(wp tp
(4)
It is important to note that we only interpolate over Wp as it restricts the dimensionality the hyper-sphere size R; the v wa term only adds a small amount of uncorrelated noise, so setting it as a constant. has negligible effect. Figure 4 shows this technique visualise the variation contained in the observation space, demonstrating that the distribution captures semantieally-relovant features, such as hair colour, skin tone and background colour for the CELEBA dataset. Interpolation between differences in such features would typically entail interpolating latent variables and forward pass through the decoder for each interval, which is not required here.
4.3 Interactive Sampling from the Observation Space
We have demonstrated that a low-rank multivariate normal observational distribution can model a range of features. However, it would be useful if we could synthesise images with scmantically-meanine input. A step in this direction entails fixing the auxiliary noise variables associated with each sample and scaling the principal components of covariance factor, P. This is demonstrated in equation 5: using
human
8

---

Published in Transactions on Machine Learning Research (10/2022)
Figure 4: Spherical linear interpolation between auxiliary noise variables Wp. The four corners are random samples from a predicted distribution with all intermediate steps as interpolations between them. The two images represent interpolations in the observation space for two observational distributions predicted from different latent codes.
the singular value decomposition (SVD) of P and introducing a A IR. (R)
diagonal matrix of scaling coefficients,
P= USAAA)
Adjusting the scaling coefficients in A allows to tune spatially correlated features in the sample image. The use of SVD often makes the effect of each coefficient separable and semantically- Images generated through this method for CELEBA are shown in Figure 5, where each row demonstrates the effect scaling a different principal component. This figure shows the effect on the first ten principal components. The effect on all components and additional results on the UKBB data are given in Appendix B. Since this manipulation is using only the observational distribution, manipulation of these samples can be achieved without performing additional forward passes on the model.
to the image domain.
4.4 Interactive Editing of Predictions
One of the benefits of modelling spatial correlations in the aoservational distribution is that this information can be leveraged to interactively edit predictions. This involves manually editing part of the prediction and calculating the conditional distribution of the remaining pixels; for an arbitrary multivariate normal distri- bution, this is expressed in equations 6 and 7, where the edited pixels are modelled by y2 and the remaining pixels are modelled by Y1- We the mean of the conditional distribution (j²) as the corrected image. The updated covariance matrix (E) is not needed, we avoid evaluating it to reduce the computational cost of this process.
y
Y1
11
μ1
Σ
Σn
Σ
(6)
y2
[2
Σn1
9

---

Published in Transactions on Machine Learning Research (10/2022)
Figure 5: The effect of scaling each of the ten most principal components (each row), from top to bottom for a fixed auxiliary noise variable. The scale factor for each component ranges from -5 to +5 with intervals of 0.5.
Interactive editing is demonstrated Figure (with additional examples in Appendix C), where a prediction from our model, trained on the CELEBA dataset, is sequentially edited to alter the hair colour and skin manually edit small number of pixels and
tone. This demonstrates the power of the method: automatically update the remainder of the image coherently with the manual edit.
ply b) b)=N(ji, E
(7)
where fa 14 Σ Dy (b 1b2)
and Σ E11 E プ En1
4.5 Observational Distribution without Deep Learning
Typical VAËs rely on their deep learning components to model the features for their output. In contrast, since we use a low-rank parameterisation of a full covariance matrix, our observational distribution can model spatially-correlated features on its own. As a result, the ability to model these features not solely left to the deep learning components of the VAE; is shared with the linear transformations that compose the low-rank multivariate normal distribution.
To understand the expressiveness of our observation space model, we carried out experiment in which the VAE architecture is replaced with the parameters for the low-rank multivariate normal distribution with no deep learning layers or latent space representation involved. We then train this simple model as described in section 3.3. This allows us to examine the capability of the model to capture the observational distribution over the dataset, without deep learning.
This modelling is reminiscent of principal component analysis (PCA), particularly given that we use low- rank parameterisation, akin to the dimensionality reduction of PCA. Figure compares samples from the
10

---

Published in Transactions on Machine Learning Research (10/2022)
Figure 6: Sequentially editing hair colour and skin tone interactively. From left to right: predicted image with a small coloured edit made to the hair, the image after the conditional distribution has been calculated, the image with a further edit to the skin tone, the image after the conditional distribution has been recalculated. N.B: The red circles highlighting the manual edits are for illustration purposes only and computational purpose.
Figure 7: A qualitative comparison between 100 samples from the learnt observation space at rank = 25 with no deep learning components (left) and 100 samples from a linear PCA model with 25 features (right). both cases, the data used for fitting is the same random subset of 10000 images from the CELEBA dataset.
distribution alongside samples from a linear PCA model of equivalent feature reduction after fitting to the CELEBA dataset (Liu et al., 2015).
There little perceivable difference between the samples of the two methods, so we conclude that our low- rank multivariate normal distribution is comparably expressive to PCA for feature reduction. Furthermore, this experiment confirms the ability learn the parameters of distribution through backpropagation.
11

---

Published in Transactions on Machine Learning Research (10/2022)
5 Discussion
This paper highlights an often-overlooked aspect of the VAE architecture- - the observational distribution. We have confirmed that pixel-wise independent observational distributions produce samples with uncorrelated pixel noise. We have introduced SOS-VAE, a low-rank multivariate normal distribution as a choice for the observational distribution of a VAE, able to model covariance between pixels and produce multiple spatially- coherent samples with a single forward pass of the decoder. Our method introduces stability issues that are otherwise not present, but that we are able to resolve with an entropy constraint. Our results indicate that our choice of observational distribution is beneficial when compared to pixcl-wise independent distribution, as well as allowing sampling to be the primary method of image synthesis without reduction in qualily, theory would entail. Our method performs favourably compared to previous work (Dorta. et al., 2018), and is compatible with many VAE architectures, including within state-of-the-art models. We find that a low- rank multivariate observational distribution can be interpolated within, and allows for semantic, interactive manipulation of samples with a single decoder forward pass from a single latent variable.
In our current experimentation we use a standard VAE architecture for ease of implementation and lower computational requirements. It would be interesting incorporate our approach into recent models such as NVAE and other deep VAE variants (Child, 2020). However, training these models is extremely resource demanding. For the original NVAE implementation', the authors state that "cight 16-GB V100 GPUs are used for training NVAE on CelebA 64" and "Iraining lakes about 92 hours". In contrast, the SOS-VAE model used in our experiments trained on a single 16-GB T4 GPU in 72 hours, which is a substantial difference of practical importance. We had considered using a pro-trained deep VAE, training only the last few layers, with and without the proposed low-rank likelihood distribution. However, we found that these models produce unsatisfactory results on lower image resolutions, and training on higher resolutions remains prohibitive with limited computational resources.
Introducing an expressive observational distribution that is able to model features on its own, as we have, promotes a discussion comparing the features modelled in the observation space to those modelled in the latent space. In section 4.5, we observe our observational distribution's 's ability to model features on its own and in section 4.1 we observe a decrease in variation of the predicted means for our model compared to a 3- VAE. Assuming a dataset containing finite uncertainty, we deduce that the modelling of this uncertainty is split between the latent space and the observation space. Understanding where this split lies and what influences this is open question left for future work.
Ethical Considerations
Generative modelling is subject to dataset-inherited bias and our contributions are also susceptible. Whilst methods, such as our own, allow us to explore the biases that exist within a dataset, which in some cases makes for a desirable tool, typical usage may expose an undesired bias, particularly after training on the CELEBA datasct, which has a larger socictal impact. We acknowledge that these biases, such as lack of diversity, are present but state these as artefacts of the chosen dataset and not of our own design, opinions or beliefs. There is a clear need for more diverse and representative datasets that would allow a more complete picture of the abilities and limitations of generative models to be obtained. The current reliance on human face dalasets for generative modelling is problematic. While some alternatives such the waterbird dalaset (Sagawa et al., 2019) exist, these still lack the fidelity for visual interpretation of generative modelling results. We would hope that better, alternative datasets would be available for future work.
Our method allows for the syuthesis of multiple plausible samples as well as manipulation of samples with a single forward pass of the model, where other methods require multiple forward passes. This reduction computational burden could provide access to machine learning models for those with low-powered devices as well as reducing energy consumption and computational costs.
hhttps://githuub.cmmm wttpp//.....o////la//////////
12

---

Published in Transactions on Machine Learning Research (10/2022)
Reproducibility
All code is made publicly available on https://githul: com//biomedda--mirr//eee------ together with clear instructions how to fully reproduce our results on CELEBA to facilitate future work and case of comparisons.
Acknowledgements
This project has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 rescarch and innovation programme (Grant Agreement No. 757173, Project MIRA).
References
Ali Borji. Pros and cons of GAN evaluation measures. Comput. Vis. Image Underst., 179:41-65, 2019. doi: htt.://..vviu221110000000 URL https://doi... wttp://...g/1111.//////1
2018.10.003
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.nec, 2019. URL https: http:////pppaaueeeeeeeeeeee/utta
Xi Chen, Diederik P. Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, 1lya Sutskever, and Picter Abbecl. Variational lossy autocncoder. In 5th International Conference on Learning Represen- lations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview. net, 2017. URL https://openrevier wttprrriiiieeeeeffrrreeee wtt.://oooooiiddd/s////sro.
Rewon Child. Very deep VA generalize autoregressive models and can outperform them on images. arXiv preprint Xrx::2001 10650, 2020.
Nicki Skafte Detlefsen, Martin Jørgensen, and Søren Hauberg. Reliable training and estimation of variance networks. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' d'Aleé-BBue, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Syslems 32: Annual Con- ference on Neural Information Processing Systems 2019, NewrIPS 2019, December 8-14, 2019, Vancou- ver, BC, Canada, 6323 6333, 2019. URL https://proceedings http://poceeddn...uri......... 07211680006009917788111111000
//aahh/
wwwtle-bbbtttt....eee html.
Laurent Dinh, David Krueger, and Yoshua. Bengio. NICE: non-linear independent components estimation. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Workshop Track Proceedings, 2015. URL http: //arxiv. warxiv.org//bar. 1410/ 01410.85166
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In 5th Interna- tional Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net net, 2017. URL https: wwwpeeeview.eee/oooruueeeee
G. Dorta, S. Vicente, L. Agapito, N. D. F. Campbell, and I. Simpson. Structured uncertainty prediction networks. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognátion, pp. 5477-5485, 2018. doi: 2ttp:///RR.22000000000
Ian J. Goodfellow, Jean Pouget-Abadie,, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. Tn NIPS, 2014.
Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taïga, Francesco Visin, David Vázquez, and Aaron C. Courville. Pixelvae: A latent variable model for natural images. In 5th International Con- ference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.ned 2017. URL httpls://oppenreviiee wtt.://toooooommiiont.tttttt
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. In Isabelle Guyon, Ulrike
13

---

Published in Transactions on Machine Learning Research (10/2022)
Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Newral Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, 6626 6637, 2017. URL https:// proceedings. neurips cc/ http://00....aeaeoooooooooooo.e. -ttp:/4449--a-aaaaaattt..
Abstract html.
I. Higgins, Loie Matthey, A. Pal, C. Burgess, Xavier Glorot, M. Botvinick, S. Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In ICLR, 2017.
Tero Karras, Samuli Laine, Milk Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of StyleGAN. In Proc. CVPR, 2020.
Diederik P. Kingma and Max Welling. Auto-cncoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL http://axxiv...ggbbbb//s/
6114.
Diederik P. Kingma and Max Welling. An introduction to variational autoencoders. Foundations and Trends@ in Machine Learning, 12(4):307-392 2019. ISSN 1935-8245. doi: 10.1561/2200000000 URL http://dd /tt.xxxxiioo/d///ooo//o.//..oo.. org/10. 1561/2200000566.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural
Information. Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https: https://proceedings. .aff
neurips.cc/papert wttp//apee/2/////tii/ie htt.:///f....1222222222220011111
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Dorde Miladinovie, Aleksandar Stanic, Stefan Bauer, Jürgen Schmidhuber, and Joachim M. Buhmann. Spa- tial dependency networks: Neural layers for improved generative image modeling. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenRe- view.nct, 2021. URL /ttp:////ppeeeieieeeeee http:///aaaaoooeeeoooooooaaaaaaa
K L Miller, F Alfaro-Almage N K Bangerter, D L Thomas, E Yacoub, J Xu, A J Bartsch, S Jbabdi, S N Sotiropoulos, J L R Andersson, L Griffanti, G Douaud, T W Okell, P Weale, I Dragonu, S Garratt, S Hudson, R Collins, M Jenkinson, P M Matthews, and S M Smith. Multimodal population brain imaging in the uk biobank prospective epidemiological study, 2016.
Miguel Monteiro, Loic Le Folgoc, Daniel Coelho de Castro, Nick Pawlowski, Bernardo Marques, Kon- stantinos Kamnitsas, Mark van der Wilk, and Ben Glocker. Stochastic segmentation networks: Mod- elling spatially correlated aleatoric uncertainly. In Hugo Larochelle, Aure'ourelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Newral Information Pro- cessing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: //pppe//procceinngs.. nwwriee.u./eppppeeeeee cwwppppscc/pappr/20222 / hash/ Abstract. html.
Charlie Nash, Jacob Menick, Sander Dieleman, and Peter W. Battaglia. Generating images with sparse representations. In Marina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 7958-7968. PMLR, 2021. URL http://proceedii htt.:///gw.......roe/////////aa/ html.
mlr pptpss/v13/naas222222
John Platt and Alan Barr. Constrained differential optimization. In D. Anderson (ed.), Neural Informa- tion Processing Systems. American Institute of Physics, 1988. URL https://proceediing: wwwp:/iiinnnnnniiis..o..
neurips. cc/
Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vg-vae-2. In H. Wallach, H. Tarochelle, A. Beygelzimer, F. d' Blch E. Fox, and
14

---

Published in Transactions on Machine Learning Research (10/2022)
R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, pp. 14866- 14876. Curran Associates, Inc., 2019. URL. http:///poceedinssneeurri www.ipu..rccoomo sherripp.ccpppee/222///f11 5f8e2fal718dtbbcafff Paceeaafccccc//554fo-----aa pdf.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 31st International Conference on International Conference Muchine Learning- Volume 32, ICML'14, pp. II 1278 II 1286. JMLR.org, 2014.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Perey Liang. Distributionally robust neural networks for group shifts: shifts: On the importance of regularization for worst-case generalization. arXiv preprinl w:Xf::1911008791, 2019.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. Pixelenn++ A pixelenn implementation with discretized logistic mixture likelihood and other modifications. In ICLR, 2017.
Maximilian Seitzer. pytorch-fid: FID Score for PyTorrh. https://githube htt.tub..cmmmeettteeee//ttrrrhth August 2020. Version 0.2.1.
Casper Kaac Sønderby, Tapani Raiko, Lars Maalec, Søren Kaac Sønderby, and Ole Winther. Ladder varia- tional autoencoders. In Proceedings of the 30th Intervational Conference on Newral Information Processing Systems, NIPS*16, pp. 1745--755 Red Hook, NY, USA, 2016. Curran Associates Inc. ISBN 9781510838819,
Yang Song and Conor Durkan. Maximum likelihood training of based generative models. CoRR, abs/2101. 210/2200009258, 2021. URL https::/ https://arxivv httpg/abs22101.. 092588.
Andrew Stirn and David A. Knowles. Variational variance: Simple and reliable predictive variance parame- terization. CoRR, abs/2006.0499 044400, 2020. URL
httg/bb//000 2006. 04910.
Lucas Theis, Aäron van den Oord, and Matthias Bethge. A note on the evaluation of generative models. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http:
//arxiv... wwexxivvorgg/aae www.abs///aa.eeeoo..o..oo..... 1111.111441
Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. In Neural Information Processing Systems (NeurIPS), 2020.
Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukeuoglu. Pixel recurrent neural networks. Tn Proceedings of the 33rd International Conference International Conference Machine Learning Volume 48, ICML'16, pp. 1747 1756. JMLRorg, 2016a.
Aäron Van Den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu. Conditional image generation with pixelenn decoders. In Proceedings of the 30th Inter- national Conference on Neural Information Processing Systems, NIPS116, pp. 4797-4805, Red Hook, NY, USA, 2016b. Curran Associates Inc. ISBN 9781510838819.
Aäron Van Den Oord, Oriol Vinyals, and Koray Kavukcuoghu. Neural discrete representation learn- ing. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, pp. 6306- 6315. Curran Associates, Inc., 2017. URL http:://poceddnnggnerrrre..ppppp heeripp..c//ppee///////eee 7a98af Teaattaeeaaaa000200000000111 Paper. .dff
Max A. Woodbury. Inverting modified matrices. Princeton University, Princeton, N. J., 1950. Statistical
Research Group, Memo. Rep. no. 42,-
15

---

Published in Transactions on Machine Learning Research (10/2022)
A
But-of-the-bal box samples without stabilising
Figure The results after training with pre-training phase and weight initialisation, but without the fixed component, D cI, the entropy constraint. The first row represents the predicted means and all other rows represent samples from the predicted distribution outputted from the probabilistic decoder. Each column is a new sample from the latent prior decoded to predict distributions over the observation space. Model trained for 100 epochs on a random subset of 10000 images from the CELEBA dataset. Latent dimensionality: = 128, rank: = 25, target KL loss: εKL = 45. The resulting average variance for cach pixel is Rs 11.76, in comparison to As 0.0322 for the equivalently trained model with the entropy const 504750.
16

---

Published in Transactions on Machine Learning Research (10/2022)
B Scaling of all individual PCA components for CELEBA and UKBB
Figure 9: The effect of scaling each of the principal components (each row), from top to bottom for a fixed auxiliary noise variable and a rank 25 parameterisation. The scale factor for each component ranges from -5 to +5 with intervals of 0.5. Observational distribution predicted by our VAE after training on the CELEBA datasct.
17

---

Published in Transactions on Machine Learning Research (10/2022)
Figure 10: The effect of scaling each of the principal components (each row), from top to bottom for fixed auxiliary noise variable and a rank 25 parameterisation. The scale factor for each component ranges from -5 to +5 with intervals of 0.5. Observational distribution predicted by our VAE after training on the UKBB dataset.
18

---

Published in Transactions on Machine Learning Research (10/2022)
C Additional Examples for Interactive Editing
Sequential interactive editing. From left to right: predicted image with small coloured edit made to the hair, the image after the conditional distribution has been calculated, the image with further edit to the skin tone, the image after the conditional distribution has been recalculated. N.B: The red circles highlighting the manual edits are for illustration purposes only and serve computational purpose.
Figure 11:
Figure 12:
Interactive editing. Left: predicted image with a single-pixel manual coloured edit over the hair. Right: the mean of the calculated conditional distribution. N.B: The red circle highlighting the manual edit is for illustration purposes only and serves no computational purpose.
Figure 13: Interactive editing. Left: predicted image with manual coloured edit over the skin. Right: the mean of the calculated conditional distribution. N.B: The red circle highlighting the manual edit is for illustration purposes only and serves no computational purpose.
19

---

Published in Transactions on Machine Learning Research (10/2022)
Figure 14: Interactive oditing. Left: predicted image with a manual coloured edit over the background. Right: the mean of the calculated conditional distribution.
D Spherical Interpolation
For completeness, we include the spherical interpolation formula used in section 4.2.
slerp(a, b,t) sin((1 - t)w) a + sin(tw)
(8)
sin w
sin w
where cos (lal [b])
20


Reviews:
[{"review_id": "Yyz9J637tN", "reviewer": "Reviewer_SxrT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper sets out a clear line of inquiry from its outset. I appreciated the focus on laying out the assumptions made among the various improvements to the standard VAE throughout the literature that highlight some limitations outsized limitations. This helps set-up the major areas of focus to guide the development of the proposed structured observation space VAE (if I may, SOSp-VAE):\n1) Pixel-wise independence assumptions limit expressivity\n2) Observation space is commonly modeled as an independent Normal distribution.\n3) Generated samples from standard VAEs are not greatly differentiated from the mean of this indep. Normal dist.\n\nThe paper makes reasonable claims about the benefit of using a structured observation space, driven by allowing the optimization process to identify spatial correlations between pixels, thus improving the coherence of the generated images. An additional benefit that is possibly downplayed, is that a more complete representation of the data distribution is generated with a single forward pass of the model, enabling more efficient means by which to generate a diversity of samples. With the constructed mean and low-rank covariance, the SOSp-VAE is able to produce a flexible range of samples without needing additional expensive executions of the trained model. This is a very impressive finding and established result (among the breadth of additional insights provided in Section 4) that I believe the generative modeling community would find useful and interesting.\n\nI am however a little surprised that the technical framing of the proposed SOSp-VAE is only covered lightly with a fairly vague section with limited formalization of the proposed modeling approach. As currently submitted, I cannot advocate for this paper to be published. There are far too few precise technical details about the proposed method. Much of the language surrounding the development of the covariance structure as well as the optimization procedure is far too vague and needs significant revision. I would recommend that the authors take the time to more fully and carefully outline their improvements to modeling the observation distribution by doing the following:\n1) Clearly define the notation at the outset of Section 3 (eg. image example X of size S = HxW, etc)\n2) Formalize the low-rank representation of the the covariance matrix, perhaps by more thoroughly describing Monteiro, et al (2020).\n\n    a) As a side, note, is there a typo for the definition of the Diagonal matrix **D** in Section 3.2?\n3) The initial description of the challenges of to the optimization of the improved obs. dist at the end of page 3 is great! This helps motivate the changes made to the objective function. But each of the adjustments made to the objective function should be clearly outlined. For example, moving to a $\\beta$-VAE type objective is not at all mentioned. This should be better justified as well as the choice for the Lagrangian form of the objective. Just a little more precision and detail behind for the formulation would be necessary here. \n\nA major omission in the technical introduction of Equation 2 is the lack of discussion around why multiple samples are needed in the first term and what M is. \n\nMore detail about the approach used in Dorta, et al (2018) is warranted. Following this, a more careful discussion that helps to differentiate this prior method from the proposed SOSp-VAE would be very helpful in understanding the technical grounding of the contributions. \n\nOverall, the paper makes unsubstantiated claims or other statements about being closer to VAE theory in the development and use of the structured observation distribution. A more thorough development in Section 3 would help solidify the claims made throughout the paper on this front. Specifically, what parts of the theory informed the development and how does the SOSp-VAE satisfy this? Please be more precise and formal here.\n\nI can understand the hesitation to expanding Section 3, but I believe it to be necessary to fully anchor the technical contributions made by the proposed SOSp-VAE. While the paper is good length, much of the space is taken by the figures and is clearly written. Of all the sections, the technical development in Section 3 is by far the weakest. Without more information and foundational development of the proposed method, I cannot envision the ML community really finding much use with this work despite the intriguing empirical results. \n\nAs an additional hesitation about the technical development of the SOSp-VAE, it's unclear how the low-rank covariance is actually represented. How are the parameters of the observation distribution actually produced by the model? How is the low-rank nature of **P** ensured?\n\nI found Figure 2 to be a very neat way to validate the claims made about the failings of standard VAE models. I wonder if there's not a way to quantify the artifacts of this analysis (percentage of image different from mean, etc?). \n\nThe choices of the loss hyperparameters is opaque and seems random. More precision here outlining how the models were trained and what types of investigations were run to come to those parameters would be helpful. As currently submitted, the method and results are not at all reproducible. In this way, the paper seems incomplete.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces an interesting alternative line of inquiry among the classes of VAE models by focusing on the representation of the observation distribution. An improved formulation of the predicted observation distribution is made by forming a full covariance matrix, removing standard assumptions of the generated pixel space being independent. To enable the execution of this approach, the covariance is estimated through a low-rank decomposition with additional adjustments to the optimization procedure to reduce instabilities while training. As a result of modeling a more complete approximation of the observation distribution, the generated images from the proposed VAE variant display greater spatial coherence and avoid uncorrelated pixel noise. Extensive qualitative analyses are performed to demonstrate the extensive capabilities of the Structured Observation Space VAE, including comparisons to related structure-inducing baselines.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Forgive my boldness in naming the proposed method. I do think that so much of the paper and discussion of the results would be greatly improved if the new modeling approach had a name. This will make it easier for future work to clearly reference this paper and the insights it contains. We see this with the $\\beta$-VAE. On top of the empirical improvements it provided, part of it's lasting impression is that it is recognizable by name. Please name your new model, feel free to take or improve upon my hastily determined SOSp-VAE if you like.\n\nPlease revise Section 3 following the recommendations given above.\n\nPlease run additional baselines that compare to the $\\beta$-VAE and possibly a standard VAE that incorporates an entropy penalty. It's important to understand how the change to the objective function affects the comparison to the change of observation distribution. Without these ablations it is difficult to fully assess the impact of the separate improvements made with the SOSp-VAE. Is it possible that adjusting the objective as is done in Eqt 2, is sufficient to cover the outlined failings of standard VAE models? A comparison of both FID and qualitative features such as those done in Figure 2 between standard VAE, SOSp-VAE (done) as well as these two ablations would be really great.\n\nOn the note of Figure 2. It would be really helpful if each row was labeled on the left of the figure with the established notation in Section 3. ", "broader_impact_concerns": "I appreciate the effort taken in the writing of the current version of the paper that lays out some ethical considerations. Perhaps something that could be taken into consideration is that generative models based on human faces and other personally identifying features should be avoided altogether? A nice example of work that has done this is Creager, et al (ICML 2021, reference below) which develops a generative model over natural images (they use the waterbirds dataset for example). Would it be admissible to reconsider the use of CelebA entirely? Would the type of results demonstrated in this work be possible with a subset of ImageNet?\n\n\n\nCreager, E., Jacobsen, J. H., & Zemel, R. (2021, July). Environment inference for invariant learning. In International Conference on Machine Learning (pp. 2189-2200). PMLR.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "h40WMNSwT2", "reviewer": "Reviewer_34NE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "It is a good observation that Gaussian VAEs with diagonal matrices correspond to modeling each pixel independently. \n\nWhile other work may be focusing on scaling up the models, this work models spatial dependencies (which is by modeling the covariance differently) more elegantly. \n\nIn general, the paper is well written and easy to follow. \n\n--- \n\nI have a few questions on Eq. (2). \n\nFirst, in Eq. (2), how are the hyperparameters tuned? For example, in CELEBA, the slack variables are 45 and -504750, respectively. How easy is it to obtain these numbers? It’s not clear how many runs are needed to get a satisfactory choice of the slack variables. It’s not clear how sensitive the results are with respect to the choice of slack variables. What should practitioners do in terms of hyperparameter tuning, if they want to use the authors’ method in the future? \n\nSecond, it is also not clear how the Lagrangian coefficients are determined.\n\nThird, the motivation for using Eq. (2) is to better constrain the variance. It'll be great if there is a table/figure that empirically describes how much variance is reduced. \n\nIn terms of baseline, it’s great that the authors compared their performance with Dorta et al. (2018), which is designed for single-channel images. First, how well tuned are the hyperparameters / how well tuned is this baseline in general? Second, did the authors attempt adapting the method on multi-channel images? What difficulties are there? \n\nIn terms of evaluation metrics, FID scores can be introduced more thoroughly. Moreover, what does a difference of x points mean? How much of a difference shows that the authors’ approach is significantly better? \n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "To use VAEs to do image generation, the authors improve the samples from the observational distribution. Specifically, Gaussian VAE’s decoder relies on a diagonal covariance matrix. The authors replace the diagonal covariance matrices with low-rank matrices (PP^T; see Section 3.2) plus diagonal covariance matrices. By doing so, the model can encode spatial dependencies among pixels better. \n\nMoreover, the authors modify the VAE objective as shown in Eq. (1). The authors also try to minimize the entropy of the predicted distribution which can constrain the variance as a byproduct, as shown in Eq. (2).  \n\nExperiments show that samples from the proposed approach capture covariances among different parts of an image. The images shown in the paper look more spatially coherent. \n", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "- Add more information on tuning details of Eq. (2). How sensitive is the performance with respect to the slack variables and the Lagrangian multipliers? What would the authors suggest, if practitioners implement and experiment on the authors' approach?\n- Discuss how well-tuned Dorta et al. (2018) is, given that it is an important baseline. \n- Discuss whether the authors have tried adapting the above baseline onto multi-channel images. \n- Quantify the effectiveness of Eq. (2) in terms of \"reducing variance.\" \n- Discuss FID scores more thoroughly. \n- Discuss how the authors' approach is \"compatible with state-of-the-art models\" -- what is a state-of-the-art model; why will the authors' approach bring further improvement on such models? ", "broader_impact_concerns": "The ethical consideration (written by the authors) is satisfactory. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "G3e7scVOgQ", "reviewer": "Reviewer_zquW", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is clearly written and easy to read overall.\n- The authors clearly tried several approaches to make their approach work in practice, including proper initialization of NN weights, fixing the covariance diagonal, and incorporating an entropy constraint on the output distribution, etc.\n- The authors provided a set of interesting experiments, such as interactive image editing.\n\nWeaknesses:\n- It's not very clear to me what the takeaway from this paper is. They seem to suggest that a structured observational distribution provides both computational (e.g. low-rank parameterization) and sample quality gains. But if that is the case, then it would have been nice to see comparisons of this method with Dorta et al. in terms of the number of model parameters or runtimes for the efficiency angle. If the main takeaway is improved spatial coherence of samples, then I have some more concerns below.\n- The authors mentioned that their goal was not to obtain state-of-the-art sample qualities, but given the current submission it’s rather hard to tell: (a) how much their approach will help on larger models (that already produce good quality samples); (b) how much their approach will help on models with more complex structure in the VAE latent space; and (c) how well they are performing relative to other baselines in terms of commonly used metrics in the VAE community. \n  - As written in the paper, the current literature has largely focused on architectural improvements and more complex models of the latent space. Such improvements (e.g. BIVA or NVAE) already produce spatially-coherent samples without additional considerations on the observational distribution. Given that the paper does not provide any way to compare across these approaches (either via models with similar architectures and different methods, or applying the authors’ method on top of these existing methods), it’s hard for me to believe that the structured observational distribution is actually providing significant gains when the baseline VAE model quality appears to be quite poor.\n  - Even though log-likelihoods are not necessarily a good proxy for high sample quality (Theis et al. 2015), reporting these bpds allows for a fair comparison of this method across other baselines. It's hard to say how \"good\" an FID score of 100+ is on CelebA.\n\n\n\nThere were also some missing references that would be helpful to include:\n- BIVA (Maaloe et al. 2019): architectural improvements to VAEs\n- Correlated VAE (Tang et al. 2020): a similar idea to inducing correlation structures across dimensions, though this operates in the latent space of a VAE\n- Bit prioritization (Shu et al. 2022)\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an improvement in sample quality (spatial coherence) to the Variational Autoencoder (VAE) by incorporating pixel-wise correlations in the continuous output space. The authors do so by using a structured output distribution: they use a low-rank parameterization of a full covariance matrix in a multivariate Normal distribution. This is in contrast to the way in which most VAEs are parameterized in practice, where the observational distribution is pixel-wise independent. They demonstrate performance improvements on the CelebA and UK Biobank Brain Imaging datasets. \n", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Major:\n- The biggest missing piece is a compelling demonstration that the structured output space in a VAE actually provides significant gains in sample quality. This can be demonstrated in a variety of ways. For example, the authors could apply their approach on top of existing models such as the NVAE, since this would lead to samples of similar quality. Or, the authors could report metrics that are used by existing methods (e.g. log-likelihoods) to demonstrate that they perform favorably. \n- Such experiments or discussions would help clarify the take-home message. Is a structured output space complementary to existing approaches? Or is it enough to have the structured output space and ignore the other components? Those are the questions that still remain for me.\n\nMinor changes:\n- It'd be helpful to be self-contained and explain what it means to \"pre-train\" the mean (Mointeiro et al. 2020).\n- Consistent notation: (nitpick) why not index z by i when talking about the ith data point x^{(i)} in Eq. 1?\n- The prior p(z) should not have \\theta in Eq. 2\n- The true posterior p(z|x) should not be parameterized by theta (which parameterized the decoder, which is being learned) in Section 3.1\n- Would the authors provide some guidance on how the hyperparameters were chosen? (e.g. e_H = -504750 doesn't seem like a very intuitive number to try as a first stab. Also, how should we go about choosing the rank of the covariance matrix relative to the dimensionality of the data, etc.?)\n", "broader_impact_concerns": "The authors included a sufficient section on ethical considerations (5.1).", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: First — thanks to the authors for submitting this interesting work, thanks to the reviewers for your insight, and thanks to all for a productive discussion that appears to have clarified the manuscript and led to concrete proposals to improve the work. 

To summarize the reviewer discussion — there was a consensus that the discussion, revisions, and discussed new experiments were satisfactory.  

One reiterated concern is that the reported improvements are due to a relatively weak VAE baseline.  One reviewer explicitly suggested incorporating another experiment with 'a simpler dataset with a baseline that's known to produce good samples' as a way to understand the effect of the new approach.  Another author proposed an additional experiment with a stronger baseline.  The concrete experiment proposal from the authors — tuning a pre-trained NVAE model with and without the proposed low-rank likelihood distribution — would improve our understanding of the new approach, and will making the final version of this work more compelling.  We suggest this experiment be added to the next revision.

Thanks again for the submission!
Decision: Accept (poster)