binary_cross_entropy_with_logits causes loss stuck at 0.6931 on large graphs, solved by using cosine similarity

[lmh1020lmh]

Hi, I encountered a problem when training a large-scale link prediction model using GraphSAGE in PyTorch Geometric. When using dot-product with binary_cross_entropy_with_logits, the training loss quickly plateaus at ~0.6931 and fails to decrease any further, especially on large graphs with many nodes and edges.

problem code link

Here is the problematic code:

pred = (h_src * h_dst).sum(dim=-1)
loss = F.binary_cross_entropy_with_logits(pred, batch.edge_label)

However, after switching to cosine similarity + binary_cross_entropy (after mapping to [0,1]), the loss starts to decrease and the model trains normally:

pred = F.cosine_similarity(h_src, h_dst)
pred = (pred + 1) / 2  # map cosine similarity from [-1, 1] to [0, 1]
loss= F.binary_cross_entropy(pred, batch.edge_label.float())

This change significantly improved convergence