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

cs.LGSep 27, 2026

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

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

Organizations: School of Mathematical Sciences, SCMS, SCAM, and CCSB, Fudan University, China · Research Institute of Intelligent Complex Systems, Fudan University, China · State Key Laboratory of Medical Neurobiology and MOE Frontiers Center for Brain Science, Institutes of Brain Science, Fudan University, China · Shanghai Artificial Intelligence Laboratory, China

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.

Figures & tables

Prediction	Data	$t$	scNODE	LGP-OT	MIOFlow	PRESCIENT	PI-SDE	VGFM	KoopCell	KoopCell-M
	ZB	3	$577.71\pm 3.69$	$577.91\pm 0.64$	$594.71\pm 32.03$	$597.49\pm 1.76$	$\mathbf{567.30\pm 1.41}$	$575.86\pm 3.22$	$572.32\pm 2.14$	$588.91\pm 6.04$
	ZB	5	$511.58\pm 5.63$	$504.45\pm 1.73$	$519.32\pm 7.48$	$540.72\pm 1.52$	$492.43\pm 2.19$	$499.66\pm 0.81$	$497.80\pm 3.83$	$\mathbf{488.84\pm 9.27}$
	ZB	7	$440.65\pm 6.47$	$411.20\pm 3.07$	$452.30\pm 0.77$	$476.68\pm 1.16$	$416.80\pm 3.36$	$434.99\pm 1.37$	$412.99\pm 1.50$	$\mathbf{408.79\pm 1.57}$
Interpolation	ZB	9	$514.24\pm 5.97$	$475.65\pm 2.60$	$539.92\pm 11.44$	$527.09\pm 1.25$	$474.41\pm 3.51$	$522.60\pm 2.83$	$474.92\pm 1.55$	$\mathbf{468.53\pm 2.72}$
	ZB	11	$653.56\pm 8.98$	$605.14\pm 1.78$	$654.64\pm 2.17$	$656.41\pm 4.17$	$612.74\pm 3.51$	$635.96\pm 2.73$	$595.82\pm 0.48$	$\mathbf{588.73\pm 0.77}$
Extrapolation	ZB	12	$719.28\pm 12.36$	$665.78\pm 4.77$	$720.67\pm 17.25$	$700.68\pm 4.58$	$679.61\pm 4.11$	$720.54\pm 10.97$	$648.91\pm 2.47$	$\mathbf{639.78\pm 0.72}$

Table 1: Hard-task Wasserstein-2 distance (

\downarrow

) at held-out time points (mean

\pm

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.

Panel of Figure 1 — 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.

Panel of Figure 2 — 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.

Figure 3: Latent representations of SC cells colored by reference state (top; stromal: green, iPSC: red) and snapshot index (bottom).

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Memory feedback sustains visible branching. Analytic densities in Example B.1 , with

a=0.6

and

\sigma_{H}=0.2

. Both initial marginals are unimodal. The solid curve is the visible mixture with

K(u)=e^{-u}

; colored dashed curves illustrate its weighted components. The gray dashed curve removes only the feedback (

C=0

), keeping the same initialization. Both models develop two modes at

t=1

, but the comparison returns to a single mode by

t=3

. Axis scales vary across panels.

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

\to

MET

\to

Epithelial

\to

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 )).

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.

$\lambda_{\star}$	Known observable	Learned observable
$-0.2$	$-0.1983\;(0.0017)$	$-0.1952\;(0.0048)$
$-1.0$	$-0.9949\;(0.0051)$	$-0.9774\;(0.0226)$
$-0.4$	$-0.3893\;(0.0107)$	$-0.4227\;(0.0227)$

Appendix

Table 4: Toy 1 spectral recovery with known and learned observables. Parentheses contain absolute eigenvalue errors.

Panel of Figure 7 — Figure 7: Predicted snapshots with KoopCell/scNODE. Both rollouts start from the observed Toy 1 population at $t=0$ . Blue circles denote independently generated target snapshots, and orange triangles denote predictions. Times after $t=1.2$ are extrapolations.

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.

Mode	Reference $\lambda$	Estimated $\widehat{\lambda}$	Absolute error
Oscillatory pair	$-0.2500\pm 1.3919i$	$-0.2979\pm 1.4811i$	$0.1012$
Basin mode	$0$	$0.0105$	$0.0105$

Appendix

Table 5: Toy 2 spectral recovery in the learned-observable setting.

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=12$ .

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=0

Times	Model	SWD	Energy $\times 10^{3}$	MMD ${}^{2}\!\times 10^{3}$	Branch $W_{1}$	Mass error
Training times	KoopCell	$0.0882\pm 0.0047$	$2.281\pm 0.561$	$0.831\pm 0.225$	$0.068\pm 0.003$	$0.0264\pm 0.0007$
Training times	KoopCell-M	$\mathbf{0.0798\pm 0.0004}$	$\mathbf{1.273\pm 0.022}$	$\mathbf{0.410\pm 0.011}$	$\mathbf{0.044\pm 0.002}$	$\mathbf{0.0122\pm 0.0017}$
Interpolation	KoopCell	$0.0994\pm 0.0083$	$2.804\pm 0.960$	$1.058\pm 0.383$	$0.105\pm 0.004$	$0.0519\pm 0.0006$
Interpolation	KoopCell-M	$\mathbf{0.0848\pm 0.0003}$	$\mathbf{1.258\pm 0.055}$	$\mathbf{0.437\pm 0.018}$	$\mathbf{0.067\pm 0.003}$	$\mathbf{0.0201\pm 0.0013}$
Extrapolation	KoopCell	$0.2120\pm 0.0212$	$10.790\pm 2.713$	$4.359\pm 0.952$	$0.344\pm 0.010$	$0.0951\pm 0.0049$
Extrapolation	KoopCell-M	$\mathbf{0.1400\pm 0.0089}$	$\mathbf{2.393\pm 0.568}$	$\mathbf{1.012\pm 0.267}$	$\mathbf{0.106\pm 0.006}$	$\mathbf{0.0016\pm 0.0002}$

Appendix

Table 6: Detailed branching prediction errors (mean

\pm

standard deviation across three optimization seeds). Lower is better; the better mean is shown in bold.

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.

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

Appendix

Table 7: Number of observed cells at each time point. Indices start at

t=1

; dashes indicate indices beyond the observed horizon.

Figure 12: Training and held-out snapshots for each prediction task.

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 $\beta$	0.01	0.01
Number of Fourier test functions	2,048	2,048

Appendix

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.

KoopCell
Task	Data	Encoder MLP	Decoder MLP	$d_{h}$	$\lambda_{z}$	$\lambda_{x}$	$\lambda_{\mathrm{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

Appendix

Table 9: Selected MLP architectures and loss weights. Encoder and decoder entries list hidden-layer widths;

d_{h}

is the memory dimension.

Method	ZB/easy task
	Interpolation
	$\bm{t=5}$	$\bm{t=7}$	$\bm{t=9}$
scNODE	$435.03\pm 7.52$	$393.17\pm 3.84$	$473.80\pm 4.00$
LGP-OT	$436.27\pm 3.20$	$362.05\pm 2.20$	$432.68\pm 1.19$
MIOFlow	$447.71\pm 5.63$	$415.37\pm 5.93$	$491.57\pm 5.63$
PRESCIENT	$475.52\pm 8.25$	$447.97\pm 6.87$	$507.70\pm 2.65$

Appendix

Table 10: Easy-task Wasserstein-2 distance (

\downarrow

) at held-out time points (mean

\pm

sample standard deviation over three seeds). The best and second-best means are shaded green (bold) and blue, respectively.

Method	ZB/medium task
	Extrapolation
	$\bm{t=11}$	$\bm{t=12}$
scNODE	$611.92\pm 5.22$	$692.19\pm 5.75$
LGP-OT	$570.90\pm 5.15$	$642.06\pm 9.06$
MIOFlow	$630.14\pm 13.26$	$755.77\pm 18.34$
PRESCIENT	$629.81\pm 3.64$	$677.10\pm 3.37$

Appendix

Table 11: Medium-task Wasserstein-2 distance (

\downarrow

) at held-out time points (mean

\pm

sample standard deviation over three seeds). The best and second-best means are shaded green (bold) and blue, respectively.

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.

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 .

Lineage	Gene	Rank	$r_{gL}$	$\Delta_{g}$
PSM	MSGN1	1	0.665	$4.388$
PSM	TBX6L	2	0.559	$1.434$
PSM	RBM38	3	0.513	$2.160$
PSM	TBX16	4	0.500	$3.844$
Hindbrain	SOX19A	1	0.568	$-3.670$
Hindbrain	SOX3	2	0.540	$-3.552$

Appendix

Table 12: Lineage-associated genes and terminal expression differences. The four highest-ranked genes are shown for each lineage.

r_{gL}

is the preterminal expression–fate correlation used for ranking;

\Delta_{g}

is mean observed log-expression in PSM minus Hindbrain at 12 hpf.

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.

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.

Explore similar work

Sep 28, 2026cs.LG

DisKO: Deep Koopman Learning in Distribution Space from Unpaired Snapshots

Many complex systems are observed only through temporally unpaired distribution snapshots, making trajectory-based dynamical learning difficult without additional assumptions. We therefore formulate the problem directly in distribution space, treating the distribution itself as the dynamical state. The challenge is that distribution space is infinite-dimensional, making compact and approximately closed representations difficult to learn from finite snapshots. We introduce DisKO, which extends deep Koopman learning to distribution dynamics by jointly learning predictive distributional observables, a finite-dimensional Koopman representation, and a generative map back to the full distribution. Across seven diverse benchmarks, DisKO achieves state-of-the-art extrapolation performance, with substantially slower error accumulation on long-horizon prediction tasks. DisKO further recovers leading Koopman eigenvalues and eigenfunctions on systems with analytic spectra, revealing meaningful dynamical structure in the learned representation.

He Ma, Xiaochen Liu, Wanfeng Lu +3

School of Mathematical Sciences, Fudan University, China. · Research Institute of Intelligent Complex Systems, Fudan University, China. · Shanghai Center for Mathematical Sciences, Fudan University, China. +2

Koopman Operator Learning Dynamics Distributions Snapshots Autonomous Hard Drive Disassembly

Ask PDF

Oct 18, 2025cs.LG

Simulation-free Structure Learning for Stochastic Population Dynamics

Modeling dynamical systems and unraveling their underlying structural dependencies is central to many domains in the natural sciences. Various physical systems, such as those arising in cell biology, are inherently high-dimensional and stochastic in nature, and admit only partial, noisy state measurements. Our primary motivating setting is single-cell biology, where destructive measurements yield unpaired population snapshots rather than longitudinal trajectories of the same cells. This poses a significant challenge for addressing the problems of modeling the underlying dynamics and inferring the network structure of these systems. Existing methods are typically tailored either for structure learning or modeling dynamics at the population level, but are limited in their ability to address both problems together. In this work, we address both problems simultaneously: we present StructureFlow, a novel and principled simulation-free training approach for jointly learning the structure and stochastic population dynamics of physical systems. We showcase the utility of StructureFlow for the tasks of structure learning from interventions and dynamical (trajectory) inference of conditional population dynamics. We empirically evaluate our approach on high-dimensional synthetic systems, a set of biologically plausible simulated systems, and an experimental single-cell dataset. We show that StructureFlow can learn the structure of underlying systems while simultaneously modeling their conditional population dynamics --- a key step toward model-based mechanistic understanding of systems behavior.

Noah El Rimawi-Fine, Adam Stecklov, Lucas Nelson +4

McGill University · Mila – Quebec AI Institute · AITHYRA +5

Population Dynamics Stochastic Dynamics Simulation-Based Reinforcement Learning

Ask PDF

Jun 3, 2026cs.LG

Deep Embedded Multiplicative DMD for Algebra-Preserving Koopman Learning

Koopman theory turns nonlinear dynamics into a linear spectral problem. In computation, however, everything depends on a hard finite-dimensional choice: the observables must be expressive, nearly invariant under the dynamics, and, ideally, compatible with composition. Deep Koopman methods learn flexible coordinates, whereas structure-preserving methods enforce operator identities on fixed dictionaries. We combine these ideas by introducing Deep Embedded Multiplicative Dynamic Mode Decomposition (DeepMDMD), a method that learns a latent space and a partition of it, while enforcing the Koopman product rule as an exact algebraic constraint. Training alternates between an exact multiplicative operator update and a differentiable latent-clustering step that promotes Koopman closure. The result is a finite transition map on learned latent cells. Its nonzero spectrum lies on the unit circle, its dictionary is shaped by the dynamics rather than by ambient geometry, and forecasts are made in latent coordinates before being decoded to physical space. Across Hamiltonian, chaotic, and fluid examples, DeepMDMD learns dictionaries that are far more compact and dynamically coherent than those produced by geometric MDMD partitions. It reduces spectral pollution, reveals richer continuous-spectrum structure, and gives stable forecasts under severe noise. In high-dimensional flows, including a 158,624-dimensional cylinder wake and a noisy

Re=20,000

lid-driven cavity, it preserves coherent structures and long-time spectral statistics where state-space MDMD fails. These results suggest a practical rule for Koopman learning: learn the coordinates, constrain the algebra.

Kelan Gray, Finlay Brown, Nicolas Boullé +1

Department of Mathematics, Imperial College London, London, SW7 2AZ, UK. · Department of Applied Mathematics and Theoretical Physics, University of Cambridge, Cambridge, CB3 0WA, UK.

Koopman Operator Dynamic Mode Decomposition Spectral State Space Models Dataset Distillation

Ask PDF