Cell2State: Learning Cell State Representations From Barcoded Single-Cell Gene-Expression Transitions 

1 Introduction . With the explosive amount of data from single-cell genomics studies , one remaining major challenge is the lack of ability to understand cell transition on the individual level . Conventional methods for analyzing single-cell dynamics are mostly based on “ ensemble ” analysis ( Kester & van Oudenaarden , 2018 ; Tanay & Regev , 2017 ; Bacher & Kendziorski , 2016 ; Stegle et al. , 2015 ) . Such analysis reveals population-level trends , but can not reveal behaviors of individual cells . Recent advances in genomic technology have enabled scientists to directly measure cell lineages and connect two cells that are far apart in the time course . This is opposed to inferring lineages from snapshots at nearby time points . The gene barcoding approach works by inserting into each cell a DNA sequence , i.e. , a barcode , that randomizes across cells so that no two cells bear the same sequence ( Woodworth et al. , 2017 ) . As the cell divides and differentiates , its descendants can be identified based on sequencing the label . This concept has been utilized recently in single-cell analysis of embryonic development ( Yao et al. , 2017 ; Wagner et al. , 2018 ) , stem cell reprogramming ( Biddy et al. , 2018 ) , and fate determination in hematopoiesis ( Weinreb et al. , 2020 ) . This genetic barcoding approach enables tracking of evolutionary trajectories across individual cell lineages . In this paper , we focus on the new type of single-cell data enabled by the single-cell gene barcoding technology . The barcode can directly connect parent cells with their descendants over a long time span . Thus , gene barcoding coupled with RNA-seq generates pairs of gene-expression transitions { ( X , X ′ ) } , where each ( X , X ′ ) is the gene expression profiles for a parent cell at an early time point and one of its descendants at the later time point . One may view the gene-expression profile X as the raw state of a cell , which is a high-dimensional vector . Such state-transition data makes it possible to learn about a cell ’ s law of state transition . In other words , the new data type can let us decode the single-cell transition law in a way similar to system identification for dynamical systems . We wish to learn mathematical abstractions of cell gene-expression states , i.e. , a map Ψ from gene-expression profile X to a vector Ψ ( X ) of lower dimension . A good cell state abstraction should be low-dimensional and predictive , compressing predictive signals from gene-expression levels in a compact vector . In other words , we hope to maximize the mutual information I ( Ψ ( X ) , X ′ ) between the embedded parent cell state and its descendant . Ideally , we hope to achieve a nearly lossless encoding of states , i.e. , I ( Ψ ( X ) , X ′ ) ≈ I ( X , X ′ ) . To learn such Ψ from noisy high-dimensional gene-expression data , we build on ideas from dynamical system theory , kernel machine learning , and low-rank optimization . In the area of molecular dynamics , Schütte et al . ( 2011 ) showed that the leading spectrum of transfer operator can be used to identify coresets for faster simulation . In reinforcement learning , computers need to figure out state abstraction of unknown transition systems in order to learn to control quickly ( Sutton et al. , 1998 ) . Sun et al . ( 2019 ) developed a state embedding method to analyze game trajectories to significantly reduce the state dimension of one-player Atari games . See Appendix A for more discussion on related works . A summary of our work : . • Building on the spectral compression ideas from dynamic systems , we develop a cell-to-state ( cell2state ) representation learning method for analyzing gene-expression transition data made available by the single-cell barcoding technology . The cell2state algorithm is trained using gene-expression transition pairs ; it finds a mapping to approximate and embed the transition distributions in a low-dimensional space . The embedding map is learned by first “ lifting ” the data ’ s dimension by random feature projection , then estimating a large matrix embedding of the transition distributions , and finally “ compressing ” the dimension by low-rank optimization . • We provide information-theoretic analysis of the learnt cell2state embedding Ψ̂ . In particular , we show that , upon appropriate quantization , the embedding map can be used to encode raw gene-expression data into a small number of bits . We show that this encoding can be nearly lossless , and we establish sample complexity bounds for preserving the mutual information between parent and descendants up to 1 bit . • We apply cell2state to a published single-cell barcoded RNA-seq dataset for studying stem cells ( Biddy et al. , 2018 ) . In this analysis , we used cell2state to map early-day cells to state vectors of dimension ≤ 8 . Via the cell2state map , the cell populations demonstrated sharp polytope structures , where distinct vertices provide early signals that predict diverse dynamics and cell fates . To evaluate the learned cell state vectors , we test them in three downstream tasks : ( i ) finding dynamically stable cell cluster ; ( ii ) early prediction of cell dynamics , such as cell descendants ’ activities or fates , based on low-dimensional state vectors ; ( iii ) subpopulation analysis to identify marker genes that signal distinct cell dynamics . Across these tasks , we observe substantially improved performance using the learned cell states , as compared with baselines that use either the raw data or features that are not dynamics-aware . In particular , our results show that cell2state achieves similar/better level of prediction accuracy using ≤7 dimensions , compared to neural networks that use raw gene expressions as input ( up to 5000 dimensions ) . • Further , we identify and examine subpopulations of cells that have the most representative low-dimensional states ( in other words , subsets of cells that are close to likely meta-states , under the assumption of a latent-state model ) . These subpopulations are used to identify biologically relevant marker genes . These marker genes identified by cell2state are known to relate to stem cell reprogramming and epigenetic regulations . 2 Markov Branching Process Model for Single-Cell Dynamics . We model the time-course dynamics of gene expressions as a branching diffusion process ( Edwards ( 1970 ) ; see Figure 1 ) . Let Xt denote the gene-expression profile of a cell at a time point t , which is a high-dimensional vector . Each Xt has a random number of descendants Xit+1 , i = 1 , . . . , N with independent and identical distributions . Definition 1 . Define p ( Xt+1|Xt ) as the transition function for the gene expression profile Xt to evolve to a collection of descendants { X ( i ) t+1 } Ni=1 in a fixed amount of time , i.e. , for any measurable set S , p ( S|Xt ) = E [ N∑ i=1 1 { X ( i ) t+1 ∈ S } | Xt ] , where N is also a random variable . Note that p is not necessarily a probability density function because a cell could have multiple descendants . If the growth rate is such that E [ N |Xt ] = ∫ p ( y|x ) dy > 1 , we say the cell is actively growing . Definition 2 . Let P be the transition operator of the branching diffusion process with transition function p , given by Pf ( X ) = E [ N∑ i=1 f { X ( i ) t+1 ∈ S } | Xt = X ] . In single-cell analysis , the transition function , p , and operator , P , are infinite-dimensional , thus estimating them is largely intractable from finite noisy data . Like many other dynamic systems , cell dynamics is often driven by a small set of marker genes , and thus , it may admit an intrinsic low-dimension structure . We make the following assumption : Assumption 1 . Let H be a space of functions . There exists a r-dimensional embedding map Φ∗ ⊂ H such that Pf ∈ Span ( Φ∗ ) , ∀f ∈ H. Here , H will be specified later . Such low-rank structure of Assumption 1 naturally exists in dynamical processes that admit latent states . Suppose that each x can be represented as a mixture over meta-states { z } such that p ( x′|x ) = ∑ z p ( x′|z ) pZ ( z|x ) . This is a common latent-state model for stochastic processes ; see Figure 2 for an illustration . In the single-cell context , a meta-state is often referred to as a “ cell type ” , which has a distinct “ pathway ” ( i.e. , future dynamics ) . The “ cell type ” is defined as a function of the gene-expression profile , but the function is unknown and to be learnt . In this latent-state model , let Φ∗ ( x ) = pZ ( ·|x ) . Then we can verify that Assumption 1 holds . In this case , finding the embedding map Φ∗ would make it possible to recover the set of meta-states { z } and aggregation distributions pZ ( ·|x ) . 3 Mapping Gene Expressions To Low-Dimensional Cell States . Recall that our goal is to find mathematical abstractions of high-dimensional expression profiles { Xt } . We will estimate an embedding map from gene-expression profiles to a low-dimensional vector space : Ψ : x ∈ Rd 7→ Ψ ( x ) ∈ Rr . Ideally , we hope Ψ ( Xt ) to be low-dimensional while still containing as much information about Xt+1 as possible . 3.1 Embedding the transition operator into a functional space . Consider a kernel mean embedding of the transition operator P : Q = ΠHPΠH , where the projections are with respect to appropriate norms . By assumption , we can verify that rank ( Q ) < r. We seek to estimate Q from cell transition data , perform singular value truncation , and then find the low-dimensional embedding map by transforming the left single functions of Q . To guarantee the function space H is sufficiently expressive , we adopt a kernel composition approach for “ lifting ” the dimensions . First , we construct an initial kernel K0 that best fits the dataset ’ s topology and preserves neighborhood relations . To find such a K0 , we can leverage existing dimension reduction methods for single-cell data analysis , such as PCA , manifold-based , and graph-based methods ( see Kester & van Oudenaarden ( 2018 ) ; Tanay & Regev ( 2017 ) ; Bacher & Kendziorski ( 2016 ) ; Stegle et al . ( 2015 ) for reviews ) . Then , we construct the kernel function K = K0 ◦ K1 by taking the composition between K0 and another kernel function K1 ( e.g. , the Gaussian kernel ) - this step would further lift the problem ’ s dimension and improve the function space ’ s expressibility . One can also take compositions of multiple kernels to mimic a multi-layer neural network ( Cho & Saul , 2009 ) . 3.2 Low-rank compression of cell states via random features . We propose a kernelized state embedding method based on random feature projection for computing an estimator Ψ̂ from transition data { ( Xi , X ′i ) } Ni=1 , which can be obtained from cell trajectories . For analyzing single-cell sequencing data and embedding transition distributions , we will choose the function space H with the kernel function K tailored to the data ’ s geometry . Suppose we have chosen a kernel function K ( for example , a Gaussian kernel ) . We perform nonparametric estimation of Ψ∗ by generating a large number of random features to approximate the kernel space in large finite dimensions . Then , we downsize the estimator by using spectral decomposition . In the case where each parent cell has a single descendant , the cell2state method works by ( informally ) : ( 1 ) Generate random Fourier functions φ ( · ) = ( φ1 ( · ) , . . . , φd ( · ) ) > by randomized decomposition of its kernel function K to approximately span H ( Rahimi & Recht , 2008 ) . ( 2 ) Estimate a finite matrix embedding of the scaled condition probability distribution 1√ p ( x′ ) p ( x′|x ) by P̂ = Σ−1/20 ( 1 N ∑N i=1 φ ( Xi ) φ ( X ′ i ) > ) Σ −1/2 1 , where Σ0 , Σ1 are covariances at the two time points . ( 3 ) Let Ψ̂ ( · ) = ( ÛrΛ̂r ) > Σ−10 Φ ( · ) where Ûr , Λ̂r are from top r truncation of the SVD of P̂ . See Algorithm 1 in the Appendix for the full description of the cell2state algorithm , which also handles the case where cells have multiple descendants . Given a cell ’ s gene expression profile x , the vector Ψ̂ ( x ) can be viewed as a low-dimensional mean embedding of the transition function p ( ·|x ) . Thus it should be predictive of this cell ’ s future dynamics . Runtime Complexity of Algorithm 1 . The algorithm uses random features and singular value truncation , both designed for maximal computation efficiency . The overall runtime for training is at most O ( n+ nD2 +D3 ) , where n is number of cells , D is number of Fourier features ( ≤ n ) . This is the same complexity as computing covariances and PCA in the random feature space . After training , querying the embedding map Ψ̂ takes only O ( rD ) time . In our experiments , Algorithm 1 runs in seconds , while training an MLP ( a deep neural network ) for cell fate prediction takes 10-15 minutes .

Summary: The authors introduce cell2state, an algorithm that incorporates both genetic barcoding coupled with single-cell sequenced data to model explicit state transitions of cell dynamics over time. Single-cell gene expression profiles are mapped to low-dimensional state vectors that are predictive of cell dynamics. Cell2state is evaluated using barcoded stem cell dataset (Biddy et al. (2018)) and simulation studies are also presented. The model demonstrates better results for cell state prediction, finding dynamically stable clusters, and reveals potential latent meta-states of the underlying cellular evolution process.