Original: S IMPLE AND E FFECTIVE VAE T RAINING WITH C ALIBRATED D ECODERS
Revised: S IMPLE AND E FFECTIVE VAE T RAINING WITH C ALIBRATED D ECODERS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Variational autoencoders (VAEs) provide an effective and simple method for modeling complex distributions.
Revised: Variational autoencoders (VAEs) provide an effective and simple method for modeling complex distributions.

Original: However, training VAEs often requires considerable hyperparameter tuning to determine the optimal amount of information retained by the latent variable.
Revised: However, training VAEs often requires considerable hyperparameter tuning to determine the optimal amount of information retained by the latent variable.

Original: 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.
Revised: 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.

Original: 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 modiﬁcations instead.
Revised: 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 modiﬁcations instead.

Original: We perform the ﬁrst comprehensive comparative analysis of calibrated decoder and provide recommendations for simple and effective VAE training.
Revised: We perform the ﬁrst comprehensive comparative analysis of calibrated decoder and provide recommendations for simple and effective VAE training.

Original: Our analysis covers a range of datasets and several single-image and sequential VAE models.
Revised: Our analysis covers a range of datasets and several single-image and sequential VAE models.

Original: We further propose a simple but novel modiﬁcation to the commonly used Gaussian decoder, which computes the prediction variance analytically.
Revised: We further propose a simple but novel modiﬁcation to the commonly used Gaussian decoder, which computes the prediction variance analytically.

Original: We observe empirically that using heuristic modiﬁcations is not necessary with our method.
Revised: We observe empirically that using heuristic modiﬁcations is not necessary with our method.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep density models based on the variational autoencoder (VAE) (Kingma & Welling, 2014; Rezende et 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.
Revised: Deep density models based on the variational autoencoder (VAE) (Kingma & Welling, 2014; Rezende et 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.

Original: 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 et al., 2020), as well as model-based reinforcement learning (Hafner et al., 2019a).
Revised: 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 et al., 2020), as well as model-based reinforcement learning (Hafner et al., 2019a).

Original: 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).
Revised: 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).

Original: This balance trades off ﬁdelity of reconstruction and quality of samples from the model: a 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).
Revised: This balance trades off ﬁdelity of reconstruction and quality of samples from the model: a 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).

Original: While a proper variational lower bound does not expose any free parameters to control this tradeoff, many prior works heuristically introduce a weight on the prior KL-divergence term, often denoted β .
Revised: While a proper variational lower bound does not expose any free parameters to control this tradeoff, many prior works heuristically introduce a weight on the prior KL-divergence term, often denoted β .

Original: Usually, β needs to be tuned for every dataset and model variant as a hyperparameter, which slows down development and can lead to poor performance as ﬁnding the optimal value is often prohibitively computationally expensive.
Revised: Usually, β needs to be tuned for every dataset and model variant as a hyperparameter, which slows down development and can lead to poor performance as ﬁnding the optimal value is often prohibitively computationally expensive.

Original: Moreover, using β (cid:54) =precludes the appealing interpretation of the VAE objective as a bound on the data likelihood, and is undesirable for applications like density modeling.
Revised: Moreover, using β (cid:54) =precludes the appealing interpretation of the VAE objective as a bound on the data likelihood, and is undesirable for applications like density modeling.

Original: While many architectures for calibrating decoders have been proposed in the literature (Kingma &
Revised: While many architectures for calibrating decoders have been proposed in the literature (Kingma &

Original: Welling, 2014; Kingma et al., 2016; Dai & Wipf, 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;
Revised: Welling, 2014; Kingma et al., 2016; Dai & Wipf, 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;

Original: Lee et al., 2019; Babaeizadeh et al., 2018; Lee et al., 2018; Hafner et al., 2019b; Pong et al., 2019;
Revised: Lee et al., 2019; Babaeizadeh et al., 2018; Lee et al., 2018; Hafner et al., 2019b; Pong et al., 2019;

Original: Zhu et al., 2017; Pavlakos et al., 2019), or uses other ad-hoc modiﬁcations to the objective (Sohn et al., 2015; Henaff et al., 2019).
Revised: Zhu et al., 2017; Pavlakos et al., 2019), or uses other ad-hoc modiﬁcations to the objective (Sohn et al., 2015; Henaff et al., 2019).

Original: Indeed, it is well known that attempting to learn the variance in a
Revised: Indeed, it is well known that attempting to learn the variance in a

Original: Gaussian decoder may lead to numerical instability (Rezende & Viola, 2018; Dai & Wipf, 2019), and na¨ıve approaches often lead to poor results.
Revised: Gaussian decoder may lead to numerical instability (Rezende & Viola, 2018; Dai & Wipf, 2019), and na¨ıve approaches often lead to poor results.

Original: As a result, it remains unclear whether practical empirical performance of VAEs actually beneﬁts from calibrated decoders or not.
Revised: As a result, it remains unclear whether practical empirical performance of VAEs actually beneﬁts from calibrated decoders or not.

Original: To rectify this, our ﬁrst contribution is a comparative analysis of various calibrated decoder architectures and practical recommendations for simple and effective VAE training.
Revised: To rectify this, our ﬁrst contribution is a comparative analysis of various calibrated decoder architectures and practical recommendations for simple and effective VAE training.

Original: We ﬁnd that, while na¨ıve calibrated decoders often lead to worse results, a careful choice of the decoder distribution can work very well, and removes the need to tune the additional parameter β .
Revised: We ﬁnd that, while na¨ıve calibrated decoders often lead to worse results, a careful choice of the decoder distribution can work very well, and removes the need to tune the additional parameter β .

Original: Indeed, we note that the entropy of the decoding distribution controls the mutual information I ( x ; z ) .
Revised: Indeed, we note that the entropy of the decoding distribution controls the mutual information I ( x ; z ) .

Original: Calibrated decoders allow the model to control I ( x ; z ) automatically, instead of relying on manual tuning.
Revised: Calibrated decoders allow the model to control I ( x ; z ) automatically, instead of relying on manual tuning.

Original: Our second contribution is a simple but novel technique for optimizing the decoder variance analytically, without requiring the decoder network to produce it as an additional output.
Revised: Our second contribution is a simple but novel technique for optimizing the decoder variance analytically, without requiring the decoder network to produce it as an additional output.

Original: We call the resulting approach to learning the
Revised: We call the resulting approach to learning the

Original: Gaussian variance the σ -VAE.
Revised: Gaussian variance the σ -VAE.

Original: In our experiments, the σ -VAE outperforms the alternative of learning the variance through gradient descent, while being simpler to implement and extend.
Revised: In our experiments, the σ -VAE outperforms the alternative of learning the variance through gradient descent, while being simpler to implement and extend.

Original: We validate our results on several VAE and sequence VAE models and a range of image and video datasets.
Revised: We validate our results on several VAE and sequence VAE models and a range of image and video datasets.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Prior work on variational autoencoders has studied a number of different decoder parameterizations.
Revised: Prior work on variational autoencoders has studied a number of different decoder parameterizations.

Original: Kingma & Welling (2014);
Revised: Kingma & Welling (2014);

Original: Rezende et al.
Revised: Rezende et al.

Original: (2014) use the Bernoulli distribution for the binary MNIST data and Kingma & Welling (2014) use Gaussian distributions with learned variance parameter for grayscale images.
Revised: (2014) use the Bernoulli distribution for the binary MNIST data and Kingma & Welling (2014) use Gaussian distributions with learned variance parameter for grayscale images.

Original: 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).
Revised: 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).

Original: Some work has attempted to rectify this problem by using dequantization (Gregor et al., 2016), which is theoretically appealing as 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), or training a post-hoc prior (Ghosh et al., 2019).
Revised: Some work has attempted to rectify this problem by using dequantization (Gregor et al., 2016), which is theoretically appealing as 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), or training a post-hoc prior (Ghosh et al., 2019).

Original: Takahashi et al.
Revised: Takahashi et al.

Original: ; Barron (2019) proposed more expressive distributions.
Revised: ; Barron (2019) proposed more expressive distributions.

Original: Additionally, different choices for representing such variance exist, including diagonal covariance (Kingma & Welling, 2014;
Revised: Additionally, different choices for representing such variance exist, including diagonal covariance (Kingma & Welling, 2014;

Original: Sønderby et al., 2016; Rolfe, 2016), or a single shared parameter (Kingma et al., 2016; Dai & Wipf, 2019; Edwards & Storkey, 2016; Rezende & Viola, 2018).
Revised: Sønderby et al., 2016; Rolfe, 2016), or a single shared parameter (Kingma et al., 2016; Dai & Wipf, 2019; Edwards & Storkey, 2016; Rezende & Viola, 2018).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Further work has explored various parameterizations based on discretized continuous distributions, such as discretized logistic (Kingma et al., 2016).
Revised: Further work has explored various parameterizations based on discretized continuous distributions, such as discretized logistic (Kingma et al., 2016).

Original: More recent work has improved expressivity of the decoder with a mixture of discretized logistics (Chen et al., 2016; Maaløe et al., 2019).
Revised: More recent work has improved expressivity of the decoder with a mixture of discretized logistics (Chen et al., 2016; Maaløe et al., 2019).

Original: However, these models also employ powerful autoregressive decoders (Chen et al., 2016; Gulrajani et al., 2016; Maaløe et al., 2019), and the latent variables in these models may not represent all of the signiﬁcant factors of variation in the data, as some factors can instead be modeled internally by the autoregressive decoder (Alemi et al., 2017).
Revised: However, these models also employ powerful autoregressive decoders (Chen et al., 2016; Gulrajani et al., 2016; Maaløe et al., 2019), and the latent variables in these models may not represent all of the signiﬁcant factors of variation in the data, as some factors can instead be modeled internally by the autoregressive decoder (Alemi et al., 2017).

Original: While many calibrated decoders have been proposed, outside the core generative modeling community uncalibrated decoders are ubiquitous.
Revised: While many calibrated decoders have been proposed, outside the core generative modeling community uncalibrated decoders are ubiquitous.

Original: They are used in work on video prediction (Denton & Fergus, 2018;
Revised: They are used in work on video prediction (Denton & Fergus, 2018;

Original: Castrejon et al., 2019; Lee et al., 2018; Babaeizadeh et al., 2018), image segmentation (Kohl et al., 2018), image-to-image translation (Zhu et al., 2017), 3D human pose (Pavlakos et al., 2019), as well as model-based reinforcement learning (Henaff et al., 2019; Hafner et al., 2019b;a), and representation learning (Lee et al., 2019; Watter et al., 2015; Pong et al., 2019).
Revised: Castrejon et al., 2019; Lee et al., 2018; Babaeizadeh et al., 2018), image segmentation (Kohl et al., 2018), image-to-image translation (Zhu et al., 2017), 3D human pose (Pavlakos et al., 2019), as well as model-based reinforcement learning (Henaff et al., 2019; Hafner et al., 2019b;a), and representation learning (Lee et al., 2019; Watter et al., 2015; Pong et al., 2019).

Original: Most of these works utilize the heuristic hyperparameter β instead, which is undesirable both as the resulting objective is no longer a bound on the likelihood, and as β usually requires extensive tuning.
Revised: Most of these works utilize the heuristic hyperparameter β instead, which is undesirable both as the resulting objective is no longer a bound on the likelihood, and as β usually requires extensive tuning.

Original: In this work, we analyze the common pitfalls of using calibrated decoders that may have prevented the practitioners from using them, propose a simple and effective analytic way of learning such calibrated distribution, and provide a comprehensive experimental evaluation of different decoding distributions.
Revised: In this work, we analyze the common pitfalls of using calibrated decoders that may have prevented the practitioners from using them, propose a simple and effective analytic way of learning such calibrated distribution, and provide a comprehensive experimental evaluation of different decoding distributions.

Original: Alternative discussions of the hyperparameter β are presented by Zhao et al.
Revised: Alternative discussions of the hyperparameter β are presented by Zhao et al.

Original: ; Higgins et al.
Revised: ; Higgins et al.

Original: Alemi et al.
Revised: Alemi et al.

Original: ; Achille & Soatto (2018), who show that it controls the amount of information in the latent variable, I ( x ; z ) .
Revised: ; Achille & Soatto (2018), who show that it controls the amount of information in the latent variable, I ( x ; z ) .

Original: Peng et al.
Revised: Peng et al.

Original: ; Rezende & Viola (2018) further discuss constrained optimization objectives for VAEs, which also yield a similar hyperparameter.
Revised: ; Rezende & Viola (2018) further discuss constrained optimization objectives for VAEs, which also yield a similar hyperparameter.

Original: Here, we focus on β -VAEs with Gaussian decoders with constant variance, as commonly used in recent work, and show that the hyperparameter β can be incorporated in the decoding likelihood for these models.
Revised: Here, we focus on β -VAEs with Gaussian decoders with constant variance, as commonly used in recent work, and show that the hyperparameter β can be incorporated in the decoding likelihood for these models.

Original:  C ALIBRATED D ECODING D ISTRIBUTIONS
Revised:  A NALYSING D ECODING D ISTRIBUTIONS

Original: The generative model of a VAE (Kingma & Welling, 2014; Rezende et al., 2014) with parameters θ is speciﬁed with a prior distribution over the latent variable p θ ( z ) , commonly unit Gaussian, and a decoding distribution p θ ( x | z ) , which for color images is commonly a conditional Gaussian parameterized with a neural network.
Revised: The generative model of a VAE (Kingma & Welling, 2014; Rezende et al., 2014) with parameters θ is speciﬁed with a prior distribution over the latent variable p θ ( z ) , commonly unit Gaussian, and a decoding distribution p θ ( x | z ) , which for color images is commonly a conditional Gaussian parameterized with a neural network.

Original: We would like to ﬁt this generative model to a given dataset by maximizing the evidence lower bound (ELBO (Neal & Hinton, 1998; Jordan et al., 1999; Kingma & Welling, 2014; Rezende et al., 2014)), which uses an approximate posterior distribution q φ ( z | x ) , also commonly a conditional Gaussian speciﬁed with a neural network.
Revised: We would like to ﬁt this generative model to a given dataset by maximizing the evidence lower bound (ELBO (Neal & Hinton, 1998; Jordan et al., 1999; Kingma & Welling, 2014; Rezende et al., 2014)), which uses an approximate posterior distribution q φ ( z | x ) , also commonly a conditional Gaussian speciﬁed with a neural network.

Original: In this work, we focus on the form of the decoding distribution p θ ( x | z ) .
Revised: In this work, we focus on the form of the decoding distribution p θ ( x | z ) .

Original: To achieve the best results, we want a decoding distribution that represents the required probability p ( x | z ) accurately, that is, we want the decoding distribution to be calibrated in the statistical sense. In this section, we 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.
Revised: To achieve the best results, we want a decoding distribution that represents the required probability p ( x | z ) accurately In this section, we 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.

Original: 3. G AUSSIAN D ECODERS
Revised: 3. G AUSSIAN D ECODERS

Original: We ﬁrst analyse the commonly used Gaussian decoders.
Revised: We ﬁrst analyse the commonly used Gaussian decoders.

Original: We note that the commonly used MSE reconstruction loss between the reconstruction ˆ x and ground truth data x is equivalent to the negative log-likelihood objective with a Gaussian decoding distribution with constant variance:
Revised: We note that the commonly used MSE reconstruction loss between the reconstruction ˆ x and ground truth data x is equivalent to the negative log-likelihood objective with a Gaussian decoding distribution with constant variance:

Original: [Equation]
Revised: [Equation]

Original: where p ( x | z ) ∼ N (ˆ x, I ) , the prediction ˆ x is produced with a neural network ˆ x = µ θ ( z ) , and D is the dimensionality of x .
Revised: where p ( x | z ) ∼ N (ˆ x, I ) , the prediction ˆ x is produced with a neural network ˆ x = µ θ ( z ) , and D is the dimensionality of x .

Original: This demonstrates a drawback of methods that rely simply on the MSE loss (Castrejon et al., 2019;
Revised: This demonstrates a drawback of methods that rely simply on the MSE loss (Castrejon et al., 2019;

Original: Denton & Fergus, 2018; Lee et al., 2019; Hafner et al., 2019b; Pong et al., 2019; Zhu et al., 2017;
Revised: Denton & Fergus, 2018; Lee et al., 2019; Hafner et al., 2019b; Pong et al., 2019; Zhu et al., 2017;

Original: Henaff et al., 2019), as it is equivalent to assuming a particular, constant variance of the Gaussian decoding distribution.
Revised: Henaff et al., 2019), as it is equivalent to assuming a particular, constant variance of the Gaussian decoding distribution.

Original: By learning this variance, we can achieve much better performance due to better calibration of the decoder.
Revised: By learning this variance, we can achieve much better performance due to better calibration of the decoder.

Original: There are several ways in which we can specify this variance.
Revised: There are several ways in which we can specify this variance.

Original: 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;
Revised: 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;

Original: Sønderby et al., 2016; Rolfe, 2016).
Revised: Sønderby et al., 2016; Rolfe, 2016).

Original: This can be done, for example, by letting a neural network σ θ output the diagonal entries of the covariance matrix given a latent sample z :
Revised: This can be done, for example, by letting a neural network σ θ output the diagonal entries of the covariance matrix given a latent sample z :

Original: [Equation]
Revised: [Equation]

Original: This parameterization of the decoding distribution outputs one variance value per each pixel and channel.
Revised: This parameterization of the decoding distribution outputs one variance value per each pixel and channel.

Original: While powerful, we observe in Section 5.3 that this approach attains suboptimal performance, and is moreover prone to numerical instability.
Revised: While powerful, we observe in Section 5.3 that this approach attains suboptimal performance, and is moreover prone to numerical instability.

Original: Instead, we will ﬁnd experimentally that a simpler parameterization, in which the covariance matrix is speciﬁed with a single shared (Kingma et al., 2016; Dai & Wipf, 2019; Edwards & Storkey, 2016; Rezende & Viola, 2018) parameter σ as Σ = σI often works better in practice:
Revised: Instead, we will ﬁnd experimentally that a simpler parameterization, in which the covariance matrix is speciﬁed with a single shared (Kingma et al., 2016; Dai & Wipf, 2019; Edwards & Storkey, 2016; Rezende & Viola, 2018) parameter σ as Σ = σI often works better in practice:

Original: [Equation]
Revised: [Equation]

Original: The parameter σ can be optimized together with parameters of the neural network θ with gradient descent.
Revised: The parameter σ can be optimized together with parameters of the neural network θ with gradient descent.

Original: Of particular interest is the interpretation of this parameter.
Revised: Of particular interest is the interpretation of this parameter.

Original: Writing out the expression for the decoding likelihood, we obtain−
Revised: Writing out the expression for the decoding likelihood, we obtain

Original: [Equation]
Revised: [Equation]

Original: The full objective of the resulting Gaussian σ -VAE is:
Revised: The full objective of the resulting Gaussian σ -VAE is:

Original: [Equation]
Revised: [Equation]

Original: Note that σ may be viewed as a weighting parameter between the MSE reconstruction term and the
Revised: Note that σ may be viewed as a weighting parameter between the MSE reconstruction term and the

Original: KL-divergence term in the objective.
Revised: KL-divergence term in the objective.

Original: Moreover, this objective explicitly speciﬁes 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.
Revised: Moreover, this objective explicitly speciﬁes 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.

Revised: Decoder Calibration It is 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 x given that prediction (DeGroot & Fienberg, 1983; Dawid, 1982).

Revised: The calibration of a neural network can be usually improved by estimating the uncertainty of that prediction (Guo et al., 2017), such as the variance of a Gaussian (Kendall & Gal, 2017).

Revised: Since the naive MSE loss assumes a constant variance, it does not effectively represent the uncertainty of the prediction, and is often poorly calibrated.

Original:  Connection to β -VAE.
Revised: Instead, learning the variance as in Eq. 3 leads to better uncertainty estimation and better calibration. In Sec 5.1, we show that learning a good estimate of this uncertainty is crucial for the quality of the VAE generations. Connection to β -VAE.

Original: The β -VAE objective (Higgins et al., 2017) for a Gaussian decoder with unit variance is:
Revised: The β -VAE objective (Higgins et al., 2017) for a Gaussian decoder with unit variance is:

Original: [Equation]
Revised: [Equation]

Original: We see that it can be interpreted as a particular case of the objective (3), where the variance is constant and the term D ln σ can be ignored during optimization.
Revised: We see that it can be interpreted as a particular case of the objective (3), where the variance is constant and the term D ln σ can be ignored during optimization.

Original: The β -VAE objective is then equivalent to a σ -VAE with a constant variance σ = (cid:112) β/ 2 (for a particular learning rate setting).
Revised: The β -VAE objective is then equivalent to a σ -VAE with a constant variance σ = (cid:112) β/ 2 (for a particular learning rate setting).

Original: In recent work (Zhu et al., 2017; Denton & Fergus, 2018; Lee et al., 2019), β -VAE models are often used in this exact regime.
Revised: In recent work (Zhu et al., 2017; Denton & Fergus, 2018; Lee et al., 2019), β -VAE models are often used in this exact regime.

Original: By tuning the β term, practitioners are able to tune the variance of the decoder, manually producing a more calibrated decoder.
Revised: By tuning the β term, practitioners are able to tune the variance of the decoder, manually producing a more calibrated decoder.

Original: However, by re-interpreting the β -VAE objective as a special case of the VAE, we can both obtain a valid evidence lower bound, and remove the need to manually select β .
Revised: However, by re-interpreting the β -VAE objective as a special case of the VAE and introducing the missing D ln σ term, we can both obtain a valid evidence lower bound, and remove the need to manually select β .

Original: Instead, the variance σ can instead simply be learned end-to-end, reducing the need for hyperparameter tuning.
Revised: Instead, the variance σ can instead simply be learned end-to-end, reducing the need for hyperparameter tuning.

Original: An alternative discussion of this connection in the context of linear VAEs is also presented by Lucas et al.
Revised: An alternative discussion of this connection in the context of linear VAEs is also presented by Lucas et al.

Original: While the β term is not necessary for good performance if the decoder is calibrated, it can still be employed if desired, such as when the aim is to attain better disentanglement (Higgins et al., 2017) or a particular rate-distortion tradeoff (Alemi et al., 2017).
Revised: While the β term is not necessary for good performance if the decoder is calibrated, it can still be employed if desired, such as when the aim is to attain better disentanglement (Higgins et al., 2017) or a particular rate-distortion tradeoff (Alemi et al., 2017).

Original: However, we found that with calibrated decoders, the best sample quality is obtained when β
Revised: However, we found that with calibrated decoders, the best sample quality is obtained when β

Original: = 1 .
Revised: = 1 .

Original: Loss implementation details.
Revised: Loss implementation details.

Original: 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.
Revised: 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.

Original: 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 signiﬁcantly different from the number of the latent dimensions.
Revised: 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 signiﬁcantly different from the number of the latent dimensions.

Original: While this can be conveniently ignored in the β -VAE regime, where the balance term is tuned manually anyway, for the σ -VAE it is essential to compute the objective value correctly.
Revised: While this can be conveniently ignored in the β -VAE regime, where the balance term is tuned manually anyway, for the σ -VAE it is essential to compute the objective value correctly.

Original: Variance implementation details.
Revised: Variance implementation details.

Original: Since the variance is non-negative, we parameterize it logarithmically as σ 2 = e 2 λ , where λ is the logarithm of the standard deviation.
Revised: Since the variance is non-negative, we parameterize it logarithmically as σ 2 = e 2 λ , where λ is the logarithm of the standard deviation.

Original: For some models, such as per-pixel variance decoders, we observed that it is necessary to restrict the variance range for numerical stability.
Revised: For some models, such as per-pixel variance decoders, we observed that it is necessary to restrict the variance range for numerical stability.

Original: We do so by using the soft clipping operations proposed by Chua et al.
Revised: We do so by using the soft clipping operations proposed by Chua et al.

Original: [Equation]
Revised: [Equation]

Original: We observe that setting λ min = − 6 to lower bound the standard deviation to be at least half of the distance between allowed color values works well in practice.
Revised: We observe that setting λ min = − 6 to lower bound the standard deviation to be at least half of the distance between allowed color values works well in practice.

Original: We also observe that this clipping is unnecessary when learning a shared σ value.
Revised: We also observe that this clipping is unnecessary when learning a shared σ value.

Original: 3. D ISCRETE D ECODERS
Revised: 3. D ISCRETE D ECODERS

Original: It is possible to use discrete decoding distributions to generate images, as color values are commonly restricted to a ﬁxed set of integer pixel intensities (e.g. 0..255).
Revised: It is possible to use discrete decoding distributions to generate images, as color values are commonly restricted to a ﬁxed set of integer pixel intensities (e.g. 0..255).

Original: Indeed, for discrete color values, discrete distributions are arguably more appropriate.
Revised: Indeed, for discrete color values, discrete distributions are arguably more appropriate.

Original: In the most general case, a discrete decoding distribution factorized per each pixel and channel would be speciﬁed by a probability mass vector ˆ x with 256 entries, one per each possible intensity value, similarly to a per-pixel classiﬁer of the intensity value.
Revised: In the most general case, a discrete decoding distribution factorized per each pixel and channel would be speciﬁed by a probability mass vector ˆ x with 256 entries, one per each possible intensity value, similarly to a per-pixel classiﬁer of the intensity value.

Original: We can implement it with a soft-max layer, yielding the following log-likelihood loss (sometimes called the cross-entropy loss) for a true pixel with intensity i : 
Revised: We can implement it with a soft-max layer, yielding the following log-likelihood loss (sometimes called the cross-entropy loss) for a true pixel with intensity i : [Equation] θ σ z μ

Original: [Equation]
Revised: [Equation]

Original:  We will evaluate these and further choices of discrete decoders, described in Appendix D. We recommend choosing the decoder distribution that best suits the structure of the data, such as discrete decoders for discrete data and continuous decoders for continuous data.
Revised: σ z μ θ σ z μ D We will evaluate these and further choices of discrete decoders, described in Appendix D. We recommend choosing the decoder distribution that best suits the structure of the data, such as discrete decoders for discrete data and continuous decoders for continuous data.

Original:  O PTIMAL V ARIANCE E STIMATION FOR C ALIBRATED G AUSSIAN D ECODERS
Revised:  O PTIMAL V ARIANCE E STIMATION FOR C ALIBRATED G AUSSIAN D ECODERS

Original: In this section, we propose a simple but novel analytic way of obtaining a calibrated decoder for continuous distributions that further improves performance.
Revised: In this section, we propose a simple but novel analytic way of obtaining a calibrated decoder for continuous distributions that further improves performance.

Original: The Gaussian decoders with learned[Equation] θ σ z μ D variance described in Section 3.1 are calibrated and work better than na¨ıve unit variance decoders.
Revised: The Gaussian decoders with learned variance described in Section 3.1 are calibrated and work better than na¨ıve unit variance decoders.

Original: However, for σ -VAE optimized with gradient descent or Adam (Kingma & Ba, 2015), we observe that careful learning rate tuning can yield signiﬁcantly 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).
Revised: However, for σ -VAE optimized with gradient descent or Adam (Kingma & Ba, 2015), we observe that careful learning rate tuning can yield signiﬁcantly 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).

Original: A smaller learning rate often produces better performance, but slows down the training, as the likelihood values p ( x | z ) will be very suboptimal in the beginning.
Revised: A smaller learning rate often produces better performance, but slows down the training, as the likelihood values p ( x | z ) will be very suboptimal in the beginning.

Original: Instead, here we propose an analytic solution for the value of σ , which computes it analytically and does not require gradient descent.
Revised: Instead, here we propose an analytic solution for the value of σ , which computes it analytically and does not require gradient descent.

Original: The maximum likelihood estimate of the variance given a known mean is the average squared distance from the mean:
Revised: The maximum likelihood estimate of the variance given a known mean is the average squared distance from the mean:

Original: [Equation]
Revised: [Equation]

Original: ) where MSE ( x, µ ) = 1 D (cid:80) i ( x i − µ i ) 2 .
Revised: where MSE ( x, µ ) = 1 D (cid:80) i ( x i − µ i ) 2 .

Original: Eq.
Revised: Eq.

Original: 5 can be easily shown using manual differentiation, and is a generalization of the fact that the MLE estimate of the variance is the sample variance.
Revised: 5 can be easily shown using manual differentiation, and is a generalization of the fact that the MLE estimate of the variance is the sample variance.

Original: 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.
Revised: 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.

Original: We leverage this to create an optimal analytic solution for the variance.
Revised: We leverage this to create an optimal analytic solution for the variance.

Original: 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.
Revised: 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.

Original: In the mini-batch setting, we use a batchwise estimate of the variance computed for the current minibatch.
Revised: In the mini-batch setting, we use a batchwise estimate of the variance computed for the current minibatch.

Original: We analyze these approximations in Appendix C. At test time, a running average of the variance over the training data is used.
Revised: We analyze these approximations in Appendix C. At test time, a running average of the variance over the training data is used.

Original: This method, which we call optimal σ -VAE , allows us to learn very efﬁciently as we use the optimal variance estimate at every training step.
Revised: This method, which we call optimal σ -VAE , allows us to learn very efﬁciently as we use the optimal variance estimate at every training step.

Original: It is also easier to implement, as no separate optimizer for the variance parameter is needed.
Revised: It is also easier to implement, as no separate optimizer for the variance parameter is needed.

Original: If the variance is not needed at test time, it can also be simply discarded after training.
Revised: If the variance is not needed at test time, it can also be simply discarded after training.

Original: Per-image optimal σ -VAE.
Revised: Per-image optimal σ -VAE.

Original: Optimal σ -VAE uses a single variance value shared across all data points.
Revised: Optimal σ -VAE uses a single variance value shared across all data points.

Original: However, the optimal σ -VAE 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 implementation simply being the dimensions across which the averaging in Equation 5 operates.
Revised: However, the optimal σ -VAE 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 implementation simply being the dimensions across which the averaging in Equation 5 operates.

Original: This approach can be interpreted as variational variance prediction in the framework of Stirn & Knowles (2020).
Revised: This approach can be interpreted as variational variance prediction in the framework of Stirn & Knowles (2020).

Original:  E XPERIMENTAL R ESULTS
Revised:  E XPERIMENTAL R ESULTS

Revised: The parameter β is tuned manually in β -VAE and learned in σ -VAE.

Revised: σ -VAE achieves better performance, while the value of β (implicitly deﬁned via the decoder variance) automatically converges close the value found by manual tuning.

Revised: [Table]

Original: We now provide an empirical analysis of different decoding distributions, and validate the beneﬁts of our σ -VAE approach.
Revised: We now provide an empirical analysis of different decoding distributions, and validate the beneﬁts of our σ -VAE approach.

Original: We use a small convolutional VAE model on SVHN (Netzer et al., 2011), a larger hierarchical HVAE model (Maaløe et al., 2019) on the CelebA (Liu et al., 2015) and CIFAR (Krizhevsky et al., 2009) datasets, and a sequence VAE model called SVG (Denton & Fergus, 2018) on the BAIR Pushing dataset (Finn & Levine, 2017).
Revised: We use a small convolutional VAE model on SVHN (Netzer et al., 2011), a larger hierarchical HVAE model (Maaløe et al., 2019) on the CelebA (Liu et al., 2015) and CIFAR (Krizhevsky et al., 2009) datasets, and a sequence VAE model called SVG (Denton & Fergus, 2018) on the BAIR Pushing dataset (Finn & Levine, 2017).

Original: We evaluate the ELBO values as well as visual quality measured by the Fr´echet Inception Distance (FID, Heusel et al.
Revised: We evaluate the ELBO values as well as visual quality measured by the Fr ´ echet Inception Distance (FID, Heusel et al.

Original: Images are 28 ×for SVHN and 32 × 32 for CelebA and CIFAR, while video experiments were performed on 64 ×frames following Denton & Fergus (2018).
Revised: Images are 28 × 28 for SVHN and 32 × 32 for CelebA and CIFAR, while video experiments were performed on 64 × 64 frames following Denton & Fergus (2018).

Original: We do not use KL annealing as it did not improve the results in our experiments.
Revised: We do not use KL annealing as it did not improve the results in our experiments.

Original: Further experimental details are in App.
Revised: Further experimental details are in App.

Original: B.
Revised: B.

Original: 5. D O CALIBRATED DECODERS BALANCE 
Revised: 5. D O CALIBRATED DECODERS BALANCE THE VAE OBJECTIVE WITHOUT TUNING β ?

Original: [Figure]

Original: Higher values of β cause the images to lose detail, while lower values of β might make samples unrealistic.
Revised: Higher values of β cause the images to lose detail, while lower values of β might make samples unrealistic.

Original: The proposed optimal σ -VAE is able to learn the balance end-to-end, here converging to an equivalent of β -VAE with β = 0 .
Revised: The proposed optimal σ -VAE is able to learn the balance end-to-end, here converging to an equivalent of β -VAE with β = 0 .

Original: As detailed in Section 3.1, a β -VAE with a unit variance Gaussian decoder commonly used in prior work is equivalent to a σ -VAE with constant, manually tuned variance.
Revised: As detailed in Section 3.1, a β -VAE with a unit variance Gaussian decoder commonly used in prior work is equivalent to a σ -VAE with constant, manually tuned variance.

Original: There is a simple relationship between beta and the variance: σ = (cid:112) β/ 2 .
Revised: There is a simple relationship between beta and the variance: σ = (cid:112) β/ 2 .

Original: To compare the variance that the σ -VAE learns to the manually tuned variance in the case of the β -VAE, we compare the ELBO values and the corresponding values of β in Table 1.
Revised: To compare the variance that the σ -VAE learns to the manually tuned variance in the case of the β -VAE, we compare the ELBO values and the corresponding values of β in Table 1.

Original: We ﬁnd that learning the variance produces similar values of β to the manually tuned values in the β -VAE case, indicating that the σ -VAE is able to learn the balance between the two objective terms in a single training run, without hyperparameter tuning.
Revised: We ﬁnd that learning the variance produces similar values of β to the manually tuned values in the β -VAE case, indicating that the σ -VAE is able to learn the balance between the two objective terms in a single training run, without hyperparameter tuning.

Original: Moreover, the σ -VAE outperforms the best β -VAE run.
Revised: Moreover, the σ -VAE outperforms the best β -VAE run.

Original: This is because end-to-end learning produces better estimates of the variance than is possible with manual search, improving the likelihood (as measured by the lower bound) and the visual quality.
Revised: This is because end-to-end learning produces better estimates of the variance than is possible with manual search, improving the likelihood (as measured by the lower bound) and the visual quality.

Original: Figure 3 shows the qualitative results from this experiment.
Revised: Figure 3 shows the qualitative results from this experiment.

Original: The parameter β is tuned manually in β -VAE and learned in σ -VAE.

Original: σ -VAE achieves better performance, while the value of β (implicitly deﬁned via the decoder variance) automatically converges close the value found by manual tuning.

Original: [Table] We further validate our results on both singleimage and sequential VAE models on a range of datasets in Table 2 and Figure 2.
Revised: We further validate our results on both single-image and sequential VAE models on a range of datasets in Table 2 and Figure 2.

Original: Single-sample ELBO values are reported, and ELBO values on discretized data are reported for discrete distributions.
Revised: Single-sample ELBO values are reported, and ELBO values on discretized data are reported for discrete distributions.

Original: We see that learning a shared variance in a Gaussian decoders (shared σ -VAE) outperforms the na¨ıve unit variance decoder (Gaussian VAE) as well as tuning the β constant for the Gaussian VAE manually.
Revised: We see that learning a shared variance in a Gaussian decoders (shared σ -VAE) outperforms the na¨ıve unit variance decoder (Gaussian VAE) as well as tuning the β constant for the Gaussian VAE manually.

Original: 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.
Revised: 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.

Original: Using Bernoulli distribution by treating the color intensities as probabilities (Gregor et al., 2015; Watter et al., 2015) performs poorly.
Revised: Using Bernoulli distribution by treating the color intensities as probabilities (Gregor et al., 2015; Watter et al., 2015) performs poorly.

Original: Our results further improve upon the sequence VAE method of Denton & Fergus (2018), which uses a unit variance Gaussian with the β -VAE objective.
Revised: Our results further improve upon the sequence VAE method of Denton & Fergus (2018), which uses a unit variance Gaussian with the β -VAE objective.

Original: 5. H OW DOES LEARNING CALIBRATED DECODERS IMPACT THE LATENT VARIABLE INFORMATION CONTENT ?
Revised: 5. H OW DOES LEARNING CALIBRATED DECODERS IMPACT THE LATENT VARIABLE INFORMATION CONTENT ?

Original: We saw above that calibrated decoders result in higher log-likelihood bounds.
Revised: We saw above that calibrated decoders result in higher log-likelihood bounds.

Original: Are calibrated decoders also beneﬁcial for representation learning?
Revised: Are calibrated decoders also beneﬁcial for representation learning?

Original: We evaluate the mutual information I e ( x ; z ) between the data
Revised: We evaluate the mutual information I e ( x ; z ) between the data

Original: p d ( x ) and encoder samples q ( z | x ) , as well as the mismatch between the prior p ( z ) and the marginal encoder distribution m ( z ) =
Revised: p d ( x ) and encoder samples q ( z | x ) , as well as the mismatch between the prior p ( z ) and the marginal encoder distribution m ( z ) =

Original: E p d ( x ) q ( z | x ) , measured by the marginal KL
Revised: E p d ( x ) q ( z | x ) , measured by the marginal KL

Original: D KL ( m ( z ) || p
Revised: D KL ( m ( z ) || p

Original: ( z )) .
Revised: ( z )) .

Original: These terms are related to the rate term of the VAE objective as follows (Alemi et al., 2017):
Revised: These terms are related to the rate term of the VAE objective as follows (Alemi et al., 2017):

Original: [Equation]
Revised: [Equation]

Revised: SVHN in terms of mutual information I e ( x ; z ) and marginal KL divergence KL ( m ( z ) || p ( z ))

Original:  I  z  increases with lower β , yielding expressive representations and better reconstruction.
Revised: (see Sec. 5.2). I e ( x ; z ) increases with lower β , yielding expressive representations and better reconstruction.

Original: However, after a certain point, lowering β leads to a rapid increase in the marginal KL, yielding poor samples from the prior.
Revised: However, after a certain point, lowering β leads to a rapid increase in the marginal KL, yielding poor samples from the prior.

Original: The σ -VAE is able to automatically ﬁnd the inﬂection point after which the marginal KL begins to increase, capturing as much information as possible while still producing good samples.
Revised: The σ -VAE is able to automatically ﬁnd the inﬂection point after which the marginal KL begins to increase, capturing as much information as possible while still producing good samples.

Original: That is, the rate term decomposes into the true mutual information and the marginal KL term.
Revised: That is, the rate term decomposes into the true mutual information and the marginal KL term.

Original: We want to learn expressive latent variables with high mutual information.
Revised: We want to learn expressive latent variables with high mutual information.

Original: However, doing so by tuning the β value relaxes the constraint that the encoder and the prior distributions match, and leads to degraded quality of samples from the prior, which creates a trade-off between expressive representations and ability to generate good samples.
Revised: However, doing so by tuning the β value relaxes the constraint that the encoder and the prior distributions match, and leads to degraded quality of samples from the prior, which creates a trade-off between expressive representations and ability to generate good samples.

Original: To compare the β -VAE and σVAE in terms of these quantities, we estimate the marginal KL term via Monte Carlo sampling, as proposed by Rosca et al.
Revised: To compare the β -VAE and σVAE in terms of these quantities, we estimate the marginal KL term via Monte Carlo sampling, as proposed by Rosca et al.

Original: (2018), and plot the results in Figure 4.
Revised: (2018), and plot the results in Figure 4.

Original: As expected, we see that lower β values lead to higher mutual information.
Revised: As expected, we see that lower β values lead to higher mutual information.

Original: However, after a certain point, lower values of β also cause a signiﬁcant mismatch between the marginal and the prior distributions.
Revised: However, after a certain point, lower values of β also cause a signiﬁcant mismatch between the marginal and the prior distributions.

Original: By calculating the “effective” β for the σ -VAE, as per
Revised: By calculating the “effective” β for the σ -VAE, as per

Original: Section 4, we can see that the σ -VAE captures an inﬂection point in the D KL ( m ( z ) || p ( z )) term, learning a representation with the highest possible MI, but without degrading sample quality.
Revised: Section 4, we can see that the σ -VAE captures an inﬂection point in the D KL ( m ( z ) || p ( z )) term, learning a representation with the highest possible MI, but without degrading sample quality.

Original: This explains the high visual quality of the optimal σ -VAE samples: since the marginal and the prior distributions match, the samples from the prior look similar to reconstructions, while for a β -VAE with low β , the samples from the prior are poor.
Revised: This explains the high visual quality of the optimal σ -VAE samples: since the marginal and the prior distributions match, the samples from the prior look similar to reconstructions, while for a β -VAE with low β , the samples from the prior are poor.

Original: We see that, in contrast to the β -VAE, where the mutual information is controlled by a hyperparameter, the σ -VAE can adjust the appropriate amount of information automatically and is able to ﬁnd the setting that produces both informative latents and high quality samples.
Revised: We see that, in contrast to the β -VAE, where the mutual information is controlled by a hyperparameter, the σ -VAE can adjust the appropriate amount of information automatically and is able to ﬁnd the setting that produces both informative latents and high quality samples.

Original: An alternative discussion of tuning β is presented by Alemi et al.
Revised: An alternative discussion of tuning β is presented by Alemi et al.

Original: (2017), who show that β controls the rate-distortion trade-off.
Revised: (2017), who show that β controls the rate-distortion trade-off.

Original: Here, we show that the crucial trade-off also controlled by β is thetrade-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. W HAT ARE THE COMMON CHALLENGES IN LEARNING THE VARIANCE THAT PREVENT PRACTITIONERS FROM USING IT , AND HOW TO RECTIFY 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.
Revised: Here, we show that the crucial trade-off also controlled by β is the

Original: First, we ﬁnd that this method often diverges very quickly due to numerical instability, as the network is able to predict

Original: For SVG, we compare with the original method (Denton & Fergus, 2018), which uses β -VAE.
Revised: For SVG, we compare with the original method (Denton & Fergus, 2018), which uses β -VAE.

Original: We see that uncalibrated decoders such as mean-only Gaussian perform poorly.
Revised: We see that uncalibrated decoders such as mean-only Gaussian perform poorly.

Original: β -VAE allows to calibrate the decoder but needs careful hyperparameter tuning.
Revised: β -VAE allows to calibrate the decoder but needs careful hyperparameter tuning.

Original: Calibrated decoders such as categorical or σ -VAE perform best.
Revised: Calibrated decoders such as categorical or σ -VAE perform best.

Original: [1] Gregor et al. (2015), [2] Takahashi et al.
Revised: [1] Gregor et al. (2015), [2] Takahashi et al.

Original: (2018), [3] Higgins et al.
Revised: (2018), [3] Higgins et al.

Original: [Table]
Revised: [Table]

Revised: 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.

Original:  certain pixels with very high certainty, leading to degenerate variances.
Revised: 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 ﬁnd 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.

Original: In contrast, learning a shared variance is always numerically stable in our experiments.
Revised: In contrast, learning a shared variance is always numerically stable in our experiments.

Original: We can rectify this numerical instability by bounding the output variance (Section 3.1).
Revised: We can rectify this numerical instability by bounding the output variance (Section 3.1).

Original: However, even with bounded variance, we observe that learning per-pixel variances leads to poor results in Table 3.
Revised: However, even with bounded variance, we observe that learning per-pixel variances leads to poor results in Table 2.

Original: While the per-pixel variance achieves a good ELBO value, it produces very poor samples, as measured by FID and visual inspection.
Revised: While the per-pixel variance achieves a good ELBO value, it produces very poor samples, as measured by FID and visual inspection.

Revised: We see that the speciﬁc form of learned variance: a shared variance, a per-image variance, or a per-pixel variance, can lead to very different performance in practice.

Original: We hypothesizethat the per-pixel decoder allocates signiﬁcant capacity to learning the variance, which prevents it from learning to produce good samples
Revised: We hypothesize the per-pixel decoder performs poorly as it incentivizes the model to focus on particular pixels that can be predicted well, instead of focusing equally on all parts of the image.

Original: [Table]
Revised: This is consistent with prior work on denoising diffusion models which noted that likelihood-based models place too much focus on imperceptible details, which leads to deteriorated results (Ho et al., 2020). The shared and per-image variance models mitigate this issue at the cost of introducing more bias, and work better in practice.

Original: 5. C AN AN ANALYTIC SOLUTION FOR OPTIMAL VARIANCE FURTHER IMPROVE LEARNING ?
Revised: 5. C AN AN ANALYTIC SOLUTION FOR OPTIMAL VARIANCE FURTHER IMPROVE LEARNING ?

Original: surprisingly performs poorly.

Original: However, calibrated decoders such as Bernoulli or shared σ -VAE perform best, improving over

Original: [Table]

Original: We evaluate the optimal σ -VAE which uses an analytic solution for the variance (Section 4).
Revised: We evaluate the optimal σ -VAE which uses an analytic solution for the variance (Section 4).

Original: Table 2 shows that it achieves superior results in terms of log-likelihood.
Revised: Tableshows that it achieves superior results in terms of log-likelihood.

Original: We also note that the optimal σ -VAE converges to a good variance estimate instantaneously, which speeds up learning (highlighted in Figure 9 in the Appendix).
Revised: We also note that the optimal σ -VAE converges to a good variance estimate instantaneously, which speeds up learning (highlighted in Figure 9 in the Appendix).

Original: In addition, we evaluate the per-image optimal σVAE, in which a single variance is computed per image.
Revised: In addition, we evaluate the per-image optimal σ -VAE, in which a single variance is computed per image.

Original: This model achieves signiﬁcantly higher visual quality.
Revised: This model achieves signiﬁcantly higher visual quality.

Original: While producing this per-image variance with a neural network would require additional architecture tuning, optimal σ -VAE is extremely simple to implement (it can be implemented simply as changing the axes of summation), not requiring any new tunable parameters.
Revised: While producing this per-image variance with a neural network would require additional architecture tuning, optimal σ -VAE is extremely simple to implement (it can be implemented simply as changing the axes of summation), not requiring any new tunable parameters.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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.
Revised: 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.

Original: The proposed method outperforms methods that use na¨ıve unit variance Gaussian decoders and tune a heuristic weight β on the KL-divergence loss, as commonly done in prior work.
Revised: The proposed method outperforms methods that use na¨ıve unit variance Gaussian decoders and tune a heuristic weight β on the KL-divergence loss, as commonly done in prior work.

Original: Moreover, it does not use the heuristic weight β , making it easier to train than this prior work.
Revised: Moreover, it does not use the heuristic weight β , making it easier to train than this prior work.

Original: 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.
Revised: 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.

Original: A A DDITIONAL EXPERIMENTAL RESULTS
Revised: A A DDITIONAL EXPERIMENTAL RESULTS

Original: In this section we provide more qualitative results in Figures 7, 6, 8, 5 as well as a graph showing the convergence properties of the variance for different models in Fig.
Revised: In this section we provide more qualitative results in Figures 7, 6, 8, 5 as well as a graph showing the convergence properties of the variance for different models in Fig.

Original: In order to validate our method with a different architecture, we also report performance of different decoders with a small 5-layer convolutional architecture on the CelebA and CIFAR dataset in Table 4.
Revised: In order to validate our method with a different architecture we also report performance of different decoders with a small 5-layer convolutional architecture on the CelebA and CIFAR dataset in Table 3.

Original: We see that the ordering of the methods is consistent with this smaller architecture.shown
Revised: We see that the ordering of the methods is consistent with this smaller architecture.

Original: Gaussian VAE produces blurry robot arm texture and the arm often disappears towards the end of the while the is able to produce sequences with and the details gradient descent has an initial period of convergence when the variance converges to the region of the optimal value.

Original: In contrast, σ -VAE with analytical (optimal) variance quickly learns a good estimate of the variance, which leads to better performance.

Original: The unit variance Gaussian β -VAE can be interpreted as having a constant variance determined by β , shown here.

Original: Since the variance doesn’t change throughout training, it achieves suboptimal performance.

Original: [Table]

Original: B E XPERIMENTAL DETAILS
Revised: B E XPERIMENTAL DETAILS

Original: For the small convolutional network test on SVHN, the encoder has 3 convolutional layers followed by a fully connected layer, while the decoder has a fully connected layer followed by 3 convolutional layers.
Revised: For the small convolutional network test on SVHN, the encoder has 3 convolutional layers followed by a fully connected layer, while the decoder has a fully connected layer followed by 3 convolutional layers.

Original: The β was tuned from 100 to 0.0001 for β -VAE.
Revised: The β was tuned from 100 to 0.0001 for β -VAE.

Original: The number of channels in the convolutional layers starts with 32 and increases 2 times in every layer.
Revised: The number of channels in the convolutional layers starts with 32 and increases 2 times in every layer.

Original: The dimension of the latent variable is 20.
Revised: The dimension of the latent variable is 20.

Original: Adam (Kingma & Ba, 2015) with learning rate of 1e-3 is used for optimization.
Revised: Adam (Kingma & Ba, 2015) with learning rate of 1e-3 is used for optimization.

Original: Batch size ofwas used and all models were trained for 10 epochs.
Revised: Batch size ofwas used and all models were trained for 10 epochs.

Original: We additionally evaluate this small convolutional network on CelebA, CIFAR, and Frey Face 2 datasets in Table 4.
Revised: We additionally evaluate this small convolutional network on CelebA, CIFAR, and Frey Face 2 datasets in Table 3.

Original: Unit Gaussian prior and Gaussian posteriors with diagonal covariance were used.
Revised: Unit Gaussian prior and Gaussian posteriors with diagonal covariance were used.

Original: For the larger hierarchical VAE, we used the ofﬁcial pytorch implementation of (Maaløe et al., 2019).
Revised: For the larger hierarchical VAE, we used the ofﬁcial pytorch implementation of (Maaløe et al., 2019).

Original: We use the baseline hierarchical VAE with 15 layers of latent variables, without the top-down and bottom-up connections. 
Revised: 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 [Table] −

Revised: [Table]

Revised: −

Original: For the hierarchical VAE  and  the SVG-LP model, we use the default hyperparameters in the respective implementations.
Revised: [Table] − [Table] VAE blurry robot arm texture and and gradient descent has an initial period of convergence when the variance converges to the region of the optimal value.

Revised: In contrast, σ -VAE with analytical (optimal) variance quickly learns a good estimate of the variance, which leads to better performance.

Revised: The unit variance Gaussian β -VAE can be interpreted as having a constant variance determined by β , shown here.

Revised: Since the variance doesn’t change throughout training, it achieves suboptimal performance.

Revised: the SVG-LP model, we use the default hyperparameters in the respective implementations.

Original: We use the standard train-val-test split for all datasets.
Revised: We use the standard train-val-test split for all datasets.

Original: All models were trained on a single high-end GPU.
Revised: All models were trained on a single high-end GPU.

Original: We use the ofﬁcial PyTorch implementation of the Inception network to compute FID.
Revised: We use the ofﬁcial PyTorch implementation of the Inception network to compute FID.

Original: All methods are compared on the same hyperparameters.
Revised: All methods are compared on the same hyperparameters.

Original: C E MPIRICAL ANALYZIS OF APPROXIMATIONS FOR OPTIMAL σ -VAE
Revised: C E MPIRICAL ANALYZIS OF APPROXIMATIONS FOR OPTIMAL σ -VAE

Original: The optimal σ -VAE requires computing the following estimate of the variance
Revised: The optimal σ -VAE requires computing the following estimate of the variance

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: This requires computing two expectations, with respect to the data in the dataset, and with respect to the encoder distribution.
Revised: This requires computing two expectations, with respect to the data in the dataset, and with respect to the encoder distribution.

Original: We use MC sampling with one sample per data point to approximate both expectations.
Revised: We use MC sampling with one sample per data point to approximate both expectations.

Original: Inspired by common practices in VAEs, we use one sample per data point to approximate the inner expectation.
Revised: Inspired by common practices in VAEs, we use one sample per data point to approximate the inner expectation.

Original: On SVHN, the standard error of this approximation is 0 .
Revised: On SVHN, the standard error of this approximation is 0 .

Original: 26% of the value of sigma.
Revised: 26% of the value of sigma.

Original: We further approximate the outer expectation with a single batch instead of the entire dataset.
Revised: We further approximate the outer expectation with a single batch instead of the entire dataset.

Original: On SVHN, the standard error of this approximation is 2% of the value of sigma.
Revised: On SVHN, the standard error of this approximation is 2% of the value of sigma.

Original: We see that both approximations are accurate in practice.
Revised: We see that both approximations are accurate in practice.

Original: 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.
Revised: 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.

Original: However, this bias can be corrected by using a different batch, or with a running average of the variance with an appropriate decay.
Revised: However, this bias can be corrected by using a different batch, or with a running average of the variance with an appropriate decay.

Original: This running average can also be used to reduce the variance of the estimate and to achieve convergence guarantees, but we did not ﬁnd it necessary in our experiments.
Revised: This running average can also be used to reduce the variance of the estimate and to achieve convergence guarantees, but we did not ﬁnd it necessary in our experiments.

Original: D A LTERNATIVE D ECODER C HOICES
Revised: D A LTERNATIVE D ECODER C HOICES

Original: have scalar performs well on discretized ELBO metric performing similarly to a discrete distribution parametrized as a discretized Gaussian or discretized Logistic.
Revised: or discretized Logistic.

Original: Full categorical distribution attains highest likelihood due to having the most statistical power.
Revised: Full categorical distribution attains highest likelihood due to having the most statistical power.

Original: [Table]
Revised: [Table]

Original:  We describe the alternative decoders evaluated in Table 2: using the bitwise-categorical, and the logistic mixture distributions.
Revised: [Table] We describe the alternative decoders evaluated in Table 2: using the bitwise-categorical, and the logistic mixture distributions.

Original: Bitwise-categorical
Revised: Bitwise-categorical

Original: VAE While the 256-way categorical decoder described in Section 3.2 is very powerful due to the ability to specify any possible intensity distribution, it suffers from high computational and memory requirements.
Revised: VAE While the 256-way categorical decoder described in Section 3.2 is very powerful due to the ability to specify any possible intensity distribution, it suffers from high computational and memory requirements.

Original: Because 256 values need to be kept for each pixel and channel, simply keeping this distribution in memory for one 3-channel 1024 × 1024 image would require 3 GiB of memory, compared to 0.012 GiB for the Gaussian decoder.
Revised: Because 256 values need to be kept for each pixel and channel, simply keeping this distribution in memory for one 3-channel 1024 × 1024 image would require 3 GiB of memory, compared to 0.012 GiB for the Gaussian decoder.

Original: Therefore, training deep neural networks with this full categorical distribution is impractical for high-resolution images or videos.
Revised: Therefore, training deep neural networks with this full categorical distribution is impractical for high-resolution images or videos.

Original: The bitwise-categorical VAE improves the memory complexity by deﬁning the distribution over 256 values in a more compact way.
Revised: The bitwise-categorical VAE improves the memory complexity by deﬁning the distribution over 256 values in a more compact way.

Original: Speciﬁcally, it deﬁnes a binary distribution over each bit in the pixel intensity value, requiring 8 values in total, one for each bit.
Revised: Speciﬁcally, it deﬁnes a binary distribution over each bit in the pixel intensity value, requiring 8 values in total, one for each bit.

Original: This distribution can be thought of as a classiﬁer that predicts the value of each bit in the image separately.
Revised: This distribution can be thought of as a classiﬁer that predicts the value of each bit in the image separately.

Original: In our implementation of the bitwise-categorical likelihood, we convert the image channels to binary format and use the standard binary cross-entropy loss (which reduces to binary log-likelihood since all bits in the image are deterministically either zero or one).
Revised: In our implementation of the bitwise-categorical likelihood, we convert the image channels to binary format and use the standard binary cross-entropy loss (which reduces to binary log-likelihood since all bits in the image are deterministically either zero or one).

Original: While in our experiments the bitwise-categorical distribution did not outperform other choices, it often performs on par with our proposed method.
Revised: While in our experiments the bitwise-categorical distribution did not outperform other choices, it often performs on par with our proposed method.

Original: We expect this distribution to be useful due to its generality as it is able to represent values stored in any digital formatby converting them into binary.
Revised: We expect this distribution to be useful due to its generality as it is able to represent values stored in any digital format

Original: Logistic mixture VAE For this decoder, we adapt the discretized logistic mixture from Salimans et al.
Revised: Logistic mixture VAE For this decoder, we adapt the discretized logistic mixture from Salimans et al.

Original: To deﬁne a discrete 256-way distribution, it divides the corresponding continuous distribution into 256 bins, where the probability mass is deﬁned as the integral of the PDF over the corresponding bin.
Revised: To deﬁne a discrete 256-way distribution, it divides the corresponding continuous distribution into 256 bins, where the probability mass is deﬁned as the integral of the PDF over the corresponding bin.

Original: (Kingma et al., 2016) uses the logistic distribution discretized in this manner for the decoder.
Revised: (Kingma et al., 2016) uses the logistic distribution discretized in this manner for the decoder.

Original: Salimans et al.
Revised: Salimans et al.

Original: (2017) suggests to make all bins except the ﬁrst and the last be of equal size, whereas the ﬁrst and the last bin include, respectively, the intervals ( −∞ , 0] and [1 , ∞ ) .
Revised: (2017) suggests to make all bins except the ﬁrst and the last be of equal size, whereas the ﬁrst and the last bin include, respectively, the intervals ( −∞ , 0] and [1 , ∞ ) .

Original: Salimans et al.
Revised: Salimans et al.

Original: (2017) further suggests using a mixture of discretized logistics for improved capacity.
Revised: (2017) further suggests using a mixture of discretized logistics for improved capacity.

Original: Our implementation largely follows the one in Salimans et al.
Revised: Our implementation largely follows the one in Salimans et al.

Original: (2017), however, we note that the original implementation is not suitable for learning latent variable models, as it generates the channels autoregressively.
Revised: (2017), however, we note that the original implementation is not suitable for learning latent variable models, as it generates the channels autoregressively.

Original: This will cause the latent variable to lose color information since it can be represented by the autoregressive decoder.
Revised: This will cause the latent variable to lose color information since it can be represented by the autoregressive decoder.

Original: We therefore adapt the mixture of discretized logistics to the pure latent variable setup by removing the mean-adjusting coefﬁcients from (Salimans et al., 2017).
Revised: We therefore adapt the mixture of discretized logistics to the pure latent variable setup by removing the mean-adjusting coefﬁcients from (Salimans et al., 2017).

Original: In our experiments, the logistic mixture outperformed other discrete distributions.
Revised: In our experiments, the logistic mixture outperformed other discrete distributions.