Improving the Reconstruction of Disentangled Representation Learners via Multi-Stage Modelling
Authors: Akash Srivastava
Yamini Bansal
Yukun Ding
Cole Lincoln Hurwitz
Kai Xu
Bernhard Egger
Prasanna Sattigeri
Joshua B. Tenenbaum
Dan Gutfreund
Venue:  (2021)
Keywords: disentanglement
disentangled representation learning
vae
generative model

Abstract: Current autoencoder-based disentangled representation learning methods achieve disentanglement by penalizing the (aggregate) posterior to encourage statistical independence of the latent factors. This approach introduces a trade-off between disentangled representation learning and reconstruction quality since the model does not have enough capacity to learn correlated latent variables that capture detail information present in most image data. To overcome this trade-off, we present a novel multi-stage modelling approach where the disentangled factors are first learned using a preexisting disentangled representation learning method (such as $\beta$-TCVAE); then, the low-quality reconstruction is improved with another deep generative model that is trained to model the missing correlated latent variables, adding detail information while maintaining conditioning on the previously learned disentangled factors. Taken together, our multi-stage modelling approach results in single, coherent probabilistic model that is theoretically justified by the principal of D-separation and can be realized with a variety of model classes including likelihood-based models such as variational autoencoders, implicit models such as generative adversarial networks, and tractable models like normalizing flows or mixtures of Gaussians. We demonstrate that our multi-stage model has much higher reconstruction quality than current state-of-the-art methods with equivalent disentanglement performance across multiple standard benchmarks.

Under review as a conference paper at ICLR 2021
IMPROVING THE RECONSTRUCTION OF DISENTAN- GLED REPRESENTATION LEARNERS VIA MULTI- STAGE MODELLING
Anonymous authors Paper under double-blind review
ABSTRACT
Current autoencoder- based disentangled representation learning methods achieve disentanglement by penalizing the (aggregate) posterior to encourage statistical independence of the latent factors. This approach introduces trade-off between disentangled representation learning and reconstruction quality since the model does not have enough capacity to learn correlated latent variables that capture detail information present in most image data. To overcome this trade-off, we present novel multi- modelling approach where the disentangled factors are first learned using a preexisting disentangled representation learning method (such as 3-TCVAB); then, the low-quality reconstruction is improved with another deep generative model that is trained to model the missing correlated latent variables, adding detail information while maintaining conditioning on the previously Iearned disentangled factors. Taken together, our multi-stage modelling approach results in single, coherent probabilistic model that is theoretically justified by the principal of D D-separation and can be realized with a variety of model classes including likelihood-baa models such as variational autoencoders, implicit models such as generative adversarial networks, and tractable models like normalizing flows or mixtures of Gaussians. We demonstrate that our multi-stage model has much higher reconstruction quality than current state-of- offte-of-the-art the methods with equivalent disentanglement performance across multiple standard benchmarks.
INTRODUCTION
Deep generative models (DGMs) such as variational autoencoders (VAEs) (Kingma & Welling 2014 Rezende et al. 2014) and generative adversarial networks (GANs) (Goodfellow et al.) (2014) have recently enjoyed great success modeling high dimensional data such as natural images. As the name suggests, DGMs leverage deep learning to model a data generating process. The underlying assumption is that the high dimensional observations X E R can be meaningfully described by a small set of latent factors H E IR K where K < D. More precisely, the observation (X- = x) is of low dimensional factors h from simple prior
assumed to be generated by first sampling distribution p( ()) and then sampling x ~ po X (h). DGMs realize po through a deep neural network also known the decoder or the generative network. VAE-based DGMs use another deep neural network (called the encoder or the inference network) to parameterize an approximate posterior (1 ()). Learning the variational posterior parameters is done by maximizing an evidence lower bound (ELBO) to the log- marginal likelihood of the data under the model (X).
Learning disentangled factors h qs (H (z) that are semantically meaningful representations of the observation x is highly desirable because such interpretable representations can arguably be advantageous for a variety of downstream tasks (Locatello al., 2018)), including classification, detection, reinforcement learning, transfer learning and image synthesis from textual descriptions (Bengio et al. 2013 LeCun et al.) 2015 Lake et al., 2017 van Steenkiste et al., 2019 Reed et al., 2016 Zhang et al. 2016). While formal definition of disentangled representation (DR) remains elusive, it is understood to mean that by manipulating only one of the factors while holding the rest constant, only one semantically meaningful aspect of the observation (e.g. the pose of an object in an image) changes. Prior work in unsupervised DR learning focuses on the objective of learning statistically independent latent factors means for obtaining DR. The underlying assumption is

---

Under review as a conference paper at ICLR 2021
(a) Target
(b) 3-TCVAE Mc) BS- VAEE
Figure I: Image reconstruction using 3-TCVAE (Figure lb and MS- VAE (Figure Icp. MS- VAE is able to take the blurry output of the underlying 3-TCVAE model and learn render much better approximation of the target while maintaining the pose of the original image (Figure Ta).
that the latent variables H can be partitioned into independent components C' (i.e. the disentangled factors) and correlated components Z. An observation (X x:) is assumed to be generated by
sampling the low dimensional factors h = (c,2) from p H) and then sampling x ~ Po(X 2) (Figure 2ap. series of works starting from (Higgins al. 2017 enforce statistical independence of
the latent factors H via regularization, up-weighting certain terms in the ELBO which penalize the (aggregate) posterior be factorized over all or some of the latent dimensions (Kumar al. 2017
Kim & Mnih 2018 Chen et al. 2018) (see Section for details).
While the aforementioned models show promising results, they suffer from a trade-off between DR learning and reconstruction quality. If the latent space is heavily regularized - not allowing enough capacity for the correlated variables then the reconstruction quality will be diminished, signaling that the learned representation is sub-optimal. As the correlated variables are functionally ignored with high levels of regularization, an observation (X = 2) can be thought to be generated by only sampling independent latent factors from p(C) and then sampling a departure from the original data generating hypothesis that x is sampled from a distribution dependent on both the independent and correlated latent variables. On the other hand, if the correlated variables are not well-constrained, the model can use them to achieve a high quality reconstruction while ignoring the independent variables (the disentangled latent factors). This phenomena is referred to as the "shortcut problem" " and has been discussed in previous works (Szabó et al.) 2018 Lezama 2018). Overcoming the aforementioned trade-off is essential for using these models in real world applications such as realistic, controlled image synthesis (Lee et al. 2020 Lezama 2018).
pg Po[Xee (e) (Figure 2b). This
In this paper, we propose a new graphical model for DR learning (Figure 2ch that allows for learning disentangled factors while also correctly realizing the data generating hypothesis that an observation generated from independent and correlated factors. Importantly. the graphical model Figure ic ---spaaraeed, meaning that any changes in the correlated latent variables Z will not influence the independent latent variables C. Generating an observation (X xr) from this model can then be done by sampling the independent factors from p(C), sampling a low low-quality reconstruction y ~ pY (c), sampling the correlated factors z from pl (Z), and then finally sampling x ~ po (X [z, Y). The final reconstruction depends both on and however, any regularization needed to extract the independent factors c no longer diminishes the model capacity for the correlated factors 2.
To realize our proposed graphical model, we introduce MS- VAE, a multi-stage DGM that is imple- mented as follows: first, the disentangles led representation C' is learned using an existing DR learning method such 3-TCVAE (Chen et al. 2018). Since the learned factors C' are regularized to be statis- tically independent- not allowing enough capacity for correlated factors - the final reconstruction Y will have diminished reconstruction quality. Then, we train another DGM to improve the low-quality reconstruction Y by learning the missing correlated factors Z. This is achieved during training by inputting the reconstruction
into the decoder of the second DGM and then modulating the hidden activation of each layer as function of latent factor Z (using Feature- wise Linear Modulation (Perez et al. (0188). Through this training paradigm, MS- VAE is able to preserve conditioning on the disentangled factors while dramatically improving the reconstruction quality. A schematic of MS- VAE is shown in Figure 2d and example images from each stage are shown in Figure The reconstruction from 3-TCVAE (lb) is improved by VS-VAE (Ic) to better approximate the target
while maintaining conditioning on the disentangled factors, e.g. azimuth.
To summarize contributions:
2

---

Under review as a conference paper at ICLR 2021
We propose graphical model for DR learning (Figure 2c) that, due the D D-separation between its independent and correlated latent variables, alleviates the disentanglement vs. reconstruction trade-off.
We introduce MS- VAE, multi-stage DGM that implements our proposed graphical model, achieving state-of-the-arr reconstruction quality while maintaining the same level of disen- tanglement as preexisting methods. MS-VAB can be implemented using variety of DGMs such as VAEs, GANs, or FLOW-based models and does not depend on hand-crafted loss functions or additional hyperparameters (unlike current state-of-the fi-theeart art (Lezama 2018
Following (Locatello al. 2018), we test our framework with wide-range of qualitative and quantitative tests to demonstrate the efficacy of the framework. We provide all code used for the experiments.
BACKGROUND
VAEs, first introduced in (Kingma & Welling 2014), represent a class of likelihood- based, DGMs comprised of an encoder and a decoder that are trained using amortized variational inference to maximize the expected ELBO under the data distribution,
f p(z))oogg) logpo(dd)) 2 Ep(x)) (k))))))) (aaaak/))llloooool (x))))())))) - (L(((())))]l((M)]]
(I)
The first term is the expected reconstruction loss under the variational posterior ls. The second term the Kullback-Leible (KL) divergence underpins most of the recent DR learning methods.
The KL term regularizes the variational posterior to be similar (in expectation) to a simple spherical Gaussian prior over H. |Higgins et al. (2017) attempts to factorize the variational aggregate posterior by up-weighting this KL-divergence term with Lagrange multiplier 8. Follow-up works decompose the KL-termm in equation| [ and up- up-weight only the KL-divergence between the aggregate posterior
Zdg (1)
p( [dr and the factorized aggregate posterior (Kim & Mnih 2018 Chen
jet al. 2018 Jeong & Song 2019). Kumar al. (2017) simply uses factorized prior.
These models introduce trade-off between unsupervised DR learning and reconstruction quality. When the latent variables are heavily regularized to be independent, an observation (X = x) can be thought to be generated by sampling independent latent factors from p(C.) and then sampling x ~ Po (oo (Xc) (Figure 2b). This leads to low-quality reconstructions as the correlated factors are not utilized. Lezama 2018) attempts to tackle this issue with teacher-student learning technique. First, the teacher model is trained to extract disentangled factors at the cost of reconstruction quality (similar the first stage of our model). Then, student autoencoder (AE) model with larger latent dimension size is trained with additional losses that force the Jacobian of the final reconstruction with respect to the disentangled factors to remain the same (a.k.a Jacobian supervision). While this model achieves promising results on MNIST (Deng) 2012). it suffers from few drawbacks including having to relearn the disentangled factors in the second stage and having no theoretical guarantees that changes in the correlated factors Z will not affect the independent factors C' in the student model. Other works have also been proposed to overcome the trade-off using adversarial training (Mathieu et al. 2016 Lample et al. 2017 Perarnau al. 2016 Szabó et al. 2018 Hu et al. 2018).
3
MS-VAE
3.1
MS-VAE GRAPHICAL MODEL
Lower Bound. To alleviate the trade-off between reconstruction and disentanglement, we propose a new graphical model for DR learning (Figure 2c). To understand this model, let us assume that the data come as pairs of images {yi, Z, N Let the x's be the ground truth images (generated using both the independent factors c and correlated factors 2) and let the y y's be the approximations of those images generated using only c. Given this paired data, the graphical model depicted in Figure 2c represents a single coherent model of the data that defines a joint distribution over and X. The log-marginal likelihood of the observations under this model can be lower-hounded as follows,
3

---

Under review as a conference paper at ICLR 2021
10gP(089220)(1..))) 2 -Kllgoccwllp()] + Ens log pe (ulc)) - KL (z|2, v))[((2]] + Ege [lo pe (z)22 v)].
(2)
(2) C-
(X
In practice, we do not have access to Y. However, by learning the disentangled factors C' from X, we first use ofate-of-the-art DR learning method to learn the sub-graph C' + Y by maximizing term (a) in Equation 2o produce Y. Then, we train another DGM to learn the sub-graph (Y, 2) + X by maximizing term (b) in Equation 2 This two-stage training procedure approximates the graphical model from Figure 2c
can generate Y. Therefore,
D-separation. Now that we introduced and lower-bounded the graphical model from Figure 2c it is important to understand the theoretical motivation for realizing this model. Recall that our goal is to take an existing method for learning the disentangled factors C and learn additional residual information in Z without changing the conditioning on C. In a standard VAE graphical model (Figure 2ap this could occur because there is no guarantee that p(Z, CIX) p(Z] ()()))))) By introducing an observed node Y (i.e. D-Separation) in our graphical model, we ensure that changes in Z will not influence C. This is because the joint distribution p(Z, C|X, Y) as p(C)Y) ) (2),, Y). In practice, realizing the graphical model in Figure 2c using deep neural network requires care as we want to ensure that the conditioning on Y is stronger than that of Z so that Z only captures the residual correlated information. Also, simply training a single neural network end-to-end approximate this graphical model is infeasible Y will no longer act as an observed node, leading to the entanglement of C and Z. We demonstrate this phenomenon by performing an ablation study of MSVAE (Appendix L).
now be broken down
3.2 MS-VAE IMPLEMENTATION
In the following sections, we describe how our MS-VAE implements the graphical model in Figure 2c using a multi-stage modelling approach.
Stage one: Learning a Disentangled Representation. In stage one, we train a 3-TCVAE (Chen ct al. 2018) to learn the independent latent factors C. As 3-TCVAE heavily regularizes the latent space to be statistically independent not allowing enough capacity for the correlated factors the final reconstruction Y will be low-quality. Stage one can be realized using any DR learning method, however, we chose 3-TCVAE in our work because it has been shown to perform well across all standard benchmarks and its total correlation penalty is simpler to compute in comparison to FactorVAE's (Kim & Mnih] 2018). We use the standard convolutio and transpose-coovolltionss based realization of 3-TCVAE as provided in the disentangiement _lib (Locatello et al. 2018) package.
Stage Two: Improving the Reconstruction. In stage two, we learn the correlated factors Z with another DGM and then use them to improve the reconstruction Y while maintaining the conditioning on the independent factors C. To this end, we train a conditional VAE that models the data X given Y and Z. The encoder network Røz learns the posterior distribution over Z as a function of the residual of X and Y. The observation model Ge then reconstructs X given Y and Z. Incorporating Z into Ge requires care as simply inputting the concatenation of Y and Z into the network at the start of training will result in non-linear entanglement between Y and Z. Once entangled this way, Ge will fail to condition on sufficiently and use the entangled representation as a whole towards the reconstruction of X. To overcome this, we induce an architectural bias with a Feature-wist Linear Modulation (FiLM) technique (Perez et al.) 2017) prevalent in style transfer literature (Huang & Belongie, 2017), introducing Z only through the adaptive instance normalization (AdaIN) layers of Ge- By inputting only Y into the decoder and then incorporating Z through AdaIN at each layer, the network is able to use 7 later in the generative process model the residual information (as explained in (Dumoulin et al.) 2018)). )). This allows the network to better utilize the information in Y. We demonstrate the importance of AdaIN by performing an ablation study of MSVAE (Appendix L).
conditioning tech- nique for deep neural networks (Perez et al. 2018). A FiLM layer learns an affine transformation of the intermediate statistics of a deep neural network, conditioned on an external input. FiLM has found great success in style transfer literature where the goal is to render content from a source domain in the style of a target domain. One such method, Instance Normalization (Ulyanov et al. 2016), incorpo- rates the style and the content information through the normalization layers in a feed- forward network,
FiLM and AdaIN. Recently, FiLM have been proposed as a general-purpos
4

---

Under review as a conference paper at ICLR 2021
learning style-specific affine parameters from the data that used to shift and scale the feature statistics (per instance) extracted from the source image across spatial dimensions, HHann & Belongie (2017) extended this method to address arbitrary styles by taking the affine parameters to be the fea-
ture statistics extracted from the style domain, arriving at the AdoIN (y. x) oxx)
( ) (())
which is applied for each normalization layer. Here, y are the incoming activations for normalization layer in the network that is receiving image from the source domain and æ are the style features computed from the target domain. The features statistics for y (as well as for for x) are computed
as follows: yu(y) 1. En. Onhwre and (y) = V
1 En. (Inhwe x(v))² c, where h, w, c HW stand for the height, width and channel, respectively. In MS- MSSVVAE, the equivalent of style features is the representation the residual information between X and Y captured the correlated variables Z. The goal is similar to style transfer: incorporate the information stored in Z while maintaining the semantic content stored in Y. Therefore, we slightly modify the AdaIN formulation as follows:
AdaI N (y, 2) (2)
+ 3(2). where p(y) and (y) are the same as before and 7(2) and B(2) are learned functions of parameterized by callycconnected neural network.
(=-(())
((() (vv)
MS-GAN and MS-FLOW. While the above section only described the VAE-based implementa- tion of our method, we can also realize these techniques using GANs or FLOW- -osed models. We
defer these implementations to Appendix and D
C
Z
Z
Y
X
X
hiroot
N
N
(a) VAE
(b) 3-TCVAE
(c) MS VAE
(d) MAE Architecture
Figure 2: (a) Graphical model of a standard VAE where C and Z are not independent conditioned on X. (b) Graphical model of 3-TCVAE where the reconstruction only depends the independent latent factors C. (c) MS- VAE graphical model where C' and Z are independent conditioned on Y. (d) Schematic for MS- VAE when implemented convolutional architecture. Both Y and X the reconstructions of the same image.
Deeper MS-VAE. In this work, focus generative models with only two levels of hierarchy and two sets of latent factors C and Z. However, MS- MS-VAE can be easily extended to graphical models with deeper or more structured hierarchies. We present concrete example of this extension of MS-VAE in appendix C where we use MS-V VAE to model a simple 2D pendulum.
EXPERIMENTS
Baseline Models Locatello et al. (2018) demonstrated that most trate-of-the-ae DR learning methods are able to effectively learn a factorized posterior distribution (i.e. extract independent latent variables C). Therefore, chose 3-TCVAE as the baseline DR learning method for our paper. For fair comparison with our work, we use 3-TCVAE to implement the sub-graph C -> Y MS- VAE. For our main experiments, we thoroughly evaluate three models: (i) MS- VAE with latent dimensions C' for the C' + Y network and Z for the (Y, 2) - X network, (i) 8-TCVAE with latent dimension size |C), and (iii) 3-TCVAE-L with latent dimension size |C + Z). In line with previous work (Locatello et al., 2018), we use 10-dimensional C. We use a 5-dimensional and 10-dimensional Z for Cars3D and Smal/Norb, respectively. In Appendix A we sweep over the dimensionalities of C' and Z to evaluate its effect on the reported results. We train 3-TCVAE with 3 values ranging from 1 to 10. addition to the main models, we include comparisons to a recent state-of state-of-the- of-the-art method, Jacobian Supervision (Lezama 2018). Although Jacobian supervision can be implemented with multiple student autoencoders, we only compare to the single student model version it is most directly comparable to MS- VAE. We utilize the same architecture and C, Z
5

---

Under review as a conference paper at ICLR 2021
dimensionalities VAE. For their teacher model, we also use a pre-trained B-TCVAE model as we did for VAE, but we only evaluate their method for B = 1. 10. The hyperparameters and loss functions are specified in the paper and are available in the linked code repository.
Measuring Disentanglement Performance [ocatello al. 2018) concluded that most metrics for measuring disentanglement performance correlate relatively well with each other. Therefore, we only use MIG (Chen al. 2018 quantitatively measure disentanglement performance. Please see Appendix G for how this is computed.
Mcasuring Reconstruction Quality. Following (Razavi et al. 2019), we report the Frechet Incep- tion Distance (FID) (Heusel al. 2017) between the dataset and each models' reconstructions. We choose FID instead of reconstruction enror for our main metric as L1/L2 metrics are known to not correlate well with perceived image quality (Wang et al.) 2003 2004). We still report both LI and L.2 results in Appendix ! Additionally, FID does not require a well-defined likelihood function and can, therefore, be used to evaluate MS- VAE when the conditional model is implemented using a GAN.
Mcasuring the Conditioning on C. To quantitatively measure how well the improved reconstruc- tions are conditioned on the independent latent variables C. we use the following procedure. Let the mean parameter of the encoder of the subgraph C + Y be Eg. Further, let CR = (X). Fa ()))) and Cy. E (G,(E,, Z))) where X is an input image, e ~ (0, I), and Ge is the decoder of the second stage of MAS. We report the following mutual information terms: Mj MI[CR; Cy M2 MI[CR; Cx M3 MI[[R;; C..l. and M4 = MI[CR; Cy.]- If Z is encoding meaningful information for the reconstruction process then we should see M2 >> M1 as the reconstruction should better approximate the ground truth image X. More importantly, if M1 > M3 >> M1 then the conditioning on Y (and therefore C'' can be assumed be maintained. This follows because if Ge was only using Z for the final reconstruction, M3 would be close to zero and M4 would be high (since Y is not being used in the final reconstruc- tion). In addition to this quantitative analysis, we provide latent traversal plots for all the datasets to qualitatively show that the conditioning on the disentangled factors C' is preserved.
Cy F (Y). Cx (Y, 2)), Cz.
Downstream Classification To demonstrate that Z is learning meaningful residual information that is complimentary to C. we perform a set of downstream classification tasks using both Z and C. In these tasks, we predict attributes of objects in the SmallNORB dataset using MLP that is trained with three different inputs: C, Z, and both C' and Z concatenated together. Ideally, the MLP trained with both C' and 7 would achieve the highest accuracy the inferred latent variables contain complementary information about the object attributes.
Code and Datasets. We use the disentanglement lib package (occttell et al. 2018 to train all 3-TCVAE models and to evaluate MIG. The datasets we use for our quantitative benchmarks are Cars3D (Reed et al., 2015) and SmalINORB (LeCun et al. 2004). These datasets are a subset of those used in the large scale study of (Locatello al. 2018) and are available (with ground truth factors) from the disentanglement lib package. Additionally, we use the CelebA (Liu et al.) dataset of real face images for qualitative evaluations since the other datasets consist of only simulated images (see Appendix B). For FID evaluations, we use the standard Tensorflow (Abadi et al., (016) implementation for all the models. We provide all the hyperparameter and architectural details for the experiments in the appendix. A reference Tensorflow inplementation is also available at hhttpps //github. http///oonnouuuuutrrsoooo DSD VAΣ in the form of a Jupyter notebook.
4.1 RESULTS
Disentanglement Performance. As mentioned above, MS-VAB utilizes 3-TCVAB to learn disen- tangled factors in the first stage of its training. Therefore, we report the MIG values for the 3-TCVAE used with MS-VAE and for 3-TCVAE-L (a 3-TCVAE with the same latent dimension as MS- VAE.. For the Cars3l dataset, we find that with both models, the MIG scores increase as / is increased (Figure 4a ii)). However, 3-TCVAE appears to consistently have a higher MIG than 3-TCVAE-L, potentially reflecting the entangling effect of an overly large latent space. For the SmalINORB dataset, MIG decreases as the value of 33 increases shown in Figure ti). This confirms the findings of Locatello et al. (2018) on this particular dataset. Overall, these results illustrate that MS- VAE is able to extract disentangled factors from image data at the same level state-of-the- methods and that naively increasing the latent dimensionality of 3-TCVAE may diminish disentanglement performance. For the Jacobian supervision baseline (Lezama 2018 we use 3-TCVAE at the teacher

---

Under review as a conference paper at ICLR 2021
(a) Reconstruction Cars3D. Starting from top left, each triplet consists (Target, 3-TCVAE and MS-
(b) Reconstruction SmalINORB. Starting from top left. each triplet consists (Target, 3-TCVAE and MS- VAE) images.
VAB) images.
(c) Traversal on Cars3D. Odd row 3-TCVAE, MS- VAE.
(d) Traversal SmalNNORB. Odd row 8-TCVAE, MS- VAE.
(e) 7. perturbation on Cars3D
Figure 3: Qualitative results on the Carx3D and Smal/NORB datasets for 3-TCVAE and MS- VAE.
ATTCVEE
0.23
2001
100
F10
MIC
6.2
00
0.15
01
0
(a) Cars3D dataset
(b) SmalINORB dataset
Figure 4: FID (lower is better) and MIG (higher is better) comparison of 3-TCVAE, 3-TCVAE-L. and MS-VAE models. On both datasets, MS- VAE is able to consistently improve the reconstruction quality its underlying 3-TCVAE model while achieving better MIG than 3-TCVAE-L. We also provide FID and MIG results for Lezama's 's model (Lezama 2018 with 3 1, 10 as well FID for a vanilla VAE model of the same capacity as MS- VAE (denoted Big-VAE).
model with two values of (1 and 10). After training, the MIG values for the student network are approximately the same as teacher model on both datasets (i.e. it maintains disentanglement).
Reconstruction Quality. As shown in Figures 4a and 4b MS-VAL VAE drastically improves the reconstruction quality compared to both 3-TCVAE models, achieving much lower FID across all values of B. We also compare MS- VAE to Jacobian supervision and find that AS-E has lower FID on both datasets. To ensure that the improvement in reconstruction quality over 3-TCVAE is not just because MS-VAE has more parameters, we trained several 3-TCVAE models with the same number of parameters MS-VAL on the SmalINORB dataset and reported the results in Appendix Although we found marginal improvements in the average FID for these bigger 3-TCVAEs (145.79 + 2.98), they are still far higher than MS-VAE's parameters of 3-TCVAE led to a decreased MIG 10. = 0.026) which indicates that increasing the parameters of 3-TCVAE may increase reconstruction quality at the cost of disentanglement. Finally, for reference, we provide FID scores for a standard VAE with same number of parameters MS-VAE (we refer this model as Big- VAE). Surprisingly, MS-VAE has lower FID than Big-VAE at low values of 3 suggesting that the residual modeling in MS- VAE leads to better image quality in general (see Figures 4a () and i)).
VAE's s FID (59. 16 + 4.52). Also, inereasing the
Conditioning on C Consider Figures 5a and 5b where we report the mutual information terms M1, M2. M3. and M4 for both datasets (see Section 4 for their definitions). As per our expectation,
7

---

Under review as a conference paper at ICLR 2021
find that M2 > M1 > M3 >> M1 and that My tracks M1 without dropping significantly. This suggests that Z is encoding meaningful information for the reconstruction X and that the model maintains the conditioning
C/ in the final reconstruction. We further illustrate the conditioning the C with qualitative examples. In Figures 3)) d) and
and e)), we show latent traversal Y and X pairs. Notice, MS- VAE is simply improving the quality of the generated output without diminishing the ability of the model to manipulate single factors of variation by traversing C.
3-TCVAE
MS-VAE
MS- VAE (Random 2)
MS-VAB (Random Y)
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0
10
(a) Cars3D dataset
(b) SmallNORB dataset
Figure 5: Mutual Information (MI) between inferred independent factors from the true image X (using 8-TCVAB) and independent factors from various reconstructions of X. Please note that Blue M1, Red M2, Green My. and Purple = M4 (see Section for their definitions).
Downstream Classification The results of the downstream attribute prediction tasks are shown in Figure As can be seen, the MLP that was trained with both C' and Z attains a higher accuracy on all five attribute prediction tasks than the MLPs trained with only C' or Z. Notice that as the regularization of the latent space increases (larger (i) and the accuracy of the classier trained on C' decreases, the accuracy of the classifier trained on Z increases. This highlights that Z contains complementary information to that of and sheds light on why simply increasing the size of the latent space (i.e. the 3TCVAE-L model) does not improve the reconstruction significantly.
2 CoZ
Accersey
6.0
6.2
(a) Category
(b) Instance
(c) Elevation
(d) Azimuth
(c) Lighting
Figure 6: Object property prediction for SmalINORB using inferred representations C, 2, and C Z as input to MLP. The accuracy is highest for the MLP trained with C' + Z.
MS-GAN and MS-FLOW. We implement GAN-based version our framework, MS-GAN. and carry out a large battery of qualitative experiments in the supervised setting using graphics rendering programs (Appendix EE). We also evaluate simple FLOW-based version MS-FLOW, in Appendix D that utilizes mixture of Gaussians model to realize the sub-graph C + Y, further demonstrating the method's flexibility. As we are limited for space, we include these results in the appendices.
framework,
Simple Pendulum. To demonstrate the ability of MS- VAE learn richer and more structured hierarchies of disentangled latent factors, i.c. Deeper MS- VAE, we learn to approximate a simple pendulum simulator using MS- VAF. We refer the reader to Appendix [or the results.
5 CONCLUSION
In this work, proposed novel graphical model for DR learning that, by virtue of being D- separated, overcomes the trade-off between reconstruction and disentanglement. To implement this graphical model, introduced multi-stage deep generative model, MS-VAE, that learns
8

---

Under review as a conference paper at ICLR 2021
the independent and correlated latent variables in two separate stages of training. We showed that this model dramatically improves the reconstruction of preexisting DR learning methods while maintaining the same level of disentanglement. Furthermore,
experiments for MS-GAN and MS-FLOW (Appendices E and [ respectively) demonstrated that our framework is agnostic to both the level of supervision (e.g. unsupervised, weakly supervised, fully supervised) and to the underlying model (e.g. VAEs, GANs, normalizing flows). Finally, through the pendulum experiment (Appendix C we showed how our framework can be extended to models with deeper and/or richer hierarchical structure. In future work, we seek to explore these deeper extensions of AS-VAE.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. 12th {USENIX) Symposium on Operating Systems Design and Implementation (LOSDI) 16), pp. 265-283, 2016.
Mathieu Aubry, Daniel Maturana, Alexei A Efros, Bryan C Russell, and Josef Sivic. Seeing 3d chairs: exemplar part-based 2d-3d alignment using large dataset of cad models. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3762-3769) 2014.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1798-1888 1828, 2013.
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al. Shapenet: An information-rich 3d model repository. arXiv preprint arXiv:1512.03012 2015.
Tian Qi Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentan- glement in variational autoencoders. In Advances in Neural Information Processing Systems, pp. 2610-2620, 2018
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. In Advances in neural information processing systems, pp. 2172-2180. 2016.
Li Deng. The mnist database of handwritten digit images for machine learning research [best of the web]. IEEE Signal Processing Magazine, 29(6):141-142 2012.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. CoRR, 2014. URL httl::ttt////x/e/ http://wwww..rgg/aaaa...... 6v6
Laurent Dinh, Jascha Sohl-Diekstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint /605v 100.0003, 2016.
Vincent Dumoulin, Ethan Perez, Nathan Schucher, Florian Strub, Harm de Vries, Aaron Courville, and Yoshua Bengio. Feature-wise transformations. Distill, 3(7) 2018
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. CoRR, 2019. URL htttt://paaxii>> org/ abs 10.da.../99...0. wttpaaa/.110000000122.
Bernhard Egger, Sandro Schönborn, Andreas Schneider, Adam Kortylewski, Andreas Moree-Forster, Clemens Blumer, and Thomas Vetter. Occlusion-aware 3d morphable models and an illumination prior for face image analysis. IJCV, 2018
Thomas Gerig, Andreas Morel-Forster, Clemens Blumer, Bernhard Egger, Marcel Lüthi, Sandro Schönborn, and Thomas Vetter. Morphable face models an open framework. In FG, pp. 75-82, 2018.
Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu. David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In Neural Information Processing Systems, 2014.

---

Under review as a conference paper at ICLR 2021
Michael Gutmann and Aapo Hyvärinen. for unnormalized statistical models. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics. pp. 297-304, 2010.
estimation: A new estimation principle
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Berhhard Nessler, and Sepp Hochreiter. GANS trained by two time-scale update rule converge to local nash equilibrium. In Neural Information Processing Systems, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. ICLR, 2(5):6, 2017.
Qiyang Hu, Attila Szabó, Tiziano Portenier, Paolo Favaro, and Matthias Zwicker. Dis- entangling factors of variation by mixing them. puter Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18- 22, 2018, pp. 3399-3407. IEEE Computer Society, 2018. doi: 20/10////R.2018 00358. URL http::/ppenacees.... Disentangling Baaiinn_aacctttto Factors of CVPR 2018 paper. html
2018 IEEE Conference on Com-
wpw.access..s.eeeeeopop.....s..e theevf. com/ content evpr 2018/html
Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. CoRR, 2018.
00779v1
Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normal- ization. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1501-1510, 2017.
Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp.
2017.
Yeonwoo Jeong and Hyun Oh Song. Learning discrete and continuous factors of data via alternating disentanglement, arXiv preprint arXiv:1905.094: 2019.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising, arXiv preprint 1r02iv:802.0533 2018.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. 2nd International Conference on Learning Representations, ICLR, 2014.
Ivan Kobyzev, Simon Prince, and Mareus A. Brubaker. Normalizing flows: Introduction and ideas. CoRR, 2019.
09257v1
Adam Kortylewski, Bernhard Egger, Andreas Schneider, Thomas Gerig, Andreas Morel-Forster, and Thomas Vetter. Analyzing and reducing the damage of dataset bias to face recognition with synthetic data. In CVPRW, 2019.
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disentangled latent concepts from unlabeled observations. arXiv preprint 17X1:1711100845: 2017.
Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and bruin sciences, 40, 2017.
Guillaume Lample, Neil Zeghidour, Nicolas Usunier, Antoine Bordes, Ludovic Denoyer, and Mare' Aurelio Ranzato. Fader networks: Manipulating images by sliding attributes. In Is- abelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA. USA. pp. 5967-5976, 2017. URL http:::ttppess...to.tt nng.s.iip..... cc/paper/ 1o0-faaerrs net wO sman pu luulat t iinin -ygg-mmage----ss1-1
Yann LeCun, Fu Jie Huang, and Leon Bottou. Learning methods for generic object recognition with invariance to pose and lighting. In Proceedings of the 2004 IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2004. CVPR 2004., volume 2, pp. II-104. IEFE, 2004.
10

---

Under review as a conference paper at ICLR 2021
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nanure, 521((553))436, 2015.
Wonkwang Lee, Donggyun Kim, Seunghoon Hong, and Honglak Lee. High-fidelity synthesis with disentangled representation. arXiv preprint arXiv:2001.04290 2020.
José Lezama. Overcoming the disentanglement vs reconstruction trade-off via jacobian supervision. 2018.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Francesco Locatello, Stefan Bauer, Mario Lucic, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled represen- tations. arXiv preprint arXiv: /8/ I. 12359, 2018.
Michael F Mathieu, Junbo Jake Zhao, Junbo Zhao, Aditya Ramesh, Pablo Sprechmann, and Yann LeCun. Disentangling factors of variation in deep representation using adversarial training. In Advances in Neural Information Processing Systems, pp. 5040-5048, 2016.
Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXiv preprint
1411. 2014.
Shakir Mohamed and Balaji Lakshminarayanan. Learning in implicit generative models. arXiv preprint arXiv:1610.03455 2016.
Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. In Advances in neural information processing systems, pp. 271-279, 2016.
Guim Perarnau, Joost van de Weijer, Bogdan Raducanu, and Jose M. Álvarez. Invertible conditional gans for image editing. CoRR, abs/1611 100655, 2016. URL
1611
06355
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual 2017.
reasoning with a general conditioning layer. arXiv preprint
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with general conditioning layer. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint
2015.
Ali Razavi, Aaron den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. In Advances in Neural Information Processing Systems, pp. 14837-14847 2019.
Scott E Reed, Yi Zhang, Yuting Zhang, and Honglak Lee, Deep visual analogy-making In Advances in neural information processing systems, pp. 1252 1260, 2015.
Scott E. Reed, Zeynep Akata, Xinchen Yan, Lajanugen Logeswaran, Bernt Schiele, and Honglak Lec. Generative adversarial text to image synthesis. In Maria-Florina Balcan and Kilian Q Weinberger (eds.), Proceedings of the 33nd International Conference Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 1060-1069. JMLR.org, 2016. URL 1488/eedd6..hae html
mlr .resss/
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. arXiv preprint / 2014.
Akash Srivastava, Lazar Valkov, Chris Russell, Michael U Gutmann, and Charles Sutton. Veegan: Reducing mode collapse in gans using implicit variational learning. In Advances in Neural Information Processing Systemx, pp. 3308-3318 2017.

---

Under review as a conference paper at ICLR 2021
Akash Srivastava, Kai Xu, Michael U Gutmann, and Charles Sutton. Ratio matching mmd nets: Low dimensional projections for effective deep generative models. arXiv preprint arXiv: 1806.00101 2018.
Akash Srivastava, Kristjan Greenewald, and Farzaneh Mirzazadeh. Bregmn: scaled-bregman genera- tive modeling networks. arXiv preprini arXiv:7906. 00,,00 2019.
Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density ratio estimation in machine learning. Cambridge University Press, 2012.
Attila Szabó, Qiyang Hu, Tiziano Portenier, Matthias Zwicker, and Paolo Favaro. Challenges in dis- entangling independent factors of variation. In 6th International Conference Learning Represen- tations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Workshop Track Proceedings. OpenReview.net, 2018. URL https:
httphttts//ppeeeeeeeeeee wwtp///wpweee...eteeepeee wwt. wwtp:/eeeettoomnd..........
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv:I607.0880 2016.
Sjoerd van Steenkiste, Francesco Locatello, Jürgen Schmidhuber, and Olivier Bachem. Are disen- tangled representations helpful for abstract visual reasoning? Advances Neural Information Processing Systems, pp. 14222-14235, 2019.
Zhou Wang, Eero P Simoncelli, and Alan C Bovik. Multiscale structural similarity for image quality assessment. In The Thritt-Seventh Avilomar Conference on Signals, Systems & Computers, 2003, volume 2, pp. 1398-1402. leee, 2003.
Zhou Wang, Alan Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from eπror visibility to structural similarity. IEEE transactions on image processing, 13(4):600-612, 2004.
Han Zhang, Tao Xu, Hongsheng Li, Shaoting Zhang, Xiaolei Huang, Xiaogang Wang, and Dimitris N. Metaxas. Stackgan: Text to choto-realiste image synthesis with stacked generative adversarial networks. CoRR, abs/1612.02244. 2016. URL http: /awp/vvvivv/v////////ea//aaaa/ /tgggabs//1/....o/////// 03242
12

---

Under review as a conference paper at ICLR 2021
A CHOOZING C AND Z
We study the impact of the dimensionality of C' and Z on the reported results in this section. To do so, we fix value of 3 and sweep over multiple values for C' and Z.
Choosing C: For C. we sweep over the range [1 10, 20, 50] for both the datasets. Figure 7a and 7b show the MIG and ELBO for the same. We can see that |C| 10 achieves good combination of MIG and ELBO for the Carx3D dataset. While there is high variance in the MIG for SmalINORB, [C| = 10 performs reasonably well. This is also consistent with the large scale study in (Locatello)
al. 2018), which suggests that value of 10 for the dimensionality of C' is optimal for DR learning.
Choosing Z: Similarly for Z, we sweep over a set of candidate values [5, 10, 50, 90] and evaluate the FID and MIG. For evaluating the FID, consider both reconstructions (by sampling Z from the posterior) and 'Random Z' (by sampling a random from the prior). We also look at the MI We find smaller values, 5 and 10 for Cars3D and Smal/NORB, respectively, achieve good combination of all the three quantities of interest. The results are shown in ??
0.15
1,500
-1,970
MIG
0.10
MIG
-1,95
0.05
-1,000
0.20
11,980
0.00
0.18
11988
50
100
20
10
Dim C
Dim C
Dim C
Dim C
(a) Cars3D
(b) SmalINORB
Figure MIG and ELBO for different dimensionality of C for Cars3D and SmalINORB at / = 4
6-TCVAE
MS- VAE (Random Z)
MS-VAE
150
0.16
0.7
100
0.14
0.6
0.12
0.5
50
0.4
20
40
60
80
20
60
40
80
Dim Z
Dim 2
Dim Z
(a) Cars3D
200
0.4
FID
150
MI
0.3
100
0.15
0.2
0
20
40
60
80
20
40
60
80
0
20
40
60
80
Dim
Dim Z
dim
(b) SmallNORB
Figure 8: MIG,FID and MI plots for VAE as function of the dimensionality of Z for Cars3D and SmalINORB (5 = 4).
Qualitatively, Figure 9 and Figure show the effect of varying the dimensionality of C' and Z. For a small C' FFigure.)... the underlying disentanglement method does a poor job at disentangling the meaningful latent variables, which further results in blurry images. In this case, Z has substantial control over the image- it adds various missing details to Y such as color. In contrast, Figure 10 shows traversals for high value of C. Now, each individual dimension of C controls fewer independent factors, the images are relatively sharper, and adding Z improves them further. Similar conclusions hold for SmallNORB shown in Fig. and IT
13

---

Under review as a conference paper at ICLR 2021
Figure 9: Traversals for cars for small C (c= 2, = and 3 4). Odd rows 3-TCVAE and even rows MS- VAE. We can see the Y images are blurry and various factors are entangled.
Figure 10: Traversals for cars large (c = 10, 2 and 3 == 4). Odd rows 8-TCVAE and even rows Larger C' achieves greater disentanglement and Z further refines the images.
14

---

Under review as a conference paper at ICLR 2021
Figure 11: Traversals for SmallNORB = and B 4. Odd rows 3-TCVAE and even rows MS-VAI
Figure 12: Traversals for SmalINORB C = 10, z= 10 and / = 6. Odd rows 3-TCVAE and even rows
MS- VAE.
15

---

Under review as a conference paper at ICLR 2021
B CELEBA RESULTS
N
2
Figure 13: CelebA Latent traversal C. Odd rows are B-TCVAE and even rows are MS- VAE.
Figure 14; CelebA Reconstruction plots.
16

---

Under review as a conference paper at ICLR 2021
C MS- VAE FOR SIMPLE PENDULUM
A simple, 2D pendulum is comprised of bob with mass m attached to string of length I. where the mass of the string is negligible in comparison to m. The primary forces acting upon the bob are the gravitational force, the tension force of the string, and the damping force from air resistance. Let 0 be the angular displacement of the weight (the corresponding angle with respect to the y-direction). The do sin 0, where b is the damping pendulum's motion can be completely described as ag m coefficient and g is the gravitational acceleration constant (the initial angular displacement/velcityy are constant).
d'e
In this toy example, we assume that, given the length L. the damping coefficient B, and the mass M, the angular displacement of pendulum is generated by noiseless simulator S (L, M, B) + where 0 = {0u, .,,,,, Or J. 0g E [- J.. and T = 100. Using samples from S where only L and B are known a priori. we aim to learn a hierarchical, generative model for the pendulum where all latent variables are completely disentangled. With MS-VAE, we can achieve this by learning the graphical model in Figure 17 this graphical model, we introduce two intermediate observed nodes y, and Yi that block the path between I., B. and Z. We define Yn as the undamped angular displacement and Y1 as the damped angular displacement with fixed mass (M = 1). This simple example demonstrates how the MS-VAE framework can be extended to models with more structured hierarchies while consistently satisfying the conditional independence assumptions implied by the simulator.
We train our model in three separate steps. First,
paired data {L', y of to learn the subgraph - Yo as a conditional Gaussian distribution (how the pendulum length affects the angular displacement). Second, use the paired data {b', y' o y } to learn (B, Yo) Yı conditional Gaussian distribution (how the damping coefficient affects the angular displacement for a fixed mass). Finally, we the paired data (y',,
1, 2²) to learn the subgraph (Z, y¹ - 0, where Z models the residual between Y1 and the true observations & (this subgraph is realized as VAE). Through construction and training, the latent variable L should control the length of the pendulum, the latent variable B should control the damping coefficient, and the latent variable Z should be correlated with the mass M. D Note that while we have labeled data for the latent variables L and B, we never have access to the latent variable M. Despite this, the latent variable 2 can still learn to model the effect of M by modeling the residual.
The results are shown in Figures 15 and 16 As can be seen, MS- VAE is able to learn completely disentangled latent variables using combination of supervision and hierarchical modeling. The final learned latent variable Z is shown to be heavily correlated with the mass of the pendulum, highlighting how MS- VA E's learned latent variables can be directly related with informative physical parameters. This simple toy example highlights two powerful features of MS- VAE. First, the hierarchy of latent variables can be made arbitrarily deep. Second, our method can be used to discover interpretable latent factors when only some of control variates C are known a priori (e.g. the length 7. and damping coefficient B).
For simplicity, we don't model the data as time series explicitly but rather as fixed length, ID vector.
17

---

Under review as a conference paper at ICLR 2021
F15.br10mm03 1.5
2.4, m:10.mm
1.5
Aegle Drgaeeeeee (90)
E.L.C.
E. bil.o.
2::.0.2:0..22
1:2.5. bi 2:0.52
1.3.0. 16, 0.22
Time
Time
Time
Time
Time (6)
(a) Traversing length from 1.0 3.0 while fixing decay constant b and
Time
(b) Traversing decay constant à from .1 to 2.0 while fixing the length and
F.1.5. bn me 8.1
F.1.5 & 1.0.m: m383
F.S.S 1.1.8.01..... m.
115. b.1.0.
115.0 b01.0.m:
M Traveenal
032
048
-339
2 Travensl
2 TTreeenel 1:15 b 1.D.c
Traversal
2 Traveradt 13. 211.0.14
begue Dhgueerbpppeenn t040
1:1.. 1.0.0 21-15
115.018
032-
Time
(c) Traversing the mass m from .1 to 2.0 (top) and traversing the corresponding values (-1.5 to 1.44) while fixing the length and decay constant b.
Figure 15: Latent traversals for the pendulum model. Odd rows (blue) are actual data from the simulator and even rows (red) are from MS-VAE. In this example, MS- VAE learns each of the three latent variables a disentangled manner without any supervision for the mass control variate. We choose the values that plotted by examining the posterior over and choosing values that correspond to the given masses (z is I-Dimensional so this is straightforwar)..
18

---

Under review as a conference paper at ICLR 2021
175
150
Mass
1.25
0.75
0.50
0.25
-1.0
-0.5
2.0
Zmu
Figure 16: Correlation between and the mass when modeling a simple pendulum. In this plot, we show 500 I-Dimensional posterior means for different 0 observations to illustrate how mass and the latent variable are heavily correlated after training.
L
yo
y1
Figure 17: MS- VAE graphical model of a simple pendulum. is the angular displacement; L is the length; B is the damping coefficient, After learning and by construction, the latent variable Z is correlated with the pendulum's mass.
19

---

Under review as a conference paper at ICLR 2021
D FLOW-BASED IMPLEMENTATION (MS-FLOW)
Normalizing flows are powerful and flexible parametric distributions which are also tractable Dinh et al. 2014; Huang et al. 2018 Durkan et al. 2019); for review see |Kobyzev al. (2019). The conditional model in MS- VAE Pe (3) y) can also be implemented as conditional normalizing flows (that conditions on Y) instead of a VAE.
As an illustration, in this section present the MS-FLOW model the MNIST digit dataset. Here the subgraph + Y is realized as a Gaussian mixture model and the subgraph (Y, Z) - X is implemented as conditional coupling flow (Dinh et al. 2016). We fit mixture of 10 Gaussians whose means are shown in Figure 18
Figure 19 shows the real data, samples from the Gaussian
2
6
I
9
9
0
9
3
0
7
0.0
0.2
0.6
Figure 18: The means of the Gaussian mixture part of MS-FLOW on MNIST.
mixture and the conditional flow. As expected, the log density the overall model improves from -15850066 -1356.089 once the FLOW is applied to the output of the mixture model. All while maintaining the ability of controllable sampling.
D.1 TRAINING USING MEANS OR SAMPLES
As mentioned in previous sections, when training the conditional model, we can use either the means from the likelihood model of the C' + Y sub-graph, or the samples from it. We fit the MS- FLWWW model with both strategies and show the change of log density during training in Figure 20 Using the means makes the training converge much faster than using the samples. It is not surprised at all, because using samples leads to noisy gradients during training. Thus, for faster convergence, we means rather than samples in all of our experiments, even though it make the estimates slightly biased.
20

---

Under review as a conference paper at ICLR 2021
3
8
3
3
7
6
7
6
2
3
4
O
3
5
7
6
7
2
5
2
7
O
3
1
2
2
4
3
5
4
3
5
4
6
2
5
8
0
9
S
O
2
O
2
0
3
1
3
7
2
3
2
3
7
27
0.0
0.2
0.4
0.6
0.8
1.0
Figure 19: MS-FLOW on MNIST. First rows are real data, next five rows are samples from the Gaussian mixture and last five rows are from the conditional flow.
-1,400
Log Density
CFLOW with samples
-1,500
CFLOW with means
GMM (trained)
1,600
0
200
400
600
800
1,000
1,200
Steps
Figure 20: Conditional log density during training for different types of Y (samples v.s. means).
21

---

Under review as a conference paper at ICLR 2021
E GAN-BASED IMPLEMENTATION (MS-GAN)
Equation 2 also allows for likelihood-free inference with a little manipulation. Simply by adding and subtracting the log of the true conditional density log ((g((()) in equation
we get the lower bound, 2
log Plo.0.z.. v) - -<CCss((h))llocc))
/ go (o(cly) logpo KZ wlll()))
/ 900 (2/ v) log pe (z12, dæ log p()l())))
(3)
p(zly)
GAN based (Y:2)+X
Here, the log density ratio log (a(arlz,)) can be estimated using adversarial learning with binary p(aly)
discriminator function (Sugiyama et al. 2012, Srivastava et al. 2017). In practice, since we are
conditional GAN Mirza & learning conditional distributions, the decoder is implemented
Osindero (014) This implementation is referred to as MS-GAN.
E.1 PRELIMINARIES: GENERATIVE ADVERSARIAL NETWORKS
Generative adversarial networks (GAN) (Goodfellow et al. 2014) represent the current state of the art in likelihood-free generative modeling. In GANs, a generator network Go is trained to produce samples that can fool a discriminator network D.,, that is in turn trained to distinguish samples from the true data distribution p(æ) and the generated samples Go(z)lz pz(z). Here, P= is usually a low dimensional easy-to- tamyle distribution like standard Gaussian. A variety of tricks and techniques need to be employed solve this min-max optimization problem. For models, employ architectural constraints proposed by DC-GAN (Radford et al. 2015) that have been widely successful in ensuring training stability and improving generated image quality.
Conditional GANs (CGAN) (Miza Osindero 2014) adapt the GAN framework for generating class conditional samples by jointly modeling the observations with their class labels. In CGAN, the generator network Go is fed class labels e to produce fake conditional samples and the discriminator Da, is trained discriminate between the samples from the joint distribution of true conditional and true labels p(z|c)p(c) and the fake conditional and true labels pe (z)c)p(c).
While the main focus of this paper, we present a novel information theoretic perspective on CGANs. Specifically, we show that CGAN is trained to maximize a lower- to the mutual information between the observation and its label while simultaneously minimizing an upper-bound to it. We state this formally:
Lemma (Information-theoreet interpretation CGAN). Given (æ, ~ p(x, c). CGAN learns the distribution p=(x)= Go(x) by training a discriminator D., to approximate the log-ratio of the true and generated data densities i.e. Du R pl,,( c)/po (2, c) in turn minimizing the following
min Eps(s)e)p(e) Dw]
min (In,0 c) 0 (((ce))))
max (I,, c) Epp(s) KLl( (arr, ()ll(())))
min TUB 9.8 (x, - -I (x.
where Ig,0 (x, c) is the generative mutual information and q(c|2x, f) is the posterior under the learned model.
The detailed derivation is provided in the section below. Notice that at the limit, the model learns exactly the marginal distribution of x and the posterior ((c(2r) and the KL terms vanish.
E.2 CGAN MUTUAL INFORMATION DERIVATION
Following Sugiyama et al. 2012); Gutmann & Hyvärinen 2010); Mohamed & Lakshminarayanan
(2016): Srivastava et al. (017) know that optima the logits (l...) of trained discriminator
22

---

Under review as a conference paper at ICLR 2021
approximate the log-ratio of the true data and generated data densities, i.e. Dw log Rz (a pe pe(z)c)p(c)
Following Nowozin al. (2016); Srivastava al. (2017 2018 2019) the generator Gg can therefore be trained to minimize the following I-divergence between the two sets of densities,
minnE[-Duj -wl
0
f pa po(zle)n(e) log plallcpp(c) 6
R min
d(x,.) c).
(4)
po po( (2)cc))((c)
The RHS of equation can be re-arranged into terms containing the upper and the lower bounds to the generative mutual information between X and C, i.e.
min E
Dwo
log Po (x(ec)
log p(c|2)
min E
+ max E
p(x)
0
p(c)
log (x(c)
(o(x)
Po (x)
min 0
+ log
(a(
p(x)
log g(c/2,
0)
+ log p(c|z)
max E
P(c)
q(cl2,,) 0)
log po(x]c)
po p=(x)]
min E 0
+ +og
(x(x)
p(z)
log g(c 0)
log ((cl2))) 0)
max E
p(c)
p(c]2)
min Ig,0 c) + Ey 0) (l.(p((=))l(=)))++
0
max (2, c) Ep) (L 0llp((z)
min TUB BB (x,, c) I (I,
(5)
E.3 LEARNING 2Y + X USING CGAN
Vanilla GANS can only model the marginal data distribution i.e. they learn Pa to match Pa and in doing so they use the input the generator (Gş) only as source of stochasticity. Therefore we start with a conditional GAN model instead, to preserve the correspondence between Y and X. As shown in section E this framework trains Gø such that the observation X maximally explained by the conditioning variable Y. One major deviation from the original model is that the conditioning variable in our case is the same type and dimensionality as the observation. That is, is an image, albeit blurry one. This setup has previously been used by Isola et al.] (2017) in the context of image-to- tongg-tt--mmge translation.
Incorporating Z requires careful implementation due to two challenges. First, trivially adding Z to the input along with Y invokes d-separation and as a result Y and Z can get entangled. Intuitively, Z is adding high level details to the intermediate representation Y. We leverage this insight as an inductive bias, by incorporating Z at higher layers of the network rather than just feeding it as an input to the bottom layer. A straightforward implementation of this idea does not work tough. The reason is that BatchNorm uses batch-level statistics to normalize the incoming activations of the previous layer to speed up learning. In practice, mini-batch statistics is used to approximate batch statistics. This adds internal stochasticity to the generator causing it to ignore any externally added noise, such as Z. An elegant solution to resolve this second challenge comes in the form of adaptive instance normalization (see Section 3.2) It not only removes any dependency on the batch-statistics Z in the normalization process itself. For this reason, has previously been used in style transfer tasks (Huang & Belongie 2017). We replace all instances of BatchNorm in the generator with Adaptive InstanceNorm.
but also allows for the incorporation
then introduce to the generative process using equation ??. 7(2) and ((z) are parameterized as simple feed-forward network and are applied to each layer of AdaIN in the generator.
23

---

Under review as a conference paper at ICLR 2021
F EXPERIMENTS FOR MS-GAN
In this section, we provide a comprehensive set of qualitative results to demonstrate how MS-GAN is clearly able to not only disentangle C' from Z in both supervised and unsupervised settings but also ensure that independent components of C' stay disentangled after training. Additionally, we show how in unsupervised settings MS-GAN can be used to discover disentangled latent factors when C not explicitly provided.
We evaluate MS-GAN on variety of image generation tasks which naturally involve observed attributes and unobserved attributes Z. To that end, we generate three 3D image datasets of faces, chairs, and cars with explicit control variables. Chairs and cars datasets are derived from Net (Chang et al. 2015). We sample 100k images from the full yaw variation and a pitch variation of 90 degrees. We used the straight chair subcategory with 1968 different chairs and the sedan subcategory with 559 different cars. We used Blender to render the ShapeNet meshes scripted with the Stanford ShapeNet renderer. For faces, generated 100k images from the Basel Face Model 2017 (Gerig et al.) 2018). We sample shape and color (first 50 coefficients), expressions (first 5 coefficients), pose (yaw -90 to 90 degrees uniformly, pitch and roll according to Gaussian with variance of 5 degrees) and the illumination from the Basel Illumination Prior (Egger et al. 2018 For the generation of the faces dataset, we use the software provided by Kortylewski al. 2011). For the stated datasets we have complete access to C, however, we also include unsupervised results on celebA (Liu et al. 2015) with unconstrained, real images. All our datasets are built from publicly available data and tools.
We use the DCGAN architecture (Radford et al. 2013) for all neural networks involved in all the experiments in this work and provide reference implementation with exact architecture and hyperparameter settings at https//siiiia....o///oooonnnnnss
VAE.
F.1 SUPERVISED SETTING
In the supervised setting we compare MS-GAN to CGAN qualitatively. To evaluate the level of disentanglement between C' and Z, we vary each individual dimension of C' over its range while holding Z constant. We plot the generated images for both models on car and chair datasets in Figure 21 Notice that MS-GAN allows to vary the control variates without changing the identity of the object, whereas CGAN does not. In addition, we find that for CGAN, the noise 2 provides little to no control over the identity of the chairs. This is potentially due the internal stochasticity introduced by the BatchNorm. The last rows for the MS-GAN figures provide the visualization Y. can be seen how is clearly preserving C' (pose information) but averaging the identity related details.
We also qualitatively evaluate MS-GAN on the more challenging faces dataset that includes 10 control variates. As shown in Figure 2 in the appendix, MS-GAN is not only able to model the common pose factors such as rotation and azimuth but also accurately captures the principal shape component of Basel face model that approximates the width of the forehead, the width of jaw etc. Compared to CGAN, MS-GAN does qualitatively better job at keeping the identity constant.
F.2 UNSUPERVISED SETTING
We now test the performance of MS-GAN in the unsupervised setting, where disentangled compo- nents of C' needs be discovered, using 3- as part of learning the mapping C C- For our purpose, we use a simple version of the original 3-VAE method with a very narrow bottleneck (6D) for faces and 2D for cars and chairs) to extract C.
The latent traversals for the faces dataset are presented in Figure 22 to recover rotation as well as translation variation present in the dataset. For comparison, we evaluate InfoGAN (Chen et al.) 2016) and present the results in Figure 23 where it is evident that MS-GAN clearly outperforms InfoGAN on both disentanglement and generative quality. More traversal results provided in the appendix. We further test method on the CelebA dataset (Liu et al. 2013). where pose information is not available. This traversal plot is shown in Figure 24 Traversal plots for
Unsupervised discovery is able
cars and chairs dataset are provided in the Figures 28 and 29
24

---

Under review as a conference paper at ICLR 2021
(a) MS-GAN rotation
(b) MS-GAN elevation
(c) MS-GAN chair rotation
(d) MS-GAN chair elevation
(c) CGAN car rotation
(f) CGAN car clevation
(g) CGAN chair rotation
(h) CGAN chair elevation
Figure 21: Latent traversal on cars and chairs. Third rows in MS-GAN results show Y.
Figure 22: Latent traversal on faces (unsupervised MS-GAN), The three latent variables capture the rotation, azimuth, and distance respectively.
Figure 23: Latent traversal of InfoGAN faces. The latent variables able capture some pose changes but the pose changes are highly entangled with other pose factors as well as the face shape.
F.3 ADDITIONAL EXPERIMENT RESULTS
Comparison of MS-GAN and CGAN on face dataset is shown in Figure 25 CGAN not only produces blurry faces but also shows more undesired identity changes. In order to show the shape variation clearly, we provide zoomed-in view in Figure 26
We provide additional results for supervised and unsupervised results the chair dataset from Aubry eet al. (2014) in Figure 27 and Figure 28 respectively. The observation is the same with the previous one. MS- GAN varies the control variables without changing the shape of chairs. In the first row in
25

---

Under review as a conference paper at ICLR 2021
7
7
T
L
R
S
Figure 24: Latent traversal on CelebA (unsupervised MS-GAN). The latent variables consistently capture the azimuth, hair-style, gender and hair color respectively while maintaining good image quality.
(a) MS-GAN face pose
(b) CGAN face pose
(c) GAN face shape
(d) CGAN face shape
and CGAN on faces. The pose variations are azimuth, horizontal translation, vertical translation, distance, rotation, and elevation from top to bottom. The shape variations show the difference in face height, forehead, jaw, and ear from top bottom.
Figure 25: Latent traversal of
Figure 27 the leg of the chairs are visually indistinguishable showing excellent disentanglement between C' and Z. For the results in unsupervised setting showing in Figure 28 MS-GAN is able to disentangle the rotation of chairs without any label.
Additional results of latent traversal of MS-GAN in the unsupervised setting is provided in Figure The model is able capture the rotation but the translation is not very smooth.
29
Figure 30 provides the InfoGAN result on the face dataset. Compared with unsupervised MS- GAN result in Figure 31. clearly InfoGAN discovers some control variables but the effect is highly entangled.
26

---

Under review as a conference paper at ICLR 2021
Figure 26: Zoomed-in comparison on face shape. Row 1: CGAN forehead variation; Row 2: CGAN jaw variation; Row 3: MS- GAN forehead variation; Row 4: MS-GAN jaw variation. Row 1 and Row should have a bigger forehead from left to right while Row 2 and Row 4 should have consistent forehead. CGAN and MS-GAN show good forehead variation in Row and Row respectively, but MS-GAN better at keeping the forehead the same while another factor is changing (Row 4 Row 1).
Figure 27: Latent traversal chairs of MS-GAN The first three rows show the effect of the variable C' that controls rotation. The last row the corresponding Y.
Figure 28: Latent traversal chairs of unsupervised MS-GAN.
27

---

Under review as a conference paper at ICLR 2021
Figure 29: Latent traversal of unsupervised MS-GAN on cars.
Figure 30: Latent traversal of InfoGAN faces dataset.
Figure 31: Latent traversal of unsupervised MS-GAN on face dataset.
28

---

Under review as a conference paper at ICLR 2021
G METRIC FORMULATION
The normalized mutual information (MI) reported in Figure 35a Iii) and Figure 35b iii) is calculated as follows.
[(c)ck)
Σ R. C = 11 :=1 H(cR)
(6)
where n is the dimension of C and CR and [(c'; (x) is the mutual information between of and ch-
For every ground truth factor, we compute mutual information with each one of the learned dimensions of C, normalized by its entropy. For each factor, we take the difference, in mutual information, between the top two learned variables with which it has the highest mutual information. MIG is defined to be the average of that value over all ground truth factors.
29

---

Under review as a conference paper at ICLR 2021
H COMPARISON OF MS- VAE AND BIG-B-TCVAE
Table shows the comparison of MS-VAE with Big-3-TCVAL which is 3-TCVAE that has the same capacity in terms of the number of parameters and the latent dimensionality. Our experiments reveal that MS- -ASE outperforms the Big-8-TCVAE in terms of both disentanglement representation learning quality and reconstruction quality.
Table 1: Comparion of MS-VAE with Big 3-TCVAE.
Dataset
Model
FID
MIG
Big- 1 3-TCAE
35.24:11.28
0.09
Cars
MS VAE
22.72:13.66
0.101..03
Big- 3 3-TCAE
152.19-1.668
0.16-0.02
Smal/NORB
MS- VAB
59.16
0.22
30

---

Under review as a conference paper at ICLR 2021
1 SAMPLING MULTIPLE C AND Y DURING TRAINING
In our quantitative experiments for MS-VAE, the paired training set D paired was constructed follows: x² p(xr), c (o(c)a) y² Ey (pl(y])). In this construction, only and, consequentially. only one is sampled for training. For the qualitative experiments with celebA, we sampled 10 different c's from the 3-TCVAE for each x; such that the final dataset had 10 different y/ for each s;;
The reason we sampled more 's for each I; is that there was high variability of the sampled yt given the same I,, To visualize this effect, we plot five Y samples for five X samples from trained 3-TCVAE (3 15, C = 15) in figure 32 As can be seen from in the figure, there is high variability in the sampled yA 's for each z;. highlighting the the high variance of 3-TCVAE's posterior distribution over C. Training with multiple
for each allows the MS- VAE to naturally adapt to We plan explore this effect in future work.
this variability as shown 33
Figure 32: Even rows are samples from CelebA (the same sample is plotted five times) and odd rows are corresponding samples from the 8-TCVAE. As can be seen, for the same celebA image, the reconstruction from the 3-TCVAE is highly variable, illustrating the high variance of the learned posterior for 3-TCVAE.
??
Figure 33: The top row is the same sample from the 3-TCVAE and the bottom row is the correspond- ing reconstruction for MS- VAE. As can seen the figure, despite variability in the sampled Y's. MS-VAE VAE is able reconstruct the celebA images with little variance.
31

---

Under review as a conference paper at ICLR 2021
J PIXEL-WISE RECONSTRUCTION LOSS
For all the methods, we report the LI and L2 reconstruction emrors. As shown in Figures 34 and 35 pixel-wise reconstruction loss can be misleading. the case of SmalINorb, where MS VAE clearly has significantly better image fidelity (as determined by the reconstruction FIDs and traversal plots). the LI and L2 error are only slightly better than TCVAE.
2.5
3-
1.3
2
10
0
(a) Cars3D dataset
(b) SmallNORB dataset
Figure 34: LI (lower is better) comparison for 3-TCVAE (C - Y), 3-TCVAE-L (latent dimension- ality same C + Z), and MS- VAE models.
TEVAE
TCVAE-L
TCVAEL
4
(a) Cars3l dataset
(b) Smal/NORB dataset
Figure 35: L2 (lower is better) comparison for 3-TCVAE (C - Y), 3-TCVAE-L (latent dimension-
ality same as C + 7), and MS- VAE models.
32

---

Under review as a conference paper at ICLR 2021
K D-SEPARATION IN MS- VAE: AN ILLUSTRATIVE EXAMPLE
We provide constructive example of the graphical model of MS- VAEE shown in Figure 2c to illustrate the d-separation Y as discussed Section 3.1
Suppose we are given dataset as shown in Figure 36d and we have trained (disentanglement)
model (C + Y) where C (niform((11 2,3) and CC)C)) is deterministic mapping as shown in
Figure 36a where the indices above each image are the corresponding Cs.. As it can be seen, only
1
2
(a)
1
1
(b) Z
(c)
3
5
(d)
Figure 36: A constructive example to illustrate d-separation in MS-VAE.
part of dataset (first three images) are modeled by the generative process of C- > Y. Now, with Y generated, add details recover back the entire dataset. We first sample 7 (ernoulli(0.5) and add mask to Y. The mask is uniformly sampled conditionally on Z as in Figure Tigure336bcc (Z = 0) and apply mask, uniformly chosen from Figure Y, which yield X. It's not hard to check that this generative process will yield a distribution with support shown in Figure 36d] Also, note that X and Y are in the same space. To simplify the discussion, we will represent the realizations of the random variable X and Y using the integer (C) above each image.
Figure 36C Z 1). To be more specific, e.g., given that we have sampled Z
We want to check that C and Z are not conditionally independent given X as implied by d-separation.
To do so, we pick X = 4 and check P(C, Z|X 4) _ P(C = 4) P(Z X = 4). The probability distributions involved here be represented by tables, as shown in Table It can be verified that ((C,2|X 4) # P(C = P = 4).
Now let's check the conditional independence if we additionally condition on Y, e.g.
Specifically, we check P(C, Z 4,Y 1) - P(C|X 4,y 1)) = Y, = 1). The probability distributions involved here can be represented by tables, as shown Table 3 It can be verified that PCC, X|X = 4,Y 1) P(CIX 4, Y 1)P(ZZX 4, Y = 1) holds. In other words, conditioning on X = 4 and Y d-seperates C' and Z.
33

---

Under review as a conference paper at ICLR 2021
Probability
0
C
Probability
Probability
1
3/8
3/8
2
0
1/4
0
1/4
5/8
3/8
3/4
3
3
0
0
(c) P(Z|X 4)
(b) P(C = 4)
(a) P(C, Z
4)
Table 2: Conditioning on X only.
C
Z
Probability
0
0
C
Probability
7. Probability
2
0
0
0 0
2
0
0
O
(c) P(Z 4.Y
(b) P(C = 4,Y
(a)
PCC, 7|X 4,Y
1)
Table 3: Conditioning on X and Y.
34

---

Under review as a conference paper at ICLR 2021
L ABLATION STUDIES
We now evaluate four major modeling choices of MSVAE. Specifically, performed ablation studies to understand the contribution of: 1) Our two-step training paradigm, 2) Our use of Y instead of C. for the second stage, 3) Our use of AdaIN to incorporate Z, 4) Our decision to learn C' before learning Z. To evaluate our two-step training process, we compare MSVAE to a model where both the C' and Z factors are learned together in an end-to-end manner (by penalizing only C) and Z is incorporated using AdaIN. To evaluate how using Y improves reconstruction performance over using C, we compare MSVAE to version of MSVAE where Y replaced with C. To evaluate the contribution of AdaIN for disentanglement, we compare MS VAE to a version of MSVAE where C and Z are concatenated together as input to the second stage DGM. Finally, to evaluate our decision to learn C' before learning Z, we train version of MSVAE where Z is learned in the first stage and then learned the second stage (C and Z are concatenated together to reconstruct X). All experiments are performed the Cars3D dataset.
L.1 IMPORTANCE OF TWO-STEP TRAINING
MSVAE uses two-step training to ensure that the disentangled factors are captured in C' via the underlying 3- VAE model and residual factors are captured Z. We now show that the same will not happen if the latent space (C) of the underlying 3- V-V is extended to incorporate 2 and then trained in an end-to-end fashion using the same AdaIN decoder. For this purpose train 8- 3-VAE model (BV) with the dimensionality of C set to 5. As shown in Figure 37 BV discovers disentangled factors that capture azimuth, scale and elevation. Next, we train a model (M1) which has 10 latent factors. We enforce a higher 3 penalty for 5 of these latent factors (C) and a normal VAE 3 penalty for the other 5 factors (Z). The contribution of Z is still limited using the AdaIN decoder strucutre. As can be seen in Figure 37 the disentangled factors C' in MI are now entangled with other factors (e.g. identity, color). This experiment illustrates how end-to-end training cannot easily reproduce the two-step training scheme in MSVAE which induces d-separation of C/ and Z.
L.2 IMPACT OF USING Y OVER C' AND ADAIN OVER CONCATENATION
MSVAE utilizes AdaIN to improve the reconstruction Y from the underlying disentangled repre- sentation learner with information from the correlated factors Z. To understand the importance of using AdaIN (rather than concatenation) and Y (instead of C). we perform the following experiment: Using a 3-VAF model (BV) as the C - Y sub-graph, we train three models MSVAF-C (M2), MSVAF-C-IN (M3) and MSVAE (M4). In M2, we concatenate the inferred C with Z and use that as input to a DGM to reconstruct X. In M3, in the second stage DGM. In M4, we train MSVAE as done in the paper using Y and AdaIN. In Figure 38 we show traversal plots of 4 of the disentangled factors for each of these models. For reconstruction quality, M4 outperforms both M3 and M2. This is expected as Y is produced using both C and the learned parameters of the first stage decoder 0 and, therefore, contains much more semantic information about the observation X than C' does. Using only C' requires the second stage DGM to relearn 0 while also modelling the residual between Y and X. For disentanglemennt, M4 and M3 both outperform M2. This illustrates how simply concatenating Z and C' and inputting into into the second stage DGM will result in a non-linear entanglement of the two. This entanglement causes the network to fail to condition on Y sufficiently and to use the entangled representation as a whole towards the reconstruction of X. Utilizing AdaIN incorporate Z (as done in M3 and M4), allows the network to use Z later in the generative process to model the residual information. This allows the network to better maintain the information in Y (or C)) while still incorporating Z.
still use the same C' but introduce Z via AdaIN
L.3 ORDERING OF C' AND Z
In MSVAE, we first learn the disentangled factors C' and then we learn the residual factors Z. In this ablation, we evaluate how the order in which these latent variables are learned affects the final disentangiement. For this purpose, we train MSVAE-ZC (M5) where, in the first stage, Z is learned using a standard VAE. In the second stage DGM, the disentangled factors C' are learned by enforcing KL-penalty and then concatenated with the learned Z from the first stage to reconstruct X. In Figure (3) we show traversals of 4 of the disentangled factors C. As can be seen, this new model M5 does not capture any disentangled factors in C. In contrast, MSVAH (M4) and BV capture
35

---

Under review as a conference paper at ICLR 2021
BV
M1
Figure 37: BV 3-VAE, MI 3-VAE+Z. BV is trained with a dimensionality of C 5. MI is trained with dimensionality of C = and Z = 5. For M1, the 8-penalty is only applied to C' and Z is still incorporated using AdaIN. In this figure. disentangled factors C' for both BV and MI. As can be seen, traversing C' for BV captures different disentangled factors (e.g. azimuth, scale, elevation). Traversing C' for M1, however, illustrates that the disentangled factors are now entangled with other factors (e.g. identity, color).
show traversals of 4 of the learned
variety of disentangled factors in C (e.g. Azimuth, scale, elevation). We hypothesize that M5 fails extract any meaningful disentangled factors because Z already captured both the entangled and disentangled factors in the first part of the training. This is clearly shown in Figure As such, the second stage will suffer from the shorteut problem where the decoder utilizes Z and does not learn anything meaningful in C' (due to the high KL-penalty).
36

---

Under review as a conference paper at ICLR 2021
BV
M2
M3
M4
BV
M2
M3
M4
BV
M2
M3
M4
BV
M2
M3
M4
Figure 38: BV = B-VAE, M2 MSVAE-C, M3 = MSVAE-C-IN, M4 = MSVAE. BV is trained with a dimensionality of C 5. this figure, we show traversals of 4 of the learned disentangled factors and how they are preserved in each version of MSVAE (each color box is one disentangled factor). To perform traversals for each MSVAE, we encode X using the 3-VAE to get and then we extract Z using Y; We fix Z for the traversals of each MSVAE. Clearly, M4 and M3 preserve the disentanglement from BV better than M2 (M2 changes the identity and color). Also, M4 has the best reconstruction of the three MSVAE models it utilizes Y improve the reconstruction (rather than just using C). This ablation study illustrates that AdAIN is important for preserving disentanglement and using Y is important for improving reconstruction.
BV
M5
M4
Figure 39: BV 3-VE, M5 MSVAE-ZC, M4 MSVAE. BV is trained with dimensionality of C = 5. In this figure, we show traversals of 4 of the learned disentangled factors and how they are preserved
each of the three models. To perform traversals of MSVAE-ZC. we encode X using VAE to first get Z and Y, then extract C using X. To perform traversals for MSVAE, we encode using the 3-VAE to get and Y, then we extract Z using Y. fix for the traversals in MSVAE. Clearly, M4 preserves the disentanglement from BV better than M5 (M5 expected, is not disentangled).
37

---

Under review as a conference paper at ICLR 2021
M5
Figure 40: M5 = MSVAE-ZC. In this figure, we show traversals of the Z factor in MSVAE-ZC model.
Clearly, capturing factors such as scale, identity, color, rotation an entangled fashion.
38


Reviews:
[{"review_id": "jj05VXLJEUn", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This article introduces a method for learning high-quality generative model with disentangled latent representation by splitting the learning process into two steps. The first step consists in learning a generative model on the data using a method with strong disentanglement constraints, producing a low-quality generation. As a second step, a conditional generative model is trained to turn this low-quality sample into an high quality one. This intermediate generation acts as an observed variable, and thus separates the latent spaces of the two models, effectively preventing disruptive interference in the learning of the \"independent factors\" on the one hand and the \"dependent factors\" on the other, as has been previously observed as a difficulty in the literature. The authors provide detailed empirical analysis of the performance of the model.\n\nOverall I would say that this is an interesting article, and a decent contribution to the conference. The design of the model is explained in detail, and its position among the literature is explained in detail. The method in itself is pretty general, and the authors provide explanation on how to apply it in non-VAE contexts (Flows & GANs mostly) in appendix.\n\nHowever, I have a few questions, and I think these points should be clarified in the text of the article:\n\n1. The use of FiLM structure of the conditional VAE is motivated by the risk that the model would ignore Y and solely rely on Z to generate X. However, this seems rather counter-intuitive from a theoretical point of view: the training objective of the VAE rewards storing as little information as possible in the latent space. If the model can reliably use information present in Y, we should thus expect it to do so. The previously observed issue of models ignoring their disentangled latent and storing everything in the correlated one does not applies as, this is the point of your model, Y is here treated as if it was part of the dataset. So, is this ignoring of Y something you observed experimentally before deciding to use FiLM structures?\n\n2. As a follow-up of the previous point, it would seem natural to expect that an abstract value for Y would potentially be more easily used by the conditional model than a low-quality reconstruction. Going further with that, one could consider using directly C as input to the conditional model: training β-TCVAE as you describe, and only keeping its encoder to produce a dataset of (X, C), then using this dataset to train the second stage as a conditional generative model p(X | C, Z) (with an encoder q(Z | X, C)). As per the previous argument, this second stage should be using C as much as possible. How do you think that would compare to MS-VAE?\n\n3. Regarding the simple pendulum example cited in the text, it is unclear how Y0 and Y1 are produced. Given how precisely they are defined in the appendix, I'm inferring that they are ground truth data produced by the simulator. If that is the case, it seems to me that the structure presented here is no different from a regular latent variable model with multiple observed variables, as the main defining component of your approach (learning Y with a disentanglement-inducing method) is not present at all. If on the other hand Y0 and Y1 are indeed directly learned from θ, then it should be clarified.\n\n-----------------\n\nMinor remarks and typos:\n\n- the legend for figure (2d) is mislabeled (c)\n- the legend for figure (3e) seems to invoke the wrong dataset (the images look like from Cars3D, not SmallNORB)\n- bottom of page 7, the text invokes figure (3h), which does not exist", "extra_scores": {}, "extra_text": {}}, {"review_id": "d0_-Lqc62-L", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper studies the problem of learning disentangled representations while maintaining good data reconstruction. As common modeling, the latent representation is decomposed into disentangled representation C and correlated representation Z. Then a hierachical generative process is proposed, where the first stage is to reconstruct a preliminary version of the data given the disentangled representation C, and the second step is to reconstruct a full version of the data given C and correlated representation Z. The two stages are learned separately, with the first stage using the previous β-TCVAE model to learn C, and the second stage using the Feature-wise Linear Modulation (FiLM) technique.\n\nExperiments show the approach can learn disentangled representation C as well as the previous β-TCVAE model, but improves in terms of the reconstruction quality. It is also shown that the second stage can properly maintain the conditioning on C (instead of conditioning only on Z). \n\nAs above, overall, the proposed modeling is simple and reasonably sound, and the experimental results / analyses are interesting and promising. \n\nA potential shortcoming is that the two stage components are learned separately instead of jointly, which could results in suboptimial learning. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Z9iei1L9RsW", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a multi-stage approach to learning disentangled representation that refines the low quality reconstructions of a state of the art architecture using a second deep generative model that learns how to modify these reconstructions with additional correlated latent variables.\n\nThe paper is well written and easy to follow. The introduced idea is quite straightforward, I would not be surprised at all if it was introduced before, but I do not recall seeing it in any other work so I will regard it as novel (unless the other reviewers point out similar works). Despite its simplicity, this method makes a noticeable difference in the experiments, so it could be very useful trick for other researchers working in the area of learning disentangled representations.\n\nMy main concern for this paper is the motivation of the proposed approach. The need for a multi-stage approach is justified using d-separation, saying that in this way independent and correlated latent variables (c and z) are completely separated during training. This is true, but is the multi-stage approach the best/only way to obtain this effect?\nIt would be useful to understand whether a similar effect could be obtained instead in an end to end fashion with the right architecture or training procedure, since I believe that such a procedure would have a much bigger impact for other researchers.\nAn interesting ablation study that would allow to better understand the mechanics of this technique would be for example using the same split of c and z and the same FiLM + AdaIN decoder, but training end to end placing the beta factor as in the beta-VAE only on the c latent variables, with for example a schedule for beta in which independent latent variables are more important at the beginning of training, or some alternating updates.", "extra_scores": {}, "extra_text": {}}, {"review_id": "19vUcuWuq9W", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "=======================================================================================\n\nSummary : \n\nDisentangled representation (DR) of data is useful in downstream tasks. However, VAE-based DR fundamentally suffers from a trade-off between high-quality reconstruction images and disentangling. To overcome this point, the paper approaches VAE from the multi-stage modeling (so-called MS-VAE). The proposed method starts from the other standard DR method which learns low-quality image(Y) and then, improves the quality of the image via training additional encoded representation(Z). The widely-used techniques in style transfer (FILM and AdaIN) are used to adopt 'Z' into 'Y' for a high-quality reconstructed image. The authors evaluated the proposed method through FID(high-quality image) and MIG(disentanglement). At the similar scale of complexity, the proposed method obtained high FID score and low MIG score than the baselines.\n\n=======================================================================================\n\nReasons for score: \n\nOverall, I vote for rejection. I'm not fully persuaded by the results from the novelty of this paper.\nI left my concerns (see cons). Hopefully, the authors cover my concerns in the rebuttal period. \n\n=======================================================================================\n\nStrong points:\n\n(1) Paper is well written.\n\n(2) Authors provided the justification of the proposed method via a graphical model.\n\n=======================================================================================\n\nCons :\n\nThe insight in the MS-VAE is quite overlapped in CascadeVAE [1]. CascadeVAE learns the continuous representation sequentially and shows that this sequential learning of representation helps disentanglement representation without special technique provided in FactorVAE, TC-\\betaVAE, and etc. I suggest the authors provide comparison experiments with this kind of sequential learning in representation. \n\nAlso, an ablation study on the proposed method without AdaIN and FILM would increase the persuasion of readers and provide justification for the proposed method. At a high level, the overall method is to learn the representation of an original image which helps to transform a low-quality reconstructed image, which is trained in VAE, to a high-quality reconstructed image. Also, the experiment results in figure 4 show that disentanglement was not improved and only the quality of the image was improved. \n\nFinally, if the ultimate motive of this paper is to learn the representation which can reconstruct the original image fairly well, the comparison with the VAE-based Method([2], [3])  which generated a high-quality image should be added.\nFor the disentangled representation which is used for future downstream tasks, there is no big difference with \\beta-TC VAE in the best hyper parameter setting as shown in Figure 6. \n\n[1] Learning discrete and continuous factors of data via alternating disentanglement, ICML2019.\n[2] IntroVAE: Introspective Variational Autoencoders for Photographic Image Synthesis, NIPS2018. \n[3] NVAE: A Deep Hierarchical Variational Autoencoder, NIPS 2020.\n\n=========================================================================================================\n\nAfter rebuttal (and a final question):\n\n\nThank you for your responses and the experiments I requested. I hope the authors would add discussions with M5 and quantitative results in the final version. Before I finalize my ratings, I have a final question after seeing supplementary material L. I believe M5 is enough to achieve the goal of this paper since $[Z, C]$ captures disentangled factors in $Z$ and $[Z, C]$ improves the reconstruction quality. Then, what is the advantage of MSVAE over M5?\n\n\n=========================================================================================================\n\nAfter the final question : \n\nThank you for your reply. I misunderstood the experiment M2 and propose M5. I believe M5 is just a variant of M2. \nThe authors should add quantitative results (FID and disentanglement score) comparing the baselines (in supplementary material L) in their final version. I would raise my score to 6. \n\nMinor : M2-M4 denote mutual information (paper) and baselines (supp). Please use different notations.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a two-step approach to achieve disentangled representation and good reconstruction at the same time in deep generative models: the first step focuses on good disentanglement (e.g., with beta-TCVAE) while possibly sacrificing reconstruction reconstruction, while the second step focuses on high-quality reconstruction, conditioned on the low-quality reconstruction from disentangled representation. In this paper, each step uses an existing method: beta-TCVAE is used for the first step and AdaIN is used for the second, so the paper presents an intuitive combination of two existing methods to achieve both goals. Some useful ablation studies are provided to empirically justify the specific method choices. The concern is whether the two step approach is necessary to achieve both; the authors' argument is that models learning only one set of latent variables may not have the capacity to achieve both goals, and methods jointly learning two set of variables ("disentangled" and "correlated" variables) can not guarantee they represent disjoint structures of data. Some of these statements seem somewhat handwavy (including the d-separation argument, I am not very sure if it applies when the variables are learned separately), and shall be made rigorous and justified (theoretically and/or empirically).

The reviewers rate this paper to be borderline.
Decision: Reject