Hi, thank you for releasing the KANO code.
I noticed a possible issue in the implementation of the functional prompt generator. In chemprop/models/model.py, Prompt_generator.forward() receives fg_states with shape [13 * batch_size, 133], where each molecule contributes 13 functional-group slots.
The current code applies self-attention directly on this flattened tensor:
hidden_states = self.attention_layer_1(fg_states, fg_states)
hidden_states = self.attention_layer_2(hidden_states, fg_states)
Inside AttentionLayer, this produces an attention matrix of shape:
[13 * batch_size, 13 * batch_size]
As far as I understand, this allows functional groups from different molecules in the same batch to attend to each other. The mask seems to only remove zero-padding slots, but does not block interactions across different molecules.
According to the paper, the functional prompt should be generated from the functional groups of each individual molecule, so I would expect the attention to be applied within each molecule only, e.g. by reshaping:
[13 * batch_size, 133] -> [batch_size, 13, 133]
so that the attention matrix becomes:
[batch_size, 13, 13]
Although this may not have a large impact on the reported results, I think correcting it could potentially improve performance, and more importantly, it would make the implementation more consistent with the intended logic of the method.
Could you confirm whether the cross-molecule attention in the current implementation is intended, or whether the functional prompt attention should be restricted to functional groups within the same molecule?
Thanks!
Hi, thank you for releasing the KANO code.
I noticed a possible issue in the implementation of the functional prompt generator. In chemprop/models/model.py, Prompt_generator.forward() receives fg_states with shape [13 * batch_size, 133], where each molecule contributes 13 functional-group slots.
The current code applies self-attention directly on this flattened tensor:
hidden_states = self.attention_layer_1(fg_states, fg_states)
hidden_states = self.attention_layer_2(hidden_states, fg_states)
Inside AttentionLayer, this produces an attention matrix of shape:
[13 * batch_size, 13 * batch_size]
As far as I understand, this allows functional groups from different molecules in the same batch to attend to each other. The mask seems to only remove zero-padding slots, but does not block interactions across different molecules.
According to the paper, the functional prompt should be generated from the functional groups of each individual molecule, so I would expect the attention to be applied within each molecule only, e.g. by reshaping:
[13 * batch_size, 133] -> [batch_size, 13, 133]
so that the attention matrix becomes:
[batch_size, 13, 13]
Although this may not have a large impact on the reported results, I think correcting it could potentially improve performance, and more importantly, it would make the implementation more consistent with the intended logic of the method.
Could you confirm whether the cross-molecule attention in the current implementation is intended, or whether the functional prompt attention should be restricted to functional groups within the same molecule?
Thanks!