I am trying to understand a graph neural network code which has implemented a weighted attention layer as follows:
class WeightedAttention(nn.Module):
"""
Weighted softmax attention layer
"""
def __init__(self, gate_nn, message_nn, num_heads=1):
"""
Inputs
----------
gate_nn: Variable(nn.Module)
"""
super(WeightedAttention, self).__init__()
self.gate_nn = gate_nn
self.message_nn = message_nn
self.pow = torch.nn.Parameter(torch.randn((1)))
def forward(self, fea, index, weights):
""" forward pass """
gate = self.gate_nn(fea)
gate = gate - scatter_max(gate, index, dim=0)[0][index]
gate = (weights ** self.pow) * gate.exp()
# gate = weights * gate.exp()
# gate = gate.exp()
gate = gate / (scatter_add(gate, index, dim=0)[index] + 1e-13)
fea = self.message_nn(fea)
out = scatter_add(gate * fea, index, dim=0)
return out
def __repr__(self):
return '{}(gate_nn={})'.format(self.__class__.__name__,
self.gate_nn)
I am unable to find anything that explains the working of the weighted attention network. Can someone explain what it is and how the code has implemented it using scatter_max and scatter_add? Also is there any mathematical formulae that can be related to this?