Density-Informed VAE (DiVAE): Reliable Log-Prior Probability via Density Alignment Regularization
This work addresses the issue of unreliable log-prior probabilities in VAEs for researchers and practitioners, but it is incremental as it builds on existing VAE frameworks with a new regularization technique.
The paper tackled the problem of standard VAEs overlooking data-space density structure by introducing DiVAE, a lightweight regularizer that aligns VAE log-prior probability with estimated data density, resulting in improved distributional alignment, prior coverage, and OOD uncertainty calibration on synthetic datasets and MNIST.
We introduce Density-Informed VAE (DiVAE), a lightweight, data-driven regularizer that aligns the VAE log-prior probability $\log p_Z(z)$ with a log-density estimated from data. Standard VAEs match latents to a simple prior, overlooking density structure in the data-space. DiVAE encourages the encoder to allocate posterior mass in proportion to data-space density and, when the prior is learnable, nudges the prior toward high-density regions. This is realized by adding a robust, precision-weighted penalty to the ELBO, incurring negligible computational overhead. On synthetic datasets, DiVAE (i) improves distributional alignment of latent log-densities to its ground truth counterpart, (ii) improves prior coverage, and (iii) yields better OOD uncertainty calibration. On MNIST, DiVAE improves alignment of the prior with external estimates of the density, providing better interpretability, and improves OOD detection for learnable priors.