Regularizing Variational Autoencoder with Diversity and Uncertainty Awareness

For enabling meaningful and compact representation learning in VAE model, we have two intuitions: 1) for different data samples $x_{1},x_{2}$, the posteriors $q(z_{1}|x_{1})$ and $q(z_{2}|x_{2})$ should mutually diversify from each other, which encourages posteriors to capture the characteristic or discriminative information from data; 2) given data sample $x$, the degree of uncertainty of the latent variable $z$ should be minimized, which encourages removing redundant information from $z$. Guided by those intuitions, we first analyze the diversity and uncertainty of latent space under quantitative metric, respectively.

In order to encourage high diversity and low uncertainty of latent space, we need to increase the MPD in Equation 3 and decrease the CE in Equation 5, simultaneously. Meanwhile, we also need to avoid $E_{p_{\mathcal{D}}(x)}[\delta^{2}_{x,d}]$ to be too small for ensuring the smoothing of the latent space. One extreme case is that when $E_{p_{\mathcal{D}}(x)}[\delta^{2}_{x,d}]$ convergence to 0, i.e., $\delta^{2}_{x,d}\approx 0,\forall x,d$, each data point is associated with a delta distribution in latent space and the VAEs degenerate into Autoencoders in this dimension. To accomplish these requirements together, we propose to apply Dropout Srivastava et al. (2014) to regularize posterior’s variance parameters in training as following,

where $g_{x,d}$ denotes the independent random variable generated from the normalized Bernoulli distribution $1/pB(1,p),~{}p\in(0,1)$, where $E_{B}[g_{x,d}]=1$. Then, we have the following proposition (see Supplementary for the proof.):

Proposition 8 tells us that: 1) Dropout regularization encourages the increase of $MPD_{p_{\mathcal{D}}(x)}[z]$ and the decrease of the conditional entropy $H_{q_{\phi}}(z|x)$ of the latent space while preserving the expectation of variance parameters, which is actually a simple but useful strategy what we need. 2) Dropout regularization also provides one lower bound of $MPD_{p_{\mathcal{D}}(x)}[z]$ independent on the variance parameters, which makes it possible to ensure positive MPD with further controls on the variance $\sum_{d=1}^{n}Var_{p_{\mathcal{D}}(x)}[\mu_{x,d}]$.

Inspired by Batch-Normalization (BN) Ioffe and Szegedy (2015), which is an effective approach to control the distribution of the output of neural network layer. We apply BN on the mean parameters $\mu_{x,d}$ to constrain $\sum_{d=1}^{n}Var_{p_{\mathcal{D}}(x)}[\mu_{x,d}]$. Mathematically, our BN is defined as:

where $\hat{\mu}_{x,d}$ represents the output of BN layer, and $\mu_{\mathcal{B}_{d}}$ and $\delta_{\mathcal{B}_{d}}$ denote the mean and standard deviation of $\mu_{x,d}$ estimated within each mini-batch. $\gamma_{\mu_{d}}$ and $\beta_{\mu_{d}}$ are the scale and shift parameters, which lead that the distribution of $\hat{\mu}_{x,d}$ has the variance $\gamma^{2}_{\mu_{d}}$ and mean $\beta_{\mu_{d}}$. Therefore, we can control the $\sum_{d=1}^{n}Var_{p_{\mathcal{D}}(x)}[\mu_{x,d}]$ by fixing the mean $E_{d}[\gamma^{2}_{\mu_{d}}]=\gamma^{2}$ with respect to each dimension $d$. Specifically, we regard each $\gamma_{d}$ as learnable parameters with initialization $\gamma$. Then after each training iteration, we re-scale each parameter $\gamma_{\mu_{d}}$ with coefficient $\gamma/\sqrt{E_{d}[\gamma^{2}_{\mu_{d}}]}$. In addition, all $\beta_{\mu_{d}}$ is learnable with initialization 0 and no constraint.

Overall, based on the analysis above, we propose our approach, namely DU-VAE, to encourage high diversity and low uncertainty of the latent space by applying Dropout regularizations on variance parameters and Batch-Normalization on mean parameters of approximated posteriors, simultaneously. Specifically, the training of DU-VAE is following Algorithm 1.

Here, to further examine the generalization of DU-VAE, we aim to extend our approach for other VAE variants, such as, VAE-IAF Kingma et al. (2016), one well-known normalizing flow-based VAE. Different from classic VAEs which assume the posterior distributions are diagonal Gaussian distributions, VAE-IAF can construct more flexible posterior distributions through applying one chain of invertible transformations, named the IAF chain, on an initial random variable drown from one diagonal Gaussian distribution. Specifically, the initial random variable $z^{0}$ is sampled from the diagonal Gaussian with parameters $\mu^{0}$ and $\delta^{0}$ outputted from the encoder network. Then, $T$ invertible transformations, are applied to transform $z^{0}$ into the final random variable $z^{T}$. More details can be found in Kingma et al. (2016).

Indeed, noting that the MPD and CE of the initial random variable $z^{0}$ have the same form as these for classic VAEs in Equation 2 and Equation 4, one intuitive idea is to apply Dropout on $\delta^{0}$ and Batch Normalization on $\mu^{0}$ with the guidance in Algorithm 1 to control the MPD and CE of $z^{0}$. It is surprising to find that this simple extension of DU-VAE, called DU-IAF, demonstrated comparative performance in our experiments. This may be attributed to the close connection between $z^{0}$ and $z^{T}$. In particular, we find that the CE of $z^{0}$ is the upper bound of CE of $z^{T}$. Meanwhile, $MPD_{p_{\mathcal{D}}(x)}[z^{0}]$ is closely related with $MPD_{p_{\mathcal{D}}(x)}[z^{T}]$, even they are equal to each other when each invertible transformation in IAF chain is independent on the input data. Further discussion and proof can be found in Supplementary.

Specifically, for text datasets, following Shen et al. (2017), we work with one downsampled version of Yelp sentiment dataset for binary classification. Table 2 shows the performance under varying number of labeled data. For the image dataset, noting that OMNIGLOT contains 1623 different handwritten characters from 50 different alphabets, where each character has 15 images in our training data and 5 images in our testing data, we conducted classifications on each alphabet with varying number of training samples for each character. Table 3 reports the average accuracy.

We can find that DU-VAE and DU-IAF achieve the best accuracy under all settings for classic VAEs and IAF-based models respectively. Interestingly, we also find that most baselines show inconsistent results on text and image classification. For example, Agg-VAE and BN-VAE may be better at text classification without sound accuracy in Table 3. On the contrary, $\delta$-VAE and MAE adapt to image classification better with uncompetitive performance in Table 2. Meanwhile, we note IAF chain trends to improve the classification accuracy for FB and our approach to both text and image datasets. However, IAF+BN fails to achieve competitive performance on image classification, which indicates that the applications of BN in BN-VAE may not be suitable for image again.

Here, we aim to provide one intuitive comparison of latent spaces learned by different models based on one simple synthetic dataset. Specifically, following Kim et al. (2018), we first sample 2-dimensional latent variable $z$ from one mixture of Gaussian distributions that have 5 mixture components. Then one text-like dataset can be generated from one LSTM layer conditioned on those latent variables. Based on this synthetic dataset, we trained different VAEs with 2-dimensional standard Gaussian prior and diagonal Gaussian posterior. Then, we visualize the learned latent spaces by displaying the counter plot of the aggregated approximated posteriors $q(z)=E_{p_{\mathcal{D}}(x)}[q_{\phi}(z|x)]$ and the location of approximated posterior’s mean parameters for different samples $x$.

According to the results in Figure 1, we have some interesting observations. First, due to the posterior collapse, VAE learns an almost meaningless latent space where the posterior $q(z|x)$ for all data are squeezed in the center. Actually, it is not surprising that the aggregated posterior matches the prior excessively in this case, because we almost have $q_{\phi}(z|x)=p(z)$, $\forall x$. Second, Agg-VAE, BN-VAE, and DU-VAE all tend to diverse samples in different categories, but in different manners and degrees. Intuitively, all three models force to embedding the blue category in the center around by the other four categories. However, only the average posterior learned by DU-VAE have five centers same as the true latent space. Meanwhile, DU-VAE with Dropout strategy encourages the aggregated posteriors to be more compact, while that of BN-VAE is more broad, compared with the prior. Those observations demonstrate that DU-VAE tends to guide the latent space to be more diverse and less uncertain.