Beam Search还能更快？结合优先队列的最佳优先化Beam Search

论文标题：Best-First Beam Search

论文作者：Clara Meister, Tim Vieira, Ryan Cotterell

论文链接 ： https://arxiv.org/pdf/2007.03909.pdf

---

Beam Search是当前各类文本生成任务的标配解码方式，作为一种受限的宽度优先搜索，它可以极大降低搜索复杂度。 但是，Beam Search依旧还有提高的空间！

本文提出一种结合优先队列和A*经验式搜索的Beam Search，可以显著减少调用打分函数（如负对数似然）的次数，从而能够使整个Beam Search速度大大加快，还能得到和Beam Search一样的结果。

Beam Search 概述

在文本生成中，Beam Search已经被作为一种普遍的解码方式运用在各种任务中。假设当前输入是 ，要生成的句子是 ，最大长度是 ，字典是 ，那么，所有可能的句子就是 。Beam Search就是在每一步通过打分函数只保留 个候选句子。打分函数一般是对数似然：

-最大堆，在每一步，对这 个句子中的每一个，如果它已经结束了（第5行），那么就不处理它，否则加上当前步的得分（每个句子都有 次计算得分），再把新得到的句子加入堆中，等到所有句子处理完之后，再从中取 个得分最大的。

因为Beam Search是主要的比较对象，因此我们把由Beam Search返回的 个句子叫做 -最佳集，其中的任何一个句子都称为 -最佳句子。

Beam Search 的问题

从上面的介绍讲，乍一看，Beam Search的效率很高，然而，Beam Search也隐藏着一些细节上的问题。举个例子，假设 ，词典是 ，现在保留的句子是 ，它们当前的得分分别是 ，且下一步可能的句子的得分是：

	句子	得分	句子	得分	句子	得分
acb	acba	0.30	acbb	0.31	acbc	0.25
abb	abba	0.29	abbb	0.28	abbc	0.24

看，在下一步的Beam Search中，我们当前得到的第二个句子 根本不重要，因为下一步得分最高的两个句子都来自于 ，因此，我们完全可以不去计算由 得到的三个句子。

那么，这里的关键是什么呢？是最佳优先，即让当前得分最大的句子先计算下一步，这就很像A*搜索。

A * Beam Search 

A*是一种启发式搜索，即通过预定义的得分函数，每次都首先探索得分最高的那一条路径。显然，只要得分函数定义得好，就能比DFS和BFS更快地找到正确答案。

那么在这里，所谓的正确答案就是BS得到的 -最佳结果。我们的目标是设计一种启发式的搜索算法，既能得到 -最佳结果，又比BS更快。

下面的算法就描述了一种通用的搜索算法框架。

注意观察，和BS相比，该算法引入了4个关键的成分：（1）比较符；（2）停止策略；（3） ；（4）经验函数。其中除了（3）是BS原有的之外，其他的都是新增的。有了这几个成分，就可以使用该算法实现BS，和即将要介绍的A* Beam Search。

来看一下这个算法做了什么。首先用定义了一个优先队列；然后在搜索的每一步，直接看得分最大的那个句子，如果这个句子结束了，那么就继续放入优先队列中。

否则像BS一样考虑字典中的每个字，在原来的得分（如对数似然）上再加上经验函数得分，然后再把它送到优先队列中。

注意到图中绿色的部分，它的意思是，对同一个句子长度，如果之前已经遇到了 个同样长度的句子，那么后面的同样长度的句子就不再考虑了，这是因为最先考虑的一定是得分最大的，后面打分的得分过小再考虑就没有意义。

那么这个算法怎么恢复到BS呢？其实很简单，只需要让（4）经验函数恒为0，让比较运算符（1）为：

这个式子的意思是，让那些要么长度更短要么长度相同且得分更大的句子优先被计算，这其实和BS是一致的。因为BS在每一步要么优先考虑已经结束的句子，要么长度相同但得分更大。

同时让终止条件（2）为：

也就是说要让所有句子都搜索完。

那么，怎么用这个算法实现最佳优先的A* Beam Search呢？回想一下，A*是优先探索得分较大的路径，但是它并不限制路径的长度，也就是说，如果已经探索了 条长度为 的路径，如果后面需要，它还会继续探索长度为 的其他路径。

所以它在最坏情况下仍然是 的。将A*和BS结合起来，就是要在A*算法上限制可探索的路径数，这就是上述算法的思想。

所以，最佳优先搜索就可以定义为：

• （1）比较运算符：
• （2）终止条件：
• （4）经验函数：根据得分函数的设置（对数似然设置为0）

正如上面所述，这个算法的优势在于，对于对数似然而言，一旦已经计算了长度为 的 个句子，那么后面任何长度小于 的句子都不需要再考虑了，因为它们的得分在该长度上一定更低。

本文有一节专门证明了该结论，并证明了该算法的时间复杂度，由于内容较为复杂，有兴趣的读者可以参考原文Section 4。

下图是从本算法可以导出的主要搜索算法，包括BS，最佳优先BS，A* 等算法。可以看到，通过定义4个关键因子，我们可以实现各种典型的搜索算法，当然运用在文本生成上，BS、最佳BS和A* BS是我们最关心的。

实验

空口无凭，且看实验效果。本文在多个数据集上测量各搜索算法具体调用了多少次得分函数。注意到这里没直接比较运行时间，原文认为运行时间和硬件关系很大。

下图是在本文实验的机器上，运行时间和得分函数调用次数的关系。可以看到，二者有很强的线性相关性，所以直接比较得分函数是合理的。

下表是主要实验结果。Beam Search (ES)是指使用了提前终止的BS。

从结果可以看到，相比BS和BS(ES)，BFBS调用得分函数的次数显著降低，这在 很大的时候尤其明显。这也就意味着，模型花在解码上的时候大大缩短。

下面一个问题是，BFBS真的可以取得BS的最后得到的句子吗？如下表所见，shrinking和ES都不能完全得到BS的结果（ ），反而效果也会变差。

相反，BFBS始终能得到BS一样的结果，只要BS能够取得好效果，BFBS也就能取得同样的效果。

上面的实验都是基于对数似然这个得分函数而言，对于其他得分函数，如互信息、长度归一化等，本文也都做了实验，限于篇幅，在此不再展开，有兴趣的读者可以参考原文。

小结

这是一篇偏算法理论的应用文，提出了一种通用的Beam Search算法框架，采用不同的比较运算符、经验函数、终止策略或Beam Size，算法可以复原很多常见的搜索算法。

基于这个算法框架，本文提出了一种新的，基于最佳优先策略的Beam Search，和BS相比，能够显著减少得分函数的调用次数，从而提高解码速度。

当然，本文所提出的算法其实并不算简单，在理解上也有诸多障碍，尤其是很多细节，如是否可以并行化解码，优先队列又该怎么在并行条件下实现等，都需要仔细推敲。原文Section3.1给出了实现的概述，读者可以根据此自己尝试实现一个版本。

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

关于PaperWeekly

PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域，欢迎在公众号后台点击「交流群」，小助手将把你带入 PaperWeekly 的交流群里。