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