Best-First Beam Search (2020)

Abstract: Decoding for many NLP tasks requires an effective heuristic algorithm for approximating exact search because the problem of searching the full output space is often intractable, or impractical in many settings. The default algorithm for this job is beam search—a pruned version of breadth-first search. Quite surprisingly, beam search often returns better results than exact inference due to beneficial search bias for NLP tasks. In this work, we show that the standard implementation of beam search can be made up to 10x faster in practice. Our method assumes that the scoring function is monotonic in the sequence length, which allows us to safely prune hypotheses that cannot be in the final set of hypotheses early on. We devise effective monotonic approximations to popular nonmonontic scoring functions, including length normalization and mutual information decoding. Lastly, we propose a memory-reduced variant of best-first beam search, which has a similar beneficial search bias in terms of downstream performance, but runs in a fraction of the time.

## introduction beam search is a common heuristic algorithm for decoding structured predictors (e.g., neural machine translation models and transition-based parsers). because of the widespread adoption of recurrent neural networks and other non-markov models, traditional dynamic programming solutions, such as the viterbi algorithm @xcite , are prohibitively inefficient; this makes beam search a common component of many stateof-the-art nlp systems. despite offering no formal guarantee of finding the highest-scoring hypothesis under the model, beam search yields impressive performance on a variety of tasks-unexpectedly providing a beneficial search bias over exact search for many tasks @xcite . within nlp, most research on beam search has focused on altering the log-probability scoring function to return improved results, for example, higher bleu scores @xcite @xcite or a more diverse set of outputs @xcite . however, little work has been done to speed up beam search itself. filling this gap, this paper focuses on reformulating beam search in order to make it faster. we propose best-first beam search, a prioritized version of traditional beam search that is up to an order of magnitude faster in practice while still returning the same set of results. we additionally discuss an even faster heuristic version of our algorithm that further limits the number of candidate solutions, leading to a smaller memory footprint while still finding good solutions. concretely, we offer a novel interpretation of beam search as an agenda-based algorithm where traditional beam search is recovered by utilizing a length-based prioritization scheme. we prove that a specific best-first prioritization scheme, as in classic a * search @xcite , allows for the elimination of paths that will necessarily fall off the beam; for many scoring functions, including standard log-probability scoring, we can still guarantee the same k hypotheses as traditional beam search are returned. indeed, our algorithm returns beam search's top hypothesis the first time it encounters a complete hypothesis, allowing the program to stop early. further, we discuss the application of best-first beam search to several popular scoring functions in the literature @xcite ; this demonstrates that we have a general framework for adapting a variety of rescoring methods and alternate objectives to work with our algorithm. empirically, we compare best-first beam search to ordinary beam search on two nlp sequence-tosequence tasks: neural machine translation (nmt) and abstractive summarization (as). on nmt, we find that our algorithm achieves roughly a 30% speed-up over traditional beam search with increased gains for larger beams (e.g., ≈ 10x for a beam of 500). we find similar results hold for as. finally, we show that our memoryreduced version, which limits the number of active hypotheses, leads to additional speed-ups over best-first beam search across beam sizes while maintaining similar bleu scores. ## sequence transduction a core operation in structured prediction models is the determination of the highest-scoring output for a given input under a learned scoring model. our work focuses on sequence-to-sequence transduction: predicting an output sequence given an input sequence. one such task is machine translation, wherein a source-language sentence is mapped (''transduced'') to a target-language sentence. while our exposition focuses on sequenceto-sequence prediction, our algorithms are directly applicable to any sequential structured prediction model, such as transition-based parsers @xcite and sequence taggers @xcite . notation. let @xmath0, . . . , x n x be an input sequence of length n x and, likewise, let @xmath1, . . . , y n y be an output sequence of length n y . each y t is an element of v, the set of output tokens. finally, let y(x) be the set of all valid output sequences (i.e., complete hypotheses). for the task of language generation, which we focus on experimentally, this set is defined as scoring. we consider a general additively decomposable scoring model of the form this framework covers a variety of modeling methodologies including probabilistic transducers (both globally and locally normalized) and nonprobabilistic models such as maximum-margin techniques @xcite . most importantly, (4) covers map decoding (2) of neural sequenceto-sequence models à la @xcite we note that (5) is the scoring function used for decoding many language generation models. beam search. the worst-case running time of exactly computing (1) is exponential in n max ; namely, o(|v| n max ). 3 beam search is a commonly used approximation to ( 1 ) in nmt and language generation tasks. it is used in many (if not most) state-of-the-art nlp systems @xcite @xcite . beam search may be understood as a pruned version of the classic path-search algorithm, breadth-first search (bfs), where the breadth is narrowed to the beam size k. pseudocode is given in (1). although, beam search does not solve (1) exactly, it is a surprisingly useful approximation for nlp models. in many settings, beam 1 bos and eos are typically members of v. often, eos counts towards the n max length limit while bos does not. this is reflected in (3). definition 2.1. k-optimal hypothesis. we say that a hypothesis is k-optimal if it is the top hypothesis returned by beam search with beam size k. ## a * beam search we develop a meta-algorithm that is parameterized by several choice points. our general search algorithm for decoding (alg. 2) takes an arbitrary prioritization function, stopping criterion, and search heuristic. with certain values of these attributes, we recover many common search algorithms: greedy search, beam search, best-first search @xcite , and a * search @xcite . we propose an alternate prioritization function for beam search that allows for faster decoding while still returning the same k-optimal set of hypotheses. algorithm 2 general decoding scheme. ## scoring functions even before the findings of @xcite , it was well known that the best-scoring hypothesis with respect to the traditional likelihood objective can be far from ideal in practice @xcite @xcite . for language generation tasks specifically, the results returned by neural models using the standard scoring function are often short and default to high-frequency words @xcite . to alleviate such problems, methods that revise hypothesis scores to incorporate preferences for longer, less repetitive, or more diverse options have been introduced and are often used in practice. while most such techniques change the scoring function such that it is no longer monotonic, we can still guarantee the k-optimality of the returned hypothesis for (upper) bounded scoring functions using the methods discussed in § 4.1. in the remainder of this section, we present alternate scoring schemes adapted to work with best-first beam search. additionally, we present several heuristics which, while breaking the k-optimality guarantee, provide another set of decoding strategies worth exploring. length normalization. length normalization is a widely used hypothesis scoring method that aims to counteract the propensity for shorter sequences to have higher scores under neural mod-els; this is done by normalizing scores by hypothesis length (see @xcite for more detail). for early stopping in beam search with length normalization, @xcite propose bounding the additive length reward as the minimum of a pre-determined optimal sequence length ratio r and the final sequence length n y : score ln (x, y) = score(x, y) we formally propose two methods for length normalization. we use the scoring functions in (7) or (8) with either: (1) the following heuristic: despite their similarities, these two methods are not guaranteed to return the same results. whereas the second method will return the same k-optimal hypotheses as beam search, using a heuristic during pruned search means we can no longer guarantee the k-optimality of the results with respect to the scoring function as the heuristic may push hypotheses off of the beam. we present experimental results for both methods in § 6. ## experiments we run our algorithm on several language-related tasks that typically use beam search for decoding: nmt and as. specifically, experiments are performed on iwslt'14 de-en @xcite , wmt'17 de-en @xcite , mttt fr-en @xcite using both transformers @xcite and convolutional sequence-to-sequence models @xcite . for reproducibility, we use the data pre-processing scripts provided by fairseq @xcite and follow their methods for training sequence transduction models. hyperparameters are set in accordance with previous works. specifically, on iwslt'14 and mttt tasks, we follow the recommended transformer settings for iwslt'14 in fairseq, foot_8 which are based on @xcite and @xcite . hyperparameters for models trained on the wmt task are set following version 3 of the tensor2tensor toolkit @xcite . we use byte-pair encoding (bpe; @xcite ) for all languages. vocabulary sizes for wmt and iwslt'14 are set from recommendations for the respective tasks in fairseq; for the mttt tasks, vocabulary sizes are tuned on models trained with standard label-smoothing regularization. similarly, the cnn/dailymail dataset is pre-processed and uses bpe following the same steps as @xcite ; model hyperparameters are likewise copied. details are available on fairseq's web site. 10 we use bleu @xcite (evaluated using sacrebleu @xcite ) for mt metrics and rouge-l (lin, 2004) for abstractive summarization metrics. we build our decoding framework in sgnmt. 11 ## conclusion we propose best-first beam search, an algorithm that allows for faster decoding while still guaranteeing k-optimality. we provide results on several sequence-to-sequence transduction tasks that show the speed-ups that our algorithm provides over standard beam search for decoding neural models. we adapt several popular alternate scoring functions to best-first beam search and provide a framework that can be used to adapt other scoring methods such as coverage normalization @xcite or diverse beam search @xcite . we also provide a memory-reduced version of our algorithm, which returns competitive results in a fraction of the time needed for standard beam search.