Repository: facebookresearch/ParlAI

Issue:
Add Top-k sampling (and stochastic sampling) to beam search
Currently we only support beam search in Torch Generator Agent. Top-K sampling is a popular strategy introduced by [Fan et al., 2018](https://arxiv.org/abs/1805.04833) which should be added.

We should aim to implement this as a somewhat clean API so as to decouple it from the Torch Generator Agent nicely.


Discussion:

Resolving patch:
diff --git a/parlai/core/torch_generator_agent.py b/parlai/core/torch_generator_agent.py
--- a/parlai/core/torch_generator_agent.py
+++ b/parlai/core/torch_generator_agent.py
@@ -287,10 +287,16 @@ def add_cmdline_args(cls, argparser):
         )
         agent.add_argument(
             '--inference',
-            choices={'beam', 'greedy'},
+            choices={'beam', 'greedy', 'topk', 'nucleus'},
             default='greedy',
             help='Generation algorithm',
         )
+        agent.add_argument(
+            '--topk', type=int, default=10, help='K used in Top K sampling'
+        )
+        agent.add_argument(
+            '--topp', type=float, default=0.9, help='p used in nucleus sampling'
+        )
 
         super(TorchGeneratorAgent, cls).add_cmdline_args(argparser)
         return agent
@@ -626,6 +632,28 @@ def _treesearch_factory(self, device):
                 eos_token=self.END_IDX,
                 device=device,
             )
+        elif method == 'topk':
+            return TopKSampling(
+                self.opt['topk'],
+                beam_size,
+                min_length=self.beam_min_length,
+                min_n_best=self.beam_min_n_best,
+                padding_token=self.NULL_IDX,
+                bos_token=self.START_IDX,
+                eos_token=self.END_IDX,
+                device=device,
+            )
+        elif method == 'nucleus':
+            return NucleusSampling(
+                self.opt['topp'],
+                beam_size,
+                min_length=self.beam_min_length,
+                min_n_best=self.beam_min_n_best,
+                padding_token=self.NULL_IDX,
+                bos_token=self.START_IDX,
+                eos_token=self.END_IDX,
+                device=device,
+            )
         else:
             raise ValueError(f"Can't use inference method {method}")
 
@@ -986,3 +1014,66 @@ def select_paths(self, logprobs, prior_scores):
         tok_ids = best_idxs % voc_size
 
         return (hyp_ids, tok_ids, best_scores)
+
+
+class TopKSampling(TreeSearch):
+    """
+    Top-K sampling (Fan et al., 2018).
+
+    Samples from a truncated distribution where only the most probable K words
+    are considered at each time.
+
+    Typical values of k are 2, 10, 50.
+
+    See https://arxiv.org/abs/1805.04833 for details.
+    """
+
+    def __init__(self, k, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.k = k
+
+    def select_paths(self, logprobs, prior_scores):
+        values, indices = logprobs.topk(self.k, dim=-1)
+        probs = torch.softmax(values, dim=-1)
+        choices = torch.multinomial(probs, 1)[:, 0]
+        hyp_ids = torch.arange(logprobs.size(0)).to(logprobs.device)
+        tok_ids = indices[hyp_ids, choices]
+        scores = values[hyp_ids, choices]
+        best_scores = prior_scores.expand_as(scores) + scores
+        return (hyp_ids, tok_ids, best_scores)
+
+
+class NucleusSampling(TreeSearch):
+    """
+    Nucelus, aka top-p sampling (Holtzman et al., 2019).
+
+    Samples from a truncated distribution which covers a fixed CDF proportion
+    of the original distribution.
+
+    Typical values of p are 0.3 and 0.9.
+
+    See https://arxiv.org/abs/1904.09751 for details.
+    """
+
+    def __init__(self, p, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.p = p
+
+    def select_paths(self, logprobs, prior_scores):
+        # Unlike the other treesearch methods, we have to switch to linspace
+        # for the probabilities in order to compute the CDF.
+        probs = torch.softmax(logprobs, dim=-1)
+        sprobs, sinds = probs.sort(dim=-1, descending=True)
+        # The subtraction here is so that we always include the first word to
+        # go over p. For example, if the most probable token has a prob of 0.5, and
+        # p = 0.3, then we need still need to include that first token.
+        mask = (sprobs.cumsum(dim=-1) - sprobs[:, :1]) >= self.p
+        sprobs[mask] = 0
+        sprobs.div_(sprobs.sum(dim=-1).unsqueeze(1))
+        choices = torch.multinomial(sprobs, 1)[:, 0]
+        hyp_ids = torch.arange(logprobs.size(0)).to(logprobs.device)
+        tok_ids = sinds[hyp_ids, choices]
+        # Convert back to logspace.
+        scores = sprobs[hyp_ids, choices].log()
+        best_scores = prior_scores.expand_as(scores) + scores
+        return (hyp_ids, tok_ids, best_scores)


Test patch:
diff --git a/tests/test_transformers.py b/tests/test_transformers.py
--- a/tests/test_transformers.py
+++ b/tests/test_transformers.py
@@ -4,6 +4,8 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+"""Test many variants of transformers."""
+
 import os
 import unittest
 import parlai.core.testing_utils as testing_utils
@@ -14,6 +16,7 @@ class TestTransformerRanker(unittest.TestCase):
 
     @testing_utils.retry(ntries=3)
     def test_repeater(self):
+        """Test a simple repeat-after-me model."""
         stdout, valid, test = testing_utils.train_model(
             dict(
                 task='integration_tests:candidate',
@@ -45,6 +48,7 @@ def test_repeater(self):
         )
 
     def test_resuming(self):
+        """Test saving and resuming training."""
         with testing_utils.tempdir() as tmpdir:
             model_file = os.path.join(tmpdir, 'model')
 
@@ -86,9 +90,10 @@ def test_resuming(self):
             )
 
     def test_resuming_reduce_on_plateau(self):
-        """ Reduce on Plateau can be tricky when combined
-            with warmup. See:
-            https://github.com/facebookresearch/ParlAI/pull/1812
+        """
+        Reduce on Plateau can be tricky when combined with warmup.
+
+        See: https://github.com/facebookresearch/ParlAI/pull/1812
         """
         with testing_utils.tempdir() as tmpdir:
             model_file = os.path.join(tmpdir, 'model')
@@ -125,9 +130,7 @@ def test_resuming_reduce_on_plateau(self):
             )
 
     def test_backcomp(self):
-        """
-        Tests that the transformer ranker model files continue to work over time.
-        """
+        """Tests that the transformer ranker model files continue to work over time."""
         testing_utils.download_unittest_models()
 
         stdout, valid, test = testing_utils.eval_model(
@@ -169,6 +172,7 @@ def test_backcomp(self):
 
     @testing_utils.retry(ntries=3)
     def test_xlm(self):
+        """Test --variant xlm."""
         stdout, valid, test = testing_utils.train_model(
             dict(
                 task='integration_tests:candidate',
@@ -243,6 +247,7 @@ class TestTransformerGenerator(unittest.TestCase):
 
     @testing_utils.retry(ntries=3)
     def test_greedysearch(self):
+        """Test greedy search."""
         stdout, valid, test = testing_utils.train_model(
             dict(
                 task='integration_tests:nocandidate',
@@ -277,6 +282,7 @@ def test_greedysearch(self):
 
     @testing_utils.retry(ntries=3)
     def test_beamsearch(self):
+        """Test beamsearch."""
         stdout, valid, test = testing_utils.train_model(
             dict(
                 task='integration_tests:nocandidate',
@@ -309,10 +315,50 @@ def test_beamsearch(self):
             test['bleu'], 0.95, "test bleu = {}\nLOG:\n{}".format(test['bleu'], stdout)
         )
 
+    def test_nucleus(self):
+        """Test nucleus generation."""
+        # Nucleus is inherently stochastic, just ensure no crash.
+        testing_utils.train_model(
+            dict(
+                task='integration_tests:nocandidate',
+                model='transformer/generator',
+                optimizer='adamax',
+                learningrate=7e-3,
+                batchsize=32,
+                num_epochs=20,
+                n_layers=1,
+                n_heads=1,
+                ffn_size=32,
+                embedding_size=32,
+                inference='nucleus',
+                topp=0.3,
+                beam_size=5,
+            )
+        )
+
+    def test_topk(self):
+        """Test topk generation."""
+        # Topk is inherently stochastic, just ensure no crash.
+        testing_utils.train_model(
+            dict(
+                task='integration_tests:nocandidate',
+                model='transformer/generator',
+                optimizer='adamax',
+                learningrate=7e-3,
+                batchsize=32,
+                num_epochs=20,
+                n_layers=1,
+                n_heads=1,
+                ffn_size=32,
+                embedding_size=32,
+                inference='topk',
+                topk=5,
+                beam_size=5,
+            )
+        )
+
     def test_generator_backcomp(self):
-        """
-        Tests that the generator model files work over time.
-        """
+        """Tests that the generator model files work over time."""
         testing_utils.download_unittest_models()
 
         stdout, valid, test = testing_utils.eval_model(
@@ -379,6 +425,7 @@ def test_badinput(self):
 
     @testing_utils.retry(ntries=3)
     def test_xlm(self):
+        """Test --variant xlm."""
         stdout, valid, test = testing_utils.train_model(
             dict(
                 task='integration_tests:nocandidate',
@@ -417,6 +464,7 @@ def test_xlm(self):
 
 
 def test_learning_rate_resuming(self, args):
+    """Test learning rate resumes correctly."""
     mdl = args['model']
     with testing_utils.tempdir() as tmpdir:
         model_file = os.path.join(tmpdir, 'model')
@@ -481,10 +529,10 @@ def test_learning_rate_resuming(self, args):
 
 
 class TestLearningRateScheduler(unittest.TestCase):
-    """Test learning rate scheduler for both generative and ranking
-    transformers."""
+    """Test learning rate scheduler for both generative and ranking transformers."""
 
     def test_resuming_generator(self):
+        """Test generators resume correctly."""
         GENERATOR_ARGS = dict(
             task='integration_tests:nocandidate',
             model='transformer/generator',
@@ -502,6 +550,7 @@ def test_resuming_generator(self):
         test_learning_rate_resuming(self, GENERATOR_ARGS)
 
     def test_resuming_ranker(self):
+        """Test resuming learning rate for the ranker."""
         RANKER_ARGS = dict(
             task='integration_tests:candidate',
             model='transformer/ranker',