Doubly Robust Conditional VAE via Decoder Calibration: An Implicit KL Annealing Approach
Authors: Chuanhui Liu
Xiao Wang
Venue: Accepted by TMLR ()

Abstract: Several variants of Variational Autoencoders have been developed to address inherent limitations. Specifically, $\sigma$-VAE utilizes a scaled identity matrix $\sigma^2 I$ in the decoder variance, while $\beta$-VAE introduces a hyperparameter $\beta$ to reweight the negative ELBO loss. However, a unified theoretical and practical understanding of model optimality remains unclear. For example, existing learning theories on the global optimality of VAE provide limited insight into their empirical success. Previous work showed the mathematical equivalence between the variance scalar $\sigma^2$ and the hyperparameter $\beta$ in shaping the loss landscape. While $\beta$-annealing  is widely used, how to implement $\sigma$-annealing is still unclear. This paper presents a comprehensive analysis of $\sigma$-CVAE, highlighting its enhanced expressiveness in parameterizing conditional densities while addressing the associated estimation challenges arising from suboptimal variational inference. In particular, we propose Calibrated Robust $\sigma$-CVAE, a doubly robust algorithm that facilitates accurate estimation of $\sigma$ while effectively preventing the posterior collapse of $\phi$. Our approach, leveraging functional neural decomposition and KL annealing techniques, provides a unified framework to understand both $\sigma$-VAE and $\beta$-VAE regarding parameter optimality and training dynamics.  Experimental results on synthetic and real-world datasets demonstrate the superior performance of our method across various conditional density estimation tasks, highlighting its significance for accurate and reliable probabilistic modeling.

Published in Transactions on Machine Learning Research (01/2025)
Doubly Robust Conditional VAE via Decoder Calibration: An Implicit KL Annealing Approach
Chuanhui Liu
liu25068purdue.cem edu
Deportment of Statistics, Purduc University
Xiao Wang
wangriaaSpuurdus.e edu
Department of Statistics, Purduc University
Reviewed on OpenReview: htips: // openreview. net/ formmo id- VikkcTWDDo
Abstract
Several variants of Variational Auloencoders have been developed to address inherent lim- itations. Specifically, a- V-E utilizes a scaled identity matrix d²1 in the decoder variance, while 3-VAE introducos hyperparameter B to reweight the negative ELBO loss. How- ever, unified theoretical and practical understanding of model optimality remains unclear. For example, existing learning theories the global optimality of VAE provide limited insight into their empirical success. Previous work showed the mathematical equivalence between the variance scalar o² and the hyperparameter 3 in shaping the loss landscape. While ß-annealing is widely used, how to implement o-annealing stiill unclear. This paper presents a comprehensive analysis of o-CVAE, highlighting its enhanced expressiveness in parameterizing conditional densities while addressing the associated estimation challenges arising from suboptimal variational inference. In particular, propose Calibrated Robust o-CVAE, a doubly robust algorithm that facilitates accurate estimation of while effectively preventing the posterior collapse of Φ. Our approach, leveraging functional neural decom- position and KL annealing techniques, provides a unified framework to understand both o-VAE and 3-VAE regarding parameter optimality and training dynamics. Experimental results on synthetic and real-world datascts demonstrate the superior performance of our method across various conditional density estimation tasks, highlighting its significance for accurate and reliable probabilistic modeling. The implementation is publicly available at _cvee
1 Introduction
Conditional distributions play an essential role in characterizing the complex dependence the data y E R on given labels æ E R.P Canonical methods, such as regression or density estimators, face challenges when the data-generating distribution Pa (y)lr) high-dimensional. Deep latent generative models based on amortized variational inference are widely used scalable approach to model complex distributions and scale to large datasets. Specifically, the Conditional VAF (CVAF) was introduced for conditional learning (Sohn et al.
2015), based on Variational Autoencoders (VAE) Kingma & Welling 2014 Rezende et al. 2014p.
To model the conditional distribution (yr(y)a x), Gaussian o-CVAE incorporates Gaussian decoder p (y|æ, = (pa(( o),o²L,), where parameterizes the decoder mean ( z), and o² parameter- izes the decoder variance [Kingma et al. 2016 Dai & Wipf (2018). In detail, the latent variable E Rd is sampled from dala-independent prior p p(z(0) N(O, Ia) (Doersch) 2021 Thus, the parametric form of o (2)2z) is
p9. (y|x) = NNyluo( x), (z) 1a)dz.
(1)

---

Published in Transactions on Machine Learning Research (01/2025)
Gaussian o-CVAE also includes Gaussian encoder, (0(2)y,)) x)) (N x), Eg(y, such that the log-likelihood po.o (y)a()))))) is replaced by a tractable Evidence Lower Bound (ELBO), i.e,
ELBO a, o) = Eg [logge pe,o (y/ x x)] - KL (lyy.)))) {)][[()]]
(2)
With both the decoder and the encoder being Gaussian distributions, the sampling procedurr of (s y, x) and the computation of ELBO are simplified, allowing the model to be scalable to large datasets via the reparamelerization trick (Kingma & Welling 2014 Rezende et al. 2014). In particular, the label influences
the mean of the decoder (0) 2), assuming a data-indepenta independent prior p(z) without loss of generality (Zheng et al. 2022).
The objective of o-CVAE is to minimize the expected negative ELBO on the model parameters {0, a, of with respect to the (empirical) data distribution D. Ideally, the expected negative ELBO should reach the conditional data entropy H1 [X)), representing the theoretical lower bound for the model's performance.
Ep Ep|-ELBO(f, a, +) 2 ED( [ (olpoo 2 Ep[- log x=) == H(IIX)
(3)
The inequalities in Eq provide the conditions under which 0, a, é are optimal in the separate objectives of o-CVAE generative modeling and variational inference. The first gap in Eq is referred to as inference gup (Cremer et al., 2018). which is defined as EpKL[qs (gs (z) æ. p)llpo (z|xx, y)]. It determines the optimality of ø given any 0,0 in terms of the quality of variational inference of po.o (z|a,, y). The second inequality defines the approcimation gap, which equals Epllog Par (1)lze log pa, (g) v Eg KL (yl=)))llp,d [1]]=]] Notably, the optimality of 0, determines how closely the generative model Pa.- (y(x) of a-CVAE approximates the data-generating density Por(y/<)) ()) in probabilistic modeling.
Given the same parameter 0, o, the negative ELBO(0, a, Φ) loss is connected to 3-VAE (Higgins et al. 2016 To be specific, it is proportional to the objective function of B-CVAF, assuming a fixed unit. variance (Lucas
ltt al. 2019b Rybkin et al.] 2021), i..,
Eyg [1lmo (2, 2) / + [T(glllp(2) x (1100 2) - -y13] (KL((g)lp(z2))
(4)
Therefore, the optimal / is believed to be the maximum likelihood estimate of a (Lucas et al. 2019a Rybkin et al. 2021). Note that the parameter o² and the hyperparameter 3 are distinct and not mutually exclusic exclusive.
This paper presents a comprehensive analysis of o-CVAE, with particular attention given to o. The organi- zation and contributions of this paper as follows.
(1) A less constrained approximation theory without assumptions for optimal variational in- ference. In Section 3 we analyze the approximation gap of o-CVAE for a non-fixed a and establish a novel approximation theory in Theorem
3.4, leveraging Gaussian noise outsourcing (Agrawal & Domko
2021 and
block neural decomposition (Sobol 2001). Our characterization of oracle decoders identifies and relaxes re-
strictive assumptions on optimal parameters 0, a, including o - (Dai & Wip 2018 Zheng et al.
2022 or 0
constrained to the linear VAE regime (Lucas et al. 2019b Wang & Ziyin 2022 Dang et al. 2023
Relaxing
these assumptions, however, may result in the absence of an oracle encoder within the amortized Gaussian lamily, emphasizing how suboptimal variational inference can be imaeed in existing learning theories.
(2) The duality between and a leads to a unified annealing. In Section we analyze the dual relationship between the parameter a and ó to enhance the optimization outcomes. In Section 4.1 we high- light. the importance of & in the training dynamics of a, where & controls the bias the gradient(optimumn of the ELBO w.r.t. a. On the other hand, in Section 4.2 leveraging the loss equivalence of B in 3-VAE Higgins et al. (2016), we demonstrate a provides a smoothing effect in preventing the posterior collapse of d. In particular, sequence of fixed calibrations on the decoder variance a. Leveraging this dual relationship between & and a, we argue that an accurate estimation of a requires a better 46, and an explicit KL annealing scheme could be redundant. In Section 4.3 we motivate and propose Calibrated Robust o-CVAE, a simple o-calibration algorithm.
show that KL (8)-annealing [Bowman et al.) 2016 Fu et al., 2019] is essentially a
we empirically validate that sub- optimal variational inference can be the main source of numerical instability in the estimation of More
(3) Calibrated Robust o-CVAE is doubly robust. In Section 5.1
2

---

Published in Transactions on Machine Learning Research (01/2025)
importantly, we confirm the double robustness of our algorithm, showing that it not only provides accu- rate estimates of the variance parameter a, but also prevents posterior collapse of qa- In other words, it can efficiently explore the loss landscape of 0, ø to prevent posterior collapse and provide robust variance estimation. Compared to existing KL anncaling methods, we validate the superiority and effectiveness of a calibration. Starting from Section 5.2 we compared the performance of Calibrated Robust o-CVAE in various conditional density estimation tasks, showing its superiority over other conditional learning methods.
2 Related Works
We consider Gaussian decoders extending from a o-VAE, where the variance is a scaled identity matrix Σ = و2 Although many recent applications assume fixed variance 0²_ = 1 Castrejon et al. 2019 Yamasaki et al. 2023 Loizillon et al. 2024). critics have highlighted that this assumption leads to limited generative power and a positive approximation gap. If a is not fixed, one line of research establishes the tightness of the approximation gap in the asymptotic regime o - 0 (Dai & Wipf 2018 Zheng al., 2022). The assumptions regarding o - 0 are problematic. Mattei & Frellsen (2018) shows that this can lead to unbounded likelihood and demonstrates a zero approximation gap between Gaussian VAE and a sub-model of nonparametric mixture models (Lindsay 1995). However, their results do not extend to conditional variants. Another independent line of rescarch shows that the generative model under linear VAE assumptions corresponds to probabilistic Principal Component Analysis (pPCA) models (Lucas et al.) 2019b Wang & Ziyin 2022), which also demonstrates a zero approximation gap. Linear VAE theories are rigorous but do not fully account for the success of VAE. Compared to existing learning thoories of the global optimality of VAE models, our approximation theory focuses on analyzing the optimality of decoders in recovering the data-generating distribution without. the above-mentioned assumptions.
Expressive decoder and approximation gap.
Robust decoder variance estimation. Flexible covariance parametrization in the Gaussian decoder faces several practical challenges. The common challenge is the likelihood blow-up problem, where the maximum
likelihood estimation is ill-defind as the likelihood function (y [z) becomes infinite or unbounded
when the covariance approaches zero. For instance, Mattei & Frellsen (2018) showed that the multilayer
perceptrons (MLP) parameterization of the decoder (Knggma & Welling 2014 results in unbounded likeli-
hood functions for continuous responses. To mitigate unbounded likelihood, Takahashi al. (2018); Skafte
ct al. 2019); Stirn& & Knowles (2020) assumes a conjugate gamma prior distribution for the inverse variance 1/a². It transforms the decoder po(y|æ, (z, z) into Student-t distribution. Another challenge is model non- identifiability, where different parameters of the model can yield the same likelihood or ELBO, leading to multiple (local) optima (Khemakhem et al. 2020p. In addition, a dual-step optimization (Detlefsen et al., 2019 Rybkin et al. 2021) is preferred over a joint optimization of B, é and the covariance of the decoder Dai & Wipl 2018 Takahashi el al. 2018) to improve numerical stability. We follow the standard set up of the Gaussian decoder variance, using a scaled identity matrix, and follow joint optimization (Rybkin et al., 2021). However, we still identify problems in the estimation of decoder variance. This paper aims to obtain better estimate of decoder variance due to suboptimal variational inference.
Posterior collapse and KL annealing. Posterior collapse is a prevalent issue that indicates suboptimal variational inference of VAE, where the learned encoders have a near-zero KL divergence from the prior, often due to inadequate approximation of the true posterior distribution (He et al.) 2019 Lucas et al. 2019b) or, less frequently, from perfect alignment with the posterior in a degraded latent variable model Wang et al. (2021). To address this, traditional KL annealing methods reweight the loss function using a tuned hyperparameter 3 (Higgins et al. 2016 Chen et al., 2018 Rezende & Viola 2018), or predefined monotonic or cyclical annealing schedule that involves cross-validation and intensive computation (Raiko ct al. 2007 Bowman et al. 2016 Fu et al.] 2019). The empirical success of these remedics can be explained by the smoothing effect of o² on the loss function (Dai el. al. 2021) or the information bottleneck trade-off within the KL divergence/ratt term (Alemi et al. 2018 Bozkurt et al. 2021). In this paper, we provide a unified approach that integrates anncaling tochniques with the training dynamics of the docoder variance o² Specifically, we propose an implicit a calibration approach that dynamically adjusts based on clear evidence of posterior collapse, leading to improved encoder estimations and reduced computational costs.
3

---

Published in Transactions on Machine Learning Research (01/2025)
3 How expressive o- o-CVAE can be?
A strictly positive approximation gap for all possible 0,a demonstrates the inability of the CVAE genera- tive distribution ps,a po,a(yje a)) to exactly match the dala-generating distribution (y(()) This paper aims to understand why o-CVAE powerful enough to parameterize a wide range of tata-generating distributions.
3.1 data- generating distribution and VAE generative model
We begin our analysis by defining data distributions Pye(y) as a measurable function.
AaabourriiggAggawrrr & Domke 2021 Suppose a and y are random vectors taking values in the standard Borel space X' and v, respectivcly. Let μ.x denote the probability measure on X, and let Pæ.y" :X x y ~+ IK be probobility kernel of interest. Then for any m > 1, and a standard Gaussian random vector z E R. with measure jzz that is independent of x, there erists Borel measurable function G: Xx R +. + such that
Lemma 3.1 (Gaussian noise
(x. = G(, 2)) a.s.
(5)
Lemma 3.1 states that the variability of any random variable y can be outsourced to a Gaussian random vector z that is independent of æ. When æ is deterministic or given, such G is the nested function of the quantile function of y given æ and a Gaussian cumulative distribution function. Note that the dimension m is not necessarily the same as the latent dimension d of CVAE. Hence, for any latent m E N+, there exists Borel measurable map {G X x R + y} corresponding to the data-generating measure Pgt-
Based on 1 we can characterize the Gaussian o-CVAE model of y given
form of G, i.c.,
y G (#1, x2) = [a( 21) + 022.
(6)
where o > O. The decomposition in Eq highlights the inherent restrictions of noise outsourcing in the o-CVAE model. These restrictions include two key aspects: 1) the enforced independence between the d- dimensional latent variable #1 and g-dimensional decoder's unscaled noise #2, restricting their interactions; 2) an additive relationship between the mcan and sealed variance, which limits complex interactions and constrains the flexibility of dala-dependent variances.
Lemma 3.1 reformulates the approximation gap comparison of crriaale-tranffommiie functions G, elim- inating the need to compute log- probabbbiyyltyy and KL divergence. Such comparisons offer new perspective on existing approximation theories. For instance, by extending and simplifying the findings of Dai & Wipf (2018) in Appendix A.2, we show that an asymptotic assumption on a - 0 eliminates the two aspects men- tioned above, enabling the recovery of (x) of Pyr (y)(a) through (a 21) only. For another example, when the Linear VAE assumes μı as a linear combination of æ and z, Ge is limited linear function.
3.2 Non-asymptotic approximation of the o-CVAE
In this section, we analyze the expressiveness of the o-CVAE model, leveraging the above restrictions in Eq Specifically, we use a less explored yet intuitive approach
in recovering the data-generating map in Eq
known as block neural decomposition (Sobol 2001 Märtens & Yau 2020).
Definition 3.2 (Block neural decomposition (Sobol 2001)). Suppose that for any response dimension g, there exists > 4, such that G X' x RTT y almost surely equal to the fn which can be arbitrarily complex. Then a block decomposition of G on imput dimension X' x 1p--q x R2 is
fn 21, 22) =fo + fz((() - f2 (z1) + (z2) + füz, (x, + f (x. 22)
21 (z1, x2) + faz (x., 21, 22)-
To be specific, this decomposition is a blockwisc version of mcasurable functions (Sobol 2001 making the non- parametaictic mixture model in Märtens & Yau (2020) a particular case. Below, we consider unique decompasition with additional constraints G fn (x. 21, 22)-
4

---

Published in Transactions on Machine Learning Research (01/2025)
Assumption 3.3 (Block ANOVA representation (Sobol 2001 ). Let Mz: P211 Mzz denote the Borel measure of x 21; 22, respectively, and let lr be the measure or product measure of the index subset of {æ, #1, 22]. The block ANOVA representation assumes that the functional f in Eq 7 satisfies the following integral constraints, f fr( = O, Vi E I for every index subset, I _ { 21, 22)
So far, we have obtained unique decomposition of the data-generating function G as a sum of orthogonal functional bases, which is directly related to the oracle o-CVAE in the form of G (Zz,Z2)) By establishing the equality conditions, we demonstrate below the expressiveness of Gaussian decoders, e.g-- how they can have a zcro approximation gap to an arbitrary data-generating map Gma by characterizing the oracle decoders with parameters 0" o0'
Theorem 3.4 (Non-asymptotic approximation of o-CVAE). Under Assumptions 3.3, for some m > 9, if the block neural decomposition of the ground-truth map G satisfies the following conditions almost everywhere up to zero measure of Z2: 1) [zz((2)) = =00 22 and 2) 12 (21,22)+ (x., z2) fazzza (x, 21, 22) = Co for some o² > O, and conslant Co, then there erists a a -VVA model with the decoder porameterized by (()), and decoder variance scalar such that
G (21, 22)) = Gmm(222) x) 0.8.
(8)
The proof is deferred to Appendix C The approximation theorem 3.4 states that there exists a Gaussian decoder in o-CVAE that well approximates any generating map G under proper assumptions.
We avoid relying on restrictive assumptions on optimal decoder parameters 6", o² in Theorem 3.4 compared to existing learning theories. For example, (Dai & Wipl
2018 Zheng et al. 2022 assumes asymptotic
condition σ + 0 for reaching xero approximation gap. Another line of studies (Lucas et al.. 2019b Dai
et al., 2020 Sicks et al. 2021 Wang & Ziyin 2022 Dang et al. 2023) ) assumes that the G is highly constrained and adheres to the pPCA model, which can be parameterized by a simple linear VAE. Here, we argue that these assumptions are primarily employed to guarantee optimal variational inference. Specifically, pPCA model in linear VAËs has a closed-form conjugate normal posterior (Lucas et al., 2019b Eq.6, page 4) and taking o + 0 has an asymptotic effect in reducing the inference gaps (Dai & Wipf 2018 Eq.30, page 32). In addition, block neural decomposition of G also naturally avoids the assumption that G can be approximated by finite Gaussian mixture, as used in Mattei & Frellsen (2018).
Theorem 3.4 characterizes the ideal parameters of decoders learned from optimizing ELBO, which are inde- pendent of the value of Φ. To this extent, it reflects and aligns with the practical limitations of qa such as the amortization of variational parameters and restrictions on the variational distribution family (Cremer) et al. 2018). Thus, we argue that Theorem 3.4 better explains the empirical success and expressiveness of VAE models. As evident in Section 5.1 the two-moon dataset is recovered by the o-CVAE model without 0, and it is not instance of the pPCA model. See Appendix D for quick comparison.
Nonetheless, Theorem 3.4 neither guarantees estimation accuracy nor offers guidance on improved training strategies for optimizing ELB.O Namely, if the inference gap varies with 0,a, the local optima 0, à shall balance the approximation gap and the inference gap, which may, however, lead to substantial deviations from the oracle parameters 0" Thus, developing practical approaches to improve the estimation of 6, à is crucial, particularly in cases of suboptimal variational inference, including posterior collapse.
4 What if variational inference is not optimal?
The rest of this work addresses the challenges in optimizing the ELBO objective posed by suboptimal variational inference. In particular, we provide a fine-grained analysis of the duality between the dynamics of a and o uncovering how they may contributo to the identified issues. This dual relationship the primary motivation for a unified annealing technique, which eliminates poor estimates 0, ô to enhance the CVAE's generative performance in general.
5

---

Published in Transactions on Machine Learning Research (01/2025)
4.1 Biased estimation: Suboptimal inference induced saddle points
Traditional wisdom in likelihood-based models has been optimizing the loss by updating a jointly with {0, }} (Dai et al.) 2021 Dai & Wipf, 2018) or by updating a iteratively (Rybkin et al. (021) to avoid unbounded likelihood or numoric instability at small a, where the error of 0 is magnified, and the gradient of a or p
explodes (Mattei & Frellsen 2018 Takahashi et al. 2018).
In this section, we analyze the training dynamics of o, showing how they could lead a significant source of numerical instability (saddle points) in optimizing the ELBO in different phenomenon.
Recalling Fisher's identity, the gradient of T(00, Φ. o) with respect to a is a biased approzimation of the true gradient of the negative log-likelihood function with respect to a, where qs controls the bias.
VoL(0,0,0) EDBgo 1 Vo log pa,o (z)le, x)] o p.. (=lu.,) log pe,o [y],,,)]]
(9)
Similarly, the estimate or given {0,, } in the two-step iterative method in Rybkin et al. 2021 by maximizing ELBO analytically, is
obtained
min of = arg min L(e,, oz, o) EDEEs, [yy ue, (2, [x)] * EDEpo (=lu..) [(y uo, (x, =]
(10)
Here, highlight that the suboptimality of qs is driven more by the effect of amortized inference than by the constraints of the distribution family. Furthermore, ls more easily recovers the posterior mean than approximates the posterior variance (Cremer et al. 2018). Thus, do is less peaked and less concentrated than the posterior, and the estimated a is often larger than its actual value (Lin et al. 2019). In Appendix we also prove a toy example of inflated estimates of a via Fisher's identity.
As shown in Fig.1 of [ucas et (2019b), we argue that an inflated a may lead to increased saddle points in the loss landscape, contrary to the ideal scenario where optimal variational inference is achievable, seen in the asymptotic analysis (Dai & Wipf, 2018 Zheng et al. 2022) or in a linear VAE setting (Lucas et al, 2019b; Dai et al. 2020 Wang & Ziyin 2022). Such saddle points increase the sensitivity of the optimization process to the initialization of a, as shown in Lucas el al. (2019a). Thus, additional elforts are needed for numerical stability and accurate estimation of a, especially if the decoder covariance matrix flexible.
4.2 Posterior collapse: a compelling evidence for suboptimal encoders
In this section, we demonstrate the smoothing effect of a in preventing certain de, highlighting the opposite direction influence. We argue that actions should also be taken if there is clear and compelling evidence of a poor encoder qg. Specifically, we consider a kell-nnown example of the evidence, termed posterior collapse (Lucas et al., 2019a He et. al. 2019; Razavi al., 2019) or KT. vamishing (Bowman al., 2016 Fu et al. 2019h. It refers to the problem that the approximate posterior (o(2)y, x) collapses to the standard Gaussian [(, La), resulting in vanishing KL divergence in the training loss function.
E (= x O.
(11)
The encoder ds approximating the true posterior will hardly become the prior. This is because under Bayes' rule, pe,o (z x) x p(z)) (3/l2, z) and pe (3)
(3)22, x., 2) will not be independent of z. In our setting, the decoder's mean po(x 2) will not be independent of z. If the VAE generative model ignores latent variable, it collapses to nonlinear regression with mean network pa(x), significantly reducing the model's expressive power and the complexity of the noise structure. As a result, posterior collapse often coincides with poor ELBO in complex datasets due to a lack of fitness.
Traditionally, the standard solution has been KL annealing (Raiko et al. 2007 Bowman et al. (016)) Fu
et al. (2019)) leveraging V-VAE. We refer to it as explicit A-anncaling. The idea. that the loss landscape w.r.t. 0, Φ highly non-convex when a is fixed. Thus, tuning the hyperparameter B, which reweights the KL divergence term in the loss function, results in relatively large KL divergence in Eq and significantly lower training ELBO loss. For example, the schedule of B in Fu et al. (2019) can be
= 2r/M ((SF<<M/2)()) (r) + (1)
(12)
6

---

Published in Transactions on Machine Learning Research (01/2025)
Algorithm 1: general a calibration for o-CVAE
Input: data {æi, (il,y:li=n deterministic neural networks fu.: Js, with initialized parameter fo, Φn, Initialization do, inner number of iteration K of updating 0, Calibration tolerance C Output: 0,
Φ1 5- - aV end for
Sample 2j +- qd. (2 (z) for j ,.... N:
Update of + N Σi= l1y pe. (2,, z, 12////// 1//// if Convergence then
while is training do for i = to K do
if <x 2;) 2;)pp(() < C then
Calibrate oz f- oz as Eq 15
Read batch {æ دزري from data
end if clsc Break end if end while
Sample 2j +- qá, (z Uj. mz)) for j = 1,.... B....
Compute batch loss L(0,, dr. Ot)
Compute gradient VoL; VoL
0, 1 02 oVol
where r = mod(l, M) is the remainder of time I divided by the length of the annealing cycle M. In other words, B3 is periodically annealed from 0 to 1 during the first half of a M-iteration cycle and remains at 1 throughout the rest of the cycle.
In this paper, we highlight the mathematical equivalence between a and 3 in Eq In probabilistic modeling terms, ß-anncaling can be viewed as soquence of pre-defined and fired calibration of scaled variance a. Taking Eq 12 as an example, the loss-equivalent decoder variance a, is annealed in the same r, i.e.
of = M/4r * 1 (05)<M/2)()) + 1/2 (1/22FrcM(())
(13)
Perhaps surprisingly, the equivalence revcals the subtlety of the hyperparameter B. From statistical perspective, when the decoder is taken from location-scale distribution family, scaling the KL divergence (KLD) in ELBO is nothing but scaling its unit diagonal variance.
Thus, several aspects of ß-anncaling are ad hoc. First, the assumption of a fixed a limits the expressive power of the marginal distribution po.r(yla) causing a potentially larger approximation gap. Secondly, the training dynamics of a can hinder the annealing effect in shaping the loss landscape due to such equivalence (See Appendix G.1 for detailed discussion). In addition, the 3-scaled objective lacks statistical interpretability if 3 # 1 since it can no longer be seen schedule is unclear in principle, let alone the tuning of its hyperparameter, e.g., how to choose the period M for a better model. From an optimization theory perspective, preventing posterior collapse can be easily achieved without introducing 3 because annealing 33 equivalent to a in changing the loss landscape of 0.0. This statement. remains valid even if is neither learned nor optimal. However, treating as a model paramcter is essential for achicving improved probabilistic modcling and capturing uncortainty more effectively.
approximate log-likelihood. Lastly, the insight of the annealing
To conclude, have demonstrated the whole picture of the duality of and d. Suboptimal qø leads to biased and inflated estimation of a, causing additional saddle points in optimizing the ELBO. On the other hand, the smoothing ellect of a alleviates poor Ga. thereby mitigating posterior collapse.
4.3 o-calibration: an implicit KL annealing with learned Ø
In this paper, we propose Calibrated Robust o-CVAE, simple framework incorporating calibrations of a using the block coordinate gradient descent method (Rybkin et al. 2021 We consider the optimization problem:
min Ep|-ELBO(0, o)] subject to: Ep|KL[ (=ly x) [[](=)]] C. 0,0.0
(14)
Calibrated Robust -CVA extends Rybkin et al. 2021 with two key enhancements and operates through the following three steps each itcration. A general pseudo-code can be found in Algorithm
i. Optimize 0, Perform K stochastic gradient steps for 0, Φ with fixed a, where K is the inner steps size.
7

---

Published in Transactions on Machine Learning Research (01/2025)
ii. Estimate a: Compute the best estimate of a analytically based on current 0 Rybkki et al.
2021)
iii. Calibrate o at convergence: At the convergence of ELBO, calibrate a if Ep [gL x) [pp(2)
[2]] <
C.
Specifically, if the KL divergence term in the loss (KLD, or equivalently Rate (Alemi et al. 2018 is less than the hyperparameter calibration lolerance C, the calibrated gcll is defined as:
اaعl := |1 pp x de:
(15)
where ip (O, S) is sampled from a normal distribution with variance S. The variance S starts 1 and increases by 1 after each calibration. An example of the calibration process of is shown in Figure 2 a). The constraint that KLD > C defines the feasible region of ø in the constrained optimization of ELBO. It excludes any local optima 0, in ELBO such that this constraint is not satisfied.
In recognition of the aforementioned issues, the merit of calibration is supported by the following rationales:
1) In the context of annealing, calibration an anncaling that is adaptive to the performance of o-CVAE, where the rate term a representative performance metric that is directly reflected in the optimization problem as hard const.raint (See the discussion below). Our calibration is similar to the sudden change of B at the end of cycle (Fu et al. 2019). The annealing effect in our algorithm is concentrated, and the stage of slow annealing in B-anncaling is replaced by the learning dynamics of minimizing a as a learned parameter. C is the minimal desired KLD after annealing. Note that the relat ionship between C' and hyperparameters in ß-annealing is not bijective (See Fig. 12 in Fu et al. (019))
2) To mitigate the numerical instability described in Soction 4.1 our proposed calibration follows the pertur- bation and random restarting strategies of to navigate saddle points. Specifically, & governs the calibration scaling adjustments. The adjustment strength (annealing temperature) is controlled by the variance S of p. Due to the above concentrated annealing effect, we may require larger calibrations. Its escalating feature of variance S is beneficial when larger calibration is necessary to escape the flat regions of the loss landscape.
3) Calibration is simple and straightforward and a predefined annealing schedule is no longer needed. It has minimal impact on model interpretability since our primary objective is the original ELBO loss and avoids the explicit usage of B. Due to the nature of the constraint and non-convexity, the learned 0,a, is highly likely to be a non- con-boundary solution.
calibration on a results from the local optima of ELBO,
Finally, we demonstrate the flexibility of our method and better understanding of the interpretation of KLD and the optimum values of the hyperparameter C. 1) To prevent posterior collapse, KLD is the direct metric, and the optimal C' is any value between 0 and the lowest KLD of qs that is not posteriorly collapsed. In practice, we define posterior collapse first by considering the permissible deviation of a zero KLD. Such permissible ion, which becomes the optimal C, should be slightly more than negligible and may be subjective. 2) To fine-tune the model for various tasks, the metric of KLD is the preference and/or strategy for selecting from nearly the best local optima Φ, particularly when multiple distinct local optima of the ELBO yield approximately the same minimal training loss. For example, in the rate-distortion trade- off fAlemi et al. 2018), C corresponds to the minimum value of the upper bound on mutual information, which can be tuned to improve latent disentanglement. For another example, in the approximation inffeennceeee trade-off in Eq C is chosen to prioritize one of the gaps for better generative or inference performance. Here, the optimal C is objectively unknown and must be carefully tuned for task-specific goals.
5 Experiments
In this section, we provide evidence and more practical insights into the calibration of o-CVAE through extensive numerical experiments given a finite training sample. We use the Adam [Kingmal 2014 stochas- tie gradient descent algorithm for training neural networks. The general learning rate is 0.005, and the convergence threshold is 0.001 in the average loss change. Other important details of each of the follow- ing experiments are attached in the Appendix F] In addition, we also validate our method on the MNIST (Deng 2012) and CelebA (Luu et al. 2018) datasets for conditional image generation and reconstruction in A Appendix] &
8

---

Published in Transactions on Machine Learning Research (01/202))
(b)
until convergence
3.77
0.13
0.38
0.38
0.75
31.99
367.02
1.2
Rybkin
K-5
1.0
Km25
Cop
Estimated Sigma
0.8
K-100
2500
0.6
(4
(d)
5
0.4
0.2
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
Normalized Epochs
-5 04
Figure 1: Trajectories of estimated under different values of K. Dynamic scaling on total epochs is applied. (c) reconstructed training Two-moon dataset; (b).(d) generated samples. Vanilla o-CVAE may not recover the true a, reconstruct training data, and learn the ground-truth distribution. The true is 0.1.
5.1 Two-moon dataset
Let æ E [-1,1} be the binary input and let y E I² be generated as follows,
(cos(a) + + sin - 8 + +2),
if x -1,
y (cos(a)) + c3. - sii((in())))))) + + +
if x
where latent Uniform(0, T| and 6, N (0, 1,2,3,4. Note that a; Ei,,,..,, are mutually independent. We simulated three datasets of size 5, 000 with 2, 500 for each class using 0.1, which is referred to as the true Note that this synthetic dataset does not follow the pPCA model described in Lucas et al. (2019a) and cannot. be recovered by a linear CVAF. A well-behaved encoder, in this case, is expected to yield a non-zero KL divergence to the prior. We have the following observations.
No calibration: a properly learned decoder variance is crucial for probabilistic modeling. Figure [ first explore the vanilla method that does not include the third step, including K suggested in Rybkin et al. (021)). We see that different inner steps K result in different generative models convergence, given the samne training dalaset. The contradictory result is well predicted by Theorem 3.4 although o-CVAE is not guaranteed to achieve optimal variance inference. P8.o in Eq 6 may well approximate the data-generating distribution in Eq 5 However, we also see that the vanilla a- o-CVAE under a small value of K, including K = 1, fails to recover the truth. This observation agrees with the claims that learning o can have unintended consequences (Tucas al., 2019bb Dai et al., 2021), highlighting practical difficulties parameter estimation. The observation that increasing the inner step size K improves the performance of the generative model aligns with our discussion in Section 4.1 where the bias of a can be reduced by obtaining more in informative éb While smaller K is generally preferred for computational efficiency, the appropriate K is generally unknown and differs from case to case. Finding the right balance between parameter estimation and computational efficiency remains a non-trivial challenge.
Evaluating the effect of calibration as third step. In Figure we evaluate the effect of calibration as the third step, compared to the vanilla method. Calibratod robust o-CVAE accurately estimates the decoder variance in Figure (aa). a). In Figure(2(b) it prevents posterior collapse and obtains better generative model with low negative ELBO. Compared with the vanilla CVAE that recovers the true using large K = 500, 2500, our algorithm using K = 25 requires significantly smaller wall time to converge. Table 0 in Appendix compares the wall time until convergence for different setups of K.
Compare with KL 8-annealing. In Table I we compare Monotonic (Bomman et al. 2016 and Cyclical annealing (uu al. 2019) with our method. We report the MSE of the learned and KLD at convergence, under the same setting as Figure averaging over 20 repetitions. Training loss is reported due to the inconsistency of the loss reweighted by B. We found that by using default small C 0.05, our algorithm
9

---

Published in Transactions on Machine Learning Research (01/2025)
Wall time until converging gevriiiggg.eo truth
31.99
367.02
3.54
Rybkin dt
ELBO
1.0
K-5
After sigma
0.8
K-
8-
Estimated
100
K-500
4
0.4
K=2500
2
Km25 wi Calbration
101
Calibrated estimates
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
100
750
12.00 Epochs
1500
1750
2000
Normalized Epochs
(a) o-calibration small sctting
(b) o-calibration is equivalent to KL anncaling
Figure 2: (a) The trajectory of the estimated by proposed calibration technique. Accurate estimate of is obtained by calibration with a small K (K-25 w./ Calibration). The calibrated are marked as (b) The KT. divergence in ELBO remains positive after several calibrations of a. The calibration tolerance
C 0.05. The true a is 0.1.
Table 1: Comparison of KLD and MSE of the learned across different anncaling techniques
Initialized
0.2
0.6
MSE
KLD
MSE
KLD
MSE
KLD
MSE
KLD
MSE
KLD
Monotonic
1.2e-002
1.81
6.09e-03
2.10
5.65e-03
2.19
3.41e-04
2.12
6.06e-04
2.17
Cyclical
1.76c-02
1.77
3.24e-02
1.71
1.926-02
1.84
2.17o-02
2.01
4.58c-02
1.45
Ours C=0.05
2.41e-02
1.79
2.71e-02
1.64
6.01e-03
1.92
5.38o-022
1.30
2. 2.660-02
1.70
Ours 1-96
1.26e-05
2.09
5.77e-05
2.09
1.22e-05
2.08
1.73e-05
2.10
1.76e-05
2.09
Vanilla
0.16
4.17e-01
0.20
1.40e-03
0.17
3.16e-01
0.20
2.20e-04
0.20
5.39e-04
Table 2: Comparison of wall time to convergence across different annealing techniques.
Inner step size K
25
50
100
Monotonic
24.85
36.26
58.38
Cyclical
10.76
23.01
45.59
Ours C=0.05
10.55
12.69
25.00
Vanilla
62.15
results a similar KLD and similar MSE; however, our algorithm can be 50 50%70 faster converge in Table [ defined as the maximum wall time until convergence. In addition, finc-tune algorithm by adjusting C' = 1.96, leading to significantly less biased estimation of the decoder. A detailed comparison of wall time savings and other experimental details is in the Appendix G
Small initialization strategy and ablation study. In Figure 3 we report the results of Calibrated Robust o-CVAE with various initializations of a ranging from 0.05 to 1. Figure 3 3(b), Calibrated Robust various metrics. In Figure a the vanilla CVAE ends up with very high training loss (negative ELBO), a large mean squared error (MSE) of the learned and vanishing KLD. Similar to Dai & Wipf (0018) Dang et al. (2023). we find that smaller initialization of has higher chance of avoiding posterior collapse with non-zero KLD. However, the strategy of smaller initializations is limited, and it may still fail to recover the ground-truth distribution. In Figure 3(b), Calibrated Robust o-CVAE demonstrates consistently superior performance in various metrics. the ablation study in Appendix H we consider various setups of C' from to 0.5, where C' = 0.05 is effective enough for preventing posterior collapse and providing an accurate estimate of a. Results of a fized-o CVAE under different initializations can also be found in Appendix H If the fixed happens recover the true 0.1, it yields the best result..
o-CVAE demonstrates consistently superior performance
10

---

Published in Transactions on Machine Learning Research (01/2025)
2.5
2.5
MSE sigma
EL30
KL
MSE of sigma
ELBO
KL
2.0
2.0
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
Initialization
Initialization
(a) without calibration
(b) with calibration
Figure 3: CVAE model using learned (a) without calibration and (b) without our proposed calibration on twomoon training datasets. It prevents vanishing KL divergence due to posterior collapse and significantly improves the decoder variance estimates and the ELBO at convergence. The true 0.1. The calibration tolerance C = 0.05. Inner step K = 25. Each point is averaged over 20 repetition experiments.
5.2 Comparison with nonparametric conditional density estimator
Next, compare the Wasserstein generative conditional sampling method (WGCS) (Liu al. 2021 along with three conventional nonparametric conditional density estimators, including the nearest-neighhbr kernel conditional density estimator (NNKCDF) (Dalmasso et al. 2020 the conditional kernel density estimator (CKDE) (Hall et al., (204)), and a basis expansion method (FlexCode) (Izbicki & Lee 2017) in the tasks of estimating summary statistics including the conditional mcan and conditional standard deviation. The metric for the generative performance of 0,a a the mean squared error (MSE) of the estimated mean (MEAN) and that of the estimated standard deviation (SD). A smaller MSE of the estimated SD means a more accurate estimate of variance. Three different simulated datasets are used, where M1 is a nonlinear model with additive noise; Ma is nonlinear model with multiplicative noise, and M3 is two-component Gaussian Mixture Model. The dimensions of label M₃ and M2 are 5. The dimension of label M3 is 2. The dimension of data Y is 1. Details can be found in Appendix[P F Our method produces consistent results and outperforms the other models, as shown in Table 3 The numbers in bold indicate the best results in the same row. The evaluation repeated 10 times.
Table 3: Comparison of summary statisties between non-paraet conditional density estimators
CVAE (OURS)
WGCS
NNKCDE
CKDE
FLEXCODE
MRAN
0.174 (0.004)
1.10(0.05)
2.49(0.01)
3.30(0.02)
2.30(0.01)
M1
SD
0.185(0.005)
0.21(0.01)
0.43(0.01)
0.59(0.01)
1.06(0.08)
MEAN
(.420(
3.71(0.23)
6.09(0.07)
66.76(2.06)
10.20(0.33)
M2
SD
2.071 (0.019)
3.52(0.17)
9.33().23)
18.87(0.59)
11.08(0.34)
MEAN
0.106(0.001)
0.32(0.03)
0.11(0.01)
1.55(0.03)
0.12(0.01)
M3
SD
0.421(0.001)
0.10(0.01)
0.36(0.01)
0.51(0.01)
0.33(0.01)
5.3 Comparison with Conditional Flow Based model in benchmark datasets
Finally, we compare the performance of calibrated o-CVAE with Bayesian conditional normalizing flows (Trippe & Turner 2018) on 6 UCI datasets, as listed Table ! The test mean log-likelihood in the natural unit of information (nats) evaluated with the associated variance reported. The models include Bayesian normalizing flows (NF), mixture density networks (MNM), and neural networks with latent inputs mean field Gaussian (MF), sampled from Hamiltonian Monte Carlo (HMC). A higher logarithmie likelihood
11

---

Published in Transactions on Machine Learning Research (01/2025)
Table 4: Comparison of mean log- bielliiel of test data nats between conditional models.
Dataset
boston
concrete
energy
power
wine-red
yacht
N
506
1030
768
9568
1599
308
p/g
13/1
8/1
8/22
4/1
11/1
6/1
MDN-2
-2.6 + 0.03
-3.23 + 0.03
- -1.60 + 0.04
-273 + 0.01
-0.91 + 0.04
-277
0.05
MDN-5
-2.73 0.04
-3.28 0.03
-1.63
0.06
-
+1.43
0.07
-2.54
MDN-20
- 0.03
-3.27 + 0.02
-1.48 + 0.04
-2.68
0.01
+1.21 + 0.06
-
IV-15
-2.44 0.05
-3.06 0.03
-0.74 0.03
+.01
-0.98
0.02
+
LV-5
- 0.05
-3.08 0.02
- 0.02
0.01
-0.96
-1.15
NF-2
-2.40
-3.03 0.05
-0.44 0.04
-
-0.87 = 0.02
-0.30 0.04
NF-5
-2.37 + 0.04
-2.97 + 0.03
-0.67 0.15
-.00
-0.76
-0.21 0.09
HMC
-2.72 0.03
-0.93 + 0.01
0.000
00.91
-1.62 0.02
Dropout
2.46
3.04 0.09
11.9 1
2.00
0.93
1.55 0.12
MF
2.62 4 0.06
3.00 + 0.03
0.57 0.04
0.04
2.79
0.97 0.01
-1.00 + 0.10
CVAF((ous)
-3.17 + 0.07
-3.48 + 0.04
0.10
-3.29
-1.48 +
-0.34 0.02
implies better generalization of the test set. The random train-test split is 75% to 25%. Fach experiment is repeated 5 times. Although o-CVAE optimizes the ELBO instead of the exact log- liellikel and does not have an optimal VI, our results are competitive. It shows superior performance on the Energy dataset, and near-optimal performance on the Energy dataset, and near optimal performance on the Yacht dataset.
6 Discussion
Limitation on approximation gap:
Our analysis and algorithm sctup are subject to several conventional assumptions. These assumptions can be relaxed in several aspects, including generalizing our analysis to accommodate heterogeneous, diagonal, or even full covariance matrix Σ, which offers better approximation properties for more complex datasets.
Although the KLD or Rate in Eq 11 is a direct measurement, the criteria for posterior collapse remain flexible. For example, Lucas et al. (2019b) define the dimension-wise pos- terior collapse by P(KL(go æ) [>][(= S <) > approximately correct manner. For another example, Dai et al. (2021) define implicit posterior collapse from observing large maximum mean discrepancy (Makhzani N EE= (= yi, v.) and p(z). This because marginalizing a well-behaved encoder would be similar to the prior: go (2ly. ())gg (dydar) R= (o.y(2lyy, (dyddr))))) Rs p(=).
Other calibration metrics:
for each latent dimension index in a probably
al., 2016) between the aggregated posterior
Hyperparameter selection: Calibrated Robust o-CVAE depends on two hyperparameters, namely C, K. While it eliminates the need for predefined annealing schedule (Raiko al. 2007 Bowman al. 2016 Fu et al.
2019), these hyperparameters are crucial and must be carefully tuned. We should generally use a small K and set 0 < C' < ()). We highlight that higher KLD does not necessarily mean a better or a lower inference gap, and the trade-off for the ncar-optimal solutions is task-specific and date-dependent. As a general rule of thumb, we recommend using small tolerance C and moderate K inner steps as a starting value for computational efficiency. Choosing an enormous tolerance C' will not only trigger more calibrations of but also lead to potential non-convergence of the algorithm.
Image generation tasks: If probabilistic modeling and parameter estimation not the top priorities of the task, it is expected to use a fixed a incorporate both / * and Icarned especially in image generation lasks, including crisp face generation (Vahdat & Kautz 2020) and abnormality detection (Loizilon) et al. 2024). Based on results from Appendix & o-calibration itself does not significantly improve the image generation results. The image generation process often deviates from the complete generative graph, as the decoder's mean is used directly as the generated sample. It neither aligns with the assumptions of Gaussian decoder nor accurately represents the underlying generative process.
12

---

Published in Transactions on Machine Learning Research (01/2025)
7 Conclusion
To the best of our knowledge, this work represents the first comprchensive investigation of the variance parameter a in VAE models. In this work, our methodology provides deeper insight into the approximation power of VAE models and the duality between a and Ø, which enables a novel pathway for o-annealing. We empirically demonstrate the effectiveness of our proposed Calibrated Robust o-CVAE addressing the challenges of accurate decoder variance estimation, numerical instabilitics due to suboptimal variational inference, and posterior collapse. Our future work will focus generalizing our approach to encompass broader and complex issues in conditional learning, particularly those involving unbalanced sampling and missing data.
13

---

Published in Transactions on Machine Learning Research (01/2025)
References
Abhinav Agrawal and Justin Domke. Amortized variational inference for simple hicrarchical models. Ad- vances in Newral Information. Processing Systems, 34:21388 21399, 2021.
Alexander Alemi, Ben Poole, lan Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing a broken elbo. In International conference on machine learning, pp. 159-168. PMLR, 2018
Samuel R. Bowman, Luke Vilnis, Oriol Vinyals, Andrew M. Dai, Rafal Józefowicz, and Samy Bengio. Gen- erating sentences from continuous space. In CoNLL, pp. 10 21. ACL, 2016. ISBN 978-1-945626-19-7.
Alican Bozkurt, Babak Esmaeili, Jean-Baptiste Tristan, Dana Brooks, Jennifer Dy, and Jan-Willen van de Meent. Rate-regularization and generalization in variational autoencoders. In Proceedings of The 24th In- lernational Conference on Artificial Intellignnce and Statistics, Proceedings of Machine Learning Research. PMLR, 13 15 Apr 2021.
Lluis Castrejon, Nicolas Ballas, and Aaron Courville. Improved conditional vrnns for vidco prediction. Proceedings of the IEEE/CVE internalional conference on compuler vision, pp. 7608 7617, 2019.
Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. Advances in Newral Information Processing Systems, 2018.
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders.
International Conference on Machine Learning, pp. 1078 1086. PMLR, 2018.
Bin Dai and David Wipf. Diagnosing and enhancing vae models. In International Conference on Learning Representations, 2018.
Bin Dai, Ziyu Wang, and David Wipf. The usual suspects? reassessing blame for vae posterior collapse. In International conference on machine learning, pp. 2313 -3222. PMLR, 2020.
Bin Dai, Li Wenliang, and David Wipf. On the value of infinite gradients in variational autoencoder models. Advances in Neural Information Processing Systems, 34:7180-7192, 2021.
Niccolò Dalmasso, Taylor Pospisil, Ann B Lee, Rafacl Izbicki, Peter E Freeman, and Alex I Malz. Conditional densily estimation tools in python and with applications to photometric redshifts and likelihood-freee cosmological inference. Astronomy and Computing, 30:100362, 2020.
Hien Dang, Tho Tran Hun, Tan Minh Nguyen, and Nhat Ho. Beyond vanilla variational autoencoders: Detecting posterior collapse in conditional and hierarchical variational autoencoders. The Twelfth International Conference on Learning Representations, 2023.
Li Deng. The mnist database of handwritten digit images for machine learning research [best of the web). IEEE Signal Processing Magazine, 29(6):141 (10::1411112,,111111111111111111 142, 2012. doi: 10.1009/MSP:201222211477
Nicki S. Detlefsen, Martin Jørgensen, and Søren Hauberg. Reliable training and estimation of variance networks, 2019.
Carl Docrsch. Tutorial on variational autoencoders, 2021.
Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin. Cyclical annealing schedule: A simple approach to mitigating kl vanishing. arXiv preprint orXiv:1903.101: 10145, 2019.
Peter Hall, Jeff Racine, and Qi Li. Cross-validation and the estimation of conditional probability densities. Journal of the American Statistical Association, 99(468):1015-1U 2004.
Junxian Hc, Danicl Spokoyny, Graham Neubig, and Taylor Berg- Korg-Kirkpatrick. Lagging inference networks and posterior collapse in variational autoencoders. In International Conference on Learning Representations, 2019.
14

---

Published in Transactions on Machine Learning Research (01/2025)
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Terchner. beta-vae: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2016.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(3):359- 366, 1989.
Xianxu Hou, Ke Sun, Linlin Shen, and Guoping Qiu. Improving variational autoencoder with deep feature consistent and generative adversarial training. Nearocompuling, 341:183-194 2019.
Rafael Izbicki and Ann Lee. Converting high-dimensional regression to high-dimensional conditional density estimation. Electronic Journal of Statistics, 11:2800 2831, 2017.
W JAMES. Estimation with quadratic loss. In Proc. 4th Berkley Symp. on Math. Statist, and Prob., 1961, 1961.
llees Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: A unifying framework. In International conference on artificial intelligence and statistics, pp. 2207-2217. PMLR, 2020.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6985 2014.
Diederik P. Kingma and Max Welling, Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, 2014.
Durk Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances in neural information processing systems, 29, 2016.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to docu- ment recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998.
Shuyu Lin, Stephen Roberts, Niki Trigoni, and Ronald Clark. Balancing reconstruction quality and regular- isation in elbo for vaes, 2019. URL httts://arxiv http::://axxvvvoo/e//e www.:oogggabao. abt////b//11//////// 03765
Bruce G Lindsay. Mixture models: theory, geometry, and applications. Ims, 1995.
Shiao Liu, Xingyu Zhou, Yuling Jiao, and Jian Huang. Wasserstein generative learning of conditional distribution, 2021.
Ziwci Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. Proceedings of International Conference on Computer Vision, 2015.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces altributes (celeba) dataset. Retrieved August, 15(2018):11, 2018
Sophie Loizillon, Yannick Jacob, Maire Aurélien, Didier Dormont, Olivier Colliot, Ninon Burgos, and AP- PRIMAGE Study Group. Detecting brain anomalies in clinical routine with the S\betaS-VAAL VAE: Feasibility study on age-related white matter hyperintensities. In Medical Imaging with Deep Learning, 2024.
James Lucas, George Tucker, Roger Grosse, and Mohammad Norouzi. Understanding posterior collapse in generative latent variable models, 2019a.
James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't blame the elbo! a linear vae perspective on posterior collapse. In Advances in Neural Information Processing Systems, volume 32 2019b.
Alireza Makhzani, Jonathon Shlens, Navdcep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoen- coders, 2016.
15

---

Published in Transactions on Machine Learning Research (01/2025)
Kaspar Märtens and Christopher Yau. Neural decomposition: Functional anova with variational autoen- coders. In International Conference on Artificial Intelligence and Stalistics, pp. 2917-2927. PMLR, 2020.
Pierre- Alexandre Mattei and Jes Frellsen. Leveraging the exact likelihood of deep latent variable models. Aduances in Newral Information Processing Systems, 31, 2018.
Tapani Raiko, Harri Valpola, Markus Harva, and Juha Karhunen. Building blocks for variational bayesian learning of latent variable models. Journal of Machine Learning Research, 8(1), 2007.
Ali Razavi, Aaron van den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with delta-vaes. In International Conference on Learning Representations, 2019.
Danilo Jimenez Rezende and Fabio Viola. Taming vaes. arXiv preprint arXiv:1810.00597, 2018
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic, backpropagation and approximate inference in deep generative models. In International conference on machine learning, pp. 1278 1286.
PMLR, 2014.
Olch Rybkin, Kostas Daniilidis, and Sergey Levine. Simple and effective vae training with calibrated decoders. Tn International Conference on Machine Learning, pp. 9179 9189. PMLR, 2021.
Robert Sicks, Ralf Korn, and Stefanie Schwaar. A generalised linear model framework for 8-variational autoencoders based exponential dispersion families. Journal of Machine Learning Research, 22(333): 1-41, 2021.
Nicki Skafte, Martin Jørgensen, and Soren Hauberg. Reliable training and estimation of variance networks. Advances in Newral Information Processing Systems, 32, 2019.
Nya M Sobol. Global sensitivity indices for nonlincar mathematical models and their monte carlo estimates. Mathematics and computers in simulation, 55(1-3):271 280, 2001.
IM Sobol. Scnsitivity estimates for nonlinear mathematical models. Math. Model. Comput. Exp., 1:407, 1993.
Kihyuk Sohm, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep condi- tional generative models. In Advonces in Neural Information Processing Systems, 2015.
Andrew Stirn and David A Knowles. Variational variance: Simple, reliable, calibrated heteroscedastic noise variance parameterization. arXiv preprint arXi::200600400 2020.
Hiroshi Takahashi, Tomoharu Iwata, Yuki Yamanaka, Masanori Yamada, and Satoshi Yagi. Student-t variational autocncoder for robust density cstimation. In IJCAL, pp. 2706-2702, 2018
Brian L Trippe and Richard E Turner. Conditional density estimation with bayesian normalising flows, 2018.
Arash Vahdat and Jan Kautz. NVAE: deep hierarchical variational autoencoder. In Neural Information Processing Systems (NeurIPS), 2020.
Yixin Wang, David Blei, and John P Cunningham. Posterior collapse and latent variable non-identifiabliity In Advances in Neural Information Processing Systems, 2021.
Zihao Wang and Liu Ziyin. Posterior collapse of lincar latent variable model. Advances in Neural Infor- mation Processing Systems, 35:37537 37548, 2022.
Yuki Yamasaki, Chiaki Doi, Shiori Kitagawa, Hiroyuki Seki, and Hiroshi Shigeno. Data generation with filtered 8-vae for the preoperative prediction of adverse events. IEEE Access, 11:48667-48676, 2023.
Yijia Zheng, Tong He, Yixuan Qiu, and David P Wipf. Learning manifold dimensions with conditional variational autoencoders. Advances in Newral Information Processing Systems, 35:34709 95:34709-34721, 2022.
16

---

Published in Transactions on Machine Learning Research (01/2025)
A Discussion on Theorem 3.4
The approximation thcories of this work are also motivated by explaining the expressiveness of the param- eterized density (o.o()lz) when a is not infinilesimal. Here, we show simple exlended result from Dai & Wipf 2018))
Assumption A.1 (Response continuity). Conditional on any X, there exists a ground truth proba- bility measure of Hgt such that its Radon-Nikodym derivative with respect to the standard Lebesgue mcasure is nonzero almost everywhere in V
Simply put, Pys (y æ) exists uniquely and non-zero almost everywhere up to a null set.
Proposition A.2 (Global optimal o-CVAE). Under Assumption A.1 if the latent dimension d is larger than or equal to the dimension of the response i, then for any a > 0, there erists a o-CVAE with encoder network parameterized by f,, and decoder network parometerized by de such that,
lim KL
= 0, lim KL[u,, (=lu, 1) (=lv, x²) = 0. a
(16)
Inspired by Dai & Wipf 2018), the proof argument involves two steps. First, we show the convergence of the parameterized density to the distribution of the ground truth measure. Then, show that the ELBO is asymptotically tight as a + 0 as the KT. divergence between the approximate posterior and the ground truth posterior vanishes. In particular, we do not assume that follows a low-dimensional simple Riemann manifold with dimension less than to avoid additional definitions such as active dimension (Zeng et. al. 2022 The detailed proof of Proposition A.2 is provided in Appendix
The proposition A.2 states that a-CVAE asymptotically approximates (gc(y)æ)) and the intractable postcrior (z y, 2) as + simultaneously. This proposition establishes the asymplotic expressiveness of the + O. The global optimality is evident from the fact that the lower
parameterized density (ylx) bound of Eq 3 is reached. However, this proposition is limited because it does not explain the practical success of experiments in which a nonzero a is learned. Indeed, Proposition A.2 cannot answer why a non- zero a is learned. is not clear whether this is due to the limited capacity of the encoder decoder network or the identifiability the VAE.
B Proof of proposition A.2
Following the line of proof of theorem 2 in Dai & Wipl 2018), first consider the case when the latent dimension d equals the response dimension
Step 1: Show that KL(po; (u x) 1)) + as o + 0.
Define the cumulative density function F R + [0, 1,s for given X as follows,
(y) = [Fi( F2(22: (u 1)
(17)
th - - Porlvila, U1,... Vi- ddyi
Vi..... Vi-1) F, Vy
(18)
where Pgc(si .r.... ((.- is the marginal distribution of the i-th dimension of Y|X, conditioned on the first dimensions.
By definition, we have (y) Pot (ylardyy And since Pyt is nonzero everywhere, the conditional distribu- tion function Fe is invertible. Denote its inverse by F-!
Similarly, we can define another differential and invertible function T: Rd + 1]
T(=) (0(z),, (
(19)
17

---

Published in Transactions on Machine Learning Research (01/2025)
where +() is the cumulative density function of the standard Gaussian.
Since d = 9, we consider a possibly non-identifiable decoder parameter 0 E (01m0(1,2)) F F=)))( then the corresponding density is
po(y) [ ((g]
= J 0.1)
(w),
(20)
7 /
Nlyve
5(y y ('))g(( = Psc())la),
RR
where the p.d.f Nyyy'', olg) - o y') when o +.
Thus, it follows immediately that (l. -
9eg (=ly.a)
Step Show V0, 3 s.t.
const.
Pag (=lv.a)
let weights of encoder network ø are such that
pa y) = T Fz( Ex( y) = (So, y) Sa y))
(21)
where the q Jacobian matrix S0, (2x, y) Vz (x. 2) ==)o())))
Denote pa(a, y) = pbb (x y) = =oEs0
Po P= x) (2 x pp(z)
z), omm
(22)
pa (y)
po p3(y)0) (
Let z' = (2---10)) (>==<o)// and dz(z'lu, m) pél 2' (,, a.) be the transformed pdf, then we have
(z(2'yy,)) x)
N(Voz' pai Hv. LE (2)()(
p/(2'1u y.
Niy: ua Voz' + ) o m W (Voz' + usi 0,
I|a /az
| uo (2,, ub vo2')l1
=C( y) exp(
2
2
20
1
1|
voz'113
|| uz)
pl(x us + voz'112
=C(x, y) exp(
+
2
2
20
(23)
1 z/
1| +
voz'l13
|| (o) /ol 0z 222
CCxx, y) exp(
+
(Taylor Expansion)
2
2
20
T 1/5-12/
1|4s voz'l13
z'TV 2/4 (xe, Mb) T (2)0((r,
C(x, y) exp(
+
2
2
2
|14 + Voz"!
1
+ C(x, y) exp( | 1
C(x, y) e
Therefore, without the transformation of z, lim KL[us; (zly, æ)) Ipo; (=1y, []]y
For the case where the latent dimension d is larger than the response dimension q, we use the first q latent dimensions to build a projection between Z1:g and y|æ without the remaining d q latent dimensions. To be specific, let (1 z) = F=1 o T(21:)) and using the same derivation of Eq 20 we get
(y]=)] +.
Now, we define the mean function of the encoder where ps (z)(x) =T Fa(v) and po(g)z))+11 - and the variance function of the encoder as follows,
Ex y) al(So,o (x, y), n TTT.TaaS y),
(So
nig -.. nz))
18

---

Published in Transactions on Machine Learning Research (01/2025)
where n m set of g-dimensional column vectors, e.g. orthonormal basis of mull(So,)); such that
ST S8 0, n n. li=j.
(24)
The set of n +i= always exists due to the fact that So. is the d x q Jacobian matrix with null space of at least d - q- This implics,
Ex v) o (o,
(25)
0
la-q]
(z)2gg y. x) can exactly match the first dimensions of true posterior as shown in 23 23 The remaining d - q dimensions follow a standardized Gaussian distribution that matches the posterior po; (zly, æ) due to the fact that these dimensions are not used in likelihood involving po z) = Fg o-toT (z11g) as remains follow the prior distribution. As a result, lim, -0 KL[ga (# y. a>|lpe (=ly,=)] = 0.
Thus, The first dimensions of 4s (z)
C Proof of theorem 3.4
We consider G as follows,
- f f2x)) f2 (21) f2, (2, ) + f (x, 21) + f (xx, x2) + S (21,22) + faziza (0, 21, 22) (26)
Step 1: we prove, under the assumption 3.3, that there exists a unique decomposition of G in Eq.??
up to rotation of coordinates within the block. Specifically, given G there is unique form for for fx))
#2), and f (x, æ1, x2).
Following the line of proof of Theorem Sobol 1993), with an abuse of notation μ, we can see that E[Y], G is the transformation of X, Zı, Ze to generate Y.
f = G (I, 21, (2)u(dazz
Then, the integration of G G (æ, 21, 22) over any two block dimensions, namely f G (ar, 21, (d) admits the following.
|Calee 21, 22) =SØ [3())) [.Clernttdei + 12 (Zettritz))
P2 (or f (a. sy)n(dz))))
(27)
f (1 (d)(d
far 2n, 2z)
where the cancellation to zeros direct result of assumption 3.3 Therefore,
fz(x) = / G (ar, 21, 2z (22)
88-
f2, (21) = / G (x, æ1, x2) ) (dxdzz) 5%
(28)
f2, (zz) =
(2, 2 #
19

---

Published in Transactions on Machine Learning Research (01/2025)
Using the same argument, can
21, / G 21, 22)u(dz22
/
/
az (00,20 21)=
G'm (ar, 21, =2)p((dddze)
G (2) +
ux2 (x), x2) :=
G (x. 21, =2)ul(z))
G (2r, 21, 22 (2)
2 + +
F2 (21,22) = / (z, 21, /
G (a, 21, z2) /
(2)
G (1, 22 fö:
(29)
and lastly, the residual
1Z 21, x2) = (2, (ff + (z1) + + (22) + fizzz (x,z1) fizzz (x, x2) + f2 (Z1, (2) + fo').
(30)
This shows that all fo are unique given G
Step 2: We prove that, with additional assumptions including f2, (22) = 0² o-CVAE recovers the ground truth G
By comparison of block coordinates, (o 21) + 022 G is almost surely true if the following holds for all x, 21, 22 (up to zero measure of 22).
no (1, 21) 8 (a:)
12, (2)) f2x, (2,, 21) + 50 Co
022 12 (22)
(31)
fiz (21, 22) + făzz (I, 22) fax (I, 21, x2) Co
Where Co is constant independent of x, 22.
In alignment with the identifiability assumption 3.3 12 (22) cannot have any additive constant. Specifically, S f2 (22) = (2))(22)) = 0. Assuning the universal approximation theorem of uæ(æ, 21), the results are immediately implied by considering the additional assumption. The conditions and proof of the universal approximation property can be found in Hornik et al. (1989
D Comparison between Theorem 3.4 and existing approximation theories.
Table 5: Comparison between Theorem 3.4 and existing approximation theories.
Approximation theory
Key Limitation on oracle 0,0
Existence of optimal
Enough to explain a good fit two-moon data?
Gaussian encoders
Proposition A.2
Yes. Asymptotically
No. The learned decoder recov- ers the two-moon data generative process without taking
Lincar VAE
i represents linear function between and z.
Yes. Closed-form Gaus- sian posterior within the mean-field Gaussian en- coder family.
No. Two-moon data is has non- linear mean function w.r.t Z x 1111
Lucas et al. 2019b)
Theorem 3.4
No additional limitation.
No. It is not guaranteed.
Yes. The learned decoder a re- covers the eetta-enerating vari- ance.
20

---

Published in Transactions on Machine Learning Research (01/2025)
E An toy example of inflated estimate of o² due to suboptimal inference (via Fisher's identity)
This serves as a solid example of how less concentrated ts leads to inflated a. For simplicity, we consider the ID unconditional case, which in statistics, is called the many-normal-means problem.
We given sample from the following p(z) (0, p(,l() ~ (zz Easily, we have the following:
o²
p()) ~ N(O, + ())p(())) ~ N ( 1
x
1
Using Fishor's identity, the gradiont of the marginal log- log-likelihood p(z) given x;
b Olog p(x) = Ep(2(a))
1
(a 2)2
1
+
20 +
20 Ep pp(z|/))(< =x)].
dg²
20
20
where [p(z|m][( 2 (x pz)² + oz. and Ms 14o2,00 02
It02
By setting (p(=)))) 2 log | = 0, we have t = x² - 1, which corresponds to MLE of Blogp(ar) 0 Now if (2((x) is less concentrated than p(z|2), e.g.,
O6"
(=11x) ~ N(uz
where o ko? for some k >
θELBO
(r 2)
1
Eg(= mn)
+
[(
-me)² ko
Og²
20
20
20²
20
1
0t2²
ko²
+
20
20
(1+
(1+
0 n²2
h(1 + 02) - (1 02
20 o2)2
0 + (x² k- +
Notice that the only solution of łELBO/00 0 is
22 + 1: 2 + (x2 + 2)2 + 4(k: - 1) > x.2 + 2h - 2> x²
o
mle
Thus, such do introduces a positive bias to the MLE of a in log-likelihood, but it reduces the possibility of an ill-defined likelihood. For example, | 1 leads invalid estimates of o and thus James-Star shrinkage (JAMES 1961) may be considered; it can be avoided by 9s if l: 2. In other words, we suffer from the numerical instability of suboptimal Jø. rather than likelihood blowup possibly near-zero Nonetheless, the estimate of a from a less concentrated ds is inflated.
F Experiment and algorithm details
We implemented the proposed method using Porch 1.8.2 +culll with Python 3.7 on an Ubuntu internal cluster with multiple Nvidia GPUs including A A10, A A100, A A100-40GB, A A100-80GB, and V100. We are not aware of which GPU is used in the experiments due to the task distribution service.
In Section 5.1 we used fully connected 4-layer neural networks with a hyperbolic tangent activation function for the encoding and decoding networks. The latent. dimension is set to 2. and the width of the hidden layers
21

---

Published in Transactions on Machine Learning Research (01/2025)
is [16, 8, 4. and [2, 4, 46, respectively. initialized at 1. The batch size is equal to the sample size of the training data. We generate 5000 data points with 2, 500 for each class, with a latent variable generated from a standard normal distribution.
In Section 5.2 we used 3 simulation datasets from the following:
M1 non-linear model with additive Gaussian noise:
y z + exp (a2 + 23/3) + sin (a4 + æς) + E, where N(O, 1).
M2 is A nonlinear model with multiplicative non-Gaussian noise:
- (5 + af/3 + 2z + 2 + æ4 + æs) * exp(0.5 x where E ~ 0.5N -
0.5N
M3. A Gaussian Mixture Model:
N(-1 - 21 - 0.5xz, 0.5²), N(1 + 21 0.522, 12),
U = 0, i 1," where U ~ Binomial(1,0.77
y
The predictor dimension 5 for M1 and M2, and for M3. The response y is univariate. The sample size N 5000, the test sample size is A = 2000. The mean squared error of the mean is & Ex= (Blylæx) - B[ylæx)) The mean squared error of the standard deviation 1
k E (SD[ylæz) - SD / xx))
We used fully connected 5-layer neural networks with a hyperbolic tangent activation function for the en- coding and decoding network. The latent dimension is set to 5. The width of the hidden layer of the network is [32, 16, 8, 4], and [8, 32, 16, 4]. For the WGCS method, the conditional generator G is parameter- ized using a fully connected neural network. The discriminator D) is parameterized using fully connected two-layer neural network. The noise vector is l (0, 1). For the NNKCDE method, the tuning parame- ters are chosen by cross-validattion The bandwidth of CKDE is determined based on the standard formula = 1.06on 1.06onn1/(2ttd)) where a measure spread, is the order of the kernel, and is the dimension of X. The FlexCode basis expansion-based method uses the Fourier basis. The maximum number of bases is set to 40, and the actual number of bases is selected using cross-validation. WGCS generates 10, 000 Monte Carlo samples to estimate the conditional distribution of each test 2 to calculate the conditional mean and conditional standard deviation, while our method uses only 500 samples. For other methods, the estimates are calculated by numerical integration.
G More discussion for experiments in Section 5.1
Additional experiment of learned o-CVAE using Two moon dataset with 0.1. The ground truth recovers T, which is 0.1.
Additional experiments on wall time until convergence. In Table 6 under different K epochs for updating& 0, o, we report the largest wall time out of 5 repetitions of successful experiments with the estimated a initialized at 1, converging to the true value of 0.1. Tolcrance C = 0.05. The same critcrion for convergence is used each K. NA represents no correct estimate of the true value a within a relative error. We find that the success rate of experiments for uncalibrated o-CVAE increases with K and due calibration in the sufficiently large K, our proposed algorithm might be slower than uncalibrated one. Time measured by the Python function time.perffcounter//
Table 6: Comparison of wall time for convergence
INNER STEP SIZE K
25
50
100
250
500
UNCALIBRATED
NA
NA
62.15
76. 10
177.81
CALIBRATED W./ C=0.05
10.55
12.69
25.00
33.72
269.78
Calibration as warm restart. Nonotheless, the calibration step our framework beneficial for the estimation of g-CVAE even if optimal is learned. As shown by Lucas et al. (2019b), posterior collapse O occurs even if optimal is learned, possibly due to local optima in ga- Therefore, when calibrating ø,
22

---

Published in Transactions on Machine Learning Research (01/2025)
S
S
S
-12
(a)
(b)
(c)
Figure 4: Calibrated o-CVAE on two moon dataset: (a) the training data of size 5000, (b) estimated training data, and (c) sample drawn from calibrated o-CVAE. The true sigma, T 0.1, 0.2 and 0.3 from top to bottom. Best viewed in color.
the training dynamics of the parameter leverage the rapid change of perhaps due to smoother loss landscape with larger Dai et al. (2021), taking advantage to escape the local optima and converge to a possibly better local optimum at the time when a returns to optimally small again. Even if we calibrate the a back to initialization, the optimization that was done before this calibration would be beneficial as a warm restart. To this extent, our calibration step is an inexpensive and implicit KL annealing scheme without predefined monotonic or cyclic schedule Skafte et al.] (2019): (Stinn & Knowles (2020). Therefore, if sufficiently large K used. Please
recommend that this calibration step should be used note that warm restart does not guarantee the convergence of the algorithm in the non-convex ELBO optimization; throughout the experiments, observe that cold restart is necessary for instances with very poor initialization.
G.1 More discussion with existing KL annealing methods.
The explicit KL annealing is undoubtedly easy and simple implement. However, it is less efficient and often requires more computational power in terms of preventing postcrior collapse. For example, if a is assumedd to be fixed, an explicit KL annealing schedule like [2 is essentially series of deterministic calibrations in 1 due to a-3 equivalence. Therefore, the annealing process of the loss landscape is predefined and not adaptive to the performance of encoders. Similarly, when is a learned parameter, an explicit KL annealing schedule confounds with it in shaping the loss landscape, similarly to how it would by a single "effective" de- For example, when = 0.25 added to the LHS of Fq the "effective" a is half of the current value, i.c. a. = 0.5 * The confounding between , and will not vanish unless 3 is 1. In extreme cases, if the predefined that there would be no KL annealing to the loss landscape at all because there would be no change in the effective de- Such confounding not only complicates the annealing process of the loss landscape determined by the effective de, but also hinders the accurate and efficient estimation of a under the maximum likelihood principle.
1/(2o?), the effective aa would be fixed al 1. This means
Experiment details of Table cyclical annealing schedule, tuning 3 linearly from 0 For hyperparameters in monotonic annealing, cycle period M 10 of 5 maximum eycles with balanced ratio R = 0.5. In their original notation,
We trained the Gaussian o-CVAE with a learned o using a monotonic or
23

---

Published in Transactions on Machine Learning Research (01/2025)
the number of cycles is the ratio is 0.5, and the total iterations is 50. For hyperparameters in monotonic annealing, we consider single monotonic annealing with the same length M = 50. The 3 becomes fixed value of after a maximum of 50 iterations is reached until the algorithm converges. The hyperparameters of KL anncaling are chosen reasonably for a fair comparison. The maximum wall time until convergence is recorded by the Python function perf.perf counter() out of 5 successful repelitions.
H Ablation experiments on the effect of calibration
2.5
2,5
+
MSE of sigma
+ EL00
XL
MSE of sigma
EL0O
KL
1.5
1.5
Valu
Vsb6e
1.0
0.5
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
Sigmatfixd)
Initialization
(a) fixed 8/3-CVAE
(b) without calibration
2.5
MSE of sigma
ELB0
KL
MSE of gammma
ELU0
2.0
1.5
1.0
0.5
0.0
0.20 Callbration Toleration
0.45
Initialization
(c) with calibration
(d) ablation
Figure Experiments on twomoon dataset with ground truth = 0.1. Results of the CVAE, models are reported using (a) a fixed o, equivalent 3-CVAE with a predefined (b) a learned a without calibration with tolerance 0.05, (c) learned sigma with calibration, and (d) learned sigma with difference toleration for calibration which is intialized at 1. K is set to 25. Each data point is averaged over 20 repetitions. MSE of sigma is calculated by the squared error between the lcarned sigma and the ground truth.
In this section, more experiments are performed highlight the importance of calibration using the two- moon dataset in Scction 5.1 a nontrivial dataset simulated from nonlinear generating function that causes CVAE models to have reprodueible posterior collapse problem. Throughout this section, we use 0.1 to generate our training dataset once and repeat experiment 20 times. We use small K 25.
To reiterate the importance of lcarning a in CVAE models, we report the results of those using fixed o ranging from 0.1 to in Figure a). CVAE are well specified if the predefined fixed a is close to the ground truth value. The performance of CVAE degenerates rapidly as the difference between the fixed a and the ground truth increases. After the prodefined is exhibited by vanishing KL divergence between the approximate posterior and prior. Therefore, we should not CVAE with predefined variance scalar when the ground truth in the data unapproachable.
larger than a certain point, the posterior collapse CVAE
24

---

Published in Transactions on Machine Learning Research (01/2025)
To emphasize the numerical instability in the training dynamics of and the posterior collapse of the learned model, we report the results of CVAE that iteratively update without calibration in Figure 5 (b) with various imitializations of a ranging from 0.05 to 1. The learned CVAE model ends up with a very high ELBO, fails to recover the ground truth o, and the KL divergence between the approximate posterior and prior vanishes in most. cases regardless of the initialization of o. In Figure | we can see that a is trapped around 0.5 if no calibration is provided. It is known that the non-convex landscape of the ELBO function w.r.t. 0, & leads to many local optima, but our experiment specifically reveals the difficulty in estimating o in dual-step algorithm.
To provide more empirical evidence of our proposed method, we report the results of our calibrated method in Figure 5(c) for different initialistiions of a ranging from 0.05 to 1. The tolerance that triggers our proposed calibration step is set to 0.05. The results show the consistency of our proposed method in accurately estimating the truth of the ground o, preventing posterior collapse, and thus obtaining locally optimal
CVAE models.
To ablate the calibration effect, we report the results of our proposed method with various tolerance hyper- parameters C' ranging from 0.00 to 0.5 in Figure d).. We initialized a at 1. In the case that the tolerance is O, our proposed calibration step will never be triggered due to the non-negativity of KL divergence. We observe that a small tolerance of 0.05 is empirically effective in ameliorating posterior collapse and robust estimation of o. Note that the tolerance setting is case- case-by-case by-case and may not be generalized to other datasets, and a prodigious tolerance, which goes beyond the possible value of KL divergence, leads to non-convergence in the algorithm. Therefore, we recommend relatively small tolerance hyperparameter that triggers fewer calibrations and offers less extra computational burden.
Image reconstruction: MNIST
MNIST dataset LeCun et al. (1998) contains 60, 000 gray images with size 28 x 28. Dividing each image into two parts x, y. we treat one part of the image as the predictor x and the rest of the image the response y. In the Appendix [ we use fully connected 3-layer neural networks with ReLU activation function for the encoding and decoding network. The latent dimension is set to 5. The dimensions of X are 196, 392, and 588, and the corresponding dimensions of Y are 588,392, and 196. Two fully connected hidden layers are 256 and 128 for the encoder, and we reverse the width in the decoders. The batch size is 100. We update and calibrate a per 100 itcrations with a postcrior collapse tolerance set to 0.001. In Figure 6 we consider three situations in which 1/4, 1/2, and 3/4 of the image are observed, and the goal is to learn the distribution of the rest of the image. We chose 10 images from the test dataset to evaluate our method.
We see that the generated images are similar to the truth, with reasonable variations, and the variations in reconstruction decrease as larger part of the image is observed.
J Image generation: CelebA
The CelebA dataset Liu et al. (0155) contains more than 200, 000 colored celebrity face images with 40 facial attribute annotations. In this section, we used the architect of the CVAE models from Hou et al. (2019) and embedded the attribute label of a vector in the images channels. If a label in one dimension is 1, we embed it as an image channel of size 96 x 96 that takes a value of on every pixel. The width of the hidden layers is set [32, 64, 128, 256, 512] and the convolution layer with kernel size of a stride size of 2, and a padding size of 1. We used LeakyReLU activation functions at negative slope of 0.02 followed by batch normalization layer. The number of latent dimensions is set to 32. The batch size is 16. We update and calibrate each gamma per 1809 itcrations, which is a factor of total training sample size and postcrior collapse tolerance set to 0.05. We apply the proposed method to the generation of human face images with binary label features as a predictor. Since some of the attributes are highly correlated, we validate our proposed framework using the following attributes Male, Young, Eyeglasses, Bald, Mustache, Smiling to be the predictor x, and the response y here is the center-chopped scaled images with a size of 96 x 96. We show the six types of true images and the generated images in each row of Figure 7 The attribute labels
25

---

Published in Transactions on Machine Learning Research (01/202)
Trum
Recondtructd
Riput
Neconatratedd
0
f
G
Q
D
O
O
O
O
O
O
O
O
O
O
O
/
/
/
/
/
/
/
/
/
/
/
/
/
/
/
2
;
3
2
2.
2
3
2
2
2
2
2
2
2
2
2
2
2
2
2
2
3
:
3
3
3
3
3
3
3
3
3
3
3
3
3
?
3
3
3
3
3
4
4
4
4
4
4
4
4
4
4
4
4
4
L
4
4
4
4
4
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
6
4
6
6
6
6
6
6
E
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
6
<
9
9
6
9
9
6
8
8
8
8
8
6
8
6
6
6
6
6
9 .
S
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
Figure 6: Reconstructing images in the MNIST test dataset. The left column of each panel contains the true images from the test dataset that are not used for training. The second left column contains the given part of the image that are used as predictor (i.e. upper left 1/4 in (a), left half 1/2 in (b), and upper left half 3/4 in (c)), and the rest 5 columns are the reconstructed images.
for these three types are shown in Table 7 moderate clarity.
Our generated images preserve the face attributes input with
Figure 7: (a) The true images in CelebA. Images of the same row have the same attributes. (b) Generated (not reconstructed) images with the same attributes. Each row corresponds to specific type of face identical to the same row of (a)
Limitations Although our simple calibration o-CVAE method docs not provide a stato-of-tho-art image generalion framework, we hypothesize that the conditional distribution of the face image may not satisly the assumption 3.3 Expanding the analysis in consideration of modern variants/techniquee that are specific for (conditional) imagine generation important subsoquent work. Further efforts are needed make it a state-of-thheare VAE variants for conditional image generator. One common challenge in conditional image generation is meaningful quantitative measures to generated samples. Metrics such test likelihood
26

---

Published in Transactions on Machine Learning Research (01/2025)
MSE failed in measuring the clarity of generated face images, and the ones like FID Inception Score add extra mathematical assumptions that are difficult to validate.
Table 7: Attributes for six types of face images in Figure
MALE
YOUNG
EYEGLASSES
BALD
MUSTACHE
SMILE
TYPE
+1
TYPE 2
+1
-1
+1
-1
-1
-1
TYPE +1
+1
+1
TYPE 4
-1
+1
-1
-1
-1
-1
TYPE
-1
+1
TYPE 6
-1
+1
-1
-1
-1
+1
27


Reviews:
[{"review_id": "dHAuajKKQa", "reviewer": "Reviewer_sQRH", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed algorithm is intuitive and simple to implement. It is not entirely clear how to choose the calibration tolerance $C$ and the number of steps $K$, although the experiments provide insight for the latter tuning parameter. I would agree that the method is easier to tune than annealed VAE, although I wonder if it is less flexible. The authors demonstrate the benefits of the calibration step, showcasing its ability to recover the true $\\sigma$ in a simulation setting.  I would suggest moving the comparison with KL annealing to the main body and highlighting the results in Table 5.\n\nThe theoretical motivation is nice. I enjoyed the discussion about the inference and the approximation gap (although I believe it can be improved), and the equivalence between $\\sigma$-CVAE and $\\beta-CVAE. I believe Sections 1.2 and 1.3 can be more concise.\n\nEquation (4) plays a central role but, as is, it is difficult to decipher. I would indicate that the inference gap is given by the KL divergence between $q_\\phi$ and $p(z \\mid x)$. I'd also have a preference for flipping the inequalities, and working with the ELBO rather than the negative ELBO, just to agree with the fact that the ELBO is a lower-bound and that typically we try to maximize the evidence.\n\nSection 2 seems interesting but lacks clarity. I could not understand the main point, namely that $\\sigma$-CVAE approximates arbitrary densities, even when the variational inference is suboptimal and the decoder non-identifiable. My questions are the following:\n* how restrictive is condition (2) in Theorem 2.4 (note the typo in the first term of the neural block decomposition: it should be $f^\\eta_{z_1 z_2}(z_1, z_2)$ instead of $f^\\eta_{x z_2}(z_1, z_2)$)?\n* can the authors indicate which terms depend on $q_\\phi$? Is there an optimal solution for any variational parameter $\\phi$?\n\nI'm not convinced by the title of Section 3.1. \"The necessity of $\\sigma$-calibration: bias in the gradient\". The main argument seems to be that our estimate of $\\sigma$ is biased, because we minimize the ELBO instead of the evidence. While this indeed leads to a bias estimator, I don't understand how this motivates $\\sigma$-calibration. On the other hand, I find the discussion in 3.2 and the connection to annealing more compelling motivation for $\\sigma$-calibration.\n\nIn 3.3, I disagree with the authors' claim that their objective sticks to the original ELBO loss and avoids explicit use of $\\beta$. First, as the authors note, there is an equivalence in loss and secondly, algorithm 1 does not quite define an optimization algorithm, since the calibration step is a departure from optimization. Framing it as a minimizer of the ELBO is misleading, and annealing has the merit of making the change in loss function explicit. By contrast, it is not clear which objective function, if any, the proposed algorithm minimizes. \n\nIn 4, the experiments are interesting and demonstrate the benefit of adding a calibration step, as well as comparing different calibration strategies. Still, I have several reservations:\n* How much sense does it make to compare $\\sigma$ between different models? In Figure 1, right (a) and (b), the performance of vanilla $\\sigma$ -CVAE would be worse if we reduced $\\sigma$ to 0.1. The large $\\sigma$ accounts for variability which is unexplained by the model. Moreover, $\\sigma$, on its own is not a meaningful quantity, it only exists in the context of a model. What would be more interesting is to compare uncertainty calibration, for example using a score function on held-out data.\n* Similarly, I'm not sure I understand the comparison of SD in Table 2. Why is a lower SD better?  \n* Could the figures, including Figure 1, account for the varying cost per iteration, as $K$ varies?\n* I don't understand the takeaway of Table 2. It seems $\\sigma$-CVAE underperforms. The authors write \"our method is very competitive with other methods\" but this qualitative statement is vague. What would an uncompetitive algorithm look like?\n\nAppendix A should appear in the main body.\n\nThe writing of the paper can be improved:\n* there are numerous syntax errors and typos. Please proofread the manuscript.\n* there are too many vacuous words, such as \"complex\" or \"reliable\". How do the authors define \"reliable\"? Similarly, the authors make extensive use of the word \"robust\", but usually do not specify \"robust\" to what. It is often unclear whether the authors are talking about posterior collapse, or poor estimates of $\\sigma$ (although my point above about how meaningful $\\sigma$ is on its own...).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates VAEs in which the decoder variance is a scaled identity matrix $\\sigma^2 I$. The object of the paper is to find ways to train the VAE, with particular attention given to $\\sigma$. The authors want (i) accurate estimates of $\\sigma$ and (ii) no posterior collapse (that is, the variational approximation $q$ reduces to the prior distribution $p(z)$). The authors propose an algorithm, which alternates optimizing the variational parameter $\\phi$, the decoder's weights $\\theta$, and the decoder's variance $\\sigma$. The algorithms departs from standard optimization by introducing a \"calibration\" step, which increases $\\sigma$, if there are signs of posterior collapse. The method is compared to annealing methods, which rely on an inverse-temperature $\\beta$. The authors argue that their method is easier to tune and apply it to several examples.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the recommendations above.\n\nI have some minor comments:\n- In Table 1, what does the crossed checkmark mean?\n- p7, the authors write \"it collapse to a nonlinear regression\". Do you mean \"linear\"?\n- p7, define mod.\n- p7 don't use $C$ for tuning parameter, since it is used for condition (2) of theorem 2.4\n- p10 can you show the optimization paths for multiple initializations? \n- p11 provide a summary of the models in the main body, even if it's only one sentence per model, and indicate the dimension of each model. The reader should understand the scope of the experiments, without reading the appendix.\n- p12 can you recommend starting defaults for $K$ and $C$?", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "6HjTuZ5QnB", "reviewer": "Reviewer_kWBh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The main strength of this paper is its algorithmic contribution: it proposes a specific annealing schedule for the variance parameter in $\\sigma$-CVAEs, and it demonstrates experimentally that this annealing schedule---which is not fixed in advance but triggered by evidence of posterior collapse---leads to better learning. There is also a theoretical contribution--a proof that under certain assumptions, $\\sigma$-CVAEs are expressive enough with a nonzero value of $\\sigma$ to parameterize an arbitrarily complex distribution. Discussions throughout the paper attempt to provide a coherent picture of inference and learning in $\\sigma$-CVAEs.\n\nThe main weakness of the paper is that it is poorly written at many levels. The following are some areas of the writing that need improvement:\n\n* The introduction does not begin by providing a high-level overview. I encourage the authors to give one or more talks on this work; I think it will force them to present the big picture first. It's unusual for a paper to have five equations by the bottom of page two. It should be possible to discuss the bigger issues of non-identifiability, posterior collapse, and inference/approximation gaps in VAEs without first introducing (nearly) all of the paper's notation.\n\n* Individual sentences are poorly balanced. For some examples, see the first sentence of the abstract, or the first sentence of 3.2. \n\n* The logical connectives between sentences do not always flow: e.g., the phrase \"in other words\" should only be used to introduce an idea that is an obvious recapitulation of what comes before.\n\n* The authors do not seem to understand the difference between in-text (\\citet) and parenthetical citations (\\citep). Some passages of text are nearly unreadable.\n\nThe above list is not exhaustive, and the cumulative effect of these issues is to place an undue burden on the reader. A paper should not require several passes to understand its main points.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "SUMMARY OF PAPER\n\n1.1 The authors study the twin problems of learning and inference in $\\sigma$-conditional variational autoencoders ($\\sigma$-CVAEs) whose latent variable has a data-independent prior.\n\n1.2 The loss function for these architectures is the negative ELBO. The authors highlight two issues when optimizing this loss function. The first is the quality of the variational approximation used for inference; the second is the expressiveness of the generative model used for learning. The authors argue that these issues must be examined in tandem to be fully understood.\n\n1.3 The authors review the relationship, previously reported, between a $\\beta$-CVAE (where $\\beta$ is an annealing parameter) and the $\\sigma$-CVAE with a fixed unit variance $\\sigma^2=1$.\n\n2.1 The authors review an earlier result (Agrawal & Domke) that the randomness of any distribution can be \"outsourced\" to a Gaussian random variable, and use it to show that a $\\sigma$-CVAE can in principle express an arbitrarily complex distribution in the limit $\\sigma\\rightarrow 0$. \n\n2.2 The authors review a previous result (Sobol) on block neural decomposition, then show how with related assumptions $\\sigma$-CVAEs can provide the same guarantee in (what they call) the *non-asymptotic* regime when $\\sigma$ is finite but nonzero (Theorem 2.4). \n\n3.1 The authors review that the gradient of the ELBO is a biased approximation to the gradient of the model's true log-likelihood, and that this bias extends to gradient-based (and other types of) estimates of $\\sigma$.\n\n3.2 The authors argue that posterior collapse is a sign of a poor encoder, and that it should be used to anneal/calibrate the variance $\\sigma$ in $\\sigma$-CVAEs (and not merely the hyperparameter $\\beta$ in $\\beta$-CVAEs.) They also show how a previous schedule for $\\beta$-annealing can be converted into a schedule for $\\sigma$-annealing.\n\n3.3. The authors argue that the annealing schedule for $\\sigma$ should be tied to evidence of posterior collapse (rather than fixed and pre-determined in advance), and they propose a new calibrated schedule that does so.\n\n4.1 The authors show that $\\sigma$-CVAEs with this calibrated annealing schedule can learn the two-moon data set with a proper batch size K and calibration hyperparameter C. They argue that their method is *doubly robust* because (1) it works with varying initializations of $\\sigma$, and (2) it is better at avoiding posterior collapse.\n\n4.2 The authors report that their $\\sigma$-CVAEs outperform four other methods (WGCS, NNKCDE, CKDE, basis expansion) on three additional data sets.\n\n4.3 The authors compare calibrated $\\sigma$-CVAEs to Bayesian normalizing flows on 6 UCI data sets. The comparisons are based on log-likelihoods, which disadvantages the $\\sigma$-CVAEs, since they only compute lower bounds on likelihoods. With this in mind, the results for calibrated $\\sigma$-CVAEs appear to be competitive.\n\n\nSUMMARY OF CONTRIBUTIONS\n\nThe main contributions, as I see them, are the specific proposal (section 3.3) to anneal the variance parameter in $\\sigma$-CVAEs based on evidence of posterior collapse, rather than with a pre-determined annealing schedule, and the experimental results (section 4) showing that this strategy is effective. Theorem 2.4 also appears to be novel, but its meaning and significance were less clear.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The writing needs improvement throughout. The authors should also address the following points.\n\n* The authors stress in sections 1 and 2 that a zero approximation gap is achievable without perfect variational inference. From the writing, it seems that the authors find this to be surprising. But it is unclear why this should be surprising. A zero approximation gap is possible when the generative model is sufficiently powerful to parameterize (express) a distribution that perfectly fits the data. The expressibility of the generative model is not related in any way to the difficulty of approximating its posterior by variational inference. (Indeed, that is why the proofs of section 2 do not make any reference to the variational approximation, $q$, for the encoder.) \n\n* There are many possible connotations that can be attached to the words *robust* and *asymptotic*. For this reason, the authors should be as explicit as possible about their meaning when they first introduce them.\n\n* The authors should not use $q$ to denote both the variational approximation in section 1.1 and the variable dimensionality in Defintion 2.2.\n\n* I hope that the authors can expand (or clarify) the discussion of identifiability in section 2.2. I found the assumptions behind Theorem 2.4 to be somewhat unintelligible. Also, is there any connection here to the relationship between posterior collapse and latent non-identifiability discussed in the NeurIPS-2021 paper by Wang, Blei, and Cunningham?\n\n* The authors stress in 3.1 that gradient-based (or coordinate-descent-based) updates of the ELBO are biased because they are designed to maximize the ELBO rather than the true likelihood. I believe this bias is well-known and widely appreciated in all models where the ELBO is maximized. It was not clear what special relevance this has for $\\sigma$-CVAEs, or why the authors devoted so much space to this point.\n\n* More generally, I often did not see the connections that the authors were attempting to make between different sections of the paper (e.g., sections 2 and 3, sections 3.1 and 3.2).\n\n* The authors provide pseudocode in Algorithm 1, but this pseudocode does not include the specific form of the calibration as given in eq. 18. It seems to be a strange omission.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "AH0GXtZvaE", "reviewer": "Reviewer_gj1z", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\n* Informative introduction to $\\sigma$ and $\\beta$-VAEs.\n* Goes beyond the usual linear VAEs assumptions, making the theory applicable for more complex, high-dimensional data.\n* Interesting insights regarding the biased gradient w.r.t $\\sigma$.\n* Analysis on the toy dataset provides interesting insights on the training dynamics when interleaved calibration of $\\sigma$.\n* Comparison with non-parametric density estimators and Flow-based models.\n* Detailed appendix.\n\n\n**Weaknesses**:\n* The claim in page 7 that the linear schedule is “computationally expensive” is unclear, and I don’t understand why the authors claim that annealing the KL divergence is a computational strain (“Obviously, these KL annealing techniques are computationally expensive and the statistical insight of the anneal schedule is obscure in principle.”). Also in the Discussion section it is referred to as “expensive”. Given the additional complexity of the proposed algorithm (see below), I don’t understand this claim. Either back this claim up with examples or relax it.\n* The proposed algorithm adds additional complexities to Rybkin et al. and introduces another hyper-parameter, the “calibration tolerance” $C$ - which requires tuning, and the calibration is done by using a heuristic (i.e., why sample $\\phi$ and increase the variance by 1 after each update?). Is this heuristic somehow related to the theory? Also, in Algorithm 1 there is a condition statement “if Convergence”, what does it mean? From section 4, “convergence threshold is 0.001 in the average loss change”, is that another hyper-parameter?\n* The performance on the image datasets is not entirely convincing (but I appreciate the discussion on that in Section 5).\n\n**Minor**:\n* Page 2, right after Eq.3, there is this sentence: “To understand the optimality of the parameters, the double inequalities in Eq. 4 is straightforward. “ – I found it very unclear.\n* It seems that the citations format is a bit undefined. For example, just after Eq. 8 on page 4, I believe the authors meant to use \\citep but the citation used something else.\n* Page 7: “A general pseudo-code can be found below.” - I recommend explicitly referring to the algorithm using \\ref{}.\n* Section 4 - Experiments - missing citations for optimizer (Adam) and datasets (MNIST, CelebA).\n* “Calibration tolerance” (section 3.3) = “posterior tolerance” (Algorithm 1) = “convergence threshold” (section 4) ? Please clarify that, because I think that according to Algorithm 1 “convergence threshold” is something else, like an additional hyper-parameter for performing the calibration.\n* Just before section 4.1 - “condition image generation” -> “conditional image generation”.\n* Section 4.1, please explicitly mention what $K$ is when it first appears (I somehow linked it to Algorithm 1 where it is the number of maximal iterations).\n* Figure 1 - “sigma” -> “$\\sigma$”.\n* Figure 3 - the legend says “EBLO” instead of “ELBO”.\n* M1-3 datasets - please refer the reader to the Appendix when you first mention them, as their details are there.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper aims to provide a theoretical analysis for two commonly used components in improving VAEs and their training dynamics is practice: (1) using $\\sigma$ as a model parameter or hyper-parameter of the decoder when calculating the likelihood, or the “reconstruction loss”; and (2) using the $\\beta$ hyper-parameter to weigh the KL-divergence term in the ELBO. The authors focus on the conditional setting–CVAE and $\\sigma$-VAEs. The paper also proposes Calibrated Robust $\\sigma$-CVAEs, an approach where $\\sigma$ is calibrated to mitigate numerical instability and prevent posterior collapse. The proposed algorithm is analyzed on 1 toy dataset, and also demonstrated on several tabular datasets, density estimation datasets and 2 image datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Maybe I missed it, but it would be helpful to explicitly mention the role of each variable, e.g. $y$ for the data distribution, $x$ for the conditioning, $z$ for the latent and etc…\n* Can the authors explain how they derived Eq. 7, and what does it mean in practice? \n* Please explain the motivation or the idea behind the proposed algorithm in practice–e.g., how the tolerance $C$ is chosen/tuned, how did you come up with the specific heuristic for the calibration update?\n* In section 4, can you please explain what it means to generate a dataset with a true $\\sigma$? “We simulated three datasets of size n = 5 000 with 2,500 for each class using $\\sigma$ = 0.1.” This will explain how you can actually recover it in the experiment.\n* Can the authors provide any insights as to what happens in non-conditional settings when using $\\sigma$? [non-critical]\n* Fix “Minor” from above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewers overall liked the idea of the paper being simple and relatively easy to implement. They appreciated the mathematical analysis regarding sub-optimality of variational inference and how that relates to posterior collapse. During discussion period the paper was improved with better presentation and additional motivating examples. 

I recommend the following revision actions for camera ready:

1. As per author-reviewer sQRH discussions: a further clarification of the approximation theory. (https://openreview.net/forum?id=VIkycTWDWo&noteId=NPhSjt6Cq8)

2. Clarify the discussion on expressiveness, and how is this related to posterior collapse.
Decision: Accept (poster)