Open rrydbirk opened 1 year ago
Thank you for the feedback. I can reproduce similar results without using the same package versions you provided, so it's not likely a problem with your package versions.
The model performs undesirably mainly because default hyperparameters do not fit some datasets well. From my experience, the learning rate plays an important role. VeloVAE has two optimizers with different learning rates for neural network weights and ODE parameters. In this example, I changed the neural networks to 1e-4 and ODE parameter learning rate to 5e-3 and the model generated expected results(See the figure below). Here's what you can try. When you call the training function, do:
config = {
'learning_rate': 1e-4,
'learning_rate_ode': 5e-3,
'learning_rate_post': 1e-4
}
vae.train(adata,
config=config,
plot=False,
gene_plot=gene_plot,
figure_path=figure_path,
embed="umap")
Here, we have an extra learning rate 'learning_rate_post' is the neural network learning rate in the second training stage. Usually we set it to be the same as learning_rate
.
Unfortunately, the model does not always perform well with default parameters. Here are some general suggestions about tuning the model:
learning_rate_ode
and decreasing learning_rate
will improve the performance if model behaves poorly.full_vb=True
when you create a VAE object. This helps regularize the ODE rate parameters by including a prior distribution on them.reverse_gene_mode
to True when you create a VAE objcet.There are still engineering efforts we need to make regarding the best default hyperparameter. I hope this helps!
To clarify, Yichen only tested the model on single-precision GPUs. Moving to double precision on CPU likely changes when the early stopping criterion is triggered or something about the relative balance of ODE and neural net losses.
Thanks for the comments.
Not sure if you could include this in the vignette/notebook, or just label this issue as good for beginners. Otherwise, feel free to close.
I just downloaded and ran the entire notebook for pancreas using the continuous model and CPUs instead of GPU. The velocities are reversed.
Complete notebook as PDF: veloVAE_panc… - JupyterLab.pdf