Best-k Search Algorithm for Neural Text Generation

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 12385–12401 July 9-14, 2023 ©2023 Association for Computational Linguistics Best-k Search Algorithm for Neural Text Generation Jiacheng Xu Caiming Xiong Silvio Savarese Yingbo Zhou Salesforce AI Research {jiacheng.xu,cxiong,ssavarese,yingbo.zhou}@salesforce.com Abstract Modern natural language generation paradigms require a decoding strategy to obtain quality se- quences out of the model. Beam search yields high-quality but low diversity outputs; stochas- tic approaches suffer from high variance and sometimes low quality. In this work, we pro- pose a deterministic search algorithm balancing both quality and diversity. We first investigate the vanilla best-first search (BFS) algorithm and then propose the best-k search algorithm. Inspired by BFS, we greedily expand the top k nodes, instead of the first node, to boost effi- ciency and diversity. Upweighting recently dis- covered nodes accompanied by heap pruning ensures the completeness of the search proce- dure. Experiments on four NLG tasks show that best-k search yields more diverse and natural outputs compared to strong baselines, while our approach maintains high text quality. The pro- posed algorithm is parameter-free, lightweight, efficient, and easy-to-use.1 1 Introduction Large-scale pre-trained language models (Devlin et al. (2019); Raffel et al. (2020); Brown et al. (2020); Nijkamp et al. (2022), inter alia) has signif- icantly advanced the field of natural language gen- eration. Despite the models’ increasing capability in fluency, expressiveness and domain generaliza- tion, the generated outputs from these models are far from perfect (Gehman et al., 2020; Kryscinski et al., 2020; Fabbri et al., 2021). The decoding strategy is another crucial piece in this paradigm. If we form text generation as a search problem, de- coding strategies are essentially search algorithms over the space composed by vocabulary V. Beam search, a heuristic search algorithm, has been the go-to choice for many years. However, the gener- ated sequences are usually repetitive because many diverse hypotheses are pruned at earlier stage of 1The code implementation is available at https:// jiacheng-xu.github.io/. 20 21 22 23 24 25 26 27 Quality (ROUGE-1) 30 40 50 Diversity (Distinct-1) BeamSearch Sample BS+Sample Ours Figure 1: Generated text diversity and quality, measured by Distinctness-1(↑) and ROUGE-1(↑), on question gen- eration. The dataset used is QuoRef and the model applied is MixQG. Our approach, best-k search with two configurations, beats baseline methods including beam search variations, sampling and BS+Sample meth- ods on diversity and quality; see Sec. 4.4 for details. search (Eikema and Aziz, 2020). Sampling-based approaches (Fan et al., 2018; Holtzman et al., 2020) can indeed generate more diverse sequences, but they are hard to control due to their stochastic na- ture. Sometimes outputs are duplicate; sometimes a sampling choice breaks the whole sequence. We are looking for a decoding algorithm with high flexibility and controllability while it could also yield diverse outputs for certain use cases. We find that best-first search (BFS) algorithm satis- fies these properties. First, it is a reproducible and deterministic algorithm. More importantly, since it theoretically does not prune hypotheses, …

Content truncated to 500 words