Loss not moving on large graphs

[davidireland3]

I am trying out a regression/classification type task on some subgraphs. The problem is essentially: given a score for a subgraph, try and predict it (for classification I use the scores to make classes, e.g. if score is in a certain range then assign to this class). The subgraphs are randomly sampled from a main graph and the score is calculated for each. My GNN is a 3 layer network where each layer first performs a SAGE convolution with ReLU activation followed by a k-pooling layer. My code has worked well on all the graphs I have tested so far, but now when I move onto the largest graph I am interested in, the loss barely moves.

I am confused by this as the subgraphs for this largest graph are not magnitudes bigger than the next largest graphs subgraphs, the avg number of nodes per subgraph in former is ~77k and edges ~200k, whilst the latter had an avg of ~23k nodes and ~90k edges. As I said, my code and training worked fine for the latter, but when testing on the larger graph the loss barely moves at all.

I have tried adjusting the number of hidden units, adding/removing a layer and playing with the learning rate (the loss seems to adjust more with a higher LR, the highest and most successful I've tried is 1). I've tried Adam and SGD for the optimisers, but I am now running out of ideas as to what could be causing this. The subgraph generation procedure and score assignment has not changed, so I am starting to be at a loss as to what is causing this issue.

has anyone else had any issues with this when dealing with large graphs? it seems like the only thing that has changed is the size of the graphs being trained with, so I am wondering if that is causing it...

if anyone has any suggestions on what to try then I'd be very grateful!

[rusty1s]

From a GNN perspective, the size of the graph does not really matter. Embeddings produced by SAGEConv should look similar. I personally think the issue might be related to the pooling layer. The larger your graph gets, the more information you will lose by pooling. How do you apply the Top_k pooling layer?

[davidireland3]

Thanks for the reply. You must be a mind read -- I tried removing the pooling layer and it improved. It still did not do as well as the previous data sets, but it was definitely an improvement (only for the regression loss, I am yet to test the classification).

I was using the pooling as follows:

x = torch.relu(self.gcn1(x, edge_index))
x, edge_index, _, batch, _, _ = self.first_pool_layer(x, edge_index, batch=batch)

On the previous datasets I was using a ratio of 0.8, i.e. we keep the top ceil(|V| * 0.8) nodes after pooling.

I am going to try tonight with the classifier loss to see if the improvement is also apparent here too.

e: on closer inspection, the removal of the pooling layer was not the issue. Also testing on the classifier loss confirms this, it attains the same results both with and without pooling.

[rusty1s]

Is this a node-level or graph-level tasks? I thought this is a graph-level task, right? How do you apply your global readout layer to obtain graph-level features?

[davidireland3]

(sub)graph level task -- for each subgraph there is an associated score/class that needs to be predicted.
the above was just an example of how I used to to do the pooling, apologies. The full code looked something like this

        x = torch.relu(self.gcn1(x, edge_index))
        x, edge_index, _, batch, _, _ = self.first_pool_layer(x, edge_index, batch=batch)
        summary_1_mean = global_mean_pool(x, batch)
        summary_1_max = global_max_pool(x, batch)
        summary_1 = torch.cat((summary_1_mean, summary_1_max), dim=1)

        x = torch.relu(self.gcn2(x, edge_index))
        x, edge_index, _, batch, _, _ = self.second_pool_layer(x, edge_index, batch=batch)
        summary_2_mean = global_mean_pool(x, batch)
        summary_2_max = global_max_pool(x, batch)
        summary_2 = torch.cat((summary_2_mean, summary_2_max), dim=1)

        x = torch.relu(self.gcn3(x, edge_index))
        x, edge_index, _, batch, _, _ = self.third_pool_layer(x, edge_index, batch=batch)
        summary_3_mean = global_mean_pool(x, batch)
        summary_3_max = global_max_pool(x, batch)
        summary_3 = torch.cat((summary_3_mean, summary_3_max), dim=1)

        summary = summary_1 + summary_2 + summary_3

and then summary would be passed through linear layer(s) to get the desired output.

However, I have found the issue -- I underestimated how much larger the original graph was (the graph that we take subgraphs from, the score would be relevant to this graph). The signal would be more difficult to detect due to this, and as I was using the same size dataset (i.e. same number of subgraphs) it was unable to learn properly. I've generated more subgraphs now so the dataset is larger and it appears to be performing as expected.

Interestingly, removing the pooling layers did not change the performance (I tested on some of the previous graphs I'd looked at) and it is notably faster without them.

Thanks for all the help and comments as usual!