Doubly Robust Conditional VAE via Decoder Calibration: An Implicit KL Annealing Approach

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.