Mixture Representation Learning with Coupled Autoencoding Agents
Authors: Yeganeh Marghi
Rohan Gala
Uygar Sümbül
Venue:  (2021)
Keywords: Multi-agent network
representation learning
collective decision making
type-preserving data augmentation

Abstract: Jointly identifying a mixture of discrete and continuous factors of variability can help unravel complex phenomena. We study this problem by proposing an unsupervised framework called coupled mixture VAE (cpl-mixVAE), which utilizes multiple interacting autoencoding agents. The individual agents operate on augmented copies of training samples to learn mixture representations, while being encouraged to reach consensus on the categorical assignments. We provide theoretical justification to motivate the use of a multi-agent framework, and formulate it as a variational inference problem. We benchmark our approach on MNIST and dSprites, achieving state-of-the-art categorical assignments while preserving interpretability of the continuous factors. We then demonstrate the utility of this approach in jointly identifying cell types and type-specific, activity-regulated genes for a single-cell gene expression dataset profiling over 100 cortical neuron types.

Under review as a conference paper at ICLR 2021
MIXTURE REPRESENTATION LEARNING WITH COU- PLED AUTOENCODING AGENTS
Anonymous authors Paper under double-blind review
ABSTRACT
Jointly identifying a mixture of discrete and continuous factors of variability can help unravel complex phenomena. We study this problem by proposing unsu- pervised framework called coupled mixture VAE (cpl-miXVAE) which utilizes multiple interacting autoencoding agents. The individual agents operate on aug- mented copies of training samples to learn mixture representations, while being encouraged to reach consensus on the categorical assignments. We provide theo- retical justification to motivate the use of a multi-agent framework, and formulate it as a variational inference problem. We benchmark our approach on MNIST and dSprites, achieving state-of-the-art categorical assignments while preserving interpretability of the continuous factors. We then demonstrate the utility of this approach in jointly identifying cell types and type-specifie, for a single-cell gene expression dataset profiling over 100 cortical neuron types.
genes
INTRODUCTION
Complex phenomena can be attributed to a mixture of discrete and continuous factors of variability, Such complexity is crucial to understand in variety of different contexts, from learning models for image datasets to identifying factors underlying neuronal identity. A common approach to study these phenomena is clustering, which produce representations that jointly capture the dependence on discrete and continuous factors. Generative models can learn such representations, which has recently received attention from the deep learning community. Deep Gaussian mixture models are among the first deep generative models to jointly represent discrete and continuous factors, in which continuous representation is decomposed into discrete clusters (Johnson et al. 2016) Dilokthanakul et al. 2016, Jiang et al. 2017). However, such models have mainly focused on clustering without regard to interpretability.. Adversarial and variational methods have been proposed to learn mixture representations that can identify interpretable continuous factors. While adversarial learning, e.g. InfoGAN (Chen et al. 2016) is susceptible to stability issues (Kim & Mnih 2018 Dupont 2018 Jeong & Song. 2019h, variational approaches, e.g. Join and Cascade VAE have produced promising and more stable results (Dupont 2018/Jeon & Song 2019). However, such variational methods utilizing single autoencoding agent rely either heuristic data-dependen embedding capacity, or on solving separate optimization problem for the discrete variable. Thus, learning interpretable and stable mixture representations remains challenging.
We introduce a multi-agent variational framework to jointly infer discrete and continuous factors through collective decision making, while sidestepping heuristic approaches used by single-agent frameworks. Coupling of autoencoding agents has been previously studied in the context of multi- modal recordings, where each agent learns a continuous latent representation for one of the data modalities (Feng et al. 2014 Gala et al., 2019). Here, we propose sairwise-couples autoencoders to learn a mixture representation for single data modality in an unsupervised fashion. Each autoencoder agent receives an augmented copy of the given sample with the same class label. To achieve this, we design novel sype-preserving augmentation that generates noisy copies of the data using within- class variabilities, while preserving its class identity. Coupling across the agents is achieved by encouraging categorical variables to be invariant under the augmentation, which regularizes the agents to learn interpretable representations. We demonstrate that such coupled multi- agent architecture can increase inference accuracy and robustness by exploiting clthin- variabilities, without requiring prior distribution the relative abundances of categories.
Our contributions can be summarized as follows: (i) We first provide theoretical justification to motivate the advantage of collective decision making for more accurate categorical assignments,

---

Under review as a conference paper at ICLR 2021
utilizing noisy copies of the same sample. To obtain such samples, we propose unsupervised sype-preserving augmentation method. (ii) We formulate collective decision making as a variational inference problem with multiple agents. In this formulation, we introduce an approximation of Aitchison distance in the simplex to compare categorical assignments of the agents, which avoids mode collapse. (iii) We benchmark our method and display its superiority over comparable approaches using the MNIST and dSprites datasets. (iv) Finally, we apply the method to challenging single cell gene expression dataset for population of neurons. We demonstrate that our method can be used to discover discrete categories referred to as neuronal types and type-specific genes regulating the continuous within-typs -tppe variability (e.g., metabolic states, disease states).
RELATED WORK
As introduced above, recent studies joint learning of discrete and continuous factors in generalized mixture models focus on variational or adversarial approaches (Dupont 2018 Jeong & Song 2019
Chen al. 2016). There is also rich literature that focus on clustering in mixture models and do not attempt to characterize the continuous variability: Dilokthanakul et al. (2016) and [iang ct al. (2017) performed variational inference in mixture models using autoencoding architectures. Tian et al. (017)) applied the alternating direction method of multipliers (ADMM) to use classical clustering algorithms in conjunction with neural network. Guo et al.. (016)) and Locatello al.) (2018b) used gradient boosting approaches (Friedman (001) to iteratively fit mixture models in variational frameworks. Moreover, the idea of improving the clustering performance through seeking a consensus across similar agents has been explored in both unsupervised (Monti et al., 20003 Kumar & Daumé 2011) and semi-supervised contexts (Blum & Mitchell 1998). However, the proposed consensus clustering approach is attempting to learn interpretable continuous variability. Here, in contrast, beyond joint disentangling, we propose a framework in which the agents seek a consensus, at the time of learning the mixture representation.
While our method does not assume any prior/supervising information, the individual agents in our approach can be considered to provide form of weak supervision each other. [Bouchacourt et al. (2017) demonstrated a multi-level variational autoencoder as a weak supervised disentanglement approach for different factors of variability by both revealing that observations within groups share the same class label where the class label variable takes values from finite set of labels and are represented by Gaussian distribution. Recently, [ocatello al. (2020) improved this framework by assuming that observation pairs share at least one underlying factor, and demonstrated disentangling of continuous (but not discrete) factors on image sets.
PRELIMINARIES
For an observation E RR,, variational autoencoder (VAE) learns generative model pe (x|z) and variational distribution ds (z) where z E M for M < D a latent variable with a parameterized distribution p(z) (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, in many real-world applications the inherent mixture distribution of continuous and discrete variations is often overlooked. This problem can be addressed within the VAE framework in an unsupervised fashion by introducing categorical latent variable E SK denoting the class label defined in a K-si Kisimplex, alongside the continuous latent variable s E R, Here, 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.) (017) for single autoencoding agent with the distributions parameterized by O and &i is given by,
L(Φ, 0) = ( [logro(x(s, c)] - DKL
DEL (1)(c)))ll((c))
(1)
Maximizing ELBO in Eq. [ to jointly learn ((s)x) and q(c)x) is challenging due to the mode collapse problem, where the network ignores subset of latent variables. Akin to 3-VAE (Higgins et al., 2017 Burgess et al., 2018), JointVAE assigns controlled information capacities to both continuous and categorical factors to prevent mode collapse (Dupont) 2018). A drawback of this method is that the capacities are dataset-dependent and need to be tuned empiricall over training iterations. As an alternative, CascadeVAL VAE (Jeong & Song 2019) maximizes the ELBO by iterating over two separate optimizations for the continuous and categorical variables after a warm-up period, instead of a fully gradient-basee optimization. Although the computational cost for the suggested optimization for the categorical variable has an approximately linear dependence the number of categories and
2

---

Under review as a conference paper at ICLR 2021
8
mean: 0.5
2000
(0,
X
S
Na
Ca
1500
5
8
pleabte)
1000
g
g
Cb
p(solca
Xb
Sb
Ca
8.0
1.0.00
Width 1.0
Sa
(a)
(b)
(c)
Figure (a) Empirical distributions for the continuous state variable representing stroke width are digit- dependent, illustrating dependence of style on type. (b) Multi-agent autocncoder framework proposed as cpl-mixVAR model. This framework uses a sype-preserving augmentation network G. Each augmented sample is obtained by first projecting the sample projecting back to the data space. Agents receive noisy copies of given samples x, i.c. {xa: [Xa,X5. -.. where they belong the same class label, to learn mixture representations, i.e. {q(Ca, 8.)) q Co, 86,.5))) 86). .......... Agents cooperate to learn the categorical assignment, p(c). Cooperation is achieved by imposing a penalty on mismatches in the categorical assignments. (c) Each autoencoder agent learns type dependence of the state variable according the graphical model.
lower dimension, then perturbing it with noise and finally
batch size, it can still be a deterrent for problems with numerous categories and unbalanced datasets requiring larger batch sizes. Thus, single-agent VAEs fall short efficiently learning interpretable mixture representations, either due to their reliance on a heuristic embedding capacity, or lacking fully variational approach.
In addition to the issues discussed above, the performance and interpretability of those VAE ap- proaches are further limited by the common assumption that the continuous variable representing the style of the data is independent of the class label. In practice, style often depends on 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 (Fig. [T). Moreover, in the discussed approaches, e.g. VAE, unsupervised clustering of only the continuous variable achieves a relatively high classification accuracy (~ 66%, see supplementary F and G), underscoring that the independence assumption is not valid.
COUPLED MIXTURE VAE FRAMEWORK
The key intuition behind multi-agent networks is cooperation for decision making to improve probabilistic estimation. Collective decision making has been studied under different contexts and popular name referring to its advantages is the "wisdom of the crowd" (Surowiecki 2005). When unanimous decisions made by crowd (multiple agents) form probability distribution, multiple estimates from the same sample increase the expected probability of true assignment. This theoretically justified by Proposition in the context of categorical decision making.
Definition 1. (A-agent VAE Framework) We define the A-agent VAE as A-tuple of independent
and architecturally identical autoencoding agents, where each agent, e.g. a-th agent, parameterizes mixture model distribution with O., Pa. While each agent has its own mixture representation with potentially non-idene non-identical parameters, all agents cooperate to learn qs_ (Cal Xa) and qo. (Ca) via cost function at the time training.
Definition 2. (Confidence) Suppose x is generated by a multivariate mixture distribution so that p(x) E. p(c))(x(e), where p(c) denotes arbirarry prior for the relative abundances of categories. The assignment confidence for category foror samples belonging to category m can be expressed as follows.
C. (k) (1k)= Exxm (logr(c- ( = kk(x)]
(2)
Proposition 1. Consider the problem of mixture representation learning in a multi-agent VAE framework with A > 2 agents using type-preserving data augmentation, where the accuracy of categorical assignment for a single agent imperfect. The confidence of the correct assignment for the multi-agent VAE is higher than that of the single agent VAE. Moreover, there exists A such that the correct category receives the highest confidence score in the A-agent framework, independent of the categorical prior (Proof in supplementary Section A)
3

---

Under review as a conference paper at ICLR 2021
While in unsupervised framework, defining the required number of agents in the absence of categorical prior and daeegory-dependent information remains a challenge, in the particular case of uniform prior distribution of categories, we have the following Corollary.
For uniform prior the discrete factors, one pair of VAE agents (A = 2) is sufficient to increase the confidence of correct categorical assignment. (see supplementary Section A)
Corollary 1.
4.1 AGENT-AGENT VAE
Using the insight obtained from Proposition 1 we formulate collective decision making for an A-agent VAE network (Fig. [b) the following constrained optimization.
max
Layle (#1, 01) + +... + -a.lo (#., Aa)
(3)
s.t. Ci = -... =
Here, La.lc, (Ru: 0o) is the variational loss for agent a as follows,
Ca_le, (Qu, = Eo(s..c [log p(xalsa, ca)l E (xte./x.) DEL (g(Sal Ca (p)su(()))))
E xls..(e.....) DL
(4)
In Eq. (4) for each agent, we use the graphical model in Fig. D and modify the loss function in
Eq. (1) by conditioning state on the categorical variable (derivation in supplementary Section B).
Not only is challenging to solve the maximization in Eq, due to the equality constraint, objective remains function of the prior p(c) which is unknown, and typically non-uniform. To
overcome this, we introduce an equivalent formulation for Eq. 3 based on the pairwise coupling paradigm follows (derivation in supplementary Section C).
Σ (A 1) (Eglen calx.) [log p(x, [(og Sa, Co)] Egce. x.) [DKL (a(salca, xa) (x)ll((s,(c.))))
max
Eggsal xa) xa) a<b
B Xe(sslosiss) [DL (q (x)) xx))lp(ca.c.))
s.l. c. = C₁ Va, b E [1, A], a <
(5)
Here, in the last term, the KL divergence across coupled agents is a function of the joint distribution p(Ca, Co), rather than p(c). relax Eq. 5 an unconstrained problem by assuming a differentiable form for p(Ca; cb) (full derivation in supplementary Section D).
A
max Σ (A 1) (Eylss p(xap(x c.)l Eggen. (x.)) DKL (q(su ca. x lp)ss c.))))
1=1
Σ H(c. ((c))x))) AlEgtes (x.) (Bgtess))))) (d (ca co)]
(6)
acb
According to the final expression in Eq. the agents try to achieve identical categorical assignments while independently learning their own style variables. For each pair of agents, there are two entropy based confidence penalty terms, which mode collapse regularizers (Pereyra al. 2017). There is also distance d(c.,, ca) between a pair of categorical variables, which encourages the consensus on the categorical assignment controlled by coupling hyperparameter X. The distance defined as d(c.,, C2) [|etr(c.) ) cdr( (,,)|2,. Vc,, C2 E S where clr(.) denotes the isometric centered-log: ratio transformation and therefore satisfies the conditions of mathematical metric according to Aitchison geometry (Aitchison, 1982, Egozcue et al. (203)). To sample from q(c, Xx,) in gradient
descent framework, we use the fumbel-softman distribution (Jang et al. 2016 Maddison et al. 2014)
with temperature parameter () < < S 1.
In the rest of this study, we refer the model in Eq. aas cpl-mixVAE (Fig. [p). Note that this formu- lation can be easily extended to include additional hyperparameter to encourage disentanglement of continuous variables as in 3-VAE Higgins et al. (2017). To train this model in an unsupervised fashion according Eq.
require augmented samples Xa for any given sample
4.2 TYPE-PRESERVINC AUGMENTATION
Augmentation can be considered generative process Antoniou al. 2017). We seek a generative model that not only leams the data distribution, but also transformations that represent within-class variations in unsupervised manner. Learning such transformations is generally not straightforward,
4

---

Under review as a conference paper at ICLR 2021
and requires prior knowledge about the underlying invariances. While conventional transformations such as rotation, scaling, or translation can serve as sype-preserving augmentations for many image datasets, they may not capture the richness of the underlying process. Moreover, such augmentation strategies cannot be used when within-class invariance are unknown. Suggested alternatives to conventional augmentations either rely class label, are specific image data (Hauberg et al.. 2016 Jaiswal et al. 2018; Antoniou et al., (0177).
To this end, inspired by DAGAN (Antoniou et 2017). we propose an unsupervised type-preserving augmentation using a VAE-GAN (Larsen et al., 2016)-like architecture, Fig. Tb. We seek network such that noisy copy, Xa can be obtained as variation of the given sample, x. based on its low dimensional representation that is concatenated with Gaussian noise n. To prevent the network from disregarding the noise, we formulate the training procedure as the following minmax optimization which uses discriminator network 'D as regularizer.
min max V (D,9) -R(9) + (9) + gg(g)
g D
While training, 9 generates two samples: Xn and Xy. The former denotes Xa. and the latter is a sam- ple generated in the absence of noise. In Eq. 7 V G) = Ex [log D(x)] + Ex [log (1 ( [)]] Ex.n [log (1 (x,))]] is the value function for the joint training of the discriminator and gen- erator; R (g) [x(x) [log [(x22)] is the reconstruction loss, which operates only over x;
(G) max (||x Xul2 ||x xnl/2 + 0) is the triplet loss that prevents network G from disregarding noise and generating identical samples; and (g) = Dx (g(z) x (l)l//////xx n)) is the distance between the latent variables in the absence and presence of noise. d (G) is regularizer to encourage original and noisy samples to be located close to one another in the latent space and is controlled by hyperparameter y < 1.
4.3 STABILIZING THE TRAINING BY MINI-BATCI VARIANCE
The solution to the maximization problem in Eq. 6 which includes minimization of d(ca, cb), has trivial local optima that result in the mode collapse issue (Lucas et al. 2019), in which during learning the network ignores subset of the discrete latent space. For instance, in the extreme case, the network learns Ca = Ce C, V n, where n denotes the sample index. In this scenario, the state variable is compelled model fails to deliver an interpretable mixture representation despite achieving an overall low loss value. We regularize the Aitchison distance d(ca... Co between the categorical assignments of the n-th samples of agents and by using mini-batch statistics to avoid mode collapse: dz (Ca., Co..
act as the latent variable of a classical variational autoencoder, while the
Σ. (aai log Came The log Cb.. where 0², indicates the variance of the k--th category of agent
2
a. In the following proposition we show that d² an approximation of the Aitchison distance in the probability simplex.
Proposition 2. Suppose Ca., Cb,, E EKK, where SK is a simplex of K parts and n the sample index. If d (Can, Co.,) denotes the Aitchison distance, then
de (Can: Cbm)) - dd² (Can, Cbn.) S K ( + To)² + K222 - -A?)
where To = max flog Caup log Cbay }. To = max (max((om)) 1) log Can,, - (0 1) log Con.g } and
Ao = > [0.1 1) log Cars (002) 1) log Cb.. (Proof in supplementary Section E)
k
Accordingly, as the Gumbel-softmax approximations of the categorical variable of the agents move closer to each other on the simplex, d, converges to d.
EXPERIMENTS
We assess the performance of epl-mix VAE for three different datasets. To facilitate comparisons with other methods, first we conducted experiments on two benchmark datasets: MNIST and dSprites. Additionally, used single cell RNA- Seuunnigg dataset (scRNA-seq) TTssic al.) 2018). to evaluate the utility of our approach in identifying neuronal cell types and sppe-sscciiic genes regulating the continuous within- variability.
MNIST: According to the approximately uniform distribution of handwritten digits in the dataset, we used 2-agent cpl-mix VAE with shared categorical variable to learn interpretable representation.
5

---

Under review as a conference paper at ICLR 2021
Original samples
9
?
4
4
/
O
8
6
4
9
9
4
8
7
9
$
/
3
0
6
1
O
Auamented samples
6
9
7
4
4
/
0
6
9
8
7
6
4
9
9
4
2
Ç
/
3
0
6
9
1
0 0 0 0
6
9
4
4
[
0
6
9
8
7
6
4
9
9
4
8
/
3
0
6
9
1
9
7
4
4
/
0
6
9
7
6
4
9
9
4
/
3
Q
6
9
9
/
7
4
4
7
0
6
9
8
7
6
DE
9
0
4
8
9
/
3
Q
6
9
2
(a)
0 0
0
0
0
9
9 5
0 0 0
0
0
0
0
0
0
0
0
0
0
0 0
0
0
6 0
0
0 0
0
0 2
/
/
//////
(
(
(
7 7
2
2
2 3
2
2277 3333%
3
2
2
2
2
2
2
2
2 2
3
2
2
2
2
2
2
2.
5 J 3
3
3 3
3
3
3
3
3
3
333332
3
3
3
3
3
3
3
3
3
44 4 4
4
4494
A
A
A
H
R
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
R
4
4 4
4
4
4
4
4
4
5 5
5
5 S
5
S
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
4 6
6
6
6
6
6
6
6
6
6
6
6
6
6
9
7
7
7 7
77777
7
7
7
7
7
7
7
7
8 8
8
8
8
1
8
8
8
8
8
8
8
8
3
8
8
8
8
8
8
8
5 8
8
8
8
8
6
9
9
9
9
9
9
9
9
9 9
4
9
9
9
(b)
(c)
(d)
(c)
Figure 2: (a) Augmented samples for the MNIST datasct generated by the tepeetyypeeyeeeeeeee preserving augmentation conserve type of the original sample. (b-e) Continuous latent traversals of 1d agent of the epl-mixVAE framework with two autoencoding agents, where each agent mixture representation is parameterized with 10-dimensional continuous and 10-dimensional categorical variables. Examples of (b) rolation angle, (c) stroke thickness, (d) character width and (c) roundness of looped features are presented. The discrete variable is constant for all reconstructions in the same row.
Each agent learned a mixture generative model including a 10-dimensional categorical variable representing digits (type), and a 10-dimensional continuous random variable representing the writing style (state). To generate noisy samples, the augmenter was trained on the MNIST dataset ahead of training the cpl-mix VAE model. Fig. 2a displays example noisy samples generated by the type-preserving augmentation for MNIST. To quantitatively evaluate the proposed type-preserving data augmentation, we used a benchmark classifier for MNIST digits, which achieves 99.54% accuracy over 10,000 test sample:" Applying the imported classifier to the augmented test samples yields 96.14% classification accuracy, which demonstrates that the augmenter preserves the label information (type) for 96.58% of the augmented samples. During training of cpl-mix VAE, each agent received an augmented copy of the original image. To interpret the role of the continuous factor, fix the discrete latent variable and change the state variable according to the conditional state distribution learned for each category. Fig- big-)( illustrates these continuous latent traversal results for four dimensions of the state variable obtained by cpl-mix VAE. Each row corresponds to different dimension of the categorical variable, and the state variable monotonically changes across columns. Panels (b), (c), and (d) represent commonly-identified continuous factors with global attributes, while panel (c) represents roundness, all in digit-dependent manner.
Table displays the classification performance of the discrete latent variable (as the predicted class label) for InfoGAN, different single-agent VAE methods including JointVAE and Cascade VAE, and cpl-mix VAEE. We report the accuracy of the categorical assignments (ACC) and the reconstruction loss across 10 random initializations. For Cascade VAE, we used the numbers reported in (Jeong Song 2019). For InfoGAN and VAE, we used the publicly available implementation and training procedure reported in (Chen et al., 2016 Dupont, 2018). All reported numbers for cpl- mixVAE models are average accuracies calculated across both agents. We reported the performance of the proposed coupled VAE for four cases: (i) cpl-mix VAE(s / c). in which the state variable is independent of the discrete variable; (ii) cpl-mix
c), which is the model in Eq, Bad includes a pair of independent architecturally identical autoencoder agents using the proposed data augmentation in Eq. 7 (iii) cpl-miXVAE" (s] c). which is also trained according to Eq. 6 and uses the proposed data augmentation, however in this setting, agents are not independent and the networks parameters are shared across agents; and (iv) cpI-mix (s c), in which we used random rotation (l-20, 20] degree) as an affine transformation for augmentation. Our results show that the cpl-mix VAE(s c) obtained the best categorical assignment among all models. Moreover, cpl-mix |c) also
Digit Recognizer, kaggle competition: http://wwwwaagggcc///////ggogggg

---

Under review as a conference paper at ICLR 2021
Table I: Clustering results for the MNIST dataset, over 10 runs with 15K training iterations. For InfoGAN, we used the same network and the same parameters reported in the original paper by [Chen et al. (016). For CascadeVAE. all Joint- VAEs, and cpl-mix models, we used c 10. |s] = 10, and r 0.67. For cpl-mix VAE models, the coupling factor is set
Method
free
ACC 1 (mean + s.d.)
InfoGAN
169.4
77.87 + 21.68
Cascade VAE
81.41 09.54
JointVAE
122.0
68.99 + 11.76
Join
126.7
62.19 1 05.73
JointVAE!
130.5
68.21 + 09.58
VAE(s
113.9
79.63 + 08.32
cpl-mixVAE"
110.1
80.25 1 05.37
cpl-mixVALE
105.9
82.92 04.64
cpl-mixVARe (AE
113.5
84.56 1 06.47
|c|=9
Icl=11
100
1.0
Digits
80
0.8
0.6
ACC
AMP
(b)
40
0.4
20
cpl-mixVAE JointVAE
0.0
(a)
6
9
10 |c|
11
12
13
(d)
(e)
Categories
Categories
Figure 3: Clustering performance for the MNIST dataset, when the number of discrete variahle ((c)) is not equal the true numher of clusters (10 in this case). (a) ACC represents the accuracy of categorical assignment and AMP denotes average maximum posterior probability for JointVAF: and 1s agent of cpl-mix VAE. Error bars indicate mean s.d. over randomly initialized runs. (b-c) Confusion matrices for Joint with AMP=0.57(( cl = 9) and AMP=0.54((cc 11): (d-e) Confusion matrices for 1s agent cpl-mix VAE with AMP=0.88( = 9) and AMP=0.80(( (0 11). Color bar indicates per-category accuracy.
achieved the second highest performance, which demonstrates that even using simple augmentation strategy enhance the representation learning. For a fair comparison, Gumble-softmax temperature, T and latent dimensionality are set to the same values as those for JointVAE and Cascade VAE. To understand whether architectural differences put JointVAE at disadvantage, report the results for JointVARt, which uses the same architecture for the basic encoder/decoder networks as the one used in cpl-mix VAE. That is, JointVAE uses the same learning procedure JointVAE, but its convolutional layers are replaced by fully-connected layers (see supplementary Section J for implementation details). Comparison of the results obtained with JointVAE and JointVAE' suggests that the superiority of cpl-mix VAE is not due to the network architecture. Additionally, to separate the impact augmentation in the training. we report the results for JointVAE in which the JointVAE model has been trained with noisy copies of the original MNIST dataset generated by the proposed augmentation method. The reported clustering performance for JointVAEt suggests that including data augmentation by itself does not enhance the categorical assignment.
also investigate the performance of cpl-miXVAE(s c) for different cardinalities of the categorical variable, Fig. of c = K E 13). Here, AMP denotes the average of maximum posterior of categories i.e., 1/KΣ k=1 max p(c+ xx). Expectedly, an insufficient number of categories results in inaccurate discrete variability encoding, which causes some dimensions being allocated to more than one digit (p and Fig. (1), which results in lower ACC, but higher AMP. On the other hand, additional C₁ leaves some categories under-utilized (Fig. 3 and Fig. 3+. which does not only lead to lower ACC, but also lower AMP. Notably, our results show that while JointVAE suffers from sensitivity to empirical choices of c (Fig. 3p and Fig. 3b), cpl-mix VAE is more robust encoding the discrete variability (Fig. 3 and Fig. Be). As the AMP measure shows, for |c| < 10, cpl-mix VAE utilizes all
3h shows the performance of cpl-mix VAE in terms of ACC and AMP function
categories, without suffering from collapse and for |c| > 10, it does not allocate unneeded categories and maintains high categorical assignment accuracy.

---

Under review as a conference paper at ICLR 2021
Table 2: Categorical assignment accuracies (ACC) and disentanglement scores (DS) for the dSprites dataset,
over 10 randomly initialized runs with cl 3. |s]
6, 7 0.67, and
JointVAE
EascadeVAEE
cpl-mix (AB c)
ACC (mean + s.d.) DS (mean + s.d.)
44.79 03.88
78.84 + 15.65
96.30 + 09.15 89.98 + 04.09
74.51 05.17
90.49 + 05.28
(a)
(b)
(c)
Figure 4: Interpretable continuous latent traversals the trained cpl-mix VAE model with 6-dimensional continuous variable and 3-dimensional categorical variable for the dSprites dataset. Examples of (a) rotation, (b) scale, and (c) position shown. The discrete variable is held fixed for all reconstructions in the same
dSprites: this dataset, due to the uniform distribution of classes, we again used a 2-agent cpl- MiXVAE model. Similar to JointVAE and CascadeVAE, used 3-dimensional categorical variable for learning the shape (type), and 6-dimensional state variable representing the style of each shape. Fig. iilluustrats these traversal results for the three dimensions of the state variable obtained from the cpl-mix VAE. Each row corresponds to different dimension of the categorical variable, and the state variable monotonically varies across columns. Table shows the degree which epl-mix VAE outperforms the other methods in terms of categorical assignment accuracy. In addition, we reported disentanglement scores for all methods. For fair comparison, we used the same disentanglement evaluation metric implemented for Cascade VAE (Jeong & Song 2019).
scRNA-seq: Compared to typical machine learning datasets, the scRNA-seq dataset is exceedingly nigh-dimensional, with over 10, 000 genes. It includes 22, 365 neurons, over 100 cell types with sizeable difference between the most and the least abundant clusters. Here, we excluded non-neuronal cells and used subset of 5, 000 most expressed genes based their peak values. While more than 115 neuronal types are suggested for this dataset (Tasic et al.) 2018). a significant challenge of representation learning in this dataset its substantial imbalance, where for the most- and the least- abundant types, there exist 1404 and 16 samples, respectively. From the perspective of neuroscience, neurons as the basic building blocks of the brain, display both significant diversity and stereotype in their shapes, gene expression, and response patterns. Individual cells inherently differ due to either their type or continuous within-type variations (Trapnell 2013 Andrews & Hemberg 2018).
We used a 1 -i--dimeniional discrete and 2-dimensional continuous variable for discrete and continuous neuronal factors representation, respectively. Fig. [5] illustrates the performance of JAin Cascade and 2-agent cpl-mix VAE model. The dendrograms on the y-axis displays the hierarchical relationship between neuron types according tofTasic et al. (2018). For many neuronal cells, whether the observed diversity corresponds to discrete variability continuum is ongoing debate. While both VAE and Cascade VAE failed to identify meaningful cell types, cpl-mix VAE
Categories
Categories
Categories
150
140
Cell Types
110
JointVAE
CascadeVAE
cpl-mixVAE
# Agents
Figure 5: Categorical assignments for the scRNA-seq dataset. (a) Confusion matrices of JointVAE, CascadeVAE. and cpl-mix VAE trained by |c| = 115, |s| = and T = 1. over 45K iterations. For each model. hyperparameters were assigned as followings: C, E [0, 71, C. E [0, 10] (over 10K iterations), and 3, = 3. 100 for VoE: 3 E [0.1, 10), and X = 0.1 for Cascade VAE; X = 1 for cpl-mixVAE. The dendrogram on the y-axis shows masker-based hierarchical classification with 115 cell types. (b) Accuracy improvement by adding more agents for cpl-mp mix VAE, over runs. A A-agent's performance A 2 compared with baseline 1-agent, JointVAE.
8

---

Under review as a conference paper at ICLR 2021
MG
IEG
HKG-COX
HKG-CC
(8)
1.0
(b)
0.2
0.0
(0)
Calb2
God2
Forc1
Lhvd
Chrall
Rbo).
Poells
Tyrobo
sesse
9I338011
Fost2
N4453
Ngtx
Matxa
ArC
Arenl
Ggn1
Eqra
Homer1
CorSa
CoxBo
Con6c
Cox
Co
Cor6al
Cox
cox24
CoNTc
Cox17
Cibp1
cont2
cont
Ciny
Conh
cond2
Rablla
Rabill
Rabla
Rao21
Figure 6: Continuous latent traversals for two excitatory cell types (a) "LS NP ALM Trhr Neff" and (b) "L6
CT Nxph2 Sla", and an inhibitory cell type (c) "Pvalh Akric18 Ntf3". For each type, the continuous latent traversal is color-mapped to normalized reconstructed gene expression value (colorbar) as function of the state variable for four gene subsets from left right: marker genes (MG), immediate early genes (IEG), and subgroups house keeping genes, cytochrome oxidase (HKG-COX)
cycle regulators (HKG-CC).
successfully identified the majority of known types. The, accuracy for categorical assignment across the entire 11. types 39% (chance level ~ 6%, based on the most abundant type). Unlike the discussed benchmark datasets, the neuronal types another experiment we applied more than two agents on the scRNA-seq dataset to investigate the accuracy improvement for categorical assignment. Fig. illustrates the accuracy improvement with respect the a single agent model, i.e. JointVAE.
uniformly distributed, Accordingly, as
To examine the role of the continuous latent variable, which can profile activity-regulated genes, we applied similar traversal analysis to that used for the MNIST and dSprites datasets. For given cell sample and its discrete type, we changed every dimension of the continuous variable using conditional distribution, and inspected gene expression changes caused by continuous variable alterations. Fig. shows the results of the state traversal experiment for two excitatory neurons belonging to the "L5 NP" (near-projecting) and "L6 CT (corticothalamic) sub-classes, and inhibitory neuron belonging to the "PV" (parvalbumin) class. In each sub-figure, the latent traversal is color-mapped to normalized reconstructed expression values, where the y-axis corresponds to dimension of the continuous variable, and the x-axis corresponds to four gene subsets, namely (i) marker genes (MG) for the two excitatory types, (ii) immediate early genes (IEG), and two house keeping gene (HKG) subgroups (iii) cytochrome e oxidase (COX), and (iv) cell cycle (CC) regulators (Hrvatin et al. 2018 Tarasenko al., 2017). MGs are normally expected to function indicators for particular cell types whose normalized expression is unaffected by the regulatory activities of the cell. Indeed, the expression of the reported excitatory MGs remains constant for excitatory traversals but not necessarily for the inhibitory traversal (i.e., Calb2, Gad2, Pdel la in Fig. 6). In contrast, the expression of IEGs and HKGs should depend strongly on the metabolic and environmental conditions. Indeed, we find that the expression changes of IEGs and HKGs are for the most part monotonically linked to the continuous variable, reaffirming that the it captures relevant, interpretable continuous variability, as in the MNIST and dSprites examples. Lastly, the expression of IEGs and HKGs dactivity-regulate genes) depends on the cell type c.g., not all IEGs are activated for all cell types. Notably, for the excitatory "L.5 NP" and "L-6 CT" cells that proximate in the hierarchy (Tasic al., 2018 state traversal is quite similar. These results suggest that the continuous variable inferred by cpl-mix VAE provides insight when deciphering the molecular mechanisms shaping the landscape of biological states, e.g. metabolic, disease.
CONCLUSION
We have proposed epi-mix VAE as multi-agent framework using type- type-preserving data augmentation to apply the power of collective decision making in unsupervised joint learning of discrete and continuous factors. This framework utilizes multiple pairwise-coupled autoencoding agents with shared categorical variable, while independently learning the continuous variables. Our experimental results for all three datasets show that epl-mix VAE outperforms comparable models. In addition, for challenging gene expression dataset, we showed that the proposed framework can identify annotated neuronal types and differentiate between type-dependent and activity-regulate genes.

---

Under review as a conference paper at ICLR 2021
REFERENCES
John Aitchison. The statistical analysis of compositional data. Journal of the Royal Statistical Sociery: Series (Methodological), 44(2):139-16)) 1982.
Tallulah S Andrews and Martin Hemberg. Identifying cell populations with sernaseq. Molecular aspects of medicine, 19:114- 2018.
Antreas Antoniou, Amos Storkey, and Harrison Edwards. Data augmentation generative adversarial networks. arXiv preprint 1rXiv: 1.11:4340 2017.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):859-877 (771))8888777711111717 2017.
Avrim Blum and Tom Mitchell. Combining labeled and unlabeled data with co-training, In Proceed- ings of the eleventh annual conference on Computational learming 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:1105.00844,, 2017.
Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling 3-vae. arXiv preprint arXiv:1804.033: 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.
Nat Dilokthanakul, Pedro AM Mediano, Marta Garnelo, Matthew CH Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture varia- tional autocncoders. arXiv preprint drXiv:1611.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 Pawlowsky-Glahn Glòria Mateu-Figueras, and Carles Barcelo-Viee 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 the 22nd ACM international conference on Multimedia, pp. 7-16, 2014.
Jerome H Friedman. Greedy function approximation: gradient boosting machine. Annals of statistics, pp. 1189-1232, 2001.
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 of 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 preprini arXiv: 16I:vv::611.06 1.055555 2016.
Søren Hauberg, Oren Freifeld, Anders Boesen Lindbo Larsen, John Fisher, and Lars Hansen. Dream- ing more data: Class-dependent distributions over diffeomorphisms for learned data augmentation. In Artificial Intelligence and Statistics, pp. 342-350, 2016.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. bae: Learning basic visual concepts with constrained variational framework. Iclr, 2(5):6, 2017.
10

---

Under review as a conference paper at ICLR 2021
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.
Ayush Jaiswal, Rex Yue Wu, Wael Abd-A And-Allageed, and Prem Natarajan. Unsupervised adversarial invariance. In Advances in Neural Information Processing Systems, pp- 5092-5102, 2018.
Eric Jang. Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax, arXiv preprint arXiv:161 1.01 144, 2016.
Yeonwoo Jeong and Hyun Oh Song. Learning discrete and continuous factors of data via alternating disentanglement. arXiv preprint UrXiv:I905.094:0 2019.
Zhuxi Jiang, Yin Zheng, Huachun Tan, Bangsheng Tang, and Hanning Zhou. Variational deep embedding: unsupervised and generative approach clustering. In Proceedings of the 26th International Joint Conference on 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 in neural information processing systems, pp.
2016.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. arXiv preprint arXiv: 100:v1180..05983, 2018.
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes. arXiv preprint
Abhishek Kumar and Hal Daumé. co-training approach for multi-view spectral clustering. In Proceedings of the 28th international conference on machine learning (ICML-11), 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 on machine learning, pp. 1558-1566. PMLR, 2016.
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: 123559, 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 disentanglement without compromises. arXiv preprint arXiv:2002.0288%, 2020.
James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't blame the elbo! a linear vae perspective on posterior collapse. In Advances in Neural Information Processing 2019.
Systems, pp.
Chris Maddison, Daniel Tarlow, and Tom Minka. A sampling. In Advances in Neural Information Processing Systems, pp. 3086-3094, 2014.
Stefano Monti, Pablo Tamayo, Jill Mesirov, and Todd Golub. Consensus clustering: a resampling- based method for class discovery and visualization of gene expression mieroarray data. Machine learning. 52(1 22:9911111111 2003.
Gabriel Pereyra, George Tucker, Jan Chorowski, Lukasz Kaiser, and Geoffrey Hinton. Regularizing neural networks by penalizing confident output distributions. arXiv preprint arXiv:1701.663: 06548, 2017.
James Surowiecki. The wisdom of crowds. Anchor, 2005.

---

Under review as a conference paper at ICLR 2021
Tatyana N Tarasenko, Susan E Pacheco, Mary Kay Koenig, Julio Gomez-Rodrigne 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, (5( 225--22222222 1268, 2017.
Bosiljka Tasic, Zizhen Yao, Lucas T Graybuck, Kimberly A Smith, Thuc Nghi Nguyen, Darren Bertagnolli, Jeff Goldy, Emma Garren, Michael N Economo, Sarada Viswanathan, et Shared and distinct transcriptomic cell types across neocortical areas. Nature, 563(7729):72-78, 2018.
Kai Tian, Shuigeng Zhou, and Jihong Guan. Deepeluster: A general clustering framework based on deep learning. Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp.
Springer, 2017.
Cole Trapnell. Defining cell types and states with single-cell genomics. Genome research, 25(10): 1491-1498, 2015.
12


Reviews:
[{"review_id": "dXg_NYYmK-y", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "====================================================================================================\n\nSummary : \n\nThe paper proposed the new disentanglement approach based on the \"wisdom of the crowd\". First, the proposed method enforces the consensus on the categorical assignments from the different agents. Each agent receives a similar image that is generated from the data augmentation method (based on the VAE-GAN technique).  The encoder of each agent first estimates the categorical distribution and estimates continuous variables from the categorical data and original image. The proposed method focuses on disentangling categorical information. The evaluation also focuses on categorical data and the proposed method outperforms the other single-agent baselines.  \n\n===================================================================================================\n\nReason for score: \n\nOverall, I'm on the borderline between acceptance and rejection. The proposed method is persuasive, but the ablation study is required to improve this paper. I will finalize my decision after reading the reviews from the other reviewers and authors' responses to my concerns.\n\n====================================================================================================\n\nStrong points: \n\n(1) I was impressed by the idea to generate a similar image with the same class label through techniques developed from VAE-GAN since this idea reverts the unsupervised setting into a weakly-supervised setting. In this regard, I have a concern (See cons). \n\n(2) I like the proposed multi-agent approach to increase the robustness of the classification label. Maybe this method largely reduces the variance from the different seeds.\n\n====================================================================================================\n\nCons :\n\nThe idea which leverages a similar image with the same class label is quite similar in the disentanglement under weakly supervised setting as in [1]. Therefore, type-preserving augmentation seems critical in the performance of the proposed method. In this sense, I believe the ablation study on the augmentation method is required to persuade readers and to provide guidelines for future research. (e.g. Adding gaussian noise, Mirror image)\n\nEven though the disentanglement score cannot be evaluated without supervision on the labels, I recommend authors to compare the disentanglement score with other baselines(Joint VAE, Cascade VAE), 'I(s;c)', and etc. \n\n=============================================================================================================\n\nAfter Rebuttal : \n\nSorry for the delay. I checked the comments of the other reviewers, responses, and the revised version. The authors address my concerns in the revised version. I vote for marginal acceptance. \n\n[1] Weakly-Supervised Disentanglement Without Compromises, ICML 2020", "extra_scores": {}, "extra_text": {}}, {"review_id": "cDhEfdZCS3", "reviewer": "AnonReviewer2", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposed a multi-agent VAE model that combines multiple copies of VAEs with coupling constraints to improve its latent representation learning (by encouraging discrete variable consistency). The experiments show that the proposed model outperforms other discrete&continuous VAE models in terms of clustering ACC. The experiment is conducted on MNIST and scRNA-seq datasets.\n\nOverall\n===\nI think this is a good submission in terms of describing its methodology. However, it is hard to justify its motivation as the model is over-complex, and the same task could be achieved by other approaches such as deep clustering. It is better if the authors could include such discussions in the paper. \n\nPros\n===\n1. This paper gives sufficient justification on why we need multiple copies of estimators to reach better estimation. \n2. The method description given in Section 3.1 is quite clear as the equations are self-explainable. \n3. Results shown in table 2 indicate significant performance improvement.\n\n\nCons\n===\n1. The approach proposed itself is very complex comparing to other generative model architecture. It is hard to justify if the model can generalize to other tasks other than simple, low-dimensional tasks. Consider maintaining multiple VAEs graphs in memory; it appears hard to take advantage of this work. Is the cost of deploying this model worth the performance improvement (as Table 1 shows limited improvement)? Also, to train such a model, the user also needs to produce a type-preserving augmentation, which is very costly for clustering.\n\n2. The proposed model jointly optimizes the main objective (as Equation 3) and also optimizes the relaxed equivalent constraints at the same time (as Equation 6 last component). Is there any justification why not alternative optimization but joint? As the DeepClustering paper mentioned, alternative optimization is better than joint in their case. Not sure if there will be a similar observation here. Isn't it comparable to the DeepClustering paper if the goal is to do the clustering and interpretation?\n\n3. The introduction of this paper gives me a hard time to follow as the terminology used is uncommon in generative model literature (single-agent, multi-agent). I was confused as it appears to be an RL paper at the beginning.\n\n\nMinors\n===\n1. Which style state shown in Figure 2 (b-e) if you have multiple agents? The description said those are four style/state dimensions, but there are two agents in the experiment setting, and each of them has 10 style/state dimensions. Please be more precise. \n2. Figure 1 (b-c) is not quite informative, and they are too close to each other. It is better to provide a more intuitive demonstrative figure. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "DmRG_JR7jow", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Quality\nThe paper proposes good ideas that are compelling and make sense. Using multiple agents and data augmentation seems directions worth pursuing to improve mixture representation learning.\nThe experiments are mostly clear and well-motivated. Unfortunately, the related work is limited and could be improved and extended. For the evaluation of the data augmentation step, there is only a qualitative evaluation in Figure 2. A more in-depth analysis regarding the type-consistency would strengthen this part of the pipeline.\nIn section 2, the authors say that previous work falls short in efficiently learning mixture distributions due to hyper-parameter tuning or additional cost due to the optimization procedure. In the proposed work, there are some (e.g. coupling hyper-parameter lambda, temperature tau or gamma in data augmentation) new hyper parameters introduced as well. This somehow contradicts the motivation of the method and it is not evident why the proposed method is more efficient. Additionally, ablation studies on critical hyper-parameters like lambda or the temperature parameter would be interesting to see and strengthen the paper. Without further evidence, the claim of a more efficient method is not really justified in my opinion.\n\nClarity\nIn general, the paper is well and clearly written. A related work section is missing which makes it more difficult to position the proposed work with respect to previous work. The authors motivate the representation learning part but sometimes use the term clustering. In my opinion, this weakens the motivation and goal of the paper.\n\nOriginality\nThe paper presents a novel view on mixture representation learning using data augmentation which itself is not a new idea nor do they use a new method to augment the data. The multi-agent view on mixture representation learning is a new idea (to my knowledge, but not aware of all related work) making the paper original.\n\nSignificance\nThe ideas of the paper are significant and worth pursuing.\n\nQuestions to the authors\n-Proof of Proposition 1, equation 7: I do not fully understand where the factor A is coming from in eq. 7. Isn’t this assuming that all agents agree on p(x_a | phi = m)? Is this valid? Thanks for explaining this a bit more in detail.\n-Is your final objective still a valid ELBO? In the proof in the appendix you use approximations to derive the objective, but not only bounds. So I am not sure if it is still an ELBO.\n-In the experiments (Table 1), why is only jointVAE used with augmented data and not cascadeVAE which seems to perform better in its vanilla version? \n\nFurther Comments\n-The work has limited comparison to previous work. Bouchacourt et al.’s “Multi-Level Variational Autoencoder: Learning Disentangled Representations from Grouped Observations” presents a similar idea of using content/class and style/state spaces (only with continuous variables). A comparison to this work showing the potential benefit of using a discrete class space would improve the quality of the paper. \n-I am happy to upgrade my score if the authors address my concerns.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "6U-U2CdfW4", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work presents a new approach to handling categorical latent variables in VAEs. The method has two key components: a *multi-agent* architecture in which categorical assignments are generated through consensus across multiple VAE models and a data-augmentation method which allows each model (termed *agent*) to be trained on a perturbed version of the original data. The authors provide some theoretical justification for their approach and evaluate on two benchm`ark data-sets and one real-world application.\n\n**Strengths**\n1. The method presented seems to be quite novel, with various technical contributions required to jointly train the separate VAE models in a way that they didn't collapse to a single model.\n2. The evaluation on MNIST and dSprites provide empirical evidence that this method out-performs a number of baseline methods. To my knowledge, the methods selected are good candidates for being SOTA methods for VAEs with discrete latent variables.\n3. The application to the scRNA-seq provides further evidence that the method works well on real-world data.\n\nHowever, I have some concerns with the presentation of the paper as is. If the authors can address these, I am happy to increase my score.\n\n**Questions / Concerns**\n1. While I did not have time to check all the proofs in detail, I wasn't convinced by the proof of Proposition 1 in Appendix A. A few specific things on which I would appreciate clarification from the authors:\n\na) It seems like the proof (as stated) would follow if all the agents were identical, but this doesn't match what one would expect. Can the authors make the impact of using different agents more explicit in the assumptions and explain how this leads to the conclusion?\n\nb) The proof assumes the augmented samples are generated from p(x | \\phi = n) but the method generates augmented samples as augmentations of a specific training example x_n. Can the authors explain why they believe this isn't an issue to the relevance of their result?\n\nc) Going from (6) to (7) appears to assume that the probability for each augmented sample is the same. We would not expect this to be the case for a single agent or multiple agents. Can the authors explain the basis for this step?\n\n2. I believe there are a number of different lines of research which could be considered as relevant related work which the authors have missed. \n\na) Boosting Variational Inference [1,2] describes techniques which use a mixture of variational distributions within the VI setup to provide a better approximation to the posterior. While the method presented here is not strictly a boosting method, it is quite possible that the gains seen are a result of being able to better approximate the true posterior by allowing the model to fit multiple inference networks. \n\nb) Consensus clustering aggregates the results of a clustering algorithm over multiple initialisations [3]. This often gives an improvement over the output of a single instance of a clustering algorithm and can be applied post-hoc to any algorithm that outputs cluster assignment. For example, a consensus version of CascadeVAE could be considered as an additional method in the experiments.\n\nc) Co-training is a method for training classifiers on multi-view data such that they predict same labels for co-occurring patterns in each view. This was originally presented by Blum and Mitchell [4] as an approach for semi-supervised learning and was subsequently applied to unsupervised learning by Kumar and Daume III [5]. The key idea in methods that leverage co-tr`aining is to use consensus across the different models, which feels similar to the consensus constraint in equation (3). As well as being relevant related work, the theoretical perpsective of co-training might also be useful to address some of the concerns stated about the theoretical result in this paper.\n\n3. I'm not sure how informative Figure 3 is. It's to be expected that the best performance will come when the dimension of the latent variable is equal to the true number of classes, but this will not be available in the fully-unsupervised case where the true classes are unknown. It would be useful if model selection using AMP was shown to lead to the correct number of classes for MNIST, but it wasn't clear to me if this could be inferred from the plot. If this is what the authors intended to show, they should add some additional text to explain why this is the case.\n\n4. In the experiments on MNIST, the authors compare to the case where m=10, S=10 (Table 3 of Jeong and Song). But Jeong and Song reported higher accuracy (with lower variance) for  the case with m=4, S=10. Why did the authors choose to compare against the lower performing of the two configurations studied in Jeong and Song?  \n\n5. It wasn't 100% clear to me whether the model uses fully independent VAEs (i.e. separate parameters for encoder and decoder in each model), but I believe this is the case. It would be informative to compare against the case where the decoder network has the same parameters in each model. I suspect the gain is primarily due to the fact that we are aggregating over different encoders / inference networks and this additional experiment would make this clear. \n\n**References**\n1. Guo et al. (2016). Boosting Variational Inference\n2. Locatello et al. (2018). Boosting Black Box Variational Inference\n3. Monti et al. (2003). Consensus Clustering: A Resampling-Based Method for Class Discovery and Visualization of Gene Expression Microarray Data\n4. Blum and Mitchell (1998). Combining labeled and unlabeled data with co-training\n5. Kumar and Daume III (2011). A co-training approach for multi-view spectral clustering", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces and analyses a method to train a population of VAEs with mixed continuous (referred to as "style") and discrete (referred to as "labels") latent-variables. The population is trained under the constraint that inferred discrete latent variables  to be the same for all models.
The paper also investigates a data augmentation mechanism inspired by  (Antoniou et al., 2017).
The presentation is overall clear and the idea is interesting, although the language of "agents" is not standard in generative model literature and is a bit confusing. The experiments also show very good clustering results of the proposed method.
Unfortunately the pipeline was determined to be quite complex while the motivation for its design choices were unclear. This, combined with multiple concerns about the experimental validation, led to a reject decision.

Decision: Reject