yousefturk/cell-trajectory-steering

Pretrained Dynamics Models for Continuous-Time Cellular Trajectory Steering

Pretrained checkpoints and publication-grade scientific artifacts for "Steering Cellular Attractor Trajectories via Continuous-Time Vectorized Reinforcement Learning" (Synthica Research Group, 2026).

These models are trained on the single-gene perturbation subset of the Norman et al. 2019 combinatorial Perturb-seq dataset (111,255 single-cell transcriptomic profiles across 237 genetic perturbation conditions in K562 chronic myelogenous leukemia cells).

Computational Architecture Flow

Figure 5. High-capacity computational architecture flow of the 2.47 Billion parameter multi-model system. The pipeline integrates a ZINB-VAE generative latent manifold ($1.32\text{B}$ parameters), a Neural SDE continuous-time transition dynamics solver ($607.4\text{M}$ parameters), and a Spatial GNN tissue steering model ($549.5\text{M}$ parameters) vectorized for parallel rollout on AMD MI300X hardware.

Checkpoints Included

1. vae_best.pt: The trained Zero-Inflated Negative Binomial Variational Autoencoder (ZINB-VAE) latent manifold (1.32B parameters).
2. sde_best.pt: The trained continuous-time Neural Stochastic Differential Equation (Neural SDE) transition dynamics solver (607.4M parameters).
3. gnn_best.pt: The trained Spatial Graph Neural Network (Spatial GNN) tissue coordination model (549.5M parameters).

Learned Action Embeddings and Non-Additivity Manifold

Figure 1. Distribution of non-additivity scores across 131 two-gene combinatorial perturbation pairs in the Norman et al. 2019 dataset.

Figure 4. UMAP projection of learned 512-dimensional action embeddings showing clear topological grouping corresponding to functional gene regulatory pathways.

Metrics Summary (Held-out Double-Knockout Pairs)

Model	STF-seq	Notes
GEARS (Roohani et al. 2023)	0.000	Architectural ceiling — cannot accept intermediate states
Discrete MLP Baseline	0.343	Compounding approximation drift
Neural SDE (Ours)	0.620	81% relative improvement over MLP

Model Convergence and Dynamics

Figure 2. Validation convergence curves for the 1.32B parameter ZINB-VAE generative manifold (left) and the 607.4M parameter Neural SDE dynamics model (right) during training on the AMD MI300X hardware.

Figure 3. Predicted vs. actual sequential perturbation trajectories in the top PCA dimensions. The solid blue line traces the empirical single-cell path (control -> A -> A+B), the dashed red line shows the continuous integration path predicted by the Neural SDE, and the discrete MLP (dashed black) drifts significantly.

Usage

These weights can be loaded directly inside the cell-steering environment to initialize the dynamics transition model:

import torch
import waddington

# Load the environment
env = waddington.make(dynamics="node", device="cuda")

# The environment automatically fetches or loads local weights from:
# checkpoints/vae_best.pt, checkpoints/sde_best.pt, checkpoints/gnn_best.pt

Citation

If you use these checkpoints in your research, please cite:

@article{turk2026steering,
  title={Steering Cellular Attractor Trajectories via Continuous-Time Vectorized Reinforcement Learning},
  author={Turk, Yousef},
  journal={Synthica Research Group Preprint},
  year={2026}
}