Edge classification in GATConv

[Haoyuan-L]

Hello. I would like to do node&edge classification by GATConv. Currently, I have implemented a model to do the node classification, but I'm not sure if GATConv can be used for edge classification(is_in_path/is_not_in_path). I saw a solution in another post, and it suggests taking a dot product between nodes connected by edges and passing it to the BCE. I deployed a similar structure, but the model accuracy is relatively low. I'm not proficient in Pytorch Geometric, so any suggestions will be appreciated!

            # x is the node matrix

            src, dst = edge_index 

            # pass to several GAT layers
            x = GATConv(x=x, edge_index=edge_index, edge_attr=edge_attr)

            score = (x[src] * x[dst ] * edge_attr).sum(dim=-1)

            x = F.dropout(x, p=0.5, training=self.training)

            score = torch.sigmoid(score)
            
            return x, score

[wsad1]

Check out this example on link prediction https://github.com/pyg-team/pytorch_geometric/blob/master/examples/link_pred.py. Some things you could try if you're not already doing so is

1. add negative sampling
2. use BCEWithLogitsLoss for stability.

[Haoyuan-L]

print("edge_index_pre")
print(edge_index[:, 0:8])
print(edge_index.size())

Sorry for the confusion, I forgot to change the configurations file (set add_self_loops=False), and now I have a consistent index for edge_attr and alpha.

[Haoyuan-L]

print("edge_index_pre")
print(edge_index[:, 0:8])
print(edge_index.size())

Sorry for the confusion, I forgot to change the configurations file, and now I have a consistent index for edge_attr and alpha.

Apparently, self-loops need to be added to make GAT more effective, is there any way I can remove the attention scores of self-loops before returning them or just use indices operations to get a consistent index? @rusty1s

[rusty1s]

If add_self_loops=False, you want to apply a skip-connection after the GAT layer instead:

x = conv(x, edge_index) + lin(x)

This lets the model preserve central node information.

[Haoyuan-L]

If add_self_loops=False, you want to apply a skip-connection after the GAT layer instead:

x = conv(x, edge_index) + lin(x)

This lets the model preserve central node information.

That's very helpful! Could you explain more about lin(x)? Is this a single linear layer with the same dim_in and dim_out?

[rusty1s]

Yes, that is correct. It is known as a "learnable" skip-connection.

[Haoyuan-L]

I got the training result following the edge embedding @rusty1s mentioned above.

The performance of my model is quite poor.
The accuracy/loss of edge and node fluctuates around a value after several epochs.
(65% for node prediction, 90% for edge prediction)

Is there any way I can improve my model performance, or maybe the GAT is not suitable for this problem?

[Haoyuan-L]

Part of the forward func:

out, (edge_index, alpha) = self.GATConv1(x=x, 
                                      edge_index=edge_index, 
                                      edge_attr=edge_attr,
                                      return_attention_weights=return_attention_weights,
                                      )
out = F.dropout(out, p=0.5, training=self.training)
out = out + self.linear(x)

edge_feat = torch.cat([out_src, alpha, out_dst], dim=-1)
edge_out = self.MLP(edge_feat)

[rusty1s]

Your model looks good. I suggest you add more layers and use non-linearities as well (not sure if you already do this). You might also wanna try to lower your LR or add some form of LR decay.