How is edge_index used in torch_geometric.nn.LGConv

[kianyewn]

I am currently learning about graph neural networks and followed this book by Maxime Labonne on using LightGCN for book recommendation.
(I have been tinkering with the notebook and reading the docs for a few days now but I am struggling to figure out how his data preparation works for pyg and hope I can get some answers here!)

The example is to use LightGCN for recommendation, where we have user and item interactions

1. The user_id starts at index 0, and item_id starts at index 0.

# Build the adjacency matrix based on user ratings
user_ids = torch.LongTensor([user_mapping[i] for i in df['User-ID']])
item_ids = torch.LongTensor([item_mapping[i] for i in df['ISBN']])
edge_index = torch.stack((user_ids, item_ids))

2. This is the LightGCN model that is indicated (I ran the whole code and the training loss decreases as desired, so I reckon the implementation is correct)

# use LGConv from pytorch geometric
from torch_geometric.nn import LGConv

class LightGCN(nn.Module):
    def __init__(self, num_users, num_items, num_layers=4, dim_h=64):
        super().__init__()
        
        self.num_users = num_users
        self.num_items = num_items
        self.num_layers = num_layers
        self.emb_users = nn.Embedding(num_embeddings=self.num_users, embedding_dim=dim_h)
        self.emb_items = nn.Embedding(num_embeddings=self.num_items, embedding_dim=dim_h)

        self.convs = nn.ModuleList(LGConv() for _ in range(num_layers))

        nn.init.normal_(self.emb_users.weight, std=0.01)
        nn.init.normal_(self.emb_items.weight, std=0.01)

    def forward(self, edge_index):
        emb = torch.cat([self.emb_users.weight, self.emb_items.weight])
        embs = [emb]

        for conv in self.convs:
           ###############################################################
           ######### I am stuck at understanding the part here!!! ##########
            emb = conv(x=emb, edge_index=edge_index)
           ###############################################################
            embs.append(emb)

        emb_final = 1/(self.num_layers+1) * torch.mean(torch.stack(embs, dim=1), dim=1)

        emb_users_final, emb_items_final = torch.split(emb_final, [self.num_users, self.num_items])

        return emb_users_final, self.emb_users.weight, emb_items_final, self.emb_items.weight

My question:

1. How does the user and item embeddings get updated individually, when user_ids ranges from (0, num_users) and item_ids ranges from (0, num_Items), which means that the adjacency matrix isn't of shape (M+N, M+N) that we see from the LightGCN paper?
   For example, given the following edge_index = [[0,0]], where user_0 interacted with item_0, how does pytorch geometric know whether index = 0 represents the user_id or the item_id?

To elaborate:

• Inside the .forward function of the LightGCN model, the node_features are:
  x=torch.cat([self.emb_users.weight, self.emb_items.weight]), where we concatenate both the user_embedding and item_embedding to have a combined embedding of shape (num_users + num_items, dim)
• The edge_index is of shape (2, num_edges), and edge_index[0] ranges from 0 to num_users while edge_index[1] ranges from 0 to num_items
• Despite the fact that edge_index and x do not have the same index representations, LGConv()(x=x, edge_index=edge_index) returns output of shape (num_users + num_items, dim), which seems a bit magical to me.
• Basically, how is the edge_index inside LGConv() being used?

edit:
I know that there is already an example of using LightGCN, https://github.com/pyg-team/pytorch_geometric/blob/master/examples/lightgcn.py, but I would like to know why this particular implementation works, thanks!

[rusty1s]

1. Our LightGCN implementation (and the one you posted) requires as well that all node IDs are unique (i.e., user IDs should range from (0, num_users) and item IDs should range from (num_users, num_users + num_items).