Trivial cluster assignments in dense_diff_pool

[mxv001]

Hello experts,
I am using dense_diff_pool to pool nodes. But the graph topology after clustering seems to be trivial and I don't know what I am doing wrong.

from numpy import random
from torch_geometric.nn import DenseGCNConv as GCNConv, dense_diff_pool

number_you_want_to_test=1
for i in range(number_you_want_to_test):
    
    #you can play with random number of nodes
    #num_nodes=random.randint(18)


    num_nodes = 6  
    reduction = 2
    node_feature_dim = 8 
    
    pool = torch.rand(node_feature_dim, num_nodes - reduction)    

    x = torch.rand(num_nodes, node_feature_dim)
    print('x=')
    print(x)

    # creating a random symmetric adjacency matrix with diagonal = 1
    adj = torch.rand(num_nodes,num_nodes).round()
    adj = torch.triu(adj, diagonal=1)
    adj = adj + adj.T
    identity = torch.eye(num_nodes)
    adj = adj + identity
    print('adj=')
    print(adj)
    
    
    # creating a selection matrix for the new clusters
    s = torch.relu(adj @ x @ pool)
    print('s=')
    print(s)

    xnew, adjnew, _, _ = dense_diff_pool(x, adj, s, mask=None)

    print('xnew=')
    print(xnew[0])

    print('adjnew=')
    print(adjnew[0])
    
    maxes = adjnew[0].argmax(1)
    maxes_matrix = torch.zeros(adjnew[0].shape).scatter(1, maxes.unsqueeze (1), 1.0)

    # The maxes are almost always the same number, like [1,1,1,1] or [0,0,0,0] or [2,2,2,2]
    # and the maxes matrix therefore also only contains one non-zero column
    # like [[1,0,0,0], [1,0,0,0], [1,0,0,0], [1,0,0,0]].  
   
    print(maxes)
    print(maxes_matrix)

My goal is to get a binary-valued adjacency matrix from the pooled matrix -- i.e. an adjacency matrix that contains entries of 0,1 determining the connectivity between the clusters of coarsened graph, based on the most probable connection. For random input, I would expect that each cluster would be connected to a random other cluster after pooling. This would to be something random like
[[1,0,0,0],
[0,0,1,0],
[0,0,1,0],
[0,0,0,1]].
Anyway I would not expect constant values.

Am I misunderstanding or misusing the dense_diff_pool? I cannot see my error here.

Thank you!

[rusty1s]

Sorry for late reply :(

Your observation is interesting. I don't have a good explanation yet, although I am wondering whether the usage of argmax is appropriate here. In the end, the output adjacency holds the sum of edge weights of all connections that point to this supernode, so it will naturally converge to a single argmax across all supernodes.

[mxv001]

No problem re late reply!

I think you are right that argmax is not the best choice. If I interpret the output of dense_diff_pool as probabilities maybe I should use a threshold to identify if two supernodes are connected.

I guess a more precise question is: how do I get the topology of the clustered graph from the output of dense_diff_pool? Is there an obvious way of getting a binary adjacency matrix from the output of dense_diff_pool?

Thanks again!

[rusty1s]

I think this is what the ASAPooling module is doing, as it can be understood as a sparse extension to DiffPool.