RTC-VAE: H ARNESSING THE PECULIARITY OF T OTAL C ORRELATION IN L EARNING D ISENTANGLED R EPRESENTATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
In the problem of unsupervised learning of disentangled representations, one of the promising methods is to penalize the total correlation of sampled latent variables.
Unfortunately, this well-motivated strategy often fails to achieve disentanglement due to a problematic difference between the sampled latent representation and its corresponding mean representation.
We provide a theoretical explanation that low total correlation of sample distribution cannot guarantee low total correlation of the mean representation.
We prove that for the mean representation of arbitrarily high total correlation, there exist distributions of latent variables of a bounded total correlation.
However, we still believe that total correlation could be a key to the disentanglement of unsupervised representative learning, and we propose a remedy, RTC-VAE, which rectiﬁes the total correlation penalty.
Experiments show that our model has a more reasonable distribution of the mean representation compared with baseline models, e.g., β -TCVAE and FactorVAE.
I NTRODUCTION
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.
; 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 deﬁned to be the KL-divergence between the joint distribution z ∼ q ( z ) and the product of marginal distributions (cid:81)
j q ( z j ) .
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 TC sample being low does not necessarily give rise to a low total correlation of the corresponding mean representation TC mean .
Conventionally, the mean representation is used as the encoded latent variables, an unnoticed high TC mean 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 TC sample and mean representation TC mean 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.
and they concluded that ”the unsupervised learning of disentangled representations is fundamentally impossible without inductive biases”.
Acknowledging the difﬁculty 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 ﬁnd 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 cannot guarantee a low total correlation of the mean representation.
(Section. 2)
• Acknowledging the issue above, we further delve into total correlation, and provide a simple 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 TC mean 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.
We point out that the method of (minibatch) estimators suffers from the curse of dimensionality and other drawbacks, making their estimation accuracy decay signiﬁcantly 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)
T HE PECULIARITY OF T OTAL C ORRELATION
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.
Deﬁnition 1.
Total correlation of random variable x ,
[Equation]
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.
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 max j σ j = c 0 .
For a ﬁxed µ , let z ∼ N ( µ, Σ (cid:48) ( µ )) , where Σ (cid:48) ( µ ) is diagonal and satisﬁes that for some R > 0 ,
[Equation]
for some constants c 1 , c 2 , c 3 , c 4 .
Then TC( z ) ≤ C for some C = C ( R, c 0 , · · · , c 4 , 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, c 0 , · · · , c 4 , l ﬁxed, 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.
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 ﬂexible to easily ﬁnd 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.
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 TC sample .
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 ( µ, Σ (cid:48) ( µ )) 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.
R ELATED W ORKS
In the study of disentanglement, Higgins et al.
(2017) proposed a modiﬁcation 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.
The works above belong to representative learning without inductive biases.
There are also works about representative learning with inductive biases, see Rolinek et al.
and references therein. As for the disentanglement metric, this will be discussed in Section.
R ECTIFIED -TCVAE
To simplify notation, let p ( n ) =
p ( x n ) , q ( z | x n ) =
q ( z | n ) .
Recall the average evidence lower bound (ELBO),
[Equation]
Chen et al.
(2018) and
independently by Kim & Mnih (2018) introduced objective function that penalizes total correlation, which can be formulated as
[Equation]
This approach unfortunately has a drawback.
It turns out that instead of being able to obtain disentangled representation, we often ﬁnd 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 deﬁne RTC-VAE,
[Equation]
where
[Equation]
Our penalty originates from the ﬁrst term of the law of total covariance
Cov q ( z ) [ z ] =
E p ( n ) Cov q ( z | n ) [ z ] + Cov p ( n ) ( E q ( z | n ) [ z ]) .
Factorized representation 1 indicates a diagonal covariance matrix
Cov q ( z ) [ z ] .
Motivated by this, Kumar et al.
(2017) penalizes the off-diagonal terms in the second term, while ignores E p ( n ) Cov q ( 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.
This is actually not surprising since the two penalty terms in DIP-VAE contribute in opposite directions, with one leading to vanishing µ ’s and another ﬁghting 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 ﬁrst sight, since penalizing a diagonal component of covariance Cov [ z ] =
Cov q ( 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.
E STIMATION OF TOTAL CORRELATION
The naive Monte Carlo method comes with an intrinsic issue of underestimating total correlation.
To avoid or resolve this, Kim & Mnih (2018) proposed a discriminator network with the help of density-ratio trick (see equation (3) and Appendix D. of Kim & Mnih (2018)).
In Chen et al.
(2018), two kinds of estimator of total correlation are proposed,
Minibatch Weighted Sampling (MWS) and
Minibatch Stratiﬁed Sampling (MSS) (see Appendix C.1 and C.2 in Chen et al.
5. M ETHOD OF M INIBATCH E STIMATORS
For instance, MSS can be described as followed.
For a minibatch of sample, B M +1 = { n 1 , . . . , n M [Equation]
[Equation]
where
[Equation]
z
[Equation]
For the convenience of readers, MWS is listed in Appendix A.1.
5. M ETHOD OF DENSITY - RATIO TRICK AND DISCRIMINATOR
Density-ratio trick Nguyen et al.
; Sugiyama et al.
(2012) can be used to estimate KLdivergence,
[Equation]
where D is discriminator that classiﬁes
z being sampled from q ( z ) or (cid:81)
j q ( z j ) .
For detail implementation, please refer to section 3 in Kim & Mnih (2018).
5. C OMPARISON OF THE TWO METHODS
For multivariate normal distribution, the total correlation can be explicitly calculated which can be used as a ground truth for our comparison.
To be speciﬁc,
Proposition 1.
Let x ∼ N (0 , Σ ) , then
[Equation] [Equation]
It’s difﬁcult to track the exact reference of Proposition 1 since it is a fundamental property in information theory.
Locatello et al.
(2018) used this proposition to approximate the total correlation of the mean representation in latent space.
In appendix, we provide a simple proof for the convenience of the readers.
We compared the performance of each method, MWS, MSS 0 and MSS 1 on the estimation of total correlation.
For µ ∼ N (0 , I ) , and z | µ ∼ N ( µ, Σ ) where Σ = diag( σ 2 ) and σ = 0 . 1 .
We choose σ small so that the distribution of z can be approximated by normal distribution.
Results are presented in Figure 1.
From Figure.
1, we can summarize the following observations: 1.
MWS tends to underestimate total correlation in general; 2.
For latent space of dimension ≤ 4 , MSS 0 and MSS 1 are quite accurate; 3.
For latent space of high dimension, both MSS 0 and MSS 1 tend to overestimate total correlation when the actual value of total correlation is small; 4.
Overall MSS 1 estimates closer to ground truth than MMS 0 does.
In the following analysis of the above observations, we use a less formal way of analyzing, which can be formalized to be rigorous, in order to convey our idea directly.
To interpret the third observation, let µ ∼ N (0 , Id) and z | µ ∼ N (0 , Σ ) , where Id is identity matrix and Σ = diag( σ 2 ) . Then TC ( µ ) = 0 and TC ( z ) small if σ small. Consider q ( z ( i ) k | n ( j ) ) where ( i, j, k ) are indices of a cube ( minibatch, minibatch, dimension ) with size M × M × D and n ( j ) is a sample drawn in a minibatch and z ( i ) =
z ( n ( i ) ) .
We claim this: when the ground truth total correlation of z is low (the off-diagnal values of correlation matrix is small), only the elements on the diagonal surface of the cube, namely those elements of index ( i, i, k ) , take some bounded values O (1) , and all the other elements are very small o (1) (since σ = 0 . 1 ). To see the claim, let’s ﬁrst consider 1-D case, where µ ∼ N (0 , 1) , z | µ ∼ N (0 , σ 2 ) . When σ is small, z can be approximately treated as N (0 , 1) .
z ( i ) and µ ( j ) are independent for i (cid:54) =
j , hence z ( i ) − µ ( j ) ∼ N (0 , 2) , and [Equation] See a proof in Appendix A.5. Then for D-dimension, the probability P ( | z ( i ) − µ ( j ) | < t ) would be
O ( t D ) .
Now, only if t takes value as small as σ , q ( z ( i ) k | n ( j ) ) is not small.
For example, σ = 0 .and
{+D = 
+}
and the chance of such case to happen is O (10 − 10 ) .
Compared to batchsize, usually
O (10 3 ) , the amount of such cases can be ignored.
Hence,
[Equation]
Assigning weights to elements q ( z ( i ) k | n ( j ) ) such as MSS does not make essential change to the analysis above.
In addition, β -TCVAE (trained with MSS in our experiments and with MWS in Locatello et al.
(2018)) seems to have an increasing total correlation of mean representation as regularization strength increases (higher β ’s), as observed by Locatello et al.
Here, we provide an explanation to the cause of this problem:
First, MSS and MWS prefer to shut down latent dimensions, meaning that distributions with fewer active dimensions can score lower estimated total correlation.
Consider that µ 0 ∼ N (0 , 0 . 01) , µ 0 − ∼ N (0 , Id) and also z | µ ∼ N (0 , Σ ) , where Σ = diag( σ 2 ) .
Then all elements of index ( i, j, 0) , i.e., q ( z ( i )
| n ( j ) ) are not small, say O (1) , for the same reason as previous in D − 1 dimension latent space.
Thus, 1 M (cid:80) j q ( z ( i ) [-
-]0 | n ( j ) )
≈ 1 M (cid:80)
j O (1)
≈ O (1) , and TC
[Equation]
compared to O (( D − 1) log M ) when µ ∼ N (0 , Id) in previous analysis.
Now, consider any strongly correlated z ’s (e.g., ( z 1 , z 2 ) ∼ N (0 , Σ ) , where Σ = (cid:18) 0 .− 0 .− 0 .
(cid:19) , see Figure 2).
Then the Gaussian (ground truth) total correlation is arbitrarily large (TC ( z 1 , z 2 ) = ∞ ).
This kind of distribution can score a relatively low TC value (for instance lower than z ) with estimators such as MSS and MWS by the analysis above.
Hence, as β increases, VAE trained with these estimators will be encouraged to obtain some dimensions of very low variance, and these dimensions are easily trapped in a strong correlation with other dimensions (like z 1 and z 2 ).
Shutting down dimension is not preferable because latent dimensions should not be fewer than ground truth.
Moreover, considering datasets such as dSprites, though shape is labelled as a single dimension, models can learn to represent complex geometry with multiple dimensions, hence more active dimensions are learned than the number of ground-truth dimensions.
Based on the reasons above, we opt for the method of discriminators (density-ratio trick) in our implementation.
Experiment shows that density-ratio trick provides a more stable estimation of total correlation when training VAE.
E XPERIMENTS
The datasets we use include dSprites Matthey et al.
(2017), Shapes3D Burgess & Kim (2018) and 3D faces Paysan et al.
At the time of writing, the scale of our experiments is limited, but there are already some evidence to deliver our arguments.
We have scheduled further experiments and tests on larger scale in future works.
For every model, we trained with 10 different initialization.
Hyperparameter β , also the regularization strength, takes 2 , 4 , · · · , 10 .
For hyperparameter η , we ﬁx η = 10 for a simple reason. Since the variance term in equation 4 becomes small (close to 0) shortly after training begins, this term will not contribute much compared with L beta − TC term. So, η = 10 is enough to strengthen the penalty at the beginning of training. In our experiments, higher value of η cannot bring further improvement since z is close to µ already.
And lower values may not guarantee z being close to µ . From experiments, we observe that RTC-VAE has much lower TC mean with different regularization strength than FactorVAE does (Figure 4).
And on different datasets, this is also the case (see Appendix).
The TC mean behaves almost identically as TC sample in RTC-VAE (see Figure 3 (b)).
The problem of contradictory behaviors of TC mean and TC sample is evidently remedied by RTCVAE.
In addition, the ELBO of RTC-VAE seems to converge faster than FactorVAE as a byproduct (see Figure 5 and Figure 6).
Examining the distributions of latent dimensions (mean representation),
FactorVAE tends to have some strongly correlated latent dimensions (see Figure 8), and RTC-VAE shows well factorized latent distributions (see Figure 7).
6. M ETRICS OF DISENTANGLEMENT
We wish there were a widely accepted metric of disentanglement to compare our model RTCVAE with other models.
Unfortunately, it is still an open question, how we can measure disentanglement.
Various attempts have been made so far, but Locatello et al.
(2018) challenged most of them, indicating that the score under any metric varies due to different initializations and datasets.
Here, [Figure] we analyse several important metrics and attempt to point out some blind spots that have not being considered by these metrics.
Higgins et al.
(2017) proposed using a classiﬁer to measure each dimension of latent space and each ground truth factor, e.g. ( x, y ) coordinates, scale, rotation, etc.
Kim & Mnih (2018) revised this [Figure] [Figure] approach yet kept the idea.
Chen et al.
(2018) introduced mutual information gap (MIG), which estimates the mutual information between each latent dimension and each ground truth factor.
Note that this is a classiﬁer-free metric.
Eastwood & Williams (2018) proposed a framework of disentan- glement metric that considers modularity, compactness and explicitness .
Then Ridgeway & Mozer (2018) made analysis on compactness , and compactness mean that each ground truth factor associates with only one or a few latent dimensions.
They pointed out that in some situation a perfectly disentangled representation may not be compact (see Section.
3 in Ridgeway & Mozer (2018)).
Here we argue that modularity also should be reconsidered.
A modular representation means that each dimension of latent space conveys information of at most one ground truth factor.
This is exactly the goal attempted by Higgins et al.
Kim & Mnih (2018); Chen et al.
(2018), etc.
However, multiple latent dimensions can work together to represent multiple ground truth factors meanwhile these latent dimensions are disentangled.
For instance, x and y coordinates can be represented by r and θ in polar coordinate system (or any coordinate system under rotation, i.e., ( x (cid:48) , y (cid:48) )
T =
A ( x, y ) T where A is any orthogonal matrix).
These coordinate systems are perfectly disentangled but r (or x (cid:48) )
conveys information of both x and y .
C ONCLUSION
In this work, we demonstrated that our RTC-VAE, which rectiﬁes the total correlation penalty can remedy its peculiar properties (disparity between total correlation of the samples and the mean representations).
Our experiments show that our model has a more reasonable distribution of the mean representation compared with baseline models including β -TCVAE and FactorVAE.
We also provide several theoretical proofs which could help diagnose several speciﬁc symptoms of entangle- ment.
Hopefully, our contributions could add to the explainability of the unsupervised learning of disentangled representations. A A PPENDIX A. M INIBATCH W EIGHTED S AMPLING (MWS)
See Chen et al.
[Equation]
[Equation]
A. MMS 0 AND MMS
There is no mathematical difference between MSS 0 and MSS 1 (same formulation as equation 5), only a difference in implementation. We replace this chunk of code ( https://github.com/ rtqichen/beta-tcvae/blob/master/vae_quant.py#L199-L201 ) to f
o r i in range ( b a t c h s i z e ) : [Equation]
A. P ROOF OF T HEOREM
In the following proof, we use the convention of mathematical analysis that the meaning of C can change through lines to eliminate some redundant work of tracking.
Theorem (Theorem 1 restated) .
Let µ ∼ N (0 , Σ ) and σ j be the standard deviation of µ j , j = 1 , · · · , D , and max j σ j = c 0 .
For a ﬁxed µ , let z ∼ N ( µ, Σ (cid:48) ( µ )) , where Σ (cid:48) ( µ ) is diagonal and satisﬁes that for some R > 0 , [Equation] some
[Equation]
(12) for some constants c 1 , c 2 , c 3 , c 4 .
Then TC( z ) ≤ C for some C = C ( R, c 0 , · · · , c 4 , l ) > 0 .
Proof.
Let
[Equation]
then
[Equation]
Since KL-divergence is non-negative, if TC( z ) + is bounded, then TC( z ) must be bounded.
In the following, we work on S + , i.e., we assume p ( z ) ≥ (cid:81)
j p ( z j ) .
For | z | < R ,
[Equation]
and (cid:89)
[Equation]
For | z | > 2 R , p ( z ) =
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where r 0 = max( c 0 , c 2 ) and r 1 = min( c 1 , c 4 ) , and since l ≥ 1 and D ≥ 1 , Dl + D − 2 ≥ 0 .
For | z | ∈ ( R 1 , 2 R 1 ) , it is easy to see that p ( z ) < C .
And for | z | > R , (cid:89)
j p ( z j ) ≥ (cid:89) j (cid:32)(cid:90)
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Equation]
e −
[Equation]
Hence,
[Equation]
A. P ROOF OF P ROPOSITION
Proposition (Proposition 1 restated) .
Let x ∼ N (0 , Σ ) [Equation] then
[Equation] [Equation]
Proof.
First [Equation] recall that the KL-divergence between two distributions P and Q is deﬁned as
[Equation]
Also
[Equation]
the density function for a multivariate Gaussian distribution N ( µ, Σ ) is
[Equation]
have
[Equation]
Let P be a multivariate Gaussian N ( µ, Σ 1 ) , and then the product of the marginal distribution (cid:81)
i p i ( x ) is also Gaussian N ( µ, Σ 2 ) , where Σ 2 = diag( Σ 1 ) .
Thus, the total correlation of multivariate Gaussian distribution is [Equation] A. P ROOF OF EQUATION
Proof.
P ( | z ( i ) − µ ( j ) | < t ) = P ( | x | < t ) where x ∼ N (0 , 2)
[Equation]
A. E XPERIMENTS
[Figure]