Any PyG GNN models/layers that actually update the "edge_attr" too?

[alzaia]

Hello, I'm trying to see if there are any PyG layers out there that also update the edge_attr and do not only use it as info for the aggregation and update of the node embeddings? In most layers I experimented with so far, it seems that they learn new node embeddings, but always keep the same initial edge_attr values even in the hidden layers.

[rusty1s]

Yes, most if not all GNN layers proposed in research do not update edge features. One exception that comes to my mind is the MetaLayer in PyG. Nonetheless, it should be straightforward to allow something like this in PyG via the edge_updater functionality:

class MyEdgeConv(MessagePassing):
    def __init__(self):
        super().__init__(aggr='add')

    def forward(self, x: Tensor, edge_index: Adj) -> Tensor:
        edge_attr = self.edge_updater(edge_index, x=x)

        out = self.propagate(edge_index, edge_attr=edge_attr,
                              size=(x.size(0), x.size(0)))

        return out, edge_attr

    def edge_update(self, x_j: Tensor, x_i: Tensor) -> Tensor:
        return x_j - x_i

    def message(self, edge_attr: Tensor) -> Tensor:
        return edge_attr

[alzaia]

Thanks for the info @rusty1s . So let's say I want to be able to do edge updates on my GAT conv. layers. Should I extend the original GATConv and modify it, or should I rather create a new class (e.g. EdgeConv) like the one you put above as example and then in my overall model stack GATConv and EdgeConv layers in each block? Or should I rather go with the MetaLayer where I set the standard GATConv as my node updater model and use the EdgeConv as my edge updater module?

[rusty1s]

This is tricky since GATConv with edge_attr only uses them to compute attention scores, and it does not return new edge features. What you can do is to modify GATConv by letting edge_update return a tuple of attention weights and new edge features. The returned edge features can then be returned in addition.

[alzaia]

@rusty1s That's basically what I did so far. I updated the returns of the edge_update method of the GATConv:
return alpha, edge_attr.view(-1, self.heads * self.out_channels)
And then also return them in the edge_updater:
alpha, edge_attr = self.edge_updater(edge_index, alpha=alpha, edge_attr=edge_attr)
It does work like this, but I guess the Linear layer for the edge features may be pretty basic here and I may need to try other things?

[rusty1s]

Yes, I think you should define a separate Linear layer or MLP for this, not re-use the one that is responsible for computing attention weights.