ST training with and without SE embeddings

[ekinda]

Dear maintainers,

The published tutorials for state tx train use data.kwargs.embed_key=X_hvg \ so the ST model is trained using 2000 HVGs.

But in the preprint, it is mentioned that ST can be trained using a combined loss on both SE embeddings and the gene expression values, where loss in the gene expression space is downweighted by 0.1.

I searched through the code and couldn't find a way to train ST from scratch using this technique. I have some questions related to this:

1. If I use state tx train data.kwargs.embed_key=X_state is the loss only on the SE embeddings?
2. If yes, how do I implement the combined loss and the latent-to-gene decoder? output_space? There are also many decoder-related arguments which are not clearly explained.
3. The default setting of decoder_loss_weight is set to 1.0 rather than 0.1 as in the preprint, which confused me further.

In the end I want to train ST models with and without SE embeddings, and this is currently not explained in the tutorial. I would really appreciate your help in this.

[Leafaeolian]

Same. Looking at the code, it seems the nb_decoder parameter needs to be added to activate self.gene_decoder. However, the preprint states that the entire model's input is log-normalized data, but the authors only provided NB_decoder.