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 extends the Wasserstein Autoencoder (WAE) work by splitting the divergence on the variational marginal into 2 terms, akin to what was done in TC-VAE. This enables directly controlling the explicit contribution of the total correlation term, which is likely to contribute to disentanglement more directly. They explore 2 variations of their model, based on different estimators of the TC term (TCWAE-MWS, using minibatch-weighted sampling; TCWAE-GAN, using a density ratio trick).