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.