Efficient Soft-KNN Graph Construction

[andreasbinder]

Hi all! For my research, I want to create a differentiable graph construction method. Basically what I want to achieve is a differentiable version of radius_graph or knn_graph.

One naive implementation I thought about is to just have a fully connected graph with the inverse distance (low distance = high weight). The problem I have is that I am unsure what is the most efficient way to get the distances between each combination of points.
My current approach involves creating a (NxN,2) matrix that matches a corresponding edge_list and looks like this: [[x_1,x_1],...,[x_1,x_N],...,[x_N,x_1],...,[x_N,x_N]].

That is the code:

  batch = next(iter(test_loader)) # batch_size = 4

  pdist = nn.PairwiseDistance(p=2, keepdim=True)

  x = batch.x
  y = batch.y

  x = torch.tensor(x, requires_grad=True)

  n_nodes = x.shape[0] # n_nodes = 4096

  # get dense edge indices (full connected): [[0,0],...,[0,N-1],...,[N-1,0],...,[N-1,N-1]].
  edge_index_dense = torch.tensor([[n,i] for n in range(n_nodes) for i in range(n_nodes)]).t().contiguous()       

  # bottleneck 
  stacked_x = torch.stack ([torch.cat ([x[n], x[i]]) for n in range (len (x)) for i in range (len (x))])

  stacked_x_reshaped = stacked_x.reshape((n_nodes * n_nodes, 2,3))
  x0 ,x1 = stacked_x_reshaped[:,0,:], stacked_x_reshaped[:,1,:]

  distances = pdist(x0.unsqueeze(dim=2), x1.unsqueeze(dim=2))

For now, I am wondering how I can skip the creation of stacked_x altogether. It would be even better if there is a concise, differentiable way to arrive from x at edge_index / edge_weight.

[rusty1s]

radius_graph and knn_graph are discrete operators, so gradients can only go through edge weights, distances. So what you can do is

edge_index = knn_graph(x, batch)
dist = (x[edge_index[0]] - x[edge_index[1]]).norm(dim=-1)

Let me know if this helps you.

[andreasbinder]

Thank you, this is definitely helpful! I need to test if it has the desired effect (your solution is definitely efficient since it restricts the number of numbers).
One follow-up question: if I have some edge_weights and I decide to apply ReLU(x - threshold) so that edges under some threshold have edge_weight = 0, is there some in-built function that removes links in edge_index for edges where the edge_weight = 0.
I appreciate your help :)

[rusty1s]

edge_index = edge_index[:, edge_weight <= 0]