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

Abstract: Predicting how a cell will change its transcriptional state under a developmental signal or a genetic perturbation is the computational core of in-silico biology and the AI Virtual Cell program. Existing approaches either fit static control-to-treated maps that discard time, or solve multi-step ODE / Schrödinger-bridge problems on each dataset independently. We introduce Chreode, a one-step cell world model that predicts action-conditioned cell-state transitions through a structured residual transition operator. It shifts distributional evolution from inference time to training time, enabling single-pass generation while preserving a Waddington-inspired decomposition into downhill landscape flow, rotational in-tangent dynamics, and stochastic spread. The model is pretrained with a shared scVI encoder and a DiT-based dynamics backbone on a 2.4M-cell mouse embryonic atlas spanning 7 datasets. As a fine-tuning initialization, Chreode improves per-target Sinkhorn distance on Weinreb hematopoiesis and Veres islet differentiation over matched scratch models, PI-SDE, and PRESCIENT. As a transferable gene-state embedding for GEARS, the pretrained dynamics representation reduces shared-vocabulary DE20 mean squared error on Norman Perturb-seq from 0.2121 to 0.1858, a 12.4% relative improvement, without changing the GEARS training procedure. We interpret this transfer to perturbation prediction as evidence that pretrained developmental-trajectory dynamics encode differentiation primitives transferable to CRISPR-induced state shifts, since both involve cell-state transitions in a shared latent geometry. The pretrained backbone additionally produces zero-shot clonal fate scores on Weinreb that are competitive with strong dynamic-OT baselines.

• The paper introduces Chreode, a transferable model that uses a single-pass update integrating gradient flow, rotational dynamics, and stochastic noise for cell-state transitions.
• It employs a DiT-based backbone with decoupled time encodings and biologically inspired priors from Waddington’s landscape to ensure robust prediction.
• Empirical evaluations demonstrate significant improvements in transition prediction accuracy and transferability in both developmental and gene perturbation tasks.

Chreode: A Structured World Model for Action-Conditioned One-Step Cell-State Transitions

Motivation and Context

Large-scale in-silico perturbation screens necessitate predictive models that can efficiently infer how cellular transcriptional states evolve following genetic or pharmacological interventions. Conventional approaches for modeling such dynamics—static perturbation predictors, iterative ODE/Schrödinger-bridge solvers, or task-specific neural networks—either ignore temporal axes, lack explicit action conditioning, or entail prohibitive computational costs. Chreode addresses these deficiencies by introducing a single-pass, action- and time-conditioned transition model. It is designed around biologically motivated priors, enabling rapid amortized inference and robust transfer across developmental and perturbational regimes.

Model Architecture and Biological Prior

Chreode is conceptually framed as a cell world model—a direct analogy to RL world models—predicting the evolution of single-cell latent states under controlled interventions and elapsed time. The model encapsulates the Waddington landscape paradigm, decomposing the transition operator into three interpretable biological effects:

1. Downhill Potential Gradient: $-\nabla_z U_\theta(z, \Delta, a)$ constitutes the drift towards the attractors of the deterministic differentiation landscape.
2. Antisymmetric (Rotational) Flow: $S_\theta(z, \Delta, a) z$, where $S_\theta$ is constructed to guarantee antisymmetry, encodes non-gradient (rotational) dynamics, reflective of fluxes and cyclic processes.
3. Stochastic Spread: $\sigma_\theta(z, \Delta, a) \odot \epsilon$ models population-level variability and biological noise.

The model computes the future latent state $\hat z_{t+\Delta}$ via a single additive update, governed by a time gate $\alpha(\Delta)$ parameterized as $1 - e^{-\Delta/\tau_0}$ with learnable $\tau_0$. This ensures trivial transitions for $\Delta=0$ and biologically informed scaling for nonzero $\Delta$.

Figure 1: Chreode’s two-stage pipeline couples a pretrained scVI encoder over a cross-species ortholog gene vocabulary with a Waddington-inspired DiT backbone that predicts distributional evolution in one pass, leveraging a structured residual composed of gradient flow, rotation, and noise.

The DiT-based backbone accepts distinct time encodings for the potential and antisymmetric components—Time2Vec for smooth interpolation in the former, and bounded low-frequency Fourier features in the latter to ensure stability at large time extrapolations—thus realizing fine-grained controllability and extrapolation.

Training Paradigm and Loss Construction

Chreode is pretrained on a harmonized 2.4M-cell mouse embryonic atlas, aggregating several datasets and spanning 88 timepoints. The pretraining protocol is strictly population-level, eschewing cell-pair information, which is infeasible in destructive scRNA-seq experiments. The population-matching loss combines:

• Kernel MMD and entropic Sinkhorn $S_\theta(z, \Delta, a) z$0 distances for distributional alignment,
• A drifting-field loss (adapted from the Drifting Models framework) to ensure the model’s pushforward matches target populations at equilibrium,
• A downhill regularizer promoting consistency with the Waddington landscape prior.

A crucial aspect of the training schedule is uniform sampling over all ordered time pairs, not just maximal time differences, enabling proper identification of temporal scaling parameters.

Empirical Results and Numerical Evaluation

Chreode demonstrates pronounced empirical advantages over baselines across a range of downstream evaluations:

• Time-Transition Prediction (Weinreb, Veres datasets): Fine-tuning the pretrained backbone yields lower per-target Sinkhorn $S_\theta(z, \Delta, a) z$1 distances compared to PI-SDE, PRESCIENT, and scratch models. For example, on the Weinreb hematopoiesis task, Chreode attains $S_\theta(z, \Delta, a) z$2 of $S_\theta(z, \Delta, a) z$3 at day 4 and $S_\theta(z, \Delta, a) z$4 at day 6; the corresponding scratch and best baseline models lag by substantial margins.
• Gene-State Embedding Transfer (Norman Perturb-seq via GEARS): Substituting the gene-state embedding in the GEARS predictor with the pretrained Chreode backbone improves DE20 mean squared error from $S_\theta(z, \Delta, a) z$5 to $S_\theta(z, \Delta, a) z$6—a $S_\theta(z, \Delta, a) z$7 relative reduction—without altering the GEARS training regime, robustly isolating the effect of transferable dynamics representation.
• Zero-Shot Clonal Fate Prediction: The pretrained model, with no further tuning, produces clonal fate scores competitive with leading dynamic-OT methods (moscot, WOT, scDiffEq), substantiating the generality of the learned transition operator.

Ablation analysis rigorously confirms the necessity of each architectural and training design. The Waddington residual structure, multi-$S_\theta(z, \Delta, a) z$8 temporal sampling, and decoupled time encodings are all indispensable: simplification or improper tying of these components leads to large degradations in transfer performance.

Implications, Limitations, and Future Directions

Theoretical and practical implications are multifold:

• Amortized, Efficient Transition Modeling: Chreode decouples transition computation from integration steps, supporting fast $S_\theta(z, \Delta, a) z$9 inference suitable for high-throughput screening. Per-query latency is $S_\theta$0 lower than comparable neural-ODE baselines.
• Biology-Aligned, Structured Priors: The decomposition into gradient, rotation, and noise instantiates a mechanistically interpretable transition map reflecting empirical differentiation trajectories and dynamical landscape theory.
• Transferability: Pretraining on developmental dynamics produces state-space representations that aid not only time-based prediction but also gene perturbation tasks, suggesting that both manipulation regimes navigate coherent latent geometry.
• Modular Pretraining: The two-stage pipeline is amenable to future integration with more expressive gene-aware perturbation encoders, and its explicit ortholog mapping enables cross-species applications.

However, several limitations persist:

• Pretraining is restricted to mouse embryonic data with cross-species transfer mediated solely by ortholog mapping, without explicit adult tissue or human pretraining.
• Fine-tuning, rather than full zero-shot inference, underpins time-transition transfer results.
• The perturbation embedding is mediated by existing architectures (e.g., GEARS); Chreode does not currently instantiate a standalone perturbation operator.

Future directions include scaling the pretraining corpus, incorporating fate-aware criteria, and integrating gene-aware perturbation encoders end-to-end with the dynamics backbone.

Conclusion

Chreode sets a new paradigm for cell-state transition prediction by pretraining a transferable, biologically structured world model that supports amortized, single-pass forecasting over both developmental and perturbational axes. Its architecture, grounded in Waddington landscape theory and formalized via a DiT-based residual decomposition, realizes strong empirical performance and transfer—underscoring the suitability of structured, amortized models for scalable and interpretable in-silico biology.