Considerations of minibatch training in a hierarchical setting

[jasperhyp]

Hi! I am now trying to revise my link minibatching (#3745 ) to update only the nodes sampled by NeighborLoader. Currently, all nodes in the lower-level graph are being updated because in our model all nodes are used for "pooling" (essentially some GATs) to update a node in a higher-level graph (it's a two-layer hierarchical graph), although only batched links are used for calculating link prediction loss. In other words, we are training all nodes using batched links and considering all nodes for pooling (GAT) in each batch. We also use a second loss function that measures similarity between the embedding of the node (say, A_i) in the higher-level graph and the embeddings of nodes in the lower-level graph corresponding to A_i (let's call it simloss).

My questions are:

• Considering the ultimate embedding learning goal, the current training strategy seems to work just fine. Is it more correct to train only corresponding nodes using batched links and consider only corresponding nodes in pooling (GAT);
• If the minibatch should be revised, I assume I would simply use the nodes NeighborLoader sampled with links together (I don't notice a mask feature in GATconv or other simple conv layers), but maybe that simloss might not be that robust with changing nodes (disconnected subgraphs) each batch.

I guess my understanding of minibatch training mainly from traditional CNNs needs some refreshments here. Would greatly appreciate your suggestions.

[rusty1s]

I hope my understanding is correct. I'm not totally sure what "corresponding" nodes actually mean.

1. It looks like you use all nodes for the pooling step/the hierarchical message passing step, so the current training strategy looks correct to me. If you only want to train on a subset of nodes, this might mean that you actually approximate the global pooling step.
2. It's not clear to me how you incorporate both high-level and low-level graph into NeighborLoader. If you rely on global representations, then NeighborLoader might not be the right choice.

[jasperhyp]

Hi Matthias, Thank you for your thoughtful reply! By "corresponding nodes" I mean only the nodes NeighborLoader sampled along with its neighbors, as is the minibatch update rule in the original GraphSAGE paper ("for u \in B^k do" as the screenshot below shows).

Also, sorry for the confusions in my wordings, but I actually meant that we are only NeighborLoader-ing lower-level graph. In other words, how the minibatching goes in the lower-level graph does not affect the update in the higher-level graph, i.e., in each batch-round, all nodes (meta-nodes) in the upper-level graph are always being updated. The reason is that the higher-level graph is small and does not really need to be minibatched.

1. I guess you are right that if I only train on a subset of nodes, I would actually be approximating the global pooling. I am thinking about another update rule now. Maybe update only corresponding (i.e., NeighborLoader sampled & neighbors) nodes using batched links but consider all nodes in pooling (GAT). I am not completely sure but this seems to be more correct than either one I suggested in the original post.
2. Could you expand a bit on why using global representations might make NeighborLoader incorrect? We are exactly using some pooling mechanism to get a global representation of the lower-level graph and use that to update the node embedding in the higher-level graph.

[rusty1s]

Thanks for the clarification. It is now much clearer to me since I wasn't aware that you do not do any mini-batching on the higher-level graph. Therefore, the points from my previous reply are largely invalid. I think the argument about approximating global pooling still holds, but it might not be a big deal and might in fact heavily outweigh the scalability issues of considering all nodes in pooling.

[jasperhyp]

Thank you Matthias! I totally agree that in terms of scalability and common practice, I should be updating batched nodes only, and only with the batched links in pooling. Will try it out.