Technical Summary: Chreode – A Cell World Model for One-Step Temporal Dynamics and Perturbation Prediction

Problem Statement

Predicting how a cell's transcriptional state changes under developmental signals or genetic perturbations is central to in-silico biology and the AI Virtual Cell program. Existing approaches face four critical limitations:

1. Underutilization of Temporal Data: Most methods collapse destructive scRNA-seq snapshots into static control-to-treated maps, discarding the rate-of-change information inherent in time-series data.
2. Improper Action Modeling: Many models treat interventions (drugs, knockouts) as latent contexts to be concatenated rather than as causal do-calculus interventions, failing to resolve compositional or trajectory-dependent perturbations.
3. High Inference Costs: Multi-step solvers (ODEs, Schrödinger bridges) require iterative integration for every query. Scaling to screen-level queries ($10^8$ tuples) results in prohibitive computational costs ($10^{10}$ forward passes).
4. Lack of Biological Inductive Bias: Current flow-matching and optimal transport methods parameterize velocity fields as generic neural networks, ignoring the biological structure of Waddington's landscape (downhill potential flow, rotational currents, and stochastic spread).

Methodology: Chreode

Chreode is introduced as a one-step cell world model that predicts action-conditioned cell-state transitions $p(z_{t+\Delta} | z_t, \text{do}(a))$ via a structured residual transition operator.

1. Core Architecture: The Waddington Residual

Instead of iterative rollout, Chreode predicts the future latent state $\hat{z}_{t+\Delta}$ in a single forward pass by adding a learned residual $R_\theta$ to the input state $z_t$:
$$\hat{z}_{t+\Delta} = z_t + \alpha(\Delta) \cdot R_\theta(z_t, \Delta, a, \epsilon)$$
where $\alpha(\Delta) = 1 - e^{-\Delta/\tau}$ is a learnable time gate ensuring zero change at $\Delta=0$. The residual $R_\theta$ is decomposed into three biologically interpretable components:

• Downhill Potential Flow ($-\nabla_z U_\theta$): The negative gradient of a learned potential function $U_\theta$, driving cells toward low-potential fate basins.
• Antisymmetric Flow ($S_\theta z_t$): A rotational component modeling non-equilibrium currents (e.g., cell cycle). $S_\theta$ is parameterized as $P_\theta Q_\theta^\top - Q_\theta P_\theta^\top$ to guarantee exact antisymmetry ($S_\theta^\top = -S_\theta$).
• Stochastic Spread ($\sigma_\theta \odot \epsilon$): Element-wise noise scaling to capture intrinsic biological variability.

2. Training Pipeline

Chreode employs a two-stage pretraining recipe on a 2.4M-cell mouse embryonic atlas (7 datasets, 88 timepoints):

• Stage 1 (Encoder): A shared scVI encoder is trained over a fixed cross-species ortholog vocabulary (16,520 mouse-human 1:1 orthologs) to produce a frozen latent representation.
• Stage 2 (Dynamics Backbone): A DiT-based (Diffusion Transformer) backbone learns the one-step transition operator. It uses a population-level drifting-field loss adapted from Deng et al. [2026], combining:
  • Kernel MMD and Sinkhorn $W_2$ for distribution matching.
  • A drifting-field term ($L_{drift}$) to enforce equilibrium consistency.
  • A downhill regularizer ($L_{down}$) to prevent uphill movement.
• Time Embedding Decoupling: To ensure stability during long-horizon extrapolation, the potential branch uses a flexible Time2Vec embedding, while the antisymmetric branch uses a bounded low-frequency Fourier code to prevent rotation rates from drifting into unseen frequencies.

3. Transfer Modes

The pretrained backbone supports two distinct downstream transfer modes:

1. Fine-tuning Initialization: The backbone initializes models for time-transition prediction on held-out datasets (e.g., Weinreb hematopoiesis, Veres islet differentiation).
2. Gene-State Embedding Injection: The pretrained dynamics representation is injected as the gene-state embedding into the existing GEARS perturbation predictor, without modifying GEARS's training procedure.

Key Results

Temporal Dynamics Prediction

• Weinreb Hematopoiesis: Fine-tuned Chreode achieves the lowest Sinkhorn $W_2$ distance at days 4 and 6 compared to matched scratch models, PI-SDE, and PRESCIENT.
• Veres Islet Differentiation: Chreode outperforms all baselines (including PI-SDE and PRESCIENT) across all seven target timepoints ($t=1$ to $7$).
• Zero-Shot Clonal Fate: Using the frozen backbone without fine-tuning, Chreode produces clonal fate scores (Pearson correlation $r=0.468$) competitive with strong dynamic-OT baselines (moscot, WOT) and neural-SDE baselines (scDiffEq).

Perturbation Prediction

• Norman Perturb-seq: When used as a gene-state embedding within GEARS, the pretrained dynamics representation reduces the shared-vocabulary DE20 Mean Squared Error (MSE) from 0.2121 to 0.1858, a 12.4% relative improvement. This suggests that developmental trajectory dynamics encode primitives transferable to CRISPR-induced state shifts.

Efficiency

• Inference Speed: Chreode resolves queries in a single forward pass (1 NFE), making it approximately 3.0$\times$ faster than PRESCIENT and 6.3$\times$ faster than CellFlow on a single NVIDIA A100. This amortization is critical for scaling to $10^8$ query screens.

Significance and Claims

The paper claims that Chreode is the first cell foundation model to combine:

1. Single-step inference (amortizing transition costs).
2. A pretrained backbone shared across downstream datasets.
3. Joint time and action conditioning.
4. A structured residual prior explicitly decomposing velocity into potential gradients, antisymmetric flows, and stochastic spread.

The authors interpret the successful transfer of developmental dynamics to perturbation prediction as evidence that both processes traverse cell-state manifolds along coherent trajectories. By shifting distributional evolution from inference time to training time, Chreode enables single-pass generation while preserving a Waddington-inspired decomposition. The work positions itself as a structural contribution to the "AI Virtual Cell" agenda, offering a scalable, biologically grounded alternative to multi-step solvers and static perturbation maps.

Limitations noted by the authors:

• Pretraining is restricted to mouse embryonic data; cross-species transfer relies solely on ortholog mapping.
• Time-transition transfer requires fine-tuning (not strict zero-shot) for optimal performance.
• The perturbation transfer reuses the GEARS architecture rather than proposing a standalone perturbation operator.
• The scale (2.4M cells) is smaller than concurrent large cellular world models trained on proprietary data.