Scooby training question

[wangyb97]

Hi @johahi,

I am trying to train and evaluate SCOOBY, but I am seeing an unusual issue and wanted to ask for your advice.

I trained the model using the default initial config (with the cell neighbours all set to 0) from the GitHub repository on a machine with 2× Quadro RTX 6000 GPUs (each with 24GB). During training, the validation Pearson metrics briefly spike and then collapse close to 0.

I then tried to evaluate the model using the Evaluate_Model.ipynb notebook you suggested. However, the evaluation Pearson is also NaN.

Could you please give me some advice? Thanks a lot for your help!

Best,
Yubo

[johahi]

Sorry for the late reply, it was sent from the wrong email so github did not register. Here’s my answer I originally sent:
Hi Yubo,

Thanks for your interest!
This sounds to me like your batch size is 2, in which case the learning rate is probably too high. Can you either get more gpus, increase the batch size or try a lower lr? You can also set the “transform_borzoi_emb to false for more a bit more stability.

Best,
Johannes

[MarwaneBourdim]

Hey @johahi,

Thank you for this tool.

I've got a similar problem training the rna_only model on 8 H100 GPUs using the default config (with scVI rna_only embedding). I've tried fiddling with the learning rate (1e-4, 2e-4, 5e-5), the LR schedule (both the shipped num_steps and one divided by num_processes), batch size, and tried both transform_borzoi_emb True and False, but the problem seems to be persistent. I find it weird since I should be using the same effective batch size as you.

I've also managed to get an atac_only variant work with scVI rna_only embedding scooby variant to work by using the default config on the 100NN averaged targets but it collapses in a similar way with the single cell targets "no neighbor" configuration. I've added a fraction of zeroes metric to monitor when the model collapses during training as you can see on the wandb plots.

One thing I noticed, the shipped config_rna_only.yaml points to embedding_no_val_genes_new.pq (the MultiVI embedding), not the RNA-only scVI embedding that the paper Methods describe. When I train with the MultiVI embedding, early training looks slightly different (shown below), more stable, but I haven't confirmed whether it eventually stabilizes. Could just be a template mistake in the shipped config, but figured I'd flag it.

rna_only with scVI embedding:

Start of a training run for rna_only with MultiVI embedding:

Would you be willing to share the exact configuration that produced the uploaded neurips-scooby-rna checkpoint?

Any hint would really help unblock me. Thanks!

[johahi]

Hi Marwane,
We did use the rna-only embedding for the RNA-only model, we only cleaned up this repository (and the config) afterwards, so this should indeed point to embedding_no_val_genes_rna_only.pq. There are 2 learning rates to be set, and you're setting the one for the LoRA only most likely. Could you please try to decrease the cellstate_lr (see here) ?
Can you print the min/max values of your embeddings? Which version of borzoi_pytorch are you using?
Best,
Johannes