Best-k Search Algorithm for Neural Text Generation (2023)

Abstract: Modern natural language generation paradigms require a decoding strategy to obtain quality sequences out of the model. Beam search yields high-quality but low diversity outputs; stochastic approaches suffer from high variance and sometimes low quality. In this work, we propose 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 efficiency and diversity. Upweighting recently discovered nodes accompanied by heap pruning ensures the completeness of the search procedure. 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 proposed algorithm is parameter-free, lightweight, efficient, and easy-to-use.

## introduction large-scale pre-trained language models @xcite ; @xcite ; @xcite ; @xcite , inter alia) has significantly advanced the field of natural language generation. despite the models' increasing capability in fluency, expressiveness and domain generalization, the generated outputs from these models are far from perfect @xcite @xcite . the decoding strategy is another crucial piece in this paradigm. if we form text generation as a search problem, decoding 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 generated sequences are usually repetitive because many diverse hypotheses are pruned at earlier stage of 1 the code implementation is available at https:// jiacheng-xu.github.io/ . search @xcite . sampling-based approaches @xcite can indeed generate more diverse sequences, but they are hard to control due to their stochastic nature. 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 satisfies these properties. first, it is a reproducible and deterministic algorithm. more importantly, since it theoretically does not prune hypotheses, it preserves a more diverse set of options and allows simultaneous expansion of hypotheses with different lengths. despite these intriguing features, we identified two challenges, efficiency and completeness, of directly applying it to text generation. in this work, we propose the best-k search for diverse and high-quality text generation. our approach re-invents bfs with a few design changes to overcome the issues mentioned before. parallel exploration is designed to explore the top k nodes from the search frontier each time instead of one in bfs. we also add a temporal decay mechanism to the algorithm to encourage search completions. a simple yet effective stateless scoring function as an alternative to more complicated length-adjusted counterparts is devised, and we show that it works well and helps further in finding diverse texts. to verify the proposed algorithm, we conduct comprehensive experiments on four tasks, question generation, commonsense generation, text summarization and machine translation. our results show that the proposed algorithm works well with a wide range of models on six datasets. our approach yields high-fidelity, diverse and natural outputs while maintaining quality. our contributions are (1) investigation of best-first search for text generation; ## revisiting best-first search in this section, we will introduce the vanilla bestfirst search in the context of natural language generation as a decoding algorithm, and cover the first research question: is bfs adequate in searching hypotheses in text generation? setup text generation can be formulated as a sequence generation process given input x and a probabilistic language model foot_0 parameterized by θ. traditionally, maximum a posteriori (map) decoding strategy is deployed to elicit highest-scoring output sequences arg max y * p θ (y * |x). most previous work uses the log-likelihood of the sequence as the proxy for assessing the (partial) sequence quality. however, recent studies found discrepancies between model likelihood and quality assessed by humans @xcite @xcite . various approaches including length normalization @xcite , quality-aware decoding @xcite , and regularized decoding @xcite have attempted to modify the objective to mitigate the gap. in this work, we adopt h(•) as the scoring function, and h(y 1•••t ) is the score of a hypothesis y 1•••t . ## our approach: best-k search in this section, we will introduce best-k search, a novel search algorithm inspired by the vanilla best-first search. it features a few components: (1) parallel exploration enables batch-wise exploration in the search graph; (2) temporal decay yields a higher completion rate and fewer dangling nodes; ## application: reranking diverse outputs machine translation is typically considered as a precision-oriented task, where typically only a few translations are considered as correct. in this section, we would like to answer the rq: do we benefit by selecting from a pool of high-quality diverse outputs, even when the task does not necessarily require such? setup we use a popular machine translation dataset with multiple references @xcite , based on wmt'14 en-fr and en-de test sets @xcite . the model for this task is the mbart foot_3 model @xcite . in order to rerank decoded outputs, we adopt a state-of-the-art quality estimation model for mt, comet-qe @xcite . the quality estimation model we use to rerank all the outputs is wmt21-comet-qe-da. the qe model is a referenceless model q(s, t) which judges whether the source input s and the hypothesis translation t form a matched pair based on regression metrics. result we present the result on mt en-de and en-fr in ## related works best-first search bfs was widely used in structural prediction @xcite , statistical mt @xcite , and for searching hypotheses @xcite . recent work in decoding strategies @xcite @xcite conceptualized best-first search as part of their paradigm, but it was not the dominant component of any of these systems. text decoding algorithms stochastic decoding algorithms have gained popularity in the past few years @xcite @xcite . rolloutbased algorithms are capable of satisfying certain utility functions or constraints at the cost of efficiency @xcite @xcite . recombination-based search algorithm @xcite can find thousands of hypotheses despite complicatedness. ## conclusion in this work, we propose best-k search, a novel decoding algorithm for text generation based on bestfirst search. the algorithm features a few technical components, and generates natural and diverse text while maintaining high quality. we conduct comprehensive experiments on four tasks to verify the approach. the algorithm is orthogonal to sampling methods and it is parameter-free, lightweight, and efficient.