Improving the Reconstruction of Disentangled Representation Learners via Multi-Stage Modelling

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 β-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 . 1 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 at 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 ∈ RD can be meaningfully described by a small set of latent factors H ∈ RK , where K < D. More precisely , the observation ( X = x ) is assumed to be generated by first sampling a set of low dimensional factors h from a simple prior distribution p ( H ) and then sampling x ∼ pθ ( X|h ) . DGMs realize pθ through a deep neural network also known as 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 qφ ( H|x ) . 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 pθ ( X ) . Learning disentangled factors h ∼ qφ ( H|x ) 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 et 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 a 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 as means for obtaining DR . The underlying assumption is 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 , z ) from p ( H ) and then sampling x ∼ pθ ( X|c , z ) ( Figure 2a ) . A series of works starting from ( Higgins et al. , 2017 ) enforce statistical independence of the latent factors H via regularization , up-weighting certain terms in the ELBO which penalize the ( aggregate ) posterior to be factorized over all or some of the latent dimensions ( Kumar et al. , 2017 ; Kim & Mnih , 2018 ; Chen et al. , 2018 ) ( see Section 2 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 = x ) can be thought to be generated by only sampling independent latent factors c from p ( C ) and then sampling x ∼ pθ ( X|c ) ( Figure 2b ) . This is 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 ) . In this paper , we propose a new graphical model for DR learning ( Figure 2c ) that allows for learning disentangled factors while also correctly realizing the data generating hypothesis that an observation is generated from independent and correlated factors . Importantly , the graphical model in Figure 2c is D-separated , meaning that any changes in the correlated latent variables Z will not influence the independent latent variables C. Generating an observation ( X = x ) from this model can then be done by sampling the independent factors c from p ( C ) , sampling a low-quality reconstruction y ∼ pθ ( Y |c ) , sampling the correlated factors z from p ( Z ) , and then finally sampling x ∼ pθ ( X|z , Y ) . The final reconstruction x depends both on z and c , however , any regularization needed to extract the independent factors c no longer diminishes the model capacity for the correlated factors z . To realize our proposed graphical model , we introduce MS-VAE , a multi-stage DGM that is implemented as follows : first , the disentangled representation C is learned using an existing DR learning method such as β-TCVAE ( Chen et al. , 2018 ) . Since the learned factors C are regularized to be statistically 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 Y 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. , 2018 ) ) . 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 1 . The reconstruction from β-TCVAE ( 1b ) is improved by MS-VAE ( 1c ) to better approximate the target ( 1a ) while maintaining conditioning on the disentangled factors , e.g . azimuth . To summarize our contributions : • We propose a new graphical model for DR learning ( Figure 2c ) that , due to the D-separation between its independent and correlated latent variables , alleviates the disentanglement vs. reconstruction trade-off . • We introduce MS-VAE , a multi-stage DGM that implements our proposed graphical model , achieving state-of-the-art reconstruction quality while maintaining the same level of disentanglement as preexisting methods . MS-VAE can be implemented using a 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-art ( Lezama , 2018 ) ) • Following ( Locatello et al. , 2018 ) , we test our framework with a wide-range of qualitative and quantitative tests to demonstrate the efficacy of the framework . We provide all code used for the experiments . 2 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 , ∫ p ( x ) log pθ ( x ) dx ≥ Ep ( x ) [ Eqφ ( h|x ) [ log pθ ( x|h ) ] −KL [ qφ ( h|x ) ‖p ( h ) ] ] ( 1 ) The first term is the expected reconstruction loss under the variational posterior qφ . The second term – the Kullback–Leibler ( 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 a Lagrange multiplier β . Follow-up works decompose the KL-term in equation 1 and up-weight only the KL-divergence between the aggregate posterior qφH ( H ) = ∫ X p ( x ) qφ ( H|x ) dx and the factorized aggregate posterior ( Kim & Mnih , 2018 ; Chen et al. , 2018 ; Jeong & Song , 2019 ) . Kumar et al . ( 2017 ) simply uses a factorized prior . These models introduce a 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 c from p ( C ) and then sampling x ∼ pθ ( X|c ) ( Figure 2b ) . This leads to low-quality reconstructions as the correlated factors are not utilized . Lezama ( 2018 ) attempts to tackle this issue with a teacher-student learning technique . First , the teacher model is trained to extract disentangled factors at the cost of reconstruction quality ( similar to the first stage of our model ) . Then , a student autoencoder ( AE ) model with a 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 a 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 et al. , 2016 ; Szabó et al. , 2018 ; Hu et al. , 2018 ) .

Summary: 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.