Simple and Effective VAE Training with Calibrated Decoders
Authors: Oleh Rybkin
Kostas Daniilidis
Sergey Levine
Venue:  (2021)
Keywords: variational autoencoders
β-VAE
representation learning

Abstract: Variational autoencoders (VAEs) provide an effective and simple method for modeling complex distributions. However, training VAEs often requires considerable hyperparameter tuning to determine the optimal amount of information retained by the latent variable. We study the impact of calibrated decoders, which learn the uncertainty of the decoding distribution and can determine this amount of information automatically, on the VAE performance. While many methods for learning calibrated decoders have been proposed, many of the recent papers that employ VAEs rely on heuristic hyperparameters and ad-hoc modifications instead. We perform the first comprehensive comparative analysis of calibrated decoder and provide recommendations for simple and effective VAE training. Our analysis covers a range of datasets and several single-image and sequential VAE models. We further propose a simple but novel modification to the commonly used Gaussian decoder, which computes the prediction variance analytically. We observe empirically that using heuristic modifications is not necessary with our method.

Under review as a conference paper at ICLR 2021
SIMPLE AND EFFECTIVE VAE TRAINING WITH CALIBRATED DECODERS
Anonymous authors Paper under double-blind review
ABSTRACT
Variational autoencoders (VAEs) provide an effective and simple method for mod- eling complex distributions. However, training VAEs often requires considerable hyperparameter tuning to determine the optimal amount of information retained by the latent variable. We study the impact of calibrated decoders, which learn the uncertainty of the decoding distribution and can determine this amount of information automatically, on the VAE performance. While many methods for learning calibrated decoders have been proposed, many of the recent papers that employ VAEs rely on heuristic hyperparameters and ad ad-hoc modifications instead. We perform the first comprehensive comparative analysis of calibrated decoder and provide recommendations for simple and effective VAE training. Our analysis covers a range of datasets and several single-imag and sequential VAE models. We further propose simple but novel modification to the commonly used Gaus- sian decoder, which computes the prediction variance analytically. We observe empirically that using heuristic modifications is not necessary with our method.
INTRODUCTION
Deep density models based on the variational autoencoder (VAE) (Kingma & Welling 2014 Rezende al. 2014) have found ubiquitous use in probabilistic modeling and representation learning as they are both conceptually simple and are able to scale to very complex distributions and large datasets.
These VAE techniques are used for tasks such as future frame prediction (Castrejon et al., 2019). image segmentation (Kohl et al. 2018), generating speech (Chung et al. 2015) and music (Dhariwal al. 2020), as well as model- based reinforcement learning (Hafner al.) 201 9a). However, in practice, many of these approaches require careful manual tuning of the balance between two terms that correspond to distortion and rate from information theory (Alemi et al., 2017). This balance trades off fidelity of reconstruction and quality of samples from the model: model with low rate would not contain enough information to reconstruct the data, while allowing the model to have high rate might lead to unrealistic samples from the prior as the KL-divergence constraint becomes weaker (Alemi et al. 2017| Higgins et al.) 2017). While a proper variational lower bound does not expose any free parameters control this tradeoff, many prior works heuristically introduce weight on the prior KL--divergence term, often denoted B. Usually, 8 needs to be tuned for every dataset and model variant hyperparameter, which slows down development and can lead poor performance as finding the optimal value is often prohibitively computationally expensive. Moreover, using 3 +f 1 precludes the appealing interpretation of the VAE objective as a bound on the data likelihood, and is undesirable for applications like density modeling.
While many architectures for calibrating decoders have been proposed in the literature (Kingma & Welling, 2014] Kingma et al. 2016]] Dai & Wipi 2019). more applied work typically employs VAES with uncalibrated decoding distributions, such as Gaussian distributions without a learned variance, where the decoder only outputs the mean parameter (Castrejon et al. 2019 Denton & Fergus 2018
Lee et al. 2019
Babaeizadeh et al. 2018 Lee et al. 2018 | Hafner et al. 2019b Pong et al. 2019
Zhu et al. 2017 Pavlakos et al.) 2019h, uses other ad-hoc modifications to the objective (Shnn
et al.] 2013 Henaff et al. 2019). Indeed, it is well known that attempting to learn the variance in a
Gaussian decoder may lead to numerical instability (Rezende & Viola 2018 Dai & Wipf 2019).
and naive approaches often lead to poor results. As a result, it remains unclear whether practical empirical performance of VAEs actually benefits from calibrated decoders not.

---

Under review as a conference paper at ICLR 2021
To rectify this, our first contribution is comparative analysis of various calibrated decoder architec- tures and practical recommendations for simple and effective VAE training. We find that, while naive calibrated decoders often lead to worse results, careful choice of the decoder distribution can work very well, and removes the need to tune the additional parameter 3. Indeed, we note that the entropy of the decoding distribution controls the mutual information 7 (ar; z). Calibrated decoders allow the model to control / (z; z) automatically, instead of relying on manual tuning. Our second contribution a simple but novel technique for optimizing the decoder variance analytically, without requiring the decoder network to produce it as an additional output. We call the resulting approach to learning the Gaussian variance the a- VAE. In our experiments, the o- v-E outperforms the alternative of learning the variance through gradient descent, while being simpler to implement and extend. We validate our results on several VAE and sequence VAE models and a range of image and video datasets.
2 RELATED WORK
Prior work on variational autoencoders has studied a number of different decoder parameterizations. Kingma& & Welling (2014); Rezende et al. (2014) use the Bernoulli distribution for the binary MNIST data and KKingma& & Welling (2014) use Gaussian distributions with learned variance parameter for grayscale images. However, modeling images with continuous distributions is prone to instability as
the variance can converge to zero (Rezende & Viola 2018 Mattei & Frellsen 2018: Dai & Wipf
2019). Some work has attempted to rectify this problem by using dequantization (Gregor et al.) 2016)
which is theoretically appealing it is tightly related to the log-likelihood of the original discrete data (Theis et al. 2016). optimizing the variance in a two-stage procedure (Arvanitidis et al. 2017). [ more expressive distributions. Additionally ly, different choices for representing such variance exist, including diagonal covariance (Kingma & Welling, 2014 Sonderby al., 2016 Rolfe (0166), or single shared parameter (Kingma et al. [016 Dai & Wipt 2019 Edwards & Storkey [016 Rezende & Viola 2018). We analyze these and notice that learning a single variance parameter shared across images leads to stable training and good performance, without the use of dequantization or even clipping the variance, although these techniques can be used with our decoders; and further improve the estimation of this variance with an analytic solution.
or training a post-hoc prior (Ghosh al. 2019). Takahashi et al. (2018); [Barron (2019) proposed
Early work on discrete VAE decoders for color images modeled them with the Bernoulli distribution, treating the color intensities as probabilities (Gregor et al. 2015). Further work has explored various parameterizations based on discretized continuous distributions, such as discretized logistic (Kingmal et al. 2016). More recent work has improved expressivity of the decoder with mixture of discretized
logistics (Chen et al.) 2016; Maalse et al. 2019). However, these models also employ powerful autoregressive decoders (Chen et al. [0166] Guirajani et al. 2016; Maaløe et al. 2019). and the latent variables in these models may not represent all of the sigmificant factors of variation in the data, as some factors can instead be modeled internally by the autoregressive decoder (Alemi et al. 2017
While many calibrated decoders have been proposed, outside the core generative modeling community uncalibrated decoders ubiquitous. They are used in work video prediction (Denton & Fergus
2018 Castrejon et al., 2019 Lee et al. 2018 |Baaeezzadhh et al., 2018). image segmentation (Kohl
al. 2018), image-to-image translation (Zhu et al. (2017), 3D human pose (Pavlakos al. 2019).
as well as mrodel-based reinforcement learning (Henaff et al. 2019 Hafner et al.) 2019b a). and
representation learning (Lee et al.. 2019 Watter al. 2015 Pong et al. 2019). Most of these works utilize the heuristic hyperparameter i instead, which is undesirable both as the resulting objective is no longer a bound on the likelihood, and as 3 usually requires extensive tuning. In this work, we analyze the common pitfalls of using calibrated decoders that may have prevented the practitioners from using them, propose simple and effective analytic way of learning such calibrated distribution, and provide comprehensive experimental evaluation of different decoding distributions.
Alternative discussions of the hyperparameter 3 are presented by Zhao et al. (2017) 7p: Higgins et al. (2017); Alemi et al.) (2017); Achille & Soatto| (2018), who show that it controls the amount of information in the latent variable, I [xx; z). Peng et al. (2018); Rezende & Viola 2018) further discuss constrained optimization objectives for VAEs, which also yield similar hyperparameter. Here, we focus on B- VAES with Gaussian decoders with constant variance, as commonly used in recent work, and show that the hyperparameter /3 can be incorporated in the decoding likelihood for these models.
'BIA Maalee al. 2019) uses the Mixture of Logistics decoder proposed in (Salimans al. 2017) that produces the channels for each pixel autoregressively, see also App
2

---

Under review as a conference paper at ICLR 2021
3 ANALYSING DECODING DISTRIBUTIONS
The generative model of a VAE (Kingma & Welling 2014 Rezende et al. 2014) with parameters
0 is specified with a prior distribution over the latent variable p2(2), commonly unit Gaussian, and a decoding distribution pe (2|2), which for color images is commonly a conditional Gaussian parameterized with neural network. We would like to fit this generative model given dataset by
maximizing the evidence lower bound (ELBO (Nell & Hinton 1998 Jordan et al. 1999 Kingma & Welling 2014 Rezende al. 2014)), which uses an approximate posterior distribution qs
(s(222)) [2) also
commonly conditional Gaussian specified with a neural network. In this work, we focus on the form of the decoding distribution po (zr]z). To achieve the best results, we want a decoding distribution that represents the required probability p(z|z) accurately In this section, will review and analyze various choices of decoding distributions that enable better decoder calibration, including expressive decoding distributions that can represent both the prediction of the image and the uncertainty about such prediction, or even multimodal predictions.
3.1 GAUSSIAN DECODERS
We first analyse the commonly used Gaussian decoders. We note that the commonly used MSE reconstruction loss between the reconstruction i and ground truth data is equivalent to the negative log-likelihood objective with a Gaussian decoding distribution with constant variance:
1 ق z=|2+ + Dn 2 2
1 || 2
D MSE( a:) +c,
- In p(x])))
æ| 2 + c
2
where p(z|2) ~ N( 12, the prediction è is produced with neural network a po(z), and D is the dimensionality
This demonstrates drawback of methods that rely simply on the MSE loss (Castrejon al. 2019
Denton & Fergus 2018 Lee et al. 2019 Hafner et al. 2019b Pong et al., 2019] Zhu et al. 2017 Henaff al. 2019), as it is equivalent to assuming particular, constant variance of the Gaussian decoding distribution. By learning this variance, we can achieve much hetter performance due to better calibration of the decoder. There several ways in which can specify this variance. An expressive way to specify the variance is to specify a diagonal covariance matrix for the image, with one value per pixel (Kingma & Welling 2014; Sønderby et 20166] Rolfe 2016). This can be done, for example, by letting neural network dg output the diagonal entries of the covariance matrix given a latent sample 2:
(1) This parameterization of the decoding distribution outputs one variance value per each pixel and channel. While powerful, we observe in Section $ec this approach attains suboptimal performance, and is moreover prone to numerical instability. Instead, we will find experimentally that a simpler parameterization, in which the covariance matrix is specified with a single shared (Kingma al. 2016, Dai & Wipf, [219 Edwards & Storkey 2016 Rezende & Viola 2018) parameter a as E = al often works better in practice:
(=]z)) - (l)())
Po.- (xl2)) N (Ho(=),
The parameter a can be optimized together with parameters of the neural network with gradient descent. Of particular interest is the interpretation of this parameter. Writing out the expression for the decoding likelihood, we obtain
1 202
1 _ [e===l+DDoott = Dinot MSE(i, +c. 20
D
- x|| -Dlov
202
The full objective of the resulting Gaussian a-VAE
D Ca.s. Dlno MSE(E, + DxL(((z))))pp())
(3)
20
Note that a may be viewed as a weighting parameter between the MSE reconstruction term and the KL-divergence term in the objective. Moreover, this objective explicitly specifies how to select the optimal variance: the variance should be selected to minimize the (weighted) MSE loss while also minimizing the logarithm of the variance.
3

---

Under review as a conference paper at ICLR 2021
Decoder Calibration It important that the decoder distribution be calibrated in the statistical sense, that is, the predicted probabilities should correspond to the frequencies of seeing a particular value of z given that prediction (DeGroot & Fienberg 1983
Dawid 1982 The calibration of
neural network can be usually improved by estimating the uncertainty of that prediction (Guo et al. 2017). such the variance of a Gaussian (Kendall & Gal 2017). Since the naive MSE loss assumes a constant variance, it does not effectively represent the uncertainty of the prediction, and is often poorly calibrated. Instead, learning the variance as in Eq. leads to better uncertainty estimation and better calibration. In 5.1] we show that learning a good estimate of this uncertainty is crucial for the quality of the VAE generations.
Connection to 3-VAE. unit variance is:
The 8-VAE objective (Higgins et al. 2017) for a Gaussian decoder with
D C -=-2MMs M) + (T
We see that it can be interpreted as a particular case of the objective (3). where the variance is constant and the term D In a can be ignored during optimization. The 3-VAE objective is then equivalent to Ø- -AE with a constant variance work (Zhu et al. 2017 Denton & Fergus 2018 Lee et al. 2019), 3-VAE models are often used in this exact regime. By tuning the 33 term, practitioners
/322 (for a particular learning rate setting). In recent
able to tune the variance of the decoder, manually producing more calibrated decoder. However, by re-interpreting the 3-VAE objective as a special case of the VAE and introducing the missing D In or term, we can both obtain a valid evidence lower bound, and remove the need to manually select B. Instead, the variance can instead simply be learned end-to-end, reducing the need for hyperparameter tuning.
An alternative discussion of this connection in the context of linear VAEs is also presented by Lucas ct al. (2019). While the 3 term is not necessary for good performance if the decoder calibrated, can still be employed if desired, such as when the aim is to attain better disentanglement (Higgins) jet al., (017) or a particular rate-distortion tradeoff (Alemi et al. 2017). However, we found that with calibrated decoders, the best sample quality is obtained when 3
Loss implementation details. For the correct evidence lower bound computation, it is necessary to add the values of the MSE loss and the KL. divergence across the dimensions. We observe that common implementations of these losses (Denton & Fergus 2018 Abadi et al., 2016; Paszke et al., 2019) use averaging instead, which will lead to poor results if the number of image dimensions is significantly different from the number of the latent dimensions. While this can be conveniently ignored in the V-V regime, where the balance term is tuned manually anyway, for the a-VA it is essential to compute the objective value correctly.
Variance implementation details. Since the variance is non-negative, we parameterize it logarith- mically as o² = e²), where À is the logarithm of the standard deviation. For some models, such as per-pixel variance decoders, we observed that it is necessary to restrict the variance range for numerical stability. We do so by using the soft clipping operations proposed by CChua et al. (2018):
A = Amax - softplus Amax A);
== Amin + softplus(A - >miin).
We observe that setting Amin 6 to lower bound the standard deviation to be at least half of the distance between allowed color values works well in practice. We also observe that this clipping is unnecessary when learning a shared a value.
3.2 DISCRETE DECODERS
It is possible use discrete decoding distributions to generate images, as color values are commonly restricted to a fixed set of integer pixel intensities (c.g. 0..255). Indeed, for discrete color values, discrete distributions are arguably more appropriate. In the most general case, discrete decoding distribution factorized per each pixel and channel would be specified by a probability mass vector f with 256 entries, one per each possible intensity value, similarly to a per-pixel classifier of the intensity value. We can implement it with a soft-max layer, yielding the following log-likelihood loss (sometimes called the cross-entropy loss) for true pixel with intensity
exp($;;
[n)((>)> = - In Σ, Exxp(1;))
4

---

Under review as a conference paper at ICLR 2021
o
H
θ
D
Z
Z
Z
Figure I: Different types of calibrated decoders Gaussian VAE, model parameters are denoted with enclosing squares. Left: both the mean and the variance are output by neural network with parameters 0. Center: o-VAE with shared variance, the mean is output by a neural network with parameters 0, but the variance it iself global parameter. Right: the proposed optimal or- o-VAE, the mean output by a neural network with parameters and the variance is computed analytically from the training data D.
We will evaluate these and further choices discrete decoders, described in Appendix D We recommend choosing the decoder distribution that best suits the structure of the data, such discrete decoders for discrete data and continuous decoders for continuous data.
4 OPTIMAL VARIANCE ESTIMATION FOR CALIBRATED GAUSSIAN DECODERS
In this section, we propose simple but novel analytic way of obtaining a calibrated decoder for continuous distributions that further improves performance. The Gaussian decoders with learned variance described in Section 3.1 are calibrated and work better than naïve unit variance decoders. However, for a VAE optimized with gradient descent or Adam (Kingma & Ba 2015). we observe that careful learning rate tuning yield significantly better performance, which is in line with prior work that reported poor performance of gradient descent for optimizing Gaussian distributions (Amari 1998 Peters & Schaal (2008). A smaller learning rate often produces better performance, but slows down the training, as the likelihood values p(z[=) ) will be very suboptimal in the beginning. Instead, here propose analytic solution for the value of a, which computes it analytically and does not require gradient descent.
The maximum likelihood estimate of the variance given known mean is the average squared distance from the mean:
arg max (xruu, o2)) = MSE p),
(5)
where MSE(2, (4) + (x
Eq. 5 can be easily shown using manual differentiation, and a generalization of the fact that the MLE estimate the variance is the sample variance.
2.) -0)02-4.2--0
The optimal variance for the decoder distribution under the maximum likelihood criterion is then simply the average MSE loss over the data and the encoder distribution. We leverage this create an optimal analytic solution for the variance. In the batch setting, the optimal variance would be simply the MSE loss, and can be updated after every gradient update for the other parameters of the decoder. In the mini-batch setting, we use batchwise estimate of the variance computed for the current minibatch. We analyze these approximations in Appendix C At test time, a running average of the variance over the training data is used. This method, which call optimal VAE, allows us to learn very efficiently as we use the optimal variance estimate at every training step. is also easier to implement, needed at test time, it can also simply discarded after training.
separate optimizer for the variance parameter is needed. If the variance is not
Per-image optimal a- VAE. Optimal a- VAE uses single variance value shared across all data points. However, the optimal o- v-E also allows more powerful variance estimates, such as learning a variance value per each pixel, or even a variance value per each image, the difference in implemen- tation simply being the dimensions across which the averaging in Equation 5 operates. This approach can be interpreted as variational variance prediction in the framework of Stirn Knowles (2020).
5 EXPERIMENTAL RESULTS
5

---

Under review as a conference paper at ICLR 2021
O-VAE
Gausian VAE
Goussinn VAE
Figure 2: Images or videos (bottom right) sampled from the proposed optimal a- VAE and unit variance Gaussian VAE models. The Gaussian VAE does not have means to control the expressivity of the latent variable and produces suboptimal, blurry samples. The o-VAE controls the expressivity by learning calibrated decoder, and produces higher quality sequences all datasets.
We now provide an empirical analysis of dif-
Table 1: Analysis of learned variance SVHN. The parameter 3 is tuned manually in 3- VAE and learned in o o-VAE. o o-VAE achieves better perfor- mance, while the value of 3 (implicitly defined via the decoder variance) automatically converges close the value found by manual tuning.
ferent decoding distributions, and validate the benefits of our o- o-VAE approach. We use small convolutional VAE model on SVHN (Net- zer et al., 2011), a larger hierarchical HVAE model (Maalpe et al. 2019) on the CelebA (Liu al. 2015) and CIFAR (Krizhevsky et al. 2009 datasets, and a sequence VAE model called SVG (Denton & Fergus| 2018) the BAIR Pushing dataset (Finn & Levine (017). We evaluate the ELBO values well as visual quality measured by the Fréchet Inception Distance (FID, Heusel al. (2017). Images are 28 x 28 for SVHN and 32 32 for CelebA and CIFAR, while video experiments were performed on 64 64 frames
3
log
FID L
8-VAE
0.001
21.43
44.54
8-VAE
0.01
< -318
27.93
8-VAE
0.1
< --223
28.3
B-VAE
1381
70.39
3-VAE
4056
219.3
--VAE
0.006
< -3333
22.25
following Denton & Fergus (2018). We do not use KL annealing as it did not improve the results in our experiments. Further experimental details are in App.
5.1 DO CALIBRATED DECODERS BALANCE THE VAE ORJECTIVE WITHOUT TUNING B??
As detailed in Section 3.1, B- 3-VAE with unit variance Gaussian decoder commonly used in prior work is equiv- alent σ VAE with constant, manu- ally tuned variance. There is simple relationship between beta and the vari- ance: VB/2. To compare the variance that the v-VAE learns to the manually tuned variance in the case of the 3-VAE, we compare the ELBO values and the corresponding values of 3 in [ We find that learning the variance produces similar values of / to the manually tuned values in the 3-VAE case, indicating that the --VAE is able to learn the balance between the two objective terms in single training run, without hyperpa- rameter tuning. Moreover, the o- VAE outperforms the best 3- 3-VAE run. This is because end-to-end learning produces better estimates of
Optimal o-VAE
A. VAE, B-10
B-VAE, B=1
B-VAE, B=0.1
2.
B-VAE, B-0.01
p V-VAE B-0.001
Figure 3: Analysis of learned variance on SVHN. The param- eter 3 is tuned manually in 3-VAE and learned in o-VAE. Higher values of 3 cause the images to lose detail, while lower values of 3 might make samples unrealistic. The pro- posed optimal a- VAE is able to learn the balance end-to-end, here converging to an equivalent of 3-VAE with , = 0.006.

---

Under review as a conference paper at ICLR 2021
the variance than is possible with manual search, improving the likelihood (as measured by the lower bound) and the visual quality. Figure shows the qualitative results from this experiment.
We further validate our results both single-image and sequential VAE models on range of datasets in SableDaad Figure 2, Single-sample ELBO values are reported, and ELBO values on discretized data are reported for discrete distributions. We see that learning shared variance in Gaussian decoders (shared o- V-VA outperforms the naïve unit variance decoder (Gaussian VAE) as well as tuning the 3 constant for the Gaussian VAE manually. We also see that calibrated discrete decoders, such as full categorical distribution or mixture of discretized logistics, perform better than the naïve Gaussian VAE. Using Bernoulli distribution by treating the color intensities as probabilities (Gregor et al. (015 Watter et 2015) performs poorly. Our results further improve upon the sequence VAE method of Denton & Fergus 2018), which uses unit variance Gaussian with the BAV objective.
/
5.2 How DOES LEARNING CALIBRATED DECODERS IMPACT THE LATENT VARIABLE INFORMATION CONTENT?
We saw above that calibrated decoders result in higher log-likelihood bounds. Are calibrated decoders also beneficial for representation learning? We evaluate the mutual information 1e (x;;a) between the data pa(z)) and encoder samples (=][z), as well as the mismatch between the prior p(z) and the marginal encoder distribution m(2) (p)(z) measured by the marginal KL
DKL (m(() (m=x))lp(z)). These terms are related the rate term of the VAE objective follows (Alemi
ct al. 2017):
Pya(s) (KL(g(<(=))lo<e)) = Epp(()) [D (2))((()))))lm())) (xz(m())llo())
(6)
= Le) (m
That is, the rate term decomposes into the true mutual information and the marginal KL term. We want to learn expressive latent variables with high mutual information. However, doing so by tuning the 3 value relaxes the constraint that
80
(x:z)
KL(m(z) || p(z))
60
the encoder and the prior distributions match, and leads to degraded quality of samples from
the prior, which creates trade-off between ex- pressive representations and ability to generate good samples. To compare the V-VAE and o- VAE in terms of these quantities, we estimate the marginal KL term via Monte Carlo sampling, as proposed by Rosca et al. (2018) and plot the results Figure + As expected, we see that lower 3 values lead to higher mutual informa- tion. However, after a certain point, lower values of] B also cause significant mismatch between the marginal and the prior distributions. By cal- culating the "effective" 3 for the v-VAE, as per Section[4] we can see that the o- VAE captures an inflection point in the (xz(m(=)]pp())) term, learning a representation with the highest pos- sible MI, but without degrading sample quality. This explains the high visual quality of the op- timal a- VAE samples: since the marginal and the prior distributions match, the samples from the prior look similar to reconstructions, while for a 3-VAE with low 3, the samples from the prior poor. We see that, in contrast the 3-VAE, where the mutual information is controlled by hyperparameter, the o-VAE can adjust the appropriate amount of information automatically and is able to find the setting that produces both informative latents and high quality samples.
20
O-VAE (Ours)
10a
10"
10-1
10-2
10-a
Effective beta (log scale)
Figure 4:
Comparison of 3-VAE and o- VAE on
SVHN in terms of mutual information L (z; 2) and marginal KL divergence K {L(m())llp(a) (see Sec. 5.2). Te(:; =) increases with lower B, yielding expressive representations and better reconstruc- tion. However, after a certain point, lowering 8 leads to rapid increase in the marginal KL, yield- ing poor samples from the prior. The o-VAE is able to automatically find the inflection point after which the marginal KL begins to increase, cap- turing much information as possible while still producing good samples.
An alternative discussion of tuning 3 is presented by Alemi et al. (0177). who show that 3 controls the rate-distortion off Here, we show that the crucial trade-off also controlled by / is the
7

---

Under review as a conference paper at ICLR 2021
Table 2: Generative modeling performance of the proposed o v-V on different models and datasets. For SVG, we compare with the original method (Denton & Fergus, 2018), which uses 3- VAE. We see that uncalibrated decoders such as mean-only Gaussian perform poorly. 3-VAE allows to calibrate the decoder but needs careful hyperparameter tuning. Calibrated decoders such as categorical or
a --A perform best. [1] Gregor et al. (2015). [2] Takahashi et al.
2018), [3] Higgins al. 2017).
CelebA HVAE
SVHN VAE
CIFAR HVAE
BAIR SVG
logp
FID 4
logp
FID
logp
FID
logp
FID
Bemoulli VAF [1]
177.6
43.26
284.5
122.6
Categorical VAE
6359
71.5
9179
46.13
7179
101.7
N/A
N/A
cittooc-attgoriial VAE
9067
66.61
< 10800
33.84
9390
91.2
48744
46.13
Logistic mixture VAB
7932
65.3
< 9085
43.19
< 8443
143.1
40616
42.94
Gaussian VAE
7173
186.5
< 2184
112.5
7186
293.7
10379
35.64
Per-pixel d-VAE
-78814
159.3
< 2184
114.7
< 7222
131
-44051
41.98
Student-t VAE [2]
< -8401
71.06
< - 3659
70.4
< -7119
123.6
8-VAE [3]
< -2113
61.6
< -3186
27.93
< --31
103
133472
34.64
Shared V-VAE
6377
60.7
< -3349
22.25
< 54435
116.1
-33974
34.24
Optimal VAE
< 8446
60.3
< -3333
27.25
5677
101.4
14173
34.13
Opt. per-image G- VAB
66.01
26.28
104.0
33,21
trade-off between two components of the rate itself, which control expressivity of representations and the match between the variational and the prior distributions, respectively,
5.3 WHAT ARE THE COMMON CHALLENGES IN LEARNING THE VARIANCE THAT PREVENT PRACTITIONERS FROM USING IT, AND HOW TO RECTIRY THEM?
If learning the decoder variance improves generation, why are learned variances not used more often? In this section, we discuss how the naïve approach to learning variances, where the decoder outputs a variance for each pixel along with the mean, leads to poor results. First, we find that this method often diverges very quickly due to numerical instability, as the network is able to predict certain pixels with very high certainty, leading to degenerate variances. In contrast, learning shared variance always numerically stable in our experiments. We can rectify this numerical instability by bounding the output variance (Section D.Tp. However, even with bounded variance, we observe that learning per-pixel variances Icads to poor results in Table 2 While the per-pixel variance achieves good ELBO value, it produces very poor samples, as measured by FID and visual inspection.
We see that the specific form of learned variance: a shared variance, per-image variance, per-pixel variance, can lead to very different performance in practice. We hypothesize the per-pixel decoder performs poorly it incentivizes the model to focus on particular pixels that can be predicted well, instead of focusing equally on all parts of the image. This is consistent with prior work on denoising diffusion models which noted that likelihood-based models place too much focus imperceptible details, which leads to deteriorated results (Ho et al. 2020). The shared and per-image variance models mitigate this issue the cost of introducing more bias, and work better practice.
5.4 CAN AN ANALYTIC SOLUTION FOR OPTIMAL VARIANCE FURTHER IMPROVE LEARNING?
We evaluate the optimal a VAF which uses analytic solution for the variance (Section 4.. Table shows that achieves superior results in terms log-likell likeliil.oood. We also note that the optimal a- v-E converges to good variance estimate instantancously, which speeds up learning (highlighted in Figure 9 in the Appendix). In addition, we evaluate the per-image optimal a VAE,, in which a single variance is computed per image. This model achieves significantly higher visual quality. While producing this per-image variance with a neural network would require additional architecture tuning. optimal v-VAE is extremely simple implement (it can be implemented simply as changing the axes of summation), not requiring any new tunable parameters.
6 CONCLUSION
We presented a simple and effective method for learning calibrated decoders, as well as an evaluation of different decoding distributions with several VAE and sequential VAE models. The proposed
8

---

Under review as a conference paper at ICLR 2021
method outperforms methods that use naive unit variance Gaussian decoders and tune heuristic weight 3 on the KL-divergence loss, as commonly done in prior work. Morcover, it does not use the heuristic weight B, making it easier to train than this prior work. We expect that the simple techniques for learning calibrated decoders can allow practitioners to speed up the development cycle, obtain better results, and reduce the need for manual hyperparameter tuning.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In 12th {USENIX) Symposium on Operating Systems Design and Implementation ({OSDI) 16), pp. 265-283, 2016.
Alessandro Achille and Stefano Soatto. Information dropout: Learning optimal representations through noisy computation. IEEE transactions on pattern analysis and machine intelligence, 40 (12):2897-2905, 2018
Alexander A Alemi, Ben Poole, Ian Fischer, Joshua V Dillon, Rif A Saurous, and Kevin Murphy, Fixing broken elbo. arXiv preprint arXiv:1711.00464 2017.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-276. 1998.
Georgios Arvanitidis, Lars Kai Hansen, and Spren Hauberg. Latent space oddity: on the curvature of deep generative models. arXiv preprint 1rXiv:1710.111 11379, 2017.
Mohammad Babaeizadeh, Chelsea Finn, Dumitru Erhan, Roy H. Campbell, and Sergey Levine. Stochastic variational video prediction. 2018.
Jonathan T Barron. A general and adaptive robust loss function. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4331-4339, 2019.
Lluis Castrejon, Nicolas Ballas, and Aaron Courville. Improved conditional vrnns for video prediction. In Proceedings of the IEEE International Conference on Computer Vision, pp. 7608 7617, 2019.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint arXiv: 1eee111111.00 1.02731. 2016.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, pp. 3754-476s 2018.
Junyoung Chung, Kyle Kastner, I Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data. 2015.
Bin Dai and David Wipf. Diagnosing and enhancing vae models. arXiv preprint arXiv: 1903.0579, 2019.
A Philip Dawid. The well-calibrated bayesian. Journal of the American Statistical Association, 77 (379):605 610, 1982.
Morris H DeGroot and Stephen E Fienberg. The comparison and evaluation of forecasters. Journal of the Royal Statistical Society: Series D) (The Statistician), 32(1-2): 12-22, 1983.
E. Denton and R. Fergus. Stochastic video generation with learned prior. 2018.
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: A generative model for music. arXiv preprint arXiv: 2020.
Harrison Edwards and Amos Storkey. Towards a neural statistician. arXiv preprint arxiv: 1wX.Xiv:66660222 600.02218. 2016.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. 2017.

---

Under review as a conference paper at ICLR 2021
Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. arXiv preprint arXiv: 903. 12436, 2019.
Karol Gregor, Ivo Danihelka, Alex Graves, Danilo Jimenez Rezende, and Daan Wierstra. Draw: A recurrent neural network for image generation. arXiv preprint arXiv: 152.v:150..04623, 2015.
Karol Gregor, Frederie Besse, Danilo Jimenez Rezende, Ivo Danihelka, and Daan Wierstra. Towards conceptual compression. 2016.
Ishaan Gulrajani, Kundan Kumar, Faruk Ahmed, Adrien Ali Taiga, Francesco Visin, David Vazquez, and Aaron Courville. Pixelvae: latent variable model for natural images. arXiv preprint arXiv: 01116611100.3. 2016.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. arXiv preprint MtXXvv1706.00599) 2017.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1XXv:1912.00100 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. 2019b.
Mikael Henaff, Alfredo Canziani, and Yann LeCun. Model-predictive policy learning with uncertainty regularization for driving in dense traffic. arXiv preprint arXiv: wtt.119111.20.... 2019.
Martin Heusel, Hubert Ramsauer. Thomas Unterthiner, Berhaard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge a local nash equilibrium. In Advances in neural information processing systems, pp. 6626-6637, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
Neural Information Processing Systems, 33, 2020.
Michael I Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, pp. 5574-5584, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 2015.
Diederik Kingma and Max Welling. Auto-encoding variational Bayes. 2014.
Durk Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances in neural information processing systems, pp. 4743-4751, 2016.
Simon Kohl, Bernardino Comera-Pareded Clemens Meyer, Jeffrey De Fauw, Joseph R Ledsam, Klaus Maier-Hein, SM Ali Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. probabilistic u-net for segmentation of ambiguous images. In Advances in Neural Information Processing Systems, pp. 6965-6975, 2018.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
A. X. Lee, R. Zhang, Ebert, P. Abbeel, C. Finn, and S. Levine. Stochastic adversarial video 2018.
prediction. arXiv: 180.0080
Alex X Lee, Anusha Nagabandi, Picter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. arXiv preprint 1Xii:1907.00053, 2019.
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.
10

---

Under review as a conference paper at ICLR 2021
James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't 't blame the elbo! a linear vae perspective on posterior collapse. In Advances in Neural Information Processing Systems, pp. 9403-9413, 2019.
Lars Maaløe, Marco Fraccaro, Valentin Liévin, and Ole Winther. Biva: very deep hierarchy of latent variables for generative modeling. In Advances in neural information processing systems, pp. 0558,6558,, 2019.
Fierre-Alexander Mattei and Jes Frellsen. Leveraging the exact likelihood of deep latent variable models. In Advances in Neural Information Processing Systems, pp. 3855-3866, 2018.
Radford M Neal and Geoffrey E Hinton. A view of the em algorithm that justifies incremental, sparse, and other variants. In Learning in graphical models, pp. 355-366 355-368. Springer, 1998.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-perfore high-performance deep learning library. In Advances in Neural Information Processing Systems, pp. 8024 8024-8035, 2019.
Georgios Pavlakos, Vasileios Choutas, Nima Ghorbani, Timo Bolkart, Ahmed AA Osman, Dimitrios Tzionas, and Michael J Black. Expressive body capture: 3d hands, face, and body from a single image. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 10975-10985, 2019.
Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine. Variational discriminator bottleneck: Improving imitation learning, inverse rl, and gans by constraining information flow. arviv preprint arXiv: 18100001,,,, 2018.
Jan Peters and Stefan Schaal. Reinforcement learning of motor skills with policy gradients. Neural networks, 21(4):682-697, 2008
Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine, Skew-fit: State-covering sele-super reinforcement learning. arXiv preprint arXiv: 0et.:190330.999 2019.
Danilo Jimenez Rezende and Fabio Viola. Taming vaes. arXiv preprint arXiv: 2018.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. 2014.
Jason Tyler Rolfe. Discrete variational autoencoders. arXiv preprint 1eXiv11609002000 2016.
Mihaela Rosca, Balaji Lakshminarayanan, and Shakir Mohamed. Distribution matching in variational
inference. arXiv preprint arXiv: 1rxiv:1802.06844 2018
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv: 701.0170 2017.
Kihyuk Sohn, Honglak Lee, and Xinehen Yan. Learning structured output representation using deep conditional generative models. 2015.
Casper Kaae Sønderby, Tapani Raiko, Lars Maalve, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. In Advances in neural information processing systems, pp. 3738-3746. 2016.
Andrew Stim and David A Knowles. Variational variance: Simple and reliable predictive variance parameterization. arXiv preprint arXiv:2006.049) 2020.
Hiroshi Takahashi, Tomoharu Iwata, Yuki Yamanaka, Masanori Yamada, and Satoshi Yagi. Student-t variational autoencoder for robust density estimation. In UJAAI pp. 2696-2702, 2018.
11

---

Under review as a conference paper at ICLR 2021
3
1
8
2
7
3
2
3
6
0
6
5 18
6
0
5
6
Figure 5: Samples from the a- VAE (left) and the Gaussian VAE (right) the SVHN dataset. The Gaussian VAE produces blurry results with muted colors, while the o-VAE is able to produce accurate images of digits.
Lucas Theis, Aäron van den Oord, and Matthias Bethge. A note on the evaluation of generative models. ICLR, 2016.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in neural information processing systems, pp. 2746-2754, 2015.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Information maximizing variational autoencoders. arXiv preprint Xiv: 7et.10066..2222 2017.
Jun-Yan Zhu, Richard Zhang, Deepak Pathak, Trevor Darrell, Alexei A Efros, Oliver Wang, and Eli Shechtman. Toward multimodal image-to-image translation. Advances in neural information processing systems, pp. 465-476, 2017.
A ADDITIONAL EXPERIMENTAL RESULTS
In this section, we provide more qualitative results in Figures 7 convergence properties of the variance for different models Fi. In order to validate our method with different architecture, we also report performance of different decoders with small 5- convolutional architecture on the CelebA and CIFAR dataset Table 3 We see that the ordering of the methods is consistent with this smaller architecture.
6 8 5 as well as a graph showing the
B EXPERIMENTAL DETAILS
For the small convolutional network test SVHN, the encoder has 3 convolutional layers followed by a fully connected layer, while the decoder has a fully connected layer followed by convolutional layers. The was tuned from 100 0.0001 for 3-VAE. The number of channels in the convolutional layers starts with 32 and increases 2 times in every layer. The dimension of the latent variable is 20. Adam (Kingma & Ba 2015) with learning rate of le-3 is used for optimization. Batch size of 128 was used and all models were trained for 10 epochs. We additionally evaluate this small convolutional network on CelebA, CIFAR, and Frey datasess in Table Unit Gaussian prior and Gaussian posteriors with diagonal covariance were used. For the larger hicrarchical VAE, we used the official pytorch implementation of (Maaløe et al. 2019). We use the baseline hierarchical VAE with 15 layers of latent variables, without the top-down and bottom-up connections. For the hierarchical VAE and
datasets
Available https: //tps:///s...e nyu. edu/ rowei 8 //data..... htmml
12

---

Under review as a conference paper at ICLR 2021
Figure 6: Samples from the V-E (left) and the Gaussian VAE (right) on the CelebA dataset, images cropped to the face for clarity. The Gaussian VAE produces blurry results with indistinct face features, while the VAE is able produce accurate images faces.
Table 3: Generative modeling performance of the proposed a- VAE on CelebA, CIFAR, and Frey Face with smaller model. We see that uncalibrated decoders such as mean-only Gaussian perform poorly. 3-VAE allows to calibrate the decoder but needs careful hyperparameter tuning. Calibrated decoders such as categorical or 3-VEE perform best.
CelebA VAE
CIFAR VAE
Frey Face VAE
log 4
FID
logg
FID 4
log
FID
Bernoulli VAE Gregor et al. (2015)
102.7
165.1
47.7
Categorical VAE
10199
50.45
106773
124.1
2454
50.16
bitwise-can categorcaat VAE
11019
56.36
11604
99.65
< 3173
66.77
Logistic mixture VAE
110154
61.81
100648
100.2
2562
50.28
Gaussian VAE
2201
144.8
1409
205.8
< 726.4
80.17
8-VAE Higgins et al. 2017
< 1942
58.73
< -1318
117.9
< -420.0
37.61
Shared VAE (Ours)
< 1939
73.27
-1830
137.8
< -49.78
42.86
Optimal a- VAE (Ours)
1191
61.27
-1832
80.9
< -1622
53.36
Opt. per-image V-E (Ours)
53.13
89.88
56.07
13

---

Under review as a conference paper at ICLR 2021
Figure 7: Samples from the a- VAE (top) and the Gaussian VAE (bottom) on the BAIR dataset. Sampled sequences conditioned on two initial frames are shown, and the ground truth sequence shown at the top. The Gaussian VAE produces blurry robot arm texture and the arm often disappears towards the end the sequence, while the VAE is able to produce sequences with realistic motion and model the details of the arm texture, such as the gripper.
14

---

Under review as a conference paper at ICLR 2021
Figure 8: Samples from the a- VAE (left) and the Gaussian VAE (right) on the challenging CIFAR dataset. The Gaussian VAE produces blurry results with muted colors, while the a --AE models the distribution of shapes in the CIFAR data more faithfully.
-0.5
optimal o-VAE shared o-VAE B-VAE
-1.
1.5
b
-
-2.0
-25
330
0
2
4
6
Epochs
Figure 9: Variance convergence speed on SVHN. We see that the shared o- VAE which optimizes the variance with gradient descent has initial period of convergence when the variance converges to the region of the optimal value. In contrast, o-VAE with analytical (optimal) variance quickly learns good estimate of the variance, which leads better performance. The unit variance Gaussian 3-VAE can be interpreted as having a constant variance determined by 3, shown here. Since the variance doesn' 't change throughout training, it achieves suboptimal performance.
15

---

Under review as a conference paper at ICLR 2021
the SVG-LP model, we use the default hyperparameters in the respective implementations. We use the standard train-val-test split for all datasets. All models were trained on a single high-end GPU. We the official Py Torch implementation of the Inception network to compute FID. All methods are compared on the same hyperparameters.
C EMPIRICAL ANALYZIS OF APPROXIMATIONS FOR OPTIMAL 0- VAE
The optimal o- o-VAE requires computing the following estimate of the variance
o² arg max Dan [Inp(zluo(2), of]]] o²
Ee (SE,
(7)
This requires computing expectations, with respect to the data in the dataset, and with respect to the encoder distribution. We use MC sampling with one sample per data point to approximate both expectations. Inspired by common practices VAEs, we use one sample per data point to approximate the inner expectation. On SVHN, the standard error of this approximation is 0.26% of the value of sigma. We further approximate the outer expectation with single batch instead of the entire dataset. On SVHN, the standard error of this approximation is 2% of the value of sigma. see that both approximations are accurate practice. The second approximation yields a biased estimate of the evidence lower bound because the same batch is used to approximate the variance and compute the lower bound estimate. However, this bias can be corrected by using a different batch, or with running average of the variance with appropriate decay. This running average also be used to reduce the variance of the estimate and to achieve convergence guarantees, but we did not find it necessary in experiments.
D ALTERNATIVE DECODER CHOICES
We describe the alternative decoders evaluated in Table 2 using the
Table 4: ELBO on discretized data. All distributions ex- cept categorical have scalar scale parameters. The a- --V performs well on the discretized ELBO metric, performing similarfy to a discrete distribution parametrized a dis- cretized Gaussian or discretized Logistic. Full categorical distribution attains highest likelihood due to having the most statistical power.
citwise-ategorrica, and the logistic mixture distributions.
Bitwise-categorical VAE While the 256-way categorical decoder described in Section 3.2 very powerful due to the ability to specify any possible intensity distribution,
CIFAR VAE
it suffers from high computational and memory requirements. Because 256 values need to be kept for each pixel and channel, simply keeping this distribution in memory for one 3-channel 1024 x 1024 image would require 3 GiB of memory, compared to 0.012 GiB for the Gaussian decoder. Therefore, training deep neural networks with this full categorical distribution is impractical for high-resolution images or videos. The stwise-categoriee VAE improves the memory complexity by defining the distribution over 256 values in more compact way. Specifically, it defines binary distribution over each bit in the pixel intensity value, requiring values in total, one for each bit. This distribution can thought of as classifier that predicts the value of each bit in the image separately. In our implementation of the bitwiseccategorical likelihood, we convert the image channels to binary format and use the standard binary cross-entropy loss (which reduces to binary logelikelihood since all bits in the image are deterministically either zero or one). While our experiments the bitwise- catriie-cateooiicall distribution did not outperform other choices, often performs par with our proposed method. We expect this distribution to be useful due to its generality as it is able to represent values stored in any digital format by converting them into binary.
log pdf
log
FID $
Categorical VAE Gaussian VAE
10673
137.6
740.5 <8966.1
15131
212.7
Gaussian a- VAE Disc. Gaussian a P-VAE Disc. Logistic a- e-V
11120
136.7
11117
136.9
11103
136.7
16

---

Under review as a conference paper at ICLR 2021
Logistic mixture VAE For this decoder, we adapt the discretized logistic mixture from Salimans et al. (0017). To define discrete 256-way distribution, it divides the corresponding continuous distribution into 256 bins, where the probability mass is defined the integral of the PDF over the corresponding bin. (Kingma et al. (016) uses the logistic distribution discretized in this manner for the decoder. [Salimans al. (20T)) suggests to make all bins except the first and the last be of equal size, whereas the first and the last bin include, respectively, the intervals (-o0, 0J and [L, 00). Salimans al. (2017) further suggests using mixture of discretized logistics for improved capacity. Our implementation largely follows the one in Salimans et al. (2017). however, we note that the original implementation is not suitable for learning latent variable models, as it generates the channels autoregressively. This will cause the latent variable to lose color information since it can be represented by the autoregressive decoder. We therefore adapt the mixture of discretized logistics to the pure latent variable setup by removing the mean-adjusting coefficients from (Salimans et al. 2017). In our experiments, the logistic mixture outperformed other discrete distributions.
17


Reviews:
[{"review_id": "FSUibfFYtlE", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**GENERAL**\n\nThe paper presents a more in-sight analysis of using learnable variance in Gaussian decoders in the Variational Auto-Encoder framework. The authors follow up on other papers in the literature and try to answer the following research question:\n- If a Gaussian decoder is used, is it better to fix the variance (e.g., \\sigma=1), or learn it?\nIn general, the paper is fine, however, it lacks novelty and it does not highlight that the VAE framework is a probabilistic framework, and a Gaussian decoder is not appropriate for modeling images. Nevertheless, I must admit that the analysis is properly performed.\n\n**Strengths:**\n\nS1: An in-depth analysis of using a shared variance across dimensions in Gaussian decoders.\n\nS2: Comparing both NLL and FID is a good indication of showing importance of learning variance (or, more generally speaking, modeling uncertainty) in decoders in VAEs.\n\n\n**Remarks:**\n\nR1: VAEs constitute a sub-class of latent variables models with prescribed distributions (namely, all distributions are known in advance). In opposition to implicit models, the distributions must be picked accordingly to observed quantities (e.g., images). I am totally aware that many authors use Gaussian decoders for modeling images, however, doing it naively is simply wrong. The support of a normal distribution is [-\\infty, +\\infty], while images are typically represented by integers in [0, ..., 255]. Therefore, even if pixel values are normalized to [0, 1], they take one of 256 possible values. As a result, using Gaussian decoders is inappropriate.\n\nIt is possible to use dequantization as it is typically done in flow-based models, see:\n- Theis, L., van den Oord, A., and Bethge, M. A note on the evaluation of generative models. ICLR 2016\n- Ho, J., Chen, X., Srinivas, A., Duan, Y., and Abbeel, P. Flow++: Improving flow-based generative models with variational dequantization and architecture design. ICML 2019\n- Winkler, C., Worrall, D., Hoogeboom, E., and Welling, M. Learning Likelihoods with Conditional Normalizing Flows. arXiv preprint, 2019\n- Hoogeboom, E., Cohen, T. S., and Tomczak, J. M. Learning Discrete Distributions by Dequantization. arXiv preprint, 2020\n\nHowever, this is not the case in this paper.\nVAEs have a strong probabilistic foundation, and it should be treated as a probabilistic model. However, taking inappropriate distributions to model data, it results in:\n(i) a wrong model;\n(ii) propagating a wrong message in the deep learning community that any loss function works for VAEs.\n\nI would highly appreciate if the authors would make it very clear in the paper, and start with a remark about choosing appropriate distribution for observed data. Afterwards, it could be explained that we need some sort of dequantization to utilize Gaussian decoders.\n\nR2: The authors indicate a connection between a VAE with a Gaussian decoder and a \\beta-VAE framework. The connection is very clear from the optimization perspective, but it is not the case from the modeling standpoint. If we take a look at the objective and consider the optimization process, then indeed, there is no difference in potential optima, because multiplying by \\sigma or \\beta results in the same objective. However, this has different consequences from the modeling perspective. \\beta-VAE is a class of stochastic Auto-Endoders where the objective, the reconstruction error, is expanded by adding a Lagrangian multiplier. In other words, adding the equality constraint KL(q(z)||p(z)) = 0. However, the authors propose a valid VAE, i.e., the objective is a valid lower-bound to the log-likelihood function. I find it very confusing to indicate the connection without being very precise in what sense these two approaches are related to each other. Similarly to the remark R1, it could be very confusing in the DL community, and it could propagate a message that the VAE framework is a DL framework, while neural networks are an important (or even crucial) component of a probabilistic framework.\n\nR3: As indicated by the authors, a similar approach was already discussed in the literature. Sharing a single variance across all dimensions was discussed in other papers. However, I must admit, not in-depth as here.\n\nR4: It would be beneficial to discuss the following paper:\n- Ghosh, P., Sajjadi, M. S., Vergari, A., Black, M., & Schölkopf, B. From variational to deterministic autoencoders. ICLR 2020\n\nIt is closely related and it could be see as an important other perspective on the problem.\n\n*AFTER REBUTTAL*\nI would like to thank the authors for their hard work! I increase my score to 6.", "extra_scores": {}, "extra_text": {}}, {"review_id": "5kwV_3Leffo", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper claims two contributions: 1. Proposes a connection between beta-VAE with fixed variance Gaussian decoder and VAE with variable variance Gaussian decoder 2. Proposes to optimize the variance of a Gaussian decoder\n\nThis paper uses a simple but useful method, but I have some concerns about novelty and soundness of the experiments. I will give a score on the low side for now. \n\nPro: \n\nThe connection between beta-VAE and variable variance Gaussian decoder is somewhat interesting. \n\nThe proposal that the variance can be adjusted for each individual image seems novel, but currently there does not seem to be a good theoretical/empirical justification for this. \n\nThe writing is quite good. Easy to follow and clear. \n\nCon: \n\nI think the main claimed contribution has limited novelty. The main difference between the current work and the prior common practice is to learn a single variance shared by all the pixels (instead of learning a variance per pixel). This by itself is fine, the lack of novelty can be made up by very convincing experiments and good practical guidance. My major concern is that the paper does not sufficiently justify its choice with experiments:\n\nThe baselines are surprisingly bad. For example, the paper says it uses the model in Maaloe et al, 2019 for celebA; comparing the samples in the current paper and the original paper, there is a very big difference. Granted, the original paper has a somewhat different setup, but the sample quality is too unreasonable to be convincing (I think VAE samples looked like that in 2013). In addition, the FID score of baselines also seem to be much worse than typically reported in the literature. For example, a FID score of around 50 is already typical for VAE models in 2019 (e.g. see Dai et al, 2019), while the current paper report an FID of 186 for baselines. Such a big discrepancy does not instill confidence that the baselines are properly trained with current recommended practices. \n\nFollowing up on the above point, typically some kind of annealing of the beta parameter is the standard practice to avoid converging to a poor local minimum (which seems to be what is happening to the baselines here). It would be nice to have some discussion / comparison. In fact, I think it is completely fine if the argument is to show that reasoning effort to choose an annealing schedule do not lead to optimal results. However, I think this point does not come out from the experiments. \n\nThe criticized alternative (learning a variance per pixel) has been widely used in the literature and the samples look quite reasonable (e.g. see http://ruishu.io/2018/03/14/vae/) unlike the baselines in the paper. Granted there are a few implementation tricks (such as bounding the output to avoid numerical instability), but these are quite mild. It is difficult to conclude that the proposed approach uses less practical implementation techniques to achieve these results. \n\nMinor comments:\n\nThe terminology calibration can have many different but precise meanings. For example, when people say a probability forecaster is calibrated, there is a specific property the forecaster must satisfy. I don’t think the usage of the terminology calibration matches any standard usage, and hence can be confusing. \n\nThe title is not very informative. In particular, with the current interpretation of \"calibration\" which is to \"learn better probabilities\", I think any VAE paper could have used that title. \n\n\n--------- \n\nThank you for the detailed reply and revisions. I have increased my review score because several of my concerns have been addressed. I am not entirely convinced that the baselines use current best practices, and several claims in the paper regarding the inductive biases and calibration. Nevertheless I think the algorithms proposed in the paper is practically useful. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "-pLD2RNXRr9", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper discusses a well-known problem of VAE training that decoder produces blurry reconstruction with constant variance. While much existing work addressed this problem by introducing independent variance training (as of the original VAE model) or additional hyper-parameters, those approaches usually come with additional training/tuning difficulty and even break the ELBO assumption. This paper proposed a simple $\\sigma$-VAE that addresses the above problem by optimizing a single variance variable. This also could be easily connected to the well known $\\beta$-VAE works. The experiment results in Tables 2 and 3 show the proposed model obtains a better FID score than the existing works on multiple datasets.\n\nIn general, I am not surprised that shared variance would work well in practice as it is somewhat obvious as lots of works with beta-VAE to get better performance. However, it is surprising that no one aligns beta-VAE with the variance of observation before. \n\nThe most valuable knowledge I learned from this paper is written in one small section about variance implementation details where a new variable $lambda$ is introduced to avoid numeric problems. It is more helpful than the main claim of the paper somewhat in practice.\n\nPros\n===\n1. The paper is well written, and the authors provide sufficient evidence to show the proposed model works well in practice. The experiment design is also intuitive.\n2. Mutual information analysis is great for reading.\n\nCons\n===\n1. The proposed approach is too simple to validate its novelty since it seems a rewritten of the existing formula by assigning the beta hyper-parameter with another meaning. \n2. Figure 8 is hard to believe as Gaussian VAE on CIFAR-10 should perform much better than it is shown here. Many GitHub implementation shows far better-generated images.\n\nAgain, it is quite surprising that the method is not introduced before the year 2020. I have an impression where some of the VAE implementations did include things like proposed in this paper (but fail to find them now).\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to (re-)examine VAEs with calibrated uncertainties for the likelihood, which is say VAEs in which the variance is learned rather than chosen as a fixed hyperparameter. The authors argue that doing so provides a reasonable means of automatically navigating the tradeoff between minimizing the distortion (the reconstruction loss) and the rate (the KL loss) in the variational objective. In particular, the authors propose to use a diagonal covariance  Σ = σ^2 Ι that is shared across pixels, and note that it is trivial to define  σ(z) = MSE(x, μ(z)) on a per-image basis to minimize the reconstruction loss. 

This is very much a borderline paper. Reviewers appreciate that the writing is clear, and acknowledge that revisiting the idea of learning calibrated is of interest to the community. At the same time, the reviewers note that the proposed approach has very limited technical novelty, and note problems with the experimental evaluation. 

The metareviewer has read the paper, and is critical of the framing of this work. The manuscript in its current form does not do a sufficiently good job of discussing the large and detailed literature that exists on this topic. Learning calibrated decoders is by no means new, which this submission could and should acknowledge much more clearly. The two seminal papers on VAEs both considered learning calibrated decoders. Moreover there is a lack of thoughtful discussion of the reasons why learning a pixel-wise σ(z) is not common practice. The authors note that this can lead to problems with training stability, but fail to note that this problem is mathematically ill-posed; A well-known property of VAEs is that high-capacity models will memorize the training data, in the sense that the optimal learned marginal likelihood is equal to the empirical distribution over the training set (i.e. a mixture over delta peaks). 

The metareviewer would expect to see a more thoughtful discussion of  the long line of work on navigating the trade-off between rate and distortion, as well as the role of model capacity. A good place to start would be a more careful discussion of the autoencoding and autodecoding limits (Alemi et al 2018) and the GECO paper (Rezende et al 2018). More broadly, the metareviewer would expect some discussion of approaches that improve the quality of generation such as [1], and work that considers effect of model capacity on generalization, such as [2].  

In terms of experimental evaluation, this paper also somewhat falls short. As R4 notes, some of the results look worryingly bad, which may be due to the fact that the authors train for only 10 epochs (as indicated in  Appendix B). Moreover, what is once again lacking in experiments is a systematic consideration of the role of model capacity. Some comparison to more recent baselines than the β-VAE (e.g. GECO) would also be helpful here.
 
The metareviewer is sympathetic to the basic premise of this paper, which is the claim that learning a σ that is shared across pixels is a pretty good best practice in terms of finding a reasonable balance between rate and distortion. There is certainly room for a paper that communicates this idea. However, such a paper should (a) more explicitly position itself as revisiting this idea rather than introducing this idea, (b) include a more thoughtful discussion of related work, and (c) include a more robust empirical evaluation. 

[1] Engel, J., Hoffman, M. & Roberts, A. Latent Constraints: Learning to Generate Conditionally from Unconditional Generative Models. arXiv:1711.05772 [cs, stat] (2017).

[2] Shu, R., Bui, H. H., Zhao, S., Kochenderfer, M. J. & Ermon, S. Amortized inference regularization. in Proceedings of the 32nd International Conference on Neural Information Processing Systems 4398–4407 (Curran Associates Inc., 2018).
Decision: Reject