I have an undirected weighted graph where the edge weights represent probabilities. The majority of the edge weights are 1 (which are 7 times more frequent than the second major group of weights). I'm using this graph to learn node embeddings for an edge weight prediction task, but the model (with MSE as the loss function) performs poorly on edges with weight values other than 1.
I've come across some strategies to handle this imbalance, such as:
- Using under/oversampling techniques.
- Implementing weighted MSE loss.
- I also considered dividing the prediction task into two parts: first predicting whether the weight is 1 or not, and then predicting the exact weight for non-1 edges to get more precise predictions.
However, I'm uncertain whether these approaches are the right direction or if I’m "cherry-picking" by applying these techniques. I want to make sure I'm improving my model's generalization ability and not overfitting to specific cases. In addition to that, I'm concerned about whether this imbalance in edge weights could also be impacting the quality of the node embeddings themselves. Should I be addressing this imbalance during the node embedding process as well, or is it primarily a concern for the edge weight prediction task?
I'll appreciate any insights or suggestions! Thanks!