Unsupervised Model Selection for Variational Disentangled Representation Learning

1 INTRODUCTION . Happy families are all alike ; every unhappy family is unhappy in its own way . — Leo Tolstoy , Anna Karenina Despite the success of deep learning in the recent years ( Hu et al. , 2018 ; Espeholt et al. , 2018 ; Silver et al. , 2018 ; Lample et al. , 2018 ; Hessel et al. , 2017 ; Oord et al. , 2016 ) , the majority of state of the art approaches are still missing many basic yet important properties , such as fairness , data efficient learning , strong generalisation beyond the training data distribution , or the ability to transfer knowledge between tasks ( Lake et al. , 2016 ; Garnelo et al. , 2016 ; Marcus , 2018 ) . The idea that a good representation can help with such shortcomings is not new , and recently a number of papers have demonstrated that models with disentangled representations show improvements in terms of these shortcomings ( Higgins et al. , 2017b ; 2018b ; Achille et al. , 2018 ; Steenbrugge et al. , 2018 ; Nair et al. , 2018 ; Laversanne-Finot et al. , 2018 ; van Steenkiste et al. , 2019 ; Locatello et al. , 2019 ) . A common intuitive way to think about disentangled representations is that it should reflect the compositional ∗Equal contribution . 1We have released the code for our method as part of disentanglement_lib Roof height . structure of the world . For example , to describe an object we often use words pertaining to its colour , position , shape and size . We can use different words to describe these properties because they relate to independent factors of variation in our world , i.e . properties which can be compositionally recombined . Hence a disentangled representation of objects should reflect this by factorising into dimensions which correspond to those properties ( Bengio et al. , 2013 ; Higgins et al. , 2018a ) . The ability to automatically discover the compositional factors of complex real datasets can be of great importance in many practical applications of machine learning and data science . However , it is important to be able to learn such representations in an unsupervised manner , since most interesting datasets do not have their generative factors fully labelled . For a long time scalable unsupervised disentangled representation learning was impossible , until recently a new class of models based on Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) was developed . These approaches ( Higgins et al. , 2017a ; Burgess et al. , 2017 ; Chen et al. , 2018 ; Kumar et al. , 2017 ; Kim & Mnih , 2018 ) scale reasonably well and are the current state of the art in unsupervised disentangled representation learning . However , so far the benefits of these techniques have not been widely exploited because of two major shortcomings : First , the quality of the achieved disentangling is sensitive to the choice of hyperparameters , however , model selection is currently impossible without having access to the ground truth generative process and/or attribute labels , which are required by all the currently existing disentanglement metrics ( Higgins et al. , 2017a ; Kim & Mnih , 2018 ; Chen et al. , 2018 ; Eastwood & Williams , 2018 ; Ridgeway & Mozer , 2018 ) . Second , even if one could apply any of the existing disentanglement metrics for model selection , the scores produced by these metrics can vary a lot even for models with the same hyperparameters and trained on the same data ( Locatello et al. , 2018 ) . While a lot of this variance is explained by the actual quality of the learnt representations , some of it is introduced by the metrics themselves . In particular , all of the existing supervised disentanglement metrics assume a single “ canonical ” factorisation of the generative factors , any deviation from which is penalised . Such a “ canonical ” factorisation , however , is not chosen in a principled manner . Indeed , for the majority of datasets , apart from the simplest ones , multiple equally valid disentangled representations may be possible ( see Higgins et al . ( 2018a ) for a discussion ) . For example , the intuitive way that humans reason about colour is in terms of hue and saturation . However , colour may also be represented in RGB , YUV , HSV , HSL , CIELAB . Any of the above representations are as valid as each other , yet only one of them is allowed to be “ canonical ” by the supervised metrics . Hence , a model that learns to represent colour in a subspace aligned with HSV will be penalised by a supervised metric which assumes that the canonical disentangled representation of colour should be in RGB . This is despite the fact that both representations are equal in terms of preserving the compositional property at the core of what makes disentangled representations useful ( Higgins et al. , 2018a ) . Hence , the field finds itself in a predicament . From one point of view , there exists a set of approaches capable of reasonably scalable unsupervised disentangled representation learning . On the other hand , these models are hard to use in practice , because there is no easy way to do a hyperparameter search and model selection . This paper attempts to bridge this gap . We propose a simple yet effective method for unsupervised model selection for the class of current state-of-the-art VAE-based unsupervised disentangled representation learning methods . Our approach , Unsupervised Disentanglement Ranking ( UDR ) , leverages the recent theoretical results that explain why variational autoencoders disentangle ( Rolinek et al. , 2019 ) , to quantify the quality of disentanglement by performing pairwise comparisons between trained model representations . We evaluate the validity of our unsupervised model selection metric against the four best existing supervised alternatives reported in the large scale study by Locatello et al . ( 2018 ) : theβ-VAE metric ( Higgins et al. , 2017a ) , the FactorVAE metric ( Kim & Mnih , 2018 ) , Mutual Information Gap ( MIG ) ( Chen et al. , 2018 ) and DCI Disentanglement scores ( Eastwood & Williams , 2018 ) . We do so for all existing state of the art disentangled representation learning approaches : β-VAE ( Higgins et al. , 2017a ) , CCI-VAE ( Burgess et al. , 2017 ) , FactorVAE ( Kim & Mnih , 2018 ) , TC-VAE ( Chen et al. , 2018 ) and two versions of DIP-VAE ( Kumar et al. , 2017 ) . We validate our proposed method on two datasets with fully known generative processes commonly used to evaluate the quality of disentangled representations : dSprites ( Matthey et al. , 2017 ) and 3D Shapes ( Burgess & Kim , 2018 ) , and show that our unsupervised model selection method is able to match the supervised baselines in terms of guiding a hyperparameter search and picking the most disentangled trained models both quantitatively and qualitatively . We also apply our approach to the 3D Cars dataset ( Reed et al. , 2014 ) , where the full set of ground truth attribute labels is not available , and confirm through visual inspection that the ranking produced by our method is meaningful ( Fig . 1 ) . Overall we evaluate 6 different model classes , with 6 separate hyperparameter settings and 50 seeds on 3 separate datasets , totalling 5400 models and show that our method is both accurate and consistent across models and datasets . Finally , we validate that the model ranking produced by our approach correlates well with the final task performance on two recently reported tasks : a classification fairness task ( Locatello et al. , 2019 ) and a model-based reinforcement learning ( RL ) task ( Watters et al. , 2019 ) . Indeed , on the former our approach outperformed the reported supervised baseline scores . 2 OPERATIONAL DEFINITION OF DISENTANGLING . Given a dataset of observationsX= { x1 , ... , xN } , we assume that there exist a number of plausible generative processes gi that produce the observations from a small set of correspondingKi independent generative factors ci . For each choice of i , g : cn 7→xn , where p ( cn ) = ∏K j=1p ( c j n ) . For example , a dataset containing images of an object , which can be of a particular shape and colour , and which can be in a particular vertical and horizontal positions , may be created by a generative process that assumes a ground truth disentangled factorisation into shape x colour x position , or shape x hue x saturation x position X x position Y . We operationalise a model as having learnt a disentangled representation , if it learns to invert of one of the generative processes gi and recover a latent representation z∈RL , so that it best explains the observed data p ( z , x ) ≈p ( ci , x ) , and factorises the same way as the corresponding data generative factors ci . The choice of the generative process can be determined by the interaction between the model class and the observed data distribution p ( x ) , as discussed next in Sec . 3.1 . 3 VARIATIONAL UNSUPERVISED DISENTANGLING . The current state of the art approaches to unsupervised disentangled representation learning are based on the Variational Autoencoder ( VAE ) framework ( Rezende et al. , 2014 ; Kingma & Welling , 2014 ) . VAEs attempt to estimate the lower bound on the joint distribution of the data and the latent factors p ( x , z ) by optimising the following objective : LV AE=Ep ( x ) [ Eqφ ( z|x ) [ log pθ ( x|z ) ] −KL ( qφ ( z|x ) || p ( z ) ) ] ( 1 ) where , in the usual case , the prior p ( z ) is chosen to be an isotropic unit Gaussian . In order to encourage disentangling , different approaches decompose the objective in Eq . 1 into various terms and change their relative weighting . In this paper we will consider six state of the art approaches to unsupervised disentangled representation learning that can be grouped into three broad classes based on how they modify the objective in Eq . 1 : 1 ) β-VAE ( Higgins et al. , 2017a ) and CCI-VAE ( Burgess et al. , 2017 ) upweight the KL term ; 2 ) FactorVAE ( Kim & Mnih , 2018 ) and TC-VAE ( Chen et al. , 2018 ) introduce a total correlation penalty ; and 3 ) two different implementations of DIP-VAE ( -I and -II ) ( Kumar et al. , 2017 ) penalise the deviation of the the marginal posterior from a factorised prior ( see Sec . A.4.1 in Supplementary Material for details ) . 3.1 WHY DO VAES DISENTANGLE ? . In order to understand the reasoning behind our proposed unsupervised disentangled model selection method , it is first important to understand why VAEs disentangle . The objective in Eq . 1 does not in itself encourage disentangling , as discussed in Rolinek et al . ( 2019 ) and Locatello et al . ( 2018 ) . Indeed , any rotationally invariant prior makes disentangled representations learnt in an unsupervised setting unidentifiable when optimising Eq . 1 . This theoretical result is not surprising and has been known for a while in the ICA literature ( Comon , 1994 ) , however what is surprising is that disentangling VAEs appear to work in practice . The question of what makes VAEs disentangle was answered by Rolinek et al . ( 2019 ) , who were able to show that it is the peculiarities of the VAE implementation choices that allow disentangling to emerge ( see also discussion in Burgess et al . ( 2017 ) ; Mathieu et al . ( 2019 ) ) . In particular , the interactions between the reconstruction objective ( the first term in Eq . 1 ) and the enhanced pressure to match a diagonal prior created by the modified objectives of the disentangling VAEs , force the decoder to approximate PCA-like behaviour locally around the data samples . Rolinek et al . ( 2019 ) demonstrated that during training VAEs often enter the so-called “ polarised regime ” , where many of the latent dimensions of the posterior are effectively switched off by being reduced to the prior qφ ( zj ) =p ( zj ) ( this behaviour is often further encouraged by the extra disentangling terms added to the ELBO ) . When trained in such a regime , a linear approximation of the Jacobian of the decoder around a data sample xi , Ji= ∂Decθ ( µφ ( xi ) ) ∂µφ ( xi ) , is forced to have orthogonal columns , and hence to separate the generative factors based on the amount of reconstruction variance they induce . Given that the transformations induced by different generative factors will typically have different effects on the pixel space ( e.g . changing the position of a sprite will typically affect more pixels than changing its size ) , such local orthogonalisation of the decoder induces an identifiable disentangled latent space for each particular dataset . An equivalent statement is that for a well disentangled VAE , the SVD decomposition J=UΣV > of the Jacobian J calculated as above , results in a trivial V , which is a signed permutation matrix .

Summary: The paper proposes a metric for unsupervised model (and hyperparameter) selection for VAE-based models. The essential basis for the metric is to rank the models based on how much disentanglement they provide. This method relies on a key observation from this paper [A] viz., disentangled representations by any VAE-based model are likely to be similar (upto permutation and sign).