Repository navigation
Question: Simulating gene knockout on novel datasets (zero-shot inference) & Implementation feedback #5
Description
Dear Authors,
First of all, thank you for your outstanding work on PerturbDiff! The concept and the methodology are truly inspiring.
I am currently trying to apply your pre-trained model (specifically the finetuned_tahoe100m_fixed.ckpt) to a novel, independent single-cell dataset (Colorectal Cancer data). My goal is to perform pure inference: simulating the knockout of a specific gene (e.g., TP53) on this unseen dataset, without having any actual ground-truth perturbation data or paired control cells.
My primary question is: Does the current theoretical framework and the pre-trained checkpoint support this kind of "zero-shot" simulation on a completely novel dataset?
While attempting to implement this inference pipeline, I ran into several engineering challenges. It seems that the current codebase is heavily optimized for training and reproducing benchmark metrics, which makes it quite difficult to decouple for pure, out-of-distribution inference. I would love to share a brief summary of the roadblocks I encountered, in hopes it might be helpful for future updates or an inference-only API:
Tight Coupling in DataLoader & Sampler: The dataset_core.py and sampler.py strictly expect paired "control" and "perturbed" cells to calculate metrics. Bypassing this to feed a simple .h5ad file of raw cells requires heavily modifying the dictionary mappings (e.g., grouped_num_cell, data_indices) to prevent KeyErrors and AssertionErrors.
Hardcoded Dataset Names & Metadata: The codebase heavily relies on predefined dataset names (pbmc, tahoe100m, etc.). When feeding novel data, the framework automatically assigns names like dummy_plate_9, which later causes AssertionErrors in functions like get_short_dsname and embedder mapping.
Strict Checkpoint Loading & Embedder Dimensions: When adapting the model to accept my dataset's dimensions (e.g., 2000 HVGs), modifying the nn.Linear layers causes Unexpected key(s) in load_state_dict because the checkpoint contains hardcoded dataset-specific embedders (e.g., x_embedder.pbmc.weight). This required setting strict=False to force initialization.
Shape Assertions in Diffusion Core: During the forward pass, the Transformer blocks often output a 3D tensor [Batch, 1, Dim], but the diffusion_core.py strictly asserts x_t.shape == eps.shape (expecting [Batch, Dim]). This required manual squeeze/reshape operations at the model's output to prevent runtime crashes.
I wanted to ask if you have any plans to release a simplified predict.py script for users who just want to input an .h5ad and a perturbation condition to get the simulated results.
Thank you again for your time, your amazing research, and for making this repository open-source!
Activity
KatarinaYuan commentedon Aug 23, 2026
Thanks for the detailed report. I looked into this and there are two separate points here.
First, the currently released checkpoints/code are intended for benchmark-style inference on the supported datasets and perturbation vocabularies.
They are not yet a general zero-shot
.h5ad -> perturbation -> predictionAPI for arbitrary out-of-distribution datasets. In particular,finetuned_tahoe100m.ckptis a Tahoe100M drug-perturbation checkpoint, so a gene knockout such asTP53is outside that checkpoint's perturbationspace. The Replogle checkpoint is the gene-perturbation checkpoint, but arbitrary CRC zero-shot prediction is still not something we have validated as
a supported mode.
We also found and fixed several release-sampling issues for the Tahoe checkpoint path in
cab9298:sample_*_onlyfiltering now preserves the correct selected-gene list entries and filters cached split dictionaries to active dataset names;sample_tahoe100m_only/sample_replogle_only/sample_pbmc_onlyflags now take priority when choosing the control AnnData.We verified the Tahoe release sampling path on a Trixie GPU node with the released
finetuned_tahoe100m.ckpt. A one-batch DDIM smoke run producednormal output ranges:
So the release checkpoint itself is not producing the abnormal 4000+ expression values in this path.
Second, we added an experimental inference-only API in
153e557.It supports the use case of providing control cells in an
.h5adfile and sampling a supported Tahoe100M drug condition without requiring pairedperturbed cells or benchmark labels.
CLI example:
A Colab-style notebook is available at:
The API is intentionally strict: the perturbation and cell type must exist in the checkpoint vocabulary, and genes are aligned to the checkpoint selected-gene file. The Tahoe checkpoint is a drug-perturbation checkpoint, so this is not a general gene-KO zero-shot interface. Novel .h5ad inputs should be treated as out-of-distribution inference unless they match the supported checkpoint setting.