Interpretable factor models of single-cell RNA-seq via variational autoencoders

1The study of the regulatory architecture of cells has revealed numerous examples of co-regulation of transcription of large numbers of genes, and this has been used to link the organization of cells to their distinct functions in response to developmental or external stimuli. While studies of cells in bulk have led to interesting population-level insights about the relationships between genes, the study of individual cells via single-cell RNA-seq has led to questions about the dependence of relationships between genes on cell type.

2Principal component analysis (PCA) is a popular linear method for dimensionality reduction in single-cell RNA-seq. As a result of its efficiency, PCA has been used for exploratory data analysis to visualize the structure of high-dimensional data in two or three dimensions. PCA also provides a linear model of the data; a key feature of the method that can be used for prediction. In the case of single-cell RNA-seq, datapoints correspond to cells and the coordinates of each cell represent the gene expression levels for each gene in the transcriptome. Thus, PCA can be used to study structured variation between cells by revealing differences along axes of greatest variation. In PCA, linear weight parameters (loadings) are used to predict gene expression in each cell, conditional on the latent variables (coordinates) per cell. The loadings corresponding to the principal component axes can be interpreted as ‘meta-genes’: sets of genes which tend to be expressed together. Thus, PCA of gene expression provides a formal mathematical framework for studying the biological idea of ‘gene programs’ by simultaneously explaining structured variation between cells and genes.

3While PCA is easy to use and is often applied to single-cell RNA-seq data, the method has some drawbacks. PCA models data as arising from a continuous multivariate Gaussian distribution, and thus optimizes a Gaussian likelihood. This model assumption is at odds with the count data measured in single-cell RNA-seq, and leads to interpretation problems. To address this issue, a number of methods define factor methods tailored to single-cell transcriptomics data. For example, ZINB-WaVE defines a linear factor model where gene weights are parameters, cell factor values are latent variables and data arise from a zero-inflated negative binomial distribution. However, as single-cell transcriptomics datasets have grown in size to hundreds of thousands of observations, efficiency and scalability considerations have become paramount and inference with parametric models can be intractable. To address scalability requirements, new methods based on variational autoencoders have been developed; these leverage the large amounts of available data to learn non-linear maps, and crucially scale well thanks to efficient algorithms for inference that leverage the structure of autoencoders.

4Autoencoders consist of a pair of functions: a representation function and a reconstruction function, which are typically parameterized as neural networks. The two autoencoder functions can be seen as a non-linear generalization of the two projections associated with PCA. By optimizing the pair of neural networks, efficient low-dimensional representations of data can be identified. A variational autoencoder (VAE) uses a similar strategy but with latent variable models. Each datapoint is represented by a set of latent variables which can be decoded by neural networks to produce parameters for a probability distribution, thus defining a generative model. To infer the latent variable values (the representation), a neural network is used to find per-datapoint parameters for a probability distribution in the representation space. This defines an ‘inference model’ which attempts to approximate the posterior distribution of the latent variables given the observed data with a variational distribution.

5Inference using VAEs scales to arbitrarily large data since mini-batches of data can be used to train the parameters for both the inference model and the decoder function. However despite these efficiency advantages, the representations inferred with VAEs are not directly interpretable. While efforts have been made to develop interpretable VAE’s, the difficulty in interpreting VAE representations continues to be a major drawback of VAE’s. We show that using a flexible non-linear inference model along with a linear reconstruction function makes it possible to benefit from the efficiency of VAEs, while retaining the interpretability provided by factor models. Specifically, by adapting the method of scVI, we demonstrate a scalable approach to learning a latent representation of single-cell RNA-seq data, that identifies the relationship between cell representation coordinates and gene weights via a factor model. Our approach results in a tradeoff: whereas typically autoencoder models are designed with the same network topology in the inference functions and the reconstruction functions, what we propose is a restricted reconstruction function that leads to an increase in reconstruction error. However, by virtue of being linear, our reconstruction function provides an interpretable link between gene programs and cellular molecular phenotypes (Fig. 1a).

6The generative model of scVI, when data are from a single batch and zero-inflation is deactivated, is

7In this model, is a random variable for the exposure or count depth of a cell, with priors and. The random variable provides a -dimensional representation of cells. The parameter represents the overdispersion of a gene, and the Gamma is parameterized by its shape and mean. We replace the neural network with a linear function:

8This way the expression level of a gene in a cell is affected by the weights depending on the coordinate of a cell, giving a direct link between cell representation and gene expression.

9We also considered a variation of this model where the latent variables are distributed as a logistic normal (ln) distribution. In this case, each has positive values and sums to, making it similar to semi-non-negative matrix factorization. Such a model adds a further layer of interpretability; the cells are embedded in a simplex, where the nodes of the simplex represent archetypal cell types. In addition, we investigated the effect of applying a batch-norm transformation of the linearly decoded parameters.

10To explore the potential for interpretability in the VAE framework, we implemented a linearly decoded variational autoencoder (LDVAE) in scVI. The model was applied to two datasets of single-cell RNA-sequencing from a large number of developing mouse embryos in different stages of development. The first dataset consists of 125 775 cells from 411 mouse embryos undergoing gastrulation measured using the commercial 10× Genomics platform and sequenced relatively deeply (11% non-zero values).

11A comparison of the VAE with the LDVAE methods showed that VAE indeed has a smaller reconstruction error than the LDVAE methods (Fig. 1b). Among the LDVAE method variants, using a normal latent distribution and batch norm has the smaller reconstruction error (on held-out data). Between the LDVAE models with ln distributed latent space, the comparison was inconclusive for the ) data, but batch norm performed better for the ) data (Supplementary Fig. S1). With either VAE or LDVAE, the representation can be used to learn which cells are similar to each other and can be used for clustering. For example, erythroid cells, extraembryonic endoderm cells and epiblast cells annotated by the original authors can be separated by factors and. However, the axes of representation learned by the LDVAE model can be directly related to axes of co-expressed genes (Fig. 1c). For example, variation along the axis is related to simultaneous variation in expression of Pou5f1 and Tdgf1, two genes important for epiblast development. Variation along the axis is related to co-variation in beta globin (Hbb) genes which are key components of erythroid cells. Additionally, variation between epiblast and erythroid cells along is orthogonal to variation between epiblast and extraembryonic endoderm cells along, two independent lineages in embryonic development.

12While the ln latent distribution results in higher reconstruction error than the normal distribution, it has benefits for interpretation. Since a factor is restricted to non-negative values, genes with negative weights can only decrease in expression as a function of. This way cells using a particular regulatory program can more effectively be grouped to specific factors (Supplementary Fig. S2). We also found that using batch-norm transformation improved model performance.

13The learned representations from the different models can be compared by investigating the covariance matrix (where is a centered and scaled version of ). This illustrates that LDVAE learns representations with fewer covarying factors (Supplementary Fig. S3). Unlike linear methods, the VAE is not constrained by covarying factors since the non-linear neural network can produce vastly different gene expressions along a linear path in the representation. Comparing the proposed alternative LDVAE models, using a normal latent distribution induces less correlation between factors. By performing eigen decomposition on a covariance matrix the proportion of variance explained by each factor can be quantified. This allows ordering of factors which can be used to identify the regulatory programs with the most variation across the dataset. It also illustrates the simplicial structure of ln distributed latent variables since one factor is always linearly dependent on the other factors (Supplementary Fig. S3).

14The second dataset consists of 1 949 131 cells from 61 embryos in total using the sci-RNA-seq method at shallow sequencing (2% non-zero values). This dataset is, to our knowledge, the largest scRNA-seq study published to date. To illustrate the scalability of our model, we fit a 10-dimensional LDVAE to the data which allows identification of cells similar to each other and for the determination of covarying genes (Supplementary Fig. S4).

15Cells were also subsampled to different numbers before fitting LDVAE models. We found that inference runs in linear time, with 5 s per 1000 cells to reach 10 epochs using a CPU (Intel Core i7-7800X). Using a consumer-grade GPU (NVIDIA GeForce RTX 2070), inference only requires 2 s per 1000 cells to reach 10 epochs, with a total time of less than an hour for the full dataset. The inference times did not depend on the sparsity of the data, as the neural architecture is fixed, and operations are invariant to observed values (Supplementary Fig. 5a). Investigating the reconstruction error curves per epoch, the models converged after 2–3 epochs for datasets larger than 100 000 cells (Supplementary Fig. 5b). Determining a minimal number of epochs is a difficult general problem, but our results suggest a rule of thumb of ‘1 million divided by the number of cells in the dataset’ epochs for first pass analysis.

16Jupyter notebooks to produce the results are available at https://github.com/pachterlab/SGYP_2019 as well as Figshare at https://doi.org/10.6084/m9.figshare.11725920.v1. For convenience, the embryo data from ) and ) are also available in an H5AD object on the Figshare accession and on Google Cloud Storage at gs://h5ad/2019-02-Pijuan-Sala-et-al-Nature/pijuan_sala_atlas.h5ad and gs://h5ad/2019-02-Cao-et-al-Nature/cao_atlas.h5ad. A general tutorial on how to use the LDVAE model is available in the scVI Github repository at https://github.com/YosefLab/scVI/blob/master/tests/notebooks/linear_decoder.ipynb.

17Our results show that interpretable non-Gaussian factor models can be linked to variational autoencoders to enable interpretable, efficient and multivariate analysis of large datasets. This is useful for the investigation of gene co-expression in large scRNA-seq datasets, and the approach we have outlined should be applicable in other settings where interpretability is paramount.