Property Controllable Variational Autoencoder via Invertible Mutual Dependence

1 INTRODUCTION . Important progress has been made towards learning the underlying low-dimensional representation and generative process of complex high dimensional data such as images ( Pu et al. , 2016 ) , natural languages ( Bowman et al. , 2016 ) , chemical molecules ( Kadurin et al. , 2017 ; Guo et al. , 2019 ) and geo-spatial data ( Zhao , 2020 ) via deep generative models . In recent years , a surge of research has developed new ways to further enhance the disentanglement and independence of the latent dimensions , creating models with better robustness , improved interpretability , and greater generalizability with inductive bias ( see Figures 1 ( a ) and 1 ( b ) ) ( Kingma et al. , 2014 ; Kulkarni et al. , 2015 ; Creswell et al. , 2017 ) or without any bias ( Higgins et al. , 2017 ; Chen et al. , 2018 ; Kumar et al. , 2018 ) . Although it is generally assumed that the complex data is generated from the latent representations , their latent dimensions are typically not associated with physical meaning and hence can not reflect real data generation mechanisms such as the relationships between structural and functional characteristics . A critical problem that remains unsolved is how to best identify and enforce the correspondence between the learned latent dimensions and key aspects of the data , such as the bio-physical properties of a molecule . Knowing such properties is crucial for many applications that depend on being able to interpret and control the data generation process with the desired properties . In an effort to achieve this , several researchers ( Klys et al. , 2018 ; Locatello et al. , 2019b ) have suggested methods that enforce a subset of latent dimensions correspond to targeted categorical properties , as shown in Figure 1 ( c ) . Though the initial results have been encouraging , critical challenges remain unsolved such as : ( 1 ) Difficulty in handling continuous-valued properties . The control imposed on data generation limits existing techniques to categorical ( typically binary ) properties , to enable tractable model inference and sufficient coverage of the data . However , continuous-valued properties ( e.g. , the scale and light level of images ) are also common in real world data , while their model inference usually can be easily intractable . Also , many cases require to generate data ∗Corresponding author : liang.zhao @ emory.edu with properties of which the values are unseen during training process . This can not be achieved by conventional techniques such as conditional models without making strong assumption on the model distributions . ( 2 ) Difficulty in efficiently enhancing mutual independence among latent variables relevant and irrelevant to the properties . This problem requires to ensure that each property is only correlated to its corresponding latent variable ( s ) and independent of all the others . Directly enforcing such mutual independence inherently between all pairs of latent variables incurs quadratic number of optimization efforts . Hence an efficient way is imperative . ( 3 ) Difficulty in capturing and controlling correlated properties . It is feasible that several independent latent variables can capture multiple independent properties . But when the properties are correlated , they can not be “ one-on-one ” mapped to corresponding independent latent variables anymore . However , correlated properties are commonly found in formatting a real world data . To solve the above challenges , we propose a new model , Property-controllable VAE ( PCVAE ) , where a new Bayesian model is proposed to inductively bias the latent representation using explicit data properties via novel group-wise and property-wise disentanglement terms . Each data property is seamlessly linked to the corresponding latent variable by innovatively enforcing an invertible mutual dependence between them , as shown in Figure 1 ( d ) . Hence , when generating data , the corresponding latent variables are manipulated to simultaneously control multiple desired properties without influencing the others . We have also further extended our model to handle inter-correlated properties . Our key contributions are summarized as follows : • A new Bayesian model that inductively biases the latent representation using explicit real data properties is proposed . A variational inference strategy and inference model have been customized to ensure effective Bayesian inference . • Group-wise and property-wise disentanglement terms are proposed to enhance the mutual independence among property , relevant and irrelevant latent variables . • The invertible mutual dependence between property-latent variable pair is achieved by enforcing an invertibility constraint over a residual-based decoder . • The quantitative and qualitative evaluation performed for this study revealed our PCVAE outperforms existing methods by up to 28 % in capturing and 65 % in manipulating the desired properties . 2 RELATED WORKS . Disentanglement Representation Learning . An important relevant area of research is disentangled representation learning ( Alemi et al. , 2017 ; Chen et al. , 2018 ; Higgins et al. , 2017 ; Kim & Mnih , 2018 ) , which structures the latent space by minimizing the mutual information between all pairs of latent variables . The goal here is to learn representations that separate out the underlying explanatory factors that are responsible for variations in the data , as these have been shown to be relatively resilient with respect to the complex variants involved ( Bengio et al. , 2013 ; Ma et al. , 2019 ; Guo et al. , 2020 ) , and thus can be used to enhance generalizability as well as improve robustness against adversarial attack . As noted by Locatello et al . ( 2019a ) , it is impossible for disentangled representation learning to capture the desired properties without supervision and inductive biases . Learning latent representations via supervision . This ensures that the latent variables capture the desired properties though supervision , generally by directly defining properties as latent variables in the model ( Locatello et al. , 2019b ) . Unfortunately , apart from providing an explicit variable for the labelled property , this yields no other easily interpretable structures , such as discovering latent variables that are correlated to the properties , as the model proposed in the current study does . This is also an issue with other methods of structuring latent space that have been explored , such as batching data according to labels ( Kulkarni et al. , 2015 ; Zhang et al. , 2020 ) or using a discriminator network in a non-generative model ( Lample et al. , 2017 ) . Some researchers addressed this problem by introducing the architecture bias through a two-way factored autoencoder and realize the supervision based on a pair-wise contrastive loss ( Gyawali et al. , 2019 ) . Other researchers addressed this problem by linking latent variables with observed labels through adversarial learning ( Creswell et al. , 2017 ; Edwards & Storkey , 2015 ; Ganin et al. , 2016 ; Mathieu et al. , 2016 ) . The most relevant work for our purpose is CSVAE ( Klys et al. , 2018 ) , where a subset of latent variables are correlated with binary properties via an adversarial learning . All the above works can not handle multiple continuous-valued properties due to their strict assumptions on the distribution of properties . Data manipulation and generation . Here , trained machine learning models are utilized to manipulate and generate data in a controllable way with the desired properties , which is especially useful for applications in the image domain . Several works have specifically considered transferring attributes in images , which is the same goal as that in the CASVE . These earlier works ( Zhou et al. , 2017 ; Xiao et al. , 2017 ; 2018 ) all transfer attributes from a source image onto a target image . These models can only perform categorical attribute transformation between images ( e.g. , “ splice the beard style of image A onto image B ” ) , but only through interpolation between existing images . Once trained , our proposed model can generate an objects with any value of a certain property ( either observed or unobserved during training ) that can be encoded in the subset of latent variables . 3 PROPERTY CONTROLLABLE VAE . 3.1 PROBLEM FORMULATION . Suppose we are given a dataset D where each data instance is ( x , y ) with x ∈ Rn and y = { yk ∈ R } Kk=1 to represent K properties of interest of x . For example , if x is a molecule , then we may have properties of interest , such as cLogP and cLogS . We assume that the data ( x , y ) are generated by some random process from continuous latent random variables ( z , w ) . Each variable in w controls one of the properties of interest in y , while the variables in z control all the other aspects of x . Our goal is to learn such a generative model involving ( x , y ) and ( z , w ) , where the subset of variables ( i.e. , z ) are disentangled from subset w , and the variables inside w are disentangled from each other . Once this model has been learned , then we can expect different elements of variables in w to control different properties of interest , which is a highly desirable goal for many data generation downstream tasks . For example , we may want to decrease the value of a specific property ( e.g. , protein energy ) by changing the value of the corresponding element in w. It is also possible to directly set a desired property value ( e.g. , the mass of a molecule ) and then generate the corresponding x with this target value ( i.e. , a molecule with the target mass value ) . 3.2 OVERALL OBJECTIVE . In this section , we first introduce the Bayesian variational inference of PCVAE . Then we introduce the group-wise and property-wise disentanglement terms as part of the overall objective . Following this , an invertibility constraint is introduced to enforce mutual dependence between each propertylatent variable pair . At last , PCVAE is extended to capture and control multiple correlated properties . 3.2.1 BAYESIAN VARIATIONAL INFERENCE OF PCVAE . The goal in Section 3.1 requires us to not only model the dependence between x and ( w , z ) for latent representation learning and data generation , but also model the dependence between y and w for property manipulation . We propose to achieve this by maximizing a form of variational lower bound on the joint log likelihood p ( x , y ) of our model . Given an approximate posterior q ( z , w∣x , y ) , we can use the Jensen ’ s equality to obtain the variational lower bound of p ( x , y ) as : log p ( x , y ) = logEq ( z , w∣x , y ) [ p ( x , y , w , z ) /q ( z , w∣x , y ) ] ≥ Eq ( z , w∣x , y ) [ log p ( x , y , w , z ) /q ( z , w∣x , y ) ] . ( 1 ) The joint likelihood log p ( x , y , w , z ) can be decomposed as log p ( x , y∣z , w ) + log p ( z , w ) . We have two assumptions : ( 1 ) w only encodes the information from y , namely , x and y are conditionally independent given w ( i.e. , x ⊥ y∣w ) ; ( 2 ) z is independent from w and y , namely z ⊥ w and z ⊥ y , which is equal to y ⊥ z∣w ( see derivation in Appendix A.3 ) . First , based on the two assumptions , we can get x ⊥ y∣ ( z , w ) ( see derivation in Appendix A.4 ) . Thus , we have log p ( x , y∣z , w ) = log p ( x∣z , w ) + log p ( y∣z , w ) . Then , based on the assumption y ⊥ z∣w , we can have log p ( y∣z , w ) = log p ( y∣w ) . Then we get log p ( x , y∣z , w ) = log p ( x∣z , w ) + log p ( y∣w ) . To explicitly represent the dependence between x and ( z , w ) as well as the dependence between y and w , we can parameterize the joint log-likelihood as log pθ , γ ( x , y , w , z ) with θ and γ as : log pθ , γ ( x , y , w , z ) = log pθ ( x∣z , w ) + log p ( z , w ) + log pγ ( y∣w ) . ( 2 ) Given the condition that a parameterized qφ ( z , w∣x , y ) = qφ ( z , w∣x ) = qφ ( z∣x ) qφ ( w∣x ) ( since the information on y is included in x ) , by taking Eq . 2 into the above variational lower bound term in Eq . 1 , we obtain the negative part as an upper bound on − log pθ , γ ( x , y ) ( as shown in the right sub-figure of Figure 1 ( d ) ) : L1 = − Eqφ ( z , w∣x ) [ log pθ ( x∣z , w ) ] − Eqφ ( w∣x ) [ log pγ ( y∣w ) ] +DKL ( qφ ( z , w∣x ) ∣∣p ( z , w ) ) ( 3 ) This gives us the proposed Bayesian variational inference of PCVAE . The detailed derivation of Eq . 3 can be found in Appendix A.1 . As there are K properties of interest in y which are assigned and disentangled by the latent variables in w , the second term in Eq . 3 can be detailed as ∑Kk Eqφ ( w∣x ) [ log pγ ( yk∣wk ) ] .

Summary: To encourage disentanglement in the latent space of a variational autoencoder (VAE), the authors propose to learn two sets of latent z and w: the dimensions of w are independent of each other and each dimension w_i maps to a known ground truth generating factor y_i. Latent z captures all the other factors. The well studied Total Correlation regularisation is used to enforce the independence of z and w, and the same is used to enforce the independence of the dimensions of w. Each dimension is learned to predict a corresponding ground truth factor. The key difference from the previous approach is the use of invertible and Lipschitz smooth mapping to learn monotonic mappings from w to y.