V ARIATIONAL I NFERENCE OF D ISENTANGLED L ATENT C ONCEPTS FROM U NLABELED O BSERVATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Disentangled representations, where the higher level data generative factors are reﬂected in disjoint latent dimensions, offer several beneﬁts such as ease of deriving invariant representations, transferability to other tasks, interpretability, etc.
We consider the problem of unsupervised learning of disentangled representations from large pool of unlabeled observations, and propose a variational inference based approach to infer disentangled latent factors.
We introduce a regularizer on the expectation of the approximate posterior over observed data that encourages the disentanglement.
We evaluate the proposed approach using several quantitative metrics and empirically observe signiﬁcant gains over existing methods in terms of both disentanglement and data likelihood (reconstruction quality).
I NTRODUCTION
Feature representations of the observed raw data play a crucial role in the success of machine learning algorithms.
Effective representations should be able to capture the underlying (abstract or high-level) latent generative factors that are relevant for the end task while ignoring the inconsequential or nuisance factors.
Disentangled feature representations have the property that the generative factors are revealed in disjoint subsets of the feature dimensions, such that a change in a single generative factor causes a highly sparse change in the representation.
Disentangled representations offer several advantages – (i) Invariance: it is easier to derive representations that are invariant to nuisance factors by simply marginalizing over the corresponding dimensions, (ii) Transferability: they are arguably more suitable for transfer learning as most of the key underlying generative factors appear segregated along feature dimensions, (iii) Interpretability: a human expert may be able to assign meanings to the dimensions, (iv) Conditioning and intervention: they allow for interpretable conditioning and/or intervention over a subset of the latents and observe the effects on other nodes in the graph.
Indeed, the importance of learning disentangled representations has been argued in several recent works (Bengio et al., 2013; Lake et al., 2016; Ridgeway, 2016).
Recognizing the signiﬁcance of disentangled representations, several attempts have been made in this direction in the past (Ridgeway, 2016).
Much of the earlier work assumes some sort of supervision in terms of: (i) partial or full access to the generative factors per instance (Reed et al., 2014; Yang et al., 2015;
Kulkarni et al., 2015; Karaletsos et al., 2015), (ii) knowledge about the nature of generative factors (e.g, translation, rotation, etc.)
(Hinton et al., 2011; Cohen & Welling, 2014), (iii) knowledge about the changes in the generative factors across observations (e.g., sparse changes in consecutive frames of a Video)
(Goroshin et al., 2015; Whitney et al., 2016;
Fraccaro et al., 2017; Denton &
Birodkar, 2017; Hsu et al., 2017), (iv) knowledge of a complementary signal to infer representations that are conditionally independent of it 1 Cheung et al.
Mathieu et al.
; Siddharth et al.
, However, in most real scenarios, we only have access to raw observations without any supervision about the generative factors.
It is a challenging problem and many of the earlier attempts have not been able to scale well for realistic settings (Cohen & Welling, 2015)
(see also, Higgins et al.
Recently, Chen et al.
(2016) proposed an approach to learn a generative model with disentangled factors based on Generative Adversarial Networks (GAN) (Goodfellow et al., 2014), however implicit generative models like GANs lack an effective inference mechanism 2 , which hinders its applicability to the problem of learning disentangled representations.
More recently, Higgins et al.
(2017) proposed an approach based on Variational AutoEncoder (VAE) Kingma & Welling (2013) for inferring disentangled factors.
The inferred latents using their method (termed as β -VAE ) are empirically shown to have better disentangling properties, however the method deviates from the basic principles of variational inference, creating increased tension between observed data likelihood and disentanglement.
This in turn leads to poor quality of generated samples as observed in (Higgins et al., 2017).
In this work, we propose a principled approach for inference of disentangled latent factors based on the popular and scalable framework of amortized variational inference (Kingma & Welling, 2013;
Stuhlm¨uller et al., 2013; Gershman & Goodman, 2014; Rezende et al., 2014) powered by stochastic optimization (Hoffman et al., 2013; Kingma & Welling, 2013; Rezende et al., 2014).
Disentanglement is encouraged by introducing a regularizer over the induced inferred prior .
Unlike β -VAE (Higgins et al., 2017), our approach does not introduce any extra conﬂict between disentanglement of the latents and the observed data likelihood, which is reﬂected in the overall quality of the generated samples that matches the VAE and is much better than β -VAE.
This does not come at the cost of higher entanglement and our approach also outperforms β -VAE in disentangling the latents as measured by various quantitative metrics.
F ORMULATION
We start with a generative model of the observed data that ﬁrst samples a latent variable z ∼ p ( z ) , and an observation is generated by sampling from p θ ( x | z ) .
The joint density of latents and observations is denoted as p θ ( x , z )
= p ( z ) p θ ( x | z ) .
The problem of inference is to compute the posterior of the latents conditioned on the observations, i.e., p ( z | x ) .
We assume that we are given a ﬁnite set of samples (observations) from the true data distribution p ( x ) .
In most practical scenarios involving high dimensional and complex data, this computation is intractable and calls for approximate inference.
Variational inference takes an optimization based approach to this, positing a family
D of approximate densities over the latents and reducing the approximate inference problem to ﬁnding a member density that minimizes the Kullback-Leibler divergence to the true posterior, i.e., q ∗ x = min q ∈D KL ( q ( z ) (cid:107) p ( z | x ))
(Blei et al., 2017).
The idea of amortized inference (Kingma &
Welling, 2013; Stuhlm ¨ uller et al., 2013; Gershman & Goodman, 2014; Rezende et al., 2014) is to explicitly share information across inferences made for each observation.
One successful way of achieving this for variational inference is to have a so-called recognition model , parameterized by φ , that encodes an inverse map from the observations to the approximate posteriors (also referred as variational autoencoder or VAE) (Kingma & Welling, 2013; Rezende et al., 2014).
The recognition model parameters are learned by optimizing the problem min φ E x KL ( q φ ( z | x ) (cid:107) p ( z | x )) .
This can be shown as equivalent to maximizing what is termed as evidence lower bound (ELBO): arg min
[Equation]
The ELBO (the objective at the right side of Eq. 1) lower bounds the log-likelihood of observed data, and the gap vanishes at the global optimum.
Often, the density forms of p ( z ) and q φ ( z | x ) are chosen such that their KL-divergence can be written analytically in a closed-form expression (e.g., p ( z ) is N (0 , I ) and q φ ( z | x ) is N ( µ φ ( x ) , Σ φ ( x )) )
(Kingma & Welling, 2013).
In such cases, the
ELBO can be efﬁciently optimized (to a stationary point) using stochastic ﬁrst order methods where both expectations are estimated using mini-batches.
Further, in cases when q φ ( · ) can be written as a continuous transformation of a ﬁxed base distribution (e.g., the standard normal distribution), a low variance estimate of the gradient over φ can be obtained by coordinate transformation (also referred as reparametrization)
(Fu, 2006; Kingma & Welling, 2013; Rezende et al., 2014).
2. G ENERATIVE STORY : DISENTANGLED PRIOR
Most VAE based generative models for real datasets (e.g., text, images, etc.)
already work with a relatively simple and disentangled prior p ( z ) having no interaction among the latent dimensions (
.g.
, the standard Gaussian N (0 , I ) )
(Bowman et al., 2015;
Miao et al., 2016;
Hou et al., 2017; Zhao et al., 2017).
The complexity of the observed data is absorbed in the conditional distribution p θ ( x | z ) which encodes the interactions among the latents.
Hence, as far as the generative modeling is concerned, disentangled prior sets us in the right direction.
2. I NFERRING DISENTANGLED LATENTS
Although the generative model starts with a disentangled prior, our main objective is to infer disentangled latents which are conducive for various goals mentioned in Sec.
1 (e.g., invariance, transferability, interpretability).
To this end, we consider the density over the inferred latents induced by the approximate posterior inference mechanism,
[Equation]
which we will subsequently refer to as the inferred prior .
For inferring disentangled factors, this should be factorizable along the dimensions, i.e., q φ ( z ) =
(cid:81)
i q i ( z i ) , or equivalently q i | j
( z i | z j ) =
q i ( z i ) , ∀ i, j .
This can be achieved by minimizing a suitable distance between the inferred prior q φ ( z ) and the disentangled generative prior p ( z ) .
If we take KL-divergence as our choice of distance, by relying on its pairwise convexity (i.e., KL ( λp 1 + (1 − λ ) p 2 (cid:107) λq 1 + (1 − λ ) q 2 ) ≤ λ KL ( p 1 (cid:107) q 1 ) + (1 − λ ) KL ( p 2 (cid:107) q 2 ) ) (Van Erven & Harremos, 2014), we can show that this distance is upper bounded by the objective of the variational inference (the ELBO in Eq.
[Equation]
Hence, variational posterior inference of latent variables with disentangled prior naturally encourages inferring factors that are disentangled.
We think this is the reason that the original VAE (Eq.
(1) has also been observed to exhibit some disentangling behavior on simple datasets such as MNIST (Kingma & Welling, 2013).
However, this behavior does not carry over to more complex datasets (Aubry et al., 2014; Liu et al., 2015; Higgins et al., 2017), unless extra supervision on the generative factors is provided (Kulkarni et al., 2015; Karaletsos et al., 2015).
This can be due to non-convexity of the ELBO objective which prevents us from achieving the global minimum of E x KL ( q φ ( z | x ) (cid:107) p ( z | x ))
=(which would imply KL ( q φ ( z ) (cid:107)
p ( z )) = 0 ).
In other words, maximizing the ELBO (Eq.
(1)) might also result in reducing the value of KL ( q ( z ) (cid:107) p ( z )) , however, due to the non-convexity of the loss surface of the ELBO, the gap between KL ( q φ ( z ) (cid:107) p ( z )) and E x KL ( q φ ( z | x ) (cid:107) p ( z | x )) could be large at the stationary point of convergence.
Hence, minimizing KL ( q φ ( z ) (cid:107) p ( z )) explicitly will give us better control on the disentanglement.
This motivates us to add KL ( q φ ( z ) (cid:107) p ( z )) as part of the objective to encourage disentanglement during inference, i.e.,
[Equation]
where λ controls its contribution to the overall objective.
Optimizing (4) directly is not tractable due to the presence of KL ( q φ ( z ) (cid:107) p ( z )) which does not have a closed-form expression.
One possibility is use the variational formulation of the KL-divergence (Nguyen et al., 2010; Nowozin et al., 2016) that needs only samples from q φ ( z ) and p ( z ) to estimate a lower bound to KL ( q φ ( z ) (cid:107) p ( z )) .
However, this would involve optimizing for a third set of parameters ψ for the KL-divergence estimator, and would also change the optimization to a saddlepoint (min-max) problem which has its own challenges (e.g., gradient vanishing as encountered in training generative adversarial networks with KL or JS divergences (Goodfellow et al., 2014;
Arjovsky & Bottou, 2017)).
Replacing KL ( q φ ( z ) (cid:107) p ( z )) with another suitable distance between q φ ( z ) and p ( z ) (e.g., integral probability metrics like Wasserstein distance (Sriperumbudur et al., 2009)) might alleviate some of these issues (Arjovsky et al., 2017) but will still involve complicating the optimization to a saddle point problem in three set of parameters 3 .
It should also be noted that using these variational forms of the distances will still leave us with an approximation to the actual distance.
We adopt a simpler alternative of matching the moments of the two distributions.
In particular, we match the covariance of the two distributions which will amount to decorrelating the dimensions of z ∼ q φ ( z ) if p ( z ) is N (0 , I ) .
Let us denote Cov q ( z ) ( z ) :
=
E q ( z ) (cid:2) ( z − E q ( z ) ( z ))( z − E q ( z ) ( z )) (cid:62) (cid:3) .
By the law of total covariance, the covariance of z ∼ q φ ( z ) is given by
[Equation]
where E q φ ( z | x ) ( z ) and Cov q φ ( z | x ) ( z ) are random variables that are functions of the random variable x
( z is marginalized over).
Most existing work on the VAE models uses q φ ( z | x ) having the form
N ( µ φ ( x ) , Σ φ ( x )) , where µ φ ( x ) and Σ φ ( x ) are the outputs of a deep neural net.
In this case
Eq.
(5) reduces to Cov q φ ( z ) ( z ) =
E p ( x ) Σ φ ( x ) + Cov p ( x ) ( µ φ ( x )) , which we want to be close to an identity matrix.
For simplicity, we choose entry-wise squared (cid:96) 2 -norm as the measure of proximity.
However, as the entanglement is mainly reﬂected in the off-diagonal entries of this matrix, we opt for two separate hyperparameters controlling the relative importance of the loss on the diagonal and off-diagonal entries.
This gives rise to the following optimization problem for inference:
[Equation]
The regularization terms involving Cov p ( x ) ( µ φ ( x )) in the above objective (6) can be efﬁciently optimized using SGD.
We maintain a running estimate of Cov p ( x ) ( µ φ ( x )) which is updated with every minibatch of x ∼ p ( x ) .
The gradient for the current minibatch can be computed by treating the previous estimate of Cov p ( x ) ( µ φ ( x )) as constant.
2. C OMPARISON WITH β -VAE
Recently proposed β -VAE (Higgins et al., 2017) proposes to modify the ELBO by upweighting the
KL ( q φ ( z | x ) (cid:107) p ( z )) term in order to encourage the inference of disentangled factors:
[Equation]
where β is taken to be great than 1 .
Higher β is argued to encourage disentanglement at the cost of reconstruction error (the likelihood term in the ELBO).
Authors report empirical results with β ranging from 4 to 250 depending on the dataset.
As already mentioned, most VAE models proposed in the literature, including β -VAE, work with N ( 0 , I ) as the prior p ( z ) and N ( µ φ ( x ) , Σ φ ( x )) with diagonal Σ φ ( x ) as the approximate posterior q φ ( z | x ) .
This reduces the objective (7) to
[Equation]
For high values of β , β -VAE would try to pull µ φ ( x ) towards zero and Σ φ ( x ) towards the identity matrix (as the minimum of x − ln x for x > 0 is at x = 1 ), thus making the approximate posterior q φ ( z | x ) insensitive to the observations.
This is also reﬂected in the quality of the generated samples which is worse than VAE ( β = 1 ), particularly for high values of β .
Our proposed method does not have such increased tension between the likelihood term and the disentanglement objective, and the sample quality with our method is on par with the VAE.
Finally, we note that both β -VAE and our proposed method encourage disentanglement of inferred factors by pulling Cov q φ ( z ) ( z ) in Eq.
(5) towards the identity matrix: β -VAE attempts to do it by making Cov q φ ( z | x ) ( z ) close to I and E q φ ( z | x ) ( z ) close to 0 individually for all observations x , while the proposed method directly works on Cov q φ ( z ) ( z ) (marginalizing over the observations x ) which retains the sensitivity of q φ ( z | x ) to the conditioned-upon observation.
E XPERIMENTS
We evaluate our proposed method, referred as DIP-VAE subsequently (Disentangled Inferred Prior), on three datasets – (i) CelebA
(Liu et al., 2015): It consists of 202 , 599 RGB face images of celebrities.
We use 64 × 64 × 3 cropped images as used in several earlier works, using 90% for training and 10% for test.
(ii) 3D Chairs
(Aubry et al., 2014):
It consists of 1393 chair CAD models, with each model rendered from 31 azimuth angles and 2 elevation angles.
Following earlier work (Yang et al., 2015; Dosovitskiy et al., 2015) that ignores near-duplicates, we use a subset of
[Table]
chair models in our experiments.
We use the binary masks of the chairs as the observed data in our experiments following (Higgins et al., 2017).
First 80% of the models are used for training and the rest are used for test.
(iii) 2D Shapes (Higgins et al., 2017):
This is a synthetic dataset of binary 2D shapes generated from the Cartesian product of the shape (heart, oval and square), x -position ( values), y -position (32 values), scale (6 values) and rotation (40 values).
We consider two baselines for the task of unsupervised inference of disentangled factors: (i) VAE (Kingma & Welling, 2013;
Rezende et al., 2014), and (ii) the recently proposed β -VAE (Higgins et al., 2017).
To be consistent with the evaluations in (Higgins et al., 2017), we use the same CNN network architectures (for our encoder and decoder), and same latent dimensions as used in (Higgins et al., 2017) for CelebA, 3D
Chairs, 2D Shapes datasets.
Hyperparameters.
For the proposed DIP-VAE, in all our experiments we vary λ
d in the set { 1 , 2 , 5 , 10 , 20 , 50 } while ﬁxing λ d = 10 λ od .
For β -VAE, we experiment with β = { 1 , 2 , 4 , 8 , 16 , 25 , 32 , 64 , 100 , 128 , 200 , 256 } (where β = 1 corresponds to VAE).
For both CelebA and 2D Shapes, we show the results for the best performing models in terms of the disentanglement metric score which was introduced in (Higgins et al., 2017).
For 3D Chairs data, this metric is at 100% for almost all models and we pick the models based on our subjective evaluation of the reconstruction quality and disentanglement.
Disentanglement metric score and reconstruction error.
Higgins et al.
(2017) propose a metric to evaluate the disentanglement performance of the inference mechanism.
It assumes that the ground truth generative factors are available.
It works by ﬁrst sampling a generative factor y , followed by sampling L pairs of examples such that for each pair, the sampled generative factor takes the same value.
Given the inferred z x := µ φ x for each example x , they compute the absolute difference of these vectors for each pair, followed by averaging these difference vectors.
This average difference vector is assigned the label of y .
By sampling n such minibatches of L pairs, we get n such averaged difference vectors for the factor y .
This process is repeated for all generative factors.
A low capacity multiclass classiﬁer is then trained on these vectors to predict the identity of the corresponding generative factor.
In all our experiments, we use a one-vs-rest linear SVM with weight on the hinge loss C set to 0 .
01 and weight on the regularizer set to 1 .
Higgins et al.
(2017) argue that this
[Table] [Table]
[Table] [Table]
[Table] [Table]
metric captures the disentangled property of the inferred latents reasonably well.
Table 1 shows the disentanglement metric scores along with reconstruction error (which directly corresponds to the data likelihood) for the test sets for CelebA and 2D Shapes data.
It is evident that the proposed DIP-VAE outperforms β -VAE both in terms of the disentanglement metric score and reconstruction error.
Further we also show the plot of how the disentanglement metric changes with the reconstruction error as we vary the hyperparameter for both methods ( β and λ od , respectively).
It is clear that the proposed method gives much higher disentanglement metric score at little to no cost on the reconstruction error when compared VAE ( β = 1 ).
The reconstruction error for β -VAE gets much worse as β is increased.
Binary attribute classiﬁcation for CelebA. We also experiment with predicting the binary attribute values for each test example in CelebA from the inferred µ φ ( x ) .
For each attribute k , we compute the attribute vector w k =
[Equation]
) from the training set,
and project the µ φ ( x ) along these vectors.
A bias is learned on these scalars (by minimizing hinge loss) which is then used for classifying the test examples.
Table 2 shows the results for the attribute which show the highest change across various methods (most other attribute accuracies do not change).
The proposed DIP-VAE outperforms both VAE and β -VAE for most attributes.
The performance of β -VAE gets worse as β is increased further.
Correlations in the inferred latent space.
We visualize the Pearson’s correlations between dimensions of the inferred latents µ φ ( x ) .
We also visualize the correlations between inferred latent dimensions and the ground truth attributes which can be taken as a proxy for true generative factors.
Tables 3 and 4 show these correlations for VAE, β -VAE and the proposed DIP-VAE.
We observe less correlations between inferred latents for DIP-VAE.
Latent correlations for β -VAE are even higher than those for VAE which seems to be going against the objective of disentanglement.
This can be due to the fact that β -VAE gives less weight to the ELBO which implies the KL ( q φ ( z ) (cid:107) p ( z )) does not get minimized as well (see Eq.
Indeed Eq.
(8) indicates that β -VAE gives more importance to minimizing the (cid:96) 2 -norm of individual posterior means µ φ ( x ) and does not minimize their correlations as DIP-VAE.
To further analyze the latent correlations with the ground truth attributes, we pick one attribute at a time, and pick the latent dimension that has highest correlation with that attribute.
Then we plot the correlations of this latent dimension with rest of the attributes (picking top 4 correlations).
This will indicate the entanglement of a latent dimension with the ground truth attributes.
It should be noted that many attributes in CelebA are naturally entangled (e.g., lipstick, heavy-makeup and gender) which can be quite difﬁcult to disentangle.
These plots for two of the attributes are shown in Fig.
Finally we also show the generated samples by varying one latent dimension at a time for CelebA and 3D Chairs data for all three methods in Tables 5 and Tables 6.
[Table] [Table]
[Figure]
[Table] [Table]
D ISCUSSION
We proposed a principled variational framework with infer disentangled latents from unlabeled observations.
Unlike β -VAE our variational objective does not have any conﬂict between the data log-likelihood and the disentanglement of the inferred latents, which is reﬂected in our empirical results that outperform β -VAE.
Directions for future work include tackling the sampling bias in the generative process which makes the problem challenging (e g., sampling the male gender makes it likely to sample beard ), and effective use of disentangled representations in transfer learning.
Appendix A C ORRELATIONS BETWEEN INFERRED LATENT FACTORS AND GROUND TRUTH ATTRIBUTES FOR C ELEB A
[Figure]
[Figure] [Figure]
L ABEL CORRELATIONS FOR C ELEB A AND 2D S HAPES DATASETS