How to concatenate head embeddings with corresponding tail and relation embedding efficiently

[icedpanda]

Hi,

I'm new to GNN and PyG.

I want to concatenate a head embedding with corresponding tails and relations. And then calculate its attention.

How to utilise PyG and sparse tensor to achieve this efficiently?

This is what I have so far.

# preproce edge index
edge_sets = torch.as_tensor(edge, dtype=torch.long)  # edge (head, tail, relation)
edge_idx = edge_sets[:, :2].t().contiguous()
edge_type = edge_sets[:, 2]

entity_embeddings = nn.Embedding(n_entity, e_dim)
relation_embeddings = nn.Embedding(n_relation, r_dim)

# example forward function
def forward(self, x, edge_index, edge_attr):  # x: entity_embeddings, edge_attr: relation_embeddings
    head = x[edge_index[0]]
    tail = x[edge_index[1]]
    concat_embedding = torch.cat([head, tail, edge_attr], dim=1)
    embedding = self.linear(concat_embedding)  # nn.Linear()
    
    # prepare for attention
    alpha = self.activation(self.linear2(embedding))  # nn.LeakyRelu()
    # use softmax to get a normalized alpha
    # here I am not sure how to normalize the coefficients across all the triplets that connect to HEAD_A for example
    
    # assume we have the attention
    neighbour_embeddings = torch.matmul(attention, embedding)
    
    # do aggregation
    # a simple add for now
    new_embedding = self.linear3(head + neighbour_embeddings)
    # but this would give a shape of (n_edge, n_dim) rather than (n_entity, n _dim)
    # how to make sure we have the correct shape that represents n_entity?

So basically, for each head_entity, I want to gather all the triplets that are connected to it:
For head_1:

1. get all connected triplets
   head_1_triptes = [(head_1, tail_1, relation_1), (head_1, tail_2, relation_2), (head_1, tail_3, relation_3)]
2. each triplet in head_1_triptes will be further used to calculate the attention.
3. normalized_attention_1 = attention_1 / sum_attentions for (head_1, tail_1, relation_1)
4. Multiplication on attention and then aggregate it to get a new head_1 embeddings

I can loop and filter the index, but I guess there will be a more efficient way to do this using PyG and spare tensosr.

Thanks in advance

[rusty1s]

I think this is what you want :)

def forward(self, x, edge_index, edge_attr):  # x: entity_embeddings, edge_attr: relation_embeddings
    head = x[edge_index[0]]
    tail = x[edge_index[1]]
    concat_embedding = torch.cat([head, tail, edge_attr], dim=1)
    embedding = self.linear(concat_embedding)  # nn.Linear()
    
    # prepare for attention
    alpha = self.activation(self.linear2(embedding))  # nn.LeakyRelu()
    # use softmax to get a normalized alpha
    alpha = torch_geometric.utils.softmax(alpha, edge_index[0])
    
    # assume we have the attention
    neighbour_embeddings = alpha * embedding
    
    # do aggregation
    out = torch_scatter.scatter_add(neighbour_embeddings, index=edge_index[1])

[icedpanda]

Hi @rusty1s,

I have 6370 unique heads, 13119 unique tails and a total of 19546 nodes. The total number of unique nodes in the edge_index is 19489. (57 nodes are isolated nodes that are not in the edge_index)

However, I will get a shape of (n_head, n_dim) when passing the head index (e.g. edge_index[0]) and (n_entities, n_dim) if I pass tail_index (edge_index[1) to the scatter_add.

Is it normal? I would expect (n_tails, n_dim) (13319, n_dim) when passing (edge_index[1]).

# preproce edge index
edge_sets = torch.as_tensor(edge, dtype=torch.long)  # edge (head, tail, relation)
# 6370 heads, 13119 tails, 215919 edges, 19489 nodes
edge_idx = edge_sets[:, :2].t().contiguous() 
edge_type = edge_sets[:, 2]

entity_embeddings = nn.Embedding(n_entity, e_dim) # (19546, n_dim), 57 entities are not in edge_set
relation_embeddings = nn.Embedding(n_relation, r_dim)

# example forward function
def forward(self, x, edge_index, edge_attr):  # x: entity_embeddings.weights, edge_attr: relation_embeddings
    head = x[edge_index[0]]  # (215919, n_dim)
    tail = x[edge_index[1]] # (215919, n_dim)
    concat_embedding = torch.cat([head, tail, edge_attr], dim=1)
    embedding = self.linear(concat_embedding)  # nn.Linear()
    
    # prepare for attention
    alpha = self.activation(self.linear2(embedding))  # nn.LeakyRelu()
    alpha = torch_geometric.utils.softmax(alpha, edge_index[0])  # (215919, 1)
    
    neighbour_embeddings = alpha * embedding
    
    # do aggregation
    # this would give me a shape of (n_head, n_dim),  (6370, n_dim)
    out = torch_scatter.scatter_add(neighbour_embeddings, index=edge_index[0]) 
    # this would give me a shape of (n_enties, n_dim) (19546, n_dim) rather than
    # (n_tails, n_dim) (13119, n_dim) or (n_head_tails, n_dim)(19489, n_dim)
    out = torch_scatter.scatter_add(neighbour_embeddings, index=edge_index[1]) 

So, should I append reversed head and tail to the edge_set if I wanna aggregate information from heads to tails and tails to heads? So I can apply softmax normalization and aggregation across the same indices and get embeddings for all nodes.

Thanks.

[rusty1s]

I think you would just need to specify the dim_size=x.size(0) argument of scatter_add if you group both head and tail into the same node feature matrix. Let me know if I am missing something.

[icedpanda]

Hi @rusty1s.

Sorry, this may be a dumb question. Isn't specify the im_size=x.size(0) argument of scatter_add would only aggregate tails to target and keep the same embeddings if it's not in the edge_set[0]?

Just wanna make sure I am on the right track.

I append reversed head and tail to the edge_set so it can aggregate information from both directions. And then apply softmax and aggregation across the same indices on edge_index[0] .Further use dim_size=x.size(0) so I will still get embedding for those nodes that are not in the edge_set. (n_entity, n_dim)

# preproce edge index
edge_sets = torch.as_tensor(edge, dtype=torch.long)  # edge (head, tail, relation)
# 6370 heads, 13119 tails, 215919 edges, 19489 nodes
edge_idx = edge_sets[:, :2].t().contiguous() 
edge_type = edge_sets[:, 2]
        
target_to_src = torch.flip(edge_idx, dims=[0])
new_edge_idx = torch.cat([edge_idx, target_to_src], dim=1)
edge_types = torch.cat([edge_type, edge_type], dim=0)  # may change relation here


entity_embeddings = nn.Embedding(n_entity, e_dim) # (19546, n_dim), 57 entities are not in edge_set
relation_embeddings = nn.Embedding(n_relation, r_dim)

# example forward function
def forward(self, x, edge_index, edge_attr):  # x: entity_embeddings.weights, edge_attr: relation_embeddings
    head = x[edge_index[0]]  # (215919, n_dim)
    tail = x[edge_index[1]] # (215919, n_dim)
    concat_embedding = torch.cat([head, tail, edge_attr], dim=1)
    embedding = self.linear(concat_embedding)  # nn.Linear()
    
    # prepare for attention
    alpha = self.activation(self.linear2(embedding))  # nn.LeakyRelu()
    alpha = torch_geometric.utils.softmax(alpha, edge_index[0])  # (215919, 1)
    
    neighbour_embeddings = alpha * embedding
    
    # do aggregation
    # (n_entity, n_dim)
    out = torch_scatter.scatter_add(neighbour_embeddings, index=edge_index[0], dim=0, dim_size=x.size(0)) 

Thanks

[rusty1s]

Yes, this looks good. Importantly, I think you would want to to increase the edge_type for reverse edges (but it does not seem to be used anyway).

[icedpanda]

Thanks @rusty1s. I use edge_types to get edge_attr, or you can call itrelation_embeddings.

Right now I just use torch.cat() to stack edges_types which make direction irrelevant. Are you suggesting adding new edge_types for those reverse edges so that direction matters? I am planning to create new edge_types for these reverse edges but I will see how it goes with the current one.

# pass edge_attr and entity embeddings to the model
edge_attr = relation_embeddings[edge_types]

Again, thanks for your help