Trouble Training GNN for Binary Node Classification Task

Question

I am using a GNN to solve a problem in which I have a query target and an undirected graph. My goal is to emit a subset of nodes in the graph (via a node-wise binary prediction) whose features sum to the target query. I figured this would be relatively easy to learn but I have been struggling quite a lot.

To simplify things to the most basic, I currently have a model comprised of 2 GATConv (hidden size 16 down to 1-dim output) layers from pytorch geometric separated by a ReLU activation. I don't use edge weights and nodes only have two features- the feature of interest, and the query target (so this second target feature is identical for all nodes in the graph). I believed this set up would be the simplest way to introduce the query but perhaps this is causing issues. I also wonder if this lack of features could be an issue as well- perhaps there isn't enough information to meaningfully differentiate nodes, though I'm not convinced why that would interfere with the simple objective.

Finally, I use the BCEWithLogitsLoss since the target subgraph is a binary mask as well as an Adam optimizer. I know this loss function is not ideal for my real goal and I do have a custom loss function ready to use but to minimize potential sources of bugs, I reverted to using the CE loss.

I have a very large dataset at my disposal but in order to diagnose the issues, I've been working with as few as 10 or even 1 graph which happens to be mirrored to force the model to overfit. When plotting the evolution of logits for this 1 graph scenario, I notice that values move consistently in a mirrored fashion, whether I train for 10 epochs or 100, which makes sense given how message passing works, and I believe eliminates the possibility of homogenization of node representations. Loss also consistently drops but it seems like the signal is not strong enough to emit anything but all 0s, perhaps because target node masks are sparse.

Any help or advice on how to proceed? I am primarily hoping to find out if I am fundamentally misunderstanding something about GNNs such that my set-up will always fail or if there are other suggested next steps for things to try. I will note I have tried the typical suggestions- modulating learning rate, diff loss functions (tried Focal Loss), deeper/wider architectures, GCNConv layers, out-of-the-box GCN model in pytorch-geometric, dataset simplification.

Answer 1

It sounds like you've tried many different approaches and have a good understanding of the problem you're trying to solve. One thing that stands out to me is the lack of node features beyond the feature of interest and the query target. You mentioned that you don't believe this should interfere with the simple objective, but it's possible that more features could help the model differentiate nodes better and make more meaningful predictions.

Although you have tried different loss functions already an idea worth trying would be with a different loss function that is more tailored to your goal of emitting a subset of nodes whose features sum to the target query. One such loss function could be a modified version of the binary cross-entropy that penalizes the model for emitting nodes whose features do not sum to the target query, rather than just penalizing incorrect binary predictions.

You could also try experimenting with different hyperparameters, such as the number of layers, hidden size, and activation functions. It's possible that a deeper or wider architecture could help the model learn more complex relationships between nodes and more accurately predict the target subgraph.

Lastly, with your large dataset, you could try using a subset of the data for training and a different subset for validation to ensure that the model is not just overfitting to a small set of graphs. It's also possible that the sparsity of the target node masks is making it difficult for the model to learn, so you could try generating synthetic masks with a higher density to see if that improves performance.