EXoN: EXplainable encoder Network

Let ${\bf x}\in\mathcal{X}\subset\mathbb{R}^{m}$ and ${\bf z}\in\mathcal{Z}\subset\mathbb{R}^{d}$ with $d<m$ be the observed data and the latent variable. $p({\bf x})$ and $p({\bf z})$ denote the probability density functions (pdf) of ${\bf x}$ and ${\bf z}$. Let $p({\bf x}|{\bf z})$ be the conditional pdf of ${\bf x}$ for a given ${\bf z}$ and ${\bf x}|{\bf z}\sim N(D({\bf z}),\beta\cdot{\bf I}),$ where $D:\mathcal{Z}\mapsto\mathcal{X}$, ${\bf I}$ is the $m\times m$ identity matrix and the observation noise $\beta>0$. $D({\bf z})$ is a neural network with parameter $\theta$. To emphasize dependence on parameters, denote $D({\bf z})$ and $p({\bf x}|{\bf z})$ by $D({\bf z};\theta)$ and $p({\bf x}|{\bf z};\theta,\beta)$, respectively. $p({\bf x}|{\bf z};\theta,\beta)$ is referred to as the decoder of VAE. The term of decoder is also used as the map ${\bf z}\mapsto(D({\bf z};\theta),\beta)$, because $p({\bf x}|{\bf z};\theta,\beta)$ is fully parameterized by $(D({\bf z};\theta),\beta)$.

The parameter of VAE $(\theta,\beta)$ is estimated by Variational method (Kingma and Welling 2013; Rezende, Mohamed, and Wierstra 2014) that maximizes ELBO. The ELBO of $\log p({\bf x};\theta,\beta)$ is obtained from the inequality

			$\displaystyle\log p({\bf x};\theta,\beta)$
		$\displaystyle\geq$	$\displaystyle\mathbb{E}_{q({\bf z}|{\bf x})}[\log p({\bf x}|{\bf z};\theta,\beta)]-\mathcal{KL}\left(q({\bf z}|{\bf x})\|p({\bf z})\right),$

where $\mathcal{KL}(q\|p)$ denotes Kullback-Leibler divergence from $p$ to $q$. (Kingma and Welling 2013; Rezende, Mohamed, and Wierstra 2014) employ a neural network with parameter $\phi$ as $q({\bf z}|{\bf x};\phi)$, and obtain the objective function under finite samples $x_{1},\cdots,x_{n}$,

$\displaystyle\sum_{i=1}^{n}\mathbb{E}_{q}[\log p(x_{i}|{\bf z};\theta,\beta)]-\mathcal{KL}\left(q({\bf z}|x_{i};\phi)\|p({\bf z})\right).$

(2)

Here, $q({\bf z}|{\bf x};\phi)$ is referred to as the encoder of VAE or the posterior distribution over the latent variables. Multivariate Gaussian distribution is a popular choice for $q({\bf z}|{\bf x};\phi)$ in which the mean and covariance are given by neural network models. Thus, the parameters of VAE consist of $\phi$ in the encoder and $(\theta,\beta)$ in the decoder.

In practice, the VAE is fitted by maximizing a stochastically approximated ELBO for (2) with respect to $(\theta,\beta,\phi)$. The approximated ELBO is given by

	$\displaystyle\sum_{i=1}^{n}-\frac{1}{2}\|x_{i}-D(z_{i};\theta)\|^{2}-\beta\cdot\mathcal{KL}(q({\bf z}|x_{i};\phi)\|p({\bf z}))$
	$\displaystyle-\frac{nm}{2}\log 2\pi\beta$		(3)

where $z_{i}$ is a sample from $q({\bf z}|x_{i};\phi)$ and $\|\cdot\|$ is $l_{2}$-norm. The VAE is fitted by maximizing (2.1) with respect to $(\theta,\beta,\phi)$ (Lucas et al. 2019). The first term in (2.1) is the precision of generated samples. Since the KL-divergence is always non-negative, (Higgins et al. 2016; Mathieu et al. 2019; Li et al. 2019) explain $\beta$ as a tuning parameter regularizing $\theta$ and $\phi$, and the last term of (2.1) is considered as constant.

The prior latent distribution for each label is assumed as ${\bf z}|{\bf y}=k\sim N(m_{k},diag\{s_{k}^{2}\})$ and $p({\bf y}=k)=w_{k}$ for $k\in\mathcal{Y}$ where $diag\{a\},a\in\mathbb{R}^{d}$ indicates $d\times d$ matrix whose diagonal elements are $a$. The prior distribution marginalized by ${\bf y}$ is

$\displaystyle p({\bf z})=\sum_{k=1}^{K}w_{k}\cdot\mathcal{N}({\bf z}|m_{k},diag\{s_{k}^{2}\}).$

Here, $m_{k}\in\mathbb{R}^{d}$ and $s_{k}^{2}\in\mathbb{R}_{+}^{d}$ are pre-determined parameters denoting the conceptual center and dispersion of the latent variable for each label. The choice of $m_{k}$ and $s_{k}^{2}$ is the customization of the latent space. Since $m_{k}$ and $w_{k}$ for all $k$ are fixed, the notation of $m_{k}$ and $diag\{s_{k}^{2}\}$ is omitted in $p({\bf z})$.

The encoder $p({\bf z}|{\bf x})$ is assumed as a mixture distribution, $p({\bf z}|{\bf x})=\sum_{k=1}^{K}p({\bf y}=k|{\bf x})\cdot p({\bf z}|{\bf y}=k,{\bf x})$. Since using $p({\bf z}|{\bf x})$ is computationally prohibitive, the proposal distribution $w({\bf y}=k|{\bf x};\eta)$ for $p({\bf y}=k|{\bf x})$ and $g({\bf z}|{\bf y}=k,{\bf x};\xi)=\mathcal{N}({\bf z}|\mu_{k}({\bf x};\xi),diag\{\sigma^{2}_{k}({\bf x};\xi)\})$ for $p({\bf z}|{\bf y}=k,{\bf x})$ are introduced. The two proposal distributions are multinomial and Gaussian, whose parameters are modeled by neural networks. The posterior distribution is approximated by

$\displaystyle q({\bf z}|{\bf x};\eta,\xi)=\sum_{k=1}^{K}w({\bf y}=k|{\bf x};\eta)\cdot g({\bf z}|{\bf y}=k,{\bf x};\xi),$

(4)

where the parameters of the neural networks are $\eta$ and $\xi$. Note that the posterior mixture weight $w({\bf y}=k|{\bf x};\eta)$ decompose the latent space according to labels. Given image ${\bf x}$ and label ${\bf y}$, the posterior mixture weight assigns a latent variable to the mixture component corresponding to the given label ${\bf y}$ and the latent space is separated by labels.

The decoder is assumed by

$\displaystyle p({\bf x}|{\bf z};\theta,\beta)=\mathcal{N}({\bf x}|D({\bf z};\theta),\beta\cdot{\bf I}),$

where the true label ${\bf y}$ is not included. Even though a natural decoder can be considered as a mixture Gaussian $p({\bf x}|{\bf z};\theta,\beta)=\sum_{k=1}^{K}p({\bf y}=k|{\bf z})\cdot\mathcal{N}({\bf x}|D_{k}({\bf z};\theta),\beta\cdot{\bf I})$, the mixture distribution is approximated by a uni-modal Gaussian distribution especially when $p({\bf y}=k|{\bf z})$ is assumed to be close to 1 for some $k$. The latent space is separated by $w({\bf y}|{\bf x};\eta)$ in our encoder and we can obtain such $p({\bf y}|{\bf z})$. Thus, the validity of the assumption mainly depends on the classification performance of $w({\bf y}|{\bf x};\eta)$. Although the decoder variance $\beta$ is trainable, we fixed it due to computational issues.

Our proposed VAE model is derived from the joint pdf $p({\bf x},{\bf y})$. $p({\bf x},{\bf y})$ is decomposed into $p({\bf x})$ and $p({\bf y}|{\bf x})$, and the derivation of the ELBO only for $\log p({\bf x})$ leads to (5). In the derivation, $p({\bf y}|{\bf x})$ is approximated by the proposal classification model $w({\bf y}|{\bf x};\eta)$.

			$\displaystyle\log p({\bf x},{\bf y};\theta,\beta)$		(5)
		$\displaystyle=$	$\displaystyle\log p({\bf x};\theta,\beta)+\log p({\bf y}|{\bf x};\theta,\beta)$
		$\displaystyle\geq$	$\displaystyle\mathbb{E}_{q}[\log p({\bf x}|{\bf z};\theta,\beta)]-\mathcal{KL}\left(q({\bf z}|{\bf x};\eta,\xi)\|p({\bf z})\right)$
			$\displaystyle+\log p({\bf y}|{\bf x};\theta,\beta)$
		$\displaystyle\simeq$	$\displaystyle\mathbb{E}_{q}[\log p({\bf x}|{\bf z};\theta,\beta)]-\mathcal{KL}\left(q({\bf z}|{\bf x};\eta,\xi)\|p({\bf z})\right)$
			$\displaystyle+\log w({\bf y}|{\bf x};\eta).$

The first two terms in (5) are the typical ELBO used in the unsupervised VAE learning, and the remained term is the classification loss. Note that these terms are coupled with the shared parameter $\eta$. The classification loss term plays the role of training the posterior mixture weights. Therefore, $w({\bf y}|{\bf x};\eta)$ with lower classification error can separate the latent space more clearly by $q({\bf z}|{\bf x};\eta,\xi)$.

Whenever $p({\bf y}|{\bf x};\theta,\beta)=w({\bf y}|{\bf x};\eta)$,

$\displaystyle\log p({\bf x},{\bf y};\theta,\beta)\geq\mathcal{L}({\bf x};\theta,\beta,\xi,\eta)+\log w({\bf y}|{\bf x};\eta),$

where

	$\displaystyle\mathcal{L}({\bf x};\theta,\beta,\xi,\eta)$	$\displaystyle=$	$\displaystyle\mathbb{E}_{q}[\log p({\bf x}|{\bf z};\theta,\beta)]$
			$\displaystyle-\mathcal{KL}^{u}\left(q({\bf z}|{\bf x};\eta,\xi)\|p({\bf z})\right)$

and $\mathcal{KL}^{u}\left(q({\bf z}|{\bf x};\eta,\xi)\|p({\bf z})\right)$ is the upper bound of $\mathcal{KL}\left(q({\bf z}|{\bf x};\eta,\xi)\|p({\bf z})\right)$ (Wang et al. 2019; Guo et al. 2020) which is written in closed-form as

			$\displaystyle\mathcal{KL}(w({\bf y}|{\bf x};\eta)\|p({\bf y}))$
		$\displaystyle+$	$\displaystyle\sum_{k=1}^{K}w({\bf y}=k|{\bf x};\eta)$
			$\displaystyle\times\mathcal{KL}\left(q({\bf z}|{\bf x},{\bf y}=k;\eta,\xi)\|p({\bf z}|{\bf y}=k)\right).$

Therefore, the lower bound on the joint likelihood for the entire dataset is

$\displaystyle\sum_{i\in I_{L}\cup I_{U}}\mathcal{L}(x_{i};\theta,\beta,\xi,\eta)+\sum_{i\in I_{L}}\log w(y_{i}|x_{i};\eta).$

(6)

(Kingma et al. 2014) first employed the classification loss (the second term in (6)) as a penalty function of VAE and (Maaløe et al. 2016; Li et al. 2019; Siddharth et al. 2017) use the same penalty function in the subsequent papers. In our semi-supervised VAE model, the penalty function is applied with a similar idea. However, the derivation of the regularized objective function stems from (5) unlike the existing studies.

The derivation of (6) is mathematically plausible, but it does not guarantee state-of-the-art semi-supervised classification performance. To improve the performance of the classification task, we propose a new loss function for the pseudo-labeling semi-supervised classification method (Iscen et al. 2019; Berthelot et al. 2019; Arazo et al. 2020) called SCI (Soft-label Consistency Interpolation) loss.

The SCI loss consists of three parts, 1) an interpolated new image with a pair of unlabeled images, 2) a pair of pseudo-label of the unlabeled images, and 3) a convex combination of the cross entropy. Let $({\bf x}^{1},{\bf x}^{2})$ be a pair of images and let $\tilde{{\bf x}}=\rho\cdot{\bf x}^{1}+(1-\rho)\cdot{\bf x}^{2}$ be the interpolated image. Denote the discrete probability $w({\bf y}|{\bf x};\hat{\eta})$ by $f({\bf x};\hat{\eta})$. The pseudo-labels of ${\bf x}^{1}$ and ${\bf x}^{2}$ are defined by $\tilde{{\bf y}}^{1}=f({\bf x}_{1};\hat{\eta})$ and $\tilde{{\bf y}}^{2}=f({\bf x}_{2};\hat{\eta})$ for a given estimate $\hat{\eta}$ (true label is used for labeled dataset). Note that the pseudo-label is not a function of $\hat{\eta}$ but a non-trainable quantity determined by $\hat{\eta}$. Then the SCI loss for $\eta$ with $({\bf x}^{1},{\bf x}^{2})$ is defined by

			$\displaystyle\mbox{SCI}(({\bf x}^{1},\tilde{{\bf y}}^{1}),({\bf x}^{2},\tilde{{\bf y}}^{2});\eta)$
		$\displaystyle=$	$\displaystyle\rho\cdot H(\tilde{{\bf y}}^{1},f(\tilde{{\bf x}},\eta))+(1-\rho)\cdot H(\tilde{{\bf y}}^{2},f(\tilde{{\bf x}},\eta)),$

where $H(p,q)$ is the cross entropy of the distribution $q$ relative to a distribution $p$.

The SCI loss is motivated by the assumption of consistency interpolation.

The consistency interpolation assumes the existence of a linear map $f$ from an image to a pseudo-label. It is well known that such a mix-up strategy can improve the generalization error (Zhang et al. 2017).

Interestingly, the estimation of $f(\cdot;\eta^{*})$ in our VAE model is also used in other existing semi-supervised learning methods (Feng et al. 2021; Arazo et al. 2020) and the following algorithm provides a general framework to estimate such a $f(\cdot;\eta^{*})$ in the VAE model. Let $\eta^{(t)}$ be an estimate of $\eta$ obtained by the $t$th step while training the VAE, and $\eta^{(t+1)}$ be the solution of following optimal interpolation problem (Feng et al. 2021)

	$\displaystyle\min_{\eta}$		$\displaystyle\rho\cdot\mathcal{KL}\left(f({\bf x}^{1};\eta^{(t)})\|f(\tilde{{\bf x}};\eta)\right)$		(8)
		$\displaystyle+$	$\displaystyle(1-\rho)\cdot\mathcal{KL}\left(f({\bf x}^{2};\eta^{(t)})\|f(\tilde{{\bf x}};\eta)\right).$

Then, it is easily shown that

			$\displaystyle f(\rho\cdot{\bf x}^{1}+(1-\rho)\cdot{\bf x}^{2};\eta^{(t+1)})$		(9)
		$\displaystyle=$	$\displaystyle\rho\cdot f({\bf x}^{1};\eta^{(t)})+(1-\rho)\cdot f({\bf x}^{2};\eta^{(t)}).$

If there exists $\eta^{*}$ such that $\eta^{(t)}\rightarrow\eta^{*}$ as $t\rightarrow\infty$ given two data point ${\bf x}^{1},{\bf x}^{2}$, then (9) finally implies the consistency interpolation. Since (8) is equivalent with (3.3) up to constant, (3.3) is introduced in our proposed objective function. We also found that using stochastically augmented images for SCI loss helps our classifier achieve higher accuracy (see Appendix A.4 for detailed augmentation techniques).

Finally, the objective function is given by

		$\displaystyle-$	$\displaystyle\sum_{i\in I_{L}\cup I_{U}}\mathcal{L}(x_{i};\theta,\beta,\xi,\eta)-\sum_{i\in I_{L}}\log w(y_{i}|x_{i};\eta)$		(10)
		$\displaystyle-$	$\displaystyle\frac{\lambda_{1}}{\beta}\sum_{i\in I_{L}}\log w(y_{i}|x_{i};\eta)$
		$\displaystyle+$	$\displaystyle\frac{\lambda_{1}}{\beta}\sum_{i\in I_{L}}\mbox{SCI}((x^{1}_{i},y^{1}_{i}),(x^{2}_{i},y^{2}_{i});\eta)$
		$\displaystyle+$	$\displaystyle\frac{\lambda_{2}(t)}{\beta}\sum_{i\in I_{U}}\mbox{SCI}((x^{1}_{i},\tilde{y}^{1}_{i}),(x^{2}_{i},\tilde{y}^{2}_{i});\eta),$

where $x_{i}^{1}=x_{i}$, $x_{i}^{2}$ is a randomly chosen sample, and $\lambda_{2}(t)$ is a ramp-up function. If $\lambda_{1}$ and $\lambda_{2}(t)$ are set to be large enough, the mixture components of $q({\bf z}|{\bf x};\eta,\xi)$ are shrunk toward those of $p({\bf z})$ according to the true matching labels, and the latent space is well separated. The shared parameter $\eta$ leads to this adaption of $q({\bf z}|{\bf x};\eta,\xi)$ to $p({\bf z})$.

In this section, we investigated the meaning of the explainability of the latent space. Denote the random vector associated with the distribution of the $k$th component of the mixture distribution (prior or posterior distribution) by ${\bf z}^{k}=({\bf z}_{1}^{k},\cdots,{\bf z}_{d}^{k})$. Then, what is the role of the posterior conditional variance $\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)$ in the encoding process?

To answer above question, consider an extreme case where $\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)$ goes to zero for given $k$. If $\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)\rightarrow 0$, then ${\bf z}_{j}^{k}$ becomes constant. ${\bf z}_{j}^{k}$ has a deterministic relationship with given ${\bf x}$, implying that ${\bf z}_{j}^{k}$ can explain the given data point. Therefore, the value of $\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)$ represents the explainability of the latent space, and we call $j$ coordinate is activated.

Interestingly, there is a theorem which shows a connection between our objective function (10) and the posterior conditional variance $\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)$.

Theorem 1 means that the KL-divergence upper bound in our objective function is bounded by the ratio of the prior and posterior variances. The lower bound of Theorem 1 can be re-written as $\sum_{k=1}^{K}\sum_{j=1}^{d}w_{k}/2\cdot\mbox{Var}_{{\bf x}|{\bf y}=k}[\mathbb{E}({\bf z}_{j}^{k}|{\bf x};\xi)]/\mbox{Var}({\bf z}_{j}^{k})$, so it can be interpreted as the coefficient of determination $R^{2}$ which measures the proportion of latent variable variation explained by an observation.

On the other hand, a small $\beta$ increases the weight of the generated sample precision term and indirectly decreases the relative weight of the KL-divergence term in (10). It implies that the KL-divergence term is relaxed (not fully minimized). So, the relaxation of the KL-divergence term induces an increase in the upper bound of Theorem 1. It means that there exists an activated coordinate $j$ such that $\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)$ shrinks to zero for some $k$, as the prior mixture weight $w_{k}$ and variance $\mbox{Var}({\bf z}_{j}^{k})$ are fixed due to pre-design. Thus, Theorem 1 shows that the relaxation of the KL-divergence obtained by tuning $\beta$ is closely related to controlling the explainability of the latent space for ${\bf x}$.

Based on Theorem 1, we propose a statistics V-nat (VAE-natural unit of information) which screens such activated latent coordinates,

$\displaystyle\log\mathbb{E}_{{\bf x}|{\bf y}=k}\left[\mbox{Var}({\bf z}_{j}^{k})/\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)\right].$

Additionally, we define the set of activated latent coordinates as activated latent subspace for $k$-labeled dataset for some $\delta>0$,

$\displaystyle{\cal A}_{k}(\delta)=\bigg{\{}j\in[d]:\log\mathbb{E}_{{\bf x}|{\bf y}=k}\left[\frac{\mbox{Var}({\bf z}_{j}^{k})}{\mbox{Var}({\bf z}_{j}^{k}|{\bf x};\xi)}\right]>\delta\bigg{\}}$

where $[d]=\{1,\cdots,d\}$ Our numerical study found that the subspace represents the informative characteristics of generated samples, and the subspace can be effectively used to produce a high-quality image (see Section 4.2). These results are consistent with those of VQ-VAE (Van Den Oord, Vinyals et al. 2017) that an image generation process mainly depends on a nearly deterministic encoding map, and the refinement of the map can improve the quality of generated images.

We have used the MNIST dataset (LeCun, Cortes, and Burges 2010) to consider the 2-dimensional latent space in our VAE model. The values were scaled in the range of -1 to 1. The encoder returns the 10-component Gaussian mixture distribution parameters, mixing probability, mean vector, and diagonal covariance elements. Thus, the encoder maps ${\cal X}$ to $\left((0,1)\times\mathbb{R}^{2}\times\mathbb{R}_{+}^{2}\right)^{10}$. Especially, the mixing probabilities are produced by the classifier in the encoder. Gumbel-Max trick (Gumbel 1954; Jang, Gu, and Poole 2016) is used for sampling discrete latent variables. Detailed network architectures of the encoder, decoder and classifier are described in Appendix A.2 Table 4, 5.

The customized $p({\bf z})$ is illustrated in Figure 1. Each label is assigned counterclockwise from 3 O’clock on the component of the mixture. Note that Figure 1 illustrates an example of conceptual centers. The $k$th component of the mixture distribution corresponds to the distribution of the digit $(k-1)$ in the MNIST dataset, for $k=1,\cdots,10$ (see Appendix A.4 for detailed pre-design settings).

Our model is compared with (Kingma et al. 2014; Maaløe et al. 2016; Li et al. 2019; Hajimiri, Lotfi, and Baghshah 2021; Feng et al. 2021), and the simulation results show that our fitted model achieves a competitive classification performance, 3.12% error with 59,900 unlabeled and 100 labeled images. (see Appendix A.2 Table 6 for comparison results). For implementation details, see Appendix A.4.

We apply our VAE model to the CIFAR-10 dataset (Krizhevsky, Hinton et al. 2009) of which images have 10 labels. We evaluated our model with 400 labeled images per class and 46,000 unlabeled images. All values of the images are scaled to the range of $(-1,1)$. 256-dimensional latent space and the 10-mixture Gaussian distribution are employed. Each component of the prior mixture distribution has a separate mean vector, and all components share the same covariance (see Appendix A.4 for detailed pre-designed priors, hyper-parameters, and implementation settings). Gumbel-Max trick (Gumbel 1954; Jang, Gu, and Poole 2016) is used for sampling discrete latent variables. The network architectures of the encoder, decoder and classifier are shown in Appendix A.3 Table 8, 9.