Replicating RGCN + DistMult for link prediction.

[sbonner0]

I am trying to use RGCN for link prediction using the DistMult interaction model as was done in the original paper and just wanted to sanity check it.

DGl have an implementation using this approach which I have based my own on and is available here: https://github.com/dmlc/dgl/blob/master/examples/pytorch/rgcn/link_predict.py

I have implemented the DistMult score function in PyG as follows:

out = model(data.edge_index, data.edge_type) # Create the embeddings

# DistMult
s = out[data.edge_index[0,:]]
r = model.conv2.weight[data.edge_type.long(), :, 0]
o = out[data.edge_index[1,:]]

pos_score = torch.sum(s * r * o , dim=1)

I also compute the score for some negative edges and then calculate the loss as follows:

labels = torch.cat((torch.ones_like(pos_score), torch.zeros_like(neg_score)),dim=0)
preds = torch.cat((pos_score, neg_score), dim=0)
loss = F.binary_cross_entropy_with_logits(preds, labels)

Where the model is defined as:

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = RGCNConv(data.num_nodes, args.hidden_units, data.num_relations)
        self.conv2 = RGCNConv(args.hidden_units, args.hidden_units, data.num_relations)

    def forward(self, edge_index, edge_type):

        x = F.relu(self.conv1(None, edge_index, edge_type))
        x = self.conv2(x, edge_index, edge_type)

        return x

Does this look roughly correct, is it indeed conv2.weight I should be using to get the relation embedding? If I were to persist with this, could it be an example I submitted as a pull request?

[rusty1s]

This looks correct to me. Instead of using model.conv2.weight, I suggest to use a distinct parameter for this (like DGL is doing, too):

rel_weight = torch.nn.Parameter(torch.Tensor(num_relations, hidden_channels))
# Initialize `rel_weight`

score = torch.sum(s * rel_weight[edge_type] * o, dim=-1)

[sbonner0]

Hey @rusty1s, thanks so much for the response. Is there any particular reason you would suggest using a new parameter matrix just for this? Are we not just duplicating the relation matrix from the final RGCN layer?

[sbonner0]

Thanks @rusty1s, just to be clear, do you mean that it would have the shape [num_relation, num_features], rather than num_nodes?

[rusty1s]

Oh sorry, my reply was meant as an answer to a different issue. I'm hiding it to avoid confusion.

I think it is more natural to treat this as its own parameter matrix, as the task of model.conv2.weight is different from the parameters computing the final score (here,model.conv2.weight is used to guide message passing, not to learn a relational embedding). In the end, both approaches should work just fine. Re-using model.conv2.weight has the advantage of utilizing a lower amount of parameters, which could reduce over-fitting. I would just test both approaches and see which works best.

[moritzblum]

@rusty1s
I have a question that probably fits well here. I'm rebuilding the original approach to link prediction on FB-237 from Schlichtkrull, which uses 5x5 as the block dimension in block decomposition. Is num_blocks=5 the right parameter for 5x5 block decomposition? Or num_blocks=100 since I use 500 dimensional embeddings.

[moritzblum]

adj_t = SparseTensor(row=train_edge_index_batch[0], col=train_edge_index_batch[1], value=train_edge_types_batch, sparse_sizes=(data.num_nodes, data.num_nodes)).t()

For some reason, this doesn't make a difference. I already tried batch processing without a decreasing in memory consumption, too.

[rusty1s]

Do you know which model part consumes more memory? It is the RGCNEncoder or the DistMultDecoder?

[moritzblum]

The memory consumption goes up during encoding, so it's the RGCNEncoder.

# encode for true edges (without negative samples) train_edge_index_batch
z = model.encode(x=x.to(device), edge_index=train_edge_index_batch.to(device), edge_type=train_edge_types_batch.to(device))

# decode for all edges (including negative samples) train_edge_index_batch_neg
p = model.decode(z, train_edge_index_batch_neg.to(device), train_edge_types_batch_neg.to(device))

I tried batch size of 1000, 30000 edges and all edges in a single batch for train_edge_index_batch. For some reason, the batch size has zero impact.

[rusty1s]

Ok, I will look into it. Thanks for digging into this.

[moritzblum]

I uploaded my code here.
The input x stays always the same, whether batch processing or not. Would x also need to be adjusted to batch processing?