Isolating Latent Structure with Cross-population Variational Autoencoders

A significant body of recent work has examined variational autoencoders as a powerful approach for tasks which involve modeling the distribution of complex data such as images and text . In this work , we present a framework for modeling multiple data sets which come from differing distributions but which share some common latent structure . By incorporating architectural constraints and using a mutual information regularized form of the variational objective , our method successfully models differing data populations while explicitly encouraging the isolation of the shared and private latent factors . This enables our model to learn useful shared structure across similar tasks and to disentangle cross-population representations in a weakly supervised way . We demonstrate the utility of our method on several applications including image denoising , sub-group discovery , and continual learning . 1 INTRODUCTION . Unsupervised learning of latent representations is widely used for dimensionality reduction , density estimation , and structure or sub-group discovery among other applications . Methods for recovering such representations typically rely on the assumption that the observed data is a manifestation of only a limited number of factors of variation ( Locatello et al. , 2019 ; Bengio et al. , 2013 ) . The variational autoencoder ( VAE ) ( Kingma & Welling , 2013 ) , a combination of a non-linear latent variable model and an amortized inference scheme ( Dayan et al. , 1995 ) , is a popular method for recovering such latent structure . VAEs and their extensions have received considerable attention in recent years and have been shown useful for modeling text ( Miao et al. , 2016 ) , images ( Gulrajani et al. , 2016 ) , and other data exhibiting complex correlations ( Gómez-Bombarelli et al. , 2018 ) . However , barring a few exceptions ( Bouchacourt et al. , 2018 ; Severson et al. , 2019 ) , this line of work has assumed the data to be independent and identically distributed . In this work , we consider the task of modeling independent but not identically distributed data . We are particularly interested in analyzing data comprising two or more distinct but related subpopulations . Such data arise frequently in practice . For instance , patients suffering from an ailment may exhibit symptomatic heterogeneity based on gender or environmental factors , documents in a corpus may exhibit semantic or syntactic similarities based on genre or authorship , images may cluster depending on the image subject . Our goal is to provide rich descriptions of such data by recovering latent representations that disentangle factors of variation common to all populations from those that are unique to a particular population . Motivated by this challenge , we propose non-linear latent variable models that explicitly account for the heterogeneity in the data . Coupled with an inference procedure that encourages high ( low ) mutual information between the latent representation and the data within ( across ) a population , we show that our models are indeed able to recover population-specific representations that are salient and disentangled across differing populations as well as shared representations that isolate commonalities between populations . Through careful experiments , we vet the effectiveness of our approach and demonstrate that the learned representations are useful for a diverse set of applications including image denoising , unsupervised sub-group discovery , and continual learning . 2 METHOD . 2.1 CROSS-POPULATION VARIATIONAL AUTOENCODER . We propose a generative model of a data instance xki belonging to population k using two latent variables zki and tki . The latent variables are mapped to the observed space using non-linear mappings , fθs ( zki ) and fθk ( tki ) , each parameterized by a neural network . While we share the parameters θs among all populations , θk are only shared among instances belonging to the population k. This construction encourages the model to capture common latent structure in θs , while allowing θk to focus on factors of variation unique to the particular population k. We combine the contributions from the two mappings using a function g. The generative procedure can be summarized as , zki ∼ N ( 0 , I ) , tki ∼ N ( 0 , I ) , xki | zki , tki ∼ p ( g ( fθs ( zi ) , fθk ( tki ) ) ) , ∀i ∈ { 1 . . . nk } , ∀k ∈ { 1 , . . . , K } , ( 1 ) where p is an appropriately chosen distribution for modeling the observed data . In our experiments , we use an additive function g ( a , b ) = a+ b , and select p to be the Gaussian distribution with mean parameterized by g and an isotropic diagonal covariance matrix , Ψ . We emphasize that other choices of function can easily be incorporated into the model . Exploring the space of aggregation functions is planned future work . 2.2 AMORTIZED VARIATIONAL INFERENCE . Amoritized variational inference ( Dayan et al. , 1995 ; Gershman & Goodman , 2014 ) with the aid of reparameterized gradients ( Rezende et al. , 2014 ; Kingma & Welling , 2013 ) is straightforward to implement for the model described in equation 1 . We assume that the variational approximation factorizes conditioned on the observation , qφ ( zki , tki | xki ) = qφs ( zki | xki ) qφk ( tki | xki ) . ( 2 ) We parameterize the variational distribution with a single , shared inference network with two distinct outputs , one for each latent variable . Finally , the model and variational parameters are jointly learned by optimizing the evidence lower bound ( ELBO ) , L ( θ , φ ) = K∑ k=1 nk∑ i=1 Eqφ ( zki , tki|xki ) [ log p ( xki | zki , tki ; θs , θk ) ] −DKL ( qφ ( zki , tki | xki ) ||p ( zki , tki ) ) . ( 3 ) Importantly , in the above setup , the per-population model and variational parameters θk and φk are learned only from the data in that population , allowing them to learn population specific representations . We refer to this combination of the model described in equation 1 and the inference network defined above as the cross-population variational auto-encoder ( CPVAE ) , owing to its similarity with the VAE ( Kingma & Welling , 2013 ) . See Figure 1 for a graphical representation of the CPVAE and comparison to the VAE . 2.3 MUTUAL INFORMATION REGULARIZED INFERENCE . By employing population-specific mappings , CPVAE prevents latent factors of variation unique to one population from ‘ leaking ’ into the private representation of another population . However , the model structure does not prevent leakage between the shared and population specific representations , since every data instance is generated by a combination of the shared and population specific representation . In fact , when CPVAE is trained by maximizing the ELBO in equation 3 we find that the private representations of a population often exhibit latent features from the shared space and vice versa . Inspired in part by the InfoVAE ( Zhao et al. , 2017 ) , we discourage such inferences by minimizing the mutual information between population-specific latent variables and data from non-corresponding populations while maximizing the mutual information within a population . Let xk = { xk1 , . . . , xknk } denote the set of all data instances belonging to population k , let tk , zk be the set of corresponding latent variables , and x−k = { xj ; ∀j 6= k ∈ K } denote the set of all non-corresponding populations . We learn the CPVAE by maximizing , J ( θ , φ ) = L ( θ , φ ) + K∑ k nk ( Iq ( xk ; tk ) − Iq ( x−k ; t̃k ) ) , ( 4 ) where t̃k = { { t̃ji } nj i=1 } j 6=k and t̃ji ∼ qφk ( t̃ji | xji ) where xji ∈ x−k , i.e. , the result of encoding members of the non-corresponding populations , x−k with the parameters φk of the current population ’ s inference network . This function can be simplified to a tractable objective by re-writing the mutual information term : Iq ( xk ; tk ) =− Eqφk ( xk , tk ) log qφk ( tk ) qφk ( tk|xk ) =− Eqφk ( xk , tk ) log qφk ( tk ) p ( tk ) qφk ( tk|xk ) p ( tk ) =− EpD ( xk ) Eqφk ( tk|xk ) log p ( tk ) qφk ( tk|xk ) − Eqφk ( tk ) log qφk ( tk ) p ( tk ) = 1 nk nk∑ i=1 [ DKL ( qφk ( tki|xki ) ||p ( tki ) ) ] −DKL ( qφk ( tk ) ||p ( tk ) ) , ( 5 ) where we assume qφk ( xk , tk ) = pD ( xk ) qφk ( tk | xk ) and pD ( xk ) is the empirical distribution . Note that when we take Iq ( xk ; tk ) − Iq ( x−k ; tk ) as in our objective , the intractable marginal KL term conveniently cancels out . The complete objective function is J ( θ , φ ) = K∑ k=1 [ nk∑ i=1 [ Eqφ ( zki , tki|xki ) [ log p ( xki | zki , tki ; θs , θk ) ] −DKL ( qφs ( zki|xki ) ||p ( zki ) ) ] − nk K − 1 K∑ j 6=k 1 nj nj∑ i=1 DKL ( qφk ( t̃ji|xji ) ||p ( tji ) ) ] , ( 6 ) This objective differs from the standard variational objective only in that the divergence of the private latent prior for any population tk is minimized for its contrasting populations x−k and not its corresponding population xk . That is , we minimize DKL ( qφk ( t̃k|x−k ) ||p ( tk ) ) rather than DKL ( qφk ( tk|xk ) ||p ( tk ) ) . As we will show , this modification reduces entanglement between private latent vectors and enables the model to learn more salient representations of each population . Moreover , since Kullback Leibler divergence is always non-negative , J ( θ , φ ) in equation 6 remains a lower bound to the marginal likelihood . In some scenarios , it may desirable to increase the importance of the cross mutual information term Iq ( x−k ; tk ) to our objective , further discouraging shared features from leaking into populationspecific representations . To this end , we can simply add a scaling constant α ≥ 1 to the final term in equation 6 . In our experiments , we often find substantially improved results by beginning with α = 1 and gradually annealing its value over training . A summary of the complete training procedure can be found in Algorithm 1 . Algorithm 1 Training procedure of CPVAE Initialize conditional parameters θ = { θs } ∪ { θk ; ∀k ∈ K } ; Initialize variational parameters φ = { φs } ∪ { φk ; ∀k ∈ K } ; repeat Sample mini-batch from each population { { xki } mki=1 } Kk=1 for k ∈ 1 . . .K do Sample shared codes zk ∼ qφs ( zk|xk ) ; Sample private codes tk ∼ qφk ( tk|xk ) ; for j 6= k ∈ 1 . . .K do Sample fictitious codes t̃j ∼ qφk ( t̃j |xj ) ; end end Calculate J ( θ , φ ) as in ( 6 ) ; Update θt+1 , φt+1 ← θt , φt according to ascending gradient estimate of J ( θ , φ ) ; until convergence ;

Summary: This paper introduces a novel model called "Cross-Population Variational Autoencoder (CPVAE), which is designed to model data from different distributions sharing some common structure. The proposed generative model utilizes both shared and private per-population latent variables. In order to restrict shared latent variables from "leaking" into private representations, the authors introduce an information-theoretic regularizer. This regularizer forces private population representations to: (a) maximize mutual information with input samples from their population, and (b) minimize mutual information with input samples from other populations. In other words, private representations are forced to be "meaningful" on the corresponding population alone.