RTC-VAE: HARNESSING THE PECULIARITY OF TOTAL CORRELATION  IN LEARNING DISENTANGLED REPRESENTATIONS

1 INTRODUCTION . VAEs ( Variational AutoEncoders ) Kingma & Welling ( 2013 ) ; Bengio et al . ( 2007 ) follow the common assumption that the high-dimensional real world observations x can be re-generated by a lowerdimension latent variable z which is semantically meaningful . Recent works Kim & Mnih ( 2018 ) ; Chen et al . ( 2018 ) ; Kumar et al . ( 2017 ) suggest that decomposing the ELBO ( Evidence Lower Bound ) could lead to distinguishing the factor of disentanglement . In particular , recent works Kim & Mnih ( 2018 ) ; Chen et al . ( 2018 ) focused on a term called total correlation ( TC ) . The popular belief Chen et al . ( 2018 ) is that by adding weights to this term in objective function , a VAE model can learn a disentangled representation . This approach appears to be promising since the total correlation of a sampled representation should describe the level of factorising since total correlation is defined to be the KL-divergence between the joint distribution z ∼ q ( z ) and the product of marginal distributions ∏ j q ( zj ) . In this case , a low value suggests a less entangled joint distribution . However , Locatello et al . ( 2018 ) pointed out that the total correlation of sampled distribution TCsample being low does not necessarily give rise to a low total correlation of the corresponding mean representation TCmean . Conventionally , the mean representation is used as the encoded latent variables , an unnoticed high TCmean is usually the culprit behind the undesirable entanglement . Moreover , Locatello et al . ( 2018 ) found that as regularization strength increases , the total correlation of sampled representation TCsample and mean representation TCmean are actually negatively correlated . Locatello et al . ( 2018 ) put doubts on most methods of disentanglement including penalizing the total correlation term Kim & Mnih ( 2018 ) ; Chen et al . ( 2018 ) , and they concluded that ” the unsupervised learning of disentangled representations is fundamentally impossible without inductive biases ” . Acknowledging the difficulty in learning disentangled representation , we provide a detailed explanation of the seemingly contradictory behaviors of the total correlations of sampled and mean representation in previous works on TC penalizing strategy . Moreover , we find that this problem described above can be remedied simply with an additional penalty term on the variance of a sampled representation . Our contributions : • In Theorem 1 , we prove that for all mean representations , there exists a large class of sample distributions with bounded total correlation . Particularly , a mean representation with arbitrarily large total correlation can have a corresponding sample distribution with low total correlation . This implies that a low total correlation of sample distribution can not guarantee a low total correlation of the mean representation . ( Section . 2 ) • Acknowledging the issue above , we further delve into total correlation , and provide a sim- ple remedy by adding an additional penalty term on the variance of sample distribution . The penalty term forces a sampled representation to behave similar to the corresponding mean representation . Such penalty term is necessary for the strategy of penalizing TCmean in the view of Theorem 1 . ( Section . 4 ) • We study several different methods of estimating total correlation . They are compared and benchmarked against the ground truth value on the multivariate Gaussian distribution Locatello et al . ( 2018 ) . We point out that the method of ( minibatch ) estimators suffers from the curse of dimensionality and other drawbacks , making their estimation accuracy decay significantly with the increase of the dimension of the latent space , and some strong correlated distributions can be falsely estimated to have low total correlation . ( Section . 5 ) 2 THE PECULIARITY OF TOTAL CORRELATION . In information theory , total correlation is one of the generalizations of mutual information , which measures the difference between the joint distribution of multiple random variables and the product of their marginal distributions . A high value means the joint distribution is far from an independent distribution , and hence it suggests high entanglement among these random variables . Definition 1 . Total correlation of random variable x , TC ( x ) : = KL p ( x ) ||∏ j p ( xj ) = Ep ( x ) [ log p ( x ) ∏ j p ( xj ) ] . Naturally , people seek the solution of disentanglement in the form of low total correlation of the latent variables , e.g . Kim & Mnih ( 2018 ) ; Chen et al . ( 2018 ) . However , there can be large difference between the total correlations of sample representation and mean representation . Forcing the former to be small does not guarantee the latter being small . In fact , given a mean representation of arbitrarily large total correlation , we can construct a family of distribution of sample representation that have a bounded total correlation , where the bound does not rely on the total correlation of the mean . Theorem 1 . Let µ ∼ N ( 0 , Σ ) and σj be the standard deviation of µj , j = 1 , · · · , D , and maxj σj = c0 . For a fixed µ , let z ∼ N ( µ , Σ′ ( µ ) ) , where Σ′ ( µ ) is diagonal and satisfies that for some R > 0 , c2 > σ ′ j ( µ ) > c1 > 0 , if |µ| < R , c3 > σ ′ j ( µ ) > c4 |µ|l , for some l ≥ 1 , if |µ| > R. ( 1 ) for some constants c1 , c2 , c3 , c4 . Then TC ( z ) ≤ C for some C = C ( R , c0 , · · · , c4 , l ) > 0 . The details of the proof are presented in Appendix A.3 . Here ’ s another way to interpret Theorem 1 : with C and parameters R , c0 , · · · , c4 , l fixed , one can make TC ( µ ) arbitrarily large , since TC ( µ ) depends only on the correlation matrix of µ ( see Proposition 1 ) . Theorem 1 provides an explanation to the contradiction observed by Locatello et al . ( 2018 ) that TC ( z ) is low does not mean TC ( µ ) is low ( actually much higher than TC ( z ) ) . Since there exist such a large class of distributions of z that all have bounded TC ( z ) . When the objective function only penalizes TC ( µ ) , neural networks are so flexible to easily find a distribution with low TC ( z ) , and total correlation estimators like MSS can encourage shutting down latent dimensions ( see Section 5.3 ) , which together cause the disparity of TC ( µ ) and TC ( z ) . This fact was not noticed until Locatello et al . ( 2018 ) , and our investigation gives an explanation of the peculiar property of total correlation . Hence , Theorem 1 leads to the necessity of a regularizer of the difference between the distributions of µ and z when penalizing TCsample . In Section . 4 , we propose a simple regularizer that serves this goal . It is an interesting question whether there exists a distribution of z N ( µ , Σ′ ( µ ) ) with arbitrarily small TC ( z ) given µ . If not , what ’ s the lower bound of TC ( z ) ? These questions remain open to us for now , and we leave them to future work . 3 RELATED WORKS . In the study of disentanglement , Higgins et al . ( 2017 ) proposed a modification of the VAE framework and introduced an adjustable hyperparameter β that balances latent channel capacity and independence constraints with reconstruction accuracy . One drawback of β-VAE is the trade-off between the reconstruction quality and disentanglement . Motivated to alleviate this trade-off of β-VAE , Kim & Mnih ( 2018 ) proposed FactorVAE which decomposes the evidence lower bound and penalize a term measuring the total correlation between latent variables . Around the same time , Chen et al . ( 2018 ) proposed a similar ELBO decomposition method called β-TCVAE . The major difference between FactorVAE and β-TCVAE lies in their different strategies of estimating total correlation . Chen et al . ( 2018 ) used formulated estimators while Kim & Mnih ( 2018 ) utilized the density-ratio trick which requires an auxiliary discriminator network and an inner optimization loop . We will discuss these two strategies more in details in Section . 5 . The works above belong to representative learning without inductive biases . There are also works about representative learning with inductive biases , see Rolinek et al . ( 2019 ) and references therein . As for the disentanglement metric , this will be discussed in Section . 6 . 4 RECTIFIED-TCVAE . To simplify notation , let p ( n ) = p ( xn ) , q ( z|xn ) = q ( z|n ) . Recall the average evidence lower bound ( ELBO ) , ELBO : = Ep ( n ) [ Eq ( z|n ) [ log p ( n|z ) ] −KL ( q ( z|n ) ‖p ( z ) ) ] . ( 2 ) Chen et al . ( 2018 ) and independently by Kim & Mnih ( 2018 ) introduced objective function that penalizes total correlation , which can be formulated as Lβ−TC : = ELBO− βTC ( z ) . ( 3 ) This approach unfortunately has a drawback . It turns out that instead of being able to obtain disentangled representation , we often find a sample representation appears to be disentangled while the mean representation is still entangled . In fact , when we are maximizing Lβ−TC , we could totally end up learning a distribution of z that makes TC ( z ) goes low , while the total correlation of its mean µ is still high . To resolve this , we define RTC-VAE , LRTC : = Lβ−TC − η · tr ( Ep ( n ) Covq ( z|n ) [ z ] ) , ( 4 ) where tr ( Ep ( n ) Covq ( z|n ) [ z ] ) = D∑ k Ep ( n ) [ σ2k ( n ) ] . Our penalty originates from the first term of the law of total covariance Covq ( z ) [ z ] = Ep ( n ) Covq ( z|n ) [ z ] + Covp ( n ) ( Eq ( z|n ) [ z ] ) . Factorized representation1 indicates a diagonal covariance matrix Covq ( z ) [ z ] . Motivated by this , Kumar et al . ( 2017 ) penalizes the off-diagonal terms in the second term , while ignores Ep ( n ) Covq ( z|n ) [ z ] since it is diagonal . Their penalty term leads to a vanishing µ , which is the mean representation . The remedy to this is to add another penalty term on the distance between µ ’ s and 1 . DIP-VAE Kumar et al . ( 2017 ) employs this remedy , however , DIPVAE does not outperform other VAE ’ s when measured by various disentanglement metrics , e.g. , FactorVAE score , see Fig . 3 & 14 in Locatello et al . ( 2018 ) . This is actually not surprising since the 1Factorized representation means each latent dimension is independent . two penalty terms in DIP-VAE contribute in opposite directions , with one leading to vanishing µ ’ s and another fighting against it . This formulation can easily get the model stuck in saddle points . Our objective , on the other hand , does not penalize directly on µ . Instead , it penalizes on σ , the standard deviation of the distribution q ( z|n ) . This may seem little counter-intuitive at first sight , since penalizing a diagonal component of covariance Cov [ z ] = Covq ( z ) [ z ] seems not helpful to factorising . However , in the view of Theorem 1 , our objective will force the distribution of z to be similar to the distribution of µ . Hence , it pushes us away from the situation of large TC ( µ ) and low TC ( z ) . Consequently , by minimizing TC ( z ) we get a model that has low TC ( µ ) , a disentangled mean representation .

Summary: This paper considers the extensions of variational autoencoders (VAEs), which take into account the total correlation of sampled distribution of latent variables. Proving a theorem that a family of distributions of sample representations with a bounded total correlation can have a mean representation of arbitrarily large total correlation, the authors propose RTC-VAE, which additionally penalizes total covariance of sampled latent variables. The authors demonstrate that RTC-VAE produces less correlated distributions of mean representation compared with baselines.