Beam Search

Beam Search

We use a linear layer to transform the Decoder’s output into a score for each word in the vocabulary.

The straightforward – and greedy – option would be to choose the word with the highest score and use it to predict the next word. But this is not optimal because the rest of the sequence hinges on that first word you choose. If that choice isn’t the best, everything that follows is sub-optimal. And it’s not just the first word – each word in the sequence has consequences for the ones that succeed it.

It might very well happen that if you’d chosen the third best word at that first step, and the second best word at the second step, and so on… that would be the best sequence you could generate.

It would be best if we could somehow not decide until we’ve finished decoding completely, and choose the sequence that has the highest overall score from a basket of candidate sequences.

Beam Search does exactly this.

*At the first decode step, consider the top k candidates.
*Generate k second words for each of these k first words.
*Choose the top k [first word, second word] combinations considering additive scores.
*For each of these k second words, choose k third words, choose the top k [first word, second word, third word] combinations.
*Repeat at each decode step.
*After k sequences terminate, choose the sequence with the best overall score.

As you can see, some sequences (striked out) may fail early, as they don’t make it to the top k at the next step. Once k sequences (underlined) generate the token, we choose the one with the highest score.

beam search

beam search只在预测的时候需要。训练的时候因为知道正确答案，并不需要再进行这个搜索。
预测的时候，假设词表大小为3，内容为a，b，c。beam size是2，decoder解码的时候：
1： 生成第1个词的时候，选择概率最大的2个词，假设为a,c,那么当前的2个序列就是a和c。
2：生成第2个词的时候，我们将当前序列a和c，分别与词表中的所有词进行组合，得到新的6个序列aa ab ac ca cb cc，计算每个序列的得分并选择得分最高2个序列，作为新的当前序列，假如为aa cb或aa ab。
3：后面会不断重复这个过程，直到遇到结束符或者达到最大长度为止。最终输出得分最高的2个序列。