Providing weights to NegativeSampling in LinkNeighborLoader

[aaronzberger]

Hi! I'm doing a link prediction task, but my graph is pretty highly unbalanced on node degree (some nodes have far more links than others). I'm passing the NegativeSampling class to the neg_sampling argument in the construction of a LinkNeighborLoader. I'm hypothesizing that weighting the negative sampling will lead to better performance (since, right now, I'm seeing that links between highest-degree nodes more often get predicted as always positive and ones between the lowest-degree ones always negative). However, I can't seem to figure out what the weight parameter should actually be.

It seems to need to be provided to the sample method at sample-time, since it sometimes needs to be of shape torch.Size([NUMBER OF SOURCE NODES IN THE GRAPH]) and sometimes of shape torch.Size([NUMBER OF TARGET NODES IN THE GRAPH]).

Can I provide this weight parameter at LinkNeighborLoader construction, or would I need to call the functional method at each training iteration or something else? Specifically, I imagined I'd be able to provide simply the node degrees for all the source and/or entity nodes when constructing the LinkNeighborLoader (perhaps a 1-d Tensor with size torch.Size([num_source_nodes]) or torch.Size([num_target_nodes]), or perhaps a 2-d Tensor with size torch.Size([num_source_nodes, num_target_nodes])).

Wondering if @rusty1s or someone else has a minute to help out. Thanks!

[rusty1s]

Do you mind to share a small code snippet to reproduce your error? Given that you have a node-level weight vector, you should be able to do

loader = LinkNeighborLoader(
    data,
    neg_sampling=dict(mode='binary', amount=1, weight=weight_vector),
)

[aaronzberger]

Hi, here's a reproducible snippet.

This code creates a small graph with 10 apple nodes, 20 orange nodes, and 20 apple, eats, orange links. Then, it uses RandomLinkSplit on the graph, and creates LinkNeighborLoader objects for each split, passing in NegativeSampling with a node-level weight vector. The issue is that neither the apple_eating_frequency weight vector (node-level on the apple dimension) nor the orange_eating_frequency weight vector (node-level in the orange dimension) work as arguments to weight in NegativeSampling. Both give ValueError: The 'weight' attribute in 'NegativeSampling' needs to match the number of nodes [X] (got [Y]).

Note that if you get an error about multinomial, just run it again: this is rare, but once in a while, the random sampling/shuffling makes it such that the frequency is 0.

Would appreciate any guidance here!

import logging

import torch
from torch_geometric import transforms
from torch_geometric.data import HeteroData
from torch_geometric.loader import LinkNeighborLoader
from torch_geometric.sampler import NegativeSampling


def make_basic_graph():
    graph = HeteroData()

    graph["apple"].num_nodes = 10
    graph["apple"].x = torch.zeros(10, 5)
    graph["apple"].node_id = torch.arange(10)
    graph["orange"].num_nodes = 20
    graph["orange"].x = torch.zeros(20, 5)
    graph["orange"].node_id = torch.arange(20)

    graph["apple", "eats", "orange"].edge_index = torch.tensor(
        [[0, 0, 0, 0, 1, 2, 2, 2, 3, 5, 5, 5, 5, 6, 6, 6, 7, 8, 9, 9],
         [1, 2, 9, 17, 0, 2, 4, 19, 2, 2, 8, 9, 14, 0, 2, 6, 7, 2, 10, 12]]
    )
    graph = transforms.ToUndirected()(graph)

    return graph


def random_link_split(data: HeteroData, add_negative_train_samples: bool = False) -> tuple[HeteroData, HeteroData, HeteroData]:
    logging.info("Generating data splits")
    transform = transforms.RandomLinkSplit(
        num_val=0.1,
        num_test=0.1,
        disjoint_train_ratio=0.3,
        neg_sampling_ratio=1.0,
        add_negative_train_samples=add_negative_train_samples,
        edge_types=("apple", "eats", "orange"),
        rev_edge_types=("orange", "rev_eats", "apple"),
        is_undirected=False,
    )

    train_data, val_data, test_data = transform(data)

    return train_data, val_data, test_data


def get_data_loader(data: HeteroData, num_neighbors: list[int] = [20, 10], shuffle: bool = True) -> LinkNeighborLoader:
    edge_label_index = data["apple", "eats", "orange"].edge_label_index
    edge_label = data["apple", "eats", "orange"].edge_label

    apple_eating_frequency = []
    for i in range(data["apple"].num_nodes):
        apple_eating_frequency.append((edge_label_index[1] == i).sum().item())
    apple_eating_frequency = torch.tensor(apple_eating_frequency, dtype=torch.float32)

    orange_eating_frequency = []
    for i in range(data["orange"].num_nodes):
        orange_eating_frequency.append((edge_label_index[0] == i).sum().item())
    orange_eating_frequency = torch.tensor(orange_eating_frequency, dtype=torch.float32)

    return LinkNeighborLoader(
        data=data,
        num_neighbors=num_neighbors,
        neg_sampling=NegativeSampling(
            mode="binary",
            amount=1.0,
            weight=apple_eating_frequency,
        ),
        edge_label_index=(("apple", "eats", "orange"), edge_label_index),
        edge_label=edge_label,
        batch_size=12,
        shuffle=shuffle,
    )


if __name__ == "__main__":
    graph = make_basic_graph()
    train_data, val_data, test_data = random_link_split(graph)

    train_loader = get_data_loader(train_data)
    val_loader = get_data_loader(val_data, shuffle=False)
    test_loader = get_data_loader(test_data, shuffle=False)

    for sampled_data in train_loader:
        pass

[aaronzberger]

Hi @rusty1s, wondering if you might have had a chance to look at this briefly? Thanks for your help!

[rusty1s]

Sorry for late reply. I fixed this in #9316. From PyG master, you should be able to now pass src_weight and dst_weight to NegativeSampling