Max-pooling with complex masks in PyTorch - pytorch

Question

Suppose I have a matrix src with shape (5, 3) and a boolean matrix adj with shape (5, 5) as follow,
src = tensor([[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8],
[ 9, 10, 11],
[12, 13, 14]])
and
adj = tensor([[1, 0, 1, 1, 0],
[0, 1, 1, 1, 0],
[1, 1, 0, 1, 1],
[1, 1, 1, 0, 0],
[0, 0, 1, 0, 1]])
We can take each row in src as one node embedding, and regard each row in adj as the indicator of which nodes are the neighborhood.
My goal is to operate a max-pooling among all neighborhood node embeddings for each node in src.
For example, as the neighborhood nodes (including itself) for the 0-th node is 0, 2, 3, thus we compute a max-pooling on [0, 1, 2], [6, 7, 8], [ 9, 10, 11] and lead an updated embedding [ 9, 10, 11] to update 0-th node in src_update.
A simple solution I wrote is
src_update = torch.zeros_like(src)
for index in range(adj.size(0)):
list_of_non_zero = adj[index].nonzero().view(-1)
mat_non_zero = torch.index_select(src, 0, list_of_non_zero)
src_update[index] = torch.sum(mat_non_zero, dim=0)
And src_update is updated as:
tensor([[ 9, 10, 11],
[ 9, 10, 11],
[12, 13, 14],
[ 6, 7, 8],
[12, 13, 14]])
Although it works, it runs very slowly and doesn't look elegant!
Any suggestions to improve it for better efficiency?
In addition, if both src and adj are appended with batches ((batch, 5, 3), (batch, 5, 5)), how to make it works?

Answer 1

I was experimenting with your code:
output = torch.zeros_like(src)
for index in range(adj.size(0)):
nz = adj[index].nonzero().view(-1)
output[index] = src.index_select(0, nz).max(0).values
The bottleneck is of course the for loop. What first comes to mind is to use some kind of scatter function. However, the main issue here is the fact that the number of neighbors can vary from row to row. This means we will be unable to construct a tensor containing the candidate nodes before max pooling.
One possible solution is to create a helper tensor similar to src where the first node would contain placeholder values (these should not get chosen by the max-pooling, i.e. we can use -inf). We can index this tensor using a tensor containing indices: compared to your method, instead of removing the zeros with torch.nonzero(), we will place an index value of 0 (referring to the placeholder row in the first position of modified-src).
In practice, here how it looks like:
For the helper tensor src_, I placed -1s as placeholder values.
>>> src_ = torch.cat((-torch.ones_like(src[:1]), src))
tensor([[-inf, -inf, -inf],
[ 0., 1., 2.],
[ 3., 4., 5.],
[ 6., 7., 8.],
[ 9., 10., 11.],
[ 12., 13., 14.]])
We can convert the adj matrix into a tensor of indices:
>>> index = torch.arange(1, adj.size(1) + 1)*adj
tensor([[1, 0, 3, 4, 0],
[0, 2, 3, 4, 0],
[1, 2, 0, 4, 5],
[1, 2, 3, 0, 0],
[0, 0, 3, 0, 5]])
For easier indexing we will flatten index, index src_ on the first axis, and reshape right after:
>>> indexed = src_[index.flatten(), :].reshape(*adj.shape, 3)
tensor([[[ 0., 1., 2.],
[-inf, -inf, -inf],
[ 6., 7., 8.],
[ 9., 10., 11.],
[-inf, -inf, -inf]],
...
[[-inf, -inf, -inf],
[-inf, -inf, -inf],
[ 6., 7., 8.],
[-inf, -inf, -inf],
[ 12., 13., 14.]]])
Finally you can max-pool:
>>> indexed.max(dim=1).values
tensor([[ 9., 10., 11.],
[ 9., 10., 11.],
[12., 13., 14.],
[ 6., 7., 8.],
[12., 13., 14.]])

Answer 2

Ivan gave a pretty smart solution. The key idea is to transform mask to index. I have tested it and wrapped it up to the function below
def mask_max_pool(embeddings, mask):
'''
Inputs:
------------------
embeddings: [B, D, E],
mask: [B, R, D], 0s and 1s, 1 indicates membership
Outputs:
------------------
max pooled embeddings: [B, R, E], the max pooled embeddings according to the membership in mask
max pooled index： [B, R, E], the max pooled index
'''
B, D, E = embeddings.shape
_, R, _ = mask.shape
# extend embedding with placeholder
embeddings_ = torch.cat([-1e6*torch.ones_like(embeddings[:, :1, :]), embeddings], dim=1)
# transform mask to index
index = torch.arange(1, D+1).view(1, 1, -1).repeat(B, R, 1) * mask# [B, R, D]
# batch indices
batch_indices = torch.arange(B).view(B, 1, 1).repeat(1, R, D)
# retrieve embeddings by index
indexed = embeddings_[batch_indices.flatten(), index.flatten(), :].view(B, R, D, E)# [B, R, D, E]
# return
return indexed.max(dim=-2)