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 presents a consensus model over a base-level VAE model that has categorical, $\mathbf c$, and continuous, $\mathbf s$ latent components. The base-level model is instantiated a number of times (arms) where given a training instance $\mathbf x$ , its arm is fed with perturbations $\mathbf x_i$ of the instance. The different arms are constrained so that they  assign the same category to the different $\mathbf x_i$ perturbations. The main motivation for such a consensus constraint is, if I am not mistaken, that it avoids the mode collapse problem. In addition, accroding to the paper, the suggested approach can handle situations in which we have very different number of training instances in the different categories.