Mixture Representation Learning with Coupled Autoencoders
Authors: Yeganeh Marghi
Rohan Gala
Uygar Sümbül
Venue: ICLR 2022 Submitted (2022)
Keywords: Mixture representation
high-dimensional categorical variable
unsupervised learning
constrained variational framework
neuronal diversity
single-cell RNA sequencing
cell types

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.

Under review as a conference paper at ICLR 2022
MIXTURE REPRESENTATION LEARNING WITH COU- PLED AUTOENCODERS
Anonymous authors Paper under double-bling review
ABSTRACT
Latent representations help unravel complex phenomena. While continuous la- tent variables can be efficiently inferred, fitting mixed discrete-continuuus 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 outper- forms state-of-the-ar in high-dimensional discrete settings, cpl-mixVA VAE introduces 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 hundred cortical neuron types.
INTRODUCTION
Fitting mixed discrete-continuond 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) jet 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 single-cell omic studies. The datasets produced by these studies document a previously neuroscience, this poses the identification of cell types
high-throughhut and high-dimension
unappreciated diversity of gene expression. and cell states as 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 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 tailored for problems with relatively few and equally-abundan discrete components, and accurate inference remains out of reach for these applications.
learning methods
Deep generative models have previously been applied 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 al. 2018). Deep Gaussian mixture models (Dilokthanakul 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 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 9s and select the member closest to the true model p. Popular choices in VAE implementations include (1) using KL divergence compute discrepancy between qs 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

---

Under review as a conference paper at ICLR 2022
underestimating the posterior variance (Minka et al. 2005 Blei et al. 2017). Solutions to resolve this
issue are mainly applicable in now-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 autoen- coding arms to jointly infer interpretable finite discrete (categorical) and continuous factors in the presence of high-dimenn nigh-dimensional discrete space. coupled-auutoencoders have been previously studied in the context of multi-modal recordings, where each arm learns only 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-couples autoencoder framework for 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 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 clustering in mixture models (Dilokthanakul et al. 2016 Jiang et al.] 2017 Tian et al. 2017 Guo et al. 2016 Locatello et al. (018b). The idea of improving the clustering performance through seeking consensus and co-training and ensembling
across multiple observations has been explored in both unsupervised (Monti et al. 2003 Kumar & Daumé (01 1 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 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) demon- strated 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 disen- tangled 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 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 the discrete space, which will be addressed in our A A-arm VAE framework.
SINGLE MIXTURE VAE FRAMEWORK
For an observation x E TR. a VAE learns a generative model pe (x|z) and a variational distribution 96 (z)x), where E R is latent variable with 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 al. 2018a). However, the interplay between continuous and discrete variabilities present in many real 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 e 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 e are independent random variables, the evidence lower bound (ELBO) (Blei et al. (017)) for a single mixture VAE with the distributions parameterized by 0 and Φ is given by.
L( 0) = E as c|x) [log po (x|s, c)] - Del (s)xx)))l)(ss) - DxL (c)(c))))P(()). (1) Maximizing ELBO in Eq. I |imposes characteristics on (ss)x) and q(c(x) that can result in underesti- mation of posterior probabilities such as the mode collapse problem, where the network ignores
2

---

Under review as a conference paper at ICLR 2022
Figure I: (a) Multi-arm autoencoder framework pro- posed as the cpl-mi mpl-mix VAE model. Individual arms receive non-identical noisy copies of given samples x, i.e. (xe, X5, ...). where they all belong to the category, learn mixture representations, i.e. {g(Ca, Sa), q(C,, 8s).... ..... }. VAE arms cooperate to learn the categorical assignment, p(c). (b) Graphical model of each autoencoder to learn type dependence of state variable.
g
X
Sa Ca
Xa
x
(tca(xa)
g
g(salcg.
{
Xb
Cb
Sb
Xx
Ca
Sp
(b)
subset of latent variables (Minka et al. 2005; Blei et al. 2017). Recently, VAE-basee solutions were proposed by imposing uniform structure on p(c): akin to 3-VAE (Higgins et al. 2017 Burgess et al. 2018), JAi (Dupont 2018) modifies the ELBO by assigning a pair of controlled information capacities for each variational factor, i.e. C. E Rala and C. Rlol, The main drawback of JointVAR is that performance tied to heuristic tuning of [s] |c| capacities over training iterations so that vulnerable to mode collapse in high-dimensional settings. Another recent VAE-based mixture model solution, Cascade VAE (Jong & Song 2019), maximizes the ELBO through a cemi-gradient-baa algorithm by iterating 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 nigh-dimensional discrete variables in real-world applications.
separate optimizations for the continuous and categorical variables.
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 g(s) x) and (((]x) can be significantly violated (see Supplementary Sections H and I).
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 decision, multiple estimates can increase the expected probability of 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. Db). In this framework, individual arms receive a collection of non-identical copies, (x,, Xb.- the given sample, x, belonging the same category. While each arm has its own mixture representation with potentially non-identical parameters, all arms cooperate to learn g(c, xx), where C, Cs =... .... via a cost function at the time of training. Accordingly, crowd of VABs with A arms can be formulated as collection of constrained variational objectives as follows.
Lxyles (Ø, 01) + Caales ( 0A)
(2)
s.t. C1 = CA
where La.le. (D., 0z) (a) the variational loss for arm a,
Exyle, (0., 0a) = E Eg, [log p(xa c.)] - Ey(e, (x.) [DKl (g(salca, x. xx)llp(s./c.))
(Fg(sslle....) [D (4(c,)x,) (p))(()))) (3)
In Eq. 3, the variational loss for each arm is defined according to the graphical model in Fig Fig which
built upon the traditional ELBO in Eq. by conditioning the continuous state on the categorical variable (derivation in Supplementary Section B). Therefore, learning interpretable decomposition of the data relies on accurate assignment (inference) of the categorical latent factor. Propositions and 2 below show that the shared categorical assignment inferred from q(c "... (AA), under the
3

---

Under review as a conference paper at ICLR 2022
c1 C1 CA constraint of the multi-arm framework improves the aceuracy of the categorical assignment on expectation.
Proposition 1. Consider the problem of mixture representation learning in multi-arm VAE frame- work. For independent samples from category m, i.e. x; ~ pppp(mmm m))
Bo(x/m) [log g(c = -m((x)):A))
> [g(x(m) [lg q(c m m (X:h:B)
8.l. C1 CA
s.l. C1 cp
(4)
if (m/xx) < 1 and A > B 2 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 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 By(x/m) [log (c(((x))aa)) 8.l. c =.....c.
(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 (Bouchacouri iet al. 2017). Instead, it relies on representations that are invariant under non-identical copies of observations. Moreover, unlike (Bouchacourt et al. 2017; (hu al. 2019 Locatello 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 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 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), develop 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 (Fq. 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.
3.2 CPL-MIX VAE: PAIRWISE COUPLING IN A-ARM VAE
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. by applying the pairwise coupling paradigm as follows (details of derivation in Supplementary Section C):
A
Σ (A (A-1) Egfss.... ix.) [logp(xa (xxpp((/sa.. ca)l - Eg(e. [DKr. (g(s. Xa (llsp ca)))
max
Eg le. x.) E (. [DK (g(ca)x.)q(cc/x)))
a<b
s.t. C. cs Va, E [I, <
(6)
We relax the optimization in Eq. inno an unconstrained problem by marginalizing the joint distribu- tion over mismatch measure between categorical variables (see Supplementary Section D):
A
max
(A 1)
E Egisy a. (x)) [oog Is., cu)l - Eygce., (x.)) [DKL (g(sa /c. xxllpc(s..)))lll...)))))
1=1
Σ H(calXa) + H(c)(x))) XB (dP(ca, a<b
(7)
4

---

Under review as a conference paper at ICLR 2022
In Eq. 7 in addition to entropy- confidence penalties known as mode collapse regulariz- ers (Pereyra et al. )017), the distance measure d(ca, Co) encourages a consensus on the categorical assignment controlled by À 2 O, the coupling hyperparameter.
We refer to the model in Eq. as cpl-mixVAE (Fig. TA). In cpl-mix VAE arms try to achieve identical categorical assignments while independently learning their own style variables. In experi- ments, we set À = universally. While the bottleneck architecture already encourages interpretable continuous variables, this formulation can be easily extended include an additional hyperparameter to promote disentanglement of continuous variables as in 3-VAE (Higgins et al., 2017). Additional analyses to assess the sensitivity of the cpl-mix VAE's performance its coupling factor can be found in the Supplementary Section G.
may be instructive to cast Eq. 7 in an equivalent constrained optimization form.
Remark 1. The A-arm VAF framework is a collection of constrained variational models as follows:
A
max
c [(xal(x ca)l -Egee (x.) [D (q(sa) Xa) Xa) ca))l ((ce
S.I. Byg(calx.) (c.: cb)] <E
where denotes the strength of the consensus constraint. Here, indicates the assigned category by any one of the arms, b E Al, imposing structure on the discrete variable to approximate its prior distribution.
Distance between categorical variables. d(c,: Ct)) 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-ddimensional simplex, where K = lcl- 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, L.e. dgx (Ca:Cb) is defined dgx (Ca, Ct) = [[,r((..) clr(ca) 2. Vc,, C₁2 E S* where cdr(-) denotes the isometric centered-log the simplex. This categorical distance satisfies the conditions of mathematical metric according to Aitchison geometry.
transformation in
3.3 SEEKING CONSENSUS IN THE SIMPLEX
An instance of the mode collapse problem (Lucas et al. 2019) manifests itself in the minimization of dgx (Ca; c5) (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; C. = C = Co- In this scenario, the continuous variable is com- pelled to act as 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 dgxx (Ca, we minimize a perturbed distance
df (b) Σx (om,) (om"loog, logCoz Cos ohe log )b2)² which corresponds to the distance between addi- tively perturbed c, and c, vectors in Aitchison geometry. Here, o2 and of indicate the bini-batch variances of the k-th category, for arms and b. We next show that the perturbed distance do (-) is bounded by dg (.) and non-negative values PuzPi:
Proposition 3. Suppose cu, C E SK, where SK is a simplex of K > Parts. Ifd (Cas Cb) denotes the distance Aitchison geometry and (Ca: ct) = =E (o=_) log Car The log Cba) denotes perturbed distance, then
&gx (Ca, Cb. pt < de (Ca: Co) S dg (Ca, Co. + Pu
A
where Pu. P1 2 0, Pu K (12 + +r²) 2A,Te. P
Kr2 Te max llog Cbu
K
k
To. max Toi = max ] A Σ yro and gk (0 1) log Cas. (0 -1
&
1) log Cby- (Proof in Supplementary Section F)
5

---

Under review as a conference paper at ICLR 2022
Thus, when Ca and G₁ are similar and their spread is not small, de (c,, cb) closely approximates dgæ (Cay Co). Otherwise, it diverges from dgæ (-) to avoid mode collapse.
EXPERIMENTS
We used four datasets: dSprites, MNIST, and two single-cell RNA sequencing (scRNA-seg) datasets; Smart-seq ALM-VISp (Tasie al., 2018 and 10X MOp Yao et al. 202T). Although dSprites and MNIST datasets do not require high- gigh-dimensional settings for mixture representation, to facilitate com- parisons of cpl-mix VAE with earlier methods, first we report the results for these benchmark datasets. We trained three unsupervised VAE-baser methods for mixture modeling: JointVAE (Dupont 2018). CascadeVAF (Jeong Song, 2019), and ours (cpl-mix VAF). 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, (1) for the discrete latent factor, we report the accu- racy (ACC) of categorical assignments and the Dx (x()] || 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-mix with the available cluster labels for each dataset. Additionally, to compare the training complexity of the multi-arm framework against earlier methods (Table All reported numbers for cpl-mix VAE models are average accuracies calculated across arms. In VAE-based models, to sample from q(C,, ).), we use the Gumbel-softmax distribution (Jang et al. 2016; Maddison al. 2014]. cpl-mix VAE, 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.
report the computational efficiency (number of iterations per second)
4.1 BENCHMARK DATASETS
dSprites. dSprites is procedurally generated from discrete (3 shapes) and continuous (6 style factors: scale, rotation, and position) latent factors. Based the uniform distribution of classes, used 2-arm cpl-mix VAE with c| = and |s] = 6 to learn interpretable representations. Results in Table ] 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 [ 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 fair comparison, we used the same disentanglement metric implemented for Cascade (Jeong & Song 2019).
MNIST. Similarly, due to uniform distribution digit labels in MNIST, we again used a 2-arm cpl-mix VAE model. Following the convention (Dupont 2018 Jeong & Song 2019 Bouchacourt
Table 1: Training results for all datasets. cpl-mixVAE uses arms. |c| and |s| denote the cardinality of latent discrete and continuous spaces. Iter, ACC and DS denote number iterations, the accuracy of the categorical as- signment, and the disentanglement score, respectively. Computation denotes the training speed (iteration/secon)) on GeForce RTX 2080 Ti GPU. The computation of epl-mix VAE includes the entire execution time for training pair of coupled networks, plus data augmentation.
Dataset
Iter
|c]
|
Method
ACC (%)
Computation (iteration/sec)
(mean + s.d.)
JointVAB
44.79 03.9
52.6
74.5
dSprites
300K
6
CascadeVAE
78.84 15.7
15.4
90.5 5.3
DS (
cpl-mixVAE
96.30 09.2
20.6
89.9 + 4.1
10
2
InfoGAN
77.87 21.7
12.2
2.15 - 4.2
JointVAE
68.99 = 11.8
74.1
3.12 1 1.0
JointVAE
68.21 09.6
54.4
3.65 1 2.7
MNIST
120K
10
10
Cascade VAE
74.83 06.9
23.8
1.99 + 1.1
Cascade
72.98 = 13.4
17.8
2,44 + 1.4
cpl-mixVAE
84.56 06.5
17.5
1.48 + 1.4
Smart-seq
JointVAB
12.53 02.9
28.6
459 12
ALM-VISP
40K
115
2
CascadeVAE
02.69
03.4
54.8 0.5
cpl-mix VAE
38.78
10.1
43,9 3.5
JointVAE
02.32 00.1
14.8
542 1 17
1000
1OX MOp
220K
140
2
Cascade VAE
02.85 00.5
00.6
150 3.5
cpl-mix
23.64 00.9
08.2
135 5.2

---

Under review as a conference paper at ICLR 2022
Rotation
Stroke thickness
Width
Roundness
0
0
0
0
D
D
D
b
0
0
o
o
o
0 0
0
0
0
0
0
6 0 0
0
0
0
0
(
\
/ 2
/
/
7
2
2
2 2
2
1L7
2
2
2
2
2
2 3
1
3
2
2
2
2
5
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3 4
44
4 4 114440
414
A
A
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
5
S
5
S
5
5
5
5
5
5
5
5
5
5
5
5
5
4
66
6
6
L
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
5 8
8
8
8
8
8 1 7
8
8
8
8
8
8
8
8
8
6
8
8
8
8
9
9 9
9
9
2
9
9
9
Rotation
Scale
Position
Figure 2: Interpretable continuous latent traversals of 1-st arm of the cpl-mix VAE framework with two autoencoders, for MNIST (top) and dSprites (bottom). The discrete variable с constant for all reconstructions in the same row.
et al. 2017), each arm of cpl-mix VAE uses 10-dimensional categorical variable representing
digits (type), and a 10-dimensional continuous random variable representing the writing style (state). Table displays the accuracy of the categorical assignment and the discrepancy between g(c) and p(c) for InfoGAN, two I-arm VAE methods (JointVAE and CascadeVAE), and cpI-mixVAE with 2 arms. Additionally, isolate the impact of data augmentation in training, we trained JointVAF and Cascade VAFT where the models were trained with the same augmented copies of the original MNIST dataset as cpl-mix VAE. The results in Table T 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-mix VAE, where each row corresponds to a different dimension of the categorical variable, and the state variable monotonically changes across and Fig. 2 show that cpI-mixVAE achieved an interpretable mixture representation with the highest categorical assignment accuracy.
columns. Both results in Table and
Summary. cpl-mix VAE 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.
4.2 SINGLE-CELL RNA SEQUENCING DATA
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- -VSP (Tasic et al.. 2018) and (ii) 10X MOp Yao et al. (021). 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 sroplet-bused 10X Genomics Chromium platform. 10X-basee 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 most- and the least-abundant cell types include 17, 000 and 20 samples, respectively. Moreover, whether the observed diversity corresponds
discrete variability or a continuum 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) 2013) is a common approach define cell types, the identified genes rarely obey the idealized MG definition in practice. Here, we focus on neuronal cells and use 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, used 140- and 2-dimensional discrete and continuous latent variables. We compared the suggested cell types in (Tasic et al.) 2018 Yao et al., 202I) with the discrete representations that are inferred from VAE models. Table and Fig. 3(a-b) demonstrate the
7

---

Under review as a conference paper at ICLR 2022
JointvAE
CascadeVAE
col-mixVAE
0.00
10
6x.0)
6x.3)
Cell lipes
140
034
0.2
22
IcC
Categories
60
Categories
1.80
Categorie
-C30
(h.l)
(b:5)
120
5.06
Cell Types
110
0.00
Categgries
Caleggoiis
5.6
Categories
Figure Categorical assignments the scRNA-seq datasets. The top row shows confusion matrices of VAE (a.1), Cascade VAE (a.2), and cpl-mix (a.3) trained by |c] 115, |s| 2. for the Smart-seq ALM- VISp dataset. The dendrogram on the y-axis shows MG-based hierarchical classification with 115 cell types, suggested Tasic al. (2018). The bottom row shows confusion matrices of JointVAE (b.1), CascadeVAB. (b.2), and cpl-mix VAE (b.3) trained by c| 140, |s| = 2. for the 10X MOp dataset. Cell types on the y-axis are sorted based on hicrarchical classification suggested by Yao al. (202T). (c) Improvement of the categorical representation (ACC) of cpl-mix VAE by adding more arms to the multi-arm framework. A-arm's performance for A is compared with the baseline 1-arm, Join across 3 randomly initialized runs.
performance of 2-arm cpl-mix VAE model against Joint and Cascade VAE. In Fig. 31.1 and 3p.L. we observe that for both datasets, JointVAr succeeds in identifying (sub)classes of neurons, c.g. excitatory class, Sst subclass, but not neuronal types at leaf nodes. On the other hand, Cascade learns an almost uniform distribution over all types despite a sizeable difference between the relative abundances of neuronal types (Fig. 3.2 and 3p22. Our results in Fig. aa-b) clearly show that cpl- mix VAE outperforms JointVAE and Cascade VAE in identifying meaningful known cell types. The confusion matrices in Fig. 3.. and D..3 demonstrate that even the inaccurate categorical assignments of cpl-mix VAE are still close to the matrix diagonals, suggesting small cophenetic distance. That is, those cells are still assigned 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. 3e illustrates the accuracy improvement with respect to single autoencoder model, i.e. JointVAE, agreement with our theoretical findings.
Identifying genes regulating cell activity. To examine the role of the continuous latent variable, we applied similar traversal analysis to that used for the benchmark datasets. For 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. shows the results of the continuous traversal study for Joint VAE and cpl-mix VAE, for two excitatory neurons belonging to the "L,5 NP" (cell type (I)) and "L6 CT (cell type (II) sub-classes in ALM and MOp regions. Note that here, JAE is equivalent to 1-arm VAE, with the exception of the type dependence of the state variable. Since Cascade VAE did not learn meaningful clustering of cells, even 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 z-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 epl-mix pol-mixVee VAE (Fig. 4b), the normalized expression of the reported MGs as indicators for excitatory contrast, cell types (discrete factors) unaffected by changes identified continuous variables.
for VAE (Fig. 4),, 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-mix VAE for IEGs and HKGs 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, Join
8

---

Under review as a conference paper at ICLR 2022
MG
IEG
HKG
MG
IEG
HKG
MG
IEG
HKG
(
Joi Smart-scq VLS-VISP
apl-mixxVAF- -aaamixAA---mar--eyy AL.S-VIS
(c) mp)- 10x MOp
Figure 4: Continuous latent traversal analysis for excitatory cell types: (I) "I.5 NP" and (II) "I.6 CT", in different brain regions. For each type, the traversal is color-mapped to normalized reconstructed gene expression value (colorbar) as a function the state variable for gene subsets: marker genes (MG), immediate early genes (IEG), and housekeeping genes (HKG).
fails to reveal such activity-regulatte monotonicity for IEGs and HKGs. Furthermore, our results for cpl-mix VAE 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. (D). 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 the cell type, and therefore produces identical expressions for both types (I) and (II). These findings are consistent See Supplementary Section K.2 for more results on other cell types and gene subsets.
multiple randomly initialized runs (Supplementary Section K.1).
Summary. The cpl-mix VAE model successfully identified the majority of known excitatory and inhibitory neurons in multiple cortical regions. Our findings suggest that cpl-mix VAE, by acknowl- edging 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.
4.3 ABLATION STUDIES
To elucidate the success of the A-arm VAE framework in mixture modeling, we investigate the categorical assignment performance under different training settings. Since CascadeVAl does not learn the categorical factors by variational inference, here we mainly study JointVAE (as a I-arm VAE) and cpl-mix VAE (as 2-arm VAE). In Section 4. we show that data augmentation by itself does not enhance the categorical assignment (JointVAL trained JointVAE! (Table S1), which uses the same
To understand whether architectural
differences put Joint at disadvantage, architecture as the one used in cpl-mix VAF. JointVAFt uses the same learning procedure as JointVAF, but its convolutional layers and L for details). The result for JAit suggests that the superiority of cpl-mix VAE is not due the network architecture either. We also examined the performance changes of the proposed 2-arm cpl-mix mix VAE under three different settings: (i) cpl-mix VAE*, where coupled networks are not independent and network parameters are shared; (ii) cpl-mix
replaced by fully-connected layers (see Supplementary Section J
where only affine transformations are used for data augmentation: and (iii) cpl-mix VAE(s X /c), where the state variable is independent of the discrete variable (Table S1). Our results show that the proposed cpl-mix VAE obtained the best categorical assignments across all training settings. We also examined the accuracy of categorical assignments for the cpl-mix VAE 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 |cl. cpl-mix VAE more robust in encoding the discrete variability, without suffering from mode collapse (Fig. S9 and Fig. S10).
CONCLUSION
We have proposed epi-mix VAE 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 shared categorical variable, while independently learning the continuous variables. Our experimental results for all datasets support the theoretical findings, and show that cpl-mix VAE outperforms comparable models. Importantly, for challenging scRNA-seq dataset, showed that the proposed framework identifies biologically interpretable cell types and differentiate between type-dependent and activity-regulatee genes.

---

Under review as a conference paper at ICLR 2022
REFERENCES
John Aitchison. The statistical analysis of compositional data. Journal of the Royal Statistical Society:
Series (Methodological), 44(2):139-10) 1982.
Antreas Antoniou, Amos Storkey, and Harrison Edwards. Data augmentation generative adversarial networks. arXiv preprint arXiv:171..04444 2017.
Cornelia Bargmann, William Newsome, A Anderson, E Brown, K Deisseroth, J Donoghue, P MacLeish, E Marder, R Normann, J Sanes, et al. Brain 2025; scientific vision. Brain Research through Advancing Innovative Neurotechnologies (BRAIN) Working Group Report to the Advisory Committee to the Director, NIH, 2014.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Journal of the American statistical Association, 112(518):59-8877 2017.
Variational inference: A review for statisticians.
Avrim Blum and Tom Mitchell. Combining labeled and unlabeled data with co-training, In Proceed- ings of the eleventh annual conference Computational learning theory, pp. 92-100, 1998.
Diane Bouchacourt, Ryota Tomioka, and Sebastian Nowozin. Multi-level variational autoen- coder: Learning disentangled representations from grouped observations. arXiv preprint arXiv:1705.088411 2017.
Christopher Burgess, Irina Higgins, Arka Pal, Loic Manthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 8-vae. arXiv preprint arXiv: arXiv:/804.00:: 1804.03599 2018.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. In Advances in neural information processing systems, pp. 2172-2180, 2016.
Jacob Deasy, Nikola Simidjievski, and Pietro Liò. Constraining variational inference with geometric jensen-shannon divergence. arXiv preprint arXiv:2006.10599 2020.
Nat Dilokthanakul, Pedro AM Mediano, Marta Gareel,, Matthew CH Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture varia- tional autoencoders. arXiv preprint arXiv: 1.02648, 2016.
Emilien Dupont. Learning disentangled joint continuous and discrete representations. In Advances in Neural Information Processing Systems, pp. 710-720, 2018
Juan José Egozcue, Vera Hewlowsky-Glahn Glòria Mateu Sateu--Figuers, and Carles Barcelo- Vidal. Isometric logratio transformations for compositional data analysis. Mathematical Geology, 35(3): 279-300, 2003.
Fangxiang Feng, Xiaojie Wang. and Ruifan Li. Cross-modal retrieval with correspondence au- toencoder. In Proceedings of the 22nd ACM international conference on Multimedia, pp. 7-16, 2014.
Rohan Gala, Nathan Gouwens, Zizhen Yao, Agata Budzillo, Osnat Penn, Bosiljka Tasic, Gabe Murphy, Hongkui Zeng, and Uygar Sümbül. A coupled autoencoder approach for multi-modal analysis cell types. In Advances in Neural Information Processing Systems, pp. 9263-9272, 2019.
Fangjian Guo, Xiangyu Wang, Kai Fan, Tamara Broderick. and David Dunson. Boosting variational inference. arXiv preprint arXiv: /6/ F.05559, 2016.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. Iclr. 2(5):6, 2017.
Haruo Hosoya. Group-based learning of disentangled representations with generalizability for novel contents. In IJCAI, pp. 2506-2513, 2019.
10

---

Under review as a conference paper at ICLR 2022
Sinisa Hrvatin, Daniel R Hochbaum, M Aurel Nagy, Marcelo Cicconet, Keiramarie Robertson, Lucas Allon M Klein, et al. Single-
Cheadle, Rapolas Zilionis, Alex Ratner, Rebeca cell analysis of experiencce-dependent transcriptomic states in the mouse visual cortex. Nature neuroscience, 21(1):120-129 2018.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.011 2016.
Yeonwoo Jeong and Hyun Oh Song. Learning discrete and continuous factors of data via alternating disentanglement. arXiv preprint arXiv:/905, 09432, 2019.
Zhuxi Jiang, Yin Zheng, Huachun Tan, Bangsheng Tang, and Hanning Zhou Variational deep embedding: an unsupervised and generative approach to clustering. In Proceedings of the 26th International Joint Conference Artificial Intelligence, pp. 1965-1972, 2017.
Matthew J Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta. Composing graphical models with neural networks for structured representations and fast inference. In Advances neural information processing systems, pp. 2946-2954, 2016.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. arXiv preprint 1rXhv:1802.05983 2018.
Diederik P Kingma and Max Welling. 14,661114200 2013.
Auto-encoding variational bayes.
arXiv preprint
Diederik P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improving variational inference with inverse autoregressive flow. arXiv preprint a6Xiv:1606. 0.4444 2016.
Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):33-97, 1955.
Abhishek Kumar and Hal Daumé. co-training approach for multi-view spectral clustering. In Proceedings of the 28th international conference on machine learning (ICML-I1), pp. 393-400, 2011.
Anders Boesen Lindbo Larsen, Søren Kaae Sønderby, Hugo Larochelle, and Ole Winther. Autoen- coding beyond pixels using a learned similarity metric. In International conference machine learning, pp. 1558-1566. PMLR, 2016.
Mihee Lee and Vladimir Pavlovic.
Private-shared disentangled multimodal vae for learning of hybrid latent representations. arXiv preprint arXiv:2012.13024 2020.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentan- gled representations. arXiv preprint arXiv: 181 /2359, 2018a
Francesco Locatello, Gideon Dresdner, Rajiv Khanna, Isabel Valera, and Gunnar Rätsch. Boosting black box variational inference. In Advances in Neural Information Processing Systems, pp. 3401-3411, 2018b.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentangiement without compromises. arXiv preprint arXiv:2002.02886, 2020.
Romain Lopez, Jeffrey Regier, Michael B Cole, Michael I Jordan, and Nir Yosef. Deep generative modeling for single-cell transcriptomics. Nature methods, 15512) 2018.
James Lucas, George Tucker, Roger Grosse, and Mohammad Norouzi. Don't blame the elbo!
linear vae perspective posterior collapse. In Advances in Neural Information Processing Systems, pp. 9403-9413, 2019.
Chris Maddison, Daniel Tarlow. and Tom Minka. A sampling. Advances in Neural Information Processing Systems, pp. 3086-3094, 2014.
11

---

Under review as a conference paper at ICLR 2022
Tom Minka al. Divergence measures and message passing. Technical report, Citeseer, 2005.
Stefano Monti, Pablo Tamayo, Jill Mesirov, and Todd Golub. Consensus clustering: resampling- based method for class discovery and visualization of gene expression microarray data. Machine learning, (2(1-))91-118 2003.
Jozsef Nemeth. Adversarial disentanglement with grouped observations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 10243-10250, 2020.
Gabriel Pereyra, George Tucker, Jan Chorowski, Lukasz Kaiser, and Geoffrey Hinton. Regularizing neural networks by penalizing confident output distributions. arXiv preprint arXiv: 06548, 2017.
Matias Quiroz, David J Nott, and Robert Kohn. Gaussian variational approximation for high- dimensional state space models. arXiv preprint arXiv:1801.07873, 2018.
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In International Conference on Machine Learning, pp. 324-333. PMLR, 2016.
Federico Scala, Dmitry Kobak, Matteo Bernabucei, Yves Bernaerts, Cathryn René Cadwell, Je- sus Ramon Castro, Leonard Hartmanis, Xiaolong Jiang, Sophie Laturnus, Elanine Miranda, et al. Phenotypic variation transcriptomic cell types in mouse motor cortex. Nature, pp. 1-7, 2020.
H Sebastian Seung and Uygar Sümbül. Neuronal cell types and connectivity: lessons from the retina. Neuron, 83(6): 2014.
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disentanglement with guarantees. arXiv preprint arXiv:/9/0009772,1 2019.
James Surowiccki. The wisdom of crowds. Anchor, 2005.
Tatyana N Tarasenko, Susan E Pacheco, Mary Kay Koenig, Julio Gomez-Rodriguez Senta M Kapnick, Francisca Diaz, Patricia M Zerfas, Emanuele Barca, Jessica Sudderth, Ralph DeBerardinis, et al. Cytochrome oxidase activity is metabolic checkpoint that regulates cell fate decisions during t cell activation and differentiation. Cell metabolism. (214)-22541123 2017.
Bosiljka Tasic, Zizhen Yao, Lucas T Graybuck, Kimberly A Smith, Thuc Nghi Nguyen, Darren Bertagnolli, Jeff Goldy, Emma Garren, Michael Economo, Sarada Viswanathan, et Shared and distinct transcriptomic cell types across neocortical areas. Nature, 363(7729))72-778, 2018.
Kai Tian, Shuigeng Zhou, and Jihong Guan. Deepcluster: A general clustering framework based on deep learning. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 809- 825. Springer, 2017.
Cole Trapnell. Defining cell types and states with single-cell genomics. Genome research, 25(10): 1491-1498, 2015.
Michael Tschannen, Olivier Bachem, and Mario Lucic. Recent advances in autoencoder-based representation learning. arXiv preprint 18Xi:: 0.00669 2018
Zizhen Yao, Cindy TJ van Velthoven, Thue Nghi Nguyen, Jeff Goldy, Adriana E Sedeno-Cortes. Fahimeh Baftizadeh, Darren Bertagnolli, Tamara Casper, Megan Chiang, Kirsten Crichton, et al. A taxonomy of transcriptomic cell types across the isocortex and hippocampal formation. Cell, 2021.
Hongkui Zeng and Joshua R Sanes. Neuronal cell-type classification: challenges, opportunities and the path forward. Nature Reviews Neuroscience, 18(9):530, 2017.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE transactions puttern analysis and machine intelligence, (1(8):2008-2222 2018.
12


Reviews:
[{"review_id": "wd-oRJio0sG", "reviewer": "Reviewer_2RLY", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper compares the proposed method with a few similar methods, using MNIST but also using a dataset for single-cell RNA sequencing data (a domain with which I am not familiar). I am also not familiar with how much efforts are currently being put into this approach of factorizing Z=(S,C) by other members in the field.\n\nMy main question for the authors is to ask about the results in table 1 when it comes to MNIST: do the std measurements indicate a spread over the multiple experimental runs, and why is the accuracy so low on MNIST? MNIST is the kind of dataset where a linear classifier can get 90% accuracy, yet all these methods presented fare much worse.\n\nMinor issue: I think it would be appropriate to mention Kingma & Welling in the background section when VAE are first mentioned in the paper instead of in the later Section 2 when they are described. The same paper can be referenced twice, but at least it should be referenced in the introduction in conjunction to the other papers on VAEs that the authors wish to mention.", "summary_of_the_paper": "The authors consider the problem of factorization of the hidden space of a VAE into two separate components Z=(S,C) where S is continuous and C is discrete. They make a particular assumption about the factorization of the encoder function q(S, C|X)=q(S|X)q(C|X) and they also take a mixture of “A” experts that are made to agree with their discrete assignment C.", "summary_of_the_review": "Hard to say how novel the contribue is. I want to see the author respond to my question about table 1.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "kV7IQUwJIML", "reviewer": "Reviewer_giw4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: The idea is quiet novel and smart. The technical formulation is appropriate, and the constraint enforcing mixture component assignment consensus is sound. The experimental results are convincing, as they include a number of benchmarks, including some challenging ones, comparison to some SOTA methods, and an ablation study that offers good insights into how and why the method works.\n\nWeaknesses: I did not find any commendable weakness in the paper.", "summary_of_the_paper": "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). \n\nThe 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. \n", "summary_of_the_review": "The paper is novel enough, methodologically sound, and empirically well-validated. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eut8MyMYgU", "reviewer": "Reviewer_mBXz", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper considers the following factorisation of the the generative model $p(\\mathbf x, \\mathbf c, \\mathbf s) =  p(x|s, c)p(s|c)p(c)$ in which the continuous components are a function\nof the categorical, where the categorical variable captures different (discrete) classes within the data and the continuous component explains variations within them.\n\nThe paper is mostly easy to read, though there are some points that require clarifications. It also comes with a quite extensive, though not complete, set of experiments and analysis. \n\nMy main issue is that the baselines against which the paper compares do not exhibit the generative structure that is given above; both of them, JointVAE and CascadeVAE, assume that\nthe categorical and discrete components are independent. This makes it difficult to assess whether the merits of the paper come from the factorization of the generative model that\nassumes the dependence of the continuous from the categorical component or whether it comes from the concensus mechanism. Such factorizations have been quite extensively discussed\nelsewhere, for example in Lavda et al, Data-dependent conditional priors for unsupervised learning of multimodal data, Entropy, 2020, where the authors discuss how such a factorisation \nactually protects from mode collapse, as well as a number of other issues relevant to the discussion in the current paper, such as how the different categories are used through entropy\nregularisers that naturally appear within the initial objective. Since the base level factorisation over which the present paper builds the consencus model is the same, it would have\nbeen very easy to include it as a considerably more relevant baseline in order to demonstrate whether the benefits come from the concensus approach as well as include a discussion of\ndifferences from Lavda et al. \n\nOne more point that I have missed concerns the claim of the paper that the multitude of arms allows to deal with cases where the class distribution is not uniform, is this something that is demonstrated in section 3.2? If yes, I did not understand what is it that makes the model robust to non-uniform $p(\\mathbf c)$. Related to that are the suggestions to use 2 arms when the distribution of $p(\\mathbf c)$ is uniform. \n\n\nDetails:\n\nIn eq~3 which gives the variational loss for arm $\\alpha$ the term $\\mathbb E_{q(s_a|c_a,x_a}[D_{KL}(q(c_a|x_a)||p(c_a)]$ should probably be simplified to just the KL \nterm $D_{KL}(q(c_a|x_a)||p(c_a))$, since the distributions in the latter do not depend on the $s_a$ variable with respect to which the expectation is taken. See for example\nthe respective derivation in Lavda et al, 2020. \n\nIn figure b) the graphical model describes the generative model $p(x|s, c)p(s|c)p(c)$ however what is given just bellow seems to be the inference model even though $p$ is used instead of $q$.\n\nI have a bit of a problem conceptualising $q(c|x_1,...,x_A)$, how is this posterior instantiated concretetly? or is it only defined implicitly through equation 2? After going through the proof\nof proposition 1. in the appendix this is clear.\n\nProposition 1. says that as the number of arms/experts increases the likelihod of the true category will also increase.  \n\nProposition 1. :\n\n* what is the difference of $x_i \\sim p(x|m)$ and the $q(x|m)$ that appears in the expectation? $p(x|m)$ never appears in the derivation in the appendix. Is $q(x|m)$ meant to denote that\nwe randomly draw an $x$ and $p(x|m)$ the noisy versions of it, in which case it would probably be more appropriate to have $p(x_i|x)$.  By the way this two level sampling is missing from \nthe derivations, though I am not sure it is needed, but I find confusing the fact that there are two types of rvs here the original $x$ and its noisy versions $x_i$, even if I get the point.\nFor example in eq. 5 we marginalise over $x$ a quantity that does not contain $x$ but a random variable of $x$ the $x_i$.\n\n* appendix, eq 5: the denominator has dissapeared and repappears again in 6. \n\n\nSince the method works by operating on perturbed versions of the training instance one needs to define a perturbation model. The paper proposes a generative model based on GAN-VAE where\nthe main desiredata is to perturbe the instances while not alterning their latent categorical code, the non-observed class.\n\nWhen computing the distance over the categorical values $c_a$ it seems that the paper considers that $c_a \\in \\mathcal S^K$, i.e. they belong to the probability simplex. However $c_a$ is not a probability \nvector but a sample from such a categorical distribution, which should normally be a one hot-encoded vector. Which is the case? i.e. are the distances computed over the samples or over the \nrespective probabilities? How relevant is this?\n\n\nExperiments: \n* the paper said that when the distribution of classes is uniform we should use a 2-arm (dSprites and MNIST), how is this motivated? If I understood correctly the theoretical results, these show that when we increase the number of arms we will have a larger log posterior for the true categorical variable, but I am not sure I saw a relation to the uniformicity of the categories. \n\n* In the real world experiments, figure 3, how are the categories (i.e. the columns) ordered? I presume the ordering of the three different algorithms is not really comparable and what really matters\nis how these align with the real cell types. I am just curious is the orderning based on the largest values that one gets in the diagonal of the confusion matrix? \n\n* In the same figure, it seems that the discrete compoments that the proposed method uncovers are in agreement with the hierarchical clustering of the cells. So what we have here are two \nunsupervised methods the results of which are in agreement. I guess that the hierachical clustering results have been validated by the domain experts in (Tasic et al., 2018; Yao et al., 2021), \nI am just wondering how to approach such an evaluation where there is no real ground truth. Of course one could say it is a good thing that two rather different methods agree on how they \ncluster cells. \n\n* The evaluation in figure 4 where a given continuous latent variable varies with the discrete component makes intuitive sense. The figure shows that the profiles produced by cpl-mixVAE varies \nsmoothly within a given category and look rather different between the two categories. This is not the case for JointVAE, where the categorical component seems to bring no structure, the two different\ncategories have rather similar profiles. The later is something to be expected, as also noted by the authors, since JointVAE assumes independence of the categorical and discrete component an assumption\nthat makes it not the most appropriate baseline. On the same time, other than saying that intuitively the results of clp-mixVAE look good it is hard to make an additional comment, because this requires\nquite some expertise on the biology side.  \n", "summary_of_the_paper": "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 \nassign 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\nproblem. 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.\n\n", "summary_of_the_review": "My main issue is that as the paper is it is difficult to assess whether the strength of the method comes from the concensus mechanism or whether it comes from the factorisation. In addition in the discussion about the number of arms being a strength when the distribution of the categorical factor is not uniform I miss the justification. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "UHjP08P8BeS", "reviewer": "Reviewer_uWQ6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Post-rebuttal: \nThe authors' feedback clarify some of my concerns. but I'm still not so convinced by the experiments. The authors argue the model is not suitable for SVHN/cifar10 due to their complexity and this makes it unclear whether the model is applicable/generalizable for the real-world datasets that are either large scale or present extremely challenging backgrounds. Achieve SOTA ACC is definitely not necessary, but I would assume the proposed model has competitive results, however, the preliminary results shown in appendix seem quite weak. Therefore, I still tend to keep my original rating. \n\nFor the paper strength: \n\n-- The latent variable with mixed representation is an important model and worth exploring. \n\n-- Not familiar with the bio-literature, but the experiments and applications on RNA-data look interesting. \n\nFor weakness: \n\n-- The motivations and descriptions of the model need to be improved. Some places are unclear and confusing. For example. the paper emphasized that the work is fully unsupervised and proposes the type-preserving data augmentation. However, since each arm should receive samples that share the same underlying categorical factor, does this requirement impose some weak-supervision signal?  i.e., we have to ensure that each arm observes the samples from same category, but for fully unsupervised setting, we don't know that. Could authors further elaborate on this point? \n\n-- The experiments seems not sufficient. (a) Since more arms potentially means more encoder/decoder pairs (i.e., more parameters involved), its better to also include the model complexity comparison with the baselines. (b) The scalability of the model is not clear since the experiments are only done on synthetic/gray-scale datasets, what about on svhn or cifar10? For these slightly complicated benchmarks, how does the model perform? How to select the number of arms? Does the model capable dealing with large-scale datasets? ", "summary_of_the_paper": "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.  ", "summary_of_the_review": "The major concern goes to the experiment setting. The benchmarks are too simple to show the potential of the model and the baselines compared can be more up-to-date. I lean towards reject and may reconsider my rating based on authors' feedbacks. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CNQlK254kPj", "reviewer": "Reviewer_J8W6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Update: I have adjusted my score upwards.  The reviewers have address some of my concerns.  I still believe that the interplay between the  VAE-GAN and the clustering method is not explored well-enough, and that complex interplay can drastically impact the results of the approach.\n\nI have issues with the experiments, justifications, and theory in this manuscript.\n\nFirst, my major issue with this manuscript is that the proposed approach is not truly an unsupervised method, as is claimed.  It uses some supervised information in order to make its decisions.  To see this, consider the usage of “type-preserving data augmentation that generates independent and identically distributed copies of data while preserving its categorical identity.”  The generation process maintains the label; hence, you are telling the neural network that 2 samples from the same class belong in the same cluster.  Thus, because the data augmentation system is using full data distributions, despite not explicitly using the label, I would argue strongly that supervised information is leaking into your system.  The reason that this data augmentation does not help JointVAE$^\\dagger$ like it helps cpl-mixVAE is that in the JointVAE the system is not told that these should exactly match, so the supervised information does not leak into the JointVAE system.\n\nAs such, if you are going to use existing clusters, I need to understand the benefits of using this approach rather than a full supervised approach before any consideration of acceptance.  This is especially true as the quantitative metrics are all supervised evaluation criteria.  Since your method is getting some supervised information, it is unsurprising that it does better.  At the same time, it does much worse than basic supervised methods, including logistic regression on MNIST.\n\nIf you are going to push this as an unsupervised method, you should explain how it could be used in an unsupervised fashion.  For example, how well does it represent the data in an unsupervised fashion.  Can you use it for model selection?\n\nTheory issues:\n\nBoth proposition 1 and proposition 2 require that you know the true $m$ to estimate $m$.  This is not unsupervised theory, and it is misleading to call it that.  Second, the interpretation that the expected log posterior increases with the number of arms is misleading.  This method artificially inflates the posterior and isn’t adding information to the system.  If you are given a new sample where you don’t know the cluster, and does not improve the correct assignment rate in the Bayes optimal case.  This theoretical section needs an increased discussion about these issues.\n\nMinor complaints:\n\nIt is confusing to introduce c and s as independent in (1) and then immediately couple them in (3) without discussion.  Please refine this transition.\n\nExpectation on the last term in (3) is on the wrong distribution. \n\nIt is unclear how the distance between the distribution on c in (6) is related to the Aitchison distance used later.  Please provide the relationship.\n\nAdditional explanation of 3.3 is necessary.  It seems like the primary benefit is to push it away from trivial solutions, and the authors should expand on how they do that.\n\nIn the ablation study, I would like to understand the impact of the novel distance better.\n", "summary_of_the_paper": "This manuscript proposes cpl-mixVAE, which is a novel VAE formulation that attempts to improve clustering models by adapting ideas from consensus clustering.  The paper proposes a new model structure and a novel distance to help encourage improved model representations.", "summary_of_the_review": "I find the idea of including consensus clustering in a VAE an interesting and potentially useful idea to explore, but the theory explicitly suggests that this model only helps in the supervised case, not an unsupervised case.  Additionally, I would argue that supervised information leaks into cpl-mixVAE through the data augmentation, as is the explicit goal, and it is unfair to compare to purely unsupervised methods on only supervised metrics.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes cpl-mixVAE, a method for fitting discrete-continuous latent variable models based on mixture representations and a novel consensus clustering constraint. After extensive discussion, no one was willing to argue in favor of acceptance, and a majority of the reviewers felt another round of revision is needed. Ultimately, I concur that while the ideas are novel and potentially interesting, more effort is needed to convincingly demonstrate the efficacy of the method. Valid concerns were also raised regarding the claimed "unsupervised" nature of the proposed method, a claim which at the very least requires some additional context. At this point, these outstanding issues require an additional round of revision.
Decision: Reject