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

The data samples, originating from either experimental measurements or real-world applications, form the foundation for the analysis or model development process. Assuming that raw data is structured as follows,

$$\{\bm{y}_{n}\}\quad(n=1,\dots,N),$$

(1)

where $n$ is the number of features; $\bm{y}_{n}$ denotes the sample vector regarding each feature $n$. The dataset adopted for PH-VAE involves standardizing raw data to a consistent scale. Herein, all data are mapped to the range $[0,1]$ as follows,

$$\bm{x}_{n}=\frac{\bm{y}_{n}-\min(\bm{y}_{n})}{\max(\bm{y}_{n})-\min(\bm{y}_{n})},$$

(2)

where $\min(\cdot)$ represents the minimum value of $\cdot$ and $\max(\cdot)$ represents the maximum value. While for image data, additional steps are required, including resizing the images to a prescribed dimension and applying enhancement techniques to increase the model’s generalization.

In order to extract more information without requiring additional sampling, the dataset $\bm{x}_{n}$ undergoes a series of transformations to generate polynomial features, such as the original data $\bm{x}_{n}^{1}$, its square $\bm{x}_{n}^{2}$, its cube $\bm{x}_{n}^{3}$, and so on. These higher-order features enable the model to capture more complex relationships within the data, improving its capacity to learn non-linear patterns. The transformed data is then transposed to ensure each row represents a feature vector corresponding to a single sample, following the standard format used in PH-VAE models. Therefore, the dataset is provided as follows,

$$\{\bm{x}_{n}^{s}\}\quad(n=1,\dots,N;s=1,\dots,S),$$

(3)

where $s$ is an index to denote the $n$-th data set of the total $N$ data set, and each data set has $m$-number of feature. The index $s$ indicates $s$-th power operation on the original data.

For example, for a given input data feature set, ${\bm{x}}_{n}^{T}=(x_{1},x_{2},\cdots,x_{m})_{n}$, we can reformat it as

$${\bm{x}}_{n}~{}\to~{}\left[\begin{array}[]{c}{\bm{x}}^{1}_{n}\\ {\bm{x}}_{n}^{2}\\ \vdots\\ {\bm{x}}_{n}^{s}\\ \vdots\\ {\bm{x}}_{n}^{S}\\ \end{array}\right]$$

(4)

where ${\bm{x}}^{s}$ is defined as

$${\bm{x}}^{s}_{n}:=\left(\begin{array}[]{c}x_{n1}^{s}\\ x_{n2}^{s}\\ \vdots\\ x_{nj}^{s}\\ \vdots\\ x_{nm}^{s}\\ \end{array}\right),~{}~{}~{}~{}s=1,2,\cdots,S$$

(5)

Subsequently, the data will be processed in a certain batch size and be shuffled to introduce randomness during training. This comprehensive preparation ensures the data is well-structured and normalized, supporting the training of PH-VAE models effectively.

The proposed PH-VAE framework is an extension of the original VAE  (Kingma and Welling,, 2014; Doersch,, 2016), and the reformulated input data is a hierarchical form of parallel data sequence.

By doing so, we can process multi-scale polynomial features that share a same decoder across these inputs as shown in Fig. 1. This architecture can be particularly efficient in scenarios where the available training data is limited. Now we dive into the formulation and architecture in the context of PH-VAE.

Each exponent $s$ in dataset $\bm{x}_{n}^{s}$ ($s=1,\dots,S$) accords with an encoder. The encoder is composed of a fully connected layer followed by an activation function, resulting in the following formula,

$$\bm{h}^{s}=g\left(\bm{W}^{s}\bm{x}_{n}^{s}+\bm{b}^{s}\right),$$

(6)

in which, $\bm{g}(\cdot)$ is the activation function; $\bm{W}^{s}$, $\bm{b}^{s}$ are the undetermined weight matrix and bias vector for the s-th encoder, respectively; $\bm{h}^{s}$ represents the output vector of the encoder.

With the utilization of $\bm{h}^{s}$, two separated fully connected layers are applied to compute the mean $\bm{\mu}_{s}$ and the logarithm of the variance $\log(\bm{\sigma}_{s})^{2}$ for the latent variable $\bm{z}$ as follows,

	$\displaystyle\bm{\mu}_{s}$	$\displaystyle=$	$\displaystyle\bm{W}_{\mu}\bm{h}^{s}+\bm{b}_{\mu},$		(7)
	$\displaystyle\log\bm{\sigma}_{s}^{2}$	$\displaystyle=$	$\displaystyle\bm{W}_{\sigma}\bm{h}^{s}+\bm{b}_{\sigma},$		(8)

where $\bm{W}_{\mu},\bm{W}_{\sigma}$ are the weight vectors for the respective transformations; $\bm{b}_{\mu},\bm{b}_{\sigma}$ are the bias parameters. Then the latent variable $\bm{z}$ is sampled from the Gaussian distribution through the reparameterization trick, allowing to express $\bm{z}$ as,

$$\bm{z}=\bm{\mu}+A\cdot\bm{\epsilon}\cdot\bm{\sigma},$$

(9)

in which, $\bm{\epsilon}$ is a vector sampled from the standard normal distribution $\mathcal{N}(0,1)$, introducing randomness into the latent space; $\bm{\mu}={1\over S}\sum_{s=1}^{S}\bm{\mu}_{s}$ and $\bm{\sigma}^{2}={1\over S}\sum_{s=1}^{S}\bm{\sigma}_{s}^{2}$; as the amplitude of $\bm{\epsilon}\cdot\bm{\sigma}$, $A$ is exploited to enhance the neural network’s expressiveness and allows for better representation of complex latent space structures beyond independent scaling.

Then the latent variable $\bm{z}$ is fed into the decoder layer to reconstruct the original input $\bm{x}_{n}^{1}$, enabling the decoder to use a unified representation of the inputs while preserving their shared latent features. It can be defined as,

$$\hat{\bm{x}}_{n}^{1}=f(\bm{W}\bm{z}+\bm{b}),$$

(10)

where $f(\cdot)$ is the activation function; $\bm{W},\bm{b}$ are the weight matrix and bias vector to be determined in the decoder layer, respectively.

In the proposed PH-VAE, a structured generation process is constructed, where the observed data $\bm{x}_{n}^{1}$, and the latent variables $\bm{z}$ are jointly distributed as follows,

$$p(\bm{x},\bm{z})\approx p(\bm{x}|\bm{z})\cdot{1\over S}\sum_{s=1}^{S}p(\bm{z}|\bm{x}^{s}),$$

(11)

where $p(\bm{x}|\bm{z})$ indicates the likelihood function that models the data reconstruction from $\bm{z}$, and $p(\bm{z}|\bm{x}^{s})$ denoting the hierarchical prior, capturing the conditional dependency of $\bm{z}$ on $\bm{x}^{s}$. Eq. (11) is derived from the Importance Weighted Autoencoder (IWAE) metric. For simplicity, the subscript $n$ of $\bm{x}_{n}^{1}$ has been removed during the derivation. The exact posterior distribution $p(\bm{x}|\bm{z})$ is generally intractable due to the complexity of the hierarchical dependencies. To solve this problem, a surrogate model $q(\bm{z}|\bm{x})$ is introduced as an approximation, which is expressed as follows,

$$q(\bm{z}|\bm{x}):={1\over S}\sum_{s=1}^{S}q_{s}(\bm{z}|\bm{x}^{s}),$$

(12)

where $q_{1}(\bm{z}|\bm{x})$ encodes the dependency of the first latent layer on the input data $\bm{x}^{1}$, and $q_{s}(\bm{z}|\bm{x}^{s}),s=2,\cdots S$ capture the conditional dependency between successive latent layers. This factorization mirrors the hierarchical structure of the generative model, facilitating efficient computation of the variational posterior and hence a disentangled representation learning. Note that $q({\bf z}|{\bf x})\not=q_{1}({\bf z}|{\bf x})$.

In this work, we propose the following evidence lower bound function,

$$\mathcal{L}^{PH}_{\text{ELBO}}:=\mathbb{E}_{q(\bm{z}|\bm{x})}\left[\log p(\bm{x}|\bm{z})\right]-{1\over S}\sum_{s=1}^{S}\text{KL}\left[q_{s}(\bm{z}|\bm{x}^{s})\|p(\bm{z})\right],$$

(13)

where, by definition, the Kullback-Leibler (KL) Divergence for each layer is given as

$$\text{KL}\left[q_{s}(\bm{z}|\bm{x}^{s})\|p(\bm{z})\right]=\int_{\bm{z}}\Bigl{(}q_{s}(\bm{z}|\bm{x}^{s})\log\Bigr{(}{q_{s}(\bm{z}|\bm{x}^{s})\over p(\bm{z})}\Bigr{)}d{\bm{z}}$$

(14)

We define the average of the Kullback-Leibler (KL) Divergences of all layers as the Polynomial Hierarchical Divergence, i.e.

$${\rm PH}(q({\bm{z}}|{\bm{x}})\|p({\bm{z}})):={1\over S}\sum_{s=1}^{S}\text{KL}\left[q_{s}(\bm{z}|\bm{x}^{s})\|p(\bm{z})\right]~{}.$$

(15)

One can verify that ${\rm PH}(q({\bm{z}}|{\bm{x}})\|p({\bm{z}})$ is indeed a divergence by definition, i.e.

$${\rm PH}(q({\bm{z}}|{\bm{x}})\|q({\bm{z}}|{\bm{x}}))=0~{}~{}{\rm and}~{}~{}{\rm PH}(q({\bm{z}}|{\bm{x}})\|p({\bm{z}}))>0$$

We note in passing that by the standard derivation procedure one can show that

	$\displaystyle\int_{\bm{z}}q_{s}({\bm{z}}|{\bm{x}}^{s})\log{q_{s}({\bm{z}}|{\bm{x}}^{s})\over p({\bm{z}},{\bm{x}})}d{\bm{z}}$	$\displaystyle=$	$\displaystyle\mathbb{E}_{q_{s}({\bm{z}}|{\bm{x}}^{s})}\Bigl{[}\log{q_{s}({\bm{z}}|{\bm{x}}^{s})\over p({\bm{z}},{\bm{x}})}\Bigr{]}$		(16)
		$\displaystyle=$	$\displaystyle\mathbb{E}_{q_{s}({\bm{z}}|{\bm{x}}^{s})}\Bigl{[}\log q_{s}({\bm{z}}|{\bm{x}}^{s})-\bigl{(}\log p({\bm{x}}|{\bm{z}})+\log(p({\bm{z}}))\bigr{)}\Bigr{]}$

Hence,

	$\displaystyle\mathcal{L}^{PH}_{ELBO}$	$\displaystyle=$	$\displaystyle-{1\over S}\sum_{s=1}^{S}\left\{\int_{\bm{z}}q_{s}({\bm{z}}|{\bm{x}}^{s})\log{q_{s}({\bm{z}}|{\bm{x}}^{s})\over p({\bm{z}},{\bm{x}})}d{\bm{z}}\right\}$		(17)
		$\displaystyle=$	$\displaystyle{1\over S}\sum_{s=1}^{S}\mathbb{E}_{q_{s}({\bm{z}}|{\bm{x}}^{s})}\Bigl{[}-\log q_{s}({\bm{z}}|{\bm{x}}^{s})+\bigl{(}\log p({\bm{x}}|{\bm{z}})+\log(p({\bm{z}}))\bigr{)}\Bigr{]}$
		$\displaystyle=$	$\displaystyle\mathbb{E}_{q(\bm{z}|\bm{x})}\left[\log p(\bm{x}|\bm{z})\right]-{1\over S}\sum_{s=1}^{S}\text{KL}\left[q_{s}(\bm{z}|\bm{x}^{s})\|p(\bm{z})\right]~{}.$

In the last line, we used the definition of Eq. (12), i.e.

$$q(\bm{z}|\bm{x}):={1\over S}\sum_{s=1}^{S}q_{s}(\bm{z}|\bm{x}^{s})~{}.$$

Following the standard procedure, one can also derive an alternative expression of ELBO objective, i.e.

$$\mathcal{L}_{ELBO}^{PH}=\log p({\bm{x}})-{1\over S}\sum_{s=1}^{S}KL[q_{s}({\bm{z}}|{\bm{x}}^{s})\|p({\bm{z}}|{\bm{x}}))$$

(18)

where

$$KL[q_{s}({\bm{z}}|{\bm{x}}^{s})\|p({\bm{z}}|{\bm{x}}))=\int_{\bm{z}}q_{s}({\bm{z}}|{\bm{x}}^{s}){q_{s}({\bm{z}}|{\bm{x}}^{s})\over p({\bm{z}}|{\bm{x}})}d{\bm{z}}~{},$$

(19)

which is often called the reverse KL divergence.

From Eq. (13), one can see that $\text{PH}\left[q(\bm{z}|\bm{x})\|p(\bm{z})\right]$ penalizes deviations of the variational posterior from the prior. By assuming that each surrogate model is a Gaussian distribution, the polynomial divergence has a closed form as follows:

$$\text{PH}\left[q(\bm{z}|\bm{x})\|p(\bm{z})\right]=\frac{1}{2S}\sum_{s=1}^{S}\left(1+\log\bm{\sigma}_{s}^{2}(\bm{x}^{s})-\bm{\mu}_{s}^{2}(\bm{x}^{s})-\bm{\sigma}_{s}^{2}(\bm{x}^{s})\right)~{}.$$

(20)

It should be noted that the product of two Gaussian variables may not be a Gaussian variable, even though the product of two Gaussian PDFs is still a Gaussian PDF. Thus, Eq. (20) is intrinsically an approximation.

The negative ELBO is minimized to obtain the undetermined parameters $\bm{W}^{s},\bm{b}^{s}$, $\bm{W}_{\mu},\bm{b}_{\mu}$, $\bm{W}_{\sigma},\bm{b}_{\sigma},\bm{W}$ and $\bm{b}$, leading to the following loss function,

$$\mathcal{L}=-\mathbb{E}_{q(\bm{z}|\bm{x})}\left[\log p(\bm{x}|\bm{z})\right]+{1\over S}\sum_{s=1}^{S}\left[\frac{1}{2}\sum\left(1+\log\bm{\sigma}_{s}^{2}(\bm{x}^{s})-\bm{\mu}_{s}^{2}(\bm{x}^{s})-\bm{\sigma}_{s}^{2}(\bm{x}^{s})\right)\right],$$

(21)

In the Python code, the loss function with $S=3$ is defined and implemented as follows,

class PHVAE(nn.Module):
    def __init__(self):
        super(PHVAE, self).__init__()
        self.latent_dim = latent_dim
        self.encoder = nn.Sequential(......)
        self.mu = nn.Linear(......)
        self.logvar = nn.Linear(......)
        self.decoder = nn.Sequential(......)
    def reparameterize(self, mu, logvar, A=1):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + A * eps * std
    def encode(self, x, A=1):
        x = self.encoder(x)
        mu, logvar = self.mu(x), self.logvar(x)
        return self.reparameterize(mu, logvar, A), mu, logvar
    def forward(self, x1, x1**2, x1**3, A=1, S=3):
        z1, mu1, logvar1 = self.encode(x1, A)
        z2, mu2, logvar2 = self.encode(x1**2, A)
        z3, mu3, logvar3 = self.encode(x1**3, A)
        mu = (mu1 + mu2 + mu3) / S
        logvar = torch.logsumexp(torch.stack([logvar1, logvar2, logvar3]), \
                 dim=0) - math.log(S)
        z_combined = reparameterize(mu, logvar, A)
        recon_x1 = self.decoder(z_combined)
        return recon_x1, mu1, logvar1, mu2, logvar2, mu3, logvar3

def loss_function(recon_x1, x1, z1_mu, z1_logvar, z2_mu, z2_logvar, S=3):
    recon_loss1 = nn.functional.mse_loss(recon_x1, x1, reduction=’sum’)
    kl_loss1 = -0.5 * torch.sum(1 + z1_logvar - z1_mu.pow(2) - z1_logvar.exp())
    kl_loss2 = -0.5 * torch.sum(1 + z2_logvar - z2_mu.pow(2) - z2_logvar.exp())
    kl_loss3 = -0.5 * torch.sum(1 + z3_logvar - z3_mu.pow(2) - z3_logvar.exp())
    return recon_loss1 + (1*kl_loss1 + 1*kl_loss2 + 1*kl_loss3) / S

The training process with the loss function proceeds iteratively across epochs. Firstly, the forward propagation is exploited to produce the output predictions and compute the associated loss. Afterward, the backward propagation involves computing the gradients of the loss function with respect to the network parameters, accomplished using the chain rule of differentiation. Consequently, the Adam optimization algorithm is employed to minimize the loss function $\mathcal{L}$, and all parameters are updated step by step until $\mathcal{L}$ converges. The architecture of the PH-VAE is illustrated in Figure 1.

The proposed PH-VAE modifies the prior structure to a hierarchical polynomial form with the variational posterior. According to (Mutual Information) MI decomposition, ELBO objective in Eq. (14) can be modified as,

$$\mathcal{L}_{\text{ELBO}}^{PH}=\mathbb{E}_{q(\bm{z}|\bm{x})}[\log p(\bm{x}|\bm{z})]-I(\bm{x},\bm{z})-{1\over S}KL[q_{1}(\bm{z}|\bm{x})\|p(\bm{z})]$$

(22)

where

$$I(\bm{x},\bm{z}):={1\over S}\sum_{s=2}^{S}KL[q_{s}(\bm{z}|\bm{x}^{s})\|p(\bm{z})]$$

(23)

where $I(\bm{x},\bm{z})$ is the mutual information (MI). It contains the disentangled parts of higher order scale information. Compared to the standard VAE, PH-VAE improves MI by leveraging a hierarchical disentangled prior, which alleviates the restrictive assumptions on the latent variable $\bm{z}$. This allows PH-VAE to encode and preserve more information about the observed data $\bm{x}$, leading to a richer and more expressive latent representation. Additionally, PH-VAE mitigates the constraints imposed by the KL divergence regularization. By distributing the KL penalty across multiple hierarchical layers, the model prevents excessive compression of the latent space, thereby maintaining a more informative posterior distribution. Furthermore, the hierarchical structure of PH-VAE effectively reduces redundancy among latent dimensions, promoting greater disentanglement and ensuring that individual latent variables capture distinct generative factors more effectively.

For the training dataset, 20 sets of probabilistic patterns are pre-selected, including log-normal, normal, and uniform distributions, which are utilized to model PH-VAE. It should be noted that each set of distributions contains 50 data points. The raw data is further transformed and normalized as $\{\bm{x}_{n}^{s}\}$ ($n=1,\ldots,20;\ s=1,\ldots,S$). In this case, $S$ is chosen as 3. Thus, PH-VAE possesses three encoders, aiming to introduce diversity and facilitate shared feature learning. $\bm{x}_{n}^{s}$ are then encapsulated as TensorDataset and loaded into DataLoader with a batch size of 5.

Three encoders map the input data into separate latent space representations, producing both the mean $\bm{\mu}_{s}$ and log-variance $\log\bm{\sigma}_{s}^{2}$. The decoder receives a combined latent space representation, computed as the mean from three individual encoders. It is designated to encourage the coupling of information from multiple input streams into a unified latent representation. The dimensions of the input, hidden, and latent layers are specified as 20, 256, and 10, respectively. To be specific, the activation functions $g(\cdot)$ in Eq. (4) and $f(\cdot)$ in Eq. (8) are set to ReLU and Sigmoid, respectively.

The loss function in Eq. (14) is expressed as the sum of the reconstruction loss and polynomial divergence losses from all encoders as

$$L=-\mathbb{E}_{q(\bm{z}|\bm{x})}[\log p(\bm{x}|\bm{z})]+{1\over S}\sum_{s=1}^{3}\frac{1}{2}\sum\big{(}1+\log\bm{\sigma}_{s}^{2}-\bm{\mu}_{s}^{2}-\bm{\sigma}_{s}^{2}\big{)}$$

.

The Adam optimizer is employed with a learning rate of $5\times 10^{-4}$. This adaptive method dynamically adjusts the step size for each parameter, improving convergence. During each iteration, the model computes the gradient of the loss with respect to its parameters using back-propagation. The optimizer updates the parameters to minimize the loss. The training process iterates over a predefined number of 100 epochs. The results are provided below.

Taking log-normal and uniform distributions as the examples, the probability densities of the predicted results from PH-VAE, VAE, as well as the ground truth samples are depicted in Figure 2. The values of $A$ are 1, 3 and 5, which are intended to gradually improve the expressive capacity of VAE. To obtain the smoother curves, the program was executed 100 times to generate a total of 5,000 data points, which were applied to produce Figure 2. It is seen that, for each value of $A$, PH-VAE can better track the hidden probability information compared with VAE in the data, especially for the uniformly distributed data.

The loss function values with $A$=1,3,5 over successive epochs are comprehensively illustrated in Figure 3, providing a detailed visualization of the model’s convergence behavior during the training process. It is observed that as the values of $s$ increase, the enriched feature sets introduce additional information into latent space, enhancing its capacity to encapsulate more comprehensive data representations for each $A$. Consequently, the accuracy of the reconstructed probabilistic information improves. When executed on a standard desktop equipped with an Intel^® Core^TM i7-10750 CPU @ 2.60 GHz and 16 GB of RAM, the computational time was measured to be 6.9065 seconds.

According to (Yacoby et al.,, 2020), two specific scenarios under which VAE encounters significant pathologies in accurately learning the true data distribution $p(\bm{x})$ resort to posterior intractability and lacking of a simpler likelihood function. These conditions highlight fundamental limitations in the model’s ability to balance posterior approximation and likelihood estimation, ultimately impacting its generative performance and density estimation accuracy.

The first condition implies that VAE often relies on simple posterior approximations. However, if the true posterior has complex dependencies, the assumed posterior cannot capture these complexities. This mismatch between the true posterior and the approximating distribution leads to suboptimal learning, as ELBO cannot fully close the gap between the true data likelihood $p(\bm{x})$ and the model’s estimate. To verify that PH-VAE is still applicable under this case, Gaussian Mixture Models (GMMs) with nonlinear transformations are adopted to generate the dataset with the mean vectors

$$\begin{array}[]{l}\begin{bmatrix}3\\ 3\end{bmatrix},\begin{bmatrix}-3\\ -3\end{bmatrix},\begin{bmatrix}3\\ -3\end{bmatrix},\begin{bmatrix}-3\\ 3\end{bmatrix},\begin{bmatrix}0\\ 0\end{bmatrix}\end{array},$$

The covariance matrices are

$$\begin{array}[]{cc}\begin{bmatrix}1&0.8\\ 0.8&1\end{bmatrix},\begin{bmatrix}1&-0.6\\ -0.6&1\end{bmatrix},\begin{bmatrix}1&0.3\\ 0.3&1\end{bmatrix},\begin{bmatrix}0.5&0\\ 0&0.5\end{bmatrix},\begin{bmatrix}2&1.5\\ 1.5&2\end{bmatrix}\end{array}$$

where the weights are $(0.25,0.25,0.2,0.2,0.1)$.

Then the sampled data undergoes hyperbolic tangent $\tanh(\bm{u})$ and sinusoidal distortion $0.1\cdot\sin(2\bm{u})$ for further increasing the complexity of the dataset. The true posterior distribution for such a complex dataset is likely far from the Mean-Field Gaussian assumption commonly used in VAEs. Figure 4(a), (b) and (c) depict the prediction results obtained from VAE and PH-VAE on the aforementioned dataset under varying conditions of $A$, specifically $A$=1,3 and 5. It clearly demonstrates that while the VAE struggles to maintain recognition accuracy under the given conditions, the PH-VAE effectively preserves high recognition accuracy for $A=1,3$ and $5$.

The second condition indicates that the trade-off between posterior approximation and likelihood modeling results in a compromise that degrades the quality of the learned $p(\bm{x})$. The Normal distribution basis functions may not be expressive enough to accommodate likelihoods that yield both tractable posteriors and high-quality density estimates. In this case, the dataset includes a mix of correlated, anti-correlated, and narrow-spread Gaussian functions by applying the following mean vectors

$$\begin{bmatrix}-4\\ 0\end{bmatrix},\quad\begin{bmatrix}4\\ 0\end{bmatrix},\quad\begin{bmatrix}0\\ 4\end{bmatrix},$$

and the covariance matrices

$$\begin{bmatrix}1&0.8\\ 0.8&1\end{bmatrix},\quad\begin{bmatrix}1&-0.8\\ -0.8&1\end{bmatrix},\quad\begin{bmatrix}0.5&0.3\\ 0.3&0.5\end{bmatrix},$$

respectively, as well as a nonlinearly transformed component. The addition of noise and sinusoidal distortion $0.1\cdot\bm{\sin(u)}$ mimics real-world complexities, making it suitable for testing the robustness of generative models like VAEs.

The prediction results produced by VAE and PH-VAE under the condition of $A$=5 on the specified dataset are given in Figure 5. In contrast, the PH-VAE demonstrates superior performance, consistently achieving and preserving a high level of recognition accuracy, highlighting its robustness and effectiveness in addressing the identified limitations of the VAE.

Moreover, this paper examines the cluster example presented in reference(Yacoby et al.,, 2020). Most importantly, (Yacoby et al.,, 2020) has rigorously deduced the inability of VAE to accurately predict the data in this specific circumstance. Figure 6 illustrates the identified results of both VAE and PH-VAE concerning the dataset for $A$=5, with the caveat that the network setting parameters may differ from those used in (Yacoby et al.,, 2020). The results clearly prove that the proposed PH-VAE is more effective in identifying the cluster points with high accuracy, which emphasizes the superiority of the proposed method in handling the pathology in VAE.

The samples adopted in this experiment originate from the MNIST dataset, ten series of benchmark dataset comprising grayscale images of handwritten digits. Each image is resized to the dimension of $64\times 64$, and then converted into tensors for processing through the neural network. The selection of $S$ is 5 concerning image reconstruction, so the training dataset is denoted by $\{\bm{x}_{n}^{s}\}$ $(n=1,\cdots,64\times 64;\;s=1,\cdots,5)$. The data pipeline confirms the uniformity of input data while preserving the original features. The batch size is set to 10, ensuring that all 10 images in the subset are loaded simultaneously for both training and testing.

The PH-VAE model contains five encoders and a combined decoder. The encoder employs three convolutional layers to upscale the feature dimensions as $32$, $64$, and $128$ channels, each followed by ReLU activations. However, it introduces five distinct latent representations derived from five polynomial features $\bm{x},\bm{x}^{2},\bm{x}^{3},\bm{x}^{4},\bm{x}^{5}$. The latent vectors are aggregated as $\bm{z}$ before being passed to the decoder. The latent space dimensionality is set to $40$. A reparameterization trick is used to sample from the latent distribution, to enable stochastic differentiation. The decoder reconstructs ten original images from the latent representation through a fully connected layer, followed by three transposed convolutional layers. The convolutional layers downscale the feature maps to restore the spatial resolution as $128$, $64$, and $32$. A final Sigmoid activation ensures the output pixel values are constrained within $[0,1]$.

The reconstruction loss is calculated for the combined latent representation, while the polynomial divergence terms are computed for each input’s latent representation. It can be derived that the expression of the total loss is

$$\mathcal{L}=-\mathbb{E}_{q(\bm{z}|\bm{x})}[\log p(\bm{x}|\bm{z})]+{1\over 5}\sum_{s=1}^{5}\frac{1}{2}\sum\left(1+\log\bm{\sigma}_{s}^{2}-\bm{\mu}_{s}^{2}-\bm{\sigma}_{s}^{2}\right).$$

The PH-VAE model is trained over $200$ epochs using the same learning rate. The training process involves computing the loss function and propagating the gradients through both the encoders and decoder via backpropagation. The Adam optimizer is employed to iteratively update the model parameters, including $\bm{W}^{s},\bm{b}^{s},\bm{W}_{\mu},\bm{b}_{\mu},\bm{W}_{\sigma},\bm{b}_{\sigma},\bm{W}$, and $\bm{b}$. Throughout each epoch, the loss values are systematically logged to facilitate detailed analysis, monitor convergence, and evaluate the model’s performance.

When $A$=1,2, the superior performance of the proposed PH-VAE architecture is demonstrated by the image results presented in in Figure 7(a) and (b), respectively, with each row corresponding to the polynomial features of $\bm{x}$ (VAE), $\bm{x}+\bm{x}^{2}$, $\bm{x}+\bm{x}^{2}+\bm{x}^{3}$, $\bm{x}+\bm{x}^{2}+\bm{x}^{3}+\bm{x}^{4}$, $\bm{x}+\bm{x}^{2}+\bm{x}^{3}+\bm{x}^{4}+\bm{x}^{5}$ and the ground truth images, respectively. It can be concluded that augmenting the number of polynomial features can significantly improve the accuracy of image reconstruction for each value of $A$, demonstrating the effectiveness of PH-VAE in capturing complex data patterns.

In Figure 8, the values of loss function for each polynomial feature are depicted in case of $A$=1,2. As illustrated, when the loss function is stable, the incorporation of additional polynomial features leads to the reduction in the loss function’s value, highlighting the effectiveness of augmenting the feature set in improving model performance. For the black-and-white image identification, the total calculation time is measured to be 22.5145 seconds.

Applying a subset of 10 images from the CelebA dataset (Liu, et al., 2015), resized to $64\times 64$ pixels, PH-VAE leverages five polynomial features of the input dataset, namely, $\{\bm{x}_{n}^{s}\}$ $(n=1,\cdots,64\times 64;s=1,\cdots,5)$, to establish the neural network structure. The DataLoader with a batch size of 10 is used for efficient training. PH-VAE is characterized by five encoders and a decoder. Each encoder consists of four convolutional layers with the dimensions of $64$, $128$, $256$, and $512$, followed by ReLU activation and batch normalization. These layers progressively extract hierarchical features while reducing the spatial dimensions of the input image. The final feature maps are flattened and passed through two fully connected layers to compute the latent mean $\bm{\mu}_{s}$ and log-variance $\log\bm{\sigma}_{s}^{2}$.

Through the reparameterization trick, latent vectors $\bm{z}$ are sampled from a Gaussian distribution defined by these parameters. The decoder reverses this process, starting with a fully connected layer to expand the latent vector into a high-dimensional feature map. It then employs four transposed convolutional layers with dimensions measuring $512$, $256$, $128$, and $64$, each accompanied by ReLU activation and batch normalization in earlier stages, to upsample the feature maps back to the original image dimensions. A sigmoid activation in the final layer ensures the output image has pixel values in the range $[0,1]$.

The reconstruction loss adheres to the same expression as described in section 3.2. During the training period, the encoders process $\{\bm{x}_{n}^{s}\}$ to generate the respective latent distributions. These latent vectors are combined and passed through the shared decoder to reconstruct the input. The optimizer, typically Adam, minimizes the total loss by backpropagating gradients through both the encoder and decoder. Over $200$ epochs, the model learns to encode robust, integrated features from the input variations, improving the fidelity of the reconstructions and the generalization of the latent space.

Figure 9 depicts the image reconstruction results to evidence the effectiveness of the proposed PH-VAE architecture for $A$=1. The first row describes the results of VAE; the second row to the fifth one represent the prediction results of polynomial features $\bm{x}+\bm{x}^{2}$, $\bm{x}+\bm{x}^{2}+\bm{x}^{3}$, $\bm{x}+\bm{x}^{2}+\bm{x}^{3}+\bm{x}^{4}$, and $\bm{x}+\bm{x}^{2}+\bm{x}^{3}+\bm{x}^{4}+\bm{x}^{5}$, respectively, and the last row comprises the ground truth samples. These results highlight that increasing the number of polynomial features can enhance the accuracy of image reconstruction to some extent, thereby underscoring the effectiveness of PH-VAE in capturing intricate data patterns. The total processing time for colorful image identification is recorded as 85.3892 seconds.