Message Passing: How to write a reduce function which calculate the weighted average of neighbors?

[hellomirro]

Dear all,

I have a graph with weights on edges. My question is how to write a reduce function which gets the weighted average of neighbors where weights are on edges. I only know dgl has built-in reduce function “mean”, but how to get the weighted mean?

For example, one node has 3 neighbors with each neighbor has value 1 and edge weights: 1, 2, 3. Then the weighted average would be (1x1 + 1x2 + 1x3)/(1+2+3)

Many thanks!
Bo

[mufeili]

You can use u_mul_e for the message function to multiply the source node feature by the edge weight first and then use sum for the reduce function.

[hellomirro]

Thanks Mufei. Yes, I can use u_mul_e to get src multiplied by edge weight. But for the denominator, which should be the sum of edge weights, is there a function can do that?

[hellomirro]

Hi Mufei,

I found a way to get the sum of edge weights. In the message func, I can save both “u_mul_e” result and edges.data[‘w’] in mailbox fields ‘m0’ and ‘m1’, and then pass those information to reduce func. something like:

def message_f(edges):
return {'m0': edges.src['h0']*edges.data['w'], 'm1': edges.data['w']}

def reduce_f(nodes):
return {'h1': torch.sum(nodes.mailbox['m0'], dim=1)/torch.sum(nodes.mailbox['m1'], dim=1)}

def propagate(g, ...):
g.update_all(message_f, reduce_f)
return g

Does it look correct to you?

[mufeili]

The implementation looks correct, though you may want to pre-compute the sum of the incoming-edge weights for each node if the weights are fixed.