Clarification on attention mask "order of prediction" and location of iterative selection code

[nc1m]

Hello, and thanks for the great work on scGPT.

Context:

• In "Attention mask for generative pretraining." you state that attention scores define the order of prediction: "To address this challenge, we developed a specialized attention-masking mechanism for scGPT that defines the order of prediction based on attention scores."
• Later you state it is based on prediction confidence: "In each iteration, scGPT selects the top 1/K genes from the unknown set with the highest prediction confidence to be included as known genes in the next iteration i + 1."
• I couldn't find any iterative code in the repo.

Questions:

• How is the order of prediction specified in the current implementation: is it determined by attention scores, prediction confidence, or something else?
• Where is the relevant iterative code that selects top 1/K genes by prediction confidence and updates the known/unknown sets across iterations? Please point to the file(s), function(s), and, if possible, line ranges.

Thanks!

[l024258]

Thanks for sharing the code and paper—this is very helpful. I went through the implementation in detail, and I may be missing something, but I’m not able to clearly identify how the specialized attention mask and iterative generation described in the paper are implemented in the current codebase.

To the best of my understanding:

The implementation primarily relies on src_key_padding_mask to indicate inactive token positions. While this ensures that unknown genes do not contribute to attention, it does not appear to implement a positional or rule-based attention mask as described in the paper.

I was also unable to find an implementation of the iterative unmasking procedure. Specifically, I do not see a loop that progressively selects top-1/top-K unknown genes and updates the attention mask across multiple iterations during generation.

Could you please clarify if these components are implemented elsewhere in the code, or if there are design considerations that led to a different implementation than what is described in the paper? I’d really appreciate any guidance or pointers.

Thanks again for the great work!

[ginnyintifa]

I had the same observation. I couldn't find code associated with the K iterative generation process, which I think is the most highlighted innovation of the model.
This it yet to be implemented? As far as I can infer from line 227 in the model.py script:

   def generate(
        self,
        cell_emb: Tensor,
        src: Tensor,
        values: Optional[Tensor] = None,
        src_key_padding_mask: Optional[Tensor] = None,
        gen_iters: int = 1,
        batch_labels: Optional[Tensor] = None,  # (batch,)
    ) -> Tensor:
        """
        Args:
            cell_emb(:obj:`Tensor`): shape (batch, embsize)
            src(:obj:`Tensor`): shape (batch, seq_len)
            values(:obj:`Tensor`): shape (batch, seq_len), optional
            src_key_padding_mask(:obj:`Tensor`): shape (batch, seq_len), optional
            gen_iters(:obj:`int`): number of generation iterations
            batch_labels(:obj:`Tensor`): shape (batch,), optional
        """
        # TODO: should have a tag indicate the generation mode
        # TODO: if gen_iters > 1, should have a tag indicate the current iteration`