Repository: pyg-team/pytorch_geometric
Pull request #10136 (2025-03-21 21:12:21)

Support floating-point numbers for `num_neg_samples` in negative sampling
Fixes https://github.com/pyg-team/pytorch_geometric/discussions/8899


Patch:
diff --git a/torch_geometric/utils/_negative_sampling.py b/torch_geometric/utils/_negative_sampling.py
index bb93a9385..841c8c984 100644
--- a/torch_geometric/utils/_negative_sampling.py
+++ b/torch_geometric/utils/_negative_sampling.py
@@ -12,7 +12,7 @@ from torch_geometric.utils.num_nodes import maybe_num_nodes
 def negative_sampling(
     edge_index: Tensor,
     num_nodes: Optional[Union[int, Tuple[int, int]]] = None,
-    num_neg_samples: Optional[int] = None,
+    num_neg_samples: Optional[Union[int, float]] = None,
     method: str = "sparse",
     force_undirected: bool = False,
 ) -> Tensor:
@@ -25,10 +25,12 @@ def negative_sampling(
             If given as a tuple, then :obj:`edge_index` is interpreted as a
             bipartite graph with shape :obj:`(num_src_nodes, num_dst_nodes)`.
             (default: :obj:`None`)
-        num_neg_samples (int, optional): The (approximate) number of negative
-            samples to return.
-            If set to :obj:`None`, will try to return a negative edge for every
-            positive edge. (default: :obj:`None`)
+        num_neg_samples (int or float, optional): The (approximate) number of
+            negative samples to return. If set to a floating-point value, it
+            represents the ratio of negative samples to generate based on the
+            number of positive edges. If set to :obj:`None`, will try to
+            return a negative edge for every positive edge.
+            (default: :obj:`None`)
         method (str, optional): The method to use for negative sampling,
             *i.e.* :obj:`"sparse"` or :obj:`"dense"`.
             This is a memory/runtime trade-off.
@@ -48,6 +50,11 @@ def negative_sampling(
         tensor([[3, 0, 0, 3],
                 [2, 3, 2, 1]])
 
+        >>> negative_sampling(edge_index, num_nodes=(3, 4),
+        ...                   num_neg_samples=0.5)  # 50% of positive edges
+        tensor([[0, 3],
+                [3, 0]])
+
         >>> # For bipartite graph
         >>> negative_sampling(edge_index, num_nodes=(3, 4))
         tensor([[0, 2, 2, 1],
@@ -74,6 +81,8 @@ def negative_sampling(
 
     if num_neg_samples is None:
         num_neg_samples = edge_index.size(1)
+    elif isinstance(num_neg_samples, float):
+        num_neg_samples = int(num_neg_samples * edge_index.size(1))
     if force_undirected:
         num_neg_samples = num_neg_samples // 2
 
@@ -117,7 +126,7 @@ def negative_sampling(
 def batched_negative_sampling(
     edge_index: Tensor,
     batch: Union[Tensor, Tuple[Tensor, Tensor]],
-    num_neg_samples: Optional[int] = None,
+    num_neg_samples: Optional[Union[int, float]] = None,
     method: str = "sparse",
     force_undirected: bool = False,
 ) -> Tensor:
@@ -131,9 +140,11 @@ def batched_negative_sampling(
             node to a specific example.
             If given as a tuple, then :obj:`edge_index` is interpreted as a
             bipartite graph connecting two different node types.
-        num_neg_samples (int, optional): The number of negative samples to
-            return. If set to :obj:`None`, will try to return a negative edge
-            for every positive edge. (default: :obj:`None`)
+        num_neg_samples (int or float, optional): The number of negative
+            samples to return. If set to :obj:`None`, will try to return a
+            negative edge for every positive edge. If float, it will generate
+            :obj:`num_neg_samples * num_edges` negative samples.
+            (default: :obj:`None`)
         method (str, optional): The method to use for negative sampling,
             *i.e.* :obj:`"sparse"` or :obj:`"dense"`.
             This is a memory/runtime trade-off.
@@ -157,6 +168,11 @@ def batched_negative_sampling(
         tensor([[3, 1, 3, 2, 7, 7, 6, 5],
                 [2, 0, 1, 1, 5, 6, 4, 4]])
 
+        >>> # Using float multiplier for negative samples
+        >>> batched_negative_sampling(edge_index, batch, num_neg_samples=1.5)
+        tensor([[3, 1, 3, 2, 7, 7, 6, 5, 2, 0, 1, 1],
+                [2, 0, 1, 1, 5, 6, 4, 4, 3, 2, 3, 0]])
+
         >>> # For bipartite graph
         >>> edge_index1 = torch.as_tensor([[0, 0, 1, 1], [0, 1, 2, 3]])
         >>> edge_index2 = edge_index1 + torch.tensor([[2], [4]])


Test patch:
diff --git a/test/utils/test_negative_sampling.py b/test/utils/test_negative_sampling.py
index 709452fe6..333fa5e74 100644
--- a/test/utils/test_negative_sampling.py
+++ b/test/utils/test_negative_sampling.py
@@ -77,6 +77,16 @@ def test_negative_sampling():
     assert neg_edge_index.size(1) == 2
     assert is_negative(edge_index, neg_edge_index, (4, 4), bipartite=False)
 
+    # Test with float multiplier less than 1
+    neg_edge_index = negative_sampling(edge_index, num_neg_samples=0.5)
+    assert neg_edge_index.size(1) == 2  # 50% of 4 edges = 2 edges
+    assert is_negative(edge_index, neg_edge_index, (4, 4), bipartite=False)
+
+    # Test with float multiplier greater than 1
+    neg_edge_index = negative_sampling(edge_index, num_neg_samples=1.5)
+    assert neg_edge_index.size(1) == 6  # 150% of 4 edges = 6 edges
+    assert is_negative(edge_index, neg_edge_index, (4, 4), bipartite=False)
+
     edge_index = to_undirected(edge_index)
     neg_edge_index = negative_sampling(edge_index, force_undirected=True)
     assert neg_edge_index.size(1) == edge_index.size(1) - 1
@@ -105,6 +115,16 @@ def test_batched_negative_sampling():
     neg_edge_index = batched_negative_sampling(edge_index, batch)
     assert neg_edge_index.size(1) <= edge_index.size(1)
 
+    # Test with float multiplier less than 1
+    neg_edge_index = batched_negative_sampling(edge_index, batch,
+                                               num_neg_samples=0.5)
+    assert neg_edge_index.size(1) <= 4  # 50% of 8 edges = 4 edges
+
+    # Test with float multiplier greater than 1
+    neg_edge_index = batched_negative_sampling(edge_index, batch,
+                                               num_neg_samples=1.5)
+    assert neg_edge_index.size(1) <= 12  # 150% of 8 edges = 12 edges
+
     adj = torch.zeros(8, 8, dtype=torch.bool)
     adj[edge_index[0], edge_index[1]] = True
     neg_adj = torch.zeros(8, 8, dtype=torch.bool)