Encoding Binary Concepts in the Latent Space of Generative Models for Enhancing Data Representation

Given a dataset $\mathcal{D}=\{x_{i}\}_{i=1}^{n}$, VAEs learn an encoder $p_{\phi}(z|x)$ and a decoder $p_{\theta}(x|z)$ that maximize the likelihood $p_{\phi,\theta}(x)$ to approximate the empirical evidence $p_{\mathcal{D}}(x)$. The encoder maps $p_{\mathcal{D}}(x)$ to a simple latent prior $p(z)$ which has better sampling properties. The objective is to maximize:

where the right hand side is ELBO, served as the proxy objective we maximize to learn the latent distribution or as the objective in the following minimization problem:

where the first term is the negative likelihood of data after passing through the model. It is modeled as the reconstruction loss in VAEs. The second term is the KL-divergence between the approximate posterior and the latent prior.

The two terms in Eq. (1) conflict with each other. The first term trades disentanglement for reconstruction quality, emphasizing it results in low interpretability and control of latent representation. The second term trades reconstruction quality for disentanglement, emphasizing that it results in low variances of high-level concepts that lead to a high correlation with latent dimensions and high noise levels of low-level concepts that lead to posterior collapse. We develop a method to to mitigate these challenges.

We define a new binarized objective that instead of minimizing the distributional distance between the approximate posterior and the prior minimizes the closest distributional distance between the approximate posterior and two binarized priors which are located $r$ away from the mean of the prior. The new binarized loss can be written as:

where $r\in[0,+\infty]$ is a model hyperparameter. It decides how disentangled the two candidate distributions are (see Figure 1). Note that the loss reduces to the original VAE loss when $r=0$. A more general form can be written as:

for any loss function $\mathcal{L}(\mu_{p(z)}-\mu_{p_{\phi}(z|x)})$, where $\mu_{p(z)}$ is a location variable of a prior distribution, and $\mu_{p_{\phi}(z|x)}$ is a location variable of an approximate posterior distribution.

Gaussian Prior Closed-form Loss

To add the binarized modification to a unimodal Gaussian prior $p(z)\sim\mathcal{N}(\mu,\sigma^{2})$, we have the closed form:

We can show that this is equivalent to ELBO with a binarized $z$-dependent prior as the following:

This equivalent prior can be visualized as in Figure 3. For more details on derivation please refer to the Appendix.

We still use the original reparametrization trick since we still consider the original prior once the estimated posterior is decided. This assumption maintains the continuous property of the approximate posterior at the prior mean which is broken by the deterministic assignment. The probabilistic reparametrization and the deterministic assignment create an instability close to the prior mean, which can alleviate posterior collapse. The source of the instability is the opposing force provided by the Bernoulli-distribution-like encoding. It will have maximum information capacity when posterior collapse happens which discourages collapse.

The instability at the prior mean will force the latent to deviate from the high probability density center. This allows more samples from the underrepresented outer region of high-variance factors, and fewer samples from the biased center high-density region, achieving reduced-bias sampling. Although this will negatively impact low-variance factors, in practice, the combined effect achieves better total sampling quality. This might be due to the alleviated posterior collapse compensating for this negative impact.

Reduced-bias Sampling: For the two binarized distributions, we assign an equal probability to sample from:

which is approximately sampling from the Gaussian:

when $r$ is small. In all of our experiments, this sampling method is used due to its simplicity and nice properties.

Biased Sampling: If we want to sample from the distribution when the input dataset is imbalanced, we restore the distribution of the dataset by conditioning the sampling prior on the latent statistics of the full dataset:

This sampling would for a skewed distribution.

First, we use toyVAE as the base model to compare the performance of binarized VAE with VAE on the MNIST dataset. Reconstruction is evaluated using image-sum pixel-wise binary cross-entropy (BCE). Unsupervised clustering is evaluated by normalized mutual information (NMI). The result is presented in Table 1. We observe that binarized VAE is able to achieve better reconstruction and clustering performances with different latent emebdding sizes.

Next, we scale the model to ConvVAE. We test VAE, $\beta$-VAE, disentangling $\beta$-VAE, and their binarized modifications (6 models) on the CelebA dataset (32x32 face-aligned). We again observe that the binarized models showed lower reconstruction loss (Figures 4 and 5), lower variational loss, and higher ELBO compared to the base models.

We then further scale up the base model to L-ConvVAE. We test VAE, $\beta$-VAE, disentangling $\beta$-VAE, DIP-VAE, WAE and their binarized modifications on the CelebA (64x64 face-aligned) dataset, and visualize the reconstruction and sampling results in Figure 6. We use the same sampling latent space between the original and binarized models.

For reconstruction, the binarized versions have sharper edges and more detailed facial features in general and are more accurate in reconstructing those underrepresented categories in CelebA dataset such as old, facial hair, wearing a cap, etc. (Figure 7).

For sampling, we observe that low-level properties are similar between the samples, while high-level properties are well-distinctive and more variant in their binarized versions. This observation aligns well with our hypothesis. An interesting observation is that when binarizing WAE, the binarized model is able to generate samples with significantly more realistic facial features.

We use ConvVAE as the base model to visualize the latent space for MNIST and CelebA(32x32, 64x64 face-aligned) reduced-size(60000 instances) datasets. We tested VAE, $\beta$-VAE, disentangling $\beta$-VAE, and the binarized models. All the base models suffer from posterior collapse due to the adjusted smaller dataset compared to the large capacity of the ConvVAE decoder. We see that the binarized models are able to maintain the collapsed latent dimensions (see Figures 8 and 10), by introducing instability at the prior mean. concept (see Figure 8).

Then we look at what exact concepts are learned in the high-variance and low-variance latent dimensions. We see more variations of low-level binary concepts are learned by the binarized models in the high-variance dimensions (see Figure 9).

While high-level concepts are learned in the low-variance dimensions restored by the binarized models. While these concepts are unlearnable by the original VAE models (see Figure 10).

Generative replay requires the fit-sample-replay loop [25, 20, 19], which highly depends on the generated quality of the samples. First, we tested the compatibility of the Gaussian mixture prior to binarized regularization, since it is commonly used for conditional replay and proves to be effective. We test a GMVAE with a binarized GMVAE on MNIST (see Figure 11). We have compared state-of-the-art CL methods that use generative replay to mitigate catastrophic forgetting. We observe that binarized GMVAE can better disentangle the classes with fewer ambiguities.

Then we test the binarized generative replay models on the most challenging incremental-class learning scenario. We performed experiments using generative replay(GR) and brain-inspired generative replay (BIR). We conclude that with our modification, the forgetting effect of early tasks can be generally reduced, under single modal (Gaussian) and multimodal (GMM) latent representation, with or without conditional replays (see Table 2).

In Figure 12, we observe that the variance in earlier tasks is maintained better, which results in less forgetting of earlier tasks which is an important challenge. The worse performance in more recent tasks and better performance on average suggests that the improvement is not from the ability to reconstruct better samples in general, but from maintaining more diverse knowledge in the model. This observation indicates that we have achieved learning more general and robust factors in the latent space.

We study the effect of major hyperparameters on the quality of reconstrcuted samples.

Visualization for the latent space(dimension = 4) of VAEs and the binarized versions. Symmetry is shown in the binarized VAEs and posterior collapse does not appear in binarized disentangling $\beta$-VAE. We use MNIST training dataset.

In row 1, the binarized latent space is more symmetric. The correlation between two latent dimensions is also reduced. The data points are spread out in a larger area without requiring additional training resources.

In row 2, the binarized latent space does not suffer from posterior collapse.

We test $r=0,0.01,0.1,1$ settings on the CelebA dataset. Reconstruction loss is reduced for binarized models $r=0.01,0.1$. The running time is also slightly reduced(Table 3).

Sampling is conducted on the same models as in the last section. We use the same latent space vector to generate the image in all four models. We can see that when $r$ is small, the samples have more distinct facial features, sharper edges, and vibrant colors. When $r$ is large, the samples sampled from the original prior are from the “ambiguous region”, since the sampling prior is between the binarized centers. And the generated images are ambiguous and lack distinction as expected.

Optimal approximate posterior estimation for VAE with a Gaussian prior $p_{\phi}(z|x)\sim\mathcal{N}(\mu,\sigma^{2})$, assuming some error $\epsilon_{\mu}$ and $\epsilon_{\sigma}$

For binarized VAE:

Difference:

When $x\approx\mu$ we have $(x-\mu-\epsilon_{\mu})^{2}\approx{\epsilon_{\mu}}^{2}<<(\sigma+\epsilon_{\sigma})^{2}$, when $\sigma$ is sufficiently large, then approximate by first-order Taylor expansion:

Thus, we have

This allows the binarized approximate posterior to be less distributed around the mean. When sampled using the same prior distribution, reduced bias is thus achieved. (Also see Figure 19. The red curve denotes high-variance posterior. Column 1 is a scenario when the estimated posterior has a larger or smaller variance than the prior. Column 2 is the difference between the posteriors and the prior. Columns 3, and 4 are binarized models.)