How to update edge embeddings?

❓ Questions & Help

Hi, I have a question about how to update node embeddings and edge embeddings simultaneously. I want to update the edge embeddings for the next update of node embeddings and the edge embedding $e_{ij}^{l+1}$ is updated with the information of $x_i^l, x_j^l, e_{ij}^l$. I think it should be finished in message method, but it should only return a Tensor. Are there some examples? Thank you.

[rusty1s]

This is correct, although I suggest to apply edge_mlp in message:

class processor(MessagePassing):

    def __init__(self):
        super(processor, self).__init__(aggr='add')
    
    ...

    def forward(self, x, edge_index, edge_attr):
        return self.propagate(edge_index, x=x, edge_attr=edge_attr), edge_attr

    def message(self, x_i, x_j, edge_attr):
        return self.edge_mlp(torch.cat([x_i, edge_attr, x_j], dim=-1)

    def update(self, inputs, x):
        return self.node_mlp(torch.cat([x, inputs], dim=-1)

[rusty1s]

That’s not possible since you can only integrate one-dimensional edge features into a sparse matrix multiplication (that act as a weighting of neighbors). In this case, this needs to be implemented like the way above.