DEFT: Distilling Entangled Factors by Preventing Information Diffusion^††thanks: *Corresponding author

In variational inference, posterior $p(z|x)$ is intractable. The variational autoencoder (VAE) (Kingma and Welling, 2014) uses a neural network $q_{\phi}(z|x)$ (encoder) to approximate the posterior $p(z|x)$. The other neural network $p_{\theta}(x|z)$ (decoder) rebuilds the observations. The objective of the VAE is to optimize the evidence lower bound (ELBO):

Higgins et al. discovered the relationship between the disentanglement and the Kullback–Liebler (KL) divergence penalty strength. They proposed the $\beta$-VAE to introduce additional pressure on the KL term:

$\beta$ controls the pressure for the posterior, $q_{\phi}(z|x)$, to match the factorized unit Gaussian prior, $p(z)$. However, there is a trade-off between the quality of the reconstructed images and disentanglement.

Burgess et al. (2018) proposed the AnnealedVAE, which progressively increases the information capacity of the latent variables while training:

where $\gamma$ is a sufficiently large constant (usually 1000) to constrain the latent information, and $C$ is a value that gradually increases from zero to a large number.

The TC Watanabe (1960) quantifies the dependency among variables. $\beta$-TCVAE (Chen et al., 2018) decomposed the KL term into three parts: mutual information (MI), total correlation (TC), and dimensional-wise KL (DWKL). They proposed that the TC be penalized to achieve high reconstruction quality and disentanglement:

Jeong and Song provided another total correlation penalization through information cascading. They proved that $TC(z)=\sum_{i=2}^{d}I(z_{1:i-1};z_{i})$. CascadeVAEC, the continuous version, sequentially relieves one latent variable at a time, encouraging the model to disentangle one factor during the $i$-th stage:

where $\beta_{l}$ is a small value for opening the information flow, $\beta_{h}$ is a large beta coefficient, and $d$ is the number of dimensions.

Artificial datasets usually have ground-truth labels, such as dSprites (Matthey et al., 2017). The supervised metrics use this label information to measure the disentanglement responsibly. Several metrics have been proposed to evaluate the disentanglement. These include the BetaVAE metric (Higgins et al., 2017a), FactorVAE metric (Kim and Mnih, 2018), MI gap (Chen et al., 2018), modularity (Ridgeway and Mozer, 2018), DCI (Eastwood and Williams, 2018), and SAP score (Kumar et al., 2018). Shannon MI is an information-theoretic quantity that measures the amount of information shared between two variables. The MIG measures the gap between the top two latent variables with the highest MI (Chen et al., 2018). $j^{1}=\operatorname*{arg\,max}_{j}I(z_{j};c_{k})$ is the index of latent variables that maximize the largest MI for $c_{k}$. $j^{2}=\operatorname*{arg\,max}_{j}I(z_{j};c_{k}),j\neq j^{1}$ is the second largest MI. We define the $m$-th largest normalized MI (NMI) between $z_{j}$ and $c_{k}$ as $\text{NMI}(c_{k},m)$:

$\text{NMI}(c_{k},2)$ is the largest $\frac{I(z_{j};c_{k})}{H(c_{k})}$ except $\text{NMI}(c_{k},1)$. The MIG can be calculated as

We trained the model in the first stage with $\beta=70$ and then trained the second stage with $\beta=30$ for different values of $\gamma$. As shown in Fig. 6, all backward information is clipped when $\gamma=0$, and $I(x;z)$ will not increase; $\text{NMI}({\bm{c}},2)$ and $\gamma$ are simultaneously increasing. In contrast, $\text{NMI}({\bm{c}},1)$ has a small increment. A small value of $\gamma$ is sufficient to learn the majority information, and it also prevents that information from diffusing into another variable. In practice, $\gamma=0.1$ achieves a good balance between learning new information and preventing ID.

AnnealedVAE directly controls the capacity of the latent variables by an annealed increasing parameter, $C$. CascadeVAEC increases the capacity by relieving the pressure on the $i$-th latent variable at the $i$-th stage, opening the information flow. We found that ID is an invisible hurdle for the IB-based approaches, which can be detected by the mean of $\text{NMI}({\bm{c}}_{i},1),\text{NMI}({\bm{c}}_{i},2)$:

We conducted experiments involving AnnealedVAE, CascadeVAEC, and DEFT on dSprites. As depicted in Fig. 7 (a), AnnealedVAE expands the capacity by a controllable parameter $C$; however, the KL consists of MI, TC, and DWKL. There is an increasing gap between the MI and $C$. At each stage, CascadeVAEC activates a new latent variable to expand the capacity. DEFT expands the capacity by relieving the pressure. One can see, DEFT has a steady increment of MI. Both AnnealedVAE and CascadeVAEC suffer from the ID problem (see in Fig. 7 (b)): AnnealedVAE tends to diminish NMI1; CascadeVAEC also has a sharp decrement of NMI1. Only DEFT remits the ID problem and promotes the disentanglement steadily.

We compared DEFT with others on dSprites (Matthey et al., 2017), color (color for short), screen (screen for short), and SmallNORB (LeCun et al., 2004). The images of dSprites are strictly generated by the five factors. It has three shapes: square, ellipse, and heart; six scale values: 0.5, 0.6, 0.7, 0.8, 0.9, and 1.0, 40 orientation values in [0, 2 pi], 32 position X values, and 32 position Y values. Two variants of dSprites (color and screen), which introduce random noise, were closer to the true situation. SmallNORB is generated from 3D objects and is much more complex than 2D shapes. It contains five generic categories, namely, four-legged animals, human figures, airplanes, trucks, and cars; nine elevation values, i.e., 30, 35, 40, 45, 50, 55, 60, 65, and 70; eighteen azimuth values, 0, 20, 40,…, 340; and six lighting conditions. Fig. 8 shows the visualization of the datasets.

The ideal situation is to find a set of $\beta$ to isolate IFPs into several parts without overlaps. To obtain the distribution of IFPs with respect to a factor ${\bm{c}}_{i}$, we enumerate all possible values of factor ${\bm{c}}_{i}$ for a random sample, and calculate its ICP using the algorithm introduced in Section 4.1. Then, we repeated the above procedure 50 times to estimate the ICP distribution of ${\bm{c}}_{i}$.

We measured the IFPs of the factors on the four datasets, as shown in Fig. 9. dSprites and Color had more separable IFPs than Scream and SmallNORB. Although the three variants of dSprites have the same factors, their IFPs are different. The difference in IFP distributions explains why current approaches fail to transfer hyperparameters across different problems in Locatello et al. (2019). Note that the IFP distributions of factors are almost separable for dSprites and Color; the ground-truth factors are independent of the four datasets. In summary, they are all independent; Scream and SmallNorb are inseparable. Based on the distribution of IFPs, we summarize the optimal training settings for the DEFT in Table 2. We tune the hyperparameters of compared approaches with the highest MIG and show theses setting in Table 3.

We trained each model 50 times and compared our model with the other four disentanglement approaches on dSprites, Color, Scream, and SmallNORB in Fig. 10. All approaches (not just ours) have a lower performance on Scream and SmallNorb (inseparable). It appears understandable that DEFT fails to handle the situation having inseparable IFP distributions. However, there is a question as to why these TC-based approaches also failed in an independent but inseparable situation (Scream). We also show the distributions of the reconstruction error in Fig. 11. In general, DEFT achieved both a high image quality and disentanglement.

The distribution of MIG scores at different stages are shown in Fig. 12. All experimental results on the four datasets reveal that DEFT obtains low scores in the first stage, and gradually improves disentanglement in the following stages.

We define the failure rate as the percentage of models that fail to learn a disentangled representation if the MIG score is lower than 0.1. Table 4 shows the failure rate. It can be seen that DEFT has the lowest average failure rates. Although AnnealedVAE success to disentangle factors on three datasets, it fails to disentangle factors on Scream for most cases. Note that, it is possible to reduce the failure rate for AnnealedVAE on Scream, but we have tried six settings. Despite all approaches show low performances on SmallNORB, the failure rates are lower than others. From the IFP distributions in Fig. 9, we can see that SmallNORB has a separable factor that is easy to be disentangled. That causes SmallNORB to have a high lower bound of disentanglement but gets a low overall score. Generally, DEFT significantly decreased the failure rate compared to the other approaches.

Higgins et al. (2017a) introduced the latent traversal to visualize the generated images through the traversal of a single latent $z_{i}$. Fig. 13 shows the latent traversal of the best model with the highest MIG score. One can see the intrinsic relationship between IFP and disentanglement. Orientation has the lowerest IFP among all factors; meanwhile, it is the hardest one to be disentangled for all approaches. For SmallNORB, the lighting condition (3) is separable with others, which is easy to be disentnagled. For Scream, three factors have similar IFP distributions, and it is also a hard problem for the disentanglement approaches.

The label information is unavailable for the common situations. Therefore, the factor’s IFP distribution is hard to be obtained. Alternatively, we calculate the upper bound of IFP distribution for the unsupervision setting. Intuitively, the rate of information increment changes if there is a new factor starting to freeze. We conducted an annealing test on dSprites and 3D chairs without labels and plotted the curves of beta vs. $\Delta I(x;z)$ in Fig. 16. This method is in agreement with the upper bound of the IFP distribution for position and scaling, as shown in Fig. 16 (a). One can recognize four points where the latent information suddenly increases: 36 and 16 from Fig. 16 (b). Though this method needs human participation, we only show the potency to develop a fully unsupervised procedure for the separations. Therefore, we set $\text{K=3},\text{G=3},\beta_{j}=\{36,16,1\}$ for 3D Chairs and trained the DEFT 20 epochs per stage. We compared the performance with $\beta$-TCVAE and CascadeVAEC on 3D Chairs, as shown in Fig. 17. We notice that DEFT can learn one additional interpretable property compared with CascadeVAEC: leg orientation.