MoCoO: Momentum Contrast ODE-Regularized VAE for Single-Cell Trajectory Inference and Representation Learning

Abstract

Characterising cellular differentiation from single-cell RNA sequencing (scRNA-seq) requires representations that capture both discrete cell-type identity and continuous developmental trajectories. We present MOCOO, a modular framework integrating a Variational Autoencoder (VAE), Neural Ordinary Differential Equations (Neural ODE), and Momentum Contrast (MoCo), complemented by a systematic Phase-2 Flow Matching (FM) refinement step applicable to all model variants. Through a systematic six-configuration ablation across 20 scRNA-seq datasets evaluated with a proposed five-metric suite covering clustering geometry (ASW, DAV, CAL) and embedding quality (DRE, DREX), we demonstrate two central findings. First, the ODE+MoCo combination is the core architectural synergy: VAE+ODE+MoCo achieves four of five top-two finishes among base configurations, including the best ASW (0.225) and DAV (1.478), plus second-best DRE (0.640) and CAL. Second, FM refinement systematically improves both embedding quality and clustering geometry across all six base configurations--DREX in 92% and DRE in 88% of 120 dataset- configuration pairs ({Delta}DREX= +0.030, {Delta}DRE= +0.023), CAL in 88%, ASW in 86% ({Delta}ASW= +0.018), and DAV in 80% ({Delta}DAV= -0.072; Fig. 2). Combined, the full MoCoO pipeline (VAE+ODE+MoCo+Proto+FM) achieves the best DRE (0.678), DREX (0.660), and CAL, while VAE+ODE+MoCo+FM achieves the best ASW (0.257) and DAV (1.359). ODE smooths the latent manifold along developmental trajectories; MoCo sharpens cluster geometry; FM recovers and amplifies both embedding quality and cluster separation post-hoc. Downstream validation confirms that MoCoO latent spaces support annotation transfer, uncertainty quantification, differential expression, and branching detection. Pseudotime predictions correlate significantly with canonical marker genes across all five core developmental systems. We publicly release the MoCoO Python package (pip install mocoo) and full benchmark suite. O_FIG O_LINKSMALLFIG WIDTH=200 HEIGHT=65 SRC="FIGDIR/small/714791v1_fig2.gif" ALT="Figure 2"> View larger version (22K): org.highwire.dtl.DTLVardef@59a7feorg.highwire.dtl.DTLVardef@242f8borg.highwire.dtl.DTLVardef@1ad851dorg.highwire.dtl.DTLVardef@fe6b2c_HPS_FORMAT_FIGEXP M_FIG O_FLOATNOFig. 2.C_FLOATNO FM systematically improves all model variants across all metrics (Fig. 2). Each subplot shows absolute performance for all 12 configurations (6 base + 6 FM variants) across 20 datasets. Coloured triangle indicators above each pair show FM effect: [Figure 2] improved, [Figure 2] degraded. Subplot titles include FM win-rate (% of 120 config x dataset pairs improved). ODE-containing configurations (V+OM, Full) dominate all metrics, and FM universally raises the median across all 6 base variants (CAL 88%, DREX 92%, DRE 88%, ASW 86%, DAV 80%). C_FIG