Model for edge regression

[101AlexMartin]

Hello,

I'm trying to solve a problem where I have to predict an output feature per edge from a set of input nodes and edge features. The number of nodes and edges varies for the different datapoints. Currently, I'm using a layer of NNConv (to take into account the edge attributes), a set of SAGEConv layers and a final linear layer with output size 1. This network predicts on the nodes, so at the end of the forward method I multiply the output of the linear layer (one output per node) in between edges in order to get an output size equal to the number of edges. For the sake of clarification, this is the trainable architecture:

(conv0): NNConv(4, 48, aggr=mean, nn=Sequential(
    (0): Linear(in_features=3, out_features=25, bias=True)
    (1): ReLU()
    (2): Linear(in_features=25, out_features=192, bias=True)
  ))
  (conv1): SAGEConv(48, 48, aggr=mean)
  (conv2): SAGEConv(48, 48, aggr=mean)
  (conv3): SAGEConv(48, 48, aggr=mean)
  (lin): Linear(in_features=48, out_features=1, bias=True)

After which I add the operations to transform the output vector of the forward pass (edge_x) to the size of the number of edges:

src, dst = edge_index
edge_x = (x[src] * x[dst]).sum(dim=-1)

However, the performance with this architecture is very poor, so I was wondering which architecture would be the recommended one in order to implement such a task.

[rusty1s]

I think you can improve your model by using a better decoder (especially integrating the edge features in the decoder part), e.g.:

src, dst = edge_index
edge_x = mlp(torch.cat([x[src], x[dst], edge_attr], dim=-1))

[101AlexMartin]

Thanks for your answer. When using mlp, do you refer to multilayer perceptron? In that case, how can you include a mlp in the forward pass, if your dataset has graphs of different sizes? In other words, in the first snippet I shared in my question, the output is independent of the number of nodes (i.e. there is one output vector with size equal to the number of nodes). However, for the second snippet, the output is dependent on the number of nodes, thus I do not see how a mlp could be fitted to the forward pass.

[rusty1s]

The node dimension is the batch dimension here, the MLP just operates on fixed-size features for every example/node.

[101AlexMartin]

True, I see your point now, thanks. I think it is indeed a good idea. Regarding the rest of the architecture of the network, what is your opinion? (I'm not very familiar with the differences among the available torch geometric layers, so I'm not sure if there is a layer more tailored to this problem than NNConv or SAGEConv)

[rusty1s]

If you have edge features, there is no real point in using any GNN layer that would ignore them. TransformerConv, GENConv or GINEConv are all good options to try for edge feature-based GNNs as well.