M IXTURE
R EPRESENTATION L EARNING WITH C OUPLED A UTOENCODING A GENTS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 justiﬁcation 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-speciﬁc, activity-regulated genes for a single-cell gene expression dataset proﬁling over 100 cortical neuron types.
I NTRODUCTION
Complex phenomena can be attributed to a mixture of discrete and continuous factors of variability.
Such complexity is crucial to understand in a 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 can 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 ﬁrst deep generative models to jointly represent discrete and continuous factors, in which a 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, 2019), variational approaches, e.g. JointVAE and CascadeVAE have produced promising and more stable results (Dupont, 2018; Jeong & Song, 2019).
However, such variational methods utilizing a single autoencoding agent rely either on a heuristic data-dependent embedding capacity, or on solving a 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 multimodal 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 pairwise-coupled autoencoders to learn a mixture representation for a 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 a novel type-preserving augmentation that generates noisy copies of the data using withinclass 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 a coupled multi-agent architecture can increase inference accuracy and robustness by exploiting within-cluster variabilities, without requiring a prior distribution on the relative abundances of categories.
Our contributions can be summarized as follows: (i) We ﬁrst provide theoretical justiﬁcation to motivate the advantage of collective decision making for more accurate categorical assignments, utilizing noisy copies of the same sample.
To obtain such samples, we propose an unsupervised type-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 a challenging single cell gene expression dataset for a population of neurons.
We demonstrate that our method can be used to discover discrete categories referred to as neuronal types and type-speciﬁc genes regulating the continuous within-type variability (e.g., metabolic states, disease states).
P RELIMINARIES
For an observation x ∈ R D , a variational autoencoder (VAE) learns a generative model p θ ( x | z ) and a variational distribution q φ ( z | x ) , where z ∈ R M for M (cid:28) D is 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 et al., 2018).
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 a categorical latent variable c ∈ S K , denoting the class label deﬁned in a K -simplex, alongside the continuous latent variable s ∈ R M .
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., 2017) for a single autoencoding agent with the distributions parameterized by θ and φ is given by,
[Equation]
Maximizing ELBO in Eq.
1 to jointly learn q ( s | x ) and q ( c | x ) is challenging due to the mode collapse problem, where the network ignores a subset of latent variables.
Akin to β -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 empirically over training iterations.
As an alternative, CascadeVAE (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-based optimization.
Although the computational cost for the suggested optimization for the categorical variable has an approximately linear dependence on the number of categories and 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 of efﬁciently learning interpretable mixture representations, either due to their reliance on a heuristic embedding capacity, or lacking a fully variational approach.
In addition to the issues discussed above, the performance and interpretability of those VAE approaches 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. 1a).
Moreover, in the discussed approaches, e.g. JointVAE, unsupervised clustering of only the continuous variable achieves a relatively high classiﬁcation accuracy ( ∼ 66% , see supplementary F and G), underscoring that the independence assumption is not valid.
C OUPLED 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 a popular name referring to its advantages is the “wisdom of the crowd” (Surowiecki, 2005).
When unanimous decisions made by a crowd (multiple agents) form a probability distribution, multiple estimates from the same sample increase the expected probability of true assignment.
This is theoretically justiﬁed by Proposition 1 in the context of categorical decision making.
Deﬁnition 1.
(Conﬁdence) Suppose the sample x is generated by a multivariate mixture distribution, so that p ( x ) = (cid:80) φ p ( φ ) p ( x | φ ) , where p ( φ ) denotes an arbitrary prior for categories abundance.
[Equation]
Proposition 1.
Consider the problem of mixture representation learning in a multi-agent framework, where the accuracy of categorical assignment for a single agent is imperfect.
There exists an A , number of the agents such that the conﬁdence of correct categorical assignment for A -agent is higher than that of one agent.
(Proof in supplementary Section A)
While in an unsupervised framework, deﬁning required number of agents in the absence of the categorical prior and category-dependent information remains as a challenge, in the particular case of uniform prior distribution over categories, we have the following Corollary.
Corollary 1.
For a uniform prior on the discrete factors, one pair of coupled agents ( A = 2 ) is sufﬁcient to increase the conﬁdence of correct categorical assignment.
(see supplementary Section A)
3. M ULTI - AGENT VAE
Using the insight obtained from Proposition 1, we formulate collective decision making for an
A -agent VAE network (Fig. 1b) as the following constrained optimization.
[Equation]
Here, L s a | c a ( φ a , θ a ) is the variational loss for agent a as follows,
[Equation]
In Eq.
(4), for each agent, we use the graphical model in Fig.
1b and modify the loss function in
Eq.
(1) by conditioning state on the categorical variable (derivation in supplementary Section B).
Not only is it challenging to solve the maximization in Eq.
3 due to the equality constraint, but the objective remains a function of the prior p ( c ) which is unknown, and typically non-uniform.
To overcome this, we introduce an equivalent formulation for Eq.
based on the pairwise coupling paradigm as follows (derivation in supplementary Section C).
[Equation]
Here, in the last term
[Equation]
the KL divergence across coupled agents is a function of the joint distribution p ( c a , c b ) , rather than p ( c ) .
We relax Eq.
5 into an unconstrained problem by assuming a differentiable form for p ( c a , c b ) (full derivation in supplementary Section D)
[Equation]
According to the ﬁnal expression in Eq. 6, 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 conﬁdence penalty terms, which are mode collapse regularizers (Pereyra et al., 2017). There is also distance d ( c a , c b ) between a pair of categorical variables, which encourages the consensus on the categorical assignment controlled by coupling hyperparameter λ .
The distance is deﬁned as d ( c a , c b ) = (cid:107) clr
( c a ) − clr ( c b ) (cid:107) 2 , ∀ c a , c b ∈ S K , where clr ( . ) denotes the isometric centered-logratio transformation and therefore d satisﬁes the conditions of a mathematical metric according to Aitchison geometry (Aitchison, 1982; Egozcue et al., 2003).
To sample from q ( c a | x a ) in a gradient descent framework, we use the Gumbel-softmax distribution (Jang et al., 2016;
Maddison et al., 2014) with a temperature parameter 0 < τ ≤ 1 .
In the rest of this study, we refer to the model in Eq.
6 as cpl-mixVAE
(Fig. 1b).
Note that this formulation can be easily extended to include an additional hyperparameter to encourage disentanglement of continuous variables as in β -VAE Higgins et al.
To train this model in an unsupervised fashion according to Eq.
6, we require augmented samples x a for any given sample x .
3. T YPE - PRESERVING AUGMENTATION
Augmentation can be considered as a generative process (Antoniou et al., 2017).
We seek a generative model that not only learns the data distribution, but also transformations that represent within-class variations in an unsupervised manner.
Learning such transformations is generally not straightforward, and requires prior knowledge about the underlying invariances.
While conventional transformations such as rotation, scaling, or translation can serve as type-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 on class label, or are speciﬁc to image data (Hauberg et al., 2016; Jaiswal et al., 2018; Antoniou et al., 2017).
To this end, inspired by DAGAN (Antoniou et al., 2017), we propose an unsupervised type-preserving augmentation using a VAE-GAN (Larsen et al., 2016)-like architecture, Fig. 1.c.
We seek a network
G such that a noisy copy, x a can be obtained as a 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 a discriminator network D as a regularizer. [Equation]
[Equation] While training, G generates two samples: x n and x (cid:54) n .
The former denotes x a , and the latter is a sample generated in the absence of noise.
In Eq.
7, V ( D , G ) =
E x [log D ( x )] + E x
[log (1 − D ( x (cid:54) n ))] + E x , n [log (1 − D ( x n ))] is the value function for the joint training of the discriminator and generator; R ( G ) =
E q ( z | x ) [log p ( x | z )] is the reconstruction loss, which operates only over ˆ x ; T α ( G ) = max ( (cid:107) x − x (cid:54) n (cid:107) 2 − (cid:107) x − x n (cid:107)
+ α, 0) is the triplet loss that prevents network G from disregarding noise and generating identical samples; and d ( G )
=
D KL ( q ( z | x ) (cid:107) q
( z | x , n )) is the distance between the latent variables in the absence and presence of noise.
d ( G ) is a regularizer to encourage original and noisy samples to be located close to one another in the latent space and is controlled by hyperparameter γ (cid:28) 1 .
3. S TABILIZING THE TRAINING BY MINI - BATCH VARIANCE
The solution to the maximization problem in Eq. 6, which includes minimization of d ( c a , c b ) , has trivial local optima that result in the mode collapse issue (Lucas et al., 2019), in which during learning the network ignores a subset of the discrete latent space.
For instance, in the extreme case, the network learns c a n = c
b n = c , ∀ n , where n denotes the sample index.
In this scenario, the state variable is compelled to act as the latent variable of a classical variational autoencoder, while the model fails to deliver an interpretable mixture representation despite achieving an overall low loss value.
We regularize the Aitchison distance d ( c a n , c b n ) ,between the categorical assignments of the n -th samples of agents a and b by using mini-batch statistics to avoid mode collapse: d 2 σ ( c a n , c b n ) = (cid:80) k (cid:16) σ − 1 a k log c a nk − σ − 1 b k log c
b nk (cid:17) 2 , where σ 2 a k indicates the variance of the k -th category of agent a .
In the following proposition we show that d 2 σ is an approximation of the Aitchison distance in the probability simplex.
Proposition 2.
Suppose c a n , c b n ∈ S K , where S K is a simplex of K parts and n is the sample index.
If d ( c a n , c b n ) denotes the Aitchison distance, then
[Equation]
where τ c = max k { log c a nk − log c
b nk } , τ σ = max k { ( σ − 1 a k − 1) log c a nk − ( σ − 1 b k − 1) log
c
b nk } , and ∆ σ = (cid:88)
[Equation]
(Proof in supplementary Section E)
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 .
E XPERIMENTS
We assess the performance of cpl-mixVAE for three different datasets.
To facilitate comparisons with other methods, ﬁrst we conducted experiments on two benchmark datasets: MNIST and dSprites.
Additionally, we used a single cell RNA-sequencing dataset (scRNA-seq) (Tasic et al., 2018), to evaluate the utility of our approach in identifying neuronal cell types and type-speciﬁc genes regulating the continuous within-type variability.
MNIST : According to the approximately uniform distribution of handwritten digits in the dataset, we used a 2-agent cpl-mixVAE with shared categorical variable to learn an interpretable representation. 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-mixVAE model.
Fig. 2a displays example noisy samples generated by the type-preserving augmentation for MNIST. During training of cpl-mixVAE, each agent received an augmented copy of the original image.
To interpret the role of the continuous factor, we ﬁx the discrete latent variable and change the state variable according to the conditional state distribution learned for each category.
Fig. 2(b-e) illustrates these continuous latent traversal results for four dimensions of the state variable obtained cpl-mixVAE.
Each row corresponds to a different dimension of the categorical variable, and the state variable monotonically changes across the columns.
Panels (b), (c), and (d) represent commonly-identiﬁed continuous factors with global attributes, while panel (e) represents roundness, all in a digit-dependent manner.
[Figure] [Table] Table 1 displays the classiﬁcation performance of the discrete latent variable (as the predicted class label) for InfoGAN, different single-agent VAE methods including JointVAE and CascadeVAE, and cpl-mixVAE. We report the accuracy of the categorical assignments (ACC) and the reconstruction loss across 10 random initializations.
For CascadeVAE, we used the numbers reported in (Jeong & Song, 2019). For InfoGAN and JointVAE, 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 three cases: (i) cpl-mixVAE ( s (cid:54) | c ) , in which the state variable is independent of the discrete variable, (ii) cpl-mixVAE ( s | c ) , which is the model in Eq.
6 using the proposed data augmentation, and (iii) cpl-mixVAE a ( s | c ) , in which we used random rotation ( [ − 20 , 20] degree) as an afﬁne transformation for augmentation. Our results show that the cpl-mixVAE ( s | c ) obtained the best categorical assignment among all models. Moreover, cpl-mixVAE a ( s | c ) also achieved the second highest performance, which demonstrates that even using a simple augmentation strategy can enhance the representation learning.
For a fair comparison, Gumble-softmax temperature, τ and latent dimensionality are set to the same values as those for JointVAE and CascadeVAE.
To understand whether architecture differences put JointVAE at a disadvantage, we report the results for JointVAE † , which uses the same architecture for the basic encoder/decoder networks as the one used in cpl-mixVAE.
That is, JointVAE † uses the same learning procedure as JointVAE, but its convolutional layers are replaced by fully-connected layers (see supplementary Section H for implementation details).
Comparison of the results obtained with JointVAE and JointVAE † suggests that the superiority of cpl-mixVAE is not due to the network architecture.
Additionally, to separate the impact of 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 JointVAE ‡ suggests that including data augmentation by itself does not enhance the categorical assignment.
We also studied the performance of cpl-mixVAE ( s | c ) for different cardinalities of the categorical variable, c , to discover the true number of categories.
Fig.
3a shows the performance of cpl-mixVAE in terms of ACC and AMP as a function of | c | = K ∈ [7 , 13] .
Here, AMP denotes the average of maximum posterior of categories i.e., 1 /K (cid:80)
Kk =1 max p ( c k | x ) .
Our results show that the best performance is obtained for K = 10 .
Fig. 3b demonstrates the probabilistic assignment of c for all digits.
As illustrated, an inadequate number of categories (left panel) results in some dimensions being used for more than one digit, while extra c k leaves some categories under-utilized (right panel).
[Table] dSprites: In this dataset, due to the uniform distribution of classes, again we used a 2-agent cplmixVAE model. Similar to JointVAE and CascadeVAE, we used a 3-dimensional categorical variable for learning the shape (type), and a 6-dimensional state variable representing the style of each shape.
Fig.
4 illustrates these traversal results for the three dimensions of the state variable obtained from the cpl-mixVAE.
Each row corresponds to a different dimension of the categorical variable, and the state variable monotonically changes across the columns.
In addition, Table 2 shows the degree to which cpl-mixVAE outperforms the other two methods in clustering.
scRNA-seq: Compared to typical machine learning datasets, the scRNA-seq dataset is exceedingly high-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 a subset of 5,000 most expressed genes based on their peak values.
While more than 115 neuronal types are suggested for this scRNA-seq dataset (Tasic et al., 2018), a signiﬁcant challenge of representation learning in this dataset is its substantial imbalance, where for the mostand the least-abundant types, there exist 140,4 and 16 samples, respectively.
From the perspective of neuroscience, neurons as the basic building blocks of the brain, display both signiﬁcant diversity and stereotypy in their shapes, gene expression, and response patterns.
Individual cells inherently differ due to either their type or continuous within-type variations (Trapnell, 2015; Andrews & Hemberg, 2018).
We used a 115-dimensional discrete and a 2-dimensional continuous variable for discrete and continuous neuronal factors representation, respectively.
Fig.
5a illustrates the performance of the 2-agent cpl-mixVAE model.
The dendrograms on the y-axis displays the hierarchical relationship between neuron types according to Tasic et al.
For many neuronal cells, whether the observed diversity corresponds to discrete variability or a continuum is an ongoing debate.
While JointVAE failed to identify meaningful cell types, cpl-mixVAE successfully identiﬁed the majority of known types. The, accuracy for categorical assignment across the entire 115 types is 39% (chance level is ∼ 6% , based on the most abundant type). Unlike the discussed benchmark datasets, the neuronal types are not uniformly distributed, Accordingly, as another experiment we applied more than two agents on the scRNA-seq dataset to investigate the accuracy improvement for categorical assignment.
Fig.
5b illustrates the accuracy improvement with respect the a single agent model, i.e. JointVAE.
To examine the role of the continuous latent variable, which can proﬁle activity-regulated genes, we applied a similar traversal analysis to that used for the MNIST and dSprites datasets.
For a 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 an inhibitory neuron belonging to the “PV” (parvalbumin) class.
In each sub-ﬁgure, the latent traversal is color-mapped to normalized reconstructed expression values, where the y-axis corresponds to one dimension of the continuous variable, and the x-axis corresponds to 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 c oxidase (COX), and (iv) cell cycle (CC) regulators (Hrvatin et al., 2018;
Tarasenko et al., 2017).
MGs are normally expected to function as 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, Pde11a in Fig. 6).
In contrast, the expression of IEGs and HKGs should depend strongly on the metabolic and environmental conditions.
Indeed, we ﬁnd that the expression changes of IEGs and HKGs are for the most part monotonically linked to the continuous variable, reafﬁrming that the it captures relevant, interpretable continuous variability, as in the MNIST and dSprites examples.
Lastly, the expression of IEGs and HKGs (activity-regulated genes) depends on the cell type e.g., not all IEGs are activated for all cell types.
Notably, for the excitatory “L5 NP” and “L6 CT” cells that are proximate in the hierarchy (suggested by Tasic et al. , state traversal is quite similar.
These results suggest that the continuous variable inferred by cpl-mixVAE provides insight when deciphering the molecular mechanisms shaping the landscape of biological states, e.g. metabolic, disease.
C ONCLUSION
We have proposed cpl-mixVAE as a multi-agent framework using a 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 a shared categorical variable, while independently learning the continuous variables.
Our experimental results for all three datasets show that cpl-mixVAE outperforms comparable models.
In addition, for a challenging gene expression dataset, we showed that the proposed framework can identify annotated neuronal types and differentiate between type-dependent and activity-regulated genes.