I NFERRING P RINCIPAL C OMPONENTS IN THE S IMPLEX WITH M ULTINOMIAL V ARIATIONAL A UTOENCODERS
Anonymous authors
Paper under double-blind review
A BSTRACT
Covariance estimation on high-dimensional data is a central challenge across multiple scientiﬁc disciplines.
Sparse high-dimensional count data, frequently encountered in biological applications such as DNA sequencing and proteomics, are often well modeled using multinomial logistic normal models.
In many cases, these datasets are also compositional, presented item-wise as fractions of a normalized total, due to measurement and instrument constraints.
In compositional settings, three key factors limit the ability of these models to estimate covariance: (1) the computational complexity of inverting high-dimensional covariance matrices, (2) the non-exchangeability introduced from the summation constraint on multinomial parameters, and (3) the irreducibility of the component multinomial logistic normal distribution that necessitates the use of parameter augmentation, or similar techniques, during inference.
We show that a variational autoencoder augmented with a fast isometric log-ratio (ILR) transform can address these issues and accurately estimate principal components from multinomially logistic normal distributed data.
This model can be optimized on GPUs and modiﬁed to handle mini-batching, with the ability to scale across thousands of dimensions and thousands of samples.
I NTRODUCTION
Many scientiﬁc disciplines that collect survey data, such as economics, psychology, political science and the biological sciences routinely deal with compositional data, where only relative information can be measured.
These datasets are often in the form of counts, where the total counts within a sample are only indicative of the conﬁdence of the measured proportions.
The resulting proportions lie within a simplex and failing to account for the structure of this simplicial sample space can confound the interpretation of the measurements.
As a result, there has been wide discussion across disparate disciplines (1; 2; 3; 4) concerning the reproducibility crisis that has arisen from the misinterpretation of compositional data.
One of the obstacles to the appropriate analysis of compositional data is the difﬁculty of efﬁciently estimating the latent parameters that lie in the simplex. Accurately scaling probabilistic inference across high-dimensional count data is a major outstanding challenge (5).
This problem is apparent in the social sciences and is particularly pronounced in biological ﬁelds where datasets can obtain observations on tens of thousands of features across hundreds or millions of samples.
One major computational bottleneck with Gaussian distributed data is the inversion of a d -dimensional covariance matrix that has a runtime of O ( d 3 ) (6; 7).
As a result, probabilistic covariance estimation for high-dimensional data is a computationally challenging problem.
Recent theoretical developments (8) cementing the connection between Variational Autoencoders (VAEs) (9) and Probabilistic Principal Components Analysis (PPCA) (10) holds much promise for enabling accurate, scalable, low-rank approximations of large covariance matrices.
Variational autoencoders were originally proposed as a generative model (9), but are now commonly deployed across scientiﬁc disciplines and have made contributions to single-cell RNA sequencing (11), microbiome modeling (12), protein modeling (13; 14; 15), natural language processing (16) and image processing (9).
Following insights that connected regularized linear autoencoders and PCA (17), Lucas et al. (8) showed that carefully designed VAEs can recover the weights that are solved by PPCA.
A computational advantage of VAEs is that they do not require the inversion of a covariance matrix, and the resulting runtime is O ( ndkT ) for n samples, d dimensions, k latent dimensions and T epochs.
While it has been noted that VAEs may take tens of thousands of epochs to estimate the principal component (18), VAEs are easily parallelizable and can be accelerated with GPUs, presenting an attractive alternative to estimating principal components (17) and the resulting covariance matrix.
The connection between VAEs and PPCA is currently limited to Gaussian distributed data and not well-suited to a compositional setting.
Showing that VAEs can recover the correct principal components from count data is nontrivial due to the non-conjugacy issues between the logistic normal distribution and count distributions such as the multinomial distribution.
Furthermore, the parameters of the multinomial distribution are compositional; they are constrained within the simplex and the resulting covariance matrix is singular and non-invertible (1; 19).
Aitchison (20) showed that PCA can be adapted to compositional data through the use of the center log-ratio (CLR) transform, which maintains isometry.
However, this transformation is not isomorphic, requiring that the resulting log-ratios sum to zero, and as a result, CLR-transformed data will produce a singular covariance matrix and rank-deﬁcient principal components.
It has been shown that the isometric log-ratio (ILR) transform (21) satisﬁes both isomorphism and isometry and can handle this singularity issue (22; 23) while enabling the estimation of full-rank principal components.
Here, we show that VAEs augmented with the ILR transform can infer principal components learned from PPCA on multinomially distributed data, beginning to address these critical shortcomings.
R ELATED W ORK
In the microbiome literature, there have been a number of methods (24; 25; 26; 27; 28) that have attempted to model ecological networks through the estimation of pairwise microbe correlations or pairwise inverse-covariance, where microbes are aggregated at different taxonomical scales or ‘taxa’.
Of these tools, only Flashweave can scale across more than thousands of taxa; however, it does this by avoiding the estimation of the covariance matrix.
Methods that attempt to estimate the covariance matrix can only handle on the order of a few thousand dimensions.
Although there is no widely accepted consensus deﬁnition of Multinomial PPCA in this context, being able to efﬁciently estimate the parameters of Multinomial PPCA would be highly useful for exploratory biological analysis.
A number of studies have proposed using mixture modeling as a proxy for PCA (29; 30; 31); however, these techniques depend either on the Dirichlet distribution, whose covariance matrix is not ﬂexible, or on stick-breaking, which violates permutation invariance (32).
Lucas et al. (8) has previously shown that the following two models can obtain the same maximum likelihood estimates of principal components W :
[Equation]
Here, p ( x | z ) denotes the likelihood of observations x ∈ R d given the latent representation z ∈ R k , p ( z ) denotes the prior on z and q ( z | x ) denotes the estimated variational posterior distribution of z given an encoder parameterized by V and diagonal variances D . Both models estimate the same low dimensional representation of the data through z , and learn the same factorization of the covariance matrix through W .
While PPCA parameters are typically estimated through expectation maximization (10), linear VAEs are optimized by maximizing the Evidence Lower Bound (ELBO) given by
[Equation]
For linear VAEs with a Gaussian likelihood, the variational posterior distribution q ( z | x ) can be shown to analytically agree with the posterior distribution p ( z | x ) learned from PPCA (8).
However, deriving this connection for count-based likelihoods such as the multinomial distribution is complicated due to non-conjugacy issues (Appendix A).
This is a major obstacle for many biological applications; multiple works have shown the merits of incorporating count distributions explicitly into the model (33; 34; 35; 36).
Here, we provide directions for overcoming this issue.
M ETHODS
First, we will redeﬁne Multinomial PPCA with the ILR transform (21).
Then we will make the connection between Multinomial VAEs and Multinomial PPCA by leveraging insights from the Collapse-Uncollapse (CU) sampler (33).
We will then derive an algorithm to obtain the maximum a posteriori (MAP) estimate for the VAE parameters.
4. P ROBABILISTIC M ULTINOMIAL PCA
PPCA can be extended to multinomially distributed data with the following generative model:
[Equation]
Here W ∈ R d − 1 × k represents the PCA loading matrix, σ 2 is the variance,
Ψ ∈ R d × d − 1 is a ﬁxed contrast matrix whose columns sum to zero and φ is the softmax transform.
For a single sample, x ∈ N d are the observed d -dimensional counts, η ∈ R d − 1 are the latent logits and z ∈ R k is the latent representation.
The term φ ( Ψ η ) is distributed logistic normal, φ ( Ψ η ) ∼ LN
( W z , σ 2 I ) , as shown by Aitchison (37).
Furthermore, p ( x | z ) yields a multinomial logistic normal distribution, which is given by marginalizing out η in the following expression:
[Equation]
This integral is not tractable; as a result, this distribution does not have an analytically deﬁned expectation, variance or probability density function.
There have been multiple attempts to estimate the posterior distribution with MCMC (38; 39; 35; 40), but the complexity of this distribution requires a large number of samples, limiting the scalability of these methods.
Variational methods have been developed to estimate the logistic normal distribution, but due to conditional non-conjugacy, these methods often rely on approximations to the ELBO, further complicating estimation (41).
Recently, Silverman et al. (33) proposed to use a Laplace approximation to estimate the parameters of the multinomial logistic normal posterior distribution.
This approach relies on a two-stage optimization procedure on the factorized posterior distribution given by
[Equation]
If η can be directly estimated, then conditional independence can be used to factorize the posterior distribution.
Since the probability densities of the multinomial distribution and the normal distribution are both log-convex functions, a global optimum can be obtained for the multinomial logistic normal distribution (33) (Appendix B.3).
Furthermore, the multinomial distribution does not introduce additional local optima for estimating Multinomial PCA. Given this, in addition to recent evidence that a PPCA MAP estimator can be obtained from regularized linear autoencoders (17), we can design a new algorithm to obtain a Multinomial PPCA MAP estimator.
4. T HE ILR T RANSFORM ENFORCES IDENTIFIABLITY
The softmax function is a shift-invariant function, which introduces an identiﬁability issue that has been addressed by the compositional data analysis community (1; 42).
In order to remove the identiﬁability issues, an isomorphism between the logits η and the multinomial parameters must be maintained.
One commonly used solution is to use a degenerate softmax, also known as the inverse additive log-ratio (ALR) transform (1) (Appendix B.2).
Previous work has suggested that the isometric log-ratio transform (ILR) (21; 22) is more suitable for principal components analysis (Appendix B).
The ILR and inverse ILR are given as follows:
[Equation]
where Ψ ∈ R d × d − 1 is a basis such that Ψ T Ψ
= I d − 1 and ΨΨ T =
I d − 1 d 1 d × d .
A naive implementation of the ILR transform can be memory-intensive and computationally intensive for large d .
However, any orthonormal basis can be used to parameterize the ILR transform and some of these bases can be represented by binary trees (43; 44; 45).
A binary tree can be used to represent Ψ with O ( d log d ) elements, where the l th column vector of Ψ is given as follows:
[Equation]
where l indexes an internal node in the tree with left children r , right children s , nodes to the left k and nodes to the right t (46) (Figure S2). Due to rotation invariance, it doesn’t matter which tree is used to parameterize the ILR basis, but the choice of tree can inﬂuence the runtime of the ILR transform.
If a balanced binary tree is used, the memory requirements representing Ψ can be brought down from O ( d 2 ) to O
( d log d ) and can reduce the matrix vector multplication runtime from O ( d 2 ) to O
( d log d ) (See Appendix B.1).
This can speed up the matrix-vector multiplication operations by an order of magnitude for datasets with more than ten thousand dimensions.
4. M ULTINOMIAL V ARIATIONAL A UTOENCODER A RCHITECTURE
Our full Multinomial VAE model is given as follows:
[Equation]
where θ dec = { W , σ 2 } denotes the decoder parameters, θ enc =
{ F L , D } denotes the encoder parameters and µ ∈ R d − 1 is a bias parameter.
Here, q ( z | x ; θ enc ) denotes the variational posterior distribution of z given by the encoder represented as an L -layer dense neural network with appropriate activations.
This encoder is directly used to evaluate p ( η | z ; θ dec ) .
Furthermore, ﬂat priors are assumed for all variables except z .
It is important to note potentially challenging modeling issues when designing the encoder.
The ILR transform is not directly applicable to count data, since log(0) is undeﬁned.
A common approach to this problem is to introduce a pseudocount before applying a logarithm, which we will denote as (cid:102) log( x ) = log( x + 1) .
The choice of pseudocount is arbitrary and can introduce biases.
To alleviate this issue, we introduce the deep encoder neural network highlighted in Equation 10; we expect that the universal approximation theorem would apply here (47; 48) and that the accuracy of estimating the latent representation z will improve with more complex neural networks.
This is supported in our simulation benchmarks; more complex encoder architectures can better remove biases induced from the pseudocounts.
4. A LTERNATING MINI - BATCH OPTIMIZATION PROCEDURE
Given that our objective here is to obtain the MAP estimate of the VAE model parameters, the
VAE parameters θ = { F L , W , D , σ 2 } can be obtained by estimating the global maximum of the posterior distribution.
In the original CU sampler implementation, the parameter ( η 1 , . . . , η n ) is optimized across the entire dataset and then ( η 1 , . . .
, η n ) is ﬁxed in order to estimate the remaining parameters.
For large studies, this can be memory demanding, since ( η 1 , . . .
η n ) ∈ R n × d − 1 alone can scale to millions of parameters.
To scale this estimation procedure to large high-dimensional datasets we have devised a minibatched alternating minimization procedure. For a mini-batch X ( i ) = ( x (1) , . . . , x ( b ) ) of size b
Algorithm repeatfor
[Equation]
end for until convergence and corresponding latent variables
Z ( i ) =
( z (1) , . . .
, z ( b ) ) and H ( i ) =
( η (1) , . . .
, η ( b ) ) , the quantities log p ( X ( i ) | H
( i ) )
and log p ( Z ( i ) ) are given by Equation 1 and 3.
The variational posterior distribution of η given by q ( H ( i ) | X ( i ) ; θ ) = (cid:81) bj =1 q ( η ( j ) | x ( j ) ; θ ) can be obtained by marginalizing out z in Equations 8, 9 and 10 as follows:
[Equation]
[-whose covariance matrix W DW T +  
-]
The prior log p ( Z ( i ) ) is evaluated, given log p ( Z ( i ) )
=
log N
( Z ( i ) | 0 , I k ) and the latent encoding means given by Z ( i ) =
F L (cid:0) Ψ T ( (cid:102) log( X ( i ) ) − µ ) (cid:1) , where Z ( i ) is integrated out of Equation 11.
R ESULTS 5. M ULTINOMIAL PPCA AND VAE S AGREE IN S IMULATION B ENCHMARKS
To determine if the proposed multinomial variational autoencoder can recover principal components, we extended the benchmarks proposed in (18).
Here, we benchmarked our proposed analytical VAE The subspace distance is a measure of agreement between the ground truth correlations and the correlations estimated from W .
Procrustes measures the residual error after obtaining the best rotation and scaling factors to match W to the ground truth principal components.
See Appendix
C.5 for descriptions of the correlation, axis-alignment, subspace-distance and Procrustes metrics. to the stochastic VAEs (9) with a multinomial likelihood across multiple simulations. These methodologies were compared against the ground truth covariances of the multinomial logistic normal, in addition to the MAP and Hamiltonian Monte Carlo (HMC) estimates obtained from Stan (49).
The agreement between the ground truth principal components and the estimated principal components is measured by axis-alignment, subspace distance, correlation and Procrustes (50) (see Appendix C.4 for details). When ﬁtting against multinomial logistic normal distributed data with no zero counts, all of the ILRbased methodologies can accurately estimate the covariance matrix up to scale (Figure 1a-d). The analytical VAE MAP estimate and the posterior samples of HMC all have a correlation close to 1, suggesting a strong agreement between the ground truth covariance and the estimated covariances.
If the principal components were perfectly estimated, the axis-alignment, subspace distance and Procrustes metrics would all be close to zero.
While the subspace distance and Procrustes are notably close to zero, the axis-alignment metric is above 0.5, which would suggest disagreement between the ground-truth and the estimated principal component axes. Given that the posterior distribution of the axis-alignment metrics obtained from HMC overlaps with the analytical VAE, stochastic VAE and
MAP estimates, there is evidence that reducing this metric to zero is inherently difﬁcult for count data.
However, both the subspace distance and Procrustes metrics approach zero, supporting our claim that principal components can be estimated up to scale and rotation on fully observed count data.
Furthermore, our simulations (Figure 1i-k and Figure 2a-h) suggest that among the log-ratio transforms benchmarked, only the ILR transform can recover principal components, corroborating previous ﬁndings surrounding compositional PCA (42; 23).
Dealing with sparse data presents addi- tional challenges; the zeros in count data are indicative of missing data, which can further complicate the estimation of principal components. Our hypothesis that boosting the complexity of the encoder architecture would help alleviate issues with missing data is supported by benchmarks shown in Figure 2i-p. In both of the sparse datasets, none of the methods were able to achieve optimal accuracy across any of the benchmark metrics, but there is a clear advantage of utilizing multilayer encoders compared to single-layer encoders.
Across the simulation benchmarks, the analytical and stochastic VAEs have comparable performance, with discrepancies highlighted by the correlation metric in the dense and sparse benchmarks.
On the dense datasets, our proposed analytical VAE has better agreement with the ground truth covariance metric, whereas on the sparse datasets, it appears that the stochastic VAE has better agreement.
It is difﬁcult to explain the reason for the performance gap between our proposed analytical VAE and the stochastic VAE due the analytical intractability of the Multinomial VAE ELBO.
The challenge of accurately estimating an optimal H for each mini-batch could be one limiting factor affecting the performance of our proposed analytical VAE (Appendix B.4).
5. P RETRAINING M ULTINOMIAL VAE S ON A VERY LARGE COMPENDIUM OF MOUSE MICROBIOME OBSERVATIONS
To showcase this on real microbiome datasets, 11k mouse fecal pellet samples with a total of 8.8k features across 45 studies were analyzed.
We trained 2 models to evaluate the differences between the stochastic VAE and the analytical VAE.
We trained these models for 10k epochs; the models with the smallest validation error are reported here.
The details behind the full training procedure are given in Appendix C.3.
Visualization of the learned decoder weights shown in Figure 3 makes
T W . it clear that the analytical VAE and, to a lesser extent, the stochastic VAE are able to learn orthog- onal embeddings, since W T W approximates a diagonal matrix.
The orthogonality of the decoder weights is more apparent in our proposed analytical VAE than in the stochastic VAE decoder; there is a larger distinction between the diagonal elements and the off-diagonal elements. Unlike linear VAEs (8), the connection between the eigenvalues and the decoder weights is not as clear due to the asymmetry between the encoder and the decoder; scale identiﬁability of the decoder weights complicates the process of estimating the eigenvalues.
This is apparent in Figure 3; heatmaps of the W inner products are on vastly different scales.
One of the advantages of employing a pretraining strategy with VAEs is that the pretrained models can enable one-shot learning (51).
By estimating a low-dimensional representation from a large unlabeled high-dimensional dataset, fewer labeled samples for training downstream classiﬁers.
We showcase this property in two classiﬁcation benchmarks.
5. C LASSIFICATION AND CORRELATION BENCHMARKS
To construct benchmarks on real biological datasets, we analyzed 16S sequencing data obtained from a study conducted by Tripathi et al. investigating hypoxia that used mouse fecal pellets (52) and a study conducted by Shalapour et al. investigating cancer effects on mice (53).
Due to the connections between Multinomial VAEs and compositional PCA (20; 22), we expect that, if our VAE is performing well, we will see that the row Euclidean distances of Ψ W ∈ R d × k correspond to Lovell’s un-normalized proportionality metric (54; 23) given by
[Equation]
[Equation]
where x refers to the observed proportions and i and j refer to the two features being compared.
The proportionality metric has a simple interpretation; if this variance is small, that implies that the two features are highly co-occurring. Here, only a relative relationship between the proportionality metric and the VAE decoder row distances can be stated due to scale identiﬁability issues inherent in Multinomial PPCA.
Proportionality has been shown to be a more reliable metric for compositional data compared to standard correlation metrics such as Pearson and Spearman in single cell RNA sequencing and microbiome studies (3; 55; 56; 2), supporting the notion that our proportional comparison of our VAE decoder to Lovell’s un-normalized proportionality metric will provide a sound perspective, as well as a means for interpreting the weights of the VAE.
We compared the estimated VAE embeddings with the proportionality metric to determine if this relationship holds empirically. The pairwise decoder distances are compared to the pairwise proportionality metrics of 300 selected features (Table 1 and Figure S5).
The learned VAE representations are benchmarked using K-nearest neighbors (KNN) classiﬁcation. KNN is applied to the learned VAE encodings across the two microbiome datasets to determine how well the classiﬁers can identify the mice that were induced with hypoxia (52) and classify mice based on their experimental group (53).
These classiﬁcation models are compared against two baseline models, namely KNN trained on LDA topic proportions (31), and KNN trained on raw counts.
From the classiﬁcation benchmarks illustrated in Table 1, we can see that there is a large margin between the VAE models and the baseline models across the measured metrics.
Both VAEs had a higher F1 score and AUC, suggesting that the learned representations can better separate the experimental groups. Part of this performance boost could be attributed to the differences between distance metrics.
When the ILR transform is utilized, the Euclidean distance between the latent encodings approximates the Aitchison distance between proportions (57; 23), which measures relative differences instead of absolute differences. In the compositional data analysis literature, it has been argued that proportions only contain relative information, making relative distance metrics more suitable for analysis.
Indeed, previous work has shown that the application of the Aitchison distance has provided a substantial boost in KNN classiﬁcation accuracy compared to the use of other distance metrics in the microbiome ﬁeld (58; 59).
Furthermore, the average negative log-likelihood score (NLL) was lower for both of the VAE models than for LDA, suggesting that the VAE models generalized better on held out data than LDA. This decrease in the predictive log-likelihood could be attributed to the increased ﬂexibility of the covariance matrix in the logistic normal distribution compared to the Dirichlet distribution. [Table]
al [Table] Tripathi et al. [Table]
There are a couple of notable discrepancies between the two Multinomial VAE models.
The stochastic VAE appears to have superior classiﬁcation performance and lower reconstruction error compared to the analytical VAE. The exception to this is the HCC dataset (53), where the analytical
VAE marginally outperforms the stochastic VAE in terms of AUC. However, the analytical VAE can learn more apparent orthogonal embeddings (Figure 3) and better agrees with Lovell’s proportionality metric (Table 1).
Only the analytical VAE was able to recover the log-linear relations between
Lovell’s proportionality and the VAE embedding distances, suggesting that it can more accurately learn biologically relevant correlations (Figure S5).
C ONCLUSION
Prior work dealing with count distributed data have been largely limited to conjugate priors due to their tractability. However, these design choices can lead to models with lower explanatory power due to the rigid structure of their covariances. Disregarding the simplicial sample space and performing Pearson correlations on raw count data is common practice in scientiﬁc applications, but is a source of reproducibility issues (2). Due to the negative bias induced from the covariance on the observed count data, the estimated covariances will not agree with the ground truth covariances in the system of interest, a fact noted by Pearson in 1897 (60).
Adapting the logistic normal distribution in place of conventional conjugate priors provides a means to remedy the issue of inferring correlations on compositional data.
The covariance matrix on ILRtransformed data can be interpreted using Lovell’s proportionality metric, and can serve as a replacement for pairwise correlations.
Since log-ratios are scale-invariant, the dependence on the total counts disappears, which is critical for assuring agreement between the relative quantities measured through count data and the absolute quantities in the system that aren’t directly observable.
While there has much motivation to employ the logistic normal distribution, its application has been limited due to the challenge of estimating these distributions.
Here, we have provided a means to accurately and efﬁciently estimate these distributions and covariance matrices up to scale using the ILR transform.
To this end, we have provided a proof-of-concept that Multinomial VAEs can learn the principal components obtained from Multinomial PPCA. We have shown that ﬁtting low-rank approximations using Multinomial VAEs can provide more explanatory representations than LDA while providing a means to obtain interpretable correlations.
These methods can be used in one-shot learning and transfer learning settings, requiring fewer labeled samples for training and allowing for the use of pretrained models to extract features from smaller datasets.
Given the vast number of scientiﬁc disciplines that collect compositional data, we anticipate that models such as Multinomial VAEs will have an impact on many scientiﬁc applications.
C ODE A VAILABILITY
All of our software and analyses can be found at https://zenodo.org/record/4062719.X3apeZpME5k
A CKNOWLEDGEMENTS
We want to acknowledge Pytorch (61), Pytorch-Lightning (62), the Biom format (63), Matplotlib (64), Scipy (65), Numpy (66) and Stan (49) for providing the software foundation that this work was built upon. Appendices A C HALLENGES IN D ERIVATION OF AN ANALYTICAL M ULTINOMIAL VAE ELBO Recall that the generative model for Multinomial PPCA is given as follows: [Equation]
[Equation]
With this in mind, we wish to estimate variational distributions q ( η, z | x ) =
q ( η | z ) q
( z | x ) to approximate the posterior p ( η, z | x ) .
These variational distributions can both be chosen to be normal distributions as follows:
[Equation]
Noting that z ∼ N ( V Ψ T ( (cid:102) log( x ) − µ ) , D ) , q ( η | x ) can be derived from q ( z | x ) as follows:
[Equation]
To ﬁne-tune these variational distributions to approximate the posterior distribution
[Equation]
we can minimize the following KL divergence:
[Equation]
Since we cannot optimize this quantity directly
[Equation]
we opt instead to maximize the evidence lower bound (ELBO) given by
[Equation]
We can partition this lower bound into three parts, given as follows:
[Equation]
Since Mult( x | p ) ∝ d (cid:80) i
=1 x i log p
i , the ﬁrst term is given by
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
Estimating the above integral is equivalent to estimating the expectation of a logistic normal distribution, which does not have an analytical solution (37).
As a result, the analytical ELBO for the
Multinomial VAE is intractable.
B T HE ILR TRANSFORM
As outlined in Equation 5, any orthogonal basis Ψ can be used to perform the ILR transform.
However, there are a select few bases that can be represented by a binary tree.
To see how an orthogonal basis can be constructed from a binary tree, consider the illustration in Figure S1.
Here the rows of the matrix in Figure S1 are orthogonal and the resulting product η = Ψ T log x yields
[Equation]
The contrast matrix Ψ can be forced to be orthonormal such that Ψ T Ψ
=
I d − 1 , as highlighted in
Equation 6.
Furthermore, this construction can be scaled to large binary trees as shown in Figure
S2.
Here, η l represents the log-ratios at the internal node l , given by
[Equation]
The runtime of this operation is further discussed in Appendix B.1.
B. T HE R UNTIME OF THE ILR TRANSFORM
The naive runtime of the ILR transform of a single sample x ∈ R d is O ( d 2 ) due to the running time of dense matrix-vector multiplication.
As shown in Equation 5, the ILR transform can also be represented by log-linear transformation with a contrast matrix.
The binary tree can be used to represent a contrast matrix, as discussed in (46).
If the binary tree is balanced, each row of Ψ will have O (log d ) non-zero elements, since the tree has a height of O (log d ) .
Given that there are d rows, the matrix-vector multiplication behind the inverse ILR transform Ψ η can be done in O ( d log d ) .
For the same reason, the ILR transform has a runtime of O ( d log d ) .
B. S HIFT I NVARIANCE OF THE S OFTMAX T RANSFORM
There are two different scale identiﬁability issues.
The softmax transform that is commonly used is shift invariant, where
[Equation]
This shift invariance will cause an identiﬁability issue when identifying the decoder matrix W .
The inverse ALR transform resolves this issue by enforcing an isomorphism; one of the coordinates is set to zero as follows:
[Equation]
One of the implicit constructions of this transform is that the resulting contrast matrix is not orthogonal (42).
Furthermore, because the ALR transform is not isometric, the resulting Euclidean distances in z will not approximate the Aitchison distance.
The ILR transform provides the best of both worlds, enforcing both isometry and isomorphism.
Furthermore, there is a tight connection between the ILR transform and compositional PCA that is further discussed in the next section.
B. G LOBAL L OG -C ONVEXITY OF THE M ULTINOMIAL L OGISTIC N ORMAL D ISTRIBUTION
Since the multinomial logistic normal distribution is difﬁcult to directly evaluate, it is challenging to make statements regarding its maximum likelihood estimator.
With the posterior factorization highlighted in Equation 4, we can make concrete statements about the original posterior factors.
The log probability density of the multinomial distribution p ( x | η ) can be written as
[Equation]
where m = (cid:80) di
=1 x i is the total number of counts and the functions g ( η )
i =
(Ψ η ) i , T ( x ) i
= x i , and
A ( η ) = log (cid:18) (cid:80) dj
=1 exp( Ψ η ) i (cid:19) are the natural parameters of the exponential family distribution.
The Hessian of log p ( x | η ) is given by
[Equation]
Since A ( η ) is strictly convex, log p ( x | η ) is also strictly convex.
Similarly, the log probability density of the multivariate Gaussian distribution p ( z | η ) =
N ( µ, σ ) is also strictly convex with respect to µ and Σ .
Since the multinomial logistic normal distribution can be written as the sum of log p ( z | η ) and log p ( x | η ) , it is also strictly convex.
Therefore, there must be a unique optimal estimate for µ and Σ .
B. M ULTINOMIAL V ARIATIONAL A UTOENCODER E STIMATION
The covariance matrix W DW T + σ 2 I d − 1 ∈ R d − 1 × d − 1 in Equation 11 can be efﬁciently inverted using the Woodbury identity (67; 68).
The prior log p ( Z ( i ) ) is evaluated given log p
( Z ( i ) )
=
N ( Z ( i ) | 0 , I k ) , where Z ( i ) =
F L (cid:0) Ψ T ( (cid:102) log( X ( i ) ) − µ ) (cid:1) .
The optimal ˆ H ( i ) that maximizes log p ( X ( i ) | H
( i ) ) is obtained through gradient descent optimization.
Once ˆ H ( i ) is obtained, the remaining VAE parameters θ can be estimated via gradient descent optimization.
Like the EM algorithm proposed in (69), this procedure alternates between estimating H ( i ) and θ and repeats until convergence.
For a single-layer linear encoder, this optimization procedure will eventually reach the global maxima with respect to H ( i ) and the multivariate normal mean and covariance in q ( H ( i ) | θ , X ( i ) ) , due to the log-convexity of the multinomial and normal distributions (Appendix B.3).
On fully observed data, W and D can be estimated up to rotation and scale (Appendix B.6).
This is not guaranteed for sparse count data, but using simulations, we can empirically show that increasing the complexity encoder architectures can help.
Since ˆ H ( i ) needs to be accurately estimated before optimizing F L , W , σ , and D , multiple gradient descent steps are required for a given mini-batch.
In practice, obtaining the optimal ˆ H ( i ) for a given mini-batch i may require hundreds of gradient descent steps from a random initialization.
{+Since H ( i ) is used to approximate the multinomial parameters describing the observed counts x , we can initialize H ( i ) with ˆ H ( i ) 
+}
= Ψ T (cid:102) log( X ( i ) ) .
In practice, this can greatly reduce the number of gradient descent updates needed per mini-batch.
B. T HE C ONNECTION BETWEEN M ULTINOMIAL VAE S AND C OMPOSITIONAL PCA
With singular value decomposition, the resulting factors can be used to approximate the row and column distances.
For a singular value decomposition given by X = USV T , row distances and column distances can be approximated as follows:
[Equation]
This relationship also yields a connection to the row and column covariances; the row covariances are given by X T X =
US 2 U and the column covariances are given by XX T =
V S 2 V .
A similar relationship applies to compositional PCA, except the singular value decomposition is applied to CLR-transformed values of X .
The CLR transform is given by
[Equation]
The inverse CLR transform is equivalent to the softmax transform φ ( x ) .
The Euclidean distance on
CLR and ILR transformed values is given by the Aitchison distance (57).
As a result, the row and column distances are given by the Aitchison distance.
The covariance matrix of CLR-transformed values is singular, and as a result, the resulting singular value decomposition will have a full rank of d − 1 .
As shown by Egozcue et al.
(23), the top d − 1 components of the singular value decomposition are all in ILR coordinates.
This observation cements the connection between compositional PCA and our proposed Multinomial VAE; since all of the singular value components can be represented in ILR coordinates, this theoretically justiﬁes the use of the ILR transform within our proposed
Multinomial VAE.
Due to this connection, the distances between the learned VAE representations z across pairs of samples should be proportional to the Aitchison distance.
Furthermore, the distances between rows of Ψ W would also be given by the Aitchison distance, which is equivalent to Lovell’s un-normalized proportionality constant (23).
B. L ACK OF SCALE IDENTIFIABILITY OF PPCA
As discussed in (8), VAEs will have the same identiﬁability issues that PPCA has; namely, for a diagonal matrix A ∈ R k × k , the following equivalences hold:
[Equation]
As a result, the decoder weights W can only be identiﬁed up to scale.
Furthermore, the Multinomial
VAE architecture is asymmetric and, as a result, the Tranpose Theorem proposed by (17) is not as readily applicable.
The lack of symmetry between W and V complicates the process of obtaining eigenvalues from W .
C S IMULATION AND M ICROBIOME B ENCHMARKS C. S IMULATION DETAILS
In the dense simulation highlighted in Figure 1, there were 100 features, 1000 samples and 1 million counts per sample.
These counts were drawn from a multinomial logistic normal distribution whose covariance matrix had a rank of 10.
The sparse dataset highlighted in Figure 2a-h containeddimensions and 1000 samples, each of which contained 100 counts that were sampled from a multinomial logistic normal distribution whose covariance matrix had a rank of 10.
The sparse dataset highlighted in Figure 2i-p contained 5000 dimensions and 10000 samples, each of which contained only 1000 counts total.
These counts were sampled from a multinomial logistic normal distribution whose covariance matrix had a rank of 128.
The encoder architecture used for the stochastic and analytical VAEs consisted of fully connected dense layers with interwoven softplus activation functions.
The intermediate encoder layers all have the same input and output dimension as the number of latent dimensions speciﬁed in the model.
The same architecture was used for both the stochastic VAEs and the analytical VAEs.
A schematic of the architecture is shown in Figure S3.
Multinomial PCA was ﬁtted using both penalized likelihood estimation and Hamiltonian Monte
Carlo using Stan (Appendix C.4).
The initialization of these models was determined by the es- timated factor loadings from the analytical VAE.
This is particularly critical for ﬁtting HMC on high-dimensional datasets.
We were not able to run Stan on the sparse datasets.
C. E XPERIMENTAL DATASET DETAILS
The full 8.8k mouse pellet dataset was retrieved from qiita (70) using redbiom (71).
In the hypoxia study conducted by Tripathi et al., mice were placed in simulated conditions to induce intermittent hypoxic/hypercapnic (IHH) stress, where the mouse’s oxygen supply was reduced and its CO 2 supply was increased.
The goal of this experiment was to detect the gut microbiome differences between the IHH mice and the control mice.
We will use this information as a classiﬁcation benchmark to determine if the experimental conditions the mice were placed under can be predicted from 16S sequencing counts.
In total, there were 48 mice; 24 of these mice were placed in these
IHH conditions and the remaining 24 mice served as controls.
Fecal samples were collected twice a week for 6 weeks, resulting in 579 samples total.
There were 5775 microbial taxa that were detected in this dataset.
The goal of the hepatocellular carcinoma (HCC) study conducted by Shalapour was to investigate the interplay between immunity, diet and carcinoma.
In total, there were 478 mice with 52 phenotypes, split according to diet, immunity status and whether or not the mice were induced with HCC.
Our classiﬁcation objective is to predict the mouse phenotype from the 16S sequencing counts obtained from their fecal pellets.
In total, there was one fecal pellet collected from each mouse and there were 2794 microbial taxa detected across these samples.
This benchmark is more challenging than the
IHH benchmark due to the extreme class imbalance; some phenotypes contain upwards of 30 mice and other contain as few as 2 mice (Figure S4).
The resulting classiﬁcation benchmarks are shown in Table 1.
C. M ICROBIOME VAE TRAINING DETAILS
For the full 11k fecal pellet dataset, we focused on samples that were processed using closedreference OTU picking using vsearch (72).
This dataset was split into a 80/10/10 train/validation/test split in order measure out-of-distribution generalizability.
We trained our proposed analytical VAE and the stochastic VAE with 128 latent dimensions and a batch size of 1000 samples for 10000 epochs, with a learning rate of 10 − 3 and 100 gradient descent steps per mini-batch.
Checkpoints were recorded every epoch; the models with the best validation error are reported.
Cosine annealing with warm restarts was used as a learning rate scheduler, with the intention of easily escaping saddle points during optimization.
C. M ETRICS FOR E VALUTING ESTIMATED DECODER WEIGHTS
To measure the agreement between the ground truth simulations and the estimated decoder weights
W , three different metrics were evaluated, namely, axis alignment and subspace distance, as proposed in (18), in addition to Procrustes analysis (50).
Given the decoder weights W and the ground truth principal components U , these metrics are deﬁned as follows:
Axis alignment
[Equation]
This metric is ultimately a measure of the average cosine distance between the ground truth eigenvectors and the estimated decoder weights.
Subspace distance
[Equation]
where W ∗ denotes the left singular vectors of W , given by W = W ∗ Λ V .
Since UU T yields the ground truth correlations and W ∗ W T ∗ yields the estimated correlations, this metric can be interpreted as a measure of agreement between the ground truth correlations and the estimated correlations.
Procrustes Analysis
[Equation]
for A, R ∈ R k × k , where A is a diagonal matrix and R is a rotation matrix.
Prior to evaluating this metric, both U and W are standardized such that Tr( UU T ) =
Tr( W W T )
= 1 and W is centered around the origin.
Correlation
[Equation]
where A U ( i, j ) = (cid:107)| u i − u j (cid:107)| 2 and A W ( i, j ) =
(cid:107)| w i − w j
(cid:107)| 2 denote the pairwise distances of U and W .
The pairwise distances are rotation invariant.
Furthermore, the measure of the correlation is agnostic to scale, and will ignore the eigenvalue scale identiﬁability issues highlighted in the main text.
C. S TAN M ULTINOMIAL PCA IMPLEMENTATION data {
{+int < lower =0 > N; / / number of samples int < lower =
+}
>
D; / / number of dimensions int < lower =0 > K; / / number of l a t e n t dimensions matrix [D − 1, D]
Psi ; / / Orthonormal b a s i s i n t y
[N, D] ; / / observed counts } parameters { matrix [N, D − 1]
e t a ; / / i l r transformed abundances matrix [D − 1, K] W;
real < lower =0 > sigma ; } transformed parameters { matrix [D − 1, D − 1] Sigma ; matrix [D − 1, D − 1] I ; vector [D − 1] z ;
I = diag matrix ( rep vector ( 1 . 0 , D − 1));
Sigma =
W ∗ W’ + square ( sigma ) ∗ I ; z = rep vector (0 , D −
} model { / / g e n e r a t i n g counts
for ( n in 1:N) { e t a [ n ] ˜ multi normal ( z , Sigma ) ; y [ n ] ˜ multinomial ( softmax ( t o v e c t
o r ( e t
a [ n ] ∗ Psi ) ) )
[Table]
All Test Samples
[Table]
Due to the relationship between squared Euclidean distance and proportionality given in Equation
[Table]
we expect the proportionality metric to exhibit a log-linear relationship between the Euclidean distances obtained from the VAE embeddings.
Indeed, this is apparent with the analytical VAE embeddings in Shalapour et al., and to a lesser extent in Tripathi et al., as shown in Figure S5.