Paper deep dive

KoopCell: Koopman-Based Generative Model for Learning Single-Cell Dynamics from Distribution Snapshots

Wanfeng Lu, Yutong Zhang, Keyi Zhou, Chenxin Ge, Wei Lin, Qunxi Zhu

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 241

Intelligence

Status: not_run | Model: - | Prompt: - | Confidence: 0%

Entities (0)

No extracted entities yet.

Relation Signals (0)

No relation signals yet.

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract:Learning population dynamics from temporally sparse, unpaired distribution snapshots is a fundamental challenge in developmental biology. Recent approaches based on neural differential equations and flow matching can interpolate between observed population snapshots, but may struggle to extrapolate beyond the training horizon and often lack an explicit mechanism for modeling developmental branching. We propose KoopCell, a unified generative framework based on Koopman-Mori-Zwanzig theory that jointly learns representations and predictive linear latent dynamics. Theoretically, using the weak continuity equation, we derive a closed-form least-squares estimator for the Koopman generator from distribution snapshots and establish convergence guarantees under suitable assumptions. To model branching dynamics, we further develop KoopCell-M, which incorporates non-Markovian memory into the latent Koopman dynamics through a Markovian embedding. Experiments on synthetic systems and three scRNA-seq datasets demonstrate the ability of our framework to recover Koopman spectra, model branching through memory, and scale to predicting high-dimensional gene expression distributions, achieving state-of-the-art performance among the evaluated methods.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

238,144 characters extracted from source content.

Expand or collapse full text

KoopCell: Koopman-Based Generative Model for Learning Single-Cell Dynamics from Distribution Snapshots Wanfeng Lu1,2, Yutong Zhang2,11footnotemark: 1 Keyi Zhou2 Chenxin Ge2 Wei Lin1,2,3,4 Qunxi Zhu2,3,4, 1School of Mathematical Sciences, SCMS, SCAM, and CCSB, Fudan University, China 2Research Institute of Intelligent Complex Systems, Fudan University, China 3State Key Laboratory of Medical Neurobiology and MOE Frontiers Center for Brain Science, Institutes of Brain Science, Fudan University, China 4Shanghai Artificial Intelligence Laboratory, China †thanks: Equal contribution.†thanks: Corresponding author: qxzhu@fudan.edu.cn. Abstract Learning population dynamics from temporally sparse, unpaired distribution snapshots is a fundamental challenge in developmental biology. Recent approaches based on neural differential equations and flow matching can interpolate between observed population snapshots, but may struggle to extrapolate beyond the training horizon and often lack an explicit mechanism for modeling developmental branching. We propose KoopCell, a unified generative framework based on Koopman–Mori–Zwanzig theory that jointly learns representations and predictive linear latent dynamics. Theoretically, using the weak continuity equation, we derive a closed-form least-squares estimator for the Koopman generator from distribution snapshots and establish convergence guarantees under suitable assumptions. To model branching dynamics, we further develop KoopCell-M, which incorporates non-Markovian memory into the latent Koopman dynamics through a Markovian embedding. Experiments on synthetic systems and three scRNA-seq datasets demonstrate the ability of our framework to recover Koopman spectra, model branching through memory, and scale to predicting high-dimensional gene expression distributions, achieving state-of-the-art performance among the evaluated methods. 1 Introduction Inferring population dynamics from temporally sparse, unpaired distribution snapshots is a fundamental challenge in machine learning and the natural sciences, with applications in biology (Schiebinger et al., 2019; Moon et al., 2019), ecology (Ayala et al., 1973), and epidemiology (Wang et al., 2021; Kosma et al., 2023). In particular, advances in high-throughput single-cell RNA sequencing (scRNA-seq) (Hwang et al., 2018; Chen et al., 2019) have made this problem increasingly important in single-cell biology. By profiling population-level gene expression at multiple time points, scRNA-seq enables the study of cellular developmental dynamics, the identification of key driver genes, and the prediction of future states. However, due to the destructive nature of scRNA-seq measurements, only a cross-sectional sample from the population-level gene-expression distribution is available, without cell-level correspondence across time. Inferring the underlying continuous population dynamics from such observations therefore remains a substantial computational challenge. Neural dynamical systems have been widely used to model cellular dynamics from population snapshots. Huguet et al. (2022); Zhang et al. (2024) use neural ordinary differential equations (ODEs) (Chen et al., 2018), whereas Yeo et al. (2021); Jiang and Wan (2024) employ neural stochastic differential equations (SDEs) to capture biological stochasticity in latent space. More recently, Zhang et al. (2025); Wang et al. (2025); Ling et al. (2026) have introduced unbalanced dynamical optimal transport to model cell proliferation and death. Although these methods perform well under their respective dynamical assumptions and can interpolate missing distributions within the observed time horizon, they do not effectively extrapolate beyond the observed time points or explicitly account for branching dynamics. Both capabilities are essential for predicting future cell-population states and representing the lineage differentiation that is crucial for understanding cellular development. In this work, we propose KoopCell, a Koopman-based generative model for population dynamics. Koopman theory (Koopman, 1931; Koopman and von Neumann, 1932) provides a unifying perspective for modeling nonlinear dynamics by lifting the system state into a space of observables (functions of the system state), in which the evolution is governed by a linear operator. Neural Koopman methods (Lusch et al., 2018) approximate these observables with an encoder and represent their evolution using a finite-dimensional linear operator, but they generally rely on paired state transitions or trajectories. In contrast, KoopCell jointly learns the observable representation and its linear dynamics directly from temporally sparse, unpaired distribution snapshots. We theoretically formulate the estimation of the latent linear propagator through the weak continuity equation (Lu et al., 2024) and integrate it with a generative VAE model in an alternating learning procedure, enabling population prediction both within and beyond the observed time horizon. A single Markovian linear flow, however, may be insufficient to represent branching processes such as the differentiation of stem cells into multiple specialized cell types. We therefore extend KoopCell using the Mori–Zwanzig formalism, which accounts for the influence of unresolved variables through non-Markovian memory (Mori, 1965; Zwanzig, 1973). By embedding the suggested memory dynamics into an augmented latent space (Hijón et al., 2010; Menier et al., 2025), we develop tractable variational inference for the hidden memory and a corresponding alternating learning algorithm for the memory-augmented model, KoopCell-M. This memory mechanism enables KoopCell-M to model the divergence of a population into distinct developmental branches. Our main contributions are summarized as follows: • We develop a theoretical and computational framework for identifying a finite-dimensional Koopman representation from unpaired distribution snapshots. Based on the weak continuity equation, we characterize the closed-form estimator, establish its convergence properties, and jointly learn the associated observable representation within a generative model. • Furthermore, we introduce KoopCell-M, a Mori–Zwanzig-inspired extension that represents non-Markovian memory in an augmented latent space, providing a mechanistic formulation and a practical learning algorithm for population dynamics with branching. • We evaluate KoopCell and KoopCell-M on synthetic toy systems and three high-dimensional scRNA-seq datasets. The experiments validate Koopman spectral recovery and the ability to model branching dynamics, while demonstrating state-of-the-art extrapolation performance in predicting cellular population distributions from sparse, unpaired measurements. 2 Preliminaries Problem Formulation and Notation. We consider a time series of data distributions ρt _t (t∈⊂ℝ+t ^+), from each of which we observe ntn_t independent samples (e.g., cells) t,ii=1nt\ x_t,i\_i=1^n_t, with t,i∼ρt x_t,i _t and t,i∈=ℝG x_t,i =R^G. We aim to learn the population dynamics that explain the evolution of these distribution snapshots. A typical application is to recover cellular developmental dynamics from time-series scRNA-seq data, where cells sampled from ρt _t form a gene expression matrix t∈ℝnt×G X_t ^n_t× G, and G is the number of genes. Given the high dimensionality of the observations, a common approach is to model temporal dynamics in a D-dimensional latent space =ℝDZ=R^D, with D≪GD G. Rather than specifying a heuristic dynamical structure in the latent space, we use Koopman theory to guide the learning of latent coordinates with approximately linear evolution. Throughout, we use uppercase state symbols for random variables and lowercase symbols for their realizations. Koopman Theory and Koopman Learning. For an autonomous system ˙=⁡() x= f( x) with flow FtF_t, an observable is a function g:→ℝg:X of the system state. For non-autonomous systems, augmenting the state with time (t˙=1 t=1) yields an autonomous formulation. The Koopman operator evolves observables by [U⁡(t)g]()=g⁡(Ft)[U(t)g]( x)=g(F_t x) and is linear even when the state-space dynamics are nonlinear (Koopman, 1931; Koopman and von Neumann, 1932). We assume that the operators U⁡(t):ℋ→ℋU(t):H form a strongly continuous semigroup on a Hilbert space ℋH of observables, with generator L and domain D⁡(L)D(L). For sufficiently smooth g∈D⁡(L)g∈ D(L), Lg=⋅∇gLg= f·∇ g and dtU(t)g=U⁡(t)Lg ddtU(t)g=U(t)Lg. Let t=Ft0 X_t=F_t X_0 with 0∼ρ0 X_0 _0. If ℋD=span⁡g1,…,gD⊂D⁡(L)H_D=span\g_1,…,g_D\⊂ D(L) is invariant under L, then L=L g= A g for a matrix A, so t=⁡(t) Z_t= g( X_t) satisfies ˙t=t Z_t= A Z_t, where =(g1,…,gD)⊤ g=(g_1,…,g_D) . With the observable map g fixed, extended dynamic mode decomposition (EDMD) (Williams et al., 2015) estimates a propagator from paired states (i,i)( x_i, y_i), i=FΔi y_i=F_ x_i, by solving a regression problem min⁡∑i⁡‖⁡(i)−(i)‖22 _ K _i\| g( y_i)- K g( x_i)\|_2^2. Koopman representation learning extends EDMD by learning the observable map together with the latent dynamics. Specifically, an encoder ↦=⁡()∈ℝD x z= g( x) ^D is jointly trained with a linear latent propagator and a decoder to obtain coordinates that support both approximately linear evolution and state reconstruction (Otto and Rowley, 2019; Lusch et al., 2018). Our work extends this framework to learning generative latent linear dynamics from temporally sparse, unpaired distribution snapshots. Mori–Zwanzig Formalism and Memory Embedding. The reduced finite-dimensional linear dynamics above are exact when the span of the chosen observables is invariant under the Koopman generator L. When this linear closure fails, the evolution of the resolved coordinates t=⁡(t)∈ℝD Z_t= g( X_t) ^D generally involves unresolved observables outside the chosen observable subspace ℋDH_D. The Mori–Zwanzig formalism, originally developed in statistical mechanics, accounts for their influence through memory and forcing terms (Mori, 1965; Zwanzig, 1973). Specifically, under suitable regularity, Mori’s linear projection onto ℋDH_D yields a generalized Langevin equation (Lin et al., 2021): ˙t=t+∫0t⁡(t−s)ss+t. Z_t= A Z_t+ _0^t K(t-s) Z_s\,ds+ η_t. Here A describes the instantaneous linear contribution, K is the memory kernel, and t η_t carries the influence of unresolved initial conditions. A tractable closure approximates the kernel by ⁡(u)≈eu K(u)≈ Be u C and represents the unresolved influence through auxiliary hidden memory variables t∈ℝdh H_t ^d_h (Lei et al., 2016; Menier et al., 2025). This leads to the augmented dynamics dt=(t+t)dtd Z_t=( A Z_t+ B H_t)\,dt and dt=(t+t)dt+hdtd H_t=( C Z_t+ H_t)\,dt+ _h\,d W_t, where t W_t is standard Brownian motion and h _h controls the hidden noise. Eliminating t H_t yields the kernel ⁡(u)=eu K(u)= Be u C and a temporally correlated forcing determined by 0 H_0 and the hidden noise. The augmented state (t,t)( Z_t, H_t) thus evolves as a Markovian system, while t Z_t alone generally retains memory. This Markovian embedding provides the underlying mechanism for modeling branching dynamics in our framework. Wasserstein Flow and Optimal Transport. Time series of distribution snapshots are often modeled as probability flows in Wasserstein space. We write 2(ℝD)P_2(R^D) for the space of probability measures with finite second moment and W2W_2 for the Wasserstein-2 distance, defined by W22(μ,ν)=infπ∈Π⁡(μ,ν)∫‖−‖2π(,)W_2^2(μ,ν)= _π∈ (μ,ν) \| z- y\|^2dπ( z, y), where Π⁡(μ,ν) (μ,ν) is the set of couplings. We use μ⁡[ψ]=∫ψ⁡()μ()μ[ψ]= ψ( z)\,dμ( z) and write F#μF_\#μ for the distribution of F⁡()F( Z) when ∼μ Z μ. The Wasserstein flow t↦μt _t, with μ⋅∈AC2([0,T],2(ℝD)) _·∈ AC^2([0,T];P_2(R^D)), is an absolutely continuous curve in Wasserstein space that describes the evolution of probability measures over time. For almost every t, there exists a unique canonical velocity t∈Tμt2(ℝD) v_t∈ T_ _tP_2(R^D)11 1 Here, Tμ2=∇φ:φ∈Cc∞(ℝD)¯L2(μ)T_μP_2= \∇ : ∈ C_c^∞(R^D)\^L^2(μ) is the tangent space; the overline denotes closure in L2(μ)L^2(μ). satisfying ∂tμt+∇⋅(tμt)=0 _t _t+∇·( v_t _t)=0 in the distributional sense (Ambrosio et al., 2008). Among all velocity fields generating the same measure curve, t v_t has minimum L2(μt)L^2( _t) norm (∫‖t‖2dμt)1/2 ( \| v_t\|^2d _t )^1/2. We denote by Pμ:L2(μ,ℝD)→Tμ2(ℝD)P_μ:L^2(μ;R^D)→ T_μP_2(R^D) the L2(μ)L^2(μ)-orthogonal projection onto this tangent space. For a zero-mass distribution σ, we define the negative Sobolev norm ‖σ‖−1,μ=sup∫‖∇φ‖2μ≤1⟨σ,φ⟩\|σ\|_-1,μ= _ \|∇ \|^2dμ≤ 1 σ, . Here ⟨σ,φ⟩ σ, denotes the duality pairing of a distribution with a test function. Equivalently, when finite, ‖σ‖−1,μ\|σ\|_-1,μ is the minimum L2(μ)L^2(μ) norm of a tangent field w satisfying σ=−∇⋅(μ)σ=-∇·( wμ); see Lemma A.2 in Appendix A.1. 3 Methods 3.1 Koopman-based VAE for Generative Modeling of Population Dynamics Motivated by Koopman theory, we jointly learn an encoder that maps cellular states to observables with approximately linear evolution, a linear propagator governing their dynamics, and a decoder reconstructing the cellular states from latent representations. Specifically, we use a VAE (Kingma and Welling, 2014), widely used for single-cell gene expression modeling (Grønbech et al., 2020), with encoder q(⋅∣,t)=((,t),(,t))q_ θ(· x,t)=N( m_ θ( x,t), S_ θ( x,t)) and decoder DecϕDec_ φ, where ∈ℝD z ^D. We concatenate t with x to encode the augmented state. Including the constant observable 11 yields affine dynamics ˙t=t+z Z_t= A Z_t+ b_z. For simplicity, we use ˙t=t Z_t= A Z_t below and describe the detailed implementation in Appendix B.1. Unlike traditional settings where the Koopman model is learned from densely sampled, paired trajectories, we are given only temporally sparse, unpaired distribution snapshots. We address the challenge of learning the Koopman operator from snapshots through the following two alternating steps: (1) with the observables (latent coordinates) fixed, we estimate the linear matrix that best fits the evolution of the latent distributions, providing an analogue of EDMD estimation for population snapshots; (2) with the estimated linear propagator fixed, we refine the latent representation to better fit the dynamics. By alternating between propagator estimation and representation refinement, we ultimately learn the Koopman representation and its linear dynamics generator simultaneously. Estimation of the linear propagator. We fix a set of observables g1,…,gD\g_1,…,g_D\, which induces a Wasserstein flow of latent distributions μt=#ρt∈2(ℝD) _t= g_\# _t _2(R^D), t∈[0,T]t∈[0,T]. We aim to find a linear generator A such that (eΔt)#μt≈μt+Δt(e A t)_\# _t≈ _t+ t. Exact propagation under this linear flow satisfies the continuity equation ∂tμt+∇⋅(μt)=0 _t _t+ _ z·( A z\, _t)=0. Accordingly, we introduce the weak population residual r(t)=∂tμt+∇⋅(μt)r_ A(t)= _t _t+ _ z·( A z\, _t) to quantify the mismatch between the linear dynamics and the evolution of μt _t. We then seek an A minimizing μ()=12∫0T‖r(t)‖−1,μt2tJ_μ( A)= 12 _0^T\|r_ A(t)\|_-1, _t^2dt. However, directly minimizing the residual is intractable, as it requires the density derivative ∂tμt _t _t. To this end, we use the equivalent min–max formulation in the dual test space, as established in Theorem A.3 (Appendix A.1): ∗=arg⁡minμ()=arg⁡min⁡supφ∈D⁡ℛD(,φ)−12‖φ‖D2,where ℛD(,φ)=μT[φ⁡(T,⋅)]−μ0[φ⁡(0,⋅)]−∫0Tμt[∂tφ⁡(t,)+∇φ(t,)⊤]dt. aligned & A^*=arg _ AJ_μ( A)=arg _ A _ _D \R_D( A; )- 12\| \|_V_D^2 \,\\ where &R_D( A; )= _T[ (T,·)]- _0[ (0,·)]- _0^T _t\! [ _t (t, z)+ _ z (t, z) A z ]dt. aligned (1) Here D=L2(0,T,H˙1(μt))V_D=L^2(0,T; H^1( _t)) is the Hilbert space with ‖φ‖D2=∫0Tμt[‖∇φ‖2]t\| \|_V_D^2= _0^T _t[\| _ z \|^2]dt. The following theorem shows that minimizing μJ_μ gives the best approximation of the canonical Wasserstein velocity t v_t by projected linear fields in L2(dtdμt)L^2(dt\,d _t). Furthermore, restricting the test space DV_D to finitely many functions yields a least-squares problem which admits a closed-form solution. Theorem 3.1. Assume μ⋅∈AC2([0,T],2(ℝD)) _·∈ AC^2([0,T];P_2(R^D)), and let t v_t be its canonical velocity. Then, for every ∈ℝD×D A ^D× D, μ()=supφ∈DℛD(,φ)−12‖φ‖D2=12∫0T‖t−Pμt()‖L2(μt)2t.J_μ( A)= _ _D \R_D( A; )- 12\| \|_V_D^2 \= 12 _0^T\| v_t-P_ _t( A z)\|_L^2( _t)^2dt. (2) Hence the minimization gives the orthogonal best approximation of the canonical velocity by projected linear velocity fields. Furthermore, restrict the supremum to D,m=span⁡φ1,…,φmV_D,m=span\ _1,…, _m\ and define yj=μT[φj(T)]−μ0[φj(0)]−∫0Tμt[∂tφj(t)]dt,j=∫0Tμt[∇φj(t,)⊤]dt,Mij=∫0Tμt[∇φi(t,)⊤∇φj(t,)]dt. gatheredy_j= _T[ _j(T)]- _0[ _j(0)]- _0^T _t[ _t _j(t)]dt,\\ H_j= _0^T _t[ _ z _j(t, z) z ]dt, M_ij= _0^T _t[ _ z _i(t, z) _ z _j(t, z)]dt. gathered (3) Let =vec⁡() a=vec( A), let row j of G be vec⁡(j)⊤vec( H_j) , and let =(y1,…,ym)⊤ y=(y_1,…,y_m) . If ≻0 M 0, restricting Equation (1) to D,mV_D,m gives the weighted least-squares objective D,m()=12(−)⊤−1(−)J_D,m( A)= 12( y- G a) M^-1( y- G a). If G has full column rank, its unique minimizer is vec⁡(^D,m)=(⊤−1)−1⊤−1vec( A_D,m)=( G M^-1 G)^-1 G M^-1 y. Finally, under the test-enrichment and uniqueness assumptions in Appendix A.2, these minimizers converge to ∗ A^*, the unique population minimizer of Equation (1) as m→∞m→∞. We provide the proof of Theorem 3.1 in Appendices A.1 and A.2. In practice, we use random Fourier tests to construct D,mV_D,m, reducing the min–max optimization to a weighted least-squares problem. For irregular snapshots at 0=t0<t1<⋯<tN=T0=t_0<t_1<·s<t_N=T, we formulate weak equations on each adjacent interval Ik=[tk,tk+1]I_k=[t_k,t_k+1] and approximate the integrals by the trapezoidal rule, ∫Ikf⁡(t)t≃tk+1−tk2f⁡(tk)+f⁡(tk+1) _I_kf(t)dt t_k+1-t_k2\f(t_k)+f(t_k+1)\. The integrands f(⋅)=μ⋅[⋅]f(·)= _·[·] are expectations under the latent distributions, approximated by empirical averages over the encoded snapshots. Finally, equations over different intervals are combined to form the overall regression problem, with ridge regularization used to stabilize the solve. Implementation details are provided in Appendix B.1.2. We further show in a Koopman-invariant RKHS that the minimum weak population residual vanishes as D→∞D→∞ (Appendix A.3.2), derive Wasserstein prediction bounds (Appendix A.4.1), and prove locally uniform strong semigroup convergence (Appendix A.4.4), under the stated approximation, stability, and identifiability assumptions. We also discuss theoretical limitations in Appendix A.4.4. Updating the representation under the linear flow. Given the estimated propagator =^D,m A= A_D,m, we update the representation to better fit the linear dynamics. To regularize the VAE representation, we use a prior transported by the linear flow: the initial prior p0=⁡(,D)p_0=N( 0, I_D) evolves into pt=⁡(,ete⊤t)p_t=N( 0,e Ate A t). We thus combine squared reconstruction error with the KL divergence (Csiszár, 1975) between the encoder posterior and this dynamic prior: ℒVAE=t,[12σx2dx∼q(⋅∣,t)∥−Decϕ()∥2+βKL(q(⋅∣,t)∥pt)].L_VAE=E_t, x [ 12 _x^2d_xE_ Z q_ θ(· x,t)\| x-Dec_ φ( Z)\|^2+ (q_ θ(· x,t)\,\|\,p_t ) ]. (4) Here dxd_x is the expression dimension, σx2 _x^2 is the fixed observation variance, and t,E_t, x averages over training times and cells. To encourage accurate population prediction in both latent and expression space, we augment ℒVAEL_VAE with the population-matching loss ℒpop=s<t∈[λzε(μt,μ^s→t)+λxε(ρt,ρ^s→t)]L_pop=E_s<t [ _zS_ ( _t, μ_s→ t)+ _xS_ ( _t, ρ_s→ t)], where μt=∫q(⋅∣,t)dρt() _t= q_ θ(· x,t)d _t( x) is the encoder-induced population, μ^s→t=(e⁡(t−s))#μs μ_s→ t=(e A(t-s))_\# _s is the latent distribution propagated from μs _s to time t, and ρ^s→t=(Decϕ)#μ^s→t ρ_s→ t=(Dec_ φ)_\# μ_s→ t is the decoded prediction. The expectation averages over all forward pairs of training times, and εS_ is the debiased Sinkhorn divergence with quadratic cost. We further regularize the encoder with the weak population residual loss ℒweak=D,m(^D,m)L_weak=J_D,m( A_D,m), which penalizes violations of the weak continuity equation and encourages latent coordinates in which population evolution is better approximated by the estimated linear dynamics (Appendix B.1.2). The linear flow is evaluated by matrix exponentials, enabling prediction at arbitrary times without costly numerical ODE integration. Training. Following Zhang et al. (2024), we pretrain the VAE with a standard-normal Gaussian prior, then alternate between generator estimation and representation learning. During each representation update, we hold A fixed and minimize ℒKC=ℒVAE+ℒpop+λweakℒweakL_KC=L_VAE+L_pop+ _weakL_weak over the VAE parameters ,ϕ θ, φ. Algorithm 1 summarizes the training procedure of KoopCell. 3.2 KoopCell-M: Memory-Augmented Koopman Dynamics for Branching Learning a linear latent flow is appealing, but it constrains the modeling of branching during cellular differentiation, in which an initially homogeneous population differentiates into multiple cell types. Specifically, an ODE with a Lipschitz vector field cannot exactly transport a source distribution supported on M disjoint, compact, path-connected components to a target with N≠MN≠ M such components (Zhao et al., 2026, Proposition 3.3). Thus, when the initial population has connected latent support, the linear dynamics followed by a Lipschitz decoder cannot produce fate populations with disconnected support. Moreover, the linear flow itself cannot create new latent density modes. Although nonlinear decoding can introduce additional modes, bounding the decoder Lipschitz constant limits the amplification of latent separation (Appendix B.2.1). To model branching through an explicit dynamical mechanism, Zhao et al. (2026) introduce non-Markovian delay dynamics with distinct initial histories for the same state, allowing trajectories to diverge. Such non-Markovian dynamics also arise within Koopman theory through the Mori–Zwanzig formalism (Section 2), which describes how unresolved observables contribute memory to the resolved linear dynamics. In cellular differentiation, such unresolved influences may reflect heritable fate biases that are not fully captured by scRNA-seq (Weinreb et al., 2020). This connection motivates us to introduce memory into KoopCell, using non-Markovian dynamics to model branching within the same Koopman framework. Memory embedding and branching. To obtain a tractable formulation of this non-Markovian model, we use the Markovian embedding from Section 2 to evolve the augmented latent state (t,t)( Z_t, H_t): dt=(t+t+z)dt,dt=(t+t+h)dt+hdt.d Z_t=( A Z_t+ B H_t+ b_z)\,dt, H_t=( C Z_t+ H_t+ b_h)\,dt+ _h\,d W_t. (5) Here t∈ℝD Z_t ^D is the visible latent state encoded from cellular gene expression, and t∈ℝdh H_t ^d_h contains auxiliary memory variables. A joint Gaussian initialization of (0,0)( Z_0, H_0) yields Gaussian visible Z-marginals. To enable branching in the latent distribution, we parameterize hidden initial distributions by a discrete variable R∈ℛR∈ R. Each value of R indexes a candidate root-to-terminal path in a directed acyclic graph (DAG) of coarse biological states. We detail the state annotations and graph construction in Appendix B.2.1 and show the resulting DAG for the mouse cell reprogramming dataset (SC; Schiebinger et al., 2019) in Figure 5. With R specifying the developmental path, we define a generative model with latent variables (R,t,t)(R, Z_t, H_t). We specify the initial distribution of the augmented dynamics in Equation (5) as R∼Cat(),0∣R=r∼(,D),0∣0=,R=r∼(r+r,r).R ( π), Z_0 R=r ( 0, I_D), H_0 Z_0= z,R=r ( F_r z+ a_r, S_r). (6) Under this parametrization, the augmented process is Gaussian conditional on R. The marginal distribution of t Z_t is therefore a Gaussian mixture, pt()=∑r∈ℛπr(,t,rz,t,rzz)p_t( z)= _r∈ R _rN( z; m^z_t,r, V^z_t,r), with propagated means t,rz m^z_t,r and covariances t,rzz V^z_t,r (Appendix B.2.2). Although all paths share the same visible initial distribution p0p_0, coupling through t B H_t can turn hidden heterogeneity into visible separation, producing multiple density modes when component means become sufficiently separated relative to their spread. We construct a toy example in Section 4.2 to illustrate how introducing memory enables the dynamics to generate distinct branches from an initially unimodal population (Figure 2). Learning the representation and memory dynamics. As in KoopCell, we first pretrain the VAE, then alternate between dynamics estimation and representation learning. With the representation and memory statistics fixed, we estimate ,z A, b_z in closed form using the same least-squares formulation, after subtracting the memory contribution from the weak continuity equations (Appendix B.2.3). With ,z A, b_z fixed, we jointly update the representation and remaining memory parameters using the distribution-matching objective ℒpopL_pop over all training pairs s<ts<t, as in KoopCell. Prediction from a snapshot at s additionally requires inferring R and s H_s. Our generative model supports tractable variational inference: given ∼q(⋅∣,s) z q_ θ(· x,s), we compute the path probabilities qs(r∣,)q_s(r z, x) and Gaussian memory conditionals p⁡(s∣s=,R=r)p( H_s Z_s= z,R=r) analytically (Appendix B.2.2). We propagate the inferred augmented population under Equation (5), with transition means and covariances evaluated by matrix exponentials. The predicted t Z_t-marginal is μ^s→t μ_s→ t, and decoding gives ρ^s→t=(Decϕ)#μ^s→t ρ_s→ t=(Dec_ φ)_\# μ_s→ t. To encourage rollouts to follow the biological pathway specified by R, we introduce a state classifier c(c∣)c_ ω(c z). We jointly train the VAE and classifier using the snapshot objective ℒsnapL_snap, which combines reconstruction, KL regularization, and supervised classification on annotated training cells (Appendix B.2.1). We then evaluate c_ ω along each rollout and use the predicted state probabilities to compute the path consistency loss ℒpathL_path, which encourages progression through biological states in the order specified by R (Equation (77) in Appendix B.2.2). The overall training objective is ℒsnap+ℒpop+λRℒpath+λweakℒweakL_snap+L_pop+ _RL_path+ _weakL_weak. Appendix B.2.2 derives these objectives, and Algorithm 2 summarizes the alternating learning procedure. 4 Experiments We evaluate our models on synthetic systems and three single-cell RNA-sequencing datasets. Through these experiments, we answer the following three questions: Q1. Can KoopCell recover the underlying Koopman spectrum from sparse, unpaired snapshots, and does this recovery support stable extrapolation? Q2. How does incorporating memory improve the modeling of branching dynamics? Q3. Can our models scale to high-dimensional scRNA-seq data and learn from sparse, irregular snapshots to accurately predict gene expression distributions within and beyond the training horizon? Table 1: Hard-task Wasserstein-2 distance (↓ ) at held-out time points (mean ± sample standard deviation over three training seeds). Shaded labels distinguish interpolation from extrapolation within each dataset. The best and second-best means are shaded green (bold) and blue, respectively. Prediction Data t scNODE LGP-OT MIOFlow PRESCIENT PI-SDE VGFM KoopCell KoopCell-M ZB 3 577.71±3.69577.71± 3.69 577.91±0.64577.91± 0.64 594.71±32.03594.71± 32.03 597.49±1.76597.49± 1.76 567.30±1.41567.30± 1.41 575.86±3.22575.86± 3.22 572.32±2.14572.32± 2.14 588.91±6.04588.91± 6.04 ZB 5 511.58±5.63511.58± 5.63 504.45±1.73504.45± 1.73 519.32±7.48519.32± 7.48 540.72±1.52540.72± 1.52 492.43±2.19492.43± 2.19 499.66±0.81499.66± 0.81 497.80±3.83497.80± 3.83 488.84±9.27488.84± 9.27 ZB 7 440.65±6.47440.65± 6.47 411.20±3.07411.20± 3.07 452.30±0.77452.30± 0.77 476.68±1.16476.68± 1.16 416.80±3.36416.80± 3.36 434.99±1.37434.99± 1.37 412.99±1.50412.99± 1.50 408.79±1.57408.79± 1.57 Interpolation ZB 9 514.24±5.97514.24± 5.97 475.65±2.60475.65± 2.60 539.92±11.44539.92± 11.44 527.09±1.25527.09± 1.25 474.41±3.51474.41± 3.51 522.60±2.83522.60± 2.83 474.92±1.55474.92± 1.55 468.53±2.72468.53± 2.72 ZB 11 653.56±8.98653.56± 8.98 605.14±1.78605.14± 1.78 654.64±2.17654.64± 2.17 656.41±4.17656.41± 4.17 612.74±3.51612.74± 3.51 635.96±2.73635.96± 2.73 595.82±0.48595.82± 0.48 588.73±0.77588.73± 0.77 Extrapolation ZB 12 719.28±12.36719.28± 12.36 665.78±4.77665.78± 4.77 720.67±17.25720.67± 17.25 700.68±4.58700.68± 4.58 679.61±4.11679.61± 4.11 720.54±10.97720.54± 10.97 648.91±2.47648.91± 2.47 639.78±0.72639.78± 0.72 DR 3 443.40±5.78443.40± 5.78 433.15±2.43433.15± 2.43 446.21±8.45446.21± 8.45 475.42±1.27475.42± 1.27 432.23±2.64432.23± 2.64 455.79±6.96455.79± 6.96 429.70±3.54429.70± 3.54 430.41±3.72430.41± 3.72 DR 5 468.39±0.68468.39± 0.68 465.22±1.10465.22± 1.10 468.02±1.97468.02± 1.97 492.41±1.06492.41± 1.06 457.57±2.85457.57± 2.85 466.01±0.99466.01± 0.99 463.32±0.60463.32± 0.60 466.04±1.46466.04± 1.46 Interpolation DR 7 534.69±8.30534.69± 8.30 519.13±0.49519.13± 0.49 533.38±1.03533.38± 1.03 541.43±0.77541.43± 0.77 519.12±0.67519.12± 0.67 526.89±0.67526.89± 0.67 522.50±2.41522.50± 2.41 518.49±1.41518.49± 1.41 DR 9 602.99±5.41602.99± 5.41 567.25±0.94567.25± 0.94 617.05±5.86617.05± 5.86 604.33±0.80604.33± 0.80 577.97±0.83577.97± 0.83 592.17±5.07592.17± 5.07 558.89±2.21558.89± 2.21 556.80±1.27556.80± 1.27 DR 10 593.05±2.68593.05± 2.68 570.52±5.15570.52± 5.15 673.08±24.93673.08± 24.93 565.45±0.37565.45± 0.37 603.87±2.94603.87± 2.94 653.24±19.74653.24± 19.74 563.86±3.26563.86± 3.26 557.68±3.56557.68± 3.56 Extrapolation DR 11 726.16±5.89726.16± 5.89 713.37±4.70713.37± 4.70 860.26±87.00860.26± 87.00 716.82±0.96716.82± 0.96 751.16±5.70751.16± 5.70 851.07±22.16851.07± 22.16 694.53±4.21694.53± 4.21 699.49±4.56699.49± 4.56 SC 6 54.44±0.2854.44± 0.28 48.66±0.4748.66± 0.47 53.87±4.2953.87± 4.29 64.70±0.8864.70± 0.88 47.35±0.0447.35± 0.04 48.94±0.3448.94± 0.34 50.20±1.5850.20± 1.58 49.49±0.6349.49± 0.63 SC 8 58.45±0.9458.45± 0.94 52.79±0.2252.79± 0.22 63.28±3.7163.28± 3.71 66.35±0.5266.35± 0.52 53.03±0.2653.03± 0.26 53.68±0.8053.68± 0.80 54.00±0.6854.00± 0.68 53.44±0.2853.44± 0.28 SC 10 103.23±1.21103.23± 1.21 96.40±0.2896.40± 0.28 110.39±5.17110.39± 5.17 108.01±0.22108.01± 0.22 95.14±1.3995.14± 1.39 108.46±1.37108.46± 1.37 96.72±0.3596.72± 0.35 96.99±0.3296.99± 0.32 Interpolation SC 12 142.30±1.36142.30± 1.36 128.31±1.26128.31± 1.26 156.91±7.55156.91± 7.55 152.00±2.76152.00± 2.76 137.35±0.78137.35± 0.78 166.69±3.40166.69± 3.40 127.47±0.92127.47± 0.92 127.33±0.37127.33± 0.37 SC 16 128.25±3.90128.25± 3.90 108.53±1.37108.53± 1.37 160.34±21.35160.34± 21.35 134.07±0.64134.07± 0.64 125.81±1.99125.81± 1.99 159.34±1.88159.34± 1.88 109.78±1.73109.78± 1.73 105.61±1.48105.61± 1.48 SC 17 142.40±4.78142.40± 4.78 114.45±1.06114.45± 1.06 195.20±44.48195.20± 44.48 142.04±1.57142.04± 1.57 147.12±3.60147.12± 3.60 161.57±2.74161.57± 2.74 118.54±4.86118.54± 4.86 111.87±0.64111.87± 0.64 SC 18 131.79±5.62131.79± 5.62 118.72±3.29118.72± 3.29 203.72±60.02203.72± 60.02 131.42±1.30131.42± 1.30 140.50±5.75140.50± 5.75 169.55±4.90169.55± 4.90 109.89±5.23109.89± 5.23 108.69±2.28108.69± 2.28 Extrapolation SC 19 145.91±8.78145.91± 8.78 132.99±4.67132.99± 4.67 243.97±92.21243.97± 92.21 136.76±1.76136.76± 1.76 160.55±7.53160.55± 7.53 181.55±9.57181.55± 9.57 121.87±8.65121.87± 8.65 117.16±3.84117.16± 3.84 4.1 Spectral Recovery from Unpaired Snapshots To test whether KoopCell can recover the underlying Koopman operator from unpaired snapshots, and whether the recovered spectral structure translates into more accurate temporal prediction, we consider two toy systems with known Koopman spectra that are widely studied in the literature: (Toy I) Following Tu et al. (2014), we consider x˙1=−0.2x1 x_1=-0.2x_1 and x˙2=−(x2−x12) x_2=-(x_2-x_1^2). The system has the attracting slow manifold x2=53x12x_2= 53x_1^2 and admits the exact Koopman observables Ψ1=x1 _1=x_1, Ψ2=x2−53x12 _2=x_2- 53x_1^2, and Ψ3=x12 _3=x_1^2. Consequently, ˙= = A with =diag⁡(−0.2,−1,−0.4) A=diag(-0.2,-1,-0.4). (Toy I) Following Otto and Rowley (2019), we consider the unforced Duffing equation q¨=−δq˙−q⁡(β+αq2) q=-δ q-q(β+α q^2) with δ=0.5δ=0.5, β=−1β=-1, and α=1α=1. The system has two stable equilibria at q=±1q=± 1, with the conjugate intrawell eigenpair λ1,2=(−1±31i)/4 _1,2=(-1± 31i)/4. Across the two wells, it also admits an eigenfunction with zero eigenvalue that separates the invariant basins. We generate 2D unpaired snapshots by solving these equations over short training horizons during which the populations remain far from equilibrium. For Toy I, training comprises 9 snapshots at t∈[0,1.2]t∈[0,1.2], and extrapolation is evaluated at 0.20.2-time-unit intervals from t=1.4t=1.4 to 4.04.0. For Toy I, training comprises 10 snapshots at t=0,0.5,…,4.5t=0,0.5,…,4.5, and extrapolation is evaluated at t∈5,6,8,10,12t∈\5,6,8,10,12\. We also synthesize a 64D version of each dataset using =softplus⁡(+)+ϵ∈ℝ64 y=softplus( W x+ c)+ ε ^64, with ∈ℝ2 x ^2 and Gaussian observation noise ϵ ε. Since the underlying systems are autonomous, we omit t from the encoder input. We compare KoopCell with scNODE (Zhang et al., 2024), which parameterizes the latent dynamics with a neural ODE. Both methods use the same latent dimension of D=3D=3 across all experiments. After training, we roll out each model from the observed initial distribution beyond the training horizon and report the mean sliced Wasserstein distance (SWD) to the ground-truth snapshots. Complete protocols and full results are provided in Appendix C. Spectral recovery. Table 1 reports the recovered continuous-time spectrum of the learned matrix A, while Figure 1 shows the corresponding eigenfunctions. KoopCell consistently recovers the Koopman spectrum with low MAE, and, for Toy I, the learned observables reproduce the level-set geometry of the ground-truth observables. For Toy I, we visualize the learned complex eigenfunction through its magnitude |Ψ|| | and phase ∠Ψ (Figure 1(b), left and middle). The magnitude decays toward zero at both attractors, q=±1q=± 1, while the phase reveals the oscillatory relaxation within each basin. Meanwhile, the eigenmode associated with the zero eigenvalue acts as an indicator that separates the two basins of attraction. Overall, these results demonstrate KoopCell’s ability to recover the underlying spectral structure from unpaired, temporally sparse distribution snapshots. Spectral structure supports stable extrapolation. Across all settings, KoopCell achieves lower extrapolation SWD than scNODE (Table 1). Together with the spectral recovery results, this suggests that learning the underlying Koopman spectral structure helps capture the long-term evolution of the system and supports more stable extrapolation beyond the training horizon. (a) Toy I: learned eigenfunctions aligned with x1,x2−53x12,x12\x_1,x_2- 53x_1^2,x_1^2\. (b) Toy I: learned oscillatory and basin eigenmodes. Table 2: Spectral recovery (MAE) and extrapolation comparison with scNODE (SWD ↓ ). Toy I Toy I Recovered Spectrum −0.195,−0.423,−0.977\-0.195,-0.423,-0.977\ 0.010,−0.298±1.481i\0.010,-0.298± 1.481i\ Spectrum MAE 0.01670.0167 0.07100.0710 Model Toy I Toy I 2D 64D 2D 64D SWD scNODE 0.19360.1936 0.20110.2011 0.10760.1076 0.02770.0277 KoopCell 0.06650.0665 0.06410.0641 0.05920.0592 0.01750.0175 Figure 1: Koopman eigenmodes learned from unpaired 2D snapshots. In (a), colors show the learned modes and black contours show the analytic eigenfunctions. In (b), the complex mode separates intrawell amplitude and phase, whereas the near-zero mode separates the two attraction basins; sink and saddle markers are included only for visualization. 4.2 Modeling Branching Dynamics with KoopCell-M We compare KoopCell and KoopCell-M on a stochastic pitchfork system in which an initially unimodal population separates into two branches. A fixed nonlinear map produces 64D noisy observations at nine training times. We initialize both models with the same pretrained autoencoder and jointly train their autoencoders and dynamics using the same population-matching objective. We impose the same Lipschitz bound on both decoders to limit the amplification of latent separation during decoding. Figure 2 overlays the population rollouts, while Table 2 summarizes SWD and branch-coordinate W1W_1 over all held-out times; definitions of both metrics and experimental details are provided in Appendix D. KoopCell-M more closely reproduces the separation of the two branches and achieves lower errors on both metrics, as expected. Table 3: Distributional prediction errors averaged over held-out times (mean ± SD over three optimization seeds). Model SWD ↓ Branch W1↓W_1 KoopCell 0.1557±0.01460.1557± 0.0146 0.2245±0.00690.2245± 0.0069 KoopCell-M 0.1124±0.00460.1124± 0.0046 0.0866±0.00260.0866± 0.0026 Figure 2: Population rollouts at all nine training times in a shared PCA plane (1,024 cells per time). Colors progress from purple to yellow with time; all panels use identical axes. 4.3 Modeling High-dimensional Single-Cell Population Dynamics To assess the scalability of KoopCell and KoopCell-M, we evaluate them on three high-dimensional single-cell RNA-sequencing datasets: zebrafish embryogenesis (ZB) (Farrell et al., 2018), Drosophila embryogenesis (DR) (Calderon et al., 2022), and mouse iPSC reprogramming (SC) (Schiebinger et al., 2019). Dataset descriptions, preprocessing, and visualizations are provided in Appendix E.1. Figure 3: Latent representations of SC cells colored by reference state (top; stromal: green, iPSC: red) and snapshot index (bottom). Tasks and evaluation protocol. Following Zhang et al. (2024), we evaluate our models on three tasks of increasing difficulty to assess their ability to predict gene expression distributions at held-out time points, both within and beyond the training horizon. For each task, we train on the populations at the retained time points. At evaluation, we sample cells from the initial snapshot, propagate their encoded states forward, and decode the predicted populations at the held-out times. Generated and observed distributions are compared in gene-expression space using the Wasserstein distance W2W_2. The tasks differ in which snapshots are withheld: (1) Easy (interpolation) removes selected intermediate snapshots to assess prediction within the training horizon; (2) Medium (extrapolation) removes several final snapshots to assess prediction beyond that horizon; and (3) Hard (interpolation and extrapolation) removes both intermediate and several final snapshots to assess the ability to recover population dynamics from highly sparse, unpaired distribution snapshots. Figure 12 shows the training and held-out time points. Baselines. We compare against several representative models for temporal dynamics of single-cell populations, including scNODE (Zhang et al., 2024), LGP-OT (Balik and Lähdesmäki, 2026), MIOFlow (Huguet et al., 2022), PRESCIENT (Yeo et al., 2021), PI-SDE (Jiang and Wan, 2024), and VGFM (Wang et al., 2025). KoopCell(-M), scNODE and LGP-OT jointly learn an informative latent representation and the latent dynamics, while other models learn dynamics in a fixed latent space obtained by dimensionality reduction. We tune all models to their optimal hyperparameters (Appendix E.2) and report the mean and std over three training seeds. Interpolation and extrapolation performance. KoopCell achieves strong predictive performance across all three datasets and tasks (Table 1; Appendix Tables 10 and 11). While some baselines attain comparable or slightly lower errors at individual interpolation times, KoopCell achieves lower average W2W_2 distances than all baselines on both the medium task and the extrapolation portion of the hard task. Incorporating memory further enhances predictive performance compared to the linear model (KoopCell). Relative to the strongest baseline for each dataset, KoopCell-M reduces the average error by 3.63.6–7.3%7.3\% on the medium task and 2.02.0–6.6%6.6\% on the extrapolation portion of the hard task. These results highlight the ability of our models to learn population dynamics from highly sparse, unpaired snapshots and extrapolate reliably beyond the training horizon. Appendix Figures 13–15 provide the corresponding visualization comparisons of predicted cell distributions. Dynamics Shapes Representation. Figure 3 compares latent representations learned by KoopCell and KoopCell-M on the SC dataset under the hard-task setting. The complete comparison of latent representations, including those from PCA, scNODE, and MIOFlow, appears in Appendix Figure 16. We match cells to published Waddington-OT reference states and label the initial population as mouse embryonic fibroblasts (MEFs), leaving unmatched or ambiguous cells gray. Reprogramming involves divergence toward stromal states or mesenchymal-to-epithelial transition (MET), from which iPSCs and alternative neural- and trophoblast-like states emerge (Schiebinger et al., 2019). KoopCell-M more clearly separates stromal and iPSC populations with distinct temporal progressions, whereas these late populations overlap substantially in KoopCell. These qualitative patterns are consistent with the branching structure of the memory-augmented model. Biological Downstream Analysis. Following Zhang et al. (2024), we examine lineage-associated genes in ZB using KoopCell-M fitted to all snapshots, focusing on Presomitic Mesoderm (PSM) and Hindbrain. Using the driver-gene statistic from CellRank (Weiler et al., 2024), we correlate preterminal gene expression with predicted fate probabilities and recover TBX16/MSGN1 for PSM and SOX3/SOX19A for Hindbrain (Appendix Figure 17). We further characterize dynamic gene-expression programs along model-predicted trajectories over 8–12 hours post fertilization (hpf). Gene modules with similar predicted expression-rate profiles are enriched for muscle and neural development in the PSM and Hindbrain lineages, respectively. Appendix E.5 provides methods and additional results. Together, these analyses support the biological interpretability of our model. 5 Related Work Learning scRNA-seq dynamics from distribution snapshots. Several approaches have been proposed to infer temporal dynamics from unpaired distributions. TrajectoryNet (Tong et al., 2020) fits continuous normalizing flows regularized by dynamic OT in fixed PCA coordinates, while PRESCIENT and PI-SDE (Yeo et al., 2021; Jiang and Wan, 2024) model potential-driven and Hamilton–Jacobi-regularized SDEs, respectively. MIOFlow (Huguet et al., 2022) fits neural ODEs in a pretrained geodesic embedding; scNODE (Zhang et al., 2024) jointly trains a VAE and latent neural ODE through dynamic regularization, and LGP-OT (Balik and Lähdesmäki, 2026) learns a latent Gaussian process and a decoder using an OT objective. Additionally, proliferation and death are accommodated through unbalanced OT (Zhang et al., 2025; Sha et al., 2024; Sun et al., 2025). VGFM (Wang et al., 2025) and WFR-FM (Peng et al., 2026) further combine unbalanced OT with flow matching (Lipman et al., 2023; Tong et al., 2024a; Tong et al., 2024b) for scalable joint velocity and growth estimation. Koopman learning and memory. Koopman theory (Koopman, 1931) describes nonlinear dynamics through the linear evolution of observables, motivating the joint learning of autoencoders and latent linear propagators for dynamical prediction and time-series generation (Otto and Rowley, 2019; Lusch et al., 2018; Liu et al., 2023; Naiman et al., 2024). While these methods learn from paired time-series data, we develop theory and algorithms for generative Koopman learning from unpaired distribution snapshots. Finding observables with exact linear closure remains challenging. Recent work uses the Mori–Zwanzig formalism (Mori, 1965; Zwanzig, 1961) to learn non-Markovian memory closures to account for unresolved variables and improve dynamical prediction (Lin et al., 2021; Menier et al., 2025; Gupta et al., 2025; Lu et al., 2026). Building on this principle, we learn a memory closure through an augmented Markovian system to model branching population dynamics from snapshots. 6 Conclusions We introduced KoopCell, a unified framework for learning Koopman generative dynamics from unpaired distribution snapshots. The weak population residual in Wasserstein space yields a closed-form least-squares estimator of the linear generator with convergence guarantees, while a Markovian memory embedding further enables KoopCell-M to capture branching dynamics. Our theoretical and empirical findings demonstrate the ability to recover Koopman spectra and improve extrapolation of gene-expression distributions from sparse, unpaired measurements, supporting quantitative analysis of developmental dynamics beyond the observed time window. Future directions include applications to multi-omics data and learning stochastic Koopman operators from distribution snapshots. AI use statement In this work, we used generative AI tools to assist with code implementation and debugging, including baseline implementations. We also used these tools to help verify the correctness and technical details of our mathematical arguments and algorithms, and to support synthetic data generation, experimental visualizations, and language polishing. We did not use generative AI tools to generate research ideas, develop the theoretical framework, formulate the main mathematical results, or design the core algorithms. We reviewed all AI-assisted work and take full responsibility for the final content, including the mathematical arguments, code, data, figures, and text. Ethics statement This work develops computational methods for modeling cellular population dynamics. Our experiments use synthetic data and publicly available single-cell datasets from zebrafish, Drosophila, and mouse, and involve no new biological sample collection or human participants. We anticipate no foreseeable adverse societal effects from this work. Reproducibility statement We facilitate reproducibility by documenting the assumptions and proofs of our theoretical results in Appendix A, and the estimation and training procedures for KoopCell and KoopCell-M, including pseudocode, in Appendix B. Synthetic data generation and experimental protocols are described in Appendices C and D. Single-cell dataset sources and preprocessing are provided in Appendix E.1, with training and held-out time points, evaluation metrics in Appendix E.1.1. Baseline configurations and model-selection procedures are documented in Appendix E.2, and the biological downstream analyses are detailed in Appendix E.5. We will release our code as open source upon acceptance of the paper. References Ambrosio et al. (2008) L. Ambrosio, N. Gigli, and G. Savaré Gradient flows in metric spaces and in the space of probability measures. 2. ed edition, Lectures in Mathematics ETH Zürich, Springer. Note: OCLC: 254181287 External Links: ISBN 978-3-7643-8722-8 978-3-7643-8721-1 Cited by: §A.4.1, §2. Aronszajn (1950) N. Aronszajn Theory of reproducing kernels. Transactions of the American Mathematical Society 68 (3), p. 337–404. Cited by: §A.3.1. Ayala et al. (1973) F. J. Ayala, M. E. Gilpin, and J. G. Ehrenfeld Competition between species: theoretical models and experimental tests. Theoretical Population Biology 4 (3), p. 331–356. Cited by: §1. Balik and Lähdesmäki (2026) M. Y. Balik and H. Lähdesmäki Modeling temporal scRNA-seq data with latent Gaussian process and optimal transport. In Forty-Third International Conference on Machine Learning, External Links: Link Cited by: §E.2, §4.3, §5. Benjamini and Hochberg (1995) Y. Benjamini and Y. Hochberg Controlling the false discovery rate: a practical and powerful approach to multiple testing. Journal of the Royal Statistical Society Series B: Statistical Methodology 57 (1), p. 289–300. External Links: ISSN 1467-9868, Link, Document Cited by: §E.5. Boyle et al. (2004) E. I. Boyle, S. Weng, J. Gollub, H. Jin, D. Botstein, J. M. Cherry, and G. Sherlock GO::TermFinder—open source software for accessing Gene Ontology information and finding significantly enriched Gene Ontology terms associated with a list of genes. Bioinformatics 20 (18), p. 3710–3715. External Links: ISSN 1367-4803, Link, Document Cited by: §E.5. Calderon et al. (2022) D. Calderon, R. Blecher-Gonen, X. Huang, S. Secchia, J. Kentro, R. M. Daza, B. Martin, A. Dulja, C. Schaub, C. Trapnell, E. Larschan, K. M. O’Connor-Giles, E. E. M. Furlong, and J. Shendure The continuum of Drosophila embryonic development at single-cell resolution. Science 377 (6606), p. eabn5800. External Links: Document, Link, https://w.science.org/doi/pdf/10.1126/science.abn5800 Cited by: §B.2.1, §B.2.1, §E.1, §4.3. Chen et al. (2019) G. Chen, B. Ning, and T. Shi Single-cell RNA-Seq technologies and related computational data analysis. Frontiers in Genetics 10, p. 317. Cited by: §1. Chen et al. (2018) R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud Neural ordinary differential equations. Advances in Neural Information Processing Systems 31. Cited by: §1. Cokus et al. (2019) S. J. Cokus, M. De La Torre, E. F. Medina, J. P. Rasmussen, J. Ramirez-Gutierrez, A. Sagasti, and F. Wang Tissue-specific transcriptomes reveal gene expression trajectories in two maturing skin epithelial layers in zebrafish embryos. G3 Genes|Genomes|Genetics 9 (10), p. 3439–3452. External Links: ISSN 2160-1836, Link, Document Cited by: §E.5. Csiszár (1975) I. Csiszár I-divergence geometry of probability distributions and minimization problems. The Annals of Probability, p. 146–158. Cited by: §3.1. Domínguez Conde et al. (2022) C. Domínguez Conde, C. Xu, L. B. Jarvis, D. B. Rainbow, S. B. Wells, T. Gomes, S. K. Howlett, O. Suchanek, K. Polanski, H. W. King, L. Mamanova, N. Huang, P. A. Szabo, L. Richardson, L. Bolt, E. S. Fasouli, K. T. Mahbubani, M. Prete, L. Tuck, N. Richoz, Z. K. Tuong, L. Campos, H. S. Mousa, E. J. Needham, S. Pritchard, T. Li, R. Elmentaite, J. Park, E. Rahmani, D. Chen, D. K. Menon, O. A. Bayraktar, L. K. James, K. B. Meyer, N. Yosef, M. R. Clatworthy, P. A. Sims, D. L. Farber, K. Saeb-Parsy, J. L. Jones, and S. A. Teichmann Cross-tissue immune cell analysis reveals tissue-specific features in humans. Science 376 (6594). External Links: ISSN 1095-9203, Link, Document Cited by: §B.2.1. Engel and Nagel (2000) K. Engel and R. Nagel One-parameter semigroups for linear evolution equations. Graduate Texts in Mathematics, Vol. 194, Springer. External Links: Document Cited by: §A.3.2. Farrell et al. (2018) J. A. Farrell, Y. Wang, S. J. Riesenfeld, K. Shekhar, A. Regev, and A. F. Schier Single-cell reconstruction of developmental trajectories during zebrafish embryogenesis. Science 360 (6392), p. eaar3131. External Links: Document, Link, https://w.science.org/doi/pdf/10.1126/science.aar3131 Cited by: §B.2.1, §B.2.1, §E.1, §4.3. Gene Ontology Consortium (2026) Gene Ontology Consortium The Gene Ontology knowledgebase in 2026. Nucleic Acids Research 54 (D1), p. D1779–D1792. External Links: ISSN 1362-4962, Link, Document Cited by: §E.5. Graves et al. (2006) A. Graves, S. Fernández, F. Gomez, and J. Schmidhuber Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks. In Proceedings of the 23rd International Conference on Machine Learning – ICML ’06, ICML ’06, p. 369–376. External Links: Link, Document Cited by: §B.2.2. Grønbech et al. (2020) C. H. Grønbech, M. F. Vording, P. N. Timshel, C. K. Sønderby, T. H. Pers, and O. Winther scVAE: variational auto-encoders for single-cell gene expression data. Bioinformatics 36 (16), p. 4415–4422. Cited by: §3.1. Guan et al. (2024) V. Guan, J. Janssen, H. Rahmani, A. Warren, S. Zhang, E. Robeva, and G. Schiebinger Identifying drift, diffusion, and causal structure from temporal snapshots. arXiv preprint arXiv:2410.22729. Cited by: §A.4.2. Gupta et al. (2025) P. Gupta, P. Schmid, D. Sipp, T. Sayadi, and G. Rigas Mori–Zwanzig latent space Koopman closure for nonlinear autoencoder. Proceedings of the Royal Society A 481 (2313), p. 20240259. Cited by: §5. Hijón et al. (2010) C. Hijón, P. Español, E. Vanden-Eijnden, and R. Delgado-Buscalioni Mori–Zwanzig formalism as a practical computational tool. Faraday Discussions 144, p. 301–322. External Links: Document Cited by: §1. Huguet et al. (2022) G. Huguet, D. S. Magruder, A. Tong, O. Fasina, M. Kuchroo, G. Wolf, and S. Krishnaswamy Manifold interpolating optimal-transport flows for trajectory inference. In Advances in Neural Information Processing Systems, A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (Eds.), External Links: Link Cited by: §E.2, §1, §4.3, §5. Hwang et al. (2018) B. Hwang, J. H. Lee, and D. Bang Single-cell RNA sequencing technologies and bioinformatics pipelines. Experimental & Molecular Medicine 50 (8), p. 1–14. Cited by: §1. Ikeda et al. (2022) M. Ikeda, I. Ishikawa, and C. Schlosser Koopman and Perron–Frobenius operators on reproducing kernel Banach spaces. Chaos: An Interdisciplinary Journal of Nonlinear Science 32 (12). Cited by: §A.3.1. Jiang and Wan (2024) Q. Jiang and L. Wan A physics-informed neural SDE network for learning cellular dynamics from time-series scRNA-seq data. Bioinformatics 40 (Supplement_2), p. i120–i127. External Links: ISSN 1367-4811, Document, Link, https://academic.oup.com/bioinformatics/article-pdf/40/Supplement_2/i120/59017026/btae400.pdf Cited by: §E.2, §1, §4.3, §5. Kingma and Welling (2014) D. P. Kingma and M. Welling Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, External Links: http://arxiv.org/abs/1312.6114v10 Cited by: §3.1. Köhne et al. (2025) F. Köhne, F. M. Philipp, M. Schaller, A. Schiela, and K. Worthmann ∞ L ∞-error bounds for approximations of the Koopman operator by kernel extended dynamic mode decomposition. SIAM Journal on Applied Dynamical Systems 24 (1), p. 501–529. External Links: Document, Link, https://doi.org/10.1137/24M1650120 Cited by: §A.3.1. Koopman and von Neumann (1932) B. O. Koopman and J. von Neumann Dynamical systems of continuous spectra. Proceedings of the National Academy of Sciences 18 (3), p. 255–263. Cited by: §1, §2. Koopman (1931) B. O. Koopman Hamiltonian systems and transformation in Hilbert space. Proceedings of the National Academy of Sciences 17 (5), p. 315–318. Cited by: §1, §2, §5. Korda and Mezić (2018) M. Korda and I. Mezić On convergence of extended dynamic mode decomposition to the Koopman operator. Journal of Nonlinear Science 28 (2), p. 687–710. External Links: ISSN 1432-1467, Link, Document Cited by: §A.4.4. Kosma et al. (2023) C. Kosma, G. Nikolentzos, G. Panagopoulos, J. Steyaert, and M. Vazirgiannis Neural ordinary differential equations for modeling epidemic spreading. Transactions on Machine Learning Research. Cited by: §1. Kostic et al. (2022) V. R. Kostic, P. Novelli, A. Maurer, C. Ciliberto, L. Rosasco, and M. Pontil Learning dynamical systems via Koopman operator regression in reproducing kernel Hilbert spaces. In Advances in Neural Information Processing Systems, A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (Eds.), External Links: Link Cited by: §A.3.1. Lei et al. (2016) H. Lei, N. A. Baker, and X. Li Data-driven parameterization of the generalized Langevin equation. Proceedings of the National Academy of Sciences 113 (50), p. 14183–14188. External Links: ISSN 1091-6490, Link, Document Cited by: §B.2.1, §2. Lin et al. (2021) Y. T. Lin, Y. Tian, D. Livescu, and M. Anghel Data-driven learning for the Mori–Zwanzig formalism: a generalization of the Koopman learning framework. SIAM Journal on Applied Dynamical Systems 20 (4), p. 2558–2601. Cited by: §B.2.1, §2, §5. Ling et al. (2026) Y. Ling, P. Zhang, Z. Zhang, and P. Zhou CellStream: dynamical optimal transport informed embeddings for reconstructing cellular trajectories from snapshots data. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 40, p. 746–754. Cited by: §1. Lipman et al. (2023) Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §5. Liu et al. (2023) Y. Liu, C. Li, J. Wang, and M. Long Koopa: learning non-stationary time series dynamics with Koopman predictors. Advances in Neural Information Processing Systems 36, p. 12271–12290. Cited by: §5. Lu et al. (2024) L. Lu, Z. Zeng, Y. Jiang, Y. Zhu, and P. Hu Weak collocation regression method: fast reveal hidden stochastic dynamics from high-dimensional aggregate data. Journal of Computational Physics 502, p. 112799. Cited by: §1. Lu et al. (2026) W. Lu, H. Ma, W. Lin, and Q. Zhu Interpretable functional Koopman learning with non-Markovian closure for spatiotemporal systems. In Forty-Third International Conference on Machine Learning, External Links: Link Cited by: §5. Lusch et al. (2018) B. Lusch, J. N. Kutz, and S. L. Brunton Deep learning for universal linear embeddings of nonlinear dynamics. Nature Communications 9 (1), p. 4950. Cited by: §1, §2, §5. Manning and Kimelman (2015) A. J. Manning and D. Kimelman Tbx16 and Msgn1 are required to establish directional cell migration of zebrafish mesodermal progenitors. Developmental Biology 406 (2), p. 172–185. External Links: ISSN 0012-1606, Link, Document Cited by: §E.5. Menier et al. (2025) E. Menier, S. Kaltenbach, M. Yagoubi, M. Schoenauer, and P. Koumoutsakos Interpretable learning of effective dynamics for multiscale systems. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences 481 (2306). Cited by: §B.2.1, §1, §2, §5. Moon et al. (2019) K. R. Moon, D. Van Dijk, Z. Wang, S. Gigante, D. B. Burkhardt, W. S. Chen, K. Yim, A. v. d. Elzen, M. J. Hirn, R. R. Coifman, et al. Visualizing structure and transitions in high-dimensional biological data. Nature Biotechnology 37 (12), p. 1482–1492. Cited by: §1. Mori (1965) H. Mori Transport, collective motion, and Brownian motion. Progress of Theoretical Physics 33 (3), p. 423–455. Cited by: §1, §2, §5. Naiman et al. (2024) I. Naiman, N. B. Erichson, P. Ren, M. W. Mahoney, and O. Azencot Generative modeling of regular and irregular time series data via Koopman VAEs. In International Conference on Learning Representations, Vol. 2024, p. 50547–50573. Cited by: §5. Osborn et al. (2020) D. P. S. Osborn, K. Li, S. J. Cutty, A. C. Nelson, F. C. Wardle, Y. Hinits, and S. M. Hughes Fgf-driven Tbx protein activities directly induce myf5 and myod to initiate zebrafish myogenesis. Development 147 (8), p. dev184689. External Links: ISSN 0950-1991, Link, Document Cited by: §E.5. Otto and Rowley (2019) S. E. Otto and C. W. Rowley Linearly recurrent autoencoder networks for learning dynamics. SIAM Journal on Applied Dynamical Systems 18 (1), p. 558–593. External Links: Document, Link, https://doi.org/10.1137/18M1177846 Cited by: §C.2, §C.2, §2, item (Toy I), §5. Peng et al. (2026) Q. Peng, Z. Wang, J. Ying, Y. Sun, Q. Nie, L. Zhang, T. Li, and P. Zhou WFR-FM: simulation-free dynamic unbalanced optimal transport. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §5. Rossi et al. (2009) C. C. Rossi, T. Kaji, and K. B. Artinger Transcriptional control of Rohon–Beard sensory neuron development at the neural plate border. Developmental Dynamics 238 (4), p. 931–943. External Links: ISSN 1097-0177, Link, Document Cited by: §E.5. Ruzicka et al. (2019) L. Ruzicka, D. G. Howe, S. Ramachandran, S. Toro, C. E. Van Slyke, Y. M. Bradford, A. Eagle, D. Fashena, K. Frazer, P. Kalita, P. Mani, R. Martin, S. T. Moxon, H. Paddock, C. Pich, K. Schaper, X. Shao, A. Singer, and M. Westerfield The Zebrafish Information Network: new support for non-coding genes, richer Gene Ontology annotations and the Alliance of Genome Resources. Nucleic Acids Research 47 (D1), p. D867–D873. External Links: ISSN 1362-4962, Link, Document Cited by: §E.5. Schiebinger et al. (2019) G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, P. Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell 176 (4), p. 928–943. Cited by: Figure 5, §B.2.1, §B.2.1, §E.1, §1, §3.2, §4.3, §4.3. Sha et al. (2024) Y. Sha, Y. Qiu, P. Zhou, and Q. Nie Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence 6 (1), p. 25–39. Cited by: §5. Sun et al. (2025) Y. Sun, Z. Zhang, Z. Wang, T. Li, and P. Zhou Variational regularized unbalanced optimal transport: single network, least action. In The Thirty-Ninth Annual Conference on Neural Information Processing Systems, External Links: Link Cited by: §5. Tambalo et al. (2020) M. Tambalo, R. Mitter, and D. G. Wilkinson A single cell transcriptome atlas of the developing zebrafish hindbrain. Development 147 (6), p. dev184143. External Links: ISSN 0950-1991, Link, Document Cited by: §E.5. Tong et al. (2024a) A. Tong, K. FATRAS, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, G. Wolf, and Y. Bengio Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research. Note: Expert Certification External Links: ISSN 2835-8856, Link Cited by: §5. Tong et al. (2020) A. Tong, J. Huang, G. Wolf, D. van Dijk, and S. Krishnaswamy TrajectoryNet: a dynamic optimal transport network for modeling cellular dynamics. In Proceedings of the 37th International Conference on Machine Learning, Cited by: §5. Tong et al. (2024b) A. Y. Tong, N. Malkin, K. Fatras, L. Atanackovic, Y. Zhang, G. Huguet, G. Wolf, and Y. Bengio Simulation-free Schrödinger bridges via score and flow matching. In Proceedings of the 27th International Conference on Artificial Intelligence and Statistics, S. Dasgupta, S. Mandt, and Y. Li (Eds.), Proceedings of Machine Learning Research, Vol. 238, p. 1279–1287. External Links: Link Cited by: §5. Tu et al. (2014) J. H. Tu, C. W. Rowley, D. M. Luchtenburg, S. L. Brunton, and J. N. Kutz On dynamic mode decomposition: theory and applications. Journal of Computational Dynamics 1 (2), p. 391–421. External Links: ISSN 2158-2491, Document, Link Cited by: item (Toy I). Wang et al. (2025) D. Wang, Y. Jiang, Z. Zhang, X. Gu, P. Zhou, and J. Sun Joint velocity-growth flow matching for single-cell dynamics modeling. In Advances in Neural Information Processing Systems, Vol. 38, Main Conference, p. 160552–160588. External Links: Document, Link Cited by: §E.2, §1, §4.3, §5. Wang et al. (2021) X. Wang, Q. An, Z. He, and W. Fang A literature review of social network analysis in epidemic prevention and control. Complexity 2021 (1), p. 3816221. Cited by: §1. Weiler et al. (2024) P. Weiler, M. Lange, M. Klein, D. Pe’er, and F. Theis CellRank 2: unified fate mapping in multiview single-cell data. Nature Methods 21 (7), p. 1196–1205. External Links: ISSN 1548-7105, Link, Document Cited by: §E.5, §4.3. Weinreb et al. (2020) C. Weinreb, A. Rodriguez-Fraticelli, F. D. Camargo, and A. M. Klein Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science 367 (6479). External Links: ISSN 1095-9203, Link, Document Cited by: §B.2.1, §3.2. Williams et al. (2015) M. O. Williams, I. G. Kevrekidis, and C. W. Rowley A data–driven approximation of the Koopman operator: extending dynamic mode decomposition. Journal of Nonlinear Science 25 (6), p. 1307–1346. Cited by: §2. Yeo et al. (2021) G. H. T. Yeo, S. D. Saksena, and D. K. Gifford Generative modeling of single-cell time series with PRESCIENT enables prediction of cell trajectories with interventions. Nature Communications 12 (1), p. 3222. Cited by: §E.2, §1, §4.3, §5. Zhang et al. (2024) J. Zhang, E. Larschan, J. Bigness, and R. Singh scNODE: generative model for temporal single cell transcriptomic data prediction. Bioinformatics 40 (Supplement_2), p. i146–i154. Cited by: §E.1, §E.2, §E.5, §1, §3.1, §4.1, §4.3, §4.3, §4.3, §5. Zhang et al. (2025) Z. Zhang, T. Li, and P. Zhou Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In International Conference on Learning Representations, Vol. 2025, p. 19888–19919. Cited by: §1, §5. Zhao et al. (2026) B. Zhao, X. Zhang, Y. Dong, X. Lu, W. Lin, and Q. Zhu Delay flow matching. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §B.2.1, §3.2. Zwanzig (1961) R. Zwanzig Memory effects in irreversible thermodynamics. Physical Review 124 (4), p. 983. Cited by: §5. Zwanzig (1973) R. Zwanzig Nonlinear generalized Langevin equations. Journal of Statistical Physics 9 (3), p. 215–220. Cited by: §1, §2. Appendix Contents Appendix A Theory This appendix develops the theoretical foundations of estimating the linear generator based on the weak population residual, introduced in Section 3.1. Sections A.1–A.2 prove Theorem 3.1: minimizing the weak population residual gives the best projected linear approximation to the canonical Wasserstein velocity, and finite tests yield a least-squares regression problem whose minimizers converge under dense test enrichment and uniqueness. We then consider increasing the observable dimension in a Koopman-invariant RKHS. Section A.3 shows that the minimum weak population residual vanishes as D→∞D→∞; Section A.4 relates this consistency to Wasserstein prediction bounds and locally uniform strong convergence of the learned propagators, under the respective approximation, stability, and identifiability assumptions. Algorithmic details and pseudocode are provided in Appendix B. A.1 Dual formulation and geometry of the weak population residual Throughout Sections A.1–A.2, the observable dimension D is fixed. Following Section 3.1, we write μt∈2(ℝD) _t _2(R^D) for the latent distribution induced by the chosen observables at time t. We work over real Hilbert spaces; complex-valued tests can be handled by separating real and imaginary parts. A.1.1 Canonical velocity and the space–time test space Assume μ⋅∈AC2([0,T],2(ℝD)). _·∈ AC^2([0,T];P_2(R^D)). (7) Then there is a unique minimum-norm velocity t∈Tμt2 v_t∈ T_ _tP_2 such that ∂tμt+∇⋅(tμt)=0,∫0T‖t‖L2(μt)2t<∞. _t _t+∇\!·( v_t _t)=0, _0^T\| v_t\|_L^2( _t)^2dt<∞. (8) Here Tμ2=∇ϕ:ϕ∈Cc∞(ℝD)¯L2(μ),T_μP_2= \∇φ:φ∈ C_c^∞(R^D)\^L^2(μ), (9) and PμP_μ denotes the L2(μ)L^2(μ)-orthogonal projection onto this tangent space. We collect these spaces along the curve in μ=L2(0,T,Tμt2),‖μ2=∫0T∫‖t()‖2dμt()t. T_μ=L^2\! (0,T;T_ _tP_2 ), \| q\|_ T_μ^2= _0^T \| q_t( z)\|^2d _t( z)dt. (10) Let D consist of C1C^1 space–time potentials compactly supported in space uniformly in t. We identify two potentials whose spatial gradients agree in L2(dtdμt)L^2(dt\,d _t) and define DV_D as the completion under ‖φ‖D2=∫0T∫‖∇φ(t,)‖2dμt()t.\| \|_V_D^2= _0^T \| _ z (t, z)\|^2d _t( z)dt. (11) Lemma A.1 (Density of smooth space–time gradients). The map φ↦∇φ _ z extends to an isometric isomorphism from DV_D onto μ T_μ. In particular, finite sums ∑ℓ=1Lηℓ(t)∇ψℓ(),ηℓ∈C1([0,T]),ψℓ∈Cc∞(ℝD), _ =1^L _ (t)∇ _ ( z), _ ∈ C^1([0,T]), _ ∈ C_c^∞(R^D), (12) are dense in μ T_μ. Proof. The gradient map is an isometry from the quotient of D into μ T_μ: spatial mollification approximates each Cc1C_c^1 potential by smooth compactly supported potentials with uniformly convergent gradients. To prove density, let ∈μ q∈ T_μ be orthogonal to every field in Equation (12). For fixed ψ∈Cc∞ψ∈ C_c^∞, set Fψ(t)=∫t()⊤∇ψ()dμt().F_ψ(t)= q_t( z) ∇ψ( z)d _t( z). Cauchy–Schwarz and boundedness of ∇ψ∇ψ imply Fψ∈L1(0,T)F_ψ∈ L^1(0,T). Orthogonality to η(t)∇ψ()η(t)∇ψ( z) for every smooth η therefore implies Fψ(t)=0F_ψ(t)=0 for almost every t. For each integer m≥1m≥ 1, choose a countable C1C^1-norm dense subset of Cc∞(Bm(0))C_c^∞(B_m(0)), viewing its members as functions on ℝDR^D, and enumerate their union as ψn\ _n\. Such subsets exist by separability of C1(Bm(0)¯)C^1( B_m(0)). Every ψ∈Cc∞(ℝD)ψ∈ C_c^∞(R^D) can then be approximated by members supported in one fixed ball, with ‖∇ψn−∇ψ‖∞→0\|∇ _n-∇ψ\|_∞→ 0 along a suitable sequence. Since μt _t is a probability measure, ‖∇ψn−∇ψ‖L2(μt)≤‖∇ψn−∇ψ‖∞\|∇ _n-∇ψ\|_L^2( _t)≤\|∇ _n-∇ψ\|_∞. Thus the same countable family generates a dense gradient class for every t. Removing the union of the null sets obtained above, we have, for almost every t, ∫t⊤∇ψndμt=0for every n. q_t ∇ _n\,d _t=0 every n. Since t∈Tμt2 q_t∈ T_ _tP_2, density forces t=0 q_t=0. Hence the orthogonal complement of the range is trivial, proving density and the claimed isomorphism after completion. ∎ This lemma allows optimization over completed test potentials to be identified with optimization over all square-integrable tangent velocity fields, as used in Theorem A.3 below. A.1.2 Residual functional and minimum-flux geometry For ∈ℝD×D A ^D× D define r(t)=∂tμt+∇⋅(μt).r_ A(t)= _t _t+∇\!·( A z\, _t). (13) For a smooth test φ∈ ∈ D, define ℛD(,φ):=⟨r,φ⟩R_D( A; ):= r_ A, . Integration by parts gives ℛD(,φ)=μT[φ⁡(T)]−μ0[φ⁡(0)]−∫0Tμt[∂tφ⁡(t)+∇φ(t)⊤]t.R_D( A; )= _T[ (T)]- _0[ (0)]- _0^T _t\! [ _t (t)+ _ z (t) A z ]dt. (14) Using Equation (8), the same quantity is ℛD(,φ)=∫0T∫∇φ(t,)⊤(t()−)dμt()t.R_D( A; )= _0^T _ z (t, z) ( v_t( z)- A z)d _t( z)dt. (15) Thus it extends continuously to DV_D. For a zero-mass distribution σ, define ‖σ‖−1,μ=supϕ∈Cc∞(ℝD)⟨σ,ϕ⟩:∫‖∇ϕ‖2μ≤1.\|σ\|_-1,μ= _φ∈ C_c^∞(R^D) \ σ,φ : \|∇φ\|^2dμ≤ 1 \. (16) Lemma A.2 (Minimum-flux representation). If ∈L2(μ,ℝD) w∈ L^2(μ;R^D) and σ=−∇⋅(μ)σ=-∇\!·( wμ), then ‖σ‖−1,μ=‖Pμ‖L2(μ).\|σ\|_-1,μ=\|P_μ w\|_L^2(μ). (17) Moreover, PμP_μ w is the unique min-norm vector field among all L2(μ)L^2(μ) fields satisfying −∇⋅(μ)=σ-∇·( qμ)=σ. Conversely, if ‖σ‖−1,μ<∞\|σ\|_-1,μ<∞, a tangent field satisfying this constraint exists. Proof. For every smooth ϕφ, ⟨σ,ϕ⟩ σ,φ =∫∇ϕ⊤dμ=∫∇ϕ⊤Pμdμ, = ∇φ w\,dμ= ∇φ P_μ w\,dμ, because −Pμ w-P_μ w is orthogonal to the tangent space. Cauchy–Schwarz proves the upper bound in Equation (17). Conversely, if Pμ≠0P_μ w≠ 0, choose gradients converging in L2(μ)L^2(μ) to PμP_μ w and normalize them. Their pairings converge to ‖Pμ‖\|P_μ w\|, proving equality. If Pμ=0P_μ w=0, all pairings vanish and equality is immediate. If q has the same weighted divergence, then − q- w is orthogonal to all smooth gradients, hence Pμ=PμP_μ q=P_μ w. Orthogonal decomposition gives ‖L2(μ)2=‖Pμ‖L2(μ)2+‖(I−Pμ)‖L2(μ)2.\| q\|_L^2(μ)^2=\|P_μ w\|_L^2(μ)^2+\|(I-P_μ) q\|_L^2(μ)^2. The minimum is therefore attained uniquely at =Pμ q=P_μ w. For the converse, finiteness of ‖σ‖−1,μ\|σ\|_-1,μ makes ∇ϕ↦⟨σ,ϕ⟩∇φ σ,φ a well-defined bounded linear functional on smooth gradients. It extends continuously to Tμ2T_μP_2, so the Riesz representation theorem gives a unique ∈Tμ2 q∈ T_μP_2 with ⟨σ,ϕ⟩=∫⊤∇ϕdμ σ,φ = q ∇φ\,dμ for every ϕ∈Cc∞(ℝD)φ∈ C_c^∞(R^D). Thus σ=−∇⋅(μ)σ=-∇·( qμ), and the preceding argument applies. ∎ Theorem A.3 (Dual representation, exact geometry, and population optimum). For every ∈ℝD×D A ^D× D, D() _D( A) :=12∫0T‖r(t)‖−1,μt2t := 12 _0^T\|r_ A(t)\|_-1, _t^2dt (18) =supφ∈DℛD(,φ)−12‖φ‖D2 = _ _D \R_D( A; )- 12\| \|_V_D^2 \ (19) =12∫0T‖t−Pμt()‖L2(μt)2t. = 12 _0^T\| v_t-P_ _t( A z)\|_L^2( _t)^2dt. (20) A minimizer D† A_D exists. It satisfies ∫0T∫(t()−Pμt(D†))⊤dμt()t=0. _0^T ( v_t( z)-P_ _t( A_D z) ) z d _t( z)dt=0. (21) If Pμt()=0 in L2(μt) for a.e. t⟹=0,P_ _t( B z)=0 in L^2( _t) for a.e. t B=0, (22) then the minimizing matrix is unique. Every minimizer satisfies D()=D(D†)+12∫0T‖Pμt((−D†))‖L2(μt)2t.J_D( A)=J_D( A_D )+ 12 _0^T\|P_ _t(( A- A_D ) z)\|_L^2( _t)^2dt. (23) Proof. The canonical equation gives r=−∇⋅((t−)μt).r_ A=-∇\!·(( v_t- A z) _t). Lemma A.2 and Pμtt=tP_ _t v_t= v_t yield, for almost every t, ‖r(t)‖−1,μt=‖t−Pμt()‖L2(μt).\|r_ A(t)\|_-1, _t=\| v_t-P_ _t( A z)\|_L^2( _t). This proves Equation (20). To obtain Equation (19), set t=t−Pμt() h_t= v_t-P_ _t( A z). For smooth φ , Equation (15) and orthogonal projection give ℛD(,φ)=⟨,∇φ⟩μR_D( A; )= h, _ z _ T_μ, where the inner product integrates over both space and time. This identity extends continuously to DV_D. Lemma A.1 makes the correspondence φ↦=∇φ q= _ z an isometric isomorphism onto μ T_μ. Thus ‖φ‖D=‖μ\| \|_V_D=\| q\|_ T_μ, and taking the supremum over φ∈D _D is equivalent to taking it over all ∈μ q∈ T_μ. Completing the square gives ⟨,⟩μ−12‖μ2=12‖μ2−12‖−‖μ2. h, q _ T_μ- 12\| q\|_ T_μ^2= 12\| h\|_ T_μ^2- 12\| q- h\|_ T_μ^2. The maximum is attained at = q= h and equals 12‖μ2=D() 12\| h\|_ T_μ^2=J_D( A) by Equation (20), proving Equation (19). Define D=[t↦Pμt()]∈μ. S_D B=[t P_ _t( B z)]∈ T_μ. This is a bounded linear map because ‖D‖μ2≤‖F2∫0T∫‖2dμtt.\| S_D B\|_ T_μ^2≤\| B\|_F^2 _0^T \| z\|^2d _tdt. Its range is finite dimensional and closed, and D()=12‖−D‖μ2J_D( A)= 12\| v- S_D A\|_ T_μ^2. Projecting v onto this range therefore gives a minimizing matrix D† A_D , whose projected velocity DD† S_D A_D is unique. Orthogonality to every D S_D B gives ∫0T∫(t−Pμt(D†))⊤Pμt()dμtt=0. _0^T ( v_t-P_ _t( A_D z)) P_ _t( B z)d _tdt=0. The first factor is tangent, so Pμt()P_ _t( B z) may be replaced by B z. Since the last identity holds for every B, it is equivalent to Equation (21). The Pythagorean theorem gives Equation (23); injectivity in Equation (22) makes the matrix representative unique. ∎ A.2 Finite tests, closed-form regression, and convergence under test enrichment Let D,1⊂D,2⊂⋯⊂D,⋃m≥1D,m¯=D,V_D,1 _D,2⊂·s _D, _m≥ 1V_D,m=V_D, (24) where D,m=span⁡φ1,…,φmV_D,m=span\ _1,…, _m\. We choose the test potentials φj _j from the class D defined above; such a dense family exists by Lemma A.1. Define yj= y_j= μT[φj(T)]−μ0[φj(0)]−∫0Tμt[∂tφj(t)]t, _T[ _j(T)]- _0[ _j(0)]- _0^T _t[ _t _j(t)]dt, (25) j= H_j= ∫0Tμt[∇φj(t,)⊤]t, _0^T _t[ _ z _j(t, z) z ]dt, (26) Mij= M_ij= ∫0Tμt[∇φi(t,)⊤∇φj(t,)]t. _0^T _t[ _ z _i(t, z) _ z _j(t, z)]dt. (27) Let =vec⁡() a=vec( A), let row j of G be vec⁡(j)⊤vec( H_j) , and set =(y1,…,ym)⊤ y=(y_1,…,y_m) . Lemma A.4. If ≻0 M 0, then supφ∈D,mℛD(,φ)−12‖φ‖D2=12(−)⊤−1(−). _ _D,m \R_D( A; )- 12\| \|_V_D^2 \= 12( y- G a) M^-1( y- G a). (28) The maximizing test is φ,m⋆=∑j=1m[−1(−)]jφj. _ A,m = _j=1^m[ M^-1( y- G a)]_j _j. (29) If G has full column rank, the unique matrix minimizer is vec⁡(D,m)=(⊤−1)−1⊤−1.vec( A_D,m)=( G M^-1 G)^-1 G M^-1 y. (30) Proof. Write φ=∑j=1mcjφj = _j=1^mc_j _j, with =(c1,…,cm)⊤ c=(c_1,…,c_m) . For each basis function, the definitions of yjy_j and j H_j give ℛD(,φj) _D( A; _j) =yj−∫0Tμt[∇φj(t,)⊤]t =y_j- _0^T _t\! [ _ z _j(t, z) A z ]dt =yj−⟨j,⟩F=yj−vec⁡(j)⊤=yj−()j. =y_j- H_j, A _F=y_j-vec( H_j) a=y_j-( G a)_j. Therefore, by linearity of ℛD(,⋅)R_D( A;·), ℛD(,φ)=∑j=1mcjℛD(,φj)=⊤(−).R_D( A; )= _j=1^mc_jR_D( A; _j)= c ( y- G a). Moreover, since M is the Gram matrix of φjj=1m\ _j\_j=1^m under the DV_D inner product, ‖φ‖D2=∑i,j=1mcicj∫0Tμt[∇φi(t,)⊤∇φj(t,)]t=⊤. \| \|_V_D^2= _i,j=1^mc_ic_j _0^T _t\! [ _ z _i(t, z) _ z _j(t, z) ]dt= c M c. Completing the square gives ⊤(−)−12⊤= c ( y- G a)- 12 c M c= −12(−1(−))⊤(−1(−)) - 12( c- M^-1( y- G a)) M( c- M^-1( y- G a)) +12(−)⊤−1(−). + 12( y- G a) M^-1( y- G a). The first term is nonpositive and vanishes exactly at the coefficients in Equation (29), proving Equation (28). Differentiating the remaining quadratic in a gives ⊤−1=⊤−1 G M^-1 G a= G M^-1 y. Full column rank makes the normal matrix positive definite and yields Equation (30). ∎ Remark (Weak population residuals as regression residuals). For fixed A, ℛD(,⋅)R_D( A;·) is the continuous linear functional obtained by testing the continuity-equation residual r_ A in Equation (13). Its value ℛD(,φj)R_D( A; _j) measures the signed discrepancy between the moment evolution of μt _t and the contribution of the candidate drift A z, as expressed in Equation (14). With the moments defined above, ℛD(;φj)=yj−⟨j,⟩F=yj−()j,j=1,…,m.R_D( A; _j)=y_j- H_j, A _F=y_j-( G a)_j, j=1,…,m. Thus − y- G a collects these scalar weak residuals, and = y= G a means that the continuity equation is satisfied against every test in D,mV_D,m. When ≻0 M 0, Lemma A.4 gives the weighted least-squares objective 12(−)⊤−1(−) 12( y- G a) M^-1( y- G a) for estimating A. In Appendix B.1.2, we implement this regression by approximating y, G, and M with empirical sample averages and temporal quadrature. Lemma A.4 gives the estimator for a fixed finite test space. We next show that, as the nested test spaces become dense in DV_D, these estimators converge to the unique minimizer of the full weak population residual objective. To compare the two objectives, we use the Riesz representation of the continuous functional ℛD(,⋅)R_D( A;·): for each A, there is a unique h∈Dh_ A _D satisfying ℛD(,φ)=⟨h,φ⟩DR_D( A; )= h_ A, _V_D for every φ∈D _D. Writing =vec⁡() a=vec( A) and letting h0h_0 correspond to =0 A=0, define the linear operator D:ℝD2→D T_D:R^D^2 _D by ⟨D,φ⟩D=∫0Tμt[∇φ(t,)⊤]t,φ∈D. T_D a, _V_D= _0^T _t[ _ z (t, z) A z]dt, _D. (31) Cauchy–Schwarz and the finite integrated second moment make this definition well posed by the Riesz theorem, with ‖D‖D≤(∫0Tμt[‖2]t)1/2‖2\| T_D a\|_V_D≤( _0^T _t[\| z\|^2]dt)^1/2\| a\|_2. The residual is affine in A, so h=h0−Dh_ A=h_0- T_D a. Under the gradient identification of Lemma A.1, ∇h0(t,⋅)=t _ zh_0(t,·)= v_t, ∇(D)(t,⋅)=Pμt() _ z( T_D a)(t,·)=P_ _t( A z), and thus ∇h(t,⋅)=t−Pμt() _ zh_ A(t,·)= v_t-P_ _t( A z), with these identities understood in μ T_μ. The Riesz representation allows us to complete the square in the dual objective: ℛD(,φ)−12‖φ‖D2=12‖h‖D2−12‖φ−h‖D2.R_D( A; )- 12\| \|_V_D^2= 12\|h_ A\|_V_D^2- 12\| -h_ A\|_V_D^2. The supremum over DV_D is therefore attained at φ=h =h_ A. For the finite-test problem, let Qm:D→D,mQ_m:V_D _D,m be the orthogonal projection. Since h−Qmh_ A-Q_mh_ A is orthogonal to D,mV_D,m, every φ∈D,m _D,m satisfies ⟨h,φ⟩D=⟨Qmh,φ⟩D h_ A, _V_D= Q_mh_ A, _V_D. The same calculation with QmhQ_mh_ A in place of h_ A shows that the restricted supremum is attained at φ=Qmh =Q_mh_ A, which belongs to D,mV_D,m. Consequently, D()=12‖h‖D2,D,m()=12‖Qmh‖D2.J_D( a)= 12\|h_ A\|_V_D^2, _D,m( a)= 12\|Q_mh_ A\|_V_D^2. (32) Thus increasing the test space approximates the full residual representative by its projections QmhQ_mh_ A. To compare the minimizers, write =D∗D C= T_D^* T_D and =D∗h0 d= T_D^*h_0 for the normal-equation matrix and vector of the full problem, and m=D∗QmD C_m= T_D^*Q_m T_D and m=D∗Qmh0 d_m= T_D^*Q_mh_0 for their finite-test counterparts. Theorem A.5 (Convergence under test-space enrichment). Suppose that the nested test spaces satisfy Equation (24) and that the population minimizer D† A_D is unique. Let D,m A_D,m minimize, over ∈ℝD×D A ^D× D, the objective obtained by restricting the dual supremum to D,mV_D,m, and set †=vec⁡(D†) a =vec( A_D ) and m=vec⁡(D,m) a_m=vec( A_D,m). Then ≻0 C 0, and, for all sufficiently large m, m≻0 C_m 0 and the finite-test minimizer is unique. For these m, ‖m−†‖2≤2λmin()(‖m−‖2+‖m−‖op‖†‖2).\| a_m- a \|_2≤ 2 _ ( C) (\| d_m- d\|_2+\| C_m- C\|_op\| a \|_2 ). (33) Consequently, D,m⟶D†as m→∞. A_D,m A_D m→∞. (34) Proof. For each m, the range of QmDQ_m T_D is finite dimensional and hence closed. Projecting Qmh0Q_mh_0 onto this range therefore yields a finite-test minimizer. The quadratic representations above yield the normal equations †=,mm=m. C a = d, C_m a_m= d_m. (35) Density in Equation (24) gives Qmq→qQ_mq→ q for every q∈Dq _D. Since the domain of D T_D is finite dimensional, this convergence is uniform after composition with D T_D. To see it directly, let 1,…,D2 e_1,…, e_D^2 be an orthonormal basis of ℝD2R^D^2. For =∑qaqq a= _qa_q e_q with ‖2=1\| a\|_2=1, ‖(I−Qm)D‖ \|(I-Q_m) T_D a\| ≤∑q|aq|‖(I−Qm)Dq‖ ≤ _q|a_q|\,\|(I-Q_m) T_D e_q\| ≤Dmaxq‖(I−Qm)Dq‖⟶0. ≤ D _q\|(I-Q_m) T_D e_q\| 0. Consequently, ‖m−‖op \| C_m- C\|_op =‖D∗(Qm−I)D‖op→0, =\| T_D^*(Q_m-I) T_D\|_op→ 0, (36) ‖m−‖2 \| d_m- d\|_2 =‖D∗(Qm−I)h0‖2→0. =\| T_D^*(Q_m-I)h_0\|_2→ 0. (37) Uniqueness of D† A_D is equivalent to injectivity of D T_D. Because its domain is finite dimensional, =D∗D≻0 C= T_D^* T_D 0. Put κ=λmin()κ= _ ( C). For all sufficiently large m, ‖m−‖≤κ/2\| C_m- C\|≤κ/2, and Weyl’s inequality gives m⪰(κ/2) C_m (κ/2) I. Thus m C_m is invertible and ‖m−1‖≤2/κ\| C_m^-1\|≤ 2/κ. Subtracting the equations in Equation (35) yields m(m−†)=(m−)−(m−)†. C_m( a_m- a )=( d_m- d)-( C_m- C) a . Multiplying by m−1 C_m^-1 proves Equation (33). Its right-hand side tends to zero by Equation (36)–Equation (37), which proves Equation (34). ∎ A.3 Convergence of the weak population residual Our goal in this section is to prove that as the number of compatible observables increases, the minimum weak population residual converges to zero. This is different from Section A.2, where the observable dimension D is fixed and only the test space is enlarged. As a beginning, we first define the observable space and the infinite-dimensional feature coordinates in which its dynamics become linear. A.3.1 Preliminaries: Koopman-invariant RKHS and linear feature dynamics Koopman operators on reproducing-kernel spaces have been studied from both operator-theoretic and statistical perspectives, including boundedness, spectral approximation, and learning from data (Ikeda et al., 2022; Kostic et al., 2022). From now on, let X be a separable metric space with a flow FtF_t, and let ℋH be a separable real RKHS with kernel k (Aronszajn, 1950). The feature map Φ⁡():=k⁡(,⋅) ( x):=k( x,·) satisfies f⁡()=⟨f,Φ⁡()⟩ℋf( x)= f, ( x) _H for every f∈ℋf ; we assume that Φ is strongly measurable. Suppose the Koopman operators U⁡(t)f=f∘FtU(t)f=f F_t preserve ℋH and form a C0C_0-semigroup with generator L. Wendland native spaces provide a concrete example of Koopman invariance: these RKHSs are norm-equivalent to Sobolev spaces on bounded Lipschitz domains, and composition preserves them when the state map is a sufficiently smooth diffeomorphism with bounded derivatives and Jacobian determinant bounded away from zero (Köhne et al., 2025). Furthermore, write =ℋZ=H for the feature space. Since ℋH is reflexive, T⁡(t):=U(t)∗T(t):=U(t)^* is also a C0C_0-semigroup, with generator =L∗ A=L^* on its adjoint domain, possibly unbounded. The reproducing identity gives ⟨f,T⁡(t)Φ()⟩ℋ=⟨U⁡(t)f,Φ⁡()⟩ℋ=f⁡(Ft) f,T(t) ( x) _H= U(t)f, ( x) _H=f(F_t x). As this holds for every f∈ℋf , the features evolve linearly, following: Φ⁡(Ft)=T⁡(t)Φ(). (F_t x)=T(t) ( x). For an orthonormal basis ejj≥1\e_j\_j≥ 1 of ℋH, let PDP_D be the projector onto D=span⁡e1,…,eDZ_D=span\e_1,…,e_D\. Since ⟨ej,Φ⁡()⟩ℋ=ej() e_j, ( x) _H=e_j( x), PDΦ()=∑j=1Dej()ej⟷D()=(e1(),…,eD())⊤.P_D ( x)= _j=1^De_j( x)e_j g_D( x)= (e_1( x),…,e_D( x) ) . Thus we may view the D-dimensional encoder as a projection of the same feature map Φ⁡() ( x). For a random initial state 0 X_0, define μt:=Law⁡(Φ⁡(Ft0)) _t:=Law( (F_t X_0)) and μtD:=(PD)#μt _t^D:=(P_D)_\# _t. Here μt _t denotes the full feature distribution, while μtD _t^D, identified with a measure on ℝDR^D, corresponds to the latent distribution denoted by μt _t in Section 3.1 and Sections A.1–A.2. The superscript makes its dependence on the observable dimension explicit. These projected laws can be compared in the common space Z as D increases, with T⁡(t)T(t) providing the exact linear evolution for the subsequent convergence analysis. A.3.2 Convergence of the weak population residual Fix τ>0τ>0 and 0∼ρ0 X_0 _0 on (Ω,ℙ)( ,P). For t∈[0,2τ]t∈[0,2τ], set 0=Φ⁡(0) Z_0= ( X_0), t=T⁡(t)0=Φ⁡(Ft0) Z_t=T(t) Z_0= (F_t X_0), μt=Law⁡(t) _t=Law( Z_t), and μtD=(PD)#μt _t^D=(P_D)_\# _t. Assume that 0∈D⁡() Z_0∈ D( A) almost surely and ⁡(‖0‖2+‖0‖2)<∞E(\| Z_0\|^2+\| A Z_0\|^2)<∞. We write ∥⋅∥L2\|·\|_L^2 for the norm on L2(Ω,)L^2( ;Z) and ∥⋅∥\|·\|_L for the norm on =L2([0,2τ]×Ω,)L=L^2([0,2τ]× ;Z). The semigroup preserves D⁡()D( A) and satisfies T(t)=T⁡(t) AT(t) z=T(t) A z and T⁡(t)−=∫0tT⁡(r)rT(t) z- z= _0^tT(r) A z\,dr for ∈D⁡() z∈ D( A) (Engel and Nagel, 2000, Chapter I). Strong continuity and uniform boundedness give Mτ:=sup0≤t≤2τ‖T⁡(t)‖<∞M_τ:= _0≤ t≤ 2τ\|T(t)\|<∞, so ∫02τ⁡(‖t‖2+‖t‖2)t≤2τMτ2(‖0‖2+‖0‖2)<∞. _0^2τE(\| Z_t\|^2+\| A Z_t\|^2)\,dt≤ 2τ M_τ^2E(\| Z_0\|^2+\| A Z_0\|^2)<∞. The same bound and dominated convergence show that t↦t Z_t is continuous in L2(Ω,)L^2( ;Z). For s≤ts≤ t, coupling PDsP_D Z_s and PDtP_D Z_t and using t−s=∫strr Z_t- Z_s= _s^t A Z_r\,dr yields W2(μsD,μtD)≤‖PD(t−s)‖L2≤∫st‖r‖L2r.W_2( _s^D, _t^D)≤\|P_D( Z_t- Z_s)\|_L^2≤ _s^t\| A Z_r\|_L^2\,dr. Since the last integrand belongs to L2(0,2τ)L^2(0,2τ), μ⋅D∈AC2([0,2τ],2(D)) _·^D∈ AC^2([0,2τ];P_2(Z_D)). The weak-continuity framework of the preceding subsections therefore applies, with weak population objective D()=12∫02τ‖∂tμtD+∇⋅(μtD)‖−1,μtD2t.J_D( B)= 12 _0^2τ\| _t _t^D+∇·( B y\, _t^D)\|_-1, _t^D^2\,dt. (38) Theorem A.3 ensures that a minimizer D∈argmin∈ℒ⁡(D)D() A_D∈ *argmin_ B (Z_D)J_D( B) exists; uniqueness is unnecessary for the statements below. Lemma A.6 (Comparison with the projected path velocity). Under the assumptions above, every ∈ℒ⁡(D) B (Z_D) satisfies 2D()≤‖PD−PD‖. 2J_D( B)≤\|P_D A Z- BP_D Z\|_L. Proof. Set t=PDt Y_t=P_D Z_t and t=PDt V_t=P_D A Z_t. Conditioning V on (t,t)(t, Y_t) under the probability measure (2τ)−1dtdℙ(2τ)^-1dt\,dP gives a jointly measurable field tD()=⁡[t∣t=] w_t^D( y)=E[ V_t Y_t= y] for almost every t. Conditional Jensen’s inequality gives ‖D()‖≤‖<∞\| w^D( Y)\|_L≤\| V\|_L<∞. For ψ∈Cc∞(D)ψ∈ C_c^∞(Z_D), the chain rule along t Y_t gives ψ(t)−ψ(s)=∫st∫⟨∇ψ(),rD()⟩dμrD()r.Eψ( Y_t)-Eψ( Y_s)= _s^t ∇ψ( y), w_r^D( y) \,d _r^D( y)\,dr. Here Fubini’s theorem applies because ∇ψ∇ψ is bounded and ∈ V . Thus D w^D satisfies the projected continuity equation ∂tμtD+∇⋅(tDμtD)=0 _t _t^D+ _ y·( w_t^D _t^D)=0 in the sense of distributions on (0,2τ)×D(0,2τ)×Z_D. Writing PμtDP_ _t^D for the orthogonal projection onto its Wasserstein tangent space, the canonical velocity is PμtDtDP_ _t^D w_t^D. Theorem A.3 and contraction of this projection imply 2D()=∫02τ‖PμtD(tD−)‖L2(μtD)2t≤‖D()−‖2.2J_D( B)= _0^2τ\|P_ _t^D( w_t^D- B y)\|_L^2( _t^D)^2\,dt≤\| w^D( Y)- B Y\|_L^2. Since tD(t)−t=⁡[t−t∣t] w_t^D( Y_t)- B Y_t=E[ V_t- B Y_t Y_t], another application of conditional Jensen gives ‖D()−‖≤‖−‖\| w^D( Y)- B Y\|_L≤\| V- B Y\|_L, proving the claim. ∎ To approximate the projected velocity, we use the bounded difference quotient Ch=(T⁡(h)−I)/hC_h=(T(h)-I)/h before projecting onto DZ_D. This avoids applying the unbounded generator to PDtP_D Z_t. Theorem A.7 (Convergence of the weak population residual). Under the assumptions above, every choice of weak population minimizers D A_D satisfies D(D)→0J_D( A_D)→ 0 as D→∞D→∞. More precisely, define a⁡(h)=‖(−Ch)‖a(h)=\|( A-C_h) Z\|_L and bD=‖(I−PD)‖b_D=\|(I-P_D) Z\|_L. Then a⁡(h)→0a(h)→ 0 as h↓0h 0, bD→0b_D→ 0 as D→∞D→∞, and 2D(D)≤a⁡(h)+‖Ch‖bD(h>0). 2J_D( A_D)≤ a(h)+\|C_h\|b_D (h>0). Proof. For fixed h>0h>0, ChC_h is bounded on Z, so D,h=PDCh|D B_D,h=P_DC_h|_Z_D is an admissible comparison matrix. The velocity error decomposes as PD−PDChPD=PD(−Ch)+PDCh(I−PD).P_D A Z-P_DC_hP_D Z=P_D( A-C_h) Z+P_DC_h(I-P_D) Z. The preceding lemma, the triangle inequality in L, and ‖PD‖≤1\|P_D\|≤ 1 therefore give 2D(D)≤2D(D,h)≤a⁡(h)+‖Ch‖bD. 2J_D( A_D)≤ 2J_D( B_D,h)≤ a(h)+\|C_h\|b_D. Strong convergence of PDP_D gives (I−PD)t→0(I-P_D) Z_t→ 0 pointwise, while ‖(I−PD)t‖2≤‖t‖2\|(I-P_D) Z_t\|^2≤\| Z_t\|^2 is integrable over time and probability. Dominated convergence yields bD→0b_D→ 0. For ∈D⁡() z∈ D( A), Ch=h−1∫0hT⁡(r)r→C_h z=h^-1 _0^hT(r) A z\,dr→ A z by strong continuity. Fix h0>0h_0>0 and let M0=sup0≤r≤h0‖T⁡(r)‖<∞M_0= _0≤ r≤ h_0\|T(r)\|<∞. For 0<h≤h00<h≤ h_0, ‖Ch‖≤M0‖‖\|C_h z\|≤ M_0\| A z\|, hence ‖(−Ch)t‖2≤(1+M0)2‖t‖2\|( A-C_h) Z_t\|^2≤(1+M_0)^2\| A Z_t\|^2. The finite orbit energy established above makes this bound integrable, so dominated convergence gives a⁡(h)→0a(h)→ 0. The bound on the minimum now proves the result by first fixing h and then increasing D. Specifically, for any ε>0 >0, choose h with a⁡(h)<ε/2a(h)< /2, and then choose D sufficiently large that bD<ε/[2(1+‖Ch‖)]b_D< /[2(1+\|C_h\|)]. It follows that 2D(D)<ε 2J_D( A_D)< . ∎ A.4 From weak residual consistency to strong semigroup convergence Section A.3.2 establishes D(D)→0J_D( A_D)→ 0 for the weak population minimizers D A_D. In this section, under a uniform stability bound on etDe^t A_D, we first show that this vanishing weak residual implies uniform Wasserstein convergence of the predicted distributions on the training interval. The subsequent subsections establish the stronger conclusion of semigroup convergence at every initial feature vector: compactness provides strong operator subsequences, and finite-time linear determinacy (FT) identifies their limits with T⁡(t)T(t). Assumption A.8 (Finite-time stability of weak population minimizers). For the selected minimizers D∈argmin∈ℒ⁡(D)D() A_D∈ *argmin_ B (Z_D)J_D( B), there exists M≥1M≥ 1, independent of D, such that supDsup0≤t≤τ‖etD‖ℒ⁡(D)≤M. _D _0≤ t≤τ\|e^t A_D\|_L(Z_D)≤ M. This assumption bounds the finite-time evolution generated by the selected minimizers. It is an additional hypothesis beyond convergence of the weak residual and does not require a uniform bound on ‖D‖\| A_D\|. A.4.1 The weak residual controls probability prediction The weak residual measures the minimum correction to a candidate linear velocity needed to transport a prescribed distribution curve. Propagating this correction gives a prediction bound on any subinterval. Lemma A.9 (The weak residual controls finite-time prediction). Let α⋅∈AC2([s,t],2(ℝd)) _·∈ AC^2([s,t];P_2(R^d)) and ∈ℝd×d B ^d× d. For r∈[s,t]r∈[s,t], write r(r)=∂rαr+∇⋅(αr)r_ B(r)= _r _r+ _ y·( B y\, _r). Then W2((e(t−s))#αs,αt)≤∫st‖e(t−r)‖‖r(r)‖−1,αrr.W_2((e^(t-s) B)_\# _s, _t)≤ _s^t\|e^(t-r) B\|\,\|r_ B(r)\|_-1, _r\,dr. Proof. Let r v_r be the canonical velocity of αr _r and set r=r−Pαr() u_r= v_r-P_ _r( B y), where PαrP_ _r is the orthogonal projection onto the Wasserstein tangent space. Lemma A.2 and Theorem A.3 give −∇⋅(rαr)=r(r)- _ y·( u_r _r)=r_ B(r) and ‖r‖L2(αr)=‖r(r)‖−1,αr\| u_r\|_L^2( _r)=\|r_ B(r)\|_-1, _r for almost every r. Hence ∂rαr+∇⋅((+r)αr)=0 _r _r+ _ y·(( B y+ u_r) _r)=0 in the sense of distributions. Choose a jointly Borel representative of u. The corrected velocity +r B y+ u_r has finite squared energy because ∈L2(drdαr) v∈ L^2(dr\,d _r), tangent projection is contractive, and α⋅ _· has bounded second moments on [s,t][s,t]. The superposition principle therefore gives a probability measure on absolutely continuous paths γ with marginals αr _r and ˙r=r+r(r) γ_r= B γ_r+ u_r( γ_r) almost everywhere (Ambrosio et al., 2008, Theorem 8.2.1). Variation of constants yields t−e(t−s)s=∫ste(t−r)r(r)r. γ_t-e^(t-s) B γ_s= _s^te^(t-r) B u_r( γ_r)\,dr. The pair (e(t−s)s,t)(e^(t-s) B γ_s, γ_t) couples the two distributions in the claim. Minkowski’s integral inequality and the marginal identity Law⁡(r)=αrLaw( γ_r)= _r give (‖t−e(t−s)s‖2)1/2≤∫st‖e(t−r)‖‖r‖L2(αr)r. (E\| γ_t-e^(t-s) B γ_s\|^2 )^1/2≤ _s^t\|e^(t-r) B\|\,\| u_r\|_L^2( _r)\,dr. Substituting the weak residual norm proves the bound. ∎ To compare the learned linear propagator in the common space Z, define KD(t)=ιDetDPD,SD(t)=KD(t)+(I−PD),K_D(t)= _De^t A_DP_D, S_D(t)=K_D(t)+(I-P_D), where ιD:D↪ _D:Z_D is inclusion. KD(t)K_D(t) evolves the projected features, while SD(t)S_D(t) also retains the omitted unresolved feature space. On D⊕D⟂Z_D _D , SD(t)S_D(t) acts as etD⊕Ie^t A_D I, so it is a C0C_0-semigroup with SD(0)=IS_D(0)=I and ‖SD(t)‖≤max⁡‖etD‖,1≤M\|S_D(t)\|≤ \\|e^t A_D\|,1\≤ M for 0≤t≤τ0≤ t≤τ. We use SDS_D only for the subsequent operator-convergence argument; KDK_D remains the propagator over the finite feature subspace, with KD(0)=PDK_D(0)=P_D. The following result combines the weak residual bound with the error from projecting t Z_t. It is uniform in both the starting time and the prediction lag, as needed later to compare operators on all source distribution laws. Theorem A.10. Under Theorem A.7 and Assumption A.8, let qD=sup0≤r≤2τ‖(I−PD)r‖L2q_D= _0≤ r≤ 2τ\|(I-P_D) Z_r\|_L^2, with the L2L^2 norm defined in Section A.3.2. Then qD→0q_D→ 0 and sups,t∈[0,τ]W2(KD(t)#μs,μs+t) _s,t∈[0,τ]W_2^Z(K_D(t)_\# _s, _s+t) ≤qD+M2τD(D)⟶0, ≤ q_D+M 2 _D( A_D) 0, sups,t∈[0,τ]W2(SD(t)#μs,μs+t) _s,t∈[0,τ]W_2^Z(S_D(t)_\# _s, _s+t) ≤2qD+M2τD(D)⟶0. ≤ 2q_D+M 2 _D( A_D) 0. Prediction from the initial law over the full training interval satisfies sup0≤t≤2τW2(KD(t)#μ0,μt)≤qD+M24τD(D)⟶0. _0≤ t≤ 2τW_2^Z(K_D(t)_\# _0, _t)≤ q_D+M^2 4 _D( A_D) 0. Proof. Strong convergence of PDP_D and dominated convergence imply ‖(I−PD)‖L2→0\|(I-P_D) V\|_L^2→ 0 for every ∈L2(Ω,) V∈ L^2( ;Z). Since r↦r Z_r is continuous in this space, its image on [0,2τ][0,2τ] is compact. For any ε>0 >0, choose a finite ε -net rjj=1m\ Z_r_j\_j=1^m of this image. Contraction of I−PDI-P_D gives qD≤ε+maxj⁡‖(I−PD)rj‖L2q_D≤ + _j\|(I-P_D) Z_r_j\|_L^2; letting D→∞D→∞ and then ε↓0 0 proves qD→0q_D→ 0. For s,t∈[0,τ]s,t∈[0,τ], apply Lemma A.9 to μ⋅D _·^D on [s,s+t]⊂[0,2τ][s,s+t]⊂[0,2τ]. With rD(r)=∂rμrD+∇⋅(DμrD)r_ A_D(r)= _r _r^D+ _ y·( A_D y\, _r^D), stability and Cauchy–Schwarz yield W2((etD)#μsD,μs+tD)≤M∫s+t‖rD(r)‖−1,μrDr≤M2tD(D).W_2((e^t A_D)_\# _s^D, _s+t^D)≤ M _s^s+t\|r_ A_D(r)\|_-1, _r^D\,dr≤ M 2tJ_D( A_D). The last inequality uses ∫02τ‖rD(r)‖−1,μrD2r=2D(D) _0^2τ\|r_ A_D(r)\|_-1, _r^D^2\,dr=2J_D( A_D). Isometric inclusion into Z preserves this Wasserstein distance, and coupling PDs+tP_D Z_s+t with s+t Z_s+t adds at most qDq_D. This proves the bound for KDK_D since t≤τt≤τ. Likewise, SD(t)s−KD(t)s=(I−PD)sS_D(t) Z_s-K_D(t) Z_s=(I-P_D) Z_s, so the coupling induced by s Z_s adds at most another qDq_D, proving the bound for SDS_D. For 0≤r≤2τ0≤ r≤ 2τ, the semigroup law gives ‖erD‖=‖(e(r/2)D)2‖≤M2\|e^r A_D\|=\|(e^(r/2) A_D)^2\|≤ M^2. Applying the same argument on [0,t][0,t] therefore bounds the full-interval prediction error by qD+M22tD(D)q_D+M^2 2tJ_D( A_D), which is at most the stated bound for t≤2τt≤ 2τ. All three bounds vanish because qD→0q_D→ 0 and D(D)→0J_D( A_D)→ 0 by Theorem A.7. ∎ A.4.2 Distributional identifiability of the linear propagator Theorem A.10 shows that a vanishing weak residual yields convergent distribution predictions. To establish convergence of the linear propagators themselves, we also need the distribution dynamics to identify the underlying map. This is a fundamental issue when inferring dynamics from unpaired snapshots: distinct trajectories can produce the same distribution at every time. For linear SDEs with additive noise in finite dimensions, Guan et al. (2024) characterize the obstruction to identifiability through generalized rotational symmetries of the initial distribution. For example, a rotational flow and a stationary flow preserve the same isotropic Gaussian distribution, despite moving individual states differently. In our setting, the subsequent operator-convergence argument therefore requires the family μs:0≤s≤τ\ _s:0≤ s≤τ\ to determine T⁡(t)T(t) uniquely for each lag t∈[0,τ]t∈[0,τ]. We impose this requirement directly on bounded linear maps on Z through the following assumption. Assumption A.11 (Finite-time linear determinacy (FT)). For every lag t∈[0,τ]t∈[0,τ] and every bounded linear map B:→B:Z , B#μs=μs+tfor all s∈[0,τ]⟹B=T⁡(t).B_\# _s= _s+t all s∈[0,τ] B=T(t). (39) For a fixed lag, FT identifies the propagator from its action on the entire family of source distributions. At t=0t=0, it rules out nonidentity linear symmetries shared by all these distributions. Allowing every bounded linear candidate B, including noninvertible ones, ensures that the condition applies to the operator limits considered below. FT supplies uniqueness; stability and compactness will provide the remaining ingredients for convergence. The next proposition relates FT to the linear span of the distribution supports and gives a sufficient condition in terms of their means. Proposition A.12 (Total linear support and a sufficient criterion for FT). Let μ¯=τ−1∫0τμss μ=τ^-1 _0^τ _s\,ds. Under FT, span⁡(supp⁡μ¯)¯=. span(supp μ)=Z. (40) A sufficient condition for FT is span⁡s:0≤s≤τ¯=,s=∫dμs(). span\ m_s:0≤ s≤τ\=Z, m_s= z\,d _s( z). (41) Proof. If Equation (40) failed, there would be a nonzero h with ⟨,⟩=0 h, z =0 for μ¯ μ-almost every z. Consequently ∫0τ|⟨,s⟩|2s=0. _0^τE| h, Z_s |^2ds=0. The integrand is nonnegative and continuous in s, because s Z_s is L2L^2-continuous. It is therefore zero for every s∈[0,τ]s∈[0,τ]. The bounded linear map B=+⟨,⟩B z= z+ h, z h fixes s Z_s almost surely for each s, and hence B#μs=μsB_\# _s= _s for all s. But B≠IB≠ I, contradicting FT at t=0t=0. This proves the first claim. For the sufficient condition, take B satisfying the premise of FT for a fixed lag t. Finite second moments allow Bochner expectations, and bounded linear maps commute with them. Thus Bs=s+t=T⁡(t)s(0≤s≤τ).B m_s= m_s+t=T(t) m_s (0≤ s≤τ). The bounded operator B−T⁡(t)B-T(t) vanishes on the dense linear span in Equation (41), so it vanishes everywhere. This proves FT. ∎ A.4.3 From distributional compactness to strong operator compactness Theorem A.10 establishes W2W_2 convergence of the pushforward measures SD(t)#μsS_D(t)_\# _s to μs+t _s+t as the weak residual vanishes. To obtain convergence of the linear propagators, we next show that relative compactness of (Bn)#ν\(B_n)_\#ν\ in W2W_2 yields strongly convergent operator subsequences when (Bn)(B_n) is uniformly bounded and supp⁡νsuppν has dense linear span in Z. In our setting, stability supplies the bound, and Proposition A.12 establishes the support condition for the mixture μ¯=τ−1∫0τμss μ=τ^-1 _0^τ _s\,ds. The following lemma provides the required compactness; the next subsection uses FT to identify each operator limit with the true propagator T⁡(t)T(t). Lemma A.13 (From W2W_2 compactness to strong operator compactness). Let Z be a separable real Hilbert space and let ν∈2()ν _2(Z) satisfy span⁡(supp⁡ν)¯= span(suppν)=Z. Suppose Bn∈ℒ⁡()B_n (Z), supn‖Bn‖≤R<∞ _n\|B_n\|≤ R<∞, and (Bn)#ν\(B_n)_\#ν\ is relatively compact in W2W_2. Then every subsequence of (Bn)(B_n) has a further subsequence and a bounded linear map B, ‖B‖≤R\|B\|≤ R, such that ‖Bn−B‖⟶0for every ∈\|B_n z-B z\| 0 every z (42) along that further subsequence. Proof. Start with any subsequence and retain the index n. Choose a countable dense set j\ q_j\ in Z. The uniform operator bound allows a diagonal extraction for which ⟨Bni,j⟩ B_n q_i, q_j converges for every i,ji,j. By density and the same bound, these limits extend to a bilinear form b⁡(,)b( x, y) with |b⁡(,)|≤R‖‖|b( x, y)|≤ R\| x\|\| y\|. Riesz representation gives B∈ℒ⁡()B (Z) with b⁡(,)=⟨B,⟩b( x, y)= B x, y and ‖B‖≤R\|B\|≤ R, so BnB_n converges to B in the weak operator topology. Extracting once more using the W2W_2 relative compactness gives Bn→WOTB,(Bn)#ν→W2λfor some λ∈2().B_n WOTB, (B_n)_\#ν W_2λ some λ _2(Z). (43) For every ∈ h , weak operator convergence and dominated convergence imply ∫ei⟨,⟩λ()=limn∫ei⟨,Bn⟩ν()=∫ei⟨,B⟩ν(). e^i h, y \,dλ( y)= _n e^i h,B_n z \,dν( z)= e^i h,B z \,dν( z). The first equality follows from W2W_2 convergence, which implies weak convergence of the measures (Bn)#ν(B_n)_\#ν to λ. Thus λ and B#νB_\#ν have the same characteristic functional. Restricting h to each finite span of an orthonormal basis gives equality of all finite-coordinate distributions. These coordinates generate the Borel sigma field of Z, since it is separable, and hence λ=B#νλ=B_\#ν. Convergence in W2W_2 also gives convergence of second moments, so ∫‖Bn‖2ν()→∫‖B‖2ν() \|B_n z\|^2\,dν( z)→ \|B z\|^2\,dν( z). Indeed, the difference between the square roots of these moments is bounded by W2((Bn)#ν,B#ν)W_2((B_n)_\#ν,B_\#ν), by the reverse triangle inequality in any transport coupling. For each fixed z, weak operator convergence gives ⟨Bn,B⟩→‖B‖2 B_n z,B z →\|B z\|^2, and |⟨Bn,B⟩|≤R2‖2| B_n z,B z |≤ R^2\| z\|^2 is ν-integrable. Dominated convergence therefore gives convergence of the corresponding integrals. Expanding the square yields ∫‖(Bn−B)‖2ν()=∫(‖Bn‖2+‖B‖2−2⟨Bn,B⟩)ν()⟶0. \|(B_n-B) z\|^2\,dν( z)= (\|B_n z\|^2+\|B z\|^2-2 B_n z,B z )\,dν( z) 0. (44) To deduce pointwise convergence on supp⁡νsuppν from convergence in L2(ν,)L^2(ν;Z), fix ∈supp⁡ν z ν and δ>0δ>0. The open ball δ()B_δ( z) has positive ν-measure. For y in this ball, ‖(Bn−B)‖≤2Rδ+‖(Bn−B)‖\|(B_n-B) z\|≤ 2Rδ+\|(B_n-B) y\|. Averaging over the ball and applying Cauchy–Schwarz gives ‖(Bn−B)‖≤2Rδ+(∫‖(Bn−B)‖2ν())1/2ν(δ())1/2.\|(B_n-B) z\|≤ 2Rδ+ ( \|(B_n-B) y\|^2\,dν( y) )^1/2ν(B_δ( z))^1/2. Letting n→∞n→∞ and then δ↓0δ 0 proves convergence at every point of supp⁡νsuppν along the same subsequence. Linearity extends it to the span of this support, which is dense by assumption. Finally, for any ∈ z and any x in this span, ‖(Bn−B)‖≤2R‖−‖+‖(Bn−B)‖\|(B_n-B) z\|≤ 2R\| z- x\|+\|(B_n-B) x\|. Approximating z by such x proves Equation (42). The starting subsequence was arbitrary, so the asserted strong operator relative compactness follows. ∎ In the next subsection, we apply the lemma with ν=μ¯ν= μ and Bn=SDn(tn)B_n=S_D_n(t_n). Theorem A.10 and W2W_2 continuity of t↦μt _t establish the required compactness of (SDn(tn))#μ¯\(S_D_n(t_n))_\# μ\, even when the time lags tnt_n vary. Along subsequences with tn→t_n→ t, FT identifies the strong operator limit B through the identities B#μs=μs+tB_\# _s= _s+t for all s∈[0,τ]s∈[0,τ]. A.4.4 Locally uniform strong convergence of the learned propagators Theorem A.10 establishes convergence of the pushforward measures SD(t)#μsS_D(t)_\# _s, and Lemma A.13 converts W2W_2 compactness of these measures into strong operator subsequences. We now combine these results with the finite-time linear determinacy condition in Assumption A.11 to identify every subsequential limit with T⁡(t)T(t). Allowing the time lags to vary along each subsequence will establish convergence uniformly on [0,τ][0,τ] for every fixed ∈ z . The semigroup property then extends the result to any finite time interval. Theorem A.14 (Locally uniform strong convergence under stability and FT). In the setting of Sections A.3.1 and A.3.2, assume 0∈D⁡() Z_0∈ D( A) almost surely and ⁡(‖0‖2+‖0‖2)<∞E(\| Z_0\|^2+\| A Z_0\|^2)<∞. Let PDP_D be the increasing finite-rank orthogonal projections onto DZ_D, with PD→P_D z→ z for every ∈ z , and choose D∈argmin∈ℒ⁡(D)D() A_D∈ *argmin_ B (Z_D)J_D( B) for the weak population objective Equation (38). Suppose the selected minimizers satisfy the finite-time stability condition in Assumption A.8, and the distribution family μs0≤s≤2τ\ _s\_0≤ s≤ 2τ satisfies FT (Assumption A.11). Then, for every ∈ z and H>0H>0, sup0≤t≤H‖KD(t)−T⁡(t)‖⟶0(D→∞). _0≤ t≤ H\|K_D(t) z-T(t) z\| 0 (D→∞). (45) The same conclusion holds for SD(t)=KD(t)+(I−PD)S_D(t)=K_D(t)+(I-P_D). Proof. Step 1: Wasserstein convergence of the averaged measures. Theorem A.10 gives eD:=sups,t∈[0,τ]W2(SD(t)#μs,μs+t)⟶0.e_D:= _s,t∈[0,τ]W_2(S_D(t)_\# _s, _s+t) 0. (46) For t∈[0,τ]t∈[0,τ], define μ¯=1τ∫0τμss,μ¯t=1τ∫0τμs+ts. μ= 1τ _0^τ _s\,ds, μ_t= 1τ _0^τ _s+t\,ds. Both measures belong to 2()P_2(Z) because r↦r Z_r is continuous in L2(Ω,)L^2( ;Z) on [0,2τ][0,2τ]. Proposition A.12 gives span⁡(supp⁡μ¯)¯= span(supp μ)=Z. For r,t∈[0,τ]r,t∈[0,τ], using the same 0 Z_0 and an independent uniform time in [0,τ][0,τ] to couple the averaged measures yields W2(μ¯t,μ¯r)2≤1τ∫0τ‖s+t−s+r‖L22s⟶0(t→r).W_2( μ_t, μ_r)^2≤ 1τ _0^τ\| Z_s+t- Z_s+r\|_L^2^2\,ds 0 (t→ r). Here the convergence follows from uniform L2L^2 continuity of r Z_r on [0,2τ][0,2τ]. Averaging the bounds in Equation (46) gives W2(SD(t)#μ¯,μ¯t)2≤1τ∫0τW2(SD(t)#μs,μs+t)2s≤eD2.W_2(S_D(t)_\# μ, μ_t)^2≤ 1τ _0^τW_2(S_D(t)_\# _s, _s+t)^2\,ds≤ e_D^2. (47) To justify this inequality, first approximate each time integral by a Riemann sum on a uniform partition. Averaging optimal couplings for the finitely many pairs of measures proves the inequality for these sums. For fixed D,tD,t, the discrete averages converge in W2W_2 to SD(t)#μ¯S_D(t)_\# μ and μ¯t μ_t: for s in a partition interval with left endpoint sjs_j, use the pairs (SD(t)s,SD(t)sj)(S_D(t) Z_s,S_D(t) Z_s_j) and (s+t,sj+t)( Z_s+t, Z_s_j+t), respectively. Uniform L2L^2 continuity and boundedness of SD(t)S_D(t) make the coupling costs vanish. The integrand on the right is continuous in s, since both curves of measures are W2W_2-continuous. Passing to the limit proves Equation (47). Step 2: Strong operator subsequences for convergent time lags. Take arbitrary sequences Dn→∞D_n→∞ and tn∈[0,τ]t_n∈[0,τ]. After passing to a subsequence, compactness of [0,τ][0,τ] gives tn→t∗t_n→ t_*. By Equation (47) and continuity of t↦μ¯t μ_t, W2((SDn(tn))#μ¯,μ¯t∗)≤eDn+W2(μ¯tn,μ¯t∗)⟶0.W_2((S_D_n(t_n))_\# μ, μ_t_*)≤ e_D_n+W_2( μ_t_n, μ_t_*) 0. Assumption A.8 gives ‖SDn(tn)‖≤M\|S_D_n(t_n)\|≤ M. Since supp⁡μ¯supp μ has dense linear span, Lemma A.13 applies and provides a further subsequence, with the same notation, such that SDn(tn)⟶Bfor every ∈,B∈ℒ⁡(),‖B‖≤M.S_D_n(t_n) z B z every z , B (Z), \|B\|≤ M. (48) Step 3: Identification of subsequential limits under FT. Fix s∈[0,τ]s∈[0,τ]. The pair (SDn(tn)s,Bs)(S_D_n(t_n) Z_s,B Z_s) gives W2((SDn(tn))#μs,B#μs)2≤‖(SDn(tn)−B)s‖2⟶0.W_2((S_D_n(t_n))_\# _s,B_\# _s)^2 \|(S_D_n(t_n)-B) Z_s\|^2 0. Indeed, Equation (48) gives pointwise convergence, and the integrand is bounded by 4M2‖s‖24M^2\| Z_s\|^2, which is integrable. Combining this limit with Equation (46) and W2W_2 continuity of r↦μr _r gives W2(B#μs,μs+t∗)≤W2(B#μs,(SDn(tn))#μs)+eDn+W2(μs+tn,μs+t∗)⟶0. W_2(B_\# _s, _s+t_*)≤ W_2(B_\# _s,(S_D_n(t_n))_\# _s)+e_D_n+W_2( _s+t_n, _s+t_*) 0. The subsequence in Equation (48) is independent of s and converges at every ∈ z . Thus the same argument applies to every s∈[0,τ]s∈[0,τ] without further extraction, establishing B#μs=μs+t∗B_\# _s= _s+t_* for all such s. Assumption A.11 therefore implies B=T⁡(t∗)B=T(t_*), including B=IB=I when t∗=0t_*=0. Step 4: Uniform strong convergence on [0,τ][0,τ]. If convergence failed to be uniform for some fixed ∈ z , there would exist ε>0 >0, Dn→∞D_n→∞, and tn∈[0,τ]t_n∈[0,τ] with ‖SDn(tn)−T⁡(tn)‖≥ε\|S_D_n(t_n) z-T(t_n) z\|≥ . Steps 2–3 give a further subsequence for which tn→t∗t_n→ t_* and SDn(tn)→T⁡(t∗)S_D_n(t_n) z→ T(t_*) z. Strong continuity gives T⁡(tn)→T⁡(t∗)T(t_n) z→ T(t_*) z, a contradiction. Consequently, sup0≤t≤τ‖SD(t)−T⁡(t)‖⟶0for every ∈. _0≤ t≤τ\|S_D(t) z-T(t) z\| 0 every z . (49) Step 5: Extension to arbitrary finite time intervals. For each fixed integer k≥0k≥ 0, SD(τ)k→T(τ)kS_D(τ)^k z→ T(τ)^k z. This follows inductively from Equation (49) and ‖SD(τ)‖≤M\|S_D(τ)\|≤ M, using SD(τ)k+1−T(τ)k+1=SD(τ)(SD(τ)k−T(τ)k)+(SD(τ)−T⁡(τ))T(τ)k. S_D(τ)^k+1 z-T(τ)^k+1 z=S_D(τ) (S_D(τ)^k z-T(τ)^k z )+ (S_D(τ)-T(τ) )T(τ)^k z. For t=kτ+rt=kτ+r with 0≤r<τ0≤ r<τ, the semigroup property gives ‖SD(t)−T⁡(t)‖≤M‖SD(τ)k−T(τ)k‖+‖(SD(r)−T⁡(r))T(τ)k‖.\|S_D(t) z-T(t) z\|≤ M\|S_D(τ)^k z-T(τ)^k z\|+\|(S_D(r)-T(r))T(τ)^k z\|. For t∈[0,H]t∈[0,H], only finitely many k occur. The first term tends to zero uniformly in r, and Equation (49), applied to the finitely many fixed vectors T(τ)kT(τ)^k z, makes the second term vanish uniformly in r. Hence sup0≤t≤H‖SD(t)−T⁡(t)‖→0 _0≤ t≤ H\|S_D(t) z-T(t) z\|→ 0. Finally, SD(t)−KD(t)=I−PDS_D(t)-K_D(t)=I-P_D gives sup0≤t≤H‖KD(t)−T⁡(t)‖≤sup0≤t≤H‖SD(t)−T⁡(t)‖+|(I−PD)|⟶0, _0≤ t≤ H\|K_D(t) z-T(t) z\|≤ _0≤ t≤ H\|S_D(t) z-T(t) z\|+\|(I-P_D) z\| 0, which proves Equation (45). ∎ The theorem establishes strong operator convergence locally uniformly in time, including on finite intervals beyond [0,2τ][0,2τ]. The extension relies on the global semigroup property after stability and FT have established convergence on [0,τ][0,τ]. Korda and Mezić (2018) analyze EDMD, which fits the Koopman action on a finite observable space by least squares from paired state observations. For a fixed observable space, the infinite-sample limit projects each evolved observable back onto that space in the L2L^2 inner product induced by the sampling measure. Under their assumptions, strong operator convergence then follows as the observable space is enlarged. Here, convergence follows from the weak population objective, compactness of pushforward measures, and FT. Limitations. Our convergence analysis as D→∞D→∞ assumes an underlying autonomous vector field whose induced distributions form an absolutely continuous Wasserstein flow. It concerns minimizers of the exact weak population residual over a compatible family of observables obtained from nested projections of a fixed Koopman-invariant RKHS. Under the regularity conditions in Sections A.3.1–A.3.2, the stability condition in Assumption A.8, and finite-time linear determinacy (FT) in Assumption A.11, the feature propagators KD(t)K_D(t) and SD(t)S_D(t) converge strongly to T⁡(t)=U(t)∗T(t)=U(t)^*, locally uniformly in time. In practice, only finite samples at sparse observation times are available, so we discretize the time integrals using trapezoidal quadrature. We further use ridge regularization and exponential moving average (EMA) updates to stabilize the estimated generator during alternating optimization of neural representations and dynamics (Appendix B). The analysis does not control the resulting sampling, quadrature, regularization, or optimization errors, nor establish convergence of this training procedure or of independently learned neural representations as D increases. Strong convergence of the adjoints to U⁡(t)U(t), operator-norm convergence, and spectral convergence also remain outside these guarantees. Appendix B Algorithmic Details This appendix develops the training objectives and algorithms for KoopCell and its memory-augmented extension, KoopCell-M, introduced in Sections 3.1 and 3.2. We first describe how KoopCell alternates between VAE representation learning and closed-form estimation of the latent dynamics. We then introduce the memory embedding and developmental-path conditioning in KoopCell-M, detailing the resulting inference and training procedures. Throughout, ρtk=nk−1∑i=1nkδk,i _t_k=n_k^-1 _i=1^n_k _ x_k,i denotes an observed snapshot, with k,i∈ℝdx x_k,i ^d_x and =0=t0<⋯<tN=TT=\0=t_0<·s<t_N=T\ the training times. We retain ,ϕ θ, φ for the encoder and decoder parameters and D for the latent dimension. Held-out snapshots are used only for evaluation. B.1 KoopCell: Koopman dynamics and alternating training KoopCell alternates between learning a VAE representation for a given latent dynamics generator and estimating the linear dynamics generator in the current latent coordinates. We describe the VAE objective first, then derive the empirical generator update from the finite-test formulation proposed in Theorem 3.1. B.1.1 VAE representation learning For a cell with expression x observed at time t, we use a posterior encoder q(⋅∣,t)=((,t),(,t))q_ θ(· x,t)=N( m_ θ( x,t), S_ θ( x,t)) with diagonal covariance S_ θ. A shared MLP takes (,t)( x,t) as input and predicts the mean and log-variance. The observation model is pϕ(∣)=⁡(,Decϕ⁡(),σx2dx)p_ φ( x z)=N( x;Dec_ φ( z), _x^2 I_d_x), with fixed variance σx2 _x^2. Reconstruction and population matching use reparameterized samples =+1/2ϵ z= m_ θ+ S_ θ^1/2 ε, ϵ∼⁡(,D) ε ( 0, I_D), and predicted expression is given by Decϕ⁡()Dec_ φ( z). The resulting encoded population is μt=∫q(⋅∣,t)dρt() _t= q_ θ(· x,t)d _t( x). Given the linear latent flow ℱt()=t+tF_t( z)= F_t z+ c_t (Appendix B.1.2), we regularize the representation with the transported prior pt=(ℱt)#(,D)=⁡(t,t)p_t=(F_t)_\#N( 0, I_D)=N( c_t, _t), where t=tt⊤ _t= F_t F_t . Here μt _t denotes the encoder-induced population and ptp_t the generative prior. We write q(⋅∣,t)q_ θ(· x,t) for the encoder distribution and q(∣,t)q_ θ( z x,t) for its density at z. Within fixed-time expectations, we suppress time subscripts on random states. The resulting generative model pϕ,t(,)=pϕ(∣)pt()p_ φ,t( x, z)=p_ φ( x z)p_t( z) yields pϕ,t()=∼q(⋅∣,t)[pϕ(∣)pt()/q(∣,t)]p_ φ,t( x)=E_ Z q_ θ(· x,t)[p_ φ( x Z)p_t( Z)/q_ θ( Z x,t)]. Jensen’s inequality therefore gives the evidence lower bound (ELBO) logpϕ,t()≥ℰKC(,t):=qlogpϕ(∣)−KL(q(⋅∣,t)∥pt). p_ φ,t( x) _KC( x,t):=E_q_ θ p_ φ( x Z)-KL (q_ θ(· x,t)\,\|\,p_t ). (50) The Gaussian observation model gives a squared reconstruction error up to an additive constant. Following Equation (4), we normalize it by the expression dimension: ℒrec=−dx−1t,qlog⁡pϕ(∣)L_rec=-d_x^-1E_t, xE_q_ θ p_ φ( x Z), where t,E_t, x averages equally over training times and uniformly over cells within each time. Writing =(,t) m= m_ θ( x,t) and =(,t) S= S_ θ( x,t), we compute the Gaussian KL divergence analytically: KL((,)∥pt)=12[tr(t−1)+(−t)⊤t−1(−t)−D+logdettdet].KL (N( m, S)\,\|\,p_t )= 12 [tr( _t^-1 S)+( m- c_t) _t^-1( m- c_t)-D+ _t S ]. (51) The training objective ℒVAE=ℒrec+βt,KL(q(⋅∣,t)∥pt)L_VAE=L_rec+ _t, xKL(q_ θ(· x,t)\,\|\,p_t) thus equals the negative ELBO divided by dxd_x when β=1/dxβ=1/d_x; otherwise, β controls the strength of the variational regularization. To learn representations whose evolution under the latent flow matches the observed populations, we augment ℒVAEL_VAE with prediction losses in latent and expression space. For every forward pair (s,t)∈=(s,t)∈2:s<t(s,t) =\(s,t) ^2:s<t\, we form μ^s→t=(ℱt−s)#μs μ_s→ t=(F_t-s)_\# _s and ρ^s→t=(Decϕ)#μ^s→t ρ_s→ t=(Dec_ φ)_\# μ_s→ t, and average the prediction losses over these pairs: ℒpop=1||∑(s,t)∈[λzε(μt,μ^s→t)+λxε(ρt,ρ^s→t)].L_pop= 1|P| _(s,t) [ _zS_ ( _t, μ_s→ t)+ _xS_ ( _t, ρ_s→ t) ]. (52) Here εS_ is the debiased Sinkhorn divergence with quadratic cost. With the dynamics fixed, we update ,ϕ θ, φ by minimizing ℒKC=ℒVAE+ℒpop+λweakℒweakL_KC=L_VAE+L_pop+ _weakL_weak, where the weak-residual loss is defined below. Reconstruction and prediction losses update both the encoder and decoder; the weak-residual term updates only the encoder. To initialize the representation before alternating training, we first pretrain the VAE with a standard-normal prior at every training time, minimizing ℒpre=ℒrec+βpret,KL(q(⋅∣,t)∥(,D)),L_pre=L_rec+ _preE_t, xKL (q_ θ(· x,t)\,\|\,N( 0, I_D) ), (53) with a gradually increased βpre _pre. The pretrained encoder then provides the initial latent coordinates for the linear dynamics estimation described next. B.1.2 Empirical weak residual and closed-form dynamics estimation With the representation fixed, we now derive a closed-form estimator for the linear generator that best fits the evolution of the latent distributions. For notational simplicity, Section 3.1 presents the method using ˙t=t Z_t= A Z_t. In the implementation, we include the constant observable 11, giving ˙t=t+z Z_t= A Z_t+ b_z and coefficients =[,z]∈ℝD×(D+1) =[ A, b_z] ^D×(D+1). Writing ~=(⊤,1)⊤ z=( z ,1) and ~=(z⊤0) A= ( smallmatrix A& b_z\\ 0 &0 smallmatrix ), we evaluate the flow through exp⁡(Δ~)=(Δ⊤1),ℱΔ()=Δ+Δ. ( A)= pmatrix F_ & c_ \\ 0 &1 pmatrix, _ ( z)= F_ z+ c_ . (54) The constant coordinate stays equal to one, and only the upper D rows are estimated. Theorem 3.1 shows that restricting the dual supremum to a finite test space gives a least-squares regression objective for estimating the generator; Lemma A.4 in Appendix A.2 provides the corresponding derivation. For the linear drift ~ z, the moment definitions in Equation (3) require only replacing ⊤ z by ~⊤ z in j H_j; y and M are unchanged. Thus, with =vec⁡() a=vec( ), the residual vector remains − y- G a and the objective is 12(−)⊤−1(−) 12( y- G a) M^-1( y- G a) when ≻0 M 0. We next construct its empirical counterpart from the observed snapshots by forming these moments on adjacent time intervals. Specifically, on each interval Ik=[tk,tk+1]I_k=[t_k,t_k+1], we use a smooth spatial test ψj _j with bounded values and gradients. The residual identity Equation (14) in Appendix A.1, applied to φkj(t,)=Ik(t)ψj() _kj(t, z)=1_I_k(t) _j( z), gives22 2 The interval test is understood by smooth approximation in DV_D, using the continuous residual extension in Equation (15). ℛD(,φkj)=μtk+1[ψj]−μtk[ψj]−⟨,∫Ikμt[∇ψj()~⊤]t⟩F.R_D( ; _kj)= _t_k+1[ _j]- _t_k[ _j]- , _I_k _t[∇ _j( z) z ]dt _F. (55) Here the endpoint difference supplies ykjy_kj, and the time integral supplies kj H_kj in the finite-test construction. The Gram matrix defined in Equation (27) is block diagonal, =diag⁡(0,…,N−1) M=diag( M_0,…, M_N-1); each block is (k)ij=∫Ikμt[∇ψi⊤∇ψj]dt( M_k)_ij= _I_k _t[∇ _i ∇ _j]dt. In the implementation, we use random Fourier tests cos⁡(j⊤) ( ξ_j z) and sin⁡(j⊤) ( ξ_j z) to construct this finite test space. Frequencies are sampled once from ⁡(,s2D/D)N( 0,s^2 I_D/D) in equal-sized groups for s∈0.5,1,2,4s∈\0.5,1,2,4\, then held fixed during training. The standard bank contains 1,024 frequencies and hence 2,048 tests. Their gradients, −jsin⁡(j⊤)- ξ_j ( ξ_j z) and jcos⁡(j⊤) ξ_j ( ξ_j z), allow all required moments to be evaluated directly from encoded cells. For the empirical regression, we choose the encoder means as the deterministic observables to which the theorem is applied. At each training time tkt_k, we encode the observed cells as k,i=(k,i,tk) z_k,i= m_ θ( x_k,i,t_k), yielding the empirical latent distribution μtkmean=((⋅,tk))#ρtk _t_k^mean=( m_ θ(·,t_k))_\# _t_k. With ~k,i=(k,i⊤,1)⊤ z_k,i=( z_k,i ,1) , the endpoint statistics are ukj=1nk∑iψj(k,i),kj=1nk∑i∇ψj(k,i)~k,i⊤.u_kj= 1n_k _i _j( z_k,i), D_kj= 1n_k _i∇ _j( z_k,i) z_k,i . (56) In Equation (55), ukju_kj replaces the endpoint expectation μtk[ψj] _t_k[ _j], and kj D_kj replaces the integrand μtk[∇ψj()~⊤] _t_k[∇ _j( z) z ] at tkt_k. Approximating the time integral by the trapezoidal rule, with Δtk=tk+1−tk t_k=t_k+1-t_k, therefore gives ykj=uk+1,j−ukj,kj=Δtk2(kj+k+1,j),(k,j),:=vec(kj)⊤.y_kj=u_k+1,j-u_kj, H_kj= t_k2( D_kj+ D_k+1,j), G_(k,j),:=vec( H_kj) . (57) Following the regression interpretation in the remark after Lemma A.4, the empirical residual over interval Ik=[tk,tk+1]I_k=[t_k,t_k+1] is ℛ^kj()=ykj−⟨,kj⟩F=ykj−()(k,j) R_kj( )=y_kj- , H_kj _F=y_kj-( G a)_(k,j). This is the empirical, time-discretized approximation of ℛD(,φkj)R_D( ; _kj) in Equation (55). The construction accommodates irregular observation times; further, if the moment integrand has two bounded time derivatives, the interval quadrature error is ⁡(Δtk3)O( t_k^3). Applying the same quadrature to the Gram matrix gives (k)ij=Δtk2μtkmean[∇ψi⊤∇ψj]+μtk+1mean[∇ψi⊤∇ψj].( M_k)_ij= t_k2 \ _t_k^mean[∇ _i ∇ _j]+ _t_k+1^mean[∇ _i ∇ _j] \. (58) Finally, we stack these empirical residuals over all interval–test pairs (k,j)(k,j) to obtain − y- G a, the residual vector in the finite-test objective of Equation (28). With these empirical quantities and =(+ϵM)−1 W=( M+ _M I)^-1, we define ℒweak=(−)⊤(−)L_weak=( G a- y) W( G a- y). To stabilize the estimation when G is poorly conditioned, we additionally introduce a ridge penalty: ^=argmin⁡12(−)⊤(−)+λ2‖22=(⊤+λ)−1⊤,λ>0. a= *arg\,min_ a \ 12( G a- y) W( G a- y)+ λ2\| a\|_2^2 \=( G W G+λ I)^-1 G W y, λ>0. (59) We form the regression statistics from all training cells and solve the normal equations using preconditioned conjugate gradients. The fitted coefficients ^=reshape⁡(^) =reshape( a) are then used to update the dynamics before refining the representation. However, the pretrained latent distributions need not follow a common linear flow, so directly using in the representation objective can induce abrupt updates and distort the pretrained latent manifold. We therefore introduce the closed-form dynamics generator gradually through an exponential moving average, ←α+(1−α) ←α +(1-α) , with α=0.99α=0.99 in the real single-cell dynamics learning procedure. With the updated held fixed, we optimize the VAE with the objective in Appendix B.1.1, recomputing the weak residual on minibatches. Algorithm 1 summarizes the alternating procedure. In both algorithms, sgsg stops gradients, ‖2=⊤\| v\|_ W^2= v W v, and AdamW⁡(,∇ℒ)AdamW( w, _ wL) denotes one optimizer update of parameters w using the indicated gradient. In Algorithm 1, WeakStats⁡(k,i,,ψj)WeakStats(\ z_k,i\;T,\ _j\) returns (,,)( G, y, M) using Equation (56), Equation (57), and Equation (58); the same operation is used for full snapshots and minibatches. Transition⁡(,Δ)Transition( , ) returns (Δ,Δ)( F_ , c_ ) via Equation (54), and PCG⁡(,)PCG( Q, r) solves = Q a= r by preconditioned conjugate gradients. The loss routines evaluate the referenced objectives on the indicated minibatches ℬ=ℬkk=0NB=\B_k\_k=0^N. SampleSample draws nbn_b cells independently from each snapshot. The pairwise predictions depend only on the source samples and fixed dynamics, so they can be computed in parallel. Algorithm 1 KoopCell: alternating closed-form generator estimation and representation learning 1: Training snapshots (tk,ρtk)k=0N\(t_k, _t_k)\_k=0^N; fixed Fourier tests ψjj=1m\ _j\_j=1^m; batch size nbn_b; λ,ϵM>0λ, _M>0; α∈[0,1)α∈[0,1); loss weights 2: Encoder q_ θ, decoder DecϕDec_ φ, and linear dynamics parameters =[,z] =[ A, b_z] 3: Stage 1: VAE pretraining 4: for each pretraining iteration do 5: ℬk←Sample(ρtk,nb),k=0,…,NB_k ( _t_k,n_b), k=0,…,N ⊳ Independent sampling at each time 6: ℒpre←PretrainLoss⁡(ℬ,,ϕ,βpre)L_pre (B; θ, φ, _pre) ⊳ Standard-normal VAE; Equation (53) 7: (,ϕ)←AdamW⁡((,ϕ),∇,ϕℒpre)( θ, φ) (( θ, φ), _ θ, φL_pre) ⊳ Pretrain encoder and decoder 8: end for 9: Stage 2: Alternating dynamics estimation and representation learning 10: ← ← 0 11: for each training iteration do 12: (a) Closed-form dynamics estimation 13: ¯k,i←sg((k,i,tk)),k=0,…,N,i=1,…,nk z_k,i ( m_ θ( x_k,i,t_k)), k=0,…,N,\ i=1,…,n_k 14: (,,)←WeakStats⁡(¯k,i,,ψj)( G, y, M) (\ z_k,i\;T,\ _j\); ←(+ϵM)−1 W←( M+ _M I)^-1 15: ^←PCG⁡(⊤+λ,⊤) a ( G W G+λ I, G W y); ^←reshape⁡(^) ( a) ⊳ closed-form estimation; Equation (59) 16: ←sg⁡(α+(1−α)^) (α +(1-α) ) ⊳ EMA update 17: (b) Representation update 18: ℬk←Sample(ρtk,nb),k=0,…,NB_k ( _t_k,n_b), k=0,…,N ⊳ Fresh minibatches for VAE updates 19: ¯k,iℬ←(k,i,tk),k,i∈ℬk z^B_k,i← m_ θ( x_k,i,t_k), x_k,i _k; ϵk,i∼⁡(,D) ε_k,i ( 0, I_D) 20: k,i←¯k,iℬ+(k,i,tk)1/2ϵk,i z_k,i← z^B_k,i+ S_ θ( x_k,i,t_k)^1/2 ε_k,i ⊳ Reparameterized posterior samples 21: (Δ,Δ)←Transition⁡(,Δ),Δ∈∪t−s:(s,t)∈( F_ , c_ ) ( , ), ∪\t-s:(s,t) \ ⊳ Calculate the transition matrix 22: ptk←⁡(tk,tktk⊤)p_t_k ( c_t_k, F_t_k F_t_k ) ⊳ Transported Gaussian prior 23: ℒVAE←VAELoss⁡(ℬ,,ϕ,ptk)L_VAE (B; θ, φ,\p_t_k\) ⊳ Reconstruction + KL; Equation (4) 24: for each (s,t)∈(s,t) in parallel do 25: ^s→t,i←t−ss,i+t−s z_s→ t,i← F_t-s z_s,i+ c_t-s; ^s→t,i←Decϕ⁡(^s→t,i) x_s→ t,i _ φ( z_s→ t,i) 26: end for 27: ℒpop←PopulationLoss⁡(ℬ,k,i,^s→t,i,^s→t,i)L_pop (B,\ z_k,i\,\ z_s→ t,i, x_s→ t,i\) ⊳ Population matching; Equation (52) 28: (ℬ,ℬ,ℬ)←WeakStats⁡(¯k,iℬ,,ψj)( G_B, y_B, M_B) (\ z^B_k,i\;T,\ _j\) 29: ℬ←(ℬ+ϵM)−1 W_B←( M_B+ _M I)^-1; ℒweak←‖ℬvec⁡()−ℬ‖ℬ2L_weak←\| G_Bvec( )- y_B\|_ W_B^2 ⊳ Weak population residual loss 30: ℒKC←ℒVAE+ℒpop+λweakℒweakL_KC _VAE+L_pop+ _weakL_weak ⊳ Combined training objective 31: (,ϕ)←AdamW⁡((,ϕ),∇,ϕℒKC)( θ, φ) (( θ, φ), _ θ, φL_KC) ⊳ Update encoder and decoder 32: end for 33: return q,Decϕ,,zq_ θ,Dec_ φ, A, b_z For evaluation, we sample 0 Z_0 from the encoder distribution of the initial observed population snapshot and return Decϕ⁡(ℱt(0))Dec_ φ(F_t( Z_0)) at each requested time. B.2 KoopCell-M: memory-augmented branching dynamics and alternating training This section develops the memory extension introduced in Section 3.2. We first introduce the augmented latent dynamics for modeling branching, then derive the probabilistic model and training objectives. The resulting algorithm retains the two alternating steps of KoopCell: closed-form dynamics estimation followed by representation learning. B.2.1 Memory-augmented dynamics and developmental paths Zhao et al. (2026, Proposition 3.3) establishes that an ODE with a Lipschitz vector field cannot exactly transport between distributions supported on different numbers of disjoint compact, path-connected components. In particular, a linear flow cannot split a connected initial population into disconnected fate populations; continuous decoding also preserves connectedness. When branching is understood as the emergence of separated density modes, a linear latent flow has a further limitation: it preserves Gaussianity, so a single Gaussian population remains unimodal. Although nonlinear decoding can introduce additional density modes, a decoder with Lipschitz constant LdecL_dec limits the amplification of latent separation: ‖Decϕ⁡(ℱt())−Decϕ⁡(ℱt(′))‖≤Ldec‖et‖‖−′‖.\|Dec_ φ(F_t( z))-Dec_ φ(F_t( z ))\|≤ L_dec\|e At\|\,\| z- z \|. (60) To model branching in the latent dynamics, we draw on the Mori–Zwanzig formalism: unresolved observables influence the resolved dynamics through memory (Section 2). Specifically, we augment the visible latent state t∈ℝD Z_t ^D, which represents gene expression, with memory variables t∈ℝdh H_t ^d_h that influence the dynamics of t Z_t. The decoder reconstructs gene expression from t Z_t alone. The memory variables can represent fate-relevant influences not captured by the measured expression state (Weinreb et al., 2020). Further, a linear augmented process with a single joint Gaussian initialization would still have Gaussian marginals. To generate distinct developmental branches, we therefore introduce a discrete variable R∈ℛR∈ R indexing heterogeneous hidden initial conditions: R∼Cat(),0∣R=r∼(,D),0∣0=,R=r∼(r+r,r).R ( π), Z_0 R=r ( 0, I_D), H_0 Z_0= z,R=r ( F_r z+ a_r, S_r). (61) Here πr _r are normalized mixture weights and r≻0 S_r 0. Conditional on R=rR=r, we evolve the augmented state (t,t)( Z_t, H_t) by dt=(t+t+z)dt,dt=(t+t+h(r))dt+hdt,d Z_t=( A Z_t+ B H_t+ b_z)\,dt, H_t=( C Z_t+ H_t+ b_h(r))\,dt+ _h\,d W_t, (62) where t W_t is a standard Brownian motion. Motivated by the use of decaying memory in related models (Lei et al., 2016; Menier et al., 2025), we parameterize =(−⊤)/2−diag⁡(δ+softplus⁡()) =( L- L )/2-diag(δ+softplus( d)), with δ>0δ>0, making the homogeneous H block dissipative. Through the coupling t B H_t, cells with similar visible states but different hidden initial conditions can follow distinct latent evolution. Each component remains Gaussian, but their mixture can develop multiple modes when their separation exceeds their spread. With nondegenerate component covariances, the mixture retains connected support; branching here refers to the separation of latent density modes. To connect this augmented model with the Mori–Zwanzig formulation in Section 2, we eliminate t H_t from Equation (62), obtaining a non-Markovian generalized Langevin equation for t Z_t: ˙t Z_t =t+z+∫0te⁡(t−s)sds+r(t), = A Z_t+ b_z+ _0^t Be (t-s) C Z_s\,ds+ η_r(t), (63) r(t) η_r(t) =et0+∫0te⁡(t−s)h(r)ds+∫0te⁡(t−s)hds. = Be t H_0+ _0^t Be (t-s) b_h(r)\,ds+ B _0^te (t-s) _h\,d W_s. The kernel ⁡(u)=eu K(u)= Be u C describes dependence on the past history, while r(t) η_r(t) collects the effects of the hidden initial state, bias, and noise. This decomposition into memory and forcing is motivated by Mori–Zwanzig reduction (Lin et al., 2021). Example B.1 (Emergence of branching in the memory-augmented model). Let R∈−1,1R∈\-1,1\ have equal probabilities, Z0|R=r∼⁡(0,1)Z_0 R=r (0,1), and H0∣Z0=z,R=r∼(−z+ra,σH2)H_0 Z_0=z,R=r (-z+ra, _H^2), with 0<σH<a<1+σH20< _H<a< 1+ _H^2. Both initial marginals are unimodal: Z0Z_0 is standard Gaussian, and H0H_0 is an equal mixture of ⁡(±a,1+σH2)N(± a,1+ _H^2). Consider the augmented dynamics dZt=(−Zt+Ht)dtdZ_t=(-Z_t+H_t)\,dt and dHt=(Zt−Ht)dtdH_t=(Z_t-H_t)\,dt in Equation (62). Here A=Λ=−1A= =-1 and B=C=1B=C=1, so eliminating HtH_t gives the memory kernel K⁡(u)=e−uK(u)=e^-u. Writing κt=(1−e−2t)/2 _t=(1-e^-2t)/2, the solution is Zt=e−2tZ0+κt(Z0+H0)Z_t=e^-2tZ_0+ _t(Z_0+H_0). Conditional on R=rR=r, the initialization makes Z0+H0Z_0+H_0 independent of Z0Z_0, with law ⁡(ra,σH2)N(ra, _H^2). Thus the visible density is pt(z)=12(z,−aκt,vt)+12(z,aκt,vt),vt=e−4t+σH2κt2.p_t(z)= 12N(z;-a _t,v_t)+ 12N(z;a _t,v_t), v_t=e^-4t+ _H^2 _t^2. (64) The components coincide at t=0t=0. When a2κt2>vta^2 _t^2>v_t, pt′(0)=(a2κt2−vt)vt−2(0,aκt,vt)>0p_t (0)=(a^2 _t^2-v_t)v_t^-2N(0;a _t,v_t)>0; symmetry and decay at infinity then imply a maximum on each side of zero. This condition holds for t>12log⁡(1+2/a2−σH2)t> 12 (1+2/ a^2- _H^2). As t→∞t→∞, the component means approach ±a/2± a/2 and their variance approaches σH2/4 _H^2/4, so the two modes persist. To isolate the role of memory feedback, set C=0C=0 while keeping the other coefficients and initialization fixed. The component means then become ±ate−t± ate^-t, with common variance e−2t[(1−t)2+σH2t2]e^-2t[(1-t)^2+ _H^2t^2]. Since a2<1+σH2a^2<1+ _H^2, this comparison becomes unimodal again for sufficiently large t, as illustrated in Figure 4. The joint initialization is non-Gaussian: R encodes latent fate information, and the memory feedback sustains the resulting visible modes. Figure 4: Memory feedback sustains visible branching. Analytic densities in Example B.1, with a=0.6a=0.6 and σH=0.2 _H=0.2. Both initial marginals are unimodal. The solid curve is the visible mixture with K⁡(u)=e−uK(u)=e^-u; colored dashed curves illustrate its weighted components. The gray dashed curve removes only the feedback (C=0C=0), keeping the same initialization. Both models develop two modes at t=1t=1, but the comparison returns to a single mode by t=3t=3. Axis scales vary across panels. In this construction, R parameterizes hidden heterogeneity. For cellular development, we use R to represent candidate developmental pathways, whose progression is supervised using biological state annotations. Biological state annotation. We use coarse biological states to model developmental paths, encompassing broad cell identities and developmental-stage categories. Initial training labels are obtained by matching cell IDs to published annotations for DR (Calderon et al., 2022) and WOT cell sets for SC (Schiebinger et al., 2019). For ZB, we annotate Louvain clusters at the latest training time using marker expression, following the terminal-cluster annotation approach of Farrell et al. (2018). We also define early developmental states using sampling time. For example, the first two training snapshots are labeled Blastula progenitor in ZB and Blastoderm in DR; the initial SC population is labeled MEF. For ZB and SC, we fit a logistic regression classifier to the initially labeled training cells and predict states for the remaining training cells. Logistic regression is also used for supervised cell annotation in CellTypist (Domínguez Conde et al., 2022). We calibrate the predicted probabilities and use only assignments that pass confidence, support, and consistency filters as additional supervision. For DR, state supervision uses only labels obtained from the published annotations or assigned using sampling time. Classifier fitting and label supervision use only training cells. Candidate developmental paths. We use these annotated biological states as nodes to define a candidate directed acyclic graph (DAG). We specify its core edges using published developmental relationships (Farrell et al., 2018; Calderon et al., 2022; Schiebinger et al., 2019) and retain them as biological priors. Additional edges are selected using optimal transport between adjacent training snapshots, subject to biological compatibility and support criteria. An edge represents an admissible transition between coarse states. Enumerating root-to-terminal paths defines ℛ R, with r=(c1,…,cℓr)r=(c_1,…,c_ _r) specifying an ordered developmental pathway. Each path parameterizes a hidden initial law in Equation (61). The path identity R remains fixed along a trajectory, while the biological state progresses along the pathway. We further introduce a path consistency loss to encourage generated trajectories to follow the developmental order specified by R (Equation (77) in Appendix B.2.2). Figure 5 shows the graph for SC, whose four root-to-terminal paths define the candidate values of R. Figure 5: Candidate developmental paths for SC. Nodes represent seven coarse biological states, and edges represent the developmental transitions described by Schiebinger et al. (2019). The graph defines four root-to-terminal paths. The highlighted sequence MEF → MET → Epithelial → Neural illustrates one value of R. Each path parameterizes a hidden initial distribution and specifies the developmental order used in the path consistency loss (Equation (77)). B.2.2 Variational inference and representation learning The augmented latent dynamics define a time-dependent prior for learning the VAE representation. At a known sampling time t, the observed variable is a single-cell expression vector t X_t; the latent variables are the developmental path R, the visible state t Z_t, and the memory state t H_t. Generative model. We draw R∼Cat⁡()R ( π) and evolve t=(t⊤,t⊤)⊤ Y_t=( Z_t , H_t ) from the path-dependent initialization in Equation (61) under Equation (62). This gives the Gaussian prior pϑ,t(,∣r)p_ ,t( z, h r). Gene expression is generated from t Z_t alone through the same observation model as KoopCell: pϕ,ϑ,t(,,,r) p_ φ, ,t( x, z, h,r) =πrpϑ,t(,∣r)pϕ(∣), = _r\,p_ ,t( z, h r)\,p_ φ( x z), (65) pϕ(∣) p_ φ( x z) =⁡(,Decϕ⁡(),σx2dx). =N ( x;Dec_ φ( z), _x^2 I_d_x ). Here σx2 _x^2 is fixed, as in Appendix B.1.1, and ϑ collects memory and initial-distribution parameters; dependence of the prior on =[,z] =[ A, b_z] is left implicit. Inference model. Given expression x at time t, we approximate the joint posterior pϕ,ϑ,t(,,r∣)p_ φ, ,t( z, h,r x) by qt(,,r∣)=q(∣,t)qt(r∣,)qtH(∣,r,).q_t( z, h,r x)=q_ θ( z x,t)\,q_t(r z, x)\,q_t^H( h z,r, x). (66) The Gaussian encoder q(⋅∣,t)=((,t),(,t))q_ θ(· x,t)=N( m_ θ( x,t), S_ θ( x,t)) infers the visible latent state; qt(r∣,)q_t(r z, x) assigns path probabilities, and qtH(∣,r,)q_t^H( h z,r, x) infers the memory state conditional on z and r. Using the factorization in Equation (65), Jensen’s inequality gives log⁡pϕ,ϑ,t() p_ φ, ,t( x) ≥qt[log⁡pϕ,ϑ,t(,,,R)qt(,,R∣)] _q_t [ p_ φ, ,t( x, Z, H,R)q_t( Z, H,R x) ] (67) =qlogpϕ(∣)−KL(qt(⋅∣)∥pϑ,t), =E_q_ θ p_ φ( x Z)-KL (q_t(· x)\,\|\,p_ ,t ), where pϑ,t(,,r)=πrpϑ,t(,∣r)p_ ,t( z, h,r)= _rp_ ,t( z, h r). We next specify the memory and path posteriors in Equation (66) to evaluate this bound. Conditional on R=rR=r, the Gaussian initialization and linear augmented dynamics give a jointly Gaussian prior for (t,t)( Z_t, H_t). The conditional distribution pϑ(∣,r,t)p_ ( h z,r,t) is therefore Gaussian and can be computed in closed form. Because the observation model depends only on t Z_t, Equation (65) implies p⁡(t∣t,R,t)=p⁡(t∣t,R)p( H_t Z_t,R, X_t)=p( H_t Z_t,R). We therefore use this Gaussian conditional as the memory posterior, setting qtH(∣,r,)=pϑ(∣,r,t)q_t^H( h z,r, x)=p_ ( h z,r,t). To derive its mean and covariance, write =() K= ( smallmatrix A& B\\ C& smallmatrix ), r=(z⊤,h(r)⊤)⊤ b_r=( b_z , b_h(r) ) , and =(,h⊤)⊤ =( 0, _h ) . The transition is t+Δ|t,R=r∼⁡(Δt+Δr,Δ) Y_t+ Y_t,R=r ( _ Y_t+ J_ b_r, Q_ ), with Δ=eΔ,Δ=∫0Δess,Δ=∫0Δes⊤e⊤ss. _ =e K , J_ = _0 e Ks\,ds, Q_ = _0 e Ks e K s\,ds. (68) The initialization in Equation (61) gives 0,r=(r),0,r=(Dr⊤rr⊤+r). m_0,r= pmatrix 0\\ a_r pmatrix, V_0,r= pmatrix I_D& F_r \\ F_r& F_r F_r + S_r pmatrix. (69) Hence t|R=r∼⁡(t,r,t,r) Y_t R=r ( m_t,r, V_t,r), where t,r=t0,r+tr m_t,r= _t m_0,r+ J_t b_r and t,r=t0,rt⊤+t V_t,r= _t V_0,r _t + Q_t. Gaussian conditioning gives qtH=⁡(¯t,r(),t,rh|z)q_t^H=N( h_t,r( z), V^h z_t,r), with ¯t,r()=t,rh+t,rhz(t,rzz)−1(−t,rz),t,rh|z=t,rhh−t,rhz(t,rzz)−1t,rzh, h_t,r( z)= m^h_t,r+ V^hz_t,r( V^z_t,r)^-1( z- m^z_t,r), V^h z_t,r= V^h_t,r- V^hz_t,r( V^z_t,r)^-1 V^zh_t,r, (70) Having specified the memory conditional qtHq_t^H, we complete the inference model by defining the path probabilities qt(r∣,)q_t(r z, x). The same joint Gaussian prior gives the visible marginal pϑ,t(∣r)=⁡(,t,rz,t,rzz)p_ ,t( z r)=N( z; m^z_t,r, V^z_t,r). We combine this density with the path prior πr _r and an annotation weight for biological states, χr() _r( x), then normalize over candidate paths: qt(r∣,)=πrpϑ,t(∣r)χr()∑r′∈ℛπr′pϑ,t(∣r′)χr′().q_t(r z, x)= _rp_ ,t( z r) _r( x) _r ∈ R _r p_ ,t( z r ) _r ( x). (71) Here χr() _r( x) sums the annotation probabilities over biological states on path r: an accepted state label restricts support to paths containing that state, soft annotations give soft weights, and an unlabeled cell uses χr()=1 _r( x)=1. At least one compatible path is required for an annotated cell. ELBO and training objective. With both posteriors specified, we return to the bound in Equation (67). Write q,t(,r∣)=q(∣,t)qt(r∣,)q_ θ,t( z,r x)=q_ θ( z x,t)q_t(r z, x) and pϑ,t(,r)=πrpϑ,t(∣r)p_ ,t( z,r)= _rp_ ,t( z r). Since qtH(∣,r,)=pϑ(∣,r,t)q_t^H( h z,r, x)=p_ ( h z,r,t), the joint-density ratio in the bound becomes pϕ,ϑ,t(,,,r)qt(,,r∣) p_ φ, ,t( x, z, h,r)q_t( z, h,r x) =pϕ(∣)pϑ,t(,r)pϑ(∣,r,t)q,t(,r∣)pϑ(∣,r,t) = p_ φ( x z)\,p_ ,t( z,r)\,p_ ( h z,r,t)q_ θ,t( z,r x)\,p_ ( h z,r,t) (72) =pϕ(∣)pϑ,t(,r)q,t(,r∣). = p_ φ( x z)\,p_ ,t( z,r)q_ θ,t( z,r x). The common memory conditional cancels, so the ratio no longer depends on h. Taking its logarithm and expectation therefore integrates out H and gives log⁡pϕ,ϑ,t()≥ℰM(,t) p_ φ, ,t( x) _M( x,t) :=qlog⁡pϕ(∣)+q,tlog⁡pϑ,t(,R)q,t(,R∣) :=E_q_ θ p_ φ( x Z)+E_q_ θ,t p_ ,t( Z,R)q_ θ,t( Z,R x) (73) =qlogpϕ(∣)−KL(q,t(⋅∣)∥pϑ,t). =E_q_ θ p_ φ( x Z)-KL (q_ θ,t(· x)\,\|\,p_ ,t ). Expanding the remaining KL over the discrete paths yields KL(q,t∥pϑ,t)=∼q∑r∈ℛqt(r∣,)logq(∣,t)qt(r∣,)πr(,t,rz,t,rzz).KL(q_ θ,t\,\|\,p_ ,t)=E_ Z q_ θ _r∈ Rq_t(r Z, x) q_ θ( Z x,t)q_t(r Z, x) _rN( Z; m^z_t,r, V^z_t,r). (74) We enumerate the paths and estimate the expectation with reparameterized encoder samples. Thus the conditional hidden posterior is analytic, while the mixture KL is evaluated by sampling over Z. To encourage generated trajectories to follow the candidate developmental pathway represented by R, we introduce a path consistency loss based on the probabilities over biological states predicted by a latent state classifier c(c∣)c_ ω(c z). We train the classifier using the annotations above by minimizing the cross-entropy loss ℒstate=−t,,c∼q(⋅∣,t)logc(c∣)L_state=-E_t, x,cE_ Z q_ θ(· x,t) c_ ω(c Z), where the outer expectation averages over training cells with supervised biological state labels c, and the inner expectation is over their encoder samples. We combine ℒstateL_state with the normalized reconstruction loss and KL regularization from Appendix B.1.1 to form the snapshot objective ℒsnap=ℒrec+βt,KL(q,t(⋅∣)∥pϑ,t)+λcℒstate.L_snap=L_rec+ _t, xKL (q_ θ,t(· x)\,\|\,p_ ,t )+ _cL_state. (75) As in KoopCell, we also match predicted and observed populations over all forward training pairs s<ts<t. In KoopCell-M this requires inferring memory at the source time. The inferred augmented source population is μ~s(,,r)=∫q(∣,s)qs(r∣,)pϑ(∣s=,R=r)dρs(). μ_s(d z,d h,r)= q_ θ(d z x,s)\,q_s(r z, x)\,p_ (d h Z_s= z,R=r)\,d _s( x). (76) We propagate μ~s μ_s to time t using Equation (68), keeping r fixed along each trajectory. We denote the predicted t Z_t-marginal by μ^s→t μ_s→ t and map this latent population to gene-expression space through the decoder: ρ^s→t=(Decϕ)#μ^s→t ρ_s→ t=(Dec_ φ)_\# μ_s→ t. We compare these predictions with the corresponding populations at time t using ℒpopL_pop defined in Equation (52). Path consistency loss. For a candidate developmental pathway r=(c1,…,cℓr)r=(c_1,…,c_ _r), we sum over admissible timings of transitions between successive biological states. On a grid 0=u1<⋯<uM=T0=u_1<·s<u_M=T, let ℐrI_r contain index sequences with i1=1i_1=1, im−im−1∈0,1i_m-i_m-1∈\0,1\, and 1≤im≤ℓr1≤ i_m≤ _r. The classifier assigns the path-consistency probability r(u1:uM)=∑(i1,…,iM)∈ℐr∏m=1Mc(cim∣um),P_r( Z_u_1:u_M)= _(i_1,…,i_M) _r _m=1^Mc_ ω(c_i_m Z_u_m), (77) and we minimize ℒpath=−∑rq¯T(r)logr(u1:uM)L_path=- _r q_T(r)E _r( Z_u_1:u_M). Here trajectories start from uniformly sampled initial cells conditional on each path, and q¯T(r) q_T(r) is the detached mean path probability at the latest training time. We evaluate Equation (77) by forward dynamic programming, adapting the alignment marginalization in connectionist temporal classification (CTC; Graves et al., 2006) to the constraints in ℐrI_r. The recursion uses no blank symbol and sums over all final path positions, allowing trajectories to end before the terminal state. The path loss updates both the classifier and the parameters governing the rollout. In Algorithm 2, (r)u1:uM Z^(r)_u_1:u_M denotes a batch of simulated visible trajectories conditional on R=rR=r. Algorithm 2 KoopCell-M: memory augmented latent dynamics for modeling branching 1: Training snapshots =(tk,ρtk)k=0ND=\(t_k, _t_k)\_k=0^N with biological state annotations; candidate paths ℛ R; fixed Fourier tests ψjj=1m\ _j\_j=1^m; batch size nbn_b; λ,ϵM>0λ, _M>0; α∈[0,1)α∈[0,1); loss weights 2: Encoder q_ θ, decoder DecϕDec_ φ, state classifier c_ ω, and dynamics parameters =[,z],ϑ =[ A, b_z], 3: Stage 1: VAE and state classifier pretraining 4: for each pretraining iteration do 5: ℬk←Sample(ρtk,nb),k=0,…,NB_k ( _t_k,n_b), k=0,…,N ⊳ Independent sampling at each time 6: ℒ←PretrainLoss⁡(ℬ,,ϕ)+λcStateLoss⁡(ℬ,,)L (B; θ, φ)+ _cStateLoss(B; θ, ω) ⊳ VAE loss: Equation (53) 7: (,ϕ,)←AdamW⁡((,ϕ,),∇,ϕ,ℒ)( θ, φ, ω) (( θ, φ, ω), _ θ, φ, ωL) ⊳ Pretrain the VAE and state classifier 8: end for 9: Stage 2: Alternating dynamics estimation and representation learning 10: ← ← 0; initialize ,,,h B, C, , _h and πr,r,r,r,h(r)r∈ℛ\ _r, F_r, a_r, S_r, b_h(r)\_r∈ R 11: ←(,ϕ,ϑ,) η←( θ, φ, , ω) ⊳ Load pretrained encoder, decoder, and classifier from Stage 1 12: for each training iteration do 13: (a) Closed-form dynamics estimation (η η fixed) 14: all←sg⁡(Encode⁡(;)) z^all (Encode(D; θ)) ⊳ Reparameterized samples from all training cells 15: (all,all)←Infer⁡(,all,sg⁡(,ϑ))( q^all,C^all) (D, z^all;sg( , )) 16: ⊳ q: path probabilities (Equation (71)); C: Gaussian moments and conditional coefficients (Equation (70)) 17: (,M,)←MemoryStats⁡(all,all,all,sg⁡())( G, y^M, M) ( z^all, q^all,C^all,sg( B)) 18: ⊳ Compute ,M, G, y^M, M by sample averaging and trapezoidal quadrature; Equation (80), Equation (58) 19: ←(+ϵM)−1 W←( M+ _M I)^-1; ^←PCG⁡(⊤+λ,⊤M) a ( G W G+λ I, G W y^M) 20: ^←reshape⁡(^) ( a); ←sg⁡(α+(1−α)^) (α +(1-α) ) ⊳ EMA update 21: (b) Representation and memory update ( fixed) 22: ℬk←Sample(ρtk,nb),k=0,…,NB_k ( _t_k,n_b), k=0,…,N; ϵk,i∼⁡(,D) ε_k,i ( 0, I_D) 23: k,i←(k,i,tk)+(k,i,tk)1/2ϵk,i z_k,i← m_ θ( x_k,i,t_k)+ S_ θ( x_k,i,t_k)^1/2 ε_k,i ⊳ Reparameterized posterior samples 24: (Δ,Δ,Δ)←OUTransition⁡(,ϑ,Δ),Δ∈∪t−s:(s,t)∈( _ , J_ , Q_ ) ( , , ), ∪\t-s:(s,t) \ 25: (,)←Infer⁡(ℬ,,,ϑ)( q,C) (B, z; , ) ⊳ Path probabilities and Gaussian conditionals 26: ℒsnap←SnapshotLoss⁡(ℬ,,,,)L_snap (B, z, q,C; η) ⊳ Reconstruction, KL, state supervision; Equation (75) 27: for each (s,t)∈(s,t) in parallel do 28: s,i,r∼⁡(¯s,r(s,i),s,rh|z) h_s,i,r ( h_s,r( z_s,i), V^h z_s,r) for all i,ri,r ⊳ Infer memory at the source time s 29: ^s→t,i,r←t−s(s,i⊤,s,i,r⊤)⊤+t−sr+ϵi,r y_s→ t,i,r← _t-s( z_s,i , h_s,i,r ) + J_t-s b_r+ ε_i,r, ϵi,r∼⁡(,t−s) ε_i,r ( 0, Q_t-s) 30: μ^s→t←|ℬs|−1∑i,rqs,i,rδ[^s→t,i,r] μ_s→ t←|B_s|^-1 _i,rq_s,i,r _[ y_s→ t,i,r]_ z; ρ^s→t←(Decϕ)#μ^s→t ρ_s→ t←(Dec_ φ)_\# μ_s→ t 31: end for 32: ℒpop←PopulationLoss⁡(ℬ,,μ^s→t,ρ^s→t)L_pop (B, z,\ μ_s→ t, ρ_s→ t\) ⊳ Population matching; Equation (52) 33: u1:uM(r)←PathRollout(ℬ0,r;,,ϑ),r∈ℛ Z^(r)_u_1:u_M (B_0,r; θ, , ), r∈ R 34: ¯T←sg⁡(|ℬN|−1∑iN,i,⋅) q_T (|B_N|^-1 _i q_N,i,·); ℒpath←PathLoss(u1:uM(r),¯T;)L_path (\ Z^(r)_u_1:u_M\, q_T; ω) 35: (ℬ,ℬM,ℬ)←MemoryStats⁡(,sg⁡(),sg⁡(),sg⁡())( G_B, y^M_B, M_B) ( z,sg( q),sg(C),sg( B)) ⊳ Equation (80) 36: ℬ←(ℬ+ϵM)−1 W_B←( M_B+ _M I)^-1 37: ℒweak←‖ℬM−ℬvec⁡()‖ℬ2L_weak←\| y^M_B- G_Bvec( )\|_ W_B^2 ⊳ Retain encoder gradients through z 38: ℒM←ℒsnap+ℒpop+λRℒpath+λweakℒweakL_M _snap+L_pop+ _RL_path+ _weakL_weak ⊳ Combined objective; Equation (82) 39: ←AdamW⁡(,∇ℒM) η ( η, _ ηL_M) ⊳ Update VAE, memory dynamics, and classifier 40: end for 41: return ,Decϕ,c,,ϑ q_ θ,Dec_ φ,c_ ω, , B.2.3 Closed-form dynamics estimation and alternating training We now derive the closed-form update of =[,z] =[ A, b_z] with the representation and memory statistics fixed. Let μ~t μ_t be the inferred augmented population in Equation (76), whose Z-marginal is μt _t, and write ~=(⊤,1)⊤ z=( z ,1) and ~t=(t⊤,1)⊤ Z_t=( Z_t ,1) . For a smooth test φ⁡(t,) (t, z), substituting the visible drift ~+ z+ B h into the definition of the weak population residual in Equation (14) (Appendix A.1) gives ℛDM(,φ) _D^M( ; ) =μT[φ⁡(T)]−μ0[φ⁡(0)] = _T[ (T)]- _0[ (0)] (78) −∫0T∑r∈ℛ∫[∂tφ(t,)+∇φ(t,)⊤(~+)]μ~t(d,d,r)dt - _0^T _r∈ R [ _t (t, z)+ _ z (t, z) ( z+ B h) ]\, μ_t(d z,d h,r)\,dt =ℛD(,φ)−∫0T∑r∈ℛ∫∇φ(t,)⊤μ~t(d,d,r)dt, =R_D( ; )- _0^T _r∈ R _ z (t, z) B h\, μ_t(d z,d h,r)\,dt, where ℛD(,φ)R_D( ; ) is the population residual used in Appendix B.1.2. Using Equation (76) and the conditional mean ∫pϑ(∣,r,t)=¯t,r() h\,p_ (d h z,r,t)= h_t,r( z) from Equation (70) gives ∑r∈ℛ∫∇φ(t,)⊤μ~t(d,d,r) _r∈ R _ z (t, z) B h\, μ_t(d z,d h,r) (79) =∫⁡∇φ(t,)⊤[∑r∈ℛqt(r∣,)¯t,r()]q(d∣,t)ρt(d) = \!\! _ z (t, z) B [ _r∈ Rq_t(r z, x) h_t,r( z) ]q_ θ(d z x,t)\, _t(d x) =⁡[∇φ(t,t)⊤¯t(t,t)], =E [ _ z (t, Z_t) B h_t( Z_t, X_t) ], where ¯t(,)=∑rqt(r∣,)¯t,r() h_t( z, x)= _rq_t(r z, x) h_t,r( z) and the expectation is over t∼ρt X_t _t, t∼q(⋅∣t,t) Z_t q_ θ(· X_t,t). Applying the spatial test ψj _j on Ik=[tk,tk+1]I_k=[t_k,t_k+1], as in Equation (55), gives ℛDM(,φkj)=ykj−⟨,kj⟩F−ckjR_D^M( ; _kj)=y_kj- , H_kj _F-c_kj, where ykj y_kj =μtk+1[ψj]−μtk[ψj],kj=∫Ik[∇ψj(t)~t⊤]dt, = _t_k+1[ _j]- _t_k[ _j], H_kj= _I_kE[∇ _j( Z_t) Z_t ]dt, (80) ckj c_kj =∫Ik[∇ψj(t)⊤¯t(t,t)]dt,yMkj=ykj−ckj. = _I_kE[∇ _j( Z_t) B h_t( Z_t, X_t)]dt, y^M_kj=y_kj-c_kj. We estimate the expectations from encoded cells and approximate the time integrals using the trapezoidal rule, as in KoopCell (Equation (57)). Stacking the interval–test pairs gives the residual vector M− y^M- G a, with =vec⁡() a=vec( ) and (k,j),:=vec(kj)⊤ G_(k,j),:=vec( H_kj) . Similarly, we solve the regularized regression used in KoopCell, replacing y with M=− y^M= y- c: ^=argmin⁡12‖M−‖2+λ2‖22=(⊤+λ)−1⊤M. a= *arg\,min_ a \ 12\| y^M- G a\|_ W^2+ λ2\| a\|_2^2 \=( G W G+λ I)^-1 G W y^M. (81) We evaluate the path probabilities qt(r∣,)q_t(r z, x) and conditional memory means ¯t,r() h_t,r( z) under the current model and hold them fixed together with the memory coupling matrix B when updating . Finally, we combine this closed-form update of with gradient-based learning of the representation and memory dynamics, following the alternating training scheme adopted in KoopCell (Algorithm 1). We first jointly pretrain the VAE and state classifier by minimizing ℒpre+λcℒstateL_pre+ _cL_state, with ℒpreL_pre defined in Equation (53). In Stage 2, we alternate between the following two updates. First, we form the weak statistics, solve Equation (81), and update ←α+(1−α) ←α +(1-α) . Second, with fixed, we update =(,ϕ,ϑ,) η=( θ, φ, , ω) by minimizing ℒM=ℒsnap+ℒpop+λRℒpath+λweakℒweak,ℒweak=‖M−‖2.L_M=L_snap+L_pop+ _RL_path+ _weakL_weak, _weak=\| y^M- G a\|_ W^2. (82) We summarize the training procedure in Algorithm 2. The routine StateLossStateLoss evaluates the supervised classification loss ℒstateL_state in Equation (75). In InferInfer, we compute path probabilities q using Equation (71), and Gaussian moments and conditional coefficients C using Equation (68) and Equation (70). In MemoryStatsMemoryStats, we form (,M,)( G, y^M, M) from Equation (80) and Equation (58) using empirical averages and the trapezoidal rule. For the minibatch weak loss, we hold (,,)( q,C, B) fixed while retaining gradients through z. Rollouts and the path consistency loss follow Equation (68) and Equation (77), respectively. The notation sgsg and PCGPCG follows Algorithm 1. Appendix C Spectral Recovery Experiments This section gives the data-generation and evaluation protocols for the two deterministic toy models in Section 4.1. Let Φt _t denote the flow of the corresponding dynamical system and let ρ0 _0 be its initial population distribution. For each toy model, the dataset consists of empirical population snapshots tkk=0K\X_t_k\_k=0^K, where each tk=tk(i)i=1NkX_t_k=\ x_t_k^(i)\_i=1^N_k contains independent samples from the time marginal ρtk=(Φtk)#ρ0 _t_k=( _t_k)_\# _0. To generate the empirical snapshot at time tkt_k, we independently draw a population of initial states from ρ0 _0, solve the governing ODE from time 00 to tkt_k, and retain only the terminal states. Repeating this construction independently at every observation time removes cross-time cell correspondence and yields the unpaired population snapshots used for training and evaluation. C.1 Toy 1: finite-dimensional nonlinear flow Dynamics and snapshot data generation. We consider the two-dimensional nonlinear flow x˙1=μx1,x˙2=λ⁡(x2−x12), x_1=μ x_1, x_2=λ(x_2-x_1^2), (83) with (μ,λ)=(−0.2,−1)(μ,λ)=(-0.2,-1). The origin (0,0)(0,0) is globally attracting, and the invariant parabola x2=λ−2μx12=53x12x_2= λ-2μx_1^2= 53x_1^2 is an attracting slow manifold. Indeed, the observables ψ1()=x1,ψ2()=x2−53x12,ψ3()=x12 _1( x)=x_1, _2( x)=x_2- 53x_1^2, _3( x)=x_1^2 (84) form a finite-dimensional Koopman-invariant subspace and satisfy ψ˙j=λjψj ψ_j= _j _j with (λ1,λ2,λ3)=(−0.2,−1,−0.4)( _1, _2, _3)=(-0.2,-1,-0.4). The transverse coordinate ψ2 _2 therefore decays faster than the coordinates that parameterize the slow manifold. We define ρ0 _0 by x1(0)∼12(−1,0.22)+12(1,0.22)x_1(0) 12N(-1,0.2^2)+ 12N(1,0.2^2) and x2(0)=x1(0)2+ηx_2(0)=x_1(0)^2+η, where η∼⁡(0,0.82)η (0,0.8^2). For each selected time, 3,000 new initial states are sampled from this law and advanced with the analytic solution of Equation (83). The nine training snapshots are observed at t∈0,0.1,0.2,0.3,0.4,0.55,0.7,0.9,1.2t∈\0,0.1,0.2,0.3,0.4,0.55,0.7,0.9,1.2\; independent target snapshots are generated every 0.20.2 time units from t=1.4t=1.4 to t=4.0t=4.0 for extrapolation. Figure 6 illustrates the generated snapshots, where we observe a rapid contraction toward the slow manifold. Figure 6: Toy 1: generated population snapshots. Each panel is an independently generated sample from the corresponding time marginal. The dashed curve denotes the attracting slow manifold. Experimental settings. We consider three settings. The learned-observable and high-dimensional settings are briefly described in Section 4.1, whereas the known-observable setting is presented only in this appendix due to space constraints. The first is the known-observable setting. Here we use the analytically specified observable map =(x1,x2,x12)⊤ z=(x_1,x_2,x_1^2) , whose components span a finite-dimensional Koopman-invariant subspace in which the dynamics are exactly linear. We therefore fix z as the encoder and estimate only the linear generator, thereby isolating spectral recovery from representation learning. Second, in the learned-observable setting, the model receives only the two physical coordinates (x1,x2)(x_1,x_2) and jointly learns a three-dimensional observable and its corresponding linear propagator. Third, in the high-dimensional setting, we generate 64-dimensional observations by applying a fixed injective map: =softplus⁡(+)+ϵY∈ℝ64,ϵY∼⁡(,0.032). Y=softplus( W x+ c)+ ε_Y ^64, ε_Y ( 0,0.03^2 I). (85) Here =(x1,x2)⊤ x=(x_1,x_2) is the physical state and ∈ℝ64×2 W ^64× 2 has full column rank. Only Y is provided to either model; x, W, and c remain unknown. This setting evaluates representation learning and population prediction in the observed 64D space, and we do not report an inferred spectrum for it. The learned-observable models use 2,048 multiscale random Fourier test functions and a three-dimensional latent space. KoopCell alternates the weak generator solve with representation updates for 500 epochs. scNODE receives the same snapshots and latent dimension, replacing the linear latent generator with an autonomous neural ODE. Table 4 reports the recovered spectrum only for the two settings in which the physical observable structure is available for an unambiguous comparison. The linear observable alignment used to compare learned and analytic eigenfunctions is applied only after training and does not enter either the objective or the rollout. Table 4: Toy 1 spectral recovery with known and learned observables. Parentheses contain absolute eigenvalue errors. λ⋆ _ Known observable Learned observable −0.2-0.2 −0.1983(0.0017)-0.1983\;(0.0017) −0.1952(0.0048)-0.1952\;(0.0048) −1.0-1.0 −0.9949(0.0051)-0.9949\;(0.0051) −0.9774(0.0226)-0.9774\;(0.0226) −0.4-0.4 −0.3893(0.0107)-0.3893\;(0.0107) −0.4227(0.0227)-0.4227\;(0.0227) Rollout and evaluation. Each rollout starts from the observed population at t=0t=0. KoopCell encodes this population, advances every latent state with the learned latent linear dynamics, and decodes at the requested time; scNODE instead integrates its learned latent neural ODE. At each horizon, we compute the sliced Wasserstein distance (SWD) between the predicted population and the independently generated target snapshot. We use physical-space SWD for the learned-observable setting and native observation-space SWD for the 64D setting, and average the metric over all extrapolation times. Figure 7 compares the population rollouts of KoopCell and scNODE. Predicted distributions from KoopCell agree more closely with the ground-truth distributions. Predicted snapshots with KoopCell Predicted snapshots with scNODE Figure 7: Predicted snapshots with KoopCell/scNODE. Both rollouts start from the observed Toy 1 population at t=0t=0. Blue circles denote independently generated target snapshots, and orange triangles denote predictions. Times after t=1.2t=1.2 are extrapolations. C.2 Toy 2: damped Duffing flow Dynamics and snapshot data generation. We consider the damped, unforced Duffing system studied by Otto and Rowley (2019), q˙=p,p˙=−0.5p+q−q3. q=p, p=-0.5p+q-q^3. (86) Its energy H⁡(q,p)=12p2−12q2+14q4H(q,p)= 12p^2- 12q^2+ 14q^4 satisfies dtH(q⁡(t),p⁡(t))=−0.5p(t)2≤0 ddtH(q(t),p(t))=-0.5p(t)^2≤ 0, so the flow is dissipative. The origin is a saddle point, while (q,p)=(±1,0)(q,p)=(± 1,0) are stable spiral equilibria. The linearization at either stable equilibrium has eigenvalues (−1±31i)/4(-1± 31i)/4. The corresponding global Koopman eigenfunctions do not have closed-form expressions. As described by Otto and Rowley (2019), the magnitude and phase of the complex eigenfunction provide action–angle-like coordinates over each basin of attraction: the magnitude vanishes at the stable equilibrium, while the phase tracks the damped rotation. A further eigenfunction with eigenvalue zero is constant along trajectories and takes distinct values in the two basins of attraction. For the initial population, we assign equal mass to two connected annular sectors, one in each basin. For a sign s∈−1,+1s∈\-1,+1\, we sample θ∼⁡[−2.20,1.10]θ [-2.20,1.10] and r∼⁡[0.16,0.48]r [0.16,0.48], set (q,p)=(s⁡(1+rcos⁡θ),−1.35srsin⁡θ)(q,p)=(s(1+r θ),-1.35sr θ), and retain samples with H⁡(q,p)<−0.025H(q,p)<-0.025. This construction covers a range of intrabasin amplitudes and phases while keeping the population away from the energy barrier. Each time marginal contains 3,000 independently resampled initial states propagated with a high-accuracy numerical solver. Training uses the ten marginals at t=0,0.5,…,4.5t=0,0.5,…,4.5; evaluation uses independently generated extrapolation targets at t∈5,6,8,10,12t∈\5,6,8,10,12\. Figure 8: Toy 2 Duffing System: generated population snapshots. The first row shows training distribution marginals and the second row shows extrapolation ground truth. Each panel is generated independently from the same initial population law. Experimental settings. We use two settings. In the learned-observable setting, KoopCell receives the physical snapshots (q,p)(q,p) and jointly learns a three-dimensional observable, its linear generator, and a nonlinear decoder. Both the encoder and decoder are two-layer MLPs with hidden widths 6464–6464, and training uses 2,048 multiscale random Fourier test functions for 500 epochs. Because the relevant eigenfunctions are not available analytically, the objective of this experiment is to determine whether snapshot data recover their characteristic behavior: a complex mode encoding decaying intrabasin amplitude and phase, and a near-zero mode that remains constant within each basin and distinguishes the two long-term outcomes. Figure 1 visualizes these learned modes, and Table 5 compares their eigenvalues with the values implied by the Duffing dynamics. In the high-dimensional setting, the same physical marginals are observed through a fixed injective map =softplus⁡([q,p]⊤+)+ϵY∈ℝ64,ϵY∼⁡(,0.022). Y=softplus( W[q,p] + c)+ ε_Y ^64, ε_Y ( 0,0.02^2 I). (87) Neither method receives (q,p)(q,p) or the observation map during training. Both KoopCell and scNODE use a three-dimensional latent state and are evaluated by their population predictions in the native 64D observation space. Table 5: Toy 2 spectral recovery in the learned-observable setting. Mode Reference λ Estimated λ λ Absolute error Oscillatory pair −0.2500±1.3919i-0.2500± 1.3919i −0.2979±1.4811i-0.2979± 1.4811i 0.10120.1012 Basin mode 00 0.01050.0105 0.01050.0105 Rollout and evaluation. We follow the same protocol as in the Toy 1 experiment. Figure 9 compares the population rollouts produced by KoopCell and scNODE. KoopCell more faithfully reproduces the target population over the full prediction horizon through t=12t=12. Figure 9: Predicted Duffing snapshots with KoopCell/scNODE. Rows show independently generated ground truth, KoopCell predictions, and scNODE predictions; columns progress from training times to extrapolation through t=12t=12. C.3 Rollout-error accumulation Figure 10 compares the rollout errors of KoopCell and scNODE over time for both toy models. The vertical dotted line marks the final training snapshot, and the shaded region denotes extrapolation. On Toy 1, scNODE’s error increases steadily beyond the training window, whereas KoopCell’s error remains low throughout extrapolation. On Toy 2, KoopCell also maintains lower rollout errors than scNODE throughout the evaluated horizon. Figure 10: Accumulation of population rollout error. Left: Toy 1. Right: Toy 2. Each curve reports the per-time SWD for a rollout initialized from the observed population at t=0t=0. Appendix D Branching Dynamics Experiments This section gives the data-generation, training, and evaluation protocols for the branching experiment in Section 4.2. Dynamics and snapshot data generation. We generate a noisy pitchfork process with progression coordinate s and branch coordinate b: ds=0.21dt,db=1.5(2.4sb−b3)dt+0.15dWt,ds=0.21\,dt, =1.5(2.4sb-b^3)\,dt+0.15\,dW_t, (88) where s0∼⁡(−1.05,0.072)s_0 (-1.05,0.07^2) and b0∼⁡(0,0.042)b_0 (0,0.04^2) are independent. As s becomes positive, the central equilibrium of the branch drift loses stability and two stable branches emerge. We simulate 4,000 trajectories with an Euler–Maruyama step of 0.00250.0025, updating s before evaluating the branch drift. To simulate high-dimensional data from this two-dimensional process, we use the nonlinear observation map =2[+0.2tanh⁡(+)−tanh⁡()]+ϵ∈ℝ64 x=2[ Q u+0.2\ ( W u+ c)- ( c)\]+ ε ^64, where =(s/1.2,b/1.8)⊤ u=(s/1.2,b/1.8) , ⊤=2 Q Q= I_2, ‖2=1\| W\|_2=1, cj∼⁡[−0.5,0.5]c_j [-0.5,0.5], and ϵ∼⁡(,0.09264) ε ( 0,0.09^2 I_64) is independent observation noise. Each snapshot is independently shuffled, and training uses only unpaired populations, without trajectory correspondences or branch labels. The nine training times are 0,1.2,2.4,3.6,4.8,5.8,6.8,7.8,9.20,1.2,2.4,3.6,4.8,5.8,6.8,7.8,9.2; interpolation is evaluated at 4.2,8.34.2,8.3 and extrapolation at 9.5,109.5,10. Models and training. At each training time, 3,000 cells are used for fitting, 500 for validation, and 500 are reserved for testing. We use the means and standard deviations estimated from fitting cells to standardize observations, and rescale time as τ=t/9.2τ=t/9.2. Both models use a Gaussian encoder with a 10-dimensional latent state. The encoder and decoder each have two hidden layers of width 128, with SiLU and tanh activation functions, respectively. Decoder weights are spectrally normalized with output scale 4 to control the decoder’s Lipschitz constant in both models and thus limit the amplification of latent differences during decoding. KoopCell uses the linear latent flow ˙τ=τ+z Z_τ= A Z_τ+ b_z with a stable parameterization of A. KoopCell-M uses the augmented dynamics in Equation (62) (Appendix B.2.1), with visible state τ∈ℝ10 Z_τ ^10 and memory state τ∈ℝ24 H_τ ^24. In this experiment, h= _h= 0 and the hidden bias h b_h is shared across components. Following the conditional-mixture construction in Equation (61), we use two hidden components with R∼Cat⁡(π1,π2)R ( _1, _2) and 0∣0=,R=r∼(r(),()) H_0 Z_0= z,R=r ( m_r( z), S( z)). Here the mixture probabilities πr _r are learned, and neural networks parameterize the conditional means r() m_r( z) and diagonal covariance ⁡() S( z). To isolate the effect of the dynamical model from that of the generator-estimation procedure, we use a common gradient-based training protocol for both models. For this controlled comparison, the closed-form generator updates in Algorithms 1 and 2 are replaced by joint optimization of the representation and dynamics, using the same reconstruction and population-matching objectives. We first pretrain a shared VAE for 2,400 Adam steps (learning rate 0.0020.002, 256 cells per time), then initialize both models from this VAE and train each for 8,000 steps, with 128 cells per time and a cosine-decayed learning rate starting at 0.0010.001. Evaluation. After training, we encode the initial observed population, propagate its latent states with the learned dynamics, and decode the predicted populations at the evaluation times. We compare these predictions with the corresponding ground-truth populations using the following metrics: • Sliced Wasserstein distance (SWD). The root mean square of one-dimensional W2W_2 distances along 512 shared random directions in standardized 64D observation space. • Energy distance. The energy distance between predicted and observed populations in the same standardized space, normalized by 64=8 64=8. • Squared maximum mean discrepancy (MMD2). We use an RBF kernel with bandwidth set to the median pairwise distance among fitting cells in standardized observation space. • Branch-coordinate W1W_1. After branching, the two populations occupy positive and negative values of the branch coordinate b. We fit a linear probe b^()=⊤+c b( x)= w x+c to training observations and their known b values, then apply it to the predicted 64D observations. The W1W_1 distance between the estimated and true b distributions measures discrepancies in branch locations, widths, and relative population sizes. • Central mass error. The interval |b|<0.5|b|<0.5 defines a fixed central region around the initial, unbranched population. Using the same linear probe, we compute the absolute difference between predicted and true fractions of cells in this region. This measures whether the model reproduces the amount of population remaining between the two branches. Table 6 reports prediction errors averaged over the corresponding training, interpolation, and extrapolation times. Table 6: Detailed branching prediction errors (mean ± standard deviation across three optimization seeds). Lower is better; the better mean is shown in bold. Times Model SWD Energy ×103× 10^3 MMD2×103^2\!× 10^3 Branch W1W_1 Mass error Training times KoopCell 0.0882±0.00470.0882± 0.0047 2.281±0.5612.281± 0.561 0.831±0.2250.831± 0.225 0.068±0.0030.068± 0.003 0.0264±0.00070.0264± 0.0007 KoopCell-M 0.0798±0.00040.0798± 0.0004 1.273±0.0221.273± 0.022 0.410±0.0110.410± 0.011 0.044±0.0020.044± 0.002 0.0122±0.00170.0122± 0.0017 Interpolation KoopCell 0.0994±0.00830.0994± 0.0083 2.804±0.9602.804± 0.960 1.058±0.3831.058± 0.383 0.105±0.0040.105± 0.004 0.0519±0.00060.0519± 0.0006 KoopCell-M 0.0848±0.00030.0848± 0.0003 1.258±0.0551.258± 0.055 0.437±0.0180.437± 0.018 0.067±0.0030.067± 0.003 0.0201±0.00130.0201± 0.0013 Extrapolation KoopCell 0.2120±0.02120.2120± 0.0212 10.790±2.71310.790± 2.713 4.359±0.9524.359± 0.952 0.344±0.0100.344± 0.010 0.0951±0.00490.0951± 0.0049 KoopCell-M 0.1400±0.00890.1400± 0.0089 2.393±0.5682.393± 0.568 1.012±0.2671.012± 0.267 0.106±0.0060.106± 0.006 0.0016±0.00020.0016± 0.0002 Appendix E Single-Cell Experimental Details E.1 Single-cell datasets and preprocessing We use the processed benchmark released with scNODE (Zhang et al., 2024) on Figshare33 3 https://doi.org/10.6084/m9.figshare.25601610.v1., which contains three high-dimensional scRNA-seq datasets. ZB comprises 38,731 cells from 12 stages of zebrafish embryogenesis, spanning 3.3–12 hours post fertilization (Farrell et al., 2018). DR contains 27,386 cells from 11 overlapping collection windows spanning 0–20 hours of Drosophila embryogenesis (Calderon et al., 2022). SC follows the reprogramming of mouse embryonic fibroblasts toward induced pluripotent stem cells (Schiebinger et al., 2019); we use the reduced benchmark subset of 23,619 cells in 19 temporal groups, formed by pooling samples collected within each integer day from day 0 to day 18. Table 7 shows the number of cells at each time point in each dataset. We consider three prediction tasks of increasing difficulty: easy withholds intermediate snapshots for interpolation; medium withholds the final snapshots for extrapolation; and hard combines intermediate and final held-out snapshots, leaving sparser observations for training. Figure 12 shows the training and held-out snapshots for each dataset and task. For each task, we use the 2,000 highly variable genes (HVGs) selected in the benchmark. For ZB and DR, we normalize each cell’s counts over these genes to a total of 10410^4 and apply a log transformation: xig=log⁡(1+104cig/∑h=12000cih)x_ig= (1+10^4c_ig/ _h=1^2000c_ih), where cigc_ig is the count for gene g in cell i. For SC, we use the normalized expression values provided in the benchmark. We index the observed populations in chronological order as t=1,…,Tt=1,…,T. Figure 11 visualizes all observed cells using the hard-task gene sets. For each dataset, we reduce the expression profiles to 50 principal components and then construct a two-dimensional UMAP embedding, with cells colored by snapshot index. Figure 11: Observed populations in (a) ZB, (b) DR, and (c) SC. Each point represents a cell in a UMAP embedding of gene expression, colored by snapshot index. Table 7: Number of observed cells at each time point. Indices start at t=1t=1; dashes indicate indices beyond the observed horizon. Dataset 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ZB 311 200 1158 1467 5716 1026 4101 6178 5442 7114 1614 4404 – – – – – – – DR 3944 2801 947 3139 2074 2147 3311 1692 1856 1491 3984 – – – – – – – – SC 800 560 1371 1413 1608 1377 1153 1156 2420 886 732 712 747 708 1444 2061 2106 1622 743 Figure 12: Training and held-out snapshots for each prediction task. E.1.1 Population prediction and evaluation For a consistent evaluation across methods, KoopCell, KoopCell-M, and all baselines each generate 2,000 cells at every held-out time point. We compare these predicted populations with all observed cells at the corresponding time point in the processed 2,000-gene space. Every task retains the earliest observed population as the source for prediction. For KoopCell and KoopCell-M, we uniformly sample 2,000 cells from this population, with replacement when fewer than 2,000 cells are available, and draw their initial latent states from the encoder posterior. KoopCell propagates these states with its learned linear dynamics. KoopCell-M additionally samples a path R and memory state 0 H_0 for each initial latent state, then propagates the augmented state while keeping R fixed. We decode the visible latent states at each evaluation time to obtain the predicted gene expression distribution. In the single-cell experiments, we evaluate predictions using the quadratic optimal transport cost 12W22(μ,ν)=infπ∈Π⁡(μ,ν)∫12‖−‖2π(,). 12W_2^2(μ,ν)= _π∈ (μ,ν) 12\| x- y\|^2\,dπ( x, y). We report and summarize the results as mean ± sample standard deviation over three training seeds. E.2 Implementation of Baseline Models We compare KoopCell and KoopCell-M with six representative models for modeling temporal single-cell population dynamics. For scNODE and LGP-OT, the released repositories provide tuned configurations for the same datasets and prediction tasks, so we use those configurations directly. We select the hyperparameters of MIOFlow, PRESCIENT, PI-SDE and VGFM by training-only three-fold cross-validation, following the protocol adopted in the scNODE benchmark. Specifically, at every training time point, we randomly partition the cells into three approximately equal subsets, fit the model on two subsets, and validate on the remaining subset. We rotate the held-out subset over the three folds and select the configuration with the lowest mean validation optimal-transport distance across folds and training time points. The test time points defined by each prediction task remain inaccessible during model selection. After selecting a configuration, we retrain the model using all training cells and evaluate it on the held-out test time points. scNODE. scNODE combines a variational autoencoder with a latent neural ordinary differential equation (ODE), and jointly optimizes the latent representation and continuous dynamics using population-level distribution matching and dynamic regularization (Zhang et al., 2024). We use the task-specific settings released as tunedOurPars: the latent dimension is fixed as 50, while the encoder, decoder, and ODE-network architectures depend on the dataset and prediction task. Because these settings were tuned on the same nine dataset–task combinations considered in our real world scRNA-seq experiments (section 4.3), we do not perform an additional search. LGP-OT. LGP-OT represents population dynamics with a heteroscedastic Gaussian process in a learned latent space, approximates the process with Hilbert-space basis functions, and trains a decoder by matching generated and observed cell populations with a Sinkhorn optimal-transport objective (Balik and Lähdesmäki, 2026). We adopt the released task-specific modelParams configurations, which use a 32-dimensional latent space and select the decoder architecture and number of basis functions for each dataset and task. We therefore do not retune LGP-OT. MIOFlow. MIOFlow first learns a geometry-preserving representation with a geodesic autoencoder and then fits neural-ODE dynamics using optimal-transport and density-regularization losses (Huguet et al., 2022). We select the geodesic-autoencoder dimension from 10,50,100\10,50,100\, the partition-encoder architecture from [50,100],[50,100,100]\[50,100],[50,100,100]\, the ODE architecture from [16],[16,16],[16,32,16]\[16],[16,16],[16,32,16]\, and the density-loss coefficient from 1,25,50,75,100\1,25,50,75,100\ using three-fold cross-validation. PRESCIENT. PRESCIENT learns a neural potential-energy landscape whose gradient defines the drift of a stochastic process, thereby generating cell-population trajectories in a fixed 50-dimensional PCA space (Yeo et al., 2021). We tune the Gaussian-noise standard deviation over [0,1][0,1], the energy-regularization strength over [0,0.1][0,0.1], the gradient-clipping threshold over [0,1][0,1], and the number of hidden layers over 1,2,3\1,2,3\ by three-fold cross-validation. PI-SDE. PI-SDE learns a physics-informed stochastic differential equation in a fixed 50-dimensional PCA space by combining distribution matching with a potential-based action regularizer (Jiang and Wan, 2024). We fix the potential function to a two-layer fully connected network with 400 units per layer and Softplus activations. We tune only the diffusion parameterization among a fixed scalar (const), a learnable constant vector (const_param), and a state-dependent neural network (Mlp); the scalar initialization is fixed to 0.10.1. We also fix the learning rate to 0.0050.005, the action-regularization coefficient to 0.50.5, and the sampled-cell fraction per training time point to 0.10.1, and select the diffusion type by three-fold cross-validation. VGFM. VGFM jointly models cellular velocity and growth through flow matching and unbalanced optimal transport (Wang et al., 2025). We fit PCA using training cells only and explore raw and standardized PCA representations, dimensions in 50,100\50,100\, Tanh and LeakyReLU activation functions, entropy regularization ϵ∈0.01,0.03ε∈\0.01,0.03\, and source-marginal relaxation τ∈5,10τ∈\5,10\. The reported results use standardized PCA with 50 components for the easy and hard tasks and 100 components for the medium task. All reported configurations use Tanh activations, ϵ=0.01ε=0.01, and τ=10τ=10. The velocity and growth networks have five and three hidden layers, respectively, with 512 units per hidden layer. Training comprises 2,000 pretraining epochs followed by 50 trajectory-training epochs, with learning rates 10−310^-3 and 10−410^-4, respectively. KoopCell(-M), scNODE, LGP-OT, MIOFlow, PRESCIENT, PI-SDE and VGFM are all implemented and trained on NVIDIA Tesla V100 GPUs with 32 GB of memory. E.3 Model configurations Table 8: Settings shared across the three datasets and prediction tasks. Stage 1 and Stage 2 denote VAE pretraining and dynamics/representation alternating training stages, respectively. Parameter KoopCell KoopCell-M Visible latent dimension D 50 50 Encoder/decoder MLP activation ReLU ReLU Training epochs (Stage 1 / Stage 2) 500 / 4,000 500 / 4,000 Batch size (Stage 1 / Stage 2) 256 / 128 256 / 128 Variational coefficient β 0.01 0.01 Number of Fourier test functions 2,048 2,048 Sinkhorn batch size 512 256 State-classifier hidden width – 128 Path-loss coefficient λR _R – 0.1 Table 9: Selected MLP architectures and loss weights. Encoder and decoder entries list hidden-layer widths; dhd_h is the memory dimension. KoopCell Task Data Encoder MLP Decoder MLP dhd_h λz _z λx _x λweak _weak Easy ZB 256–256–256 256–256–256 – 1 1.5 1 DR 512–256–256 256–256–512 – 1 5 1 SC 512–256–256 256–256–512 – 1 3 1 Medium ZB 256–256–256 256–256–256 – 1 1.5 1 DR 256–256–256 256–256–256 – 1 5 1 SC 256–256–256 256–256–256 – 1 3 1 Hard ZB 256–256–256 256–256–256 – 1 1.5 1 DR 256–256–256 256–256–256 – 1 5 1 SC 256–256–256 256–256–256 – 1 3 1 KoopCell-M Task Data Encoder MLP Decoder MLP dhd_h λz _z λx _x λweak _weak Easy ZB 512–256–256 256–256–512 16 5 10 10510^5 DR 512–256–256 256–256–512 32 2 10 10410^4 SC 512–256–256 256–256–512 16 10 10 7,500 Medium ZB 512–256–256 256–256–512 32 5 10 10510^5 DR 512–256–256 256–256–512 50 5 10 10410^4 SC 512–256–256 256–256–512 32 5 5 7,500 Hard ZB 256–256–256 256–256–256 16 5 10 10510^5 DR 256–256–256 256–256–256 32 2 10 10410^4 SC 256–256–256 256–256–256 16 1 3 7,500 E.4 Additional Single-Cell Results Table 10: Easy-task Wasserstein-2 distance (↓ ) at held-out time points (mean ± sample standard deviation over three seeds). The best and second-best means are shaded green (bold) and blue, respectively. Method ZB/easy task Interpolation = t=5 = t=7 = t=9 scNODE 435.03±7.52435.03± 7.52 393.17±3.84393.17± 3.84 473.80±4.00473.80± 4.00 LGP-OT 436.27±3.20436.27± 3.20 362.05±2.20362.05± 2.20 432.68±1.19432.68± 1.19 MIOFlow 447.71±5.63447.71± 5.63 415.37±5.93415.37± 5.93 491.57±5.63491.57± 5.63 PRESCIENT 475.52±8.25475.52± 8.25 447.97±6.87447.97± 6.87 507.70±2.65507.70± 2.65 PI-SDE 422.71±1.60422.71± 1.60 369.90±2.79369.90± 2.79 435.65±3.68435.65± 3.68 VGFM 443.05±5.89443.05± 5.89 391.59±3.52391.59± 3.52 490.49±4.32490.49± 4.32 KoopCell 431.53±6.02431.53± 6.02 366.01±1.99366.01± 1.99 436.71±1.02436.71± 1.02 KoopCell-M 428.05±3.79428.05± 3.79 355.84±2.31355.84± 2.31 426.25±1.34426.25± 1.34 Method DR/easy task Interpolation = t=5 = t=7 = t=9 scNODE 348.05±3.93348.05± 3.93 410.98±3.28410.98± 3.28 499.14±5.93499.14± 5.93 LGP-OT 346.08±0.54346.08± 0.54 400.04±0.63400.04± 0.63 471.29±2.28471.29± 2.28 MIOFlow 345.69±0.55345.69± 0.55 404.59±1.80404.59± 1.80 497.52±2.97497.52± 2.97 PRESCIENT 365.36±0.65365.36± 0.65 428.44±0.67428.44± 0.67 500.78±0.58500.78± 0.58 PI-SDE 335.03±1.02335.03± 1.02 392.45±0.72392.45± 0.72 474.66±1.54474.66± 1.54 VGFM 343.37±3.20343.37± 3.20 406.06±2.83406.06± 2.83 500.79±1.07500.79± 1.07 KoopCell 339.74±2.81339.74± 2.81 396.39±6.49396.39± 6.49 464.00±2.58464.00± 2.58 KoopCell-M 340.94±0.34340.94± 0.34 389.83±1.27389.83± 1.27 462.52±1.06462.52± 1.06 Method SC/easy task Interpolation = t=6 = t=11 = t=16 scNODE 57.42±0.9657.42± 0.96 134.96±5.05134.96± 5.05 118.18±3.13118.18± 3.13 LGP-OT 51.16±0.3151.16± 0.31 124.83±0.71124.83± 0.71 100.55±1.35100.55± 1.35 MIOFlow 55.96±0.7355.96± 0.73 137.79±0.78137.79± 0.78 122.78±1.98122.78± 1.98 PRESCIENT 68.12±0.6768.12± 0.67 144.99±0.61144.99± 0.61 122.43±1.48122.43± 1.48 PI-SDE 50.83±0.1850.83± 0.18 128.23±1.44128.23± 1.44 101.86±1.09101.86± 1.09 VGFM 51.94±2.6051.94± 2.60 152.69±3.70152.69± 3.70 160.06±1.12160.06± 1.12 KoopCell 52.40±1.0952.40± 1.09 124.38±1.09124.38± 1.09 96.33±0.5196.33± 0.51 KoopCell-M 50.52±0.7050.52± 0.70 124.65±0.93124.65± 0.93 95.54±0.7095.54± 0.70 Table 11: Medium-task Wasserstein-2 distance (↓ ) at held-out time points (mean ± sample standard deviation over three seeds). The best and second-best means are shaded green (bold) and blue, respectively. Method ZB/medium task Extrapolation = t=11 = t=12 scNODE 611.92±5.22611.92± 5.22 692.19±5.75692.19± 5.75 LGP-OT 570.90±5.15570.90± 5.15 642.06±9.06642.06± 9.06 MIOFlow 630.14±13.26630.14± 13.26 755.77±18.34755.77± 18.34 PRESCIENT 629.81±3.64629.81± 3.64 677.10±3.37677.10± 3.37 PI-SDE 572.99±1.83572.99± 1.83 642.34±4.18642.34± 4.18 VGFM 607.60±4.59607.60± 4.59 665.91±12.69665.91± 12.69 KoopCell 558.09±2.51558.09± 2.51 605.61±3.34605.61± 3.34 KoopCell-M 548.66±2.12548.66± 2.12 603.10±5.57603.10± 5.57 Method DR/medium task Extrapolation = t=9 = t=10 = t=11 scNODE 561.97±16.65561.97± 16.65 573.46±13.14573.46± 13.14 715.56±24.23715.56± 24.23 LGP-OT 514.14±1.59514.14± 1.59 520.71±15.78520.71± 15.78 661.35±4.50661.35± 4.50 MIOFlow 558.93±5.41558.93± 5.41 614.22±21.61614.22± 21.61 803.22±62.71803.22± 62.71 PRESCIENT 551.35±0.75551.35± 0.75 511.82±0.22511.82± 0.22 669.97±1.78669.97± 1.78 PI-SDE 525.17±1.02525.17± 1.02 543.32±2.89543.32± 2.89 708.72±2.27708.72± 2.27 VGFM 554.89±5.43554.89± 5.43 670.18±41.27670.18± 41.27 901.81±72.49901.81± 72.49 KoopCell 505.50±4.18505.50± 4.18 521.45±12.38521.45± 12.38 644.53±10.65644.53± 10.65 KoopCell-M 495.08±1.65495.08± 1.65 496.61±3.11496.61± 3.11 643.43±5.02643.43± 5.02 Method SC/medium task Extrapolation = t=17 = t=18 = t=19 scNODE 130.31±3.95130.31± 3.95 124.75±2.33124.75± 2.33 130.34±2.47130.34± 2.47 LGP-OT 115.15±1.66115.15± 1.66 122.09±3.98122.09± 3.98 126.47±3.44126.47± 3.44 MIOFlow 149.96±14.49149.96± 14.49 154.17±23.86154.17± 23.86 173.96±39.70173.96± 39.70 PRESCIENT 145.16±2.48145.16± 2.48 135.73±2.14135.73± 2.14 136.07±2.06136.07± 2.06 PI-SDE 122.47±5.89122.47± 5.89 115.87±3.50115.87± 3.50 126.16±4.15126.16± 4.15 VGFM 164.07±1.59164.07± 1.59 168.12±1.03168.12± 1.03 174.05±1.34174.05± 1.34 KoopCell 110.59±0.65110.59± 0.65 111.27±5.20111.27± 5.20 115.10±5.88115.10± 5.88 KoopCell-M 114.75±0.39114.75± 0.39 109.09±2.17109.09± 2.17 113.48±2.22113.48± 2.22 E.4.1 Predicted cell populations at held-out time points Figures 13–15 compare the predicted populations from individual training runs. All methods share the saved PCA–UMAP projection for each dataset and task. We display up to 2,000 cells per time point, resampling VGFM particles according to their normalized growth weights. Figure 13: Easy-task predictions for ZB, DR, and SC. Colors indicate held-out time points; gray points show all observed cells. Diamonds and dark dashed frames identify ground truth; stars and orange frames mark KoopCell and KoopCell-M. All methods within each dataset share the same projection and axis ranges. Figure 14: Medium-task predictions for ZB, DR, and SC. Colors indicate held-out time points; gray points show all observed cells. Diamonds and dark dashed frames identify ground truth; stars and orange frames mark KoopCell and KoopCell-M. All methods within each dataset share the same projection and axis ranges. Figure 15: Hard-task predictions for ZB, DR, and SC, with all held-out interpolation and extrapolation time points overlaid in each panel. Colors indicate time points; gray points show all observed cells. Diamonds and dark dashed frames identify ground truth; stars and orange frames mark KoopCell and KoopCell-M. All methods within each dataset share the UMAP embedding used in Figure 11 and the same axis ranges. E.4.2 Visualization of the learned latent representations Figure 16 compares the representations of all 23,619 observed SC cells under the hard-task setting. Each method has its own UMAP projection; both rows use the same coordinates for that method. Reference states are obtained by matching cell identifiers to the published Waddington-OT cell sets and labeling the initial population as MEF. Unmatched or ambiguous cells remain gray; classifier predictions are not used to color these panels. Figure 16: Latent representations of SC cells under the hard-task setting, colored by (a) reference cell state and (b) snapshot time index. All five methods show the same observed cells. Each method uses its own UMAP coordinates, shared between the two rows. The KoopCell and KoopCell-M columns are also shown in Figure 3. E.5 Biological Downstream Analysis We investigate how the population dynamics learned by KoopCell-M relate to lineage-associated genes and developmental expression programs. Using KoopCell-M trained on all ZB snapshots, we focus on Presomitic Mesoderm (PSM) and Hindbrain development, following Zhang et al. (2024). Lineage-associated gene identification. We first estimate the probability that each cell contributes to a terminal cell type at 12 hpf. For a cell observed at time tit_i, we encode i=(i,ti) z_i= m_ θ( x_i,t_i) and infer its path probabilities and conditional memory distribution from the generative model. We then construct a directed graph between successive snapshots by propagating augmented states (,)( Z, H) with the Gaussian transition in Equation (68) and assigning the propagated mass to neighboring states. Distances are scaled by the corresponding model covariance, and the path identity is retained across transitions. Reference annotations of all terminal cell types define the endpoints. Let ℓ,r P_ ,r denote the transition matrix from snapshot ℓ to ℓ+1 +1 along path r, and let ℓ,r(j,L) F_ ,r(j,L) denote the probability of reaching terminal cell type L from augmented state j at snapshot ℓ , conditional on path r. The probability matrices satisfy ℓ,r=ℓ,rℓ+1,r,K,r(j,L)=(cj=L), F_ ,r= P_ ,r F_ +1,r, F_K,r(j,L)=1(c_j=L), (89) where K indexes the terminal snapshot and cjc_j is the reference label of terminal cell j. Averaging over the inferred memory and path distributions gives the cell-level fate probability fiLf_iL. Following the lineage-driver statistic of CellRank (Weiler et al., 2024), we rank genes by their Pearson correlation with the predicted probability of reaching lineage L: rgL=Corri:ti<T(xig,fiL),r_gL=Corr_i:t_i<T(x_ig,f_iL), (90) where xigx_ig is observed log-expression and T denotes the terminal time. To examine whether these fate-associated genes also distinguish the observed terminal populations, we compare log-normalized expression between 540 PSM and 511 Hindbrain cells at 12 hpf. Table 12 reports the mean expression difference Δg _g (PSM minus Hindbrain) for the four highest-ranked genes per lineage, selected by rgLr_gL. Positive and negative values indicate higher expression in PSM and Hindbrain, respectively. Recovery of developmental markers. MSGN1 and TBX16 rank first and fourth for PSM, with correlations of 0.6650.665 and 0.5000.500, respectively (Table 12). Both genes participate in zebrafish mesodermal morphogenesis and differentiation (Manning and Kimelman, 2015). For Hindbrain, SOX19A and SOX3 rank first and second (r=0.568r=0.568 and 0.5400.540), consistent with their expression in the developing neuroectoderm (Rossi et al., 2009). All eight genes show higher terminal expression in their corresponding lineage. APELA shows a modest terminal expression difference (Δg=−0.243 _g=-0.243). In the UMAP projection, the observed expression patterns of these genes align with the annotated PSM and Hindbrain populations (Figure 17). Table 12: Lineage-associated genes and terminal expression differences. The four highest-ranked genes are shown for each lineage. rgLr_gL is the preterminal expression–fate correlation used for ranking; Δg _g is mean observed log-expression in PSM minus Hindbrain at 12 hpf. Lineage Gene Rank rgLr_gL Δg _g PSM MSGN1 1 0.665 4.3884.388 PSM TBX6L 2 0.559 1.4341.434 PSM RBM38 3 0.513 2.1602.160 PSM TBX16 4 0.500 3.8443.844 Hindbrain SOX19A 1 0.568 −3.670-3.670 Hindbrain SOX3 2 0.540 −3.552-3.552 Hindbrain SI:CH211-152C2.3 3 0.472 −3.770-3.770 Hindbrain APELA 4 0.371 −0.243-0.243 Figure 17: Observed expression of lineage-associated genes on a UMAP projection of the latent representations learned by KoopCell-M for ZB cells. Rows show the four highest-ranked genes for (a) PSM and (b) Hindbrain. The left panels highlight the corresponding annotated populations at 12 hpf; the remaining panels show observed log-expression across all cells using a common color scale. All panels share the same UMAP coordinates. Arrows indicate the terminal populations. Dynamic gene-expression programs. We next identify groups of genes with similar predicted expression-rate profiles along developmental trajectories. Starting from cells observed at 8 hpf, we infer their path and memory distributions and propagate the augmented dynamics to 12 hpf. The decoder maps each trajectory to log-normalized gene-expression space. Because noise enters only the memory variables, the predicted expression rate for gene g follows the chain rule: ug(t)=∇Decϕ,g(t)⊤(t+t+z).u_g(t)= _ zDec_ φ,g( Z_t) ( A Z_t+ B H_t+ b_z). (91) We compute this rate using Jacobian–vector products and average it separately for each lineage. We first give each source cell equal total weight, distributed across paths according to their inferred probabilities and equally among trajectory samples along each path. Let αa _a denote the resulting weight of trajectory a, whose path is rar_a. To determine its contribution to lineage L, we compare its endpoint Ta z_T^a at T=12T=12 hpf with annotated terminal cells. We define bL(Ta,ra)b_L( z_T^a,r_a) as the distance-weighted fraction of neighboring cells labeled L, with distances scaled by the terminal covariance for path rar_a. Normalizing the products αabL(Ta,ra) _ab_L( z_T^a,r_a) gives the lineage-specific weights and mean expression rates: waL=αabL(Ta,ra)∑a′αa′bL(Ta′,ra′),u¯gL(t)=∑awaLuga(t).w_aL= _ab_L( z_T^a,r_a) _a _a b_L( z_T^a ,r_a ), u_gL(t)= _aw_aLu_g^a(t). (92) For each lineage, we select the 100 genes with the largest mean expression rates over 8–12 hpf, emphasizing net upregulation over this interval. We standardize each rate profile across time and apply Ward clustering to obtain three modules, ordered by the temporal center of their above-average rates. To identify the biological functions associated with these programs, we test enrichment of Gene Ontology (GO) biological-process terms (Gene Ontology Consortium, 2026) in each lineage’s top-100 genes and in each of its three modules. We map input gene symbols to Zebrafish Information Network (ZFIN) identifiers using official symbols and aliases (Ruzicka et al., 2019). The background comprises the 1,549 uniquely mapped input genes with GO biological-process annotations; each tested gene set is restricted to this same background. A one-sided hypergeometric test measures whether a term annotates more genes in a tested set than expected from its frequency in the background (Boyle et al., 2004). Benjamini–Hochberg correction (Benjamini and Hochberg, 1995) is applied jointly across the tested terms and gene sets in both lineages, with q<0.05q<0.05 as the significance threshold. Developmental functions of the expression programs. The PSM modules are enriched for genes involved in muscle differentiation and development (Figure 18). The first module contains MYF5, MYOD1, PRDM1A, and RBM24A, and is enriched for positive regulation of muscle cell differentiation (q=6.91×10−4q=6.91× 10^-4). This is consistent with the role of MYF5 and MYOD1 in initiating zebrafish myogenesis (Osborn et al., 2020). The second includes ACTA1A, DESMA, and DMRT2A, with enrichment for skeletal muscle organ development (q=1.43×10−2q=1.43× 10^-2). In Hindbrain, the first module includes FGF3, HER6, and NEUROG1, with enrichment for regulation of neuron differentiation (q=8.57×10−3q=8.57× 10^-3). The second contains SOX3, SOX21A, PAX6A, and GBX2, and is enriched for nervous system development (q=2.47×10−5q=2.47× 10^-5) and brain development (q=7.68×10−4q=7.68× 10^-4). These enrichments are consistent with the neural progenitor and differentiation programs described in the developing zebrafish hindbrain (Tambalo et al., 2020). The third module in each lineage is enriched for intermediate filament organization and includes the keratin genes KRT4 and KRT5 (Cokus et al., 2019). Together, the recovered markers and lineage-associated functional enrichments support the biological interpretability of the learned population dynamics. Figure 18: Dynamic gene-expression programs along KoopCell-M trajectories over 8–12 hpf. Each lineage contains the 100 genes with the largest mean predicted rates of change in log-expression, grouped into three Ward modules. Colors show within-gene temporal z-scores of expression rates, with red indicating rates above the gene’s temporal mean and blue indicating rates below it. Boxes report representative GO biological-process terms and adjusted q-values. PSM programs are enriched for muscle development and Hindbrain programs for neural development; both include a module containing keratin genes.

Instant research discovery