ISA-VAE: Independent Subspace Analysis with Variational Autoencoders

1 INTRODUCTION . Recently there has been an increased interest in unsupervised learning of disentangled representations . The term disentangled usually describes two main objectives : First , to identify each true factor of variation with a latent variable , and second , interpretability of these latent factors ( Schmidhuber , 1992 ; Ridgeway , 2016 ; Achille & Soatto , 2017 ) . Most of this recent work is inspired by the β-VAE concept introduced in Higgins et al . ( 2016 ) , which proposes to re-weight the terms in the evidence lower bound ( ELBO ) objective . In Higgins et al . ( 2016 ) a higher weight for the Kullback-Leibler divergence ( KL ) between approximate posterior and prior is proposed , and putative mechanistic explanations for the effects of this modification are studied in Burgess et al . ( 2017 ) ; Chen et al . ( 2018 ) . An alternative decomposition of the ELBO leads to the recent variant of β-VAE called βTCVAE ( Chen et al. , 2018 ) , which shows the highest scores on recent disentanglement benchmarks . These modifications of the evidence lower bound however lead to a trade-off between disentanglement and reconstruction loss and therefore the quality of the learned model . This trade-off is directly encoded in the modified objective : by increasing the β-weight of the KL-term , the relative weight of the reconstruction loss term is more and more decreased . Therefore , optimization of the modified ELBO will lead to latent encodings which have a lower KL-divergence from the prior , but at the same time lead to a higher reconstruction loss . Furthermore , we discuss in section 2.4 that using a higher weight for the KL-term amplifies existing biases of variational inference , potentially to a catastrophic extent . There is a foundational contradiction in many approaches to disentangling deep generative models ( DGMs ) : the standard model employed is not identifiable as it employs a standard normal prior which then undergoes a linear transformation . Any rotation of the latent space can be absorbed into the linear transform and is therefore statistically indistinguishable . If interpretability is desired , the modelling choices are setting us up to fail . We make the following contributions : • We show that the current state of the art approaches employ a trade-off between reconstruction loss and disentanglement of the latent representation . • In section 2.3 we show that variational inference techniques are biased : the estimated components are biased towards having orthogonal effects on the data and the number of components is underestimated . • We provide a novel description of the origin of disentanglement in β-VAE and demonstrate in section 2.4 that increasing the weight of the KL term increases the over-pruning bias of variational inference . • To mitigate these drawbacks of existing approaches , we propose a family of rotationally asymmetric distributions for the latent prior , which removes the rotational ambiguity from the model . This approach resembles independent component analysis ( ICA ) for variational autoencoders . • We propose to use a prior which allows a decomposition of the latent space using independent subspace analysis ( ISA ) and demonstrate that this prior leads to disentangled representations even for the unmodified ELBO objective . This removes the trade-off between disentanglement and reconstruction loss of existing approaches . • An even higher disentanglement of the latent space can be achieved by incorporating the proposed prior distribution into the existing approaches β-VAE and β-TCVAE . Since the prior distribution already favours a disentangled representation , the new method dominates previous in terms of the trade-off between disentanglement and model quality . 2 BACKGROUND . We briefly discuss previous work on variational inference in deep generative models and two modifications of the learning objective that have been proposed to learn a disentangled representation . We discuss characteristic biases of variational inference and how the modifications of the learning objective actually accentuate these biases . 2.1 DISENTANGLED REPRESENTATION LEARNING . Variational Autoencoder Kingma & Welling ( 2014 ) introduce a latent variable model that combines a generative model , the decoder , with an inference network , the encoder . Training is performed by optimizing the evidence lower bound ( ELBO ) averaged over the empirical distribution : LELBO = Eqφ ( z|x ) [ log pθ ( x|z ) ] −DKL ( qφ ( z|x ) ‖p ( z ) ) , ( 1 ) where the decoder pθ ( x|z ) is a deep learning model with parameters θ and each zl is sampled from the encoder zl ∼ qφ ( z|x ) with variational parameters φ . When choosing appropriate families of distributions , gradients through the samples zl can be estimated using the reparameterization trick . The approximate posterior qφ ( z|x ) is usually modelled as a multivariate Gaussian with diagonal covariance matrix and the prior p ( z ) is typically the standard normal distribution . β-VAE Higgins et al . ( 2016 ) propose to modify the evidence lower bound objective and penalize the KL-divergence of the ELBO : Lβ-ELBO = Eqφ ( z|x ) [ log pθ ( x|z ) ] − βDKL ( qφ ( z|x ) ‖p ( z ) ) , ( 2 ) where β > 1 is a free parameter that should encourage a disentangled representation . In Burgess et al . ( 2017 ) the authors provide further thoughts on the mechanism that leads to these disentangled representations . However we will show in the following that this parameter introduces a trade-off between reconstruction loss and disentanglement . Furthermore , we show in section 2.4 that this parameter amplifies biases of variational inference towards orthogonalization and pruning . β-TCVAE In Chen et al . ( 2018 ) the authors propose an alternative decomposition of the ELBO , that leads to the recent variant of β-VAE called β-TCVAE . They demonstrate that β-TCVAE allows to learn representations with higher MIG score than β-VAE ( Higgins et al. , 2016 ) , InfoGAN ( Chen et al. , 2016 ) and FactorVAE ( Kim & Mnih , 2018 ) . The authors propose to decompose the KL-term in the ELBO objective into three parts and to weight them independently : Epθ ( x ) [ DKL ( qφ ( z|x ) ‖p ( z ) ) ] = = DKL ( qφ ( z|x ) ‖qφ ( z ) pθ ( x ) ) +DKL ( qφ ( z ) ‖ ∏ j qφ ( zj ) ) + ∑ j DKL ( qφ ( zj ) ‖p ( zj ) ) . ( 3 ) The first term is the index-code mutual information , the second term is the total correlation and the third term the dimension-wise KL-divergence . Because the index-code mutual information can be viewed as an estimator for the mutual information between pθ ( x ) and qφ ( z ) , the authors propose to exclude this term when reweighting the KL-term with the β weight . In addition to the improved objective , the authors propose a quantitative evaluation score for disentanglement , the mutual information gap ( MIG ) . They propose to first estimate the mutual information between a latent factor and an underlying generative factor of the dataset . The mutual information gap is then defined as the difference of the mutual information between the highest and second highest correlated underlying factor . 2.2 RELATED WORK . Recent work has shown an increased interest into learning of interpretable representations . In addition to the work mentioned already , we briefly review some of the influential papers : Chen et al . ( 2016 ) present a variant of a GAN that encourages an interpretable latent representation by maximizing the mutual information between the observation and a small subset of latent variables . The approach relies on optimizing a lower bound of the intractable mutual information . Kim & Mnih ( 2018 ) propose a learning objective equivalent to β-TCVAE , and train it with the density ratio trick ( Sugiyama et al. , 2012 ) . Kumar et al . ( 2017 ) introduce a regularizer of the KL-divergence between the approximate posterior and the prior distribution . A parallel line of research proposes not to train a perfect generative model but instead to find a simpler representation of the data ( Vedantam et al. , 2017 ; Hinton et al. , 2011b ) . A similar strategy is followed in semi-supervised approaches that require implicit or explicit knowledge about the true underlying factors of the data ( Kulkarni et al. , 2015 ; Kingma et al. , 2014 ; Reed et al. , 2014 ; Baydin et al. , 2017 ; Hinton et al. , 2011a ; Zhu et al. , 2017 ; Goroshin et al. , 2015 ; Hsu et al. , 2017 ; Denton et al. , 2017 ) . 2.3 ORTHOGONALIZATION AND PRUNING IN VARIATIONAL INFERENCE . There have been several interpretations of the behaviour of the β-VAE ( Chen et al. , 2018 ; Burgess et al. , 2017 ) . Here we provide a complementary perspective : that it enhances well known statistical biases in VI ( Turner & Sahani , 2011 ) to produce disentangled , but not necessarily useful , representations . The form of these biases can be understood by considering the variational objective when written as an explicit lower-bound : the log-likelihood of the parameters minus the KL divergence between the approximate posterior and the true posterior LELBO = log pθ ( x ) −DKL ( qφ ( z|x ) ‖pθ ( z|x ) ) ( 4 ) From this form it is clear that VI ’ s estimates of the parameters θ will be biased away from the maximum likelihood solution ( the maximizer of the first term ) in a direction that reduces the KL between the approximate and true posteriors . When factorized approximating distributions are used , VI will therefore be biased towards settings of the parameters that reduce the statistical dependence between the latent variables in the posterior . For example , this will bias learned components towards orthogonal directions in the output space as this reduces explaining away ( e.g . in the factor analysis model , VI breaks the degeneracy of the maximum-likelihood solution finding the orthogonal PCA directions , see appendix B.8 ) . Moreover , these biases often cause components to be pruned out ( in the sense that they have no effect on the observed variables ) since then their posterior sits at the prior , which is typically factorized ( e.g . in an over-complete factor analysis model VI prunes out components to return a complete model , see appendix B.8 ) . For simple linear models these effects are not pathological : indeed VI is arguably selecting from amongst the degenerate maximum likelihood solutions in a sensible way . However , for more complex models the biases are more severe : often the true posterior of the underlying model has significant dependencies ( e.g . due to explaining away ) and the biases can prevent the discovery of some components . For example , VAEs are known to over-prune ( Burda et al. , 2015 ; Cremer et al. , 2018 ) . 2.4 β-VAE EMPHASIZES ORTHOGONALIZATION AND PRUNING What happens to these biases in the β-VAE generalization when β > 1 ? The short answer is that they grow . This can be understood by considering coordinate ascent of the modified objective . With θ fixed , optimising q finds a solution that is closer to the prior distribution than VI due to the upweighting of the KL term in 2 . With q fixed , optimization over θ returns the same solution as VI ( since the prior does not depend on the parameters θ and so the value of β is irrelevant ) . However , since q is now closer to the prior than before , the KL bias in equation 2 will be greater . These effects are shown in the ICA example in appendix B.8 . VI ( β = 1 ) learns components that are more orthogonal than the underlying ones , but β = 5 prunes out one component entirely and sets the other two to be orthogonal . This is disentangled , but arguably leads to incorrect interpretation of the data . This happens even though both methods are initialised at the true model . Arguably , the β-VAE is enhancing a statistical bug in VI and leveraging this as a feature . We believe that this can be dangerous , preventing the discovery of the underlying model .

Summary: This paper presents a methodology to bring together independent subspace analysis and variational auto-encoders. Naturally, in order to do that, the authors propose a specific family of prior distributions that lead to subspace independence the Lp-nested distribution family. This prior distribution is then used to learn disentangled and interpretable representations. The mutual information gap is taken as the measure of disentanglement, while the reconstruction loss measures the quality of the representation. Experiments on the sPrites dataset are reported, and comparison with the state of the art shows some interesting results.