Learning disentangled representations with the Wasserstein Autoencoder

1 INTRODUCTION . Learning representations of data is at the heart of deep learning ; the ability to interpret those representations empowers practitioners to improve the performance and robustness of their models ( Bengio et al. , 2013 ; van Steenkiste et al. , 2019 ) . In the case where the data is underpinned by independent latent generative factors , a good representation should encode information about the data in a semantically meaningful manner with statistically independent latent variables encoding for each factor . Bengio et al . ( 2013 ) define a disentangled representation as having the property that a change in one dimension corresponds to a change in one factor of variation , while being relatively invariant to changes in other factors . While many attempts to formalize this concept have been proposed ( Higgins et al. , 2018 ; Eastwood & Williams , 2018 ; Do & Tran , 2019 ) , finding a principled and reproducible approach to assess disentanglement is still an open problem ( Locatello et al. , 2019 ) . Recent successful unsupervised learning methods have shown how simply modifying the ELBO objective , either re-weighting the latent regularization terms or directly regularizing the statistical dependencies in the latent , can be effective in learning disentangled representation . Higgins et al . ( 2017 ) and Burgess et al . ( 2018 ) control the information bottleneck capacity of Variational Autoencoders ( VAEs , ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) ) by heavily penalizing the latent regularization term . Chen et al . ( 2018 ) perform ELBO surgery to isolate the terms at the origin of disentanglement in β-VAE , improving the reconstruction-disentanglement trade off . Esmaeili et al . ( 2018 ) further improve the reconstruction capacity of β-TCVAE by introducing structural dependencies both between groups of variables and between variables within each group . Alternatively , directly regularizing the aggregated posterior to the prior with density-free divergences ( Zhao et al. , 2019 ) or moments matching ( Kumar et al. , 2018 ) , or simply penalizing a high Total Correlation ( TC , ( Watanabe , 1960 ) ) in the latent ( Kim & Mnih , 2018 ) has shown good disentanglement performances . In fact , information theory has been a fertile ground to tackle representation learning . Achille & Soatto ( 2018 ) re-interpret VAEs from an Information Bottleneck view ( Tishby et al. , 1999 ) , re-phrasing it as a trade off between sufficiency and minimality of the representation , regularizing a pseudo TC between the aggregated posterior and the true conditional posterior . Similarly , Gao et al . ( 2019 ) use the principle of total Correlation Explanation ( CorEX ) ( Ver Steeg & Galstyan , 2014 ) and maximize the mutual information between the observation and a subset of anchor latent points . Maximizing the mutual information ( MI ) between the observation and the latent has been broadly used ( van den Oord et al. , 2018 ; Hjelm et al. , 2019 ; Bachman et al. , 2019 ; Tschannen et al. , 2020 ) , showing encouraging results in representation learning . However , Tschannen et al . ( 2020 ) argued that MI maximization alone can not explain the disentanglement performances of these methods . Building on the Optimal Transport ( OT ) problem ( Villani , 2008 ) , Tolstikhin et al . ( 2018 ) introduced the Wasserstein Autoencoder ( WAE ) , an alternative to VAE for learning generative models . Similarly to VAE , WAE maps the data into a ( low-dimensional ) latent space while regularizing the averaged encoding distribution . This is in contrast with VAEs where the posterior is regularized at each data point , and allows the encoding distribution to capture significant information about the data while still matching the prior when averaged over the whole data set . Interestingly , by directly regularizing the aggregated posterior , WAE hints at more explicit control on the way the information is encoded , and thus better disentanglement . The reconstruction term of the WAE allows for any cost function on the observation space , opening the door to better suited reconstruction terms , for example when working with continuous RGB data sets where the Euclidean distance or any metric on the observation space can result in more accurate reconstructions of the data . In this work , following the success of regularizing the TC in disentanglement , we propose to use the Kullback-Leibler ( KL ) divergence as the latent regularization function in the WAE . We introduce the Total Correlation WAE ( TCWAE ) with an explicit dependency on the TC of the aggregated posterior . Using two different estimators for the KL terms , we perform extensive comparison with succesful methods on a number of data sets . Our results show that TCWAEs achieve competitive disentanglement performances while improving modelling performance by allowing flexibility in the choice of reconstruction cost . 2 IMPORTANCE OF TOTAL CORRELATION IN DISENTANGLEMENT . 2.1 TOTAL CORRELATION . The TC of a random vector Z ∈ Z under P is defined by TC ( Z ) , dZ∑ d=1 Hpd ( Zd ) −Hp ( Z ) ( 1 ) where pd ( zd ) is the marginal density over only zd and Hp ( Z ) , −Ep log p ( Z ) is the Shannon differential entropy , which encodes the information contained in Z under P . Since dZ∑ d=1 Hpd ( Zd ) ≤ Hp ( Z ) ( 2 ) with equality when the marginals Zd are mutually independent , the TC can be interpreted as the loss of information when assuming mutual independence of the Zd ; namely , it measures the mutual dependence of the marginals . Thus , in the context of disentanglement learning , we seek a low TC of the aggregated posterior , p ( z ) = ∫ X p ( z|x ) p ( x ) dx , which forces the model to encode the data into statistically independent latent codes . High MI between the data and the latent is then obtained when the posterior , p ( z|x ) , manages to capture relevant information from the data . 2.2 TOTAL CORRELATION IN ELBO . We consider latent generative models pθ ( x ) = ∫ Z pθ ( x|z ) p ( z ) dz with prior p ( z ) and decoder network , pθ ( x|z ) , parametrized by θ. VAEs approximate the intractable posterior p ( z|x ) by introducing an encoding distribution ( the encoder ) , qφ ( z|x ) , and learning simultaneously θ and φwhen optimizing the variational lower bound , or ELBO , defined in Eq . 3 : LELBO ( θ , φ ) , E pdata ( X ) [ E qφ ( Z|X ) [ log pθ ( X|Z ) ] −KL ( qφ ( Z|X ) ‖ p ( Z ) ) ] ≤ E pdata ( X ) log pθ ( X ) ( 3 ) Following Hoffman & Johnson ( 2016 ) , we can decompose the KL term in Eq . 3 as : 1 Nbatch N∑ n=1 KL ( qφ ( Z|xn ) ‖ p ( Z ) ) = KL ( q ( Z , N ) ‖ q ( Z ) p ( N ) ) ︸ ︷︷ ︸ i index-code MI + KL ( q ( Z ) ‖ p ( Z ) ) ︸ ︷︷ ︸ ii marginal KL ( 4 ) where p ( n ) = 1N , q ( z|n ) = q ( z|xn ) , q ( z , n ) = q ( z|n ) p ( n ) and q ( z ) = ∑N n=1 q ( z|n ) p ( n ) . i refers to the index-code mutual information and represents the MI between the data and the latent under the join distribution q ( z , n ) , and ii to the marginal KL matching the aggregated posterior to the prior . While discussion on the impact of a high index-code MI on disentanglement learning is still open , the marginal KL term plays an important role in disentanglement . Indeed , it pushes the encoder network to match the prior when averaged , as opposed to matching the prior for each data point . Combined with a factorized prior p ( z ) = ∏ d pd ( zd ) , as it is often the case , the aggregated posterior is forced to factorize and align with the axis of the prior . More specifically , the marginal KL term in Eq . 4 can be decomposed the as sum of a TC term and a dimensionwise-KL term : KL ( q ( Z ) ‖ p ( Z ) ) = TC ( q ( Z ) ) + dZ∑ d=1 KL ( qd ( Zd ) ‖ pd ( Zd ) ) ( 5 ) Thus maximizing the ELBO implicitly minimizes the TC of the aggregated posterior , enforcing the aggregated posterior to disentangle as Higgins et al . ( 2017 ) and Burgess et al . ( 2018 ) observed when strongly penalizing the KL term in Eq . 3 . Chen et al . ( 2018 ) leverage the KL decomposition in Eq . 5 by refining the heavy latent penalization to the TC only . However , the index-code MI term in Eq . 4 seems to have little to no role in disentanglement ( see ablation study of Chen et al . ( 2018 ) ) , potentially arming the reconstruction performances ( Hoffman & Johnson , 2016 ) . 3 WAE NATURALLY GOOD AT DISENTANGLING ? . In this section we introduce the OT problem and the WAE objective , and discuss the compelling properties of WAEs for representation learning . Mirroring β-TCVAE decomposition , we derive the TCWAE objective . 3.1 WAE . The Kantorovich formulation of the OT between the true-but-unknown data distribution PD and the model distribution Pθ , for a given cost function c , is defined by : OTc ( PD , Pθ ) = inf Γ∈P ( PD , Pθ ) ∫ X×X c ( x , x̃ ) γ ( x , x̃ ) dxdx̃ ( 6 ) where P ( PD , Pθ ) is the space of all couplings of PD and Pθ ; namely , the space of joint distributions Γ on X × X whose densities γ have marginals pD and pθ . Tolstikhin et al . ( 2018 ) derive the WAE objective by restraining this space and relaxing the hard constraint on the marginal using a soft constraint with a Lagrange multiplier ( see Appendix A for more details ) : WD , c ( θ , φ ) , E pD ( x ) E qφ ( z|x ) E pθ ( x̃|z ) c ( x , x̃ ) + λD ( q ( Z ) ‖ p ( Z ) ) ( 7 ) where D is any divergence function and λ a relaxation parameter . The decoder , pθ ( x̃|z ) , and the encoder , qφ ( z|x ) , are optimized simultaneously by dropping the closed-form minimization over the encoder network , with standard stochastic gradient descent methods . Similarly to the ELBO , objective 7 consists of a reconstruction cost term and a latent regularization term , preventing the latent codes to drift away from the prior . However , WAE explicitly penalizes the aggregate posterior . This motivates , following Section 2.2 , the use of WAE in disentanglement learning . Rubenstein et al . ( 2018 ) have shown promising disentanglement performances without modifying the objective 7 . Another important difference lies in the functional form of the reconstruction cost in the reconstruction term . Indeed , WAE allows for more flexibility in the reconstruction term with any cost function allowed , and in particular , it allows for cost functions better suited to the data at hand and for the use of deterministic decoder networks ( Tolstikhin et al. , 2018 ; Frogner et al. , 2015 ) . This can potentially result in an improved reconstruction-disentanglement trade off as we empirically find in Sections 4.2 and 4.1 .

Summary: This paper addresses disentanglement in the latent space of autoencoders. To this end, it combines ideas from four existing papers, namely the reconstruction loss of the Wasserstein autoencoder, the regularization term decomposition from the total correlation autoencoder, and entropy estimation using minibatch-weighted sampling or the density-ratio trick. This combination certainly makes sense, as it brings together methods that have previously been shown to work well in isolation.