Discond-VAE: Disentangling Continuous Factors from the Discrete

We propose a modification to JointVAE [7] whose latent variable is composed of the discrete, public, and private variables. Since the public variable represents generative factors shared by all classes and the private variable represents variation within each class, it is natural to assume that the prior $p(\mathbf{z},\mathbf{w},\mathbf{c})$ factorizes to $p(\mathbf{z})$ and $p(\mathbf{w},\mathbf{c})$.

$$p(\mathbf{z},\mathbf{w},\mathbf{c})=p(\mathbf{z})\cdot p(\mathbf{w},\mathbf{c})=p(\mathbf{z})\cdot p(\mathbf{c})\cdot p(\mathbf{w}\mid\mathbf{c})$$

(6)

Likewise, the variational distribution $q_{\phi}(\mathbf{z},\mathbf{w},\mathbf{c}|\mathbf{x})$ is modeled as the following.

$$\begin{split}q_{\phi}\left(\mathbf{z},\mathbf{w},\mathbf{c}\mid\mathbf{x}\right)=\,\,&q_{\phi}\left(\mathbf{z}\mid\mathbf{x}\right)\cdot q_{\phi}\left(\mathbf{w},\mathbf{c}\mid\mathbf{x}\right)\\ =q_{\phi}&\left(\mathbf{z}\mid\mathbf{x}\right)\cdot\,\,q_{\phi}(\mathbf{c}\mid\mathbf{x})\cdot q_{\phi}\left(\mathbf{w}\mid\mathbf{x},\mathbf{c}\right)\end{split}$$

(7)

For our Discond-VAE model, the $\beta$-VAE objective (Eq 4) becomes

$$\mathcal{L}_{\textrm{Cond}}(\theta,\phi)=\mathbb{E}_{q_{\phi}(\mathbf{z},\mathbf{w},\mathbf{c}\mid\mathbf{x})}[\log p_{\theta}(\mathbf{x}\mid\mathbf{z},\mathbf{w},\mathbf{c})]\\ -\beta\,D_{KL}(q_{\phi}(\mathbf{z},\mathbf{w},\mathbf{c}\mid\mathbf{x})\mid\mid p(\mathbf{z},\mathbf{w},\mathbf{c}))$$

(8)

The former log-likelihood term stands for the reconstruction error as in the previous VAE models. The latter KL divergence regularizer can be decomposed by the independence assumption.

$$\begin{split}&D_{KL}(q_{\phi}(\mathbf{z},\mathbf{w},\mathbf{c}\mid\mathbf{x})\mid\mid p(\mathbf{z},\mathbf{w},\mathbf{c})\\ &=D_{KL}(q_{\phi}(\mathbf{z}\mid\mathbf{x})\mid\mid p(\mathbf{z}))+D_{KL}(q_{\phi}(\mathbf{w},\mathbf{c}\mid\mathbf{x})\mid\mid p(\mathbf{w},\mathbf{c}))\end{split}$$

(9)

Then, we can address the latter KL divergence as the following. (See appendix for proof)

$$\begin{split}D_{KL}(q_{\phi}(\mathbf{w},\mathbf{c}&\mid\mathbf{x})\mid\mid p(\mathbf{w},\mathbf{c}))\\ =\,\mathbb{E}_{q_{\phi}(\mathbf{c}\mid\mathbf{x})}&\left[D_{KL}(q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c})\mid\mid p(\mathbf{w}\mid\mathbf{c}))\right]\\ &+D_{KL}(q_{\phi}(\mathbf{c}\mid\mathbf{x})\mid\mid p(\mathbf{c}))\end{split}$$

(10)

In brief, the learning objective of Discond-VAE (Eq 8) is expressed as

$$\begin{split}\max_{\theta,\phi}\,&\mathcal{L}_{\textrm{Cond}}(\theta,\phi)=\mathbb{E}_{q_{\phi}(\mathbf{z},\mathbf{w},\mathbf{c}\mid\mathbf{x})}[\log p_{\theta}(\mathbf{x}\mid\mathbf{z},\mathbf{w},\mathbf{c})]\\ &-\beta_{\mathbf{z}}\cdot D_{KL}(q_{\phi}(\mathbf{z}\mid\mathbf{x})\mid\mid p(\mathbf{z}))\\ &-\beta_{\mathbf{w}}\cdot\mathbb{E}_{q_{\phi}(\mathbf{c}\mid\mathbf{x})}\left[D_{KL}(q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c})\mid\mid p(\mathbf{w}\mid\mathbf{c}))\right]\\ &-\beta_{\mathbf{c}}\cdot D_{KL}(q_{\phi}(\mathbf{c}\mid\mathbf{x})\mid\mid p(\mathbf{c}))\end{split}$$

(11)

Discond-VAE models the $q_{\phi}\left(\mathbf{z}|\mathbf{x}\right),q_{\phi}(\mathbf{c}|\mathbf{x})$ by the factorized Gaussian and Gumbel-Softmax as in the JointVAE [7]. Moreover, Discond-VAE introduces the Gaussian Mixture encoder to model the joint distribution of the private and discrete variables. Each mode of the Mixture represents the generative factors within a class.

	$\displaystyle p(\mathbf{w}\mid\mathbf{c})$	$\displaystyle=\prod_{i}p(\mathbf{w}\mid\mathbf{c}=e_{i})^{c_{i}}$		(12)
		$\displaystyle=\prod_{i}\mathcal{N}\left(\mu_{i},I\right)^{c_{i}}$		(13)
	$\displaystyle q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c})$	$\displaystyle=\prod_{i}q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c}=e_{i})^{c_{i}}$		(14)
		$\displaystyle=\prod_{i}\mathcal{N}\left(\mu_{i}(\mathbf{x},\mathbf{c}),\Sigma_{i}(\mathbf{x},\mathbf{c})\right)^{c_{i}}$		(15)

where $\mathbf{c}=(c_{1},c_{2},\cdots,c_{d})\in\{0,1\}^{d}$ denote a one-hot sample from the $d$-dimensional categorical distribution and $e_{i}$ denote a one-hot vector with $i$th component is one. Then, the KL divergence term of private variable $\mathbf{w}$ becomes

$$\begin{split}&\mathbb{E}_{q_{\phi}(\mathbf{c}\mid\mathbf{x})}\left[D_{KL}(q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c})\mid\mid p(\mathbf{w}\mid\mathbf{c}))\right]\\ =&\sum_{i=1}^{d}\alpha_{i}\cdot D_{KL}(q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c}=e_{i})\mid\mid p(\mathbf{w}\mid\mathbf{c}=e_{i}))\end{split}$$

(16)

where $q_{\phi}(\mathbf{c}|\mathbf{x})=(\alpha_{1},\alpha_{2},\cdots,\alpha_{d})$ denotes the variational distribution of the discrete variable. Eq 16 shows that the Discond-VAE encourages the disentanglement of private variables by regularizing the KL divergence to the prior by each mode.

We propose two methods to implement the probabilistic model of Discond-VAE. These two methods differ in how they encode and reparametrize the private variable $\mathbf{w}$.

First, we can model the posterior distribution $q_{\phi}(\mathbf{w}|\mathbf{x},\mathbf{c})$ while keeping the discreteness of the categorical variables. For each class $\mathbf{c}=e_{i}$, the private variable encoder $q_{\phi}(\mathbf{w}|\mathbf{x},\mathbf{c})$ takes a concatenation of features extracted from the data $\mathbf{x}$ and one-hot encoding of the class $e_{i}$ to infer the corresponding mode of the Gaussian Mixture.

$$q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c}=e_{i})=\mathcal{N}\left(\mu_{i}(\mathbf{x},e_{i}),\Sigma_{i}(\mathbf{x},e_{i})\right)$$

(17)

Note that the private variable encoder infers $d$-times where $d$ denotes the number of classes. For the reparametrization trick, this method takes a sample from the mode of the most likely class in $q_{\phi}(\mathbf{c}|\mathbf{x})$.

$$\mathbf{w}=\mathbf{w}_{j}$$

(18)

where $j=\arg\max_{i}q_{\phi}(\mathbf{c}=e_{i}|\mathbf{x})$ and $\mathbf{w}_{j}\sim q_{\phi}(\mathbf{w}|\mathbf{x},\mathbf{c}=e_{j})$. (Fig 2) We refer to this model by Discond-VAE-exact. Note that for the Discond-VAE-exact model with the perfect classification, the continuous variable encoder is a combination of the vanilla encoder applied to the entire dataset and the class-specific vanilla encoder applied only to the corresponding class.

Instead, the private variable encoder $q_{\phi}(\mathbf{w}|\mathbf{x},\mathbf{c})$ can take the data $\mathbf{x}$ and the discrete variable $q_{\phi}(\mathbf{c}|\mathbf{x})=\boldsymbol{\alpha}$ to encode the private variable. We refer to this model by Discond-VAE-approx.

$$q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c}=e_{i})=\mathcal{N}\left(\mu_{i}(\mathbf{x},\boldsymbol{\alpha}),\Sigma_{i}(\mathbf{x},\boldsymbol{\alpha})\right)$$

(19)

The Discond-VAE-approx model adopts a continuous approximation to the sampling from Gaussian Mixture $q_{\phi}(\mathbf{w},\mathbf{c}|\mathbf{x})$ by taking a linear combination of the samples from each mode $q_{\phi}(\mathbf{w}|\mathbf{x},\mathbf{c}=e_{i})$. (Fig 2)

$$\mathbf{w}=\sum_{i}^{d}\pi_{i}\cdot\mathbf{w}_{i}$$

(20)

where $\mathbf{w}_{i}\sim q_{\phi}(\mathbf{w}|\mathbf{x},\mathbf{c}=e_{i})$ and $\boldsymbol{\pi}=(\pi_{1},\cdots,\pi_{d})$ denotes a sample from the Gumbel-Softmax distribution. As in the Discond-VAE-exact, the Discond-VAE-approx with the perfect classification of confidence $100\%$ has a continuous variable encoder equivalent to a combination of the public vanilla encoder and the class-specific vanilla encoder.

Two Discond-VAE implementations are optimized with the capacity objective [4] as in JointVAE to prevent a discrete variable from posterior collapsing. Hence, the learning objective of Discond-VAE becomes

$$\begin{split}&\max_{\theta,\phi}\mathcal{L}_{\textrm{Cond}}(\theta,\phi)=\mathbb{E}_{q_{\phi}(\mathbf{z},\mathbf{w},\mathbf{c}\mid\mathbf{x})}[\log p_{\theta}(\mathbf{x}\mid\mathbf{z},\mathbf{w},\mathbf{c})]\\ &-\beta_{\mathbf{z}}\left|D_{KL}(q_{\phi}(\mathbf{z}\mid\mathbf{x})\mid\mid p(\mathbf{z}))-C_{\mathbf{z}}\right|\\ &-\beta_{\mathbf{w}}\left|\mathbb{E}_{q_{\phi}(\mathbf{c}\mid\mathbf{x})}\left[D_{KL}(q_{\phi}(\mathbf{w}\mid\mathbf{x},\mathbf{c})\mid\mid p(\mathbf{w}\mid\mathbf{c}))\right]-C_{\mathbf{w}}\right|\\ &-\beta_{\mathbf{c}}\left|D_{KL}(q_{\phi}(\mathbf{c}\mid\mathbf{x})\mid\mid p(\mathbf{c}))-C_{\mathbf{c}}\right|\end{split}$$

(21)

The mean vector of each mode in Mixture prior $\mu_{i}$ are hyperparameters. The Discond-VAE-approx model sets the $\mu_{i}$ by the random samples from the standard Gaussian $\mathcal{N}(0,1)$. Interestingly, the Discond-VAE-exact models show similar performance for the $\mu_{i}=0$ with a smaller variance. Therefore, we set the $\mu_{i}=0$ for the Discond-VAE-exact models in Sec 4. We tried the EM algorithm or Warm-up approach to optimize the $\mu_{i}$. However, these approaches showed inferior performance and often unstable training dynamics. (The details of these approaches are provided in the appendix.)

The Discond-VAE-exact model achieves much higher disentanglement scores in both metrics than the JointVAE on the CondSprites in Table 2. The Discond-VAE-approx model shows the higher FactorVAE metric and comparable MIG to the JointVAE. The disentanglement results demonstrate that the Discond-VAE can disentangle the private variables. Furthermore, we report the classification accuracy on the CondSprites in Table 3. The Discond-VAE outperforms the JointVAE by a significant margin on the CondSprites. By disentangling the private variables from the discrete variable, the Discond-VAE can attain a more disentangled discrete representation, which is proved by a higher classification accuracy in Table 3.

The dSprites is a synthetic dataset created with five independent generative factors. Hence, the probabilistic model of the JointVAE is a more suitable assumption for representing the dSprites compared to that of Discond-VAE. Nevertheless, the Discond-VAE-approx models show a comparable classification accuracy in Table 4, and a similar disentanglement metric of the continuous variables in Table 2.

For the case of the two-dimensional public variable, both of the Discond-VAE models show a relatively low FactorVAE score compared to the other models of the same type. In fact, the two-dimensional public variable is insufficient to model four independent continuous factors of dSprites. Nevertheless, the FactorVAE scores are higher than the theoretical limit with two public variables only. This result implies that Discond-VAE can adapt the private variables to represent public generative factors. Considering the CondSprites and dSprites results, the Discond-VAE-exact disentangles the private factors better but has less flexibility to adapt the private variables.

Table 5 shows the unsupervised classification accuracy and NLL scores of each model on the MNIST. The NLL of each model is evaluated by the standard importance weighted sampling strategy [3] with 100 samples. Both types of Discond-VAE show a similar or better accuracy and NLL scores compared to the JointVAE while representing the additional private variables. The result indicates that introducing the private variable to learn the intra-class variation provides a capability to disentangle the discrete factors and learn the data distribution better. Both types of Discond-VAE models with two-dimensional public variable show lower NLL scores compared to the other cases. We suspect this is because a two-dimensional public variable is insufficient to model the major public variations of MNIST such as Angle and Thickness in Fig 3.

We observed the latent traversals of the Discond-VAE to evaluate the disentanglement property qualitatively. For the continuous latent variables, each row corresponds to the latent traversals of an axis over a given example. For the discrete variable, each row shows the one-hot traversal of the discrete variable.

The Discond-VAE shows a smooth variation in angle, slant, thickness, and width of the decoded images as we traverse the public variable in Fig 3. In the discrete variable traversal (Fig 4(c)), we can observe a transition in digit-type of given examples. These results demonstrate that the Discond-VAE can disentangle the public and discrete generative factors from the MNIST.

Moreover, the Discond-VAE discovers the class-specific variation of the digit-type 2 and 7. The Fig 4(a) and 4(b) shows the private variable traversal of the Discond-VAE. Each private variable in Fig 4(a) and 4(b) represents the ring of digit-type 2 and the center-stroke of digit-type 7, respectively. Since these two variations are exclusive to each class, the latent traversals of these two variables show relatively minor or irrelevant variations to the other class images.

To further test the generalizability of Discond-VAE, we observed the latent traversals on the CelebA. Fig 5 shows that Discond-VAE can disentangle the public and private generative factors on the more challenging domain of CelebA. The top two rows and the bottom two rows in Fig 5 represent latent traversals of two examples from the two different classes. The traversal of the public variable (Fig 5(a)) generates variation in brightness for the both classes. For the private variable(Fig 5(b)), the Gender traversal is observed for the class 1, but such change is not observed in the class 0. (See appendix for the more public traversals representing the background and face-width.)