Pub-AI: AI in Science

JEPA · 2026-09-30

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.33350PDF

Auto-summarized: this summary was generated by a language model from the paper’s text and has not been reviewed by an editor. Check the paper before relying on it. Benchmark numbers appear only after a human has verified them.

TL;DR

KoopCell is a Koopman-based generative framework for learning continuous single-cell population dynamics from temporally sparse, unpaired distribution snapshots. It jointly learns latent coordinates and linear latent dynamics guided by Koopman–Mori–Zwanzig theory, with convergence guarantees from a weak continuity equation. KoopCell-M adds non-Markovian memory via Markovian embedding to model branching in latent dynamics.

Why it matters

The authors address learning developmental dynamics when scRNA-seq provides only cross-sectional, unpaired distribution snapshots (no cell-level correspondence). Their Koopman/continuity-equation formulation targets linear latent dynamics from snapshots, and KoopCell-M adds a memory mechanism to represent branching, which prior interpolation-focused methods may miss.

Method

  • Learn latent observables with a VAE (encoder/decoder) whose latent distribution evolution matches a linear Koopman flow using a weak continuity equation residual in Wasserstein space; estimate the Koopman generator via a closed-form least-squares problem from random Fourier test functions.
  • Alternate between estimating the linear propagator from snapshot distributions (while fixing observables) and refining the representation (while fixing the propagator), using a transported prior and population-matching loss with debiased Sinkhorn divergence.
  • Extend to branching with KoopCell-M: incorporate non-Markovian memory using Mori–Zwanzig formalism and a Markovian embedding with auxiliary hidden memory variables learned via variational inference.

Limitation

The authors state that a linear latent flow with a Lipschitz decoder cannot exactly transport an initial connected latent support to a target distribution with disconnected support when the number of components changes (citing Zhao et al., 2026, Proposition 3.3).

Abstract (from arXiv)

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.

Related papers