Is it possible to use median as the aggregation function?

[sisaman]

I want to test the median as the GNN aggregation, but I'm not sure whether it is possible to implement it within the PyG framework. Has anyone tried to do that before?

[wsad1]

I haven't tried using median for the aggregation operation.
However I think PyG doesn't support this right now, as the aggregate function in PyG GNN models uses
the scatter function to aggregate messages from different nodes, and scatter only supports sum, mean, max and min reduction.

I haven't worked on scatter so, I can't comment on how easy this is to integrate. But I'd be really interested to hear why you feel a median aggregation would help.

[rusty1s]

See rusty1s/pytorch_scatter#227.

Overall, scatter_median is a difficult method to implement. As a workaround, you can implement the aggregate method by yourself. In case your edge_index is sorted column-wise, you can do something like:

def aggregate(self, inputs, index):
    out, mask = to_dense_batch(inputs, index)
    return out.median(dim=1)

[wsad1]

Correct me if I am wrong. But the above approach will not return the right median if the nodes have a different in degree.

index = torch.tensor([0,1,1,1]) # nodes with different in degree
x = torch.ones(4,3)
out , mask = to_dense_batch(x, index)
torch.median(out, dim=1)[0] 

Outputs tensor([[0., 0., 0.], [1., 1., 1.]]) but the output should be a tensor of 1s. Maybe something like:

def aggregate(self, inputs, index):
     out = torch.zeros_like(inputs)
     for i in index.unique():
         out[i] = inputs[index ==i].median(dim=0)[0]

it's pretty inefficient as a for loop is run every forward pass.

[rusty1s]

Oh you are right. As a slightly more efficient workaround, you can group features by same node degree:

dense_inputs, mask = to_dense_batch(inputs, index)
out = inputs.new_zeros(dense_inputs.size(0), dense_inputs.size(-1))
deg = mask.sum(dim=1)
for i in deg.unique():
    deg_mask = deg == i
    out[deg_mask] = dense_inputs[deg_mask, :i].median(dim=1)