Gradients on edge in sparse mode

[kavehhassani]

Hello,

I am trying to implement a GNN that takes in a graph and learns to add or drop edges w.r.t a down-stream task. I am using pyro for differentiable sampling. Now the problem is that when I implement it in a sparse mode, the gradients wouldn't pass the edge_index as it is a tensor of indices.

`class EdgeAugment(nn.Module):
def init(self, in_dim, hid_dim, layers):
super(EdgeAugment, self).init()
self.encoder = GNN(in_dim, hid_dim, layers)
self.head = nn.Sequential(
nn.BatchNorm1d(2 * hid_dim + 1),
nn.Linear(2 * hid_dim + 1, 1),
nn.Sigmoid()
)

def forward(self, batch, temperature):
	h, __ = self.encoder(batch.x, batch.edge_index, batch.batch)
	p_edges = batch.edge_index
	n_edges = batched_negative_sampling(batch, batch.edge_index, batch.batch, p_edges.size(1))
	pos = torch.cat((h[p_edges[0]], h[p_edges[1]], torch.ones((p_edges.size(1), 1), device=batch.x.device)), dim=-1)
	neg = torch.cat((h[n_edges[0]], h[n_edges[1]], torch.zeros((n_edges.size(1), 1), device=batch.x.device)), dim=-1)
	edge_prob = self.head(torch.cat([pos, neg], dim=0))
	edge_sample = RelaxedBernoulliStraightThrough(temperature, probs=edge_prob).rsample().squeeze().bool()
	batch.edge_index = torch.cat((p_edges, n_edges), dim=1)[:, edge_sample]
	edge_prob = self.head(pos)
	edge_sample = RelaxedBernoulliStraightThrough(temperature, probs=edge_prob).rsample().squeeze().bool()
	batch.edge_index = p_edges[:, edge_sample]
	return batch`

For this purpose, I reused the dense GNN layers which well works perfectly:

`
class EdgeAugment(nn.Module):
def init(self, in_dim, hid_dim, layers):
super(EdgeAugment, self).init()
self.encoder = GNN(in_dim, hid_dim, layers)
self.head = nn.Sequential(
nn.Linear(2 * hid_dim + 1, hid_dim),
nn.PReLU(),
nn.Linear(hid_dim, 1),
nn.Sigmoid()
)

def forward(self, x, a, temperature):
    h, __ = self.encoder(x, a)
    n_batch, n_node, __ = x.size()
    h = torch.cat([h.repeat_interleave(n_node, dim=1), h.repeat(1, n_node, 1)], -1).view(n_batch, n_node, n_node, -1)
    h = torch.cat((h, a.unsqueeze(-1)), -1)
    edge_prob = self.head(h)
    mask = RelaxedBernoulliStraightThrough(temperature, probs=edge_prob).rsample().squeeze()
    a = torch.ones_like(a) * mask
    return x, a

`

I was wondering if there is a way to get this work in sparse mode as the dense mode is super inefficient?

Thanks

[rusty1s]

While you cannot propagate gradients through edge_index, you can propagate them trough an additional edge_weight vector that denotes the strengths of each edge. I'm not familar with pyro, but something like this hopefully works:

edge_sample = RelaxedBernoulliStraightThrough(temperature, probs=edge_prob).rsample().squeeze().bool()
batch.edge_index = p_edges[:, edge_sample]
batch.edge_weight = edge_prob[edge_sample]
conv(batch.x, batch.edge_index, batch.edge_weight)

[kavehhassani]

I see, so basically even there is still no gradients through batch.edge_index = p_edges[:, edge_sample] but you get them pass through batch.edge_weight = edge_prob[edge_sample]. That makes sense. Will try and update it here. Thanks for the quick reply.

[kavehhassani]

Ok, I tried using the edge_weight but somehow I am getting shape mismatch error when I pass the edge_weights to GCN. Size of the edge_weigh is set to edge_index.size(1) and the mismatch happens within the message function: edge_weight.view(-1, 1) * x_j. Thanks

[rusty1s]

Can you show me the shapes of edge_index and edge_weight when passed to the GNN?