what's the "loss_gen" during pre-training stage

[lishensuo]

I'm currently studying your model and find it fascinating. I've learned a lot, but I have a question about the pre-training stage.

In the example pre-training script dev-temp/examples/pretrain.py, the model calculates various types of loss from different perspectives. From my understanding:

• loss_mse: This represents the gene-prompt task for generative prediction.

  loss = loss_mse = criterion(
      output_values, target_values, positions_to_match
  )

https://github.com/bowang-lab/scGPT/blob/4068d67caaac1e28d56964da68e0214817e38428/examples/pretrain.py#L886C1-L890C43

• loss_mvc: This represents the cell-prompt task for generative prediction.

  if MVC:
      loss_mvc = criterion(
          output_dict["mvc_output"], target_values, positions_to_match
      )
      loss = loss + loss_mvc

https://github.com/bowang-lab/scGPT/blob/4068d67caaac1e28d56964da68e0214817e38428/examples/pretrain.py#L886C1-L890C43

However, I'm unclear about loss_gen. I see it is added to the total loss, but I'm unsure of its specific role. How does loss_gen differ from loss_mse? I've noticed differences, such as it being calculated after 1000 iterations and the output_dict["cell_emb"] gradient being detached. Could you clarify its purpose?

if USE_GENERATIVE_TRAINING and global_iter > 1000:
    previous_cell_embs = output_dict["cell_emb"].detach()
    preds = model(
        pcpt_gene,
        pcpt_expr,
        pcpt_key_padding_mask,
        gen_gene,
        gen_key_padding_mask,
        CLS=False,
        MVC=False,
        input_cell_emb=previous_cell_embs,
        generative_training=True,
    )["gen_preds"]
    loss_gen = criterion(preds, gen_expr_target, positions_to_match)
    loss = loss + loss_gen

https://github.com/bowang-lab/scGPT/blob/4068d67caaac1e28d56964da68e0214817e38428/examples/pretrain.py#L894C1-L908C39

Thank you for your help!