Questions about Transconv's attention design

[HelloWorldLTY]

Hi, I am trying to replace the transformconv's attention manchism with current fast attention design. However, the speed becomes slower, and my output is quite different.

The original code is:

        alpha = (query_i * key_j).sum(dim=-1) / math.sqrt(self.out_channels)
        alpha = softmax(alpha, index, ptr, size_i)
        self._alpha = alpha
        alpha = F.dropout(alpha, p=self.dropout, training=self.training)

        out = value_j
        if edge_attr is not None:
            out = out + edge_attr

        out = out * alpha.view(-1, self.heads, 1)
        return out

My design is:

        query_i = query_i.view(1, query_i.shape[1], -1,  query_i.shape[2],)
        key_j = key_j.view(1, key_j.shape[1], -1,  key_j.shape[2],)
        value_j = value_j.view(1, value_j.shape[1], -1,  value_j.shape[2],)
        
        out = self.fast_attn(query_i, key_j, value_j)
        out = out.reshape(-1, self.heads, self.out_channels)

[akihironitta]

I would suggest running the script with PyTorch profiler to see which part of the execution is taking so long:

• https://pytorch.org/docs/stable/autograd.html#profiler
• https://youtu.be/qRZrVNNe3gQ?t=719

[rusty1s]

Also answered you on Slack. I think @akihironitta is right. From looking at the code, I am not sure your formula is correct. It looks like you will compute attention over all edges rather than local neighborhoods. I am also not sure fast_attn is very optimized for batch_size=1 with a large number of tokens.