How to parallelly process all logits of a batch of graphs

[vietngth]

I want to process a batch of graphs, where the graphs are stacked in a big matrix following the documentation: https://pytorch-geometric.readthedocs.io/en/latest/advanced/batching.html

The differences are:

1. My label sets have the sizes of adjacency matrices. This means each entry (i, j) is predicted by concatenating/dot products of each pair of node embeddings of node (i, j). Note that the label matrices are not symmetric.
2. For each node attributes set of a graph, I need to materialize a matrix of the adjacency matrix size, where each entry is a pair-wise concatenating/dot products of node embeddings in order to compare with the labels in 1)

To illustrate the idea better, here is a figure:

The final goal is to feed all logits at-once in the block-diagonal matrix form, and compare them with the labels using binary cross entropy loss. Could you help me with the predict() function?

class GraphModel(pl.LightningModule):
    def __init__(self, **kargs): 
        ...
        # Some GNN layers like GAT
        self.encoder = ...
        
        self.predictor= torch.nn.Sequential(
            Linear(latent_size * 2, decoder_size),
            ReLU(),
            BatchNorm1d(decoder_size),
            Linear(decoder_size, 1),
            ReLU())

    def encode(self, x, edge_attr, edge_index):
        return self.encoder(x, edge_index, edge_attr)

    def predict(self, z, batch_index):
         # TODO: compute pair-wise node embeddings and stack the matrices
         # z: latent vector

[rusty1s]

I personally think the easiest way to do this would be via stacking of labels in a separate batch dimension and utilize padding (e.g., via to_dense_batch). I think you need to do this anyway inside the model, and so there doesn't exist a good reason to not do it already during batch creation. Then, you can directly compare z and y and compute BCE loss on top, e.g., via:

    z, mask = to_dense_batch(z, batch_index)
    z_t = z.transpose(1, 2)
    # Compute a prediction for every edge using `z` and `z_t`:
    pred = ...
    return bce_loss(pred[mask], y[mask])

[vietngth]

Is it a common practice to use fully-connected layers after learning the node embeddings for classification task? Normally, I would feed the each pair-wise embeddings into some MLPs to get the final predictions. Does using these embeddings directly for the prediction task hurt the performance?

[rusty1s]

This depends on your learning task. In your case, I assume your output should be permutation-invariant. An MLP on top of pair-wise predictions will not respect this permutation-invariance.

[vietngth]

Thanks for pointing out the permutation invariance point. My current (extremely slow) implementation computes pair-wise node embeddings and calculates the loss for each graph. Instead of using the whole logit matrix, only its triangular upper part is fed into an MLP. Then, the final predictions are mirrored in the inference step to form a complete matrix. I noticed the model performance was better with MLP than using the pair-wise embeddings directly with a sigmoid function.

Does this help with the permutation-invariant point? Since each pair of nodes (i, j) and (j, i) are considered to be the same by this approach.

[rusty1s]

Not totally sure how your implementation looks like but if your Linear layer depends on the number of nodes in its input, then it will not help with permutation invariance. It might work well for your data since you overfit on certain permutations, but it will likely fail if you randomly shuffle your data. Maybe I am misunderstanding though and permutation invariance is not a requirement. In this case, everything is good :)

[vietngth]

thank you for your support, I think it's clear to me now. I'm gonna close the thread here, happy holiday!