2.1 VeloVAE allows bayesian inference of cell times, cell states, and rate parameters

VeloVAE uses a Bayesian model for RNA velocity inference. We assume that we are given individual scRNA-seq profiles that measure the amounts of spliced (s) and unspliced (u) transcripts at single moments of a developmental process. Our goal is to use these observations to simultaneously infer the posterior distributions of underlying latent variables that generated the data: cell time (t), cell state (c), and rate parameters () describing the biochemical kinetics of gene expression. Our key modeling assumption is that the observed time-varying (u(t),s(t)) levels are related by a system of two ordinary differential equations (Fig 1A). As with previous RNA velocity approaches [9,10], these ODEs capture the simple insight that a gene must first be transcribed as nascent mRNA, then spliced into mature mRNA, and then subsequently degraded (Fig 1A). However, we make two important changes: (1) rather than assuming a single fixed transcription rate parameter for each gene across all cells, we allow each cell to have its own transcription rate for each gene; (2) gene counts are generated by a mixture of two ODE models. The first change removes the need for discrete induction and repression phases and models continuous changes in transcription rates, such as transcriptional boosts [16] and cell fate bifurcations. The second change considers two generative processes: (1) a gene is induced and may be repressed later; (2) a gene is repressed from the very beginning. Each gene has a probability of being generated by each of the two processes, which can be represented by a categorical random variable . This accounts for genes that only undergo repression and turns out to be important for the correctness of time reconstruction in certain dataset such as the erythroid lineage. Note that, unlike scVelo and VeloVI, which place each gene on a separate time scale, the cell time parameter t is shared across all genes within a cell.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 1. VeloVAE model.

(A) Differential equation model of transcription used by VeloVAE. Nascent mRNA molecules are transcribed at gene- and cell-specific rate . Next, nascent mRNA is spliced into mature mRNA at gene-specific rate . Finally, mature mRNA is degraded at gene-specific rate . (B) Graphical model for latent variable inference and data generative processes. Cell time and state, as well as gene-wise mixture weight, are treated as latent variables, which are inferred using variational Bayes. Note that, unlike most previous approaches, latent time (and state) is shared across all genes within each cell. The latent cell state lies on a smooth low-dimensional manifold representing cell development. Each cell state has gene-specific transcription rates . RNA count data are generated based on mixture of two ODE systems at the gene level with known analytical solutions. (C) VAE architecture. The encoder neural network learns estimates the posterior distribution of the latent variables, while the decoder learns a mapping from cell state to cell-wise transcriptional rates and simulates the generative process from an ODE system.

https://doi.org/10.1371/journal.pcbi.1014060.g001

The VeloVAE model can be viewed from either an inference or a generative perspective (Fig 1B). From an inference perspective, if we know (u(t),s(t)) values for cells, we can infer something about the parameters, (t,c,w,,), that generated them (Fig 1B, left). Each particular location c in cell state space has an associated transcription rate for each gene; nearby cell state space locations will have similar transcription rates (Fig 1B, middle). From a generative perspective, if we know the (t,c,w,,) parameters for a cell, we can predict the distribution of their (u,s) values (Fig 1B, right). We can also incorporate prior information about the latent variables and rate parameters. If cell capture times are known (e.g., if cells were isolated separately on days 7 and 14), we can use the capture times as an informative prior for cell time.

To fit this statistical model on real data, we use autoencoding variational Bayes, a powerful statistical inference method in which neural networks approximate the posterior distribution of latent variables that may be nonlinearly related to observed data. Intuitively, autoencoding variational Bayes jointly trains the inference and generative models shown in Fig 1B, so that after training we can infer latent variables given observed data or predict new data given values of the latent variables. VeloVAE implements the inference model of Fig 1B using a neural network that takes (u,s) values as input and outputs the posterior distribution of cell time and cell state parameters (Fig 1C). VeloVAE implements the generative model of Fig 1B using a neural network that predicts the gene-specific transcription rates for each cell from the cell’s time and state values (Fig 1C). The (u,s) values for each cell can then be predicted using the analytical solution to the ODE, which describes how spliced and unspliced counts vary over time (Fig 1C). We previously described the underlying theory and motivation of this approach in a machine learning conference paper [23]. Intuitively, VeloVAE is like an autoencoder whose decoder network has been replaced with the solution to a differential equation. Our previous work assumes unspliced and spliced counts as continuous-valued functions of time, which is the first moment of a stochastic process [16] also known as the chemical master equation (CME). Thus, we investigated the discrete nature of RNA velocity by proposing a discrete model. First, we assume u and s are Poisson random variables at time t = 0. Because the reactions are all uni-molecular, we apply the theorem by Jahnke et al. [24] and reach a conclusion that both u(t) and s(t) are Poisson random variables whose mean is given by the solution to the continuous-valued ODE. Thus, we can fit the VeloVAE model with minimal change to the model architecture and simply use a Poisson likelihood to infer the latent variables and fit rate parameters.

In addition to the discrete model, we also extend previously developed theory to a binary mixture of ODE models at the gene level and also take advantage of the little-known fact that autoencoding variational Bayes allows inference of posterior distributions for parameters in the generative model (decoder network). This allows us to infer distributional estimates of the ODE rate parameters as well.

We train the model by maximizing the evidence lower bound of the marginal likelihood using mini-batch stochastic gradient descent [25]. Crucially, cells are loaded in batches during training, which means that not every cell must be used in each training iteration, dramatically decreasing both the time and memory requirements.

2.2 VeloVAE significantly improves model fit and latent time accuracy

We evaluated our method on a variety of real scRNA-seq datasets of different sizes and complexity [26–32] and compared our results with scVelo [10], UniTVelo [17], DeepVelo [18], VeloVI [20] and cellDancer [19]. UniTVelo did a comprehensive benchmarking tests on a number of other datasets [8,9,33–37]. Hence, we included them in our evaluation. In addition to continuous-count models, we also evaluated the discrete-count version of VeloVAE and PyroVelocity [21], since these are the only two methods directly modeling raw count numbers. We evaluated all models in terms of data reconstruction, how accurately they recover latent time and how well velocity flows match the biological truth. We also directly looked into each individual genes and compare the fitting qualitatively. Our results show that VeloVAE achieves comparable or better performance across all performance metrics, while existing methods are less versatile (Figs 2A, 2B, and S1).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2. VeloVAE significantly improves model fit and latent time accuracy and models complex gene expression dynamics (A)-(C) Quantitative performance comparison.

We report the correlation between inferred latent time and true capture time (when available) (A), mean squared error (MSE) between the observed and predicted (u,s) counts for continuous-count models (B) and gene log likelihood for discrete-count models (C) across 14 benchmarking datasets in a boxplot. (D)-(F) Examples of individual genes fit by different RNA velocity methods for a late-induced gene (D) and early-repressed gene (E) in pancreas and a gene with branching dynamics in subsampled mouse brain (F). Gene fits are shown for both u and s values, with inferred latent time on the x-axis. Fitted values from scVelo and UniTVelo are shown as lines, with observed data values shown as points. Only fitted values are shown for the other methods, because the fit is a point cloud (rather than a line) that would completely cover the observed data. Points arve colored by published cluster assignments, with the same colors as in the corresponding UMAP plots. The arrow length and direction indicate the velocity inferred for each cell.

https://doi.org/10.1371/journal.pcbi.1014060.g002

To quantify model performance, we calculated mean squared error (MSE) for continuous-count models and log likelihood (LL) for discrete-count models (Figs 2B and 2C). For VeloVAE, VeloVI and PyroVelocity, we calculated these metrics on a held-out test set not used during training. Note that DeepVelo and UniTVelo cannot perform out-of-sample prediction, so we were not able to evaluate it on a held-out test set. ScVelo can in theory perform out of sample prediction, but this capability is not implemented in the scVelo package, so we calculated the metrics for scVelo on whole dataset in this comparison. CellDancer does not directly model count data and its training objective is cosine similarity, so we do not compute its MSE. Our results show that VeloVAE achieves comparable or better data fit than other methods in the evaluation. Furthermore, this better performance does not come from overfitting the training dataset - VeloVAE shows good performance on the held-out datasets.

To evaluate the accuracy of latent time inference, we used scRNA-seq datasets with cells sampled from multiple time points. These time points usually have rather coarse granularity (e.g., day 7 and day 14), and cells captured at the same time may span a wide range of developmental stages. Nevertheless, the inferred cell times should at least be correlated with the capture times. Thus, we computed the Spearman correlation between the cell times inferred by each method and the capture times (Fig 2A). VeloVAE consistently achieves comparable or better performance in terms of latent time accuracy, with scVelo often inferring latent time that is anticorrelated with real time (Fig 2A). Although scVelo and VeloVI infer latent time separately for each gene, they provide a post-hoc procedure for estimating a single global time for each cell. Using this global time for comparison with our methods casts these methods in the best possible light because the global time is more robust than the gene-specific latent times. The low time correlation from scVelo may be partly explained by inconsistency among the different notions of time fitted for each gene. To investigate this further, we computed the average time correlation between scVelo’s and VeloVI’s gene-specific and global latent time. As S2 Fig shows, the correlation between the global latent time and the latent time for each gene is indeed quite low; the latent time values for many genes are even anticorrelated with global latent time. We also note that cellDancer and DeepVelo do not include cell time in their models, but rather perform a post-hoc global time estimation based on velocity.

To quantitatively measure how well the velocity flow matches our prior knowledge of different tissue developmental processes, we propose a modified version of cross-boundary direction correctness (CBDir). CBDir was initially introduced in the VeloAE paper [38]. This metric quantifies how accurately each estimated velocity matches prior knowledge at the cell type level. We generalized it by considering k-step neighbors and including time ordering in the evaluation. We also subtract a background CBDir under random walk to quantify how much better the model performs than random guessing. We compare our results in both gene expression space and 2D low-dimensional embedding based on umap (except for the human bone marrow dataset which based on tsne). Our results show that VeloVAE achieves a competitive performance on the benchmarking datasets (S1 Fig).

To further assess the reliability of VeloVAE inference, we investigated the time variance and cell state uncertainty across cells. We observe low time variance in early developmental cells, such as ductal cells in the pancreas dataset and progenitor cells in the human bone marrow dataset (S3 Fig). Correspondingly, these early-time cells exhibited higher cell state uncertainty compared with cells at later latent times (S4 Fig), consistent with biological expectations. That is, undifferentiated cells possess greater developmental potential and can diverge into multiple lineages. To better visualize the inferred dynamics, we extended the 2-D UMAP embedding by adding latent time as a third axis, generating a 3-D representation of the cellular state space along developmental progression. In this space, cells display smooth temporal transitions with clear lineage branching (S5 Fig), and 3-D velocity streamlines align with these trajectories (S6 Fig), confirming that VeloVAE captures coherent and biologically meaningful differentiation dynamics.

In addition, one of our benchmark datasets was generated with metabolic labeling to profile the embryonic differentiation process [31], providing a ground-truth reference for evaluating transcription rate parameters. We compare the transcription rates predicted by the VeloVAE continuous model with those from DeepVelo and cellDancer, the two other models that estimate cell-wise transcription rates. VeloVAE accurately infers transcription rates, achieving a mean Pearson correlation coefficient (PCC) of 0.72 between predicted and ground-truth transcription rates per cell per gene. Furthermore, VeloVAE exhibits a much higher correlation with the ground truth transcription rates compared to both DeepVelo and cellDancer (S7 Fig), demonstrating its superior performance in recovering cell-wise transcription rates. To enable a biologically meaningful comparison between predicted and ground-truth transcription rates, we further compute a Transcript-per-million-like (TPM-like) normalization (see Methods) and evaluate performance in terms of mean absolute error (MAE). Consistently, VeloVAE achieves the lowest MAE of approximately 13,046 transcripts per million, outperforming DeepVelo (44,450) and cellDancer (216,155).

2.3 VeloVAE better captures qualitative properties of complex gene expression dynamics

We designed the VeloVAE model to relax several of the restrictive assumptions of previous RNA velocity approaches. Thus, we expect that the model should show increased expressiveness, allowing it to capture qualitative properties of gene expression changes that previous approaches cannot. To assess this, we fit VeloVAE and all the other models on two representative datasets–mouse pancreas and brain–and inspected the resulting model fits. We found two types of qualitative behaviors that other existing approaches cannot accurately model, while VeloVAE can.

2.3.2 Cell type bifurcations.

Many cell differentiation processes produce multiple descendant cell types from a single progenitor type, but previous RNA velocity approaches model only a single cell type. By including a cell-specific latent state, VeloVAE can model the continuous emergence of multiple cell types from a single progenitor type. For example, in the pancreas dataset, the Nnat gene is upregulated as cells differentiate towared the beta cell fate, but not in any other cell type (Fig 2D).

As another example, we fit all models on a developing mouse brain atlas [30]. For clarity, we subsampled the dataset to include only cell types arising from the neural tube and neural crest (see S8 Fig for an analysis of the full dataset).

In this system, neural tube cells develop into radial glia. Some of the radial glia cells differentiate into neuronal cell types, while the others give rise to the glial lineage, including glioblasts, oligodendrocytes, astrocytes, and ependymal cells. The neural crest cells develop to become fibroblast cells. In short, this is a complex system with many distinct lineages emerging. Nevertheless, VeloVAE can accurately model the complex, multi-lineage dynamics of genes in this system. For example, VeloVAE accurately models the behavior of the Atp1a2 gene (Fig 2F), which is turned on in the fibroblast and glial cells with distinct transcription rates and remains off in the neuronal cells differentiated from radial glia. In contrast, scVelo rearranges the latent time values of the cells in a vain attempt to fit the Atp1a2 gene into a single induction and repression cycle. VeloVI reverses the time order as fibroblast cells flow backwards to mesenchymal cells. We trained UniTVelo with unified cell time and the cell time is generally correct. However, without explicitly modeling branching, the parametric model from UniTVelo underfits the count data. Neither DeepVelo nor cellDancer accurately recovered the cell time for the fibroblast branch, and they suffer from the same problem of uninformative velocity.

Although VeloVAE accurately models bifurcations using a continuous cell state variable, the resulting parameters are not readily interpretable in terms of discrete cell types. Thus, we developed a model extension that aids in interpreting how the kinetic parameters of gene expression change across cell types. We extended the simple differential equation model shown in Fig 1 to a system of equations that we refer to as a branching ODE model (Fig 3A). Instead of a cell-specific transcription rate and fixed splicing and degradation rates, the branching ODE model assigns each cell type a unique ODE with cell-type-specific transcription, splicing, and degradation rates and an initial condition determined by the progenitor cell type (Methods). The model relies on a directed graph relationship among discrete cell types. The graph can be inferred directly from the data, as we did here; if some aspects of the cell type transition graph are known, these can also be manually specified. We constrain the branching ODE model so that each cell type emerges at a specific time and the initial conditions of each cell type match the ODE prediction of its parent cell type at the time the child cell type emerges. This provides a qualitative view of the change of kinetic rates during cell development. We use the branching ODE model to replace the decoder of the VeloVAE, giving an alternative, more interpretable generative model.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3. Cell Type Transition Graph Inference and Branching Differential Equation Model.

(A) Schematic of branching ODE solutions. The example shows the output prediction of u and s versus time for a process with a single progenitor, an intermediate and three distinct descendant cell types. Inferred cell type transition graphs from pancreas (B), iPSC (E) and subsampled mouse brain (H) datasets. In each graph, a vertex represents a cell type and a directed edge points from a progenitor cell type to its immediate descendant(s). Examples of individual genes fit by the branching ODE model for pancreas (C), iPSC (F) and subsampled mouse brain (I) datasets. Each column represents a gene and plots the unspliced count, spliced count and RNA velocity versus time from top to bottom. The branching ODE fits are shown as solid lines. Cell-type-specific rate parameters inferred by branching ODE model for pancreas (D), iPSC (G) and subsampled mouse brain (J) datasets are also shown. Each column represents a gene and plots the transcription (), splicing () and degradation () rates from top to bottom. Each point represents a cell type, the types are arranged in chronological order, and progenitor-descendant relationships are indicated with lines.

https://doi.org/10.1371/journal.pcbi.1014060.g003

Having established the accuracy of the results from VeloVAE, we used these results to infer cell type transition graphs. To infer the transition probability between cell types A and B, we used the cell times to simply count how often cells of type A and B occur in two immediately adjacent short time intervals (see Methods for details). We show three examples of these graphs (Figs 3B, 3E and 3H), inferred from the pancreas, iPS reprogramming, and mouse brain datasets from Fig 2. The inferred cell type transition graphs match closely with biological expectations for these systems, even for the complex mouse brain dataset.

To train the model, we first obtain the time and state assignments for each cell by training VeloVAE. Next, we infer a transition graph describing the progenitor-descendant relations among cell types (Methods). Finally, we fix the encoder of VeloVAE (so that latent time and cell state estimates are fixed) and estimate the parameters of the branching ODE model. We perform this parameter estimation by maximizing the Gaussian likelihood of all genes under the branching ODE model, which is equivalent to minimizing the Mahalanobis distance between model fit and observed data.

We trained branching ODEs on pancreas, iPSC and subsampled mouse brain datasets. The cell type transition graphs are all consistent with prior knowledge (Figs 3B, 3E and 3H). For example, our computational method successfully finds the expected cell differentiation path in pancreatic development, starting from ductal cells and branching into , , and cells (Fig 3B).

In addition, the branching ODE is able to infer different and asynchronous gene expression kinetics in multiple branches and successfully captures different rates in different branches in these datasets. For example, Ppp1r1a has almost all transcription activity in cells (Fig 3C), which was verified by previous studies [39,40]. The Rnf130 and 1500009L16Rik genes similarly show significant branching trends that are accurately modeled by the branching ODE (Fig 3C). The branching ODE model also accurately fits genes that show differential kinetics among lineages that emerge during induced pluripotent stem cell reprogramming. For example, Nr2f1 is strongly and specifically upregulated in neural-like cells (Fig 3F). Hsp90aa1 is upregulated moderately in Trophoblast and strongly in IPS cells. Krt7 is upregulated in epithelial-like and trophoblast-like cells with differing expression levels (Fig 3F).

Another example is Mapt from the mouse brain dataset (Fig 3I). The gene is upregulated strongly in neurons and subsequently transcribed at much lower levels in oligodendrocytes, which coheres with previous studies [41]. The Napl15 gene shows the opposite trend, with high transcription in oligodendrocytes and low but detectable transcription in glioblasts and neurons (Fig 3I). As another example, the Cnn2 gene is most highly transcribed in the early transition from neural crest to mesenchyme and then to fibroblast (Fig 3I). The cell-type-specific rate parameters inferred by fitting the branching ODE describe the differences in transcription, splicing, and degradation that cells undergo as they differentiate (Figs 3D, 3G and 3J).

2.4 VeloVAE accurately models human hematopoiesis from discrete count data

Previous papers have noted that hematopoiesis is a particularly difficult system for existing RNA velocity methods [16]. Latent time and velocity inferences often seem to point in the opposite direction of the known blood cell differentiation hierarchy. Two aspects in particular likely make this system challenging. First, many distinct cell types emerge simultaneously from the hematopoietic stem cell (HSC). Recent studies suggest that hematopoietic progenitors are more like a continuum of primed states than a set of discrete states neatly organized in a tree structure [42]. Second, blood cells are produced exceptionally rapidly compared with other cell types; a recent study estimated that about 2 million new red blood cells per second enter the bloodstream [43]. Thus, blood cells may use special gene regulatory mechanisms such as “transcriptional boosts” and other time-varying rate parameters.

To investigate whether VeloVAE can resolve these difficulties, we analyzed a recent human bone marrow dataset [29]. Our stream plot shows that VeloVAE correctly identifies HSCs as the start of differentiation and predicts that they differentiate into megakaryocytes, platelets, dendritic cells, monocytes, and B-cells (Fig 4A). As a quantitative comparison, we computed k-step CBDir (See Methods for the definition) on RNA velocity in the original gene expression space. Our results show that all variants of VeloVAE perform equally well or better than other methods in terms of velocity direction (Fig 4B). Note that both velocity stream plot and CBDir are high-level representations of cellular dynamics. In fact, the quality of gene-level velocity fitting is not embodied entirely in the stream plot. Inspecting the fits of individual genes revealed that other methods fail to capture the complexity of multiple lineages simultaneously emerging. For example, the CD36 gene is upregulated strongly in red blood cells, moderately in monocytes, and slightly in dendritic cells. VeloVAE correctly models the expression dynamics while other methods either reverse the cell order or fail to capture distinct dynamics (Fig 4C). Noticeably, the discrete VeloVAE model achieves the best performance and we show that it is capable of fitting genes with very low count numbers (Fig 4D). VeloVAE accurately infers the true differentiation trajectory in this complex system, but none of the other methods do (Fig 4E). In the latent time inference, scVelo, UniTVelo, and DeepVelo identify CD14 monocytes as having the earliest latent time, whereas CD14 monocytes are known to represent terminally differentiated cells in hematopoiesis. Similarly, cellDancer infers memory B cells as the earliest cell population, and VeloVI identifies plasmablasts as root cells. In contrast, VeloVAE correctly assigns HSCs as the population with the earliest latent time. Additionally, VeloVAE identifies naive CD4 and CD8 T cells as distinct developmental roots, which subsequently differentiate into CD8 memory and CD8 effector T cells. Accurately inferring these developmental paths is critical for evaluating method performance. As shown in Figs 4A and 4E, VeloVAE correctly reconstructs the ground truth trajectories from HSCs → LMPP → GMP → CD14/CD16 monocytes, as well as the HSC → RBC lineage. In contrast, scVelo, UniTVelo, DeepVelo, and cellDancer predict the RBC lineage in reverse, from RBCs toward HSCs, while VeloVI predicts the CD14/CD16 lineage in reverse. Overall, VeloVAE is the only model that accurately infers differentiation paths consistent with prior biological knowledge.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 4. VeloVAE Models Complex Kinetics in Bone Marrow Cells.

(A) UMAP plots colored by published cell types overlaid with velocity inferred by all methods. (B) Quantitative performance comparison. We computed k-step cross-boundary direction correctness (K-CBDir) for k = 1,2,3,4,5. (C) Examples of individual genes fit by all methods. Gene fits are shown for both u and s values, with inferred latent time on the x-axis. Similar to Fig 2, fitted values from scVelo and UniTVelo are shown as lines, with observed data values shown as points and only fitted values are plotted for other methods. Points are colored by published cluster assignments, with the same colors as in the corresponding UMAP plots. The arrow length and direction indicate the velocity inferred for each cell. (D) Examples of low-count genes fit by discrete VeloVAE. (E) UMAP plots colored by latent time inferred by different methods.

https://doi.org/10.1371/journal.pcbi.1014060.g004

2.5 VeloVAE accurately models organ development in whole mouse embryos

The challenges of studying cell differentiation with single-cell data become particularly acute at the scale of an entire organism. For example, Cao et al. performed single-cell RNA-seq on 61 entire mouse embryos sampled from E9.5-E13.5, the developmental period when mouse organogenesis occurs [4]. The dataset contains 1,191,071 cells after preprocessing, which fall into 38 major cell types and can divided into 9 major lineages. Pseudotime analysis and RNA velocity analysis have been performed on this dataset, but both analyses were highly manual processes that required separate curation of dozens of cell subsets [4,44].

In principle, the continuous cell state variable of VeloVAE is sufficiently expressive to provide a single model of the differentiation potential for an entire organism. To investigate this, we trained VeloVAE on the entire mouse organogenesis dataset. Because of the large size and cellular diversity of this dataset, we used a larger batch size of 1024 and increased the dimension of c from 5 to 30. Our results show that VeloVAE discovered a meaningful latent cell state space representing the whole developmental process (Fig 5A). To examine the results in more detail, we ran UMAP individually on the same 9 broad cell lineages that were identified in the initial paper. Importantly, we performed this subset analysis purely for visualization purposes–the cell time, cell state, and kinetic rate parameters were estimated only once using all cells jointly (S8-S14 Figs).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 5. Whole mouse embryo analysis.

(A) UMAP of entire mouse embryo dataset, colored by broad cell class. (B)-(C) UMAP plots colored by published cell types overlaid with velocity inferred by VeloVAE for neuronal (B) and mesenchymal (C) trajectories. The UMAP coordinates are computed using only cells from each trajectory, but the velocity estimates are from a single VeloVAE model fit on all cells. (D)-(E) UMAP plots colored by capture time, published pseudotime, and VeloVAE latent time values for the neuronal (D) and mensenchymal (E) trajectories. Note that latent times are from a single VeloVAE model fit on all cells. (F)-(G) t-SNE plots colored by (F) cell state uncertainty and (G) time variance. (H)-(I) VeloVAE fits for selected genes from the neuronal (H) and mesenchymal (I) trajectories.

https://doi.org/10.1371/journal.pcbi.1014060.g005

These visualizations indicate that latent time and velocity estimates are highly consistent with cell capture times and biological prior knowledge (Figs 5B-E, S9 and S10). Remarkably, the latent time estimates are even more accurate than the authors’ reported pseudotime values in several cases. For example, the VeloVAE latent time estimates for the two lineages with the most cells, mesenchyme and neural tube, both show a higher correlation with capture time than the pseudotime estimates reported by Cao et al [4]. In particular, the pseudotime estimates for the oligodendrocyte, neural progenitor and inhibitory and excitatory neuron cells are essentially uncorrelated with capture time, whereas the VeloVAE results proceed in the direction of increasing embryonic stage (Fig 5D). Similarly, the pseudotime estimates for skeletal muscle cells are actually uncorrelated with cell capture times, and the pseudotimes of developing connective tissue cells substantially underestimate their developmental progress (Fig 5E).

These results are remarkable because the pseudotime values were generated by the developers of Monocle3–one of the most popular pseudotime methods–on their own data using their own tool in a highly manual process. Cao et al. used expert knowledge to group cells into “subtrajectories” and choose root cells for each. In contrast, we obtained the VeloVAE results by simply running the model jointly on all cells with no manual curation. The accuracy of the VeloVAE latent time estimates underscores the power of this approach for studying cell differentiation in large, complex, multi-lineage single-cell datasets. This comparison also highlights the inherent difficulty of pseudotime analysis in such complex datasets, particularly when the continuous nature of differentiation and the presence of multiple “root cell” populations make it very challenging to divide the cells into discrete subtrajectories.

In comparing latent time and pseudotime values, we realized that the latent time values offer another advantage: they can be assigned real time units. Because we use the capture times as prior information, this places the inferred latent time on the same scale as the capture times. That is, with appropriate normalization, we can report the latent times in units of days or hours. Consequently, the rate parameters can also be assigned units, such as transcripts per minute. To explore this, we converted the latent time values into units of minutes and normalized the transcription rate parameters to the absolute number of transcripts in a typical mammalian cell. This shows that the transcription rates are roughly on the order of 1–10 transcripts per minute for most genes (S15 Fig), in accordance with previous estimates from single-molecule FISH experiments [45]. The splicing rates are predicted to be slightly lower, on the order of 0.1-1 transcripts per minute. This is consistent with a previous study that estimated the elongation rate in human cells at 3.8 kb/min and found that splicing begins within 10 min [46]. Our estimated degradation rates are similar to the transcription rates, on the order of 1–10 transcripts per minute. Though these estimates should not be taken as definitive, we find it reassuring that our estimated rate parameters are on roughly the correct order of magnitude compared with prior knowledge.

5.1 Problem setup

First, we formulate the computational problem of modeling cellular gene expression changes using the mathematical language. Single-cell RNA experiments produce cell-by-gene count matrices. RNA reads are further classified to be either spliced or unspliced, and correspondingly, there are two count matrices of size , representing unspliced and spliced RNA counts. Each sample (cell), indexed by i, is represented by a vector parametrized by unobserved time t_i. The trajectory X_i(t) is governed by some differential equation plus random noise. Our goal is to infer the cell time t_i and predict the velocity .

We list a set of definitions from our previous work [23] for convenience of description. We refer the reader to the paper for more details.

Definition 1. (Gu et al., 2022 [23]). Let u_g and s_g denote the unspliced and spliced mRNA count of the g-th gene. Let be a set of genes measured in an scRNA-seq experiment. The feature vector of a cell is defined as .

Definition 2. (Gu et al., 2022 [23]). The kinetic equation of gene g is defined as a system of ordinary differential equations relating changes in u and s over time. If there exists a solution to the initial value problem with , we call this solution the kinetic function for g.

Definition 3. (La Manno et al., 2018 [9]). Given a kinetic function u(t) and s(t) of a gene, the RNA velocity of the gene is defined as .

5.2 Modeling gene expression kinetics

As first proposed in 2018 [9], the kinetic equation is modeled by a system of two linear ODEs:

(1)

where is an indicator function for the condition in brackets. The model parameters and correspond to the RNA transcription, splicing and degradation rates, respectively. The model assumes that two discrete phases can occur in the gene expression process: (1) induction, when new unspliced RNA molecules are being transcribed (2) repression, when the transcription process stops and no new unspliced molecules are made. The induction phase is assumed to start at t_on=0 and the transition from induction to repression occurs at time t_off. Given an initial condition , the analytical solution to the ODE is

(2)

(3)

5.3 BasisVAE

BasisVAE [22] was orignally designed to account for feature-level clustering. It aims at clustering different features (dimensions) of a high-dimensional observation via a variational auto-encoder. A potential application is to cluster genes with similar patterns.

The model assumes each feature is generated by one of K basis functions, i.e.,

(4)

Here, is a low-dimensional latent variable generating the observation . The clustering information is stored in Categorial and Dirichlet. Thus, there are two additional latent variables, and . According to Märtens and Yau [22], we can marginalize using collapsed inference. The collapsed evidence lower bound (ELBO) is given by

(5)

Here, the second term has an analytical form of

(6)

where and is the gamma function.

5.4 Variational mixture of ODE model

5.4.1 ODE formulation.

We extend the ODE formulation (1) to broader cases. First, we adopt the idea of varying transcription rate proposed in our previous work [23]:

(7)

To account for repression-only genes, we add a second set of ODEs:

(8)

Let and be the analytical solutions to equations (7) and (8) of the g-th gene. Define and . Given t, and are vectors containing solutions to two sets of kinetics equations with different assumptions. As we will discuss below, solutions to the two sets of ODEs are two basis functions of our VAE model.

5.4.2 Generative process.

The generative process for the variational mixture of ODE model is as follows:

(9)

Here, is a neural network with parameters , is a vector indicator function, ⊙ is the elementwise product, F is the kinetic function of all genes, and is a diagonal covariance matrix.

In the context of BasisVAE, F_ind and F_rep are two basis functions representing two potential data generative processes. By introducing , we distinguish genes by different generative processes, thus promoting gene-level clustering.

5.4.5 Parameter initialization.

The neural network weights are randomly initialized with Xavier uniform distributions [47]. The only exception is the output layer of the decoder network, which is initialized with Xavier normal distributions, to prevent gradient vanishing of the sigmoid activation. VeloVAE applies the same initialization method as scVelo by applying the steady-state assumption and the dynamical model. Next, a global cell time is estimated as the median of locally initialized time across all genes.

If capture times are available, VeloVAE initializes a unique cell time by directly sampling from a distribution centered at the real time. Whether the initial global time is estimated by the steady-state model or directly sampling from capture time, rate parameters are always re-estimated based on the dynamical model. We still make the assumption that . First, switch-off time is estimated as the average global cell time of the steady-state cells. Next, we estimate and the switch-on time t_on by solving a set of equations using two cells with their initialized cell time :

(10)

Note that u₁ and u₂ are chosen as the sample average around the median and top quantile of u values to promote robustness against noise. Finally, is estimated by where s_ss is the estimated steady-state s value.

Let n_ind and n be the number of cells in the induction phase and total number of cells. A naive approach will be , i.e., assigning the proportional of cells in the induction phase. However, such a gene-wise initialization is sensitive to data noise, similar to the problem of gene-wise time fitting. Instead, we make use of a simple but critical fact: genes generated by the same underlying generative process (inductive or repressive) should have similar dynamical behavior over time. This means the count number should be positively correlated when both genes are generated by the same basis function. Using this insight, we compute a gene-by-gene weighted adjacency matrix A where , where denotes Pearson correlation coefficient and denotes unspliced and spliced counts of the i-th gene across all cells. We apply spectral clustering on genes using A. The number of clusters is determined by using an empirical threshold on the singular values of A with a Gaussian noise assumption [48]. Denote as the proportion of cells in the induction phase and be the clustering result of the g-th gene. Next, for each cluster y, we perform two Kolmogorov-Smirnov tests between and Dirichlet(5.0,5.0). The two null hypotheses are

1. is sampled from Dirichlet whose CDF is less than CDF of Dirichlet(5.0,5.0)
2. is sampled from Dirichlet whose CDF is greater than CDF of Dirichlet(5.0,5.0)

If the first null hypothesis is not rejected, genes in cluster y are likely to be inductive. We set and reinitialize by sampling from Dirichlet. If the second null hypothesis is not rejected, genes in the cluster y are likely to be repressive and we perform the same sampling but with . If both null hypotheses are rejected, we sample from a mixture of two Dirichlet distributions mentioned in the previous two cases. To prevent a bad initialization, which is likely when the dynamical model fails, we add an option to replace by . This helped to correct the time inference in some of our experiments.

5.4.6 Estimating ODE parameter uncertainty.

The VAE model can be extended [49] to account for uncertainty in the parameters of the generative model, which in our case are the kinetic rate parameters of the ODE system. Let be the set of all ODE parameters. We assume that some prior distribution generates the parameters. Similar to the (intractable) problem of inferring the latent cell time and state variables, we can use a variational approximation of the posterior . The marginal likelihood of the input features can be bounded by

(11)

For the prior , we choose a factorized log-normal distribution, i.e., each rate parameter , or is a random variable drawn from a log-normal distribution. By default, the logarithm of rate parameters has a mean of zero and a standard deviation of 1 for and 0.5 for and . The posterior mean is initialized using the same method described in the previous paragraph. Again, we can apply the reparameterization trick to take samples of and estimate the expectation of ELBO with a sample mean. The KL divergence between and can be viewed as a regularization of the ODE parameters.

This statistical approach also has a nice intuition in our case: it prevents the “vanishing velocity” problem. If we don’t regularize , optimizing the evidence lower bound may lead to large rate parameters, so that . This in turn causes the velocity to vanish so that each cell will be approximately at the steady state, with and . Because is given by a neural network g, we can achieve good reconstruction if g is expressive enough, but the velocity will be close to zero. This issue is resolved by placing a prior distribution on the rates and including them in the variational approximation. This has the effect of regularizing the posterior distribution of the rates toward their prior, avoiding the vanishing velocity problem.

5.4.7 Interpreting rate parameters.

Although the model accepts any notion of time and rates, the units can be converted to match the actual units of cell developmental time, usually in days. Suppose the cell time from the model ranges from t₀ to t₁, in a hypothetical time unit, and we have prior knowledge about the entire duration of m days. From the mathematical property of our ODE system, we know that scaling time by k and rate parameters (,,) by results in the same u and s. Thus, we can perform unit analysis to convert the rates into units of minutes as follows:

(12)

Similarly, and

We can then scale the rate parameters to account for the fact that scRNA-seq captures only a fraction of the transcripts in a cell. A reasonable estimate is about 360000 transcripts in a typical mammalian cell (assuming 10 pg total RNA per cell, 1–2% mRNA, average transcript length of 2kb). Therefore, we can scale by , where x_total is the median total mRNA count number. We analyzed rate parameters learned from our model, converted the units and computed the histograms (S15 Fig). For , we analyzed the peak transcription rate by considering the case of highest transcription (). For and , we multiplied them by u_top and s_top respectively to obtain the same units as . Here, u_top and s_top are the 95-percentile u and s values, representing cells with high expression levels.

5.5 Solving the chemical master equation with discrete VeloVAE

The biological process of RNA splicing is a chemical reaction network described by the chemical master equation. It considers the change of molecule numbers, which are discrete quantities, as a stochastic process. As Bergen et al. pointed out [10], the dynamical model of RNA velocity is justified by considering the first moment of a stochastic process.

Here, we close the gap between the chemical master equation and its moment equation by introducing the discrete VeloVAE model. We start with chemical equations of all reactions. Consider a gene with transcription, splicing and degradation. For any gene g, the reactions are described by the following chemical equations:

(13)

Suppose we have a total of G genes. Then, there are reactions and molecules in total. Since all reactions are monomolecular, the system has a closed form solution if the initial condition is a product Poisson distribution. This is formally described as the following theorem from a previous work [24]:

Definition 4. A product Poisson distribution with mean is a discrete distribution with a support of and PMF of the following form:

(14)

Theorem 1. (Jahnke, 2007). Let P(t,x) be the pmf of a feature vector (mRNA counts) at time t. If P(0, x) is a product Poisson distribution, then, P(t, x) at any time t > 0 is a product Poisson distribution whose mean is given by the following ODE:

(15)

Here the matrix A contains splicing and degradation rates while b is a vector of transcription rates. An exact solution of a single-gene system, as a special case, is presented in another work [50]. The discrete VeloVAE model is based on the Poisson initial condition assumption. By theorem 1, solving the chemical master equation is equivalent to solving an ODE describing the Poisson mean. In the discrete VeloVAE model, we again use the variational mixture of ODEs as a cell-specific ODE model for the Poisson mean. The generative process is the same as the continuous model except for the last few steps:

(16)

Here, we still have a mixture of basis functions for each gene. However, the mixture components are Poisson distributions instead of Gaussian distributions. Besides, we have an additional cell-wise library size scaling factor l to account for differences in count numbers resulting from different sequencing depths.

Because count numbers are drawn from a Poisson distribution, we define the RNA velocity of a discrete model as , where is the Poisson rate as well as the mean. This is a natural extension from the continuous model.

5.6 Branching ODE

5.6.1 Model description.

A variational mixture of ODEs provides sufficient flexibility to account for complex kinetics, but only describes such kinetics at the cell level. In order to distill qualitative knowledge about gene expression kinetics and reveal cell-type relations, we propose a new ODE model called branching ODE.

The fundamental assumptions we make about branching ODE include:

1. Each cell belongs to exactly one of the cell types .
2. At least one of the cell types is a stem cell type. Each non-stem-cell type is a descendant of exactly one other cell type. (Note that we could relax this assumption to allow multiple progenitors, but we choose not to purely for simplicity.) Each cell type has an initial time t₀ when it emerges in the differentiation process.
3. Cells of the same type, y, share identical transcription, splicing and degradation rates ().

With these assumptions, we can summarize differentiation with directed cell type transition graph , called a transition graph. Each cell type corresponds to a vertex in G, and each edge (u,v) represents the relation of u differentiating into v. By our second assumption, G is composed of one or multiple trees. The kinetic functions retain the same analytical form, except for type-specific rate parameters and initial conditions. Furthermore, the equations are defined recursively because the initial condition of any non-stem-cell type depends on its progenitor cell type.

(17)(18)(19)(20)

Here, par refers to the parent vertex in the transition graph. Note that is redefined as for each cell type y, i.e., is the time duration starting from the initial time of each cell type.

5.6.2 Inferring the transition graph.

A challenging problem for applying branching ODE is the absence of the transition graph. In many cases, we do have some prior knowledge of the transition graph, but it is more desirable to simply infer the graph directly from the data when possible. Therefore, we apply a computational method to infer the transition graph based on cell time and states. Because the transition graph is a collection of arborescences, i.e., rooted trees in a directed graph, we can solve the problem with simple graph algorithms.

First, we partition the cells into distinct cluster(s). In particular, we perform Leiden clustering with a low resolution on the UMAP coordinates of the data.

Next, we build a complete subgraph in each partition. Let be the set of all cells and be the set of all cells of type y in a partition. For each cell with time t_i, we take a time window and find k nearest neighbors based on cell state c. Denote as the set of neighbors of i. Then, for any two cell types y and z, the empirical transition probability from y to z is defined as

(21)

In other words, we group cells into the cell types and count the number of transitions from any cell type to any other type.

Finally, we apply Edmond’s algorithm [51] to find the maximum spanning arborescence in each partition. The earliest cell type is choosen to be the root. The algorithm starts by picking the parent y of each non-root vertex z greedily, i.e., . Next, it checks loops, collapses loops into super-vertices and is recursively applied to the new graph until no loop exists. Finally, it breaks the loops after each level of recursion.

5.6.3 Training the branching ODE.

We assume that the cell time has already been inferred from VeloVAE. Thus, the branching ODE is a regression model with an analytical form shown in equation (17) and (18). To train the model, we find the cell-type-specific rate parameters that maximize the Gaussian likelihood, which is equivalent to minimizing a Mahalanobis distance:

(22)

Here, is predicted expression level. We train using mini-batch stochastic gradient descent with the ADAM optimizer.

5.7 Data preprocessing

We loaded all datasets from AnnData (h5ad) format. We first used scanpy [52] to select highly variable genes, normalize and scale the gene expression counts. Next, we followed the scVelo preprocessing pipeline by performing principal component analysis on the normalized and scaled expression data, then smoothing the unspliced and spliced expression levels among k-nearest neighbors identified from the principal components.

5.8 Evaluation metrics

5.8.1 Cross-Boundary Direction Correctness (CBDir).

In a previous work [17], the authors quantitatively measured the accuracy of velocity flow on any 2D visualization. Denote as the set of all cell types. For any two cell types , we define (A,B) to be a transition pair if and only if B is a descendant cell type of A in a cell developmental process. The CBDir metric takes a set of know transition pairs as input and computes an average directional accuracy via cosine similarity. For any cell in with transition pair (A,B), its CBDir is defined as:

(23)

Here, is the projection of RNA velocity to a low-dimensional space, is the coordinates of c in the same space, and N(c) is the set of all neighbors of c in a KNN graph built from gene expression similarity. The overall CBDir of all known transition pairs is defined as

(24)

5.8.2 Generalized CBDir.

The CBDir metric is the first one directly measuring the correctness of velocity flow. It connects visual perception of velocity flow to a quantity. We think it beneficial to extend it for the following reasons:

1. The velocity in CBDir is only a projection to a low-dimensional space. Hence, it’s not a direct measure of velocity at a gene expression level.
2. Only the direct neighbors in the KNN graph are considered as targets, i.e., future transcriptomic states. In fact, such assumption might not hold due to stochastic noise in the dataset. Directly connected neighbors on a KNN graph are too similar at the transcriptomic level to be considered as a future direction.
3. CBDir doesn’t take the ordering of cells into account. Time order is important because any real-world system must proceed with increasing time.
4. It remains unclear whether velocity flow from one cell type to another is statistically significant.

Due to all the limitations described above, we devised k-step CBDir (KCBDir) metric. Given a cell with transition pair (A,B), KCBDir is defined as follows:

(25)

where . Here, we replace direct neighbors with k-step neighbors, N_k(c), on a KNN graph. Besides, and can be from either a low-dimensional embedding or the spliced counts. We further defined the generalized CBDir (GCBDir) by averaging over multiple different step sizes and removing an offset induced by taking a random walk on a KNN graph:

(26)

Here, KCBDir(c;N_k(c)) is the KCBDir between c and its k-step neighbors from its descendant cell type, and R_k(c) is the set of cells reached by randomly traversing a KNN graph starting from c.

5.9 Transcription rate comparison

For the Scifate2 dataset [31], metabolic labeling counts were measured over a 2-hour period. We defined the ground-truth transcription rate by dividing the measured labeled counts by the labeling duration (2 hours). To ensure comparability with the transcription rates predicted by continuous models (VeloVAE, DeepVelo, and cellDancer), the ground-truth transcription rates were smoothed using library-size normalization followed by K-nearest neighbor (KNN) averaging. To obtain a biologically meaningful comparison, both the real and predicted transcription rates are multiplied by 100,000 to generate transcripts-per-million-like (TPM-like) measurements. We evaluate model performance using two metrics: mean absolute error (MAE) and the Pearson correlation coefficient (PCC). MAE is computed using the TPM-like normalized transcription rates described above to compare the absolute differences between predicted and ground-truth values. For correlation analysis, the PCC is calculated between the KNN-averaged ground-truth transcription rates and the model-predicted transcription rates per gene across cells.

5.10 Time and resource consumption

We evaluate the time and computational resources required for preprocessing, model initialization, and mini-batch training on the Pancreas dataset. The preprocessing step took approximately 8 seconds to select 2,000 genes, followed by around 30 seconds for model initialization. The mini-batch training phase required roughly 5 minutes to complete. All computations were performed on an NVIDIA A40 GPU with 40 GB of RAM.