How to interpret reconstruction loss increase but elbo loss decrease on trained scvi and scanvi models

Help scvi-tools
1

Hi, I am trying to use lung cancer single cell atlas (https://luca.icbi.at/) to annotate a Xenium experiment. Given Xenium is gene-panel based (~400 genes), I retrained atlas by subsetting features to 400, first training scvi model to integrate, then using semi-supervised scanvi model with cell type as label key. For both model, I trained 300 epochs. However, when examining the model metrics, I noticed elbo loss decrease but reconstruction loss increase. Here are the plots

luca_ref_404var_lungOnly.model_metrics
luca_ref_404var_lungOnly.model_metrics1500×1000 97.6 KB

I am wondering:

  1. Is my planned workflow (subset features - reference scvi training- reference scanvi training - query-reference mapping) valid?
  2. how do we interpret “reconstruction loss increase”? I am a little suspicious that the features from Xenium gene panel are not good to represent cell type. I am wondering whether this can be the reason.
  3. With this quality of training, can i use the model to do reference mapping for Xenium data?

Thank you!!!

2

  1. yes

  2. the problem where reconstruction loss increase and KL loss keep decreasing (or vice versa), is something known for cVAE’s.. and is a result of optimal parameters being optimzie “better” in the direction of one of those losses only (i.e other directions). however, the global ELBO is what you really need to look for and it does optimize to a final plateau.

Having said that, you can decide to stop the training after certain number of epochs and it may depends on your downstream task whether you prefer low reconstruction loss (a good generative model) or low KL local (likely better integration). In your case, increase in reconstruction loss means the model relays less on encoder latent variables and more on the decoder mapping (a sort of overfit).
I would suggest to play of some of the training plan parameter: lower learning rate, increase kl-warmup steps and lower max_kl_weight. Cant say apriori what will work.

  1. I think that the figures can give us idea whether we trained enough and to the optimal point and did not overfit, but as to whether to say we can use to do reference mapping that will work then, I would just check the classification quality of the scanvi reference model validation set to get an idea.