Unable to reproduce Table 4 Replogle results using released finetuned_replogle.ckpt

Thanks for releasing the code and checkpoints!

I'm trying to reproduce the Replogle row of Table 4 (PerturbDiff Finetuned) using the released preprocessed data and checkpoint, running inference only (no retraining). I'm getting results very different from the paper and would appreciate some guidance.

Setup

• Checkpoint: finetuned_replogle.ckpt from katarinayuan/PerturbDiff_release_ckpt
• Data: katarinayuan/PerturbDiff_data, Replogle processed data
• Gene vocabulary:merged_pbmc_tahoe_rep_cellxgene_genes_mapped.pkl (12,626 genes, 12626 model mode). I confirmed the Replogle shared-gene ratio is 5,760/12,626 = 45.6%, which matches Table 3.
• Evaluation: Cell-Eval v0.6.6
• Git commit: f4e27c155be5325418c4cb3182453d4022754e91 (origin/main, 2026-04-07)
• Seed / devices: seed 42 (repo default via optimization.seed, not overridden); single GPU (trainer.devices=[0])
• Hardware: NVIDIA GB10 (aarch64), CUDA 13.0, PyTorch cu130, Python 3.10

** Command: **

python ./src/apps/run/rawdata_diffusion_sampling.py \
  run_name=replogle_finetuned_full \
  model_checkpoint_path=<path_to>/finetuned_replogle.ckpt \
  trainer.use_distributed_sampler=false \
  trainer.devices=[0] \
  data.normalize_counts=10 \
  path=trixie_path \
  cov_encoding=trixie_onehot \
  cov_encoding.batch_encoding=onehot \
  cov_encoding.celltype_encoding=llm \
  cov_encoding.replogle_gene_encoding=genept \
  model.p_drop_control=0 \
  data.keep_control_cell=false \
  sampling.use_ddim=true \
  sampling.num_sampled_batches=null \
  data=replogle_finetune \
  data.sample_replogle_only=true \
  data.selected_gene_file=<path_to>/merged_pbmc_tahoe_rep_cellxgene_genes_mapped.pkl \
  data.pad_length=12626 \
  model.hidden_num=[12626,512] \
  model.input_dim=12626 \
  data.embed_key=X \
  optimization.micro_batch_size=128 \
  data.use_cell_set=32 \
  optimization.optimizer.lr=0.002

** Note on checkpoint loading: **
the checkpoint's baked-in hyper_parameters reference the original training cluster's absolute paths (e.g. /projects/AI4D/core-132/...), which don't resolve on a different machine. I had to route checkpoint loading through the repo's own load_plmodel_checkpoint() (which already supports runtime path overrides) instead of a plain PlModel.load_from_checkpoint(...). Flagging in case it's relevant to reproducibility for others as well.

Result

Running sampling directly from the released checkpoint, I get Overall R² = -11.24, which is far from the reported 0.988.

Looking at the raw predictions, some of the predicted expression values are abnormally large (max ≈ 4000+), concentrated in a few perturbations (e.g. hepg2 DNAJA1), whereas the ground-truth values look normal (max ≈ 6.8).

Question

Running the released checkpoint as-is gives R² = -11.24 instead of the 0.988 reported in Table 4, so something in my setup clearly differs from yours. Could you help me figure out what's going wrong? In particular, has the released (refactored) code + checkpoint been verified to actually reproduce the Table 4 Replogle numbers on your side?

Any pointers on where to look first would be greatly appreciated. Thanks again for the great work!

Bắt đầu với src/apps/run/rawdata_diffusion_sampling.py và load_plmodel_checkpoint() của repo, kiểm tra cách các ghi đè đường dẫn trong thời gian chạy tương tác với hyper_parameters của checkpoint. So sánh checkpoint đã phát hành, bước tiền xử lý Replogle, từ vựng gene và các tùy chọn sampling với thiết lập được báo cáo trong Table 4. Hoàn tất khi xác định được điểm không khớp trong thiết lập hoặc tái tạo được Overall R² được báo cáo là 0.988.