chapter_recurrent-modern/beam-search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 0
   },
   "source": [
    "# Beam Search\n",
    ":label:`sec_beam-search`\n",
    "\n",
    "In :numref:`sec_seq2seq`,\n",
    "we predicted the output sequence token by token\n",
    "until the special end-of-sequence \"&lt;eos&gt;\" token\n",
    "is predicted.\n",
    "In this section,\n",
    "we will begin with formalizing this *greedy search* strategy\n",
    "and exploring issues with it,\n",
    "then compare this strategy with other alternatives:\n",
    "*exhaustive search* and *beam search*.\n",
    "\n",
    "Before a formal introduction to greedy search,\n",
    "let us formalize the search problem\n",
    "using\n",
    "the same mathematical notation from :numref:`sec_seq2seq`.\n",
    "At any time step $t'$, \n",
    "the probability of the decoder output $y_{t'}$ \n",
    "is conditional \n",
    "on the output subsequence\n",
    "$y_1, \\ldots, y_{t'-1}$ before $t'$ and \n",
    "the context variable $\\mathbf{c}$ that\n",
    "encodes the information of the input sequence.\n",
    "To quantify computational cost,\n",
    "denote by \n",
    "$\\mathcal{Y}$ (it contains \"&lt;eos&gt;\")\n",
    "the output vocabulary.\n",
    "So the cardinality $\\left|\\mathcal{Y}\\right|$ of this vocabulary set\n",
    "is the vocabulary size.\n",
    "Let us also specify the maximum number of tokens\n",
    "of an output sequence as $T'$.\n",
    "As a result,\n",
    "our goal is to search for an ideal output\n",
    "from all the \n",
    "$\\mathcal{O}(\\left|\\mathcal{Y}\\right|^{T'})$\n",
    "possible output sequences.\n",
    "Of course, \n",
    "for all these output sequences,\n",
    "portions including and after \"&lt;eos&gt;\" will be discarded\n",
    "in the actual output.\n",
    "\n",
    "## Greedy Search\n",
    "\n",
    "First, let us take a look at \n",
    "a simple strategy: *greedy search*.\n",
    "This strategy has been used to predict sequences in :numref:`sec_seq2seq`.\n",
    "In greedy search,\n",
    "at any time step $t'$ of the output sequence, \n",
    "we search for the token \n",
    "with the highest conditional probability from $\\mathcal{Y}$, i.e., \n",
    "\n",
    "$$y_{t'} = \\operatorname*{argmax}_{y \\in \\mathcal{Y}} P(y \\mid y_1, \\ldots, y_{t'-1}, \\mathbf{c}),$$\n",
    "\n",
    "as the output. \n",
    "Once \"&lt;eos&gt;\" is outputted or the output sequence has reached its maximum length $T'$, the output sequence is completed.\n",
    "\n",
    "So what can go wrong with greedy search?\n",
    "In fact,\n",
    "the *optimal sequence*\n",
    "should be the output sequence\n",
    "with the maximum \n",
    "$\\prod_{t'=1}^{T'} P(y_{t'} \\mid y_1, \\ldots, y_{t'-1}, \\mathbf{c})$,\n",
    "which is\n",
    "the conditional probability of generating an output sequence based on the input sequence.\n",
    "Unfortunately, there is no guarantee\n",
    "that the optimal sequence will be obtained\n",
    "by greedy search.\n",
    "\n",
    "![At each time step, greedy search selects the token with the highest conditional probability.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/s2s-prob1.svg)\n",
    ":label:`fig_s2s-prob1`\n",
    "\n",
    "Let us illustrate it with an example.\n",
    "Suppose that there are four tokens \n",
    "\"A\", \"B\", \"C\", and \"&lt;eos&gt;\" in the output dictionary.\n",
    "In :numref:`fig_s2s-prob1`,\n",
    "the four numbers under each time step represent the conditional probabilities of generating \"A\", \"B\", \"C\", and \"&lt;eos&gt;\" at that time step, respectively.  \n",
    "At each time step, \n",
    "greedy search selects the token with the highest conditional probability. \n",
    "Therefore, the output sequence \"A\", \"B\", \"C\", and \"&lt;eos&gt;\" will be predicted \n",
    "in :numref:`fig_s2s-prob1`. \n",
    "The conditional probability of this output sequence is $0.5\\times0.4\\times0.4\\times0.6 = 0.048$.\n",
    "\n",
    "![The four numbers under each time step represent the conditional probabilities of generating \"A\", \"B\", \"C\", and \"&lt;eos&gt;\" at that time step.  At time step 2, the token \"C\", which has the second highest conditional probability, is selected.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/s2s-prob2.svg)\n",
    ":label:`fig_s2s-prob2`\n",
    "\n",
    "\n",
    "Next, let us look at another example \n",
    "in :numref:`fig_s2s-prob2`. \n",
    "Unlike in :numref:`fig_s2s-prob1`, \n",
    "at time step 2\n",
    "we select the token \"C\"\n",
    "in :numref:`fig_s2s-prob2`, \n",
    "which has the *second* highest conditional probability.\n",
    "Since the output subsequences at time steps 1 and 2, \n",
    "on which time step 3 is based, \n",
    "have changed from \"A\" and \"B\" in :numref:`fig_s2s-prob1` to \"A\" and \"C\" in :numref:`fig_s2s-prob2`, \n",
    "the conditional probability of each token \n",
    "at time step 3 has also changed in :numref:`fig_s2s-prob2`. \n",
    "Suppose that we choose the token \"B\" at time step 3. \n",
    "Now time step 4 is conditional on\n",
    "the output subsequence at the first three time steps\n",
    "\"A\", \"C\", and \"B\", \n",
    "which is different from \"A\", \"B\", and \"C\" in :numref:`fig_s2s-prob1`. \n",
    "Therefore, the conditional probability of generating each token at time step 4 in :numref:`fig_s2s-prob2` is also different from that in :numref:`fig_s2s-prob1`. \n",
    "As a result, \n",
    "the conditional probability of the output sequence \"A\", \"C\", \"B\", and \"&lt;eos&gt;\" \n",
    "in :numref:`fig_s2s-prob2`\n",
    "is $0.5\\times0.3 \\times0.6\\times0.6=0.054$, \n",
    "which is greater than that of greedy search in :numref:`fig_s2s-prob1`. \n",
    "In this example, \n",
    "the output sequence \"A\", \"B\", \"C\", and \"&lt;eos&gt;\" obtained by the greedy search is not an optimal sequence.\n",
    "\n",
    "## Exhaustive Search\n",
    "\n",
    "If the goal is to obtain the optimal sequence, we may consider using *exhaustive search*: \n",
    "exhaustively enumerate all the possible output sequences with their conditional probabilities,\n",
    "then output the one \n",
    "with the highest conditional probability.\n",
    "\n",
    "Although we can use exhaustive search to obtain the optimal sequence, \n",
    "its computational cost $\\mathcal{O}(\\left|\\mathcal{Y}\\right|^{T'})$ is likely to be excessively high. \n",
    "For example, when $|\\mathcal{Y}|=10000$ and $T'=10$, we will need to evaluate $10000^{10} = 10^{40}$ sequences. This is next to impossible!\n",
    "On the other hand,\n",
    "the computational cost of greedy search is \n",
    "$\\mathcal{O}(\\left|\\mathcal{Y}\\right|T')$: \n",
    "it is usually significantly smaller than\n",
    "that of exhaustive search. For example, when $|\\mathcal{Y}|=10000$ and $T'=10$, we only need to evaluate $10000\\times10=10^5$ sequences.\n",
    "\n",
    "\n",
    "## Beam Search\n",
    "\n",
    "Decisions about sequence searching strategies\n",
    "lie on a spectrum,\n",
    "with easy questions at either extreme.\n",
    "What if only accuracy matters?\n",
    "Obviously, exhaustive search.\n",
    "What if only computational cost matters?\n",
    "Clearly, greedy search.\n",
    "A real-world application usually asks\n",
    "a complicated question,\n",
    "somewhere in between those two extremes.\n",
    "\n",
    "*Beam search* is an improved version of greedy search. It has a hyperparameter named *beam size*, $k$. \n",
    "At time step 1, \n",
    "we select $k$ tokens with the highest conditional probabilities.\n",
    "Each of them will be the first token of \n",
    "$k$ candidate output sequences, respectively.\n",
    "At each subsequent time step, \n",
    "based on the $k$ candidate output sequences\n",
    "at the previous time step,\n",
    "we continue to select $k$ candidate output sequences \n",
    "with the highest conditional probabilities \n",
    "from $k\\left|\\mathcal{Y}\\right|$ possible choices.\n",
    "\n",
    "![The process of beam search (beam size: 2, maximum length of an output sequence: 3). The candidate output sequences are $A$, $C$, $AB$, $CE$, $ABD$, and $CED$.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/beam-search.svg)\n",
    ":label:`fig_beam-search`\n",
    "\n",
    "\n",
    ":numref:`fig_beam-search` demonstrates the \n",
    "process of beam search with an example. \n",
    "Suppose that the output vocabulary\n",
    "contains only five elements: \n",
    "$\\mathcal{Y} = \\{A, B, C, D, E\\}$, \n",
    "where one of them is “&lt;eos&gt;”. \n",
    "Let the beam size be 2 and \n",
    "the maximum length of an output sequence be 3. \n",
    "At time step 1, \n",
    "suppose that the tokens with the highest conditional probabilities $P(y_1 \\mid \\mathbf{c})$ are $A$ and $C$. At time step 2, for all $y_2 \\in \\mathcal{Y},$ we compute \n",
    "\n",
    "$$\\begin{aligned}P(A, y_2 \\mid \\mathbf{c}) = P(A \\mid \\mathbf{c})P(y_2 \\mid A, \\mathbf{c}),\\\\ P(C, y_2 \\mid \\mathbf{c}) = P(C \\mid \\mathbf{c})P(y_2 \\mid C, \\mathbf{c}),\\end{aligned}$$  \n",
    "\n",
    "and pick the largest two among these ten values, say\n",
    "$P(A, B \\mid \\mathbf{c})$ and $P(C, E \\mid \\mathbf{c})$.\n",
    "Then at time step 3, for all $y_3 \\in \\mathcal{Y}$, we compute \n",
    "\n",
    "$$\\begin{aligned}P(A, B, y_3 \\mid \\mathbf{c}) = P(A, B \\mid \\mathbf{c})P(y_3 \\mid A, B, \\mathbf{c}),\\\\P(C, E, y_3 \\mid \\mathbf{c}) = P(C, E \\mid \\mathbf{c})P(y_3 \\mid C, E, \\mathbf{c}),\\end{aligned}$$ \n",
    "\n",
    "and pick the largest two among these ten values, say \n",
    "$P(A, B, D \\mid \\mathbf{c})$   and  $P(C, E, D \\mid  \\mathbf{c}).$\n",
    "As a result, we get six candidates output sequences: (i) $A$; (ii) $C$; (iii) $A$, $B$; (iv) $C$, $E$; (v) $A$, $B$, $D$; and (vi) $C$, $E$, $D$. \n",
    "\n",
    "\n",
    "In the end, we obtain the set of final candidate output sequences based on these six sequences (e.g., discard portions including and after “&lt;eos&gt;”).\n",
    "Then\n",
    "we choose the sequence with the highest of the following score as the output sequence:\n",
    "\n",
    "$$ \\frac{1}{L^\\alpha} \\log P(y_1, \\ldots, y_{L}) = \\frac{1}{L^\\alpha} \\sum_{t'=1}^L \\log P(y_{t'} \\mid y_1, \\ldots, y_{t'-1}, \\mathbf{c}),$$\n",
    ":eqlabel:`eq_beam-search-score`\n",
    "\n",
    "where $L$ is the length of the final candidate sequence and $\\alpha$ is usually set to 0.75. \n",
    "Since a longer sequence has more logarithmic terms in the summation of :eqref:`eq_beam-search-score`,\n",
    "the term $L^\\alpha$ in the denominator penalizes\n",
    "long sequences.\n",
    "\n",
    "The computational cost of beam search is $\\mathcal{O}(k\\left|\\mathcal{Y}\\right|T')$. \n",
    "This result is in between that of greedy search and that of exhaustive search. In fact, greedy search can be treated as a special type of beam search with \n",
    "a beam size of 1. \n",
    "With a flexible choice of the beam size,\n",
    "beam search provides a tradeoff between\n",
    "accuracy versus computational cost.\n",
    "\n",
    "\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Sequence searching strategies include greedy search, exhaustive search, and beam search.\n",
    "* Beam search provides a tradeoff between accuracy versus computational cost via its flexible choice of the beam size.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Can we treat exhaustive search as a special type of beam search? Why or why not?\n",
    "1. Apply beam search in the machine translation problem in :numref:`sec_seq2seq`. How does the beam size affect the translation results and the prediction speed?\n",
    "1. We used language modeling for generating text following  user-provided prefixes in :numref:`sec_rnn_scratch`. Which kind of search strategy does it use? Can you improve it?\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java",
   "language": "java",
   "name": "java"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "Java",
   "pygments_lexer": "java",
   "version": "14.0.2+12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}