FEA: add alpha parameter for popularity sampling distribution

recbole/config/configurator.py

@@ -408,6 +408,7 @@ def _set_default_parameters(self):
         default_train_neg_sample_args = {
             "distribution": "uniform",
             "sample_num": 1,
+            "alpha": 1.0,
             "dynamic": False,
             "candidate_num": 0,
         }
@@ -500,6 +501,7 @@ def _set_train_neg_sample_args(self):
             self.final_config_dict["train_neg_sample_args"] = {
                 "distribution": "none",
                 "sample_num": "none",
+                "alpha": "none",
                 "dynamic": False,
                 "candidate_num": 0,
             }
@@ -514,6 +516,7 @@ def _set_train_neg_sample_args(self):
                 self.final_config_dict["train_neg_sample_args"] = {
                     "distribution": "none",
                     "sample_num": "none",
+                    "alpha": "none",
                     "dynamic": False,
                     "candidate_num": 0,
                 }

recbole/data/transform.py

@@ -220,7 +220,7 @@ def __call__(self, dataset, interaction):
 
 class ReorderItemSequence:
     """
-    Random crop for item sequence.
+    Reorder operation for item sequence.
     """
 
     def __init__(self, config):

recbole/data/utils.py

@@ -158,7 +158,9 @@ def data_preparation(config, dataset):
             )
         else:
             kg_sampler = KGSampler(
-                dataset, config["train_neg_sample_args"]["distribution"]
+                dataset,
+                config["train_neg_sample_args"]["distribution"],
+                config["train_neg_sample_args"]["alpha"],
             )
             train_data = get_dataloader(config, "train")(
                 config, train_dataset, train_sampler, kg_sampler, shuffle=True
@@ -278,23 +280,35 @@ def create_samplers(config, dataset, built_datasets):
     if train_neg_sample_args["distribution"] != "none":
         if not config["repeatable"]:
             sampler = Sampler(
-                phases, built_datasets, train_neg_sample_args["distribution"]
+                phases,
+                built_datasets,
+                train_neg_sample_args["distribution"],
+                train_neg_sample_args["alpha"],
             )
         else:
             sampler = RepeatableSampler(
-                phases, dataset, train_neg_sample_args["distribution"]
+                phases,
+                dataset,
+                train_neg_sample_args["distribution"],
+                train_neg_sample_args["alpha"],
             )
         train_sampler = sampler.set_phase("train")
 
     if eval_neg_sample_args["distribution"] != "none":
         if sampler is None:
             if not config["repeatable"]:
                 sampler = Sampler(
-                    phases, built_datasets, eval_neg_sample_args["distribution"]
+                    phases,
+                    built_datasets,
+                    eval_neg_sample_args["distribution"],
+                    train_neg_sample_args["alpha"],
                 )
             else:
                 sampler = RepeatableSampler(
-                    phases, dataset, eval_neg_sample_args["distribution"]
+                    phases,
+                    dataset,
+                    eval_neg_sample_args["distribution"],
+                    train_neg_sample_args["alpha"],
                 )
         else:
             sampler.set_distribution(eval_neg_sample_args["distribution"])

recbole/properties/overall.yaml

@@ -24,6 +24,7 @@ learning_rate: 0.001            # (float) Learning rate.
 train_neg_sample_args:          # (dict) Negative sampling configuration for model training.
   distribution: uniform         # (str) The distribution of negative items.
   sample_num: 1                 # (int) The sampled num of negative items.
+  alpha: 1.0                    # (float) The power of sampling probability for popularity distribution.
   dynamic: False                # (bool) Whether to use dynamic negative sampling.
   candidate_num: 0              # (int) The number of candidate negative items when dynamic negative sampling.
 eval_step: 1                    # (int) The number of training epochs before an evaluation on the valid dataset.

recbole/sampler/sampler.py

@@ -33,8 +33,9 @@ class AbstractSampler(object):
         used_ids (numpy.ndarray): The result of :meth:`get_used_ids`.
     """
 
-    def __init__(self, distribution):
+    def __init__(self, distribution, alpha):
         self.distribution = ""
+        self.alpha = alpha
         self.set_distribution(distribution)
         self.used_ids = self.get_used_ids()
 
@@ -74,15 +75,17 @@ def _build_alias_table(self):
         self.alias = self.prob.copy()
         large_q = []
         small_q = []
-
         for i in self.prob:
             self.alias[i] = -1
-            self.prob[i] = self.prob[i] / len(candidates_list) * len(self.prob)
+            self.prob[i] = self.prob[i] / len(candidates_list)
+            self.prob[i] = pow(self.prob[i], self.alpha)
+        normalize_count = sum(self.prob.values())
+        for i in self.prob:
+            self.prob[i] = self.prob[i] / normalize_count * len(self.prob)
             if self.prob[i] > 1:
                 large_q.append(i)
             elif self.prob[i] < 1:
                 small_q.append(i)
-
         while len(large_q) != 0 and len(small_q) != 0:
             l = large_q.pop(0)
             s = small_q.pop(0)
@@ -202,7 +205,7 @@ class Sampler(AbstractSampler):
         phase (str): the phase of sampler. It will not be set until :meth:`set_phase` is called.
     """
 
-    def __init__(self, phases, datasets, distribution="uniform"):
+    def __init__(self, phases, datasets, distribution="uniform", alpha=1.0):
         if not isinstance(phases, list):
             phases = [phases]
         if not isinstance(datasets, list):
@@ -221,7 +224,7 @@ def __init__(self, phases, datasets, distribution="uniform"):
         self.user_num = datasets[0].user_num
         self.item_num = datasets[0].item_num
 
-        super().__init__(distribution=distribution)
+        super().__init__(distribution=distribution, alpha=alpha)
 
     def _get_candidates_list(self):
         candidates_list = []
@@ -306,7 +309,7 @@ class KGSampler(AbstractSampler):
         distribution (str, optional): Distribution of the negative entities. Defaults to 'uniform'.
     """
 
-    def __init__(self, dataset, distribution="uniform"):
+    def __init__(self, dataset, distribution="uniform", alpha=1.0):
         self.dataset = dataset
 
         self.hid_field = dataset.head_entity_field
@@ -317,7 +320,7 @@ def __init__(self, dataset, distribution="uniform"):
         self.head_entities = set(dataset.head_entities)
         self.entity_num = dataset.entity_num
 
-        super().__init__(distribution=distribution)
+        super().__init__(distribution=distribution, alpha=alpha)
 
     def _uni_sampling(self, sample_num):
         return np.random.randint(1, self.entity_num, sample_num)
@@ -378,7 +381,7 @@ class RepeatableSampler(AbstractSampler):
         phase (str): the phase of sampler. It will not be set until :meth:`set_phase` is called.
     """
 
-    def __init__(self, phases, dataset, distribution="uniform"):
+    def __init__(self, phases, dataset, distribution="uniform", alpha=1.0):
         if not isinstance(phases, list):
             phases = [phases]
         self.phases = phases
@@ -388,7 +391,7 @@ def __init__(self, phases, dataset, distribution="uniform"):
         self.user_num = dataset.user_num
         self.item_num = dataset.item_num
 
-        super().__init__(distribution=distribution)
+        super().__init__(distribution=distribution, alpha=alpha)
 
     def _uni_sampling(self, sample_num):
         return np.random.randint(1, self.item_num, sample_num)
@@ -451,14 +454,14 @@ class SeqSampler(AbstractSampler):
         distribution (str, optional): Distribution of the negative items. Defaults to 'uniform'.
     """
 
-    def __init__(self, dataset, distribution="uniform"):
+    def __init__(self, dataset, distribution="uniform", alpha=1.0):
         self.dataset = dataset
 
         self.iid_field = dataset.iid_field
         self.user_num = dataset.user_num
         self.item_num = dataset.item_num
 
-        super().__init__(distribution=distribution)
+        super().__init__(distribution=distribution, alpha=alpha)
 
     def _uni_sampling(self, sample_num):
         return np.random.randint(1, self.item_num, sample_num)

tests/config/test_config.py

@@ -67,6 +67,7 @@ def test_default_context_settings(self):
             {
                 "distribution": "none",
                 "sample_num": "none",
+                "alpha": "none",
                 "dynamic": False,
                 "candidate_num": 0,
             },

tests/model/test_model_auto.py

@@ -48,6 +48,7 @@ def test_bpr_with_dns(self):
             "train_neg_sample_args": {
                 "distribution": "uniform",
                 "sample_num": 1,
+                "alpha": 1.0,
                 "dynamic": True,
                 "candidate_num": 2,
             },