PH-VAE: A Polynomial Hierarchical Variational Autoencoder Towards Disentangled Representation Learning

TL;DR

PH-VAE tackles the lack of disentanglement and brittle posterior in standard VAEs by disentangling input features into a polynomial hierarchical form and regularizing with a Polynomial Divergence that generalizes KL. The model employs a multi-encoder setup over polynomial features, a shared decoder, and a reparameterization z = mu + A epsilon sigma, with an ELBO replacement that includes a hierarchical prior and a closed-form PH divergence. Empirical results across probabilistic-density reconstruction and color/gray-scale image datasets show improved density tracking, reconstruction fidelity, and latent disentanglement, particularly under limited data, validating the approach's robustness and scalability. This work provides a principled, unsupervised pathway to richer latent representations and more controllable generative processes without increasing data size.

Abstract

The variational autoencoder (VAE) is a simple and efficient generative artificial intelligence method for modeling complex probability distributions of various types of data, such as images and texts. However, it suffers some main shortcomings, such as lack of interpretability in the latent variables, difficulties in tuning hyperparameters while training, producing blurry, unrealistic downstream outputs or loss of information due to how it calculates loss functions and recovers data distributions, overfitting, and origin gravity effect for small data sets, among other issues. These and other limitations have caused unsatisfactory generation effects for the data with complex distributions. In this work, we proposed and developed a polynomial hierarchical variational autoencoder (PH-VAE), in which we used a polynomial hierarchical date format to generate or to reconstruct the data distributions. In doing so, we also proposed a novel Polynomial Divergence in the loss function to replace or generalize the Kullback-Leibler (KL) divergence, which results in systematic and drastic improvements in both accuracy and reproducibility of the re-constructed distribution function as well as the quality of re-constructed data images while keeping the dataset size the same but capturing fine resolution of the data. Moreover, we showed that the proposed PH-VAE has some form of disentangled representation learning ability.

Figures (10)

• Figure 1: Architecture of the polynomial-enriched hierarchical variational autoencoder.
• Figure 2: Probability densities of the predicted results for PH-VAE, VAE and the ground truth samples. (a1) Uniform distribution: $A$=1; (b1) Uniform distribution: $A$=3; (c1) Uniform distribution: $A$=5; (a2) Log-normal distribution: $A$=1; (b2) Log-normal distribution: $A$=3; (c2) Log-normal distribution: $A$=5.
• Figure 3: Values of the loss function in each epoch for the polynomial features $\boldsymbol{x}^s$. ($a1$) Uniform distribution: $A$=1; ($b1$) Uniform distribution: $A$=3; ($c1$) Uniform distribution: $A$=5; ($a2$) Log-normal distribution: $A$=1; ($b2$) Log-normal distribution: $A$=3; ($c2$) Log-normal distribution: $A$=5.
• Figure 4: Comparison of the ground truth data and the reconstructed results from VAE and PHVAE in the first failure case. (a) $A=1$; (b) $A=3$, and (c) $A=5$.
• Figure 5: Comparison of the ground truth data and the corresponding results obtained from VAE and PH-VAE in the second failure case.