HICAI-ZJU/KANO

Prompt Implementation - Details Question

Closed this issue · 1 comments

vthost commented

Thank you for providing all the code, it's a very interesting work!
I have three questions about the prompt implementation, just to make sure I understand it :)

  1. What is the purpose of self.cls in the prompt?
  2. For each molecule you add one all-ones (kind of dummy) FG, and this is the one you then "extract" in model.py? : cls_hiddens = torch.gather(hidden_states, 0, fg_indexs)
  3. In the same file, why do you concatenate fg_out = torch.zeros(1, self.hidden_size).cuda()? I am actually even missing why this doesn't lead to a mismatch in dimensions because you have only atom_num atoms... but I guess it's really me missing something.

Hello, thank you very much for your interest in our project!

Here are some insights into the process:

  1. fg_states stores the representations of all functional groups for each molecule in the batch. We reserve a slot for the mediator tensor at the beginning of each molecule's functional group tensors by appending a tensor of ones. The mediator vector for each molecule is then initialized using the parameters in self.cls.
  2. Following two rounds of attention, we employ the torch.gather function to retrieve the resultant mediator tensor. Note that the fg_indexsmaintains a record of the indexes for each molecule's mediator tensor within fg_states.
  3. In the next step, we map cls_hiddens to the same dimensionality as the model's hidden size by implementing a linear layer. To superimpose the mediator tensor of each molecule onto each of its constituent atoms, we replicate each molecule's mediator as many times as there are atoms within it, thereby deriving fg_hiddens. We concatenate fg_out = torch.zeros(1, self.hidden_size).cuda() with fg_hiddens. This is primarily done because the initial row of atom_hiddens also comprises a tensor of all zeroes. By doing this, we maintain consistency which is crucial for facilitating further computations.

I hope these explanations clarify the functionality of this part of the code for you! If there are any more questions or aspects you'd like to discuss, please do not hesitate to reach out at your convenience.