Softmax with top K indices mask

[HamidL]

Hi, everyone!

I'm having issues with one thing that I wish to test. I have a model that produces attention coefficients for the directed edges in a graph, like this:

alpha = self.attention(x)
alpha = softmax(
    alpha,
    dst_path_indexes
)

The issue is that, since I have a full-mesh topology, most of my alphas tend to be values near 0 and the biggest attention coefficients have a low weight.
I have tried to use a fixed temperature coefficient to fix this, but to no avail:

alpha = self.attention(x)
alpha = softmax(
    alpha * self.temperature,
    dst_path_indexes
)

I would like to generate a binary mask so that either only the top K alpha coefficients (or those coefficients that have a value above the mean) in each index in dst_path_indexes are set to a value > 0 while the rest are set to 0 after the softmax.
Sorry if this is a stupid question, but is there an efficient way to do so?? My dst_path_indexes tensor has a shape of torch.Size([1024128]) with 8128 unique values.

Thank you for your time :-)

[rusty1s]

You can use to_dense_batch(fill_value=float('-inf')) to convert alpha into shape [batch_size, num_nodes], and then call topk(dim=1) on it. Would that work in your case?

[HamidL]

Yes, this indeed works! Once I sorted the dst_path_indexes and alpha tensors accordingly, the resulting matrix is exactly what I wanted.

Thank you very much :-)