Blocked and Hierarchical Disentangled Representation From Information Theory Perspective

IB theory aims to learn a representation $Z$ that maximizes the compression of informaiton in real data $X$ while maximizing the expression of target $Y$. So we can describe it as:

$\displaystyle\min I(X;Z)-\beta I(Z;Y)$

(1)

the target $Y$ is the attribute information under supervision, and is equal to $X$ under unsupervision [36].

In the case of supervised IB theory [2], we can get the upper bound:

	$\displaystyle I_{\phi}(X;Z)-\beta I_{\theta}(Z;Y)\leq$	$\displaystyle\mathbb{E}_{p_{D}(x)}[D_{KL}(q_{\phi}(z|x)\|p(z))]$
		$\displaystyle-\beta\mathbb{E}_{p(x,y)}[q_{\phi}(z|x)\log p_{\theta}(y|z)]$		(2)

The first term represents the KL divergence between the posterior $q_{\phi}(z|x)$ and the prior distribution $p(z)$; and absolutely, the second term equals cross-entropy loss of label prediction.

And in the case of unsupevised IB theory, the we can rewrite the objective Eq. (1) as:

$\displaystyle\min I_{\phi}(X;Z)-\beta I_{\theta}(Z;X)$

(3)

Unsupervised IB theory seems like generalization of VAEs model, with an encoder to learn representation and a decoder to reconstruct. $\beta$-VAE [12] is actually the upper bound of it:

	$\displaystyle\mathcal{L}_{\beta-VAE}=$	$\displaystyle\mathbb{E}_{p(x)}[D_{KL}(q_{\phi}(z|x)\|p(z))$
		$\displaystyle-\beta\mathbb{E}_{q_{\phi}(z|x)}[\log(p_{\theta}(x|z))]]$		(4)

FactorVAE [14] and $\beta$-TCVAE [7] just add more weight on the TC term $\mathbb{E}_{q(z)}[\log\frac{q(z)}{\tilde{q}(z)}]$, which express the dependence across dimensions of variable in information theory, where $\tilde{q}(z)=\prod_{i=1}^{n}q(z_{i})$.

We build our BHiVAE model upon above works and models. We focus on information transmission and loss through the whole network, and then achieve it through different methods.

Now let us present our detailed model architecture. As shown in Fig 1, feed data $X$ into the encoder (parameterized as $\phi$), and in the first layer, we get the latent representation $z^{1}$, be divided into two parts $s^{1}$ and $h^{1}$. The part $s^{1}$ is the final representation part, which corresponds to feature $y^{1}$, and $h^{1}$ is the input of next layer’s encoder to get latent representation $z^{2}$. Then through three similar network processes, we can get three representation parts $s^{1},s^{2},s^{3}$, which are disentangled, and get the part $c^{3}$ in the last layer, that contains information other than the above attributes of the data. All of them make up the whole representation $z=(s^{1};s^{2};s^{3};c^{3})$. The representation of each part is then mapped to the same space by a different decoder (all parameterized as $\theta$) and finally concatenated together to reconstruct the raw data, which is shown in Fig 1(b). For the problem we discussed, we need to get the final disentangled representation $z$, i.e., we need the independence between each representation part $s^{1},s^{2},s^{3}$, and $c^{3}$.

Then we can separate the whole problem into two sub-problem in $i$-th layer, so the input is $h^{i-1}$(where $h^{0}=x$):

• (1)

  Information flow $h^{i-1}\rightarrow s^{i}\rightarrow y^{i}$: Encode the upper layer’s output $h^{i-1}$ to representation $z^{i}$, with one part $s^{i}$ containing sufficient information about one feature factor $y^{i}$;

• (2)

  Information separation of $s^{i}$ and $h^{i}$: Eliminate the information about $s^{i}$ in $h^{i}$ while requiring $s^{i}$ only to contain label $y^{i}$ information.

The first subproblem can be regarded as IB problem, the goal is to learn a representation of $s^{i}$, i.e. maximally expressive about feature $y^{i}$ while minimally informative about input data $h^{i-1}$. So it can described as:

$\displaystyle\min I(h^{i-1};s^{i})-\beta I(s^{i};y^{i})$

(5)

To satisfy the second subproblem is a complex issue, and it requires different methods to achieve it with different known conditions. So we will introduce these in follow conditions in detail. In summary, our representation is designed to enhance the internal correlation of each block while reducing the relationships between them to achieve the desired disentanglement goal.

When training BHiVAE model, we need the encoder and decoder (Fig 1) both in supervised and unsupervised cases. On the CelabA dataset, we build our network with both a convolutional layer and a fully connected layer. On the MNIST and dSprites datasets, the datasets are both $64\times 64$ binary images, so we design our network to consist entirely of fully connected layers.

In evaluating the experimental results, we use the Z-differ [12], SAP [16], and MIG [7] metrics to measure the quality of the disentangled representation, and observe the images generated by the traversal representation. Moreover, we use some pre-trained classifiers on attribute features to analyze the model according to the classification accuracy.

In the unsupervised case, as introduced in the previous section, the most significant novel idea is we use a different prior $\mathcal{N}(0,\Sigma)$ to guide the representation learning. Additionally, we need another one to estimate the KL divergence (18). Therefore, two extra discriminators are needed for BHiVAE in Fig 2(a). Actually, because we aim to get $D_{KL}(q_{\phi}(z^{i})\|p(z))=0$, the latent representation $\{z_{j}^{i}\}_{j=1}^{N}$ can be considered as generated from true distribution, while prior and permuted ’presentations’ $\{z^{i-perm}_{j}\}_{j=1}^{N}$ can both be considered as false. Therefore, we can simplify the network to contain only one discriminator to score these three distributions.

We want to reinforce the relationship within $s^{i}$ to retain the information and then decrease the dependency between $s^{i}$ and $h^{i}$ to separate information, so in our unsupervised experiments, we use this prior $\mathcal{N}(0,\Sigma)$, where

$\displaystyle\scriptsize\Sigma=\begin{bmatrix}1&0.5&0&\cdots&0\\ 0.5&1&0&\cdots&0\\ 0&0&1&\cdots&0\\ \vdots&\vdots&\vdots&\ddots&\vdots\\ 0&0&0&\cdots&1\end{bmatrix}$

First, we use some experiments to demonstrate the feasibility and validity of this prior setting. We train the model on the dSprites dataset first, with setting the dimension of representation $z$ to 14 ($d(z)=14$), where $d(s^{i})=2,i=1,2,3$ and $d(c^{3})=8$. Then we get a representation in each layer of the 1000 test images, while the three subfigures in Fig 3 shows a scatter plot of each layer representation respectively, and the curves in these figures both are the contour of the block target distribution $p(s)\sim\mathcal{N}(0,\begin{bmatrix}1&0.5\\ 0.5&1\end{bmatrix})$. And it is shown in Fig 3 that in the first and second layer, the distribution of $s^{1}$ and $s^{2}$ do not sufficiently match the prior $p(s)$, but as the layer forward, the KL divergence between $q_{\phi}(s^{i})$ and $p(s)$ keep decresing, and the scatter plot of $s^{i}$ fits the prior distribution more closely. In the model, we train the encoder globally, so the front layer’s representation learning can be influenced by the change of deeper representation and then yields larger KL divergence than the next layer.

Even more surprisingly, in Fig 3(c), we find that in the third layer, visualizing the ’Shape’ attribute of dSprites dataset, there is an apparent clustering effect (the different colors denote different categories). This result proves our hypothesis about the deep network’s ability: the deeper network is, the more detailed information it extracts. And it almost matches the prior perfectly. Fig 3(c) also gives us a better traversal way. In previous works, because only one dimension represents the attribute, they can simply change the representation from $a$ to $b$ ($a$ and $b$ both are constant). However, this does not fit our model, so the direction of the category transformation in Fig 3(c) inspires us to traverse the data along the diagonal line ($y=x$). Our block prior $p(s)$ also supports that (because the prior distribution’s major axis is the diagonal line too).

We perform several experiments under above architecture setting and traversal way to show the disentanglement quality on MNIST datasets. The disentanglement quantitative results of comparing with $\beta$-VAE [12], FactorVAE [14] and Guided-VAE [10] are presented in Fig 4. Here, considering the dimension of the representation and the number of parameters, other works’ bottleneck size is set to 12, i.e., $d(z)=12$. This setting helps reduce the impact of differences in complexity between model frameworks. However, for a better comparison, we only select seven dimensions that change more regularly. In our model, we change the three-block representation $\{s^{i}\}_{i=1}^{3}$ and then the rest representation $c^{3}$ changes according to two dimensions as a whole, i.e., $c^{3}=(c^{3}_{1:2},c^{3}_{3:4},c^{3}_{5:6},c^{3}_{7:8})$. And Fig 4 shows that $\beta$-VAE hardly ever gets a worthwhile disentangled representation, but FactorVAE appears to attribute change as representation varies. Moreover, Fig 4(c) and Fig 4(d) both show great disentangled images, with $h_{1}$ changing in Guided-VAE and $s^{1}$ changing in BHiVAE, the handwriting is getting thicker, and $h_{3},s^{2}$ control the angle of inclination. These all demonstrate the model capabilities of our model.

We then progress to the traversal experiments on the dSprites dataset. This dataset has clear attributes distinctions, and these allow us to better observe the disentangled representation. In these experiments, BHiVAE learns a 10-dimensional representation $z=(s^{1},s^{2},s^{3},c^{3}_{1:2},c^{3}_{3,4})$ and 8-dimensional $z=(z_{1},z_{2},\dots,z_{8})$ in other works. We present the experiments results in Fig 5 of reconstruction and traversal results. The first and second rows in four figure represent original and reconstruction images respectively. In Fig 5(d), it shows that our first three variables $s^{1},s^{2},s^{3}$ have learned the attribute characteristics (Scale, Orientation, and Position) of the data.

Moreover, we perform two quantitive experiments comparing with previous works and present our results in Table 1 and Table 2. The experiments are all based on the same experiment setting in Fig 4.

First, we compare BHiVAE with previous models with Z-differ Score [12], SAP Score [16] and MIG Score [7] and present the results in Table 1. It is clear that our model BHiVAE is at the top and that the MIG metric is better than other popular models. The high value of the Z-diff score indicates that learned disentangled representation has less variance on the attributes of generated data as corresponding dimension changing, while SAP measures the degree of coupling between data factors and representations. Additionally MIG metric uses mutual information to measure the correlation between the data factor and learned disentangled representation, and our work is just modeled from the perspective of mutual information, which makes us performs best on the MIG score.

Not only that, but we also perform transferability experiments by conducting classification tasks on the generated representation. Here we set the representation dimensions to be the same in all models. First, we have learned a pre-trained model to obtain the representation $z$ and a pre-trained classifier to predict MNIST image label from representation. We compare the classification accuracy in Table 2 with different dimension settings.

Our model appears not higher accuracy than FactorVAE and Guided-VAE in the case of $d_{z}=10$. That block representation setting causes a small number of factors it learns. However, as $d(z)$ is increased, our representation can learn more attribute factors of data, and then the classification accuracy can also be improved.

In the supervised case, we still did the qualitative and quantitative experiments to evaluate our model. The same as the unsupervised case, overall autoencoder is required, and then we need a classifier to satisfy the segmentation of information at each level, as shown in Fig 2(b). And we set the dimension of representation $z$ as 12 ($d(z)=12,d(c^{3})=6$, and $d(s^{i})=2,i=1,2,3$).

We first perform several experiments comparing with Guided-VAE [10] in two attributes(Gender and Black Hair) and present the results in Fig 6. When changing each attribute $s^{i}\in\{s^{1},s^{2},s^{3}\}$, we keep other attributes representations and content representation $c^{3}$ unchanged. We use the third layer representation $s^{3}$ to control gender attribute, while the first layers correspond to the black hair and bale, respectively. In the supervised case, compared to Guided-VAE, we use multiple dimensions to control an attribute while Guided-VAE uses only one dimension, which may lead to insufficient information to control the traversal results. And Fig 6 shows that our model has a broader range of control over attributes, especially reflected in the range of hair from pure white to pure black.

Besides, our quantitative experiment is to first pre-train the BHiVAE model and three attribute classifiers of the representation and then get the representationS of the training set, traversing the three representation blocks a,b,c from $(-3,3)$ to $(3,3)$ along with the diagonal($y=x$). Fig 7 shows that all three attributes have a transformation threshold in the corresponding representation blocks.

In the previous experiments, we are all making judgments about how well the representation is disentangled and did not prove that the block setting is beneficial, so we set up the following comparison experiments for this problem.

For the comparison experiment here, we set the dimension of the model representation $z$ to 10, 16, and 32. Then in the comparison experiment, we just changed the dimension of representation $s^{1}$ (black hair) in the first layer to 1, and therefore the dimension of $c^{3}$ is changed to 5, 11, and 27 accordingly. First we pre-train these two models under the same conditions and learn a binary classifier that predicts the black hair attributes with representation $z$. It is shown in Fig 8 that Block is better than Single in every dimension setting, and the accuracy of them has increased with increasing representation dimension. It could be that there is still some information about black hair in other representation parts of the model, and then the increasing dimension will allow more information about black hair to be preserved, getting better prediction accuracy.