Year: 2021

Abstract:
Latent representations help unravel complex phenomena. While continuous latent variables can be efficiently inferred, fitting mixed discrete-continuous models remains challenging despite recent progress, especially when the discrete factor dimensionality is large. A pressing application for such mixture representations is the analysis of single-cell omic datasets to understand neuronal diversity and its molecular underpinnings. Here, we propose an unsupervised variational framework using multiple interacting networks called cpl-mixVAE that significantly outperforms state-of-the-art in high-dimensional discrete settings. cpl-mixVAE introduces a consensus constraint on discrete factors of variability across the networks, which regularizes the mixture representations at the time of training. We justify the use of this framework with theoretical results and validate it with experiments on benchmark datasets. We demonstrate that our approach discovers interpretable discrete and continuous variables describing neuronal identity in two single-cell RNA sequencing datasets, each profiling over a hundred cortical neuron types.

heading: 1 INTRODUCTION
text: 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.

heading: 2 SINGLE MIXTURE VAE FRAMEWORK
text: 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).

heading: 3 COUPLED MIXTURE VAE FRAMEWORK
text: 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.

heading: 3.1 A-ARM VAE FRAMEWORK
text: 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.

heading: 3.2 CPL-MIXVAE: PAIRWISE COUPLING IN A-ARM VAE
text: In the A-arm VAE framework, the mixture representation is obtained through the optimization in Eq. 2. Not only is it challenging to solve the maximization in Eq. 2 due to the equality constraint, but the objective remains a function of p(c) which is unknown, and typically non-uniform. To overcome this, we use an equivalent formulation for Eq. 2 by applying the pairwise coupling paradigm as follows (details of derivation in Supplementary Section C):
max
A∑ a=1 (A− 1) ( Eq(sa,ca|xa) [log p(xa|sa, ca)]− Eq(ca|xa) [DKL (q(sa|ca,xa)‖p(sa|ca))] ) −∑
a<b
Eq(sa|ca,xa)Eq(sb|cb,xb) [DKL (q(ca|xa)q(cb|xb)‖p(ca, cb))]
s.t. ca = cb ∀a, b ∈ [1, A], a < b (6) We relax the optimization in Eq. 6 into an unconstrained problem by marginalizing the joint distribution over a mismatch measure between categorical variables (see Supplementary Section D):
max A∑ a=1 (A− 1) ( Eq(sa,ca|xa) [log p(xa|sa, ca)]− Eq(ca|xa) [DKL (q(sa|ca,xa)‖p(sa|ca))] ) +∑
a<b
H(ca|xa) +H(cb|xb)− λEq(ca,cb|xa,xb) [ d2(ca, cb) ] (7)
In Eq. 7, in addition to entropy-based confidence penalties known as mode collapse regularizers (Pereyra et al., 2017), the distance measure d(ca, cb) encourages a consensus on the categorical assignment controlled by λ ≥ 0, the coupling hyperparameter. We refer to the model in Eq. 7 as cpl-mixVAE (Fig. 1a). In cpl-mixVAE, VAE arms try to achieve identical categorical assignments while independently learning their own style variables. In experiments, we set λ = 1 universally. While the bottleneck architecture already encourages interpretable continuous variables, this formulation can be easily extended to include an additional hyperparameter to promote disentanglement of continuous variables as in β-VAE (Higgins et al., 2017). Additional analyses to assess the sensitivity of the cpl-mixVAE’s performance to its coupling factor can be found in the Supplementary Section G.
It may be instructive to cast Eq. 7 in an equivalent constrained optimization form.
Remark 1. The A-arm VAE framework is a collection of constrained variational models as follows:
max A∑ a=1 Eq(sa,ca|xa) [log p(xa|sa, ca)]− Eq(ca|xa) [DKL (q(sa|ca,xa)‖p(sa|ca))] +H(ca|xa)
s.t. Eq(ca|xa) [ d2(ca, cb) ] < (8)
where denotes the strength of the consensus constraint. Here, cb indicates the assigned category by any one of the arms, b ∈ {1, . . . , A}, imposing structure on the discrete variable to approximate its prior distribution.
Distance between categorical variables. d(ca, cb) denotes the distance between a pair of |c|dimensional un-ordered categorical variables, which are associated with probability vectors with non-negative entries and sum-to-one constraint that form a K-dimensional simplex, where K = |c|. In the real space, a typical choice to compute the distance between two vectors is using Euclidean geometry. However, this geometry is not suitable for probability vectors. Here, we utilize Aitchison geometry (Aitchison, 1982; Egozcue et al., 2003), which defines a vector space on the simplex. Accordingly, the distance in the simplex, i.e. dSK (ca, cb) is defined as dSK (ca, cb) = ‖clr(ca)− clr(cb)‖2, ∀ca, cb ∈ SK , where clr(·) denotes the isometric centered-log-ratio transformation in the simplex. This categorical distance satisfies the conditions of a mathematical metric according to Aitchison geometry.

heading: 3.3 SEEKING CONSENSUS IN THE SIMPLEX
text: An instance of the mode collapse problem (Lucas et al., 2019) manifests itself in the minimization of dSK (ca, cb) (Eq. 7): its trivial local optima encourages the network to abuse the discrete latent factor by ignoring many of the available categories. In the extreme case, the representations can collapse onto a single category; ca = cb = c0. In this scenario, the continuous variable is compelled to act as a primary latent factor, while the model fails to deliver an interpretable mixture representation despite achieving an overall low loss value. To avoid such undesirable local equilibria while training, we add perturbations to the categorical representation of each arm. If posterior probabilities in the simplex have small dispersion, the perturbed distance calculation overstates the discrepancies. Thus, instead of minimizing d2SK (ca, cb), we minimize a perturbed distance
d2σ(ca, cb) = ∑ k ( σ−1ak log cak − σ −1 bk log cbk )2
, which corresponds to the distance between additively perturbed ca and cb vectors in Aitchison geometry. Here, σ2ak and σ 2 bk
indicate the mini-batch variances of the k-th category, for arms a and b. We next show that the perturbed distance dσ(·) is bounded by dSK (·) and non-negative values ρu, ρl: Proposition 3. Suppose ca, cb ∈ SK , where SK is a simplex ofK > 0 parts. If dSK (ca, cb) denotes the distance in Aitchison geometry and d2σ(ca, cb) = ∑ k ( σ−1ak log cak − σ −1 bk log cbk )2
denotes a perturbed distance, then
d2SK (ca, cb)− ρl ≤ d 2 σ (ca, cb) ≤ d2SK (ca, cb) + ρu
where ρu, ρl ≥ 0, ρu = K ( τ2σu + τ 2 c ) + 2∆στc, ρl = ∆2σ K −Kτ2σl , τc = maxk {log cak − log cbk}, τσu = max k {gk}, τσl = max k {−gk}, ∆σ = ∑ k gk, and gk = (σ−1ak − 1) log cak − (σ −1 bk − 1) log cbk . (Proof in Supplementary Section E)
Thus, when ca and cb are similar and their spread is not small, dσ(ca, cb) closely approximates dSK (ca, cb). Otherwise, it diverges from dSK (·) to avoid mode collapse.

heading: 4 EXPERIMENTS
text: We used four datasets: dSprites, MNIST, and two single-cell RNA sequencing (scRNA-seq) datasets; Smart-seq ALM-VISp (Tasic et al., 2018) and 10X MOp (Yao et al., 2021). Although dSprites and MNIST datasets do not require high-dimensional settings for mixture representation, to facilitate comparisons of cpl-mixVAE with earlier methods, first we report the results for these benchmark datasets. We trained three unsupervised VAE-based methods for mixture modeling: JointVAE (Dupont, 2018), CascadeVAE (Jeong & Song, 2019), and ours (cpl-mixVAE). For MNIST, we additionally trained the popular InfoGAN (Chen et al., 2016) as the most comparable GAN-based model. To show the interpretability of the mixture representations, (i) for the discrete latent factor, we report the accuracy (ACC) of categorical assignments and the DKL(q(c)‖p(c))), (ii) for the continuous variable, we perform latent traversal analysis by fixing the discrete factor and changing the continuous variable according to p(s|c,x). We calculated the accuracy by using minimum weight matching (Kuhn, 1955) to match the categorical variables obtained by cpl-mixVAE with the available cluster labels for each dataset. Additionally, we report the computational efficiency (number of iterations per second) to compare the training complexity of the multi-arm framework against earlier methods (Table 1). All reported numbers for cpl-mixVAE models are average accuracies calculated across arms. In VAE-based models, to sample from q(ca|xa), we use the Gumbel-softmax distribution (Jang et al., 2016; Maddison et al., 2014). In cpl-mixVAE, each arm received an augmented copy of the original input generated by the deep generative augmenter (Supplementary Section F) during training. Details of the network architectures and training settings can be found in Supplementary Section L.

heading: 4.1 BENCHMARK DATASETS
text: dSprites. dSprites is procedurally generated from discrete (3 shapes) and continuous (6 style factors: scale, rotation, and position) latent factors. Based on the uniform distribution of classes, we used a 2-arm cpl-mixVAE with |c| = 3 and |s| = 6 to learn interpretable representations. Results in Table 1 show that our method outperforms the other methods in terms of categorical assignment accuracy. In addition to demonstrating the traversal results (Fig. 2, bottom row), we report disentanglement scores (DS in Table 1). Even though the continuous factors do not depend on the discrete factors in this synthetic dataset, we did not change the architecture and expected the network to infer this independence. For a fair comparison, we used the same disentanglement metric implemented for CascadeVAE (Jeong & Song, 2019).
MNIST. Similarly, due to the uniform distribution of digit labels in MNIST, we again used a 2-arm cpl-mixVAE model. Following the convention (Dupont, 2018; Jeong & Song, 2019; Bouchacourt
et al., 2017), each arm of cpl-mixVAE uses a 10-dimensional categorical variable representing digits (type), and a 10-dimensional continuous random variable representing the writing style (state). Table 1 displays the accuracy of the categorical assignment and the discrepancy between q(c) and p(c) for InfoGAN, two 1-arm VAE methods (JointVAE and CascadeVAE), and cpl-mixVAE with 2 arms. Additionally, to isolate the impact of data augmentation in training, we trained JointVAE† and CascadeVAE† where the models were trained with the same augmented copies of the original MNIST dataset as cpl-mixVAE. The results in Table 1 suggest that data augmentation by itself does not enhance the performance. Fig. 2 (top row) illustrates the continuous latent traversals for four dimensions of the state variable inferred by cpl-mixVAE, where each row corresponds to a different dimension of the categorical variable, and the state variable monotonically changes across columns. Both results in Table 1 and Fig. 2 show that cpl-mixVAE achieved an interpretable mixture representation with the highest categorical assignment accuracy.
Summary. cpl-mixVAE improves the discrete density approximation and infers better mixture representations. It outperforms earlier methods, without using extraneous optimization or heuristic channel capacities. Beyond performance and robustness, its computational cost is also comparable to that of the baselines.

heading: 4.2 SINGLE-CELL RNA SEQUENCING DATA
text: In this dataset the observations are individual cells and each observation consists of expressions of thousands of genes. Here, we used two scRNA-seq datasets: (i) Smart-seq ALM-VISp (Tasic et al., 2018) and (ii) 10X MOp (Yao et al., 2021). The Smart-seq dataset includes transcriptomic profiles of more than 10, 000 genes for ∼ 22, 000 cells from the mouse anterior lateral motor cortex (ALM) and the primary visual cortex (VIPs). The 10X MOp dataset profiles ∼ 123000 cells in the mouse primary motor cortex (MOp) with the droplet-based 10X Genomics Chromium platform. 10X-based data often display more gene dropouts, especially for genes with lower expression levels. scRNA-seq datasets are significantly more complex than typical benchmark datasets due to (i) large number of cell types (discrete variable), and (ii) class imbalance; in the 10X MOp dataset, for instance, the mostand the least-abundant cell types include 17, 000 and 20 samples, respectively. Moreover, whether the observed diversity corresponds to discrete variability or a continuum is an ongoing debate in neuroscience (Scala et al., 2020). While using genes that are differentially expressed in subsets of cells, known as marker genes (MGs) (Trapnell, 2015) is a common approach to define cell types, the identified genes rarely obey the idealized MG definition in practice. Here, we focus on neuronal cells and use a subset of 5, 000 highest variance genes. The original MG-based studies for each dataset suggested 115 (Smart-seq data) (Tasic et al., 2018) and 140 (10X-based data) (Yao et al., 2021) discrete neuronal types.
Neuron type identification. Based on the suggested taxonomies in (Tasic et al., 2018; Yao et al., 2021), for the Smart-seq ALM-VISp data, we used 115- and 2-dimensional discrete and continuous variables, and for the 10X MOp data, we used 140- and 2-dimensional discrete and continuous latent variables. We compared the suggested cell types in (Tasic et al., 2018; Yao et al., 2021) with the discrete representations that are inferred from VAE models. Table 1 and Fig. 3(a-b) demonstrate the
performance of a 2-arm cpl-mixVAE model against JointVAE and CascadeVAE. In Fig. 3a.1 and 3b.1, we observe that for both datasets, JointVAE succeeds in identifying (sub)classes of neurons, e.g. excitatory class, or Sst subclass, but not neuronal types at leaf nodes. On the other hand, CascadeVAE learns an almost uniform distribution over all types despite a sizeable difference between the relative abundances of neuronal types (Fig. 3a.2 and 3b.2). Our results in Fig. 3(a-b) clearly show that cplmixVAE outperforms JointVAE and CascadeVAE in identifying meaningful known cell types. The confusion matrices in Fig. 3a.3 and 3b.3 demonstrate that even the inaccurate categorical assignments of cpl-mixVAE are still close to the matrix diagonals, suggesting a small cophenetic distance. That is, those cells are still assigned to nearby cell types in the dendrogram.
Using A > 2. Unlike the discussed benchmark datasets, the neuronal types are not uniformly distributed. Accordingly, we also investigated the accuracy improvement for categorical assignment when more than two arms are used. Fig. 3c illustrates the accuracy improvement with respect to a single autoencoder model, i.e. JointVAE, in agreement with our theoretical findings.
Identifying genes regulating cell activity. To examine the role of the continuous latent variable, we applied a similar traversal analysis to that used for the benchmark datasets. For a given cell sample and its discrete type, we changed each dimension of the continuous variable using the conditional distribution, and inspected gene expression changes caused by continuous variable alterations. Fig. 4 shows the results of the continuous traversal study for JointVAE and cpl-mixVAE, for two excitatory neurons belonging to the “L5 NP” (cell type (I)) and “L6 CT” (cell type (II)) sub-classes in ALM and MOp regions. Note that here, JointVAE is equivalent to a 1-arm VAE, with the exception of the type dependence of the state variable. Since CascadeVAE did not learn meaningful clustering of cells, even at the subclass level, we did not consider it for the continuous factor analysis. In each sub-figure, the latent traversal is color-mapped to normalized reconstructed expression values, where the y-axis corresponds to one dimension of the continuous variable, and the x-axis corresponds to three gene subsets, namely (i) MGs for the two excitatory types, (ii) immediate early genes (IEGs), and (iii) housekeeping gene (HKG) subgroups (Hrvatin et al., 2018; Tarasenko et al., 2017). For cpl-mixVAE (Fig. 4b), the normalized expression of the reported MGs as indicators for excitatory cell types (discrete factors) is unaffected by changes of identified continuous variables. In contrast, for JointVAE (Fig. 4a), we observed that the normalized expression of some MGs (5 out of 10) are changed due to the continuous factor traversal. Additionally, we found that the expression changes inferred by cpl-mixVAE for IEGs and HKGs are essentially monotonically linked to the continuous variable, confirming that the expression of IEGs and HKGs depends strongly on the cell activity variations under different metabolic and environmental conditions. Conversely, JointVAE
fails to reveal such activity-regulated monotonicity for IEGs and HKGs. Furthermore, our results for cpl-mixVAE reveal that the expression of activity-regulated genes depends on the cell type, i.e. IEGs and HKGs respond differently to activation depending on their cell types (compare rows I and II in Fig. 4b). However, in Fig. 4a, since the baseline JointVAE does not take into account the dependency of discrete and continuous factors, it fails to reveal the dependence of activity-regulated expression to the cell type, and therefore produces identical expressions for both types (I) and (II). These findings are consistent over multiple randomly initialized runs (Supplementary Section K.1). See Supplementary Section K.2 for more results on other cell types and gene subsets.
Summary. The cpl-mixVAE model successfully identified the majority of known excitatory and inhibitory neurons in multiple cortical regions. Our findings suggest that cpl-mixVAE, by acknowledging the dependencies of continuous and categorical factors, captures relevant and interpretable continuous variability that can provide insight when deciphering the molecular mechanisms shaping the landscape of biological states, e.g. due to metabolism or disease.

heading: 4.3 ABLATION STUDIES
text: To elucidate the success of the A-arm VAE framework in mixture modeling, we investigate the categorical assignment performance under different training settings. Since CascadeVAE does not learn the categorical factors by variational inference, here we mainly study JointVAE (as a 1-arm VAE) and cpl-mixVAE (as a 2-arm VAE). In Section 4.1, we show that data augmentation by itself does not enhance the categorical assignment (JointVAE†). To understand whether architectural differences put JointVAE at a disadvantage, we trained JointVAE‡ (Table S1), which uses the same architecture as the one used in cpl-mixVAE. JointVAE‡ uses the same learning procedure as JointVAE, but its convolutional layers are replaced by fully-connected layers (see Supplementary Section J and L for details). The result for JointVAE‡ suggests that the superiority of cpl-mixVAE is not due to the network architecture either. We also examined the performance changes of the proposed 2-arm cpl-mixVAE under three different settings: (i) cpl-mixVAE∗, where coupled networks are not independent and network parameters are shared; (ii) cpl-mixVAEa, where only affine transformations are used for data augmentation; and (iii) cpl-mixVAE(s 6 | c), where the state variable is independent of the discrete variable (Table S1). Our results show that the proposed cpl-mixVAE obtained the best categorical assignments across all training settings. We also examined the accuracy of categorical assignments for the cpl-mixVAE model, under different dimensions of discrete latent variable, for both MNIST and scRNA-seq datasets (see Supplementary Section J). We experimentally observe that while JointVAE suffers from sensitivity to empirical choices of |c|, cpl-mixVAE is more robust in encoding the discrete variability, without suffering from mode collapse (Fig. S9 and Fig. S10).

heading: 5 CONCLUSION
text: We have proposed cpl-mixVAE as a multi-arm framework to apply the power of collective decision making in unsupervised joint representation learning of discrete and continuous factors, scalable to the high-dimensional discrete space. This framework utilizes multiple pairwise-coupled autoencoding arms with a shared categorical variable, while independently learning the continuous variables. Our experimental results for all datasets support the theoretical findings, and show that cpl-mixVAE outperforms comparable models. Importantly, for challenging scRNA-seq dataset, we showed that the proposed framework identifies biologically interpretable cell types and differentiate between type-dependent and activity-regulated genes.

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

The paper considers the problem of generative modeling for mixed discrete-continuous data. To this end, it introduces a novel variant of coupled variational autoencoders. Specifically, they reformulate the coupled variational autoencoder paradigm in such a way that, instead of each "arm" of the autoencoder modeling a different modality, the arms model different chunks of the same modality (which comprises a high-dimensional discrete part, in addition to a continuous part).   The key novelty of the paper is the idea of postulating a finite mixture model as the likelihood (decoder) pertaining to each arm of the autoencoder, and stipulating that all arms essentially infer similar categorical posteriors of mixture component assignment for the different chunks pertaining to the same observation. This consensus constraint is enforced on the grounds of the Aitchison geometry in the probability simplex, which avoids the mode collapse problem.  

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.