Mixture Representation Learning with Coupled Autoencoders

1 INTRODUCTION . Fitting mixed discrete-continuous models arises in many contexts . While continuous latent variables can be efficiently inferred with variational and adversarial formulations , inference of continuous and discrete factors in generalized mixture models remains challenging despite recent progress ( Jang et al. , 2016 ; Chen et al. , 2016 ; Dupont , 2018 ; Jeong & Song , 2019 ) . A pressing domain of application for such models is quantifying factors of biological variability in single-cell omic studies . The high-throughput and high-dimensional datasets produced by these studies document a previously unappreciated diversity of gene expression . In neuroscience , this poses the identification of cell types and cell states as a key research area to understand how neuronal circuits function ( Bargmann et al. , 2014 ) , where the notions of cell types and states can be considered as biological interpretations of discrete and continuous variability . While marker gene based studies suggest the existence of more than 100 neuronal cell types in just a single brain region , there is no agreement on such categorization and interpretation of the remaining continuous variability ( Seung & Sümbül , 2014 ; Zeng & Sanes , 2017 ; Tasic et al. , 2018 ) . Moreover , existing unsupervised , joint continuous-discrete learning methods are tailored for problems with relatively few and equally-abundant discrete components , and accurate inference remains out of reach for these applications . Deep generative models have previously been applied to single-cell datasets , where the focus is on the cluster identity and it is typically inferred by post-hoc analysis of a continuous factor ( Lopez et al. , 2018 ) . Deep Gaussian mixture models ( Dilokthanakul et al. , 2016 ; Johnson et al. , 2016 ; Jiang et al. , 2017 ) also focus on the identification of categories and do not take interpretability of the remaining continuous variability into account . To address the need for joint inference of interpretable discrete and continuous factors , various adversarial and variational methods have been proposed . While existing adversarial generative models , e.g . InfoGAN ( Chen et al. , 2016 ) , are susceptible to stability issues ( Higgins et al. , 2017 ; Kim & Mnih , 2018 ) , variational autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) emerge as efficient and more stable alternatives ( Tschannen et al. , 2018 ; Zhang et al. , 2018 ; Dupont , 2018 ; Jeong & Song , 2019 ) . VAE-based approaches approximate the mixture model by assuming a family of distributions qφ and select the member closest to the true model p. Popular choices in VAE implementations include ( 1 ) using KL divergence to compute discrepancy between qφ and p , and ( 2 ) using a multivariate Gaussian mixture distribution with uniformly distributed discrete and isotropic Gaussian distributed continuous priors . However , such choices may lead to underestimating the posterior variance ( Minka et al. , 2005 ; Blei et al. , 2017 ) . Solutions to resolve this issue are mainly applicable in low-dimensional spaces or for continuous factors alone ( Deasy et al. , 2020 ; Kingma et al. , 2016 ; Ranganath et al. , 2016 ; Quiroz et al. , 2018 ) . Inspired by collective decision making , we introduce a variational framework using multiple autoencoding arms to jointly infer interpretable finite discrete ( categorical ) and continuous factors in the presence of high-dimensional discrete space . Coupled-autoencoders have been previously studied in the context of multi-modal recordings , where each arm learns only a continuous latent representation for one of the data modalities ( Feng et al. , 2014 ; Gala et al. , 2019 ; Lee & Pavlovic , 2020 ) . Here , we develop a novel pairwise-coupled autoencoder framework for a single data modality . The proposed framework imposes a consensus constraint on the categorical posterior at the time of training and allows dependencies between continuous and categorical factors . We define the consensus constraint based on the Aitchison geometry in the probability simplex , which avoids the mode collapse problem . We show that the coupled multi-arm architecture enhances accuracy , robustness , and interpretability of the inferred factors without requiring any priors on the relative abundances of categories . Finally , on datasets profiling different cortical regions in the mammalian brain , we show that our method can be used to discover neuronal types as discrete categories and type-specific genes regulating the continuous within-type variability , such as metabolic state or disease state . Related work . There is an extensive body of research on clustering in mixture models ( Dilokthanakul et al. , 2016 ; Jiang et al. , 2017 ; Tian et al. , 2017 ; Guo et al. , 2016 ; Locatello et al. , 2018b ) . The idea of improving the clustering performance through seeking a consensus and co-training and ensembling across multiple observations has been explored in both unsupervised ( Monti et al. , 2003 ; Kumar & Daumé , 2011 ) and semi-supervised contexts ( Blum & Mitchell , 1998 ) . However , these methods do not consider the underlying continuous variabilities across observations . Moreover , unlike ensemble methods , which pool the results of different trained workers , autoencoding arms seek a consensus at the time of learning in our framework . The proposed framework does not need any supervision since the individual arms provide a form of prior or weak supervision for each other . In this regard , our paper is related to a body of work that attempts to improve representation learning by using semi-supervised or group-based settings ( Bouchacourt et al. , 2017 ; Hosoya , 2019 ; Nemeth , 2020 ) . Bouchacourt et al . ( 2017 ) demonstrated a multi-level variational autoencoder ( MLVAE ) as a semi-supervised VAE by revealing that observations within groups share the same type . Hosoya ( 2019 ) and Nemeth ( 2020 ) attempted to improve MLVAE by imposing a weaker condition to the grouped data . In recent studies ( Shu et al. , 2019 ; Locatello et al. , 2020 ) , a weakly supervised variational setting has been proposed for disentangled representation learning by providing pairs of observations that share at least one underlying factor . These studies rely on learning latent variables in continuous spaces , and have been applied only to image datasets with low-dimensional latent representations . Recent advances in structured variational methods , such as imposing a prior ( Ranganath et al. , 2016 ) or spatio-temporal dependencies ( Quiroz et al. , 2018 ) on the latent distribution parameters , allow for scaling to larger dimensions . However , these solutions are not directly applicable to the discrete space , which will be addressed in our A-arm VAE framework . 2 SINGLE MIXTURE VAE FRAMEWORK . For an observation x ∈ RD , a VAE learns a generative model pθ ( x|z ) and a variational distribution qφ ( z|x ) , where z ∈ RM is a latent variable with a parameterized distribution p ( z ) and M D ( Kingma & Welling , 2013 ) . Disentangling different sources of variability into different dimensions of z enables an interpretable selection of latent factors ( Higgins et al. , 2017 ; Locatello et al. , 2018a ) . However , the interplay between continuous and discrete variabilities present in many real-world datasets is often overlooked by existing methods . This problem can be addressed within the VAE framework in an unsupervised fashion by introducing a categorical latent variable c denoting the class label , alongside the continuous latent variable s. We refer to the continuous variable s as the state or style variable interchangeably . Assuming s and c are independent random variables , the evidence lower bound ( ELBO ) ( Blei et al. , 2017 ) for a single mixture VAE with the distributions parameterized by θ and φ is given by , L ( φ , θ ) = Eqφ ( s , c|x ) [ log pθ ( x|s , c ) ] − DKL ( qφ ( s|x ) ‖p ( s ) ) − DKL ( qφ ( c|x ) ‖p ( c ) ) . ( 1 ) Maximizing ELBO in Eq . 1 imposes characteristics on q ( s|x ) and q ( c|x ) that can result in underestimation of posterior probabilities such as the mode collapse problem , where the network ignores a subset of latent variables ( Minka et al. , 2005 ; Blei et al. , 2017 ) . Recently , VAE-based solutions were proposed by imposing a uniform structure on p ( c ) : akin to β-VAE ( Higgins et al. , 2017 ; Burgess et al. , 2018 ) , JointVAE ( Dupont , 2018 ) modifies the ELBO by assigning a pair of controlled information capacities for each variational factor , i.e . Cs ∈ R|s| and Cc ∈ R|c| . The main drawback of JointVAE is that its performance is tied to heuristic tuning of |s|× |c| capacities over training iterations so that it is vulnerable to mode collapse in high-dimensional settings . Another recent VAE-based mixture model solution , CascadeVAE ( Jeong & Song , 2019 ) , maximizes the ELBO through a semi-gradient-based algorithm by iterating over two separate optimizations for the continuous and categorical variables . While the separation of the optimization steps avoids the mode collapse problem , this separation is valid only when the categorical variable is uniformly distributed . Therefore , its performance strongly depends on the clusters having similar abundances in the dataset . Thus , earlier solutions fall short of learning interpretable mixture representations with high-dimensional discrete variables in real-world applications . In addition to the issues discussed above , the performance and interpretability of those approaches are further limited by the common assumption that the continuous variable representing the style of the data is independent of the categorical variable . In practice , style often depends on the class label . For instance , even for the well-studied MNIST dataset , the histograms of common digit styles , e.g . “ width ” , markedly vary for different digits ( Supplementary Section I ) . Moreover , further analysis of the identified continuous factor in the earlier approaches reveals that the independence assumption among q ( s|x ) and q ( c|x ) can be significantly violated ( see Supplementary Sections H and I ) . 3 COUPLED MIXTURE VAE FRAMEWORK . The key intuition behind multi-arm networks is cooperation to improve posterior estimation . While the context is different , the popular phrase “ wisdom of the crowd ” ( Surowiecki , 2005 ) can nevertheless be revealing : when a crowd ( multiple arms ) needs to make a decision , multiple estimates can increase the expected probability of a correct choice . 3.1 A-ARM VAE FRAMEWORK . We define the A-arm VAE as an A-tuple of independent and architecturally identical autoencoding arms , where the a-th arm parameterizes a mixture model distribution ( Fig . 1a ) . In this framework , individual arms receive a collection of non-identical copies , { xa , xb , . . . } of the given sample , x , belonging to the same category . While each arm has its own mixture representation with potentially non-identical parameters , all arms cooperate to learn q ( ca|xa ) , where ca = cb = · · · , via a cost function at the time of training . Accordingly , a crowd of VAEs with A arms can be formulated as a collection of constrained variational objectives as follows . max Ls1|c1 ( φ1 , θ1 ) + · · ·+ LsA|cA ( φA , θA ) s.t . c1 = · · · = cA ( 2 ) where Lsa|ca ( φa , θa ) is the variational loss for arm a , Lsa|ca ( φa , θa ) = Eq ( sa , ca|xa ) [ log p ( xa|sa , ca ) ] − Eq ( ca|xa ) [ DKL ( q ( sa|ca , xa ) ‖p ( sa|ca ) ) ] − Eq ( sa|ca , xa ) [ DKL ( q ( ca|xa ) ‖p ( ca ) ) ] . ( 3 ) In Eq . 3 , the variational loss for each arm is defined according to the graphical model in Fig . 1b , which is built upon the traditional ELBO in Eq . 1 by conditioning the continuous state on the categorical variable ( derivation in Supplementary Section B ) . Therefore , learning an interpretable decomposition of the data relies on accurate assignment ( inference ) of the categorical latent factor . Propositions 1 and 2 below show that the shared categorical assignment inferred from q ( c|x1 , · · · , xA ) , under the c = c1 = · · · = cA constraint of the multi-arm framework improves the accuracy of the categorical assignment on expectation . Proposition 1 . Consider the problem of mixture representation learning in a multi-arm VAE framework . For independent samples from category m , i.e . xi ∼ p ( x|m ) , Eq ( x|m ) [ log q ( c = m| { xi } 1 : A ) ] > Eq ( x|m ) [ log q ( c = m| { xi } 1 : B ) ] s.t . c = c1 = · · · = cA s.t . c = c1 = · · · = cB ( 4 ) if q ( m|xi ) < 1 and A > B ≥ 1 denote the number of arms . ( Proof in Supplementary Section A ) Thus , having more arms increases the expected log posterior for the true categorical latent variable unless it is already at its maximum . Proposition 2 . In the A-arm VAE framework , there exists an A that guarantees a true categorical assignment on expectation . That is , m = arg max c Eq ( x|m ) [ log q ( c| { xi } 1 : A ) ] , s.t . c = c1 = · · · = cA . ( 5 ) ( Proof in Supplementary Section A ) Accordingly , the consensus constraint is sufficient to enhance inference for mixture representations in the A-arm VAE framework . Our theoretical results show that the required number of arms satisfying Eq . 5 is a function of the categorical distribution and the likelihood ( Eq . 15 , Supplementary Section A ) . In the particular case of uniformly distributed categories , one pair of coupled arms is enough to satisfy Eq . 5 ( see Corollary 1 , Supplementary Section A ) . We emphasize that the proposed framework does not require any weak supervision as in ( Bouchacourt et al. , 2017 ) . Instead , it relies on representations that are invariant under non-identical copies of observations . Moreover , unlike ( Bouchacourt et al. , 2017 ; Shu et al. , 2019 ; Locatello et al. , 2020 ) , the multi-arm framework is not restricted to the continuous space . Arms observe non-identical copies of samples . In the A-arm VAE framework , arms receive non-identical observations that share the discrete variational factor . To achieve this in a fully unsupervised setting , we use type-preserving data augmentation that generates independent and identically distributed copies of data while preserving its categorical identity . For image datasets , conventional transformations such as rotation , scaling , or translation can serve as type-preserving augmentations . However , for non-image datasets , e.g . single-cell data , we seek a generative model that learns transformations representing within-class variability in an unsupervised manner . To this end , inspired by DAGAN ( Antoniou et al. , 2017 ) and VAE-GAN ( Larsen et al. , 2016 ) , we develop a generative model to provide collections of observations for our multi-arm framework ( Supplementary Section F ) . The proposed generative model learns to generate augmented samples in the vicinity of given samples in the latent space , without knowing their types ( Eq . 70 ) . In Supplementary Section A , Remark 2 , we further discuss an under-exploration scenario in data augmentation , in which the augmented samples are not independently distributed and are concentrated around the given sample .

Summary: The paper proposes to learn latent variable model with mixed discrete and continuous latent variables. Specifically, the framework utilizes multiple pairwise-coupled autoencoding arms to learn shared categorical variable. The experiments on dSprites, MNIST and RNA-seq datasets verify the effectiveness of the model.