Can KGE be utilized for the initialization of node vectors for downstream tasks?

[songsong0425]

Dear PyG community,

Before starting, always thank you for your kind reply.
As in the PyG document, there are various knowledge graph embedding(KGE) models such as TransE, DistMult, and RotatE, etc.
I'd like to use the KGE model to get the initial embedding vectors of all nodes in the knowledge graph for downstream tasks (Fig a).
But I'm not sure how to do this, because most cases of KGE need the training/validation/test dataset to calculate loss (Fig b).

(Sorry for the low resolution of the hand-drawing picture.)

My concern is 'If I split the entire KG into tr/val/ts and run KGE, can I get the 'all' node embedding vectors representing overall relational information?'
In summary, is it possible to use KGE models for the initialization of node embedding vectors like node2vec or random walk?

[wsad1]

My concern is 'If I split the entire KG into tr/val/ts and run KGE, can I get the 'all' node embedding vectors representing overall relational information?'

You can split the data into train and val while training the KGE. Using the val metrics to decide when to stop training. Then you could pass your whole dataset set to the KGE to get an embedding for all your nodes.

[songsong0425]

May I ask one more question? As you suggested, I got a trained optimal KGE model (i.e., RotatE) using early stopping and passed the whole dataset to the model as below:

#homo_g: Data(edge_index=[2, 2159252], edge_type=[2159252], num_nodes=47824)

rotate_model = RotatE(
    num_nodes=homo_g.num_nodes,
    num_relations = homo_g.num_edge_types,
    hidden_channels = 50,
    margin= 9.0
).to(device)

rotate_model.load_state_dict(torch.load('.../RotatE_checkpoints/checkpoint_1000epoch_lr1e-2.pt'))
rotate_model.eval()
result = rotate_model(homo_g.edge_index[0], homo_g.edge_type, homo_g.edge_index[1])

After that, it returned the value per entity in the knowledge graph with the size of a 1-dimensional vector.

print(result)
print(result.size())
#tensor([ 1.9494,  1.0245, -0.1833,  ...,  2.4121,  2.1687, -1.1437],
       device='cuda:0', grad_fn=<RsubBackward1>)
#torch.size([2159252]) 

I'm not sure which information is contained in result. If the values in result are the embedding vectors of entities, how can I regulate their dimension? In this case, I want the 128 dimensions of node embedding for the entity resulting `torch.size([2159252, 128]).

[wsad1]

RotatE.forward returns a raw-score indicating the likelihood of an edge being present, it doesn't return the embedding vector.

No RotatE internally computes a two embedding vector for each node using RotatE.node_emb and RotatE.node_im you could use these to compute node embeddings. See the forward function here for how the embedding and the score are computed: https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/kge/rotate.html#RotatE.

[rusty1s]

The output is the score you get for each entry in the (edge_index[0], edge_type, edge_index[1]) tuple. If you want to obtain the embeddings, you can just get them via model.node_emb(node_id_to_query).

[songsong0425]

Thank you so much for kind answers and sorry for rate reply. I confirmed the parameters in RotatE, but there was ambiguous thing. What is the difference between node_emb and node_im? I found that node_im is initialized with self.node_emb_im = Embedding(num_nodes, hidden_channels, sparse=sparse) and node_emb is function for head and tail as below:

...
        head_re = self.node_emb(head_index)
        head_im = self.node_emb_im(head_index)
        tail_re = self.node_emb(tail_index)
        tail_im = self.node_emb_im(tail_index)
...

Also, what is the full name of head(tail)_re and head(tail)_im?
head(tail)_re = real part and head(tail)_im = imaginary part as in Euler's identity.

[ZionDoki]

Hello! I would like to ask, is there any advantage in using the KGE Model to obtain node embedding compared to directly using GNN? I am currently exploring whether the KGE Model can be used to solve the problem of processing large-scale graph data in graph neural networks, hoping to gain some inspiration.