PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.33350v1cs.LGcs.AIq-bio.QM

TL;DR

Temporally sparse, unpaired snapshots make it difficult to infer continuous population dynamics relevant to developmental biology. KoopCell learns linear latent dynamics from distributions, and KoopCell-M adds memory for branching; experiments report spectrum recovery and improved extrapolation across single-cell datasets.

  • Problem

    Inferring continuous population dynamics from temporally sparse, unpaired distribution snapshots is challenging, including in single-cell biology where destructive measurements provide no cell-level correspondence across time.

    Source: §1 · p. 1

  • Method

    KoopCell jointly learns an observable representation and linear latent dynamics from snapshots using the weak continuity equation, while KoopCell-M adds non-Markovian memory through an augmented latent embedding.

    Source: §1 · p. 2

  • Results

    KoopCell-M reduces average error relative to the strongest baseline for each dataset by 3.6–7.3% on the medium task and 2.0–6.6% on the hard task’s extrapolation portion, while the models recover spectra and improve extrapolation of gene-expression distributions.

    Source: §4.1 · p. 7, §4.3 · p. 9, §6 · p. 10

  • Takeaways & Limitations

    The findings support quantitative analysis of developmental dynamics beyond the observed time window from sparse, unpaired measurements.

    Source: §6 · p. 10

Abstract

from arXiv · show

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

KoopCell learns linear latent population dynamics directly from temporally sparse, unpaired distributions, addressing limits in extrapolation and explicit modeling of branching. KoopCell-M adds memory to represent branching, and experiments report spectral recovery and strong extrapolation performance.

1 INTRODUCTION

  • Existing neural dynamical models can interpolate missing distributions within the observed horizon but do not effectively extrapolate or explicitly account for branching.Both capabilities matter for predicting future cell-population states and representing developmental lineage differentiation.Source: §1 · pp. 1–2
  • KoopCell jointly learns observables and linear latent dynamics from unpaired distribution snapshots, using the weak continuity equation in a generative model.The framework includes a closed-form estimator and convergence guarantees.Source: §1 · p. 2
  • KoopCell-M incorporates non-Markovian memory in an augmented latent space to model branching dynamics.The paper presents this as a mechanistic formulation with a practical learning algorithm.Source: §1 · p. 2
  • Experiments on synthetic systems and three scRNA-seq datasets report Koopman spectral recovery, branching modeling, and state-of-the-art extrapolation performance for cellular population distributions.The reported evaluations use sparse, unpaired measurements.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries frame population snapshots as distributions over high-dimensional states and motivate learning approximately linear latent coordinates. Koopman dynamics supply the linear operator view, while Mori–Zwanzig memory handles unresolved influences when a reduced linear closure fails.

Problem Formulation and Notation

  • The setup treats each time point as a distribution with independent observed samples, and uses a lower-dimensional latent space for high-dimensional gene expression data.The latent coordinates are guided toward approximately linear evolution rather than assigned a heuristic dynamical structure.Source: §2 · p. 2

Koopman Theory and Koopman Learning

  • Koopman operators evolve observables linearly even when the underlying state dynamics are nonlinear; an invariant finite-dimensional observable subspace yields a linear latent system.Traditional EDMD estimates a propagator from paired states, while learned Koopman representations jointly fit observables and dynamics.Source: §2 · pp. 2–3

Mori–Zwanzig Formalism and Memory Embedding

  • When the chosen observable subspace is not invariant, Mori–Zwanzig formalism represents unresolved influences through memory and forcing terms.A kernel approximation with auxiliary hidden variables yields an augmented Markovian state, although the resolved coordinates alone generally retain memory.Source: §2 · p. 3

Wasserstein Flow and Optimal Transport

  • Wasserstein flow describes the evolution of probability measures, with a canonical velocity satisfying the continuity equation and having minimum L2(µ) norm among velocities generating the same curve.The negative Sobolev norm can be expressed as the minimum L2(µ) norm of a tangent field generating a zero-mass distribution.Source: §2 · p. 3

3 METHODS

KoopCell couples a VAE representation with a linear latent flow estimated from distribution snapshots through weak continuity equations. KoopCell-M augments that flow with hidden memory and path-specific initial conditions, allowing a shared visible starting distribution to develop into multiple density modes.

3.1 KOOPMAN-BASED VAE FOR GENERATIVE MODELING OF POPULATION DYNAMICS

  • KoopCell uses a VAE to encode cellular states into observables with approximately linear evolution, estimate a linear propagator, and decode latent representations into cellular states.The encoder includes time in the augmented state, and the implementation describes the latent flow as ˙Zt = AZt.Source: §3.1 · p. 3
  • The method alternates between estimating the generator with observables fixed and refining the representation with the propagator fixed.This learns the Koopman representation and its generator from sparse, unpaired snapshots.Source: §3.1 · pp. 4–5
  • Restricting the weak population residual to finitely many test functions gives a weighted least-squares estimator of the linear generator, whose minimizers converge under test-enrichment and uniqueness assumptions.The estimator is closed form when the design matrix has full column rank.Source: §3.1 · p. 4
  • Under stated approximation, stability, and identifiability assumptions, the paper establishes vanishing minimum weak residual as dimension grows, Wasserstein prediction bounds, and locally uniform strong semigroup convergence.In practice, random Fourier tests, trapezoidal integration over adjacent snapshot intervals, empirical averages, and ridge regularization form the regression problem.Source: §3.1 · pp. 4–5
  • Population-matching and weak-residual losses encourage accurate predictions and latent coordinates consistent with the weak continuity equation; matrix exponentials permit predictions at arbitrary times.The linear-flow prior is transported from an initial Gaussian prior according to the learned propagator.Source: §3.1 · p. 5

3.2 KOOPCELL-M: MEMORY-AUGMENTED KOOPMAN DYNAMICS FOR BRANCHING

  • A linear flow followed by a Lipschitz decoder cannot produce disconnected fate populations from connected latent support, motivating an explicit branching mechanism.The discussion also notes that linear flow alone cannot create new latent density modes.Source: §3.2 · p. 5
  • KoopCell-M evolves visible coordinates and auxiliary memory jointly as an augmented Markovian system, while the visible coordinates alone retain non-Markovian dynamics.A discrete path variable indexes candidate root-to-terminal paths, with path-specific hidden-memory initialization.Source: §3.2 · pp. 5–6
  • Although paths share the same visible initial distribution, coupling hidden heterogeneity into visible dynamics can separate component means and produce multiple density modes.The method alternates dynamics estimation and representation learning, infers path and memory variables, and uses population matching and path consistency objectives.Source: §3.2 · p. 6

4 EXPERIMENTS

Across synthetic and single-cell experiments, KoopCell recovers spectral structure from sparse, unpaired snapshots, while its memory-augmented variant models branching and improves distribution prediction. The experiments also examine latent-state organization and lineage-associated gene programs.

4.1 SPECTRAL RECOVERY FROM UNPAIRED SNAPSHOTS

  • KoopCell consistently recovers the Koopman spectrum with low MAE, and its learned observables reproduce Toy I level-set geometry and distinguish Toy II intrawell relaxation from attraction basins.For Toy II, the complex mode encodes decaying intrabasin amplitude and phase, while the near-zero mode separates the two basins.Source: §4.1 · pp. 6–7, Figure 1 · p. 8
  • KoopCell achieves lower extrapolation SWD than scNODE across all settings, and the results suggest that learned spectral structure supports more stable predictions beyond training times.The comparison covers two toy systems with known Koopman spectra.Source: §4.1 · p. 7, Table 2 · p. 8

4.2 MODELING BRANCHING DYNAMICS WITH KOOPCELL-M

  • On a stochastic pitchfork system, KoopCell-M more closely reproduces the two-branch separation and has lower SWD and branch-coordinate W1 errors than KoopCell.The system begins unimodal and is evaluated using held-out-time errors; population rollouts are shown in a shared PCA plane.Source: §4.2 · p. 7, Table 3 · p. 8, Figure 2 · p. 8

4.3 MODELING HIGH-DIMENSIONAL SINGLE-CELL POPULATION DYNAMICS

  • KoopCell achieves lower average W2 distances than all baselines on the medium task and hard-task extrapolation portion; relative to each dataset’s strongest baseline, KoopCell-M reduces average error by 3.6–7.3% and 2.0–6.6%, respectively.Source: §4.3 · p. 9
  • On the SC hard task, KoopCell-M more clearly separates stromal and iPSC latent populations with distinct temporal progressions than KoopCell.The representation comparison labels cells by reference state and snapshot index.Source: §4.3 · p. 9, Figure 3 · p. 9
  • In zebrafish data, KoopCell-M recovers lineage-associated driver genes and predicts gene-expression programs enriched for muscle development in PSM and neural development in Hindbrain.The analyses examine predicted trajectories over 8–12 hours post fertilization.Source: §4.3 · p. 9

5 RELATED WORK

Prior work models single-cell population dynamics with neural ODEs, SDEs, optimal transport, and flow matching, while Koopman learning has largely relied on paired time-series data. This work develops generative Koopman learning from unpaired distribution snapshots and uses an augmented memory closure to model branching.

Learning scRNA-seq dynamics from distribution snapshots

  • Existing approaches include continuous normalizing flows, latent neural ODEs and SDEs, latent Gaussian processes, unbalanced optimal transport, and flow-matching methods for single-cell dynamics.Examples include TrajectoryNet (Tong et al., 2020), scNODE, and VGFM (Wang et al., 2025).Source: §5 · p. 9

Koopman learning and memory

  • Prior Koopman methods jointly learn autoencoders and latent linear propagators for prediction but learn from paired time-series data.The cited examples include Otto & Rowley (2019), Lusch et al. (2018), Liu et al. (2023), and Naiman et al. (2024).Source: §5 · p. 9
  • KoopCell develops theory and algorithms for generative Koopman learning from unpaired snapshots, with a Mori–Zwanzig-inspired augmented Markovian memory closure for branching dynamics.The related-work passage notes that exact linear closure for observables remains challenging.Source: §5 · p. 10

6 CONCLUSIONS

This section summarizes KoopCell’s learning of linear latent dynamics from unpaired distribution snapshots and KoopCell-M’s use of Markovian memory to capture branching, alongside theoretical links between weak population-residual minimization, generator estimation, and convergence. It also establishes stronger distributional and semigroup convergence from weak-residual convergence, with the semigroup results requiring an additional finite-time stability assumption.

6 CONCLUSIONS

  • KoopCell recovers Koopman spectra and improves extrapolation of gene-expression distributions from sparse, unpaired measurements; KoopCell-M captures branching through memory.The framework is presented as supporting quantitative analysis of developmental dynamics beyond the observed time window.Source: §6 · p. 10

A THEORY

  • The weak population residual yields the best projected linear approximation to the canonical Wasserstein velocity, and finite tests reduce its minimization to least-squares regression.Under dense test enrichment and uniqueness, finite-test minimizers converge to the population minimizer.Source: A THEORY · p. 17, §A.2 · p. 21

A.3 CONVERGENCE OF THE WEAK POPULATION RESIDUAL

  • As compatible observable dimension D increases in a Koopman-invariant RKHS, the minimum weak population residual vanishes under the stated assumptions.The bound uses a difference-quotient approximation to the generator and the projection error of the features.Source: A THEORY · p. 17, §A.3.2 · p. 23

A.3.2 CONVERGENCE OF THE WEAK POPULATION RESIDUAL

  • The weak population residual converges to zero by controlling approximation errors through dominated convergence, first taking h small and then increasing D.The proof uses strong convergence of the projection and finite orbit energy to show the error terms vanish.Source: §A.3.2 · pp. 23–24

A.4 FROM WEAK RESIDUAL CONSISTENCY TO STRONG SEMIGROUP CONVERGENCE

  • Vanishing weak residuals imply uniform Wasserstein convergence of predicted distributions on the training interval when the selected minimizers satisfy uniform finite-time stability.This is an additional hypothesis beyond weak-residual convergence and does not require a uniform bound on the generator norms.Source: §A.4 · p. 24
  • Under the stated stability bound, compactness yields strong operator subsequences and finite-time linear determinacy identifies their limits with T(t), giving semigroup convergence at every initial feature vector.Source: §A.4 · p. 24

A.4.1 THE WEAK RESIDUAL CONTROLS PROBABILITY PREDICTION

This appendix relates the weak population residual to distribution-prediction error, then gives conditions under which distribution snapshots identify the learned linear propagator and support operator convergence.

A.4.1 THE WEAK RESIDUAL CONTROLS PROBABILITY PREDICTION

  • A vanishing weak residual and projection error yield convergent distribution predictions, with bounds uniform over starting times and prediction lags.The full training-interval prediction error also vanishes when the residual and projection errors do.Source: §A.4.1 · p. 25

A.4.2 DISTRIBUTIONAL IDENTIFIABILITY OF THE LINEAR PROPAGATOR

  • Finite-time linear determinacy requires that only the true propagator maps every source distribution to its corresponding later distribution.This condition addresses the fact that distinct trajectories can produce identical distributions at every time.Source: §A.4.2 · pp. 25–26

A.4.3 FROM DISTRIBUTIONAL COMPACTNESS TO STRONG OPERATOR COMPACTNESS

  • Uniformly bounded operators with relatively compact pushforward distributions have strongly convergent subsequences when the distribution support spans the space.The appendix applies this result to averaged source distributions before using finite-time linear determinacy to identify each operator limit.Source: §A.4.3 · pp. 26–27

A.4.4 LOCALLY UNIFORM STRONG CONVERGENCE OF THE LEARNED PROPAGATORS

Under stability and finite-time linear determinacy, the appendix proves locally uniform strong convergence of learned propagators and describes the alternating training procedures for KoopCell and KoopCell-M. The implementation estimates latent linear dynamics from snapshot-based weak residuals while jointly refining representations.

A.4.4 LOCALLY UNIFORM STRONG CONVERGENCE OF THE LEARNED PROPAGATORS

  • Under stability and finite-time linear determinacy, learned propagators converge strongly and locally uniformly in time, including over finite intervals beyond the training interval.The result concerns minimizers of the exact weak population residual under the stated assumptions; practical sampling, quadrature, regularization, and optimization errors are not controlled.Source: §A.4.4 · pp. 27, 29

B ALGORITHMIC DETAILS

  • The appendix presents KoopCell and KoopCell-M training as alternating closed-form dynamics estimation and representation learning, with memory embedding and developmental-path conditioning added for branching.Held-out snapshots are reserved for evaluation.Source: B ALGORITHMIC DETAILS · p. 29, §B.2 · p. 32

B.1 KOOPCELL: KOOPMAN DYNAMICS AND ALTERNATING TRAINING

  • KoopCell alternates between learning a VAE representation for fixed latent dynamics and estimating the linear generator in the current latent coordinates.The VAE uses a transported prior induced by the latent flow and adds population-prediction and weak-residual losses to its objective.Source: §B.1 · p. 29, §B.1.1 · p. 30

B.1.1 VAE REPRESENTATION LEARNING

  • The VAE encodes cell expression and time into a latent posterior and decodes latent states back to expression.A transported Gaussian prior is used, with reconstruction and variational regularization in the training objective.Source: §B.1.1 · pp. 29–30

B.1.2 EMPIRICAL WEAK RESIDUAL AND CLOSED-FORM DYNAMICS ESTIMATION

  • Snapshot-based empirical weak residuals form a weighted least-squares problem for estimating the linear generator, with ridge regularization and an exponential moving average used to stabilize updates.Trapezoidal quadrature accommodates irregular observation times; its interval error is O(∆t^3) when the moment integrand has two bounded time derivatives.Source: §B.1.2 · p. 31

B.2 KOOPCELL-M: MEMORY-AUGMENTED BRANCHING DYNAMICS AND ALTERNATING TRAINING

  • KoopCell-M introduces augmented latent dynamics to represent memory for branching while retaining the two alternating training steps.Those steps are closed-form dynamics estimation followed by representation learning.Source: §B.2 · p. 32

B.2.1 MEMORY-AUGMENTED DYNAMICS AND DEVELOPMENTAL PATHS

This section presents KoopCell-M’s use of hidden memory and path-specific initial conditions to model developmental fate modes, alongside a training procedure that alternates closed-form updates of memory-augmented dynamics with gradient-based learning of representations and memory dynamics. It also covers path inference, rollouts, path-consistency training, and weak-statistics regression using inferred memory statistics.

B.2.1 MEMORY-AUGMENTED DYNAMICS AND DEVELOPMENTAL PATHS

  • A linear latent flow preserves Gaussianity, while continuous decoding preserves connectedness; KoopCell-M instead couples visible states to memory and heterogeneous hidden initial conditions.The augmented state evolves linearly with dissipative memory dynamics; separated Gaussian components can form multiple density modes, while nondegenerate covariances leave mixture support connected.Source: §B.2.1 · p. 33
  • With memory feedback, initially unimodal marginals develop two visible modes that persist; removing feedback yields a single mode again for sufficiently large t.In Example B.1, the memory kernel is K(u) = e^-u, and the component means approach ±a/2 as t → ∞. Figure 4 shows both models becoming bimodal at t = 1, but only the feedback model remains bimodal at t = 3.Source: §B.2.1 · pp. 33–34, Figure 4 · p. 34

B.2.2 VARIATIONAL INFERENCE AND REPRESENTATION LEARNING

  • The VAE inference model uses a Gaussian encoder for visible state, path probabilities conditioned on state and expression, and an analytic Gaussian conditional posterior for memory.Path probabilities combine the path prior, visible-state marginal, and annotation weights; accepted labels restrict support to compatible paths, while unlabeled cells receive unit weights. A path consistency loss supervises generated trajectories to follow pathway state order.Source: §B.2.2 · pp. 35–37

B.2.3 CLOSED-FORM DYNAMICS ESTIMATION AND ALTERNATING TRAINING

  • The closed-form update for visible dynamics Θ = [A, bz] substitutes visible drift and the inferred conditional memory mean into weak population residuals.The derivation holds the representation and memory statistics fixed and averages over observed expression, encoder samples, and inferred path probabilities.Source: §B.2.3 · p. 38
  • The Θ regression uses weak statistics estimated from encoded cells, trapezoidal time-integral approximations, and fixed path probabilities, conditional memory means, and memory coupling B.The stacked interval–test residual is y_M − Ga, with a = vec(Θ); the algorithm forms weak statistics from empirical averages and the trapezoidal rule.Source: §B.2.3 · p. 38
  • Training first pretrains the VAE and state classifier, then alternates a closed-form update of Θ with gradient-based updates of the representation and memory dynamics.The gradient step minimizes L_M = L_snap + L_pop + λ_R L_path + λ_weak L_weak, where L_weak = ∥y_M − Ga∥².Source: §B.2.3 · p. 38
  • Algorithm 2 computes path probabilities, Gaussian moments, and conditional coefficients, while minibatch weak loss holds q, C, and B fixed and retains gradients through z.Rollouts and the path-consistency loss follow the specified model equations.Source: §B.2.3 · p. 38

C SPECTRAL RECOVERY EXPERIMENTS

The toy-model protocol generates population snapshots by independently sampling initial states and evolving them to each observation time. This produces unpaired time marginals without cross-time cell correspondence.

C SPECTRAL RECOVERY EXPERIMENTS

  • Each toy dataset consists of empirical population snapshots drawn from the time marginals of the corresponding dynamical system.The marginals are defined by pushing the initial population distribution through the system flow.Source: C SPECTRAL RECOVERY EXPERIMENTS · p. 39
  • At each observation time, a new population of initial states is sampled, evolved from time 0, and retained only at its terminal time.The construction is repeated independently across observation times.Source: C SPECTRAL RECOVERY EXPERIMENTS · p. 39
  • Independent resampling at every time removes cross-time cell correspondence and yields the unpaired snapshots used for training and evaluation.Source: C SPECTRAL RECOVERY EXPERIMENTS · p. 39

C.1 TOY 1: FINITE-DIMENSIONAL NONLINEAR FLOW

Toy 1 tests spectral recovery and population prediction in a nonlinear flow with an attracting slow manifold. The experiments compare KoopCell with scNODE across learned-observable and high-dimensional settings, including extrapolation beyond the training horizon.

C.1 TOY 1: FINITE-DIMENSIONAL NONLINEAR FLOW

  • The Toy 1 flow has an attracting slow manifold, and its transverse coordinate decays faster than the coordinates parameterizing that manifold.The specified observables form a finite-dimensional Koopman-invariant subspace with eigenvalues (−0.2, −1, −0.4).Source: §C.1 · p. 39
  • Training uses nine snapshots from t = 0 to 1.2, while independent target snapshots from t = 1.4 to 4.0 evaluate extrapolation.Each selected time uses 3,000 newly sampled initial states, and generated snapshots show rapid contraction toward the slow manifold.Source: §C.1 · p. 39
  • The experiments include a known-observable setting that fixes the analytically specified Koopman-invariant map to isolate spectral recovery, and a learned-observable setting that jointly learns observables and a linear propagator.The learned-observable model uses a three-dimensional latent space; scNODE uses the same snapshots and latent dimension with an autonomous neural ODE.Source: §C.1 · pp. 39–40
  • A high-dimensional setting supplies only 64-dimensional observations through a fixed injective map, evaluating representation learning and prediction in observed space without reporting an inferred spectrum.The physical state and observation map remain unknown to both models.Source: §C.1 · p. 40
  • Table 4 reports spectral recovery for the known- and learned-observable settings, with parentheses indicating absolute eigenvalue errors.Source: Table 4 · p. 40, §C.1 · p. 40
  • KoopCell predictions agree more closely with ground-truth Toy 1 population distributions than scNODE predictions across extrapolation times.Rollouts start from the observed population at t = 0; times after t = 1.2 are extrapolations.Source: §C.1 · p. 40, Figure 7 · p. 41

C.2 TOY 2: DAMPED DUFFING FLOW

Toy 2 evaluates KoopCell on the damped Duffing flow, whose two basins have distinct long-term outcomes. It assesses learned spectral behavior and population prediction through extrapolation to t = 12, including in a high-dimensional observation setting.

C.2 TOY 2: DAMPED DUFFING FLOW

  • The damped Duffing system has stable spiral equilibria at (q, p) = (±1, 0), and a constant Koopman eigenfunction distinguishes its two basins of attraction.Its global complex eigenfunctions lack closed-form expressions; their magnitude and phase provide action–angle-like coordinates within each basin.Source: §C.2 · p. 42
  • Training uses ten population marginals from t = 0 to 4.5, with independently generated extrapolation targets at t = 5, 6, 8, 10, and 12.Each time marginal contains 3,000 independently resampled initial states in two connected annular sectors, one per basin.Source: §C.2 · p. 42
  • In the learned-observable setting, KoopCell learns a three-dimensional observable, linear generator, and nonlinear decoder from physical snapshots.The experiment examines whether learned modes recover decaying intrabasin amplitude and phase, plus a near-zero mode that distinguishes the long-term basin outcomes.Source: §C.2 · p. 42
  • The high-dimensional setting applies a fixed injective map to the physical marginals, and neither KoopCell nor scNODE receives the physical coordinates or map during training.Both use a three-dimensional latent state and are evaluated in the native 64D observation space.Source: §C.2 · p. 42
  • Table 5 compares learned-observable Toy 2 spectral recovery with values implied by the Duffing dynamics.Source: §C.2 · p. 42, Table 5 · p. 43
  • KoopCell more faithfully reproduces the target Duffing population over the full prediction horizon through t = 12 than scNODE.Figure 9 compares their rollouts against independently generated ground truth.Source: §C.2 · p. 43, Figure 9 · p. 43

C.3 ROLLOUT-ERROR ACCUMULATION

Rollout errors show KoopCell maintaining lower error than scNODE during extrapolation on both toy systems.

C.3 ROLLOUT-ERROR ACCUMULATION

  • KoopCell maintains lower rollout error than scNODE throughout the evaluated extrapolation horizon on both Toy 1 and Toy 2.On Toy 1, scNODE’s error rises steadily beyond the final training snapshot while KoopCell’s remains low.Source: §C.3 · p. 43, Figure 10 · p. 43

D BRANCHING DYNAMICS EXPERIMENTS

The branching experiment tests prediction on synthetic noisy pitchfork dynamics from unpaired population snapshots. It compares KoopCell and KoopCell-M using shared training and evaluation protocols and several population-level error measures.

D BRANCHING DYNAMICS EXPERIMENTS

  • The experiment simulates a noisy pitchfork process in which two stable branches emerge as the progression coordinate becomes positive.Snapshots are independently shuffled; training uses unpaired populations without trajectory correspondences or branch labels.Source: D BRANCHING DYNAMICS EXPERIMENTS · p. 44
  • Both models use a 10-dimensional Gaussian-encoded latent state, while KoopCell-M augments it with a 24-dimensional memory state and two hidden components.The decoder is spectrally normalized to control its Lipschitz constant and limit amplification of latent differences.Source: D BRANCHING DYNAMICS EXPERIMENTS · p. 44
  • A shared gradient-based protocol replaces the closed-form generator updates to isolate the dynamical model from the generator-estimation procedure.Both models start from a shared pretrained VAE and are then trained with the same reconstruction and population-matching objectives.Source: D BRANCHING DYNAMICS EXPERIMENTS · p. 44
  • Predictions are evaluated against ground-truth populations using SWD, energy distance, squared MMD, branch-coordinate W1, and central mass error.The branch-coordinate W1 measures discrepancies in branch locations, widths, and relative population sizes; central mass error measures the fraction remaining between branches.Source: D BRANCHING DYNAMICS EXPERIMENTS · pp. 44–45
  • Table 6 reports prediction errors averaged over training, interpolation, and extrapolation times, with lower values indicating better performance.The table caption specifies means and standard deviations across three optimization seeds.Source: D BRANCHING DYNAMICS EXPERIMENTS · p. 45, Table 6 · p. 45

E.1 SINGLE-CELL DATASETS AND PREPROCESSING

The single-cell benchmark covers three datasets and three prediction tasks of increasing difficulty. The section describes how populations are predicted and evaluated, how baselines and model settings are handled, and how held-out predictions and latent representations are visualized.

E.1 SINGLE-CELL DATASETS AND PREPROCESSING

  • The benchmark uses ZB, DR, and SC scRNA-seq datasets, with 38,731, 27,386, and 23,619 cells, respectively.The datasets cover zebrafish and Drosophila embryogenesis and mouse embryonic-fibroblast reprogramming.Source: §E.1 · p. 45
  • Easy tasks withhold intermediate snapshots, medium tasks withhold final snapshots, and hard tasks combine intermediate and final held-out snapshots.The benchmark uses 2,000 highly variable genes; ZB and DR counts are normalized and log-transformed, while SC uses benchmark-provided normalized expression.Source: §E.1 · p. 45, Figure 12 · p. 46

E.1.1 POPULATION PREDICTION AND EVALUATION

  • Each method generates 2,000 cells per held-out time point, and predictions are compared with observed cells in the processed 2,000-gene space.KoopCell and KoopCell-M start from the earliest observed population, propagate encoded latent states, and decode predicted gene-expression distributions.Source: §E.1.1 · p. 46

E.2 IMPLEMENTATION OF BASELINE MODELS

  • Six representative baselines are included, with hyperparameters selected using released tuned configurations or training-only three-fold cross-validation.The held-out test time points remain inaccessible during model selection.Source: §E.2 · pp. 46–47

E.3 MODEL CONFIGURATIONS

  • Tables 8 and 9 describe settings shared across datasets and tasks and the selected MLP architectures and loss weights.Table 8 labels Stage 1 as VAE pretraining and Stage 2 as alternating dynamics/representation training; Table 9 identifies dh as the memory dimension.Source: Table 8 · p. 48, Table 9 · p. 48

E.4 ADDITIONAL SINGLE-CELL RESULTS

  • Tables 10 and 11 report held-out Wasserstein-2 distances for the easy and medium tasks, respectively.The table captions specify means and sample standard deviations over three seeds and identify the best and second-best means.Source: Table 10 · p. 49, Table 11 · p. 50

E.4.1 PREDICTED CELL POPULATIONS AT HELD-OUT TIME POINTS

  • Figures 13–15 compare predicted and observed held-out populations across datasets, with shared projections and axis ranges within each dataset.Figure 16 shows the same observed SC cells under each method’s own UMAP coordinates, colored by reference cell state and snapshot time.Source: §E.4.1 · p. 50, Figure 13 · p. 51, Figure 14 · p. 52

E.5 BIOLOGICAL DOWNSTREAM ANALYSIS

The downstream analysis assigns cells probabilities of reaching terminal lineages, then relates those probabilities to gene expression and inferred expression-rate programs. The analyses recover lineage-associated markers and lineage-specific developmental enrichments in zebrafish.

Lineage-associated gene identification

  • The method infers each cell’s terminal-lineage fate probability by propagating augmented states along paths and averaging over inferred memory and path distributions.Terminal cell-type reference annotations define the endpoints used in the fate probabilities.Source: §E.5 · p. 54

Recovery of developmental markers

  • MSGN1 and TBX16 rank highly for PSM, while SOX19A and SOX3 rank highest for Hindbrain; all eight selected genes have higher terminal expression in their corresponding lineage.The reported correlations are 0.665 and 0.500 for MSGN1 and TBX16, and 0.568 and 0.540 for SOX19A and SOX3.Source: §E.5 · p. 55, Table 12 · p. 55

Dynamic gene-expression programs

  • The dynamic-program analysis propagates cells from 8 hpf to 12 hpf and estimates lineage-specific gene expression rates from decoded trajectories.Genes are selected by mean predicted rate of change, standardized across time, and grouped into three Ward modules per lineage.Source: §E.5 · pp. 55–56

Developmental functions of the expression programs

  • PSM modules are enriched for muscle-development functions, while Hindbrain modules are enriched for neural-development functions; both lineages also have a keratin-containing third module.Examples include positive regulation of muscle cell differentiation (q = 6.91 × 10^-4) and nervous system development (q = 2.47 × 10^-5).Source: §E.5 · p. 56, Figure 18 · p. 57
Loading 2609.33350v1…