Embeddings onto the same vector space

[farazmah]

Hi,

I have been working off one of the unsupervised bipartite examples. I was wondering if there is a way to structure the network so that the node representation for the two different nodes are on the same vector space. Currently I am able to perform cosine similarity within each node types. I would like to extract the embeddings and perform cosine similarity to find nearest neighbours across the different node types. Any help is appreciated.

Thanks in advance

class GNNEncoder(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, dropout=0.4):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)
        self.dropout = dropout
        self.lin = Linear(out_channels, out_channels)
        self.reset_parameters()
        
    def reset_parameters(self):
        self.conv1.reset_parameters()
        self.conv2.reset_parameters()
        torch.nn.init.xavier_uniform_(self.lin.weight)

    def forward(self, x, edge_index):
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.conv1(x, edge_index).relu()
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.conv2(x, edge_index).relu()
        return self.lin(x)

class EdgeDecoder(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, 1)
        self.reset_parameters()
    
    def reset_parameters(self):
        torch.nn.init.xavier_uniform_(self.lin1.weight)
        torch.nn.init.xavier_uniform_(self.lin2.weight)

    def forward(self, z_src, z_dst, edge_label_index):
        row, col = edge_label_index
        z = torch.cat([z_src[row], z_dst[col]], dim=-1)

        z = self.lin1(z).relu()
        z = self.lin2(z)
        return z.view(-1)


class Model(torch.nn.Module):
    def __init__(self, hidden_channels, data):
        super().__init__()
        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr="sum")
        self.decoder = EdgeDecoder(hidden_channels)
        
    def forward(self, x_dict, edge_index_dict, edge_label_index):
        z_dict = self.encoder(x_dict, edge_index_dict)
        return self.decoder(z_dict['user'],z_dict['item'], edge_label_index)
    
    @torch.no_grad()
    def predict(self, x_dict, edge_index_dict):
        return self.encoder(x_dict, edge_index_dict)

[mzamini92]

You can modify the model to have a single encoder that processes both node types together. This can be done by concatenating the feature tensors of the two node types along the node dimension and passing them through a shared encoder. something like this:

class Model(torch.nn.Module):
    def __init__(self, hidden_channels, data):
        super().__init__()
        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr="sum")
        self.decoder = EdgeDecoder(hidden_channels)

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        x_cat = torch.cat([x_dict['user'], x_dict['item']], dim=0)
        edge_index_cat = {
            key: torch.cat([edge_index_dict['user'], edge_index_dict['item']], dim=1)
            for key in edge_index_dict.keys()
        }
        
        z_dict = self.encoder(x_cat, edge_index_cat)

        z_user = z_dict[:x_dict['user'].size(0)]
        z_item = z_dict[x_dict['user'].size(0):]
        
        return self.decoder(z_user, z_item, edge_label_index)

    @torch.no_grad()
    def predict(self, x_dict, edge_index_dict):
        x_cat = torch.cat([x_dict['user'], x_dict['item']], dim=0)
        edge_index_cat = {
            key: torch.cat([edge_index_dict['user'], edge_index_dict['item']], dim=1)
            for key in edge_index_dict.keys()
        }
        
        z_dict = self.encoder(x_cat, edge_index_cat)

        z_user = z_dict[:x_dict['user'].size(0)]
        z_item = z_dict[x_dict['user'].size(0):]
        
        return z_user, z_item

[farazmah]

Thanks for your reply!

I am assuming this only works if the features for user and item have the same dimension? Is there a way to get around that constraint?

Thanks

[mzamini92]

Yes it assumes that the features for the "user" and "item" nodes have the same dimension. If the user and item features have different dimensions, you cannot directly concatenate them along the node dimension. However, you can still process them separately by modifying the code. One approach is to have separate encoders for the user and item features, each with its own dimensions.

self.user_encoder = GNNEncoder(user_hidden_channels, user_hidden_channels)
self.item_encoder = GNNEncoder(item_hidden_channels, item_hidden_channels)

Then, in the forward and predict methods, you can pass the user and item features separately to their respective encoders and obtain separate embeddings:

z_user = self.user_encoder(x_dict['user'], edge_index_dict['user'])
z_item = self.item_encoder(x_dict['item'], edge_index_dict['item'])

This is a rough structure of the code.

[farazmah]

If I use a different encoder would this mean that they are not on the same vector space?

[farazmah]

For any one who is interested I solved this by using a simple classifier that performs dot product rather than the edge decoder stated above.

Also linear layers to the model to handle different features dimensions for the different node types.

class Classifier(torch.nn.Module):
    def forward(self, x_user, x_item, edge_label_index):
        # Convert node embeddings to edge-level representations:
        edge_feat_user = x_user[edge_label_index[0]]
        edge_feat_item = x_item[edge_label_index[1]]
        # Apply dot-product to get a prediction per supervision edge:
        return (edge_feat_user * edge_feat_item).sum(dim=-1)


class Model(torch.nn.Module):
    def __init__(self, hidden_channels, data):
        super().__init__()
        self.lin_user = Linear(-1, hidden_channels)
        self.lin_item = Linear(-1, hidden_channels)
        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr="sum")
        self.classifier = Classifier()

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        x_dict = {
            'client': self.lin_user(x_dict['user']),
            'worker': self.lin_item(x_dict['item'])
        }

        z_dict = self.encoder(x_dict, edge_index_dict)

        pred = self.classifier(
            z_dict["user"],
            z_dict["item"],
            edge_label_index,
        )
        return pred

    @torch.no_grad()
    def predict(self, x_dict, edge_index_dict):
        x_dict = {
            'client': self.lin_user(x_dict['user']),
            'worker': self.lin_item(x_dict['item'])
        }
        return self.encoder(x_dict, edge_index_dict)