Learnable edge masking

[GiuseppeSerra93]

Hello,
suppose we have a custom layer/function that, given some edge_weight tensor, performs an edge sampling. Then, the resulting sampling (or rather, the subgraph) should be used in the subsequent layers to perform a classification task. Now, I have several questions:

1. Suppose the weight tensor is computed as edge_weight = (x[src] * x[dst]).sum(dim=-1), where x is the results of some previous GNN layer. In this way, the PyG batching procedure would generate a tensor containing all the edge weights of all the graphs in the current batch. Is there any efficient way to 'select' the weights corresponding to each graph to perform the sampling separately? Do I need to use the batch vector? I guess that, after that, everything should be merged again for creating the input of the subsequent layer?
2. When dealing with undirected graphs, edge_weight contains all the edges considering both directions (e.g. [0,1] and [1,0]). Consequently, edge_weight may contain different weights depending on the direction of the same edge (this is not the case in the shown example, but if we try to use a different strategy for computing the weights it may happen). Also, if we consider both directions, the sampling may return different outputs for the two directions too. Now, I would like to have a single result for each edge. Is there any to_directed function that I can use for the sampling part? I could not find it...
3. The edge masking is not part of the PyG computational graph. The easiest way to mask the graph for the subsequent layers would be something like (please, correct me if it's wrong):

        x = self.conv1(x, edge_index)
        src, dst = edge_index
        edge_weight = (x[src] * x[dst]).sum(dim=-1)
        edge_mask = self.custom(edge_weight)
        x = self.conv2(x, edge_index[:, edge_mask])
        ....

The problem is that the code above would not call the backward function of the custom layer (needed to improve the sampling), since we are tracing x during the backward step. I checked #2033, but the thing is that the edge_weight argument is not included in all the GNN layers in PyG. What would you suggest to make the network use and update the edge weights? Do I need to create a custom message-passing function to deal with a learnable edge masking during training?

Thanks for any clarification!

[rusty1s]

1. You might want to look at torch.split which helps you to split the edge weights into a list of individual edge weights for each graph. You can use degree(batch[edge_index[0]], dtype=torch.long).tolist() as split_size argument.
2. You can only apply the sampling for the "lower triangular part" of the edge indices via:

   mask = edge_index[0] < edge_index[1]
   edge_index = edge_index[mask]

   and convert the resulting graph to an undirected one (via to_undirected) after sampling is done.
3. Yes, gradients can only flow through the edge weights of an edge, not the actual index. This should be fixable by passing edge_weight[mask] to the GNN layer as well (which is, e.g., supported for GraphConv).

[GiuseppeSerra93]

Thanks for your suggestions. Then, I will start experimenting with some examples where the edge_weight argument is supported. Anyway, my last question was more about the case where the edge_weight argument cannot be used (e.g. GINConv). What would be a solution to use edge sampling in that case? Ideally, I would like to implement this strategy for any kind of GNN operator... Do you have any suggestions?

[rusty1s]

I'm afraid that without any edge weights, there is no way to back-propagate information through an edge sampling operation since we are solely operating on discrete structures.