A Streaming Approach For Efficient Batched Beam Search (2020)

Abstract: We propose an efficient batching strategy for variable-length decoding on GPU architectures. During decoding, when candidates terminate or are pruned according to heuristics, our streaming approach periodically “refills” the batch before proceeding with a selected subset of candidates. We apply our method to variable-width beam search on a state-of-the-art machine translation model. Our method decreases runtime by up to 71% compared to a fixed-width beam search baseline and 17% compared to a variable-width baseline, while matching baselines’ BLEU. Finally, experiments show that our method can speed up decoding in other domains, such as semantic and syntactic parsing.

## introduction while inference is often cheap compared to training in modern neural models, one may need to run inference frequently or continually. such is the case for online machine translation (mt) services: as far back as 2016, google translate already translated 100 billion words daily @xcite . largescale inference is also required for methods such as iterative backtranslation and knowledge distillation to generate training data @xcite . for such high-throughput applications, it is useful to decrease inference cost. meanwhile, we must preserve accuracy: beam search is slower than greedy decoding, but is nevertheless often preferred in mt. not only is beam search usually more accurate than greedy search, but it also outputs a diverse set of decodings, enabling reranking approaches to further improve accuracy @xcite @xcite . however, it is challenging to optimize the performance of beam search for modern neural architectures. unlike classical methods in sparse computation settings, modern neural methods typi-cally operate in dense (batched) settings to leverage specialized hardware such as gpus. in this work, we propose a streaming method to optimize gpu-batched variable-output-length decoding. our method does not use a fixed batch during inference; instead, it continually "refills" the batch after it finishes translating some fraction of the current batch. our method then continues decoding on the remaining candidates in the batch, prioritizing those least expanded. we apply our method to variable-width beam search. for variable-output-length decoding even in batched settings, variable-width beam search often modestly decreases accuracy in exchange for substantial speedups over fixed-width beam search @xcite . when decoding with fairseq's state-ofthe-art wmt'19 model @xcite , our method further improves over the speed of baseline variable-width beam search: up to 16.5% on a 32gb v100 gpu, without changing bleu @xcite . our approach also improves decoding efficiency in lightweight models for semantic and syntactic parsing. 1 in principle, our method can be applied to any task which sequentially processes variable-length data. ## background: beam search given encoder e and decoder d, our task is to convert inputs {x 1 . . . x n } into corresponding outputs {ȳ 1 . . . ȳn }, for data size n . for example, in machine translation, each x i is a source sentence consisting of a sequence of tokens and each ȳi is a translation. we assume d(e i , y i ) receives e @xmath0, constructing ȳi one token at a time. one method of constructing ȳi for a given x i is greedy search. let y lt i be the in-construction candidate with length l @xmath1. we initialize candidates become finalized after following the third (bottom-most) branch in an expansion, corresponding to the end token w eos ; they then undergo only no-op expansions thereafter. the first two rows of beams depict normal operation of variable-width beam search, including heuristic pruning in the light blue beam at @xmath2. the third row shows an important detail of our method: var-stream refills the batch after @xmath3, when only n beams remain, and the remaining purple beam halts computation until the two newly added beams reach the same l t . (this detail matters in transformer architectures; see appendix a.2.) y 1 i as the start token w sos , and at each timestep t obtain y lt+1 i by concatenating the maximumprobability token. we finalize y lt i as ȳi once we append the end token, or at some maximum length. previous work has found that greedy search often underperforms beam search in accuracy, and gains from non-greedy decoding have also been observed in many classical models @xcite @xcite . see the dark blue, green, and brown beams in figure 1 for normal operation of fixed-width beam search using beam width @xmath4. for each input x i , fixed-width beam search tracks a lengthl t , width-k beam b lt i for each time t. b lt i contains k length-l t candidates y lt i1 . . . y lt ik with maximum log-likelihood in order, denoted by the shapes inside the rounded rectangles (beams) in the figure. at each step, beam search considers all k|v| possible candidate expansions (one-token extensions of existing candidates), where v is the vocabulary. the top k expansions become the expanded beam b lt+1 i . figure 1 shows these expansions at each timestep for |v| = 3, with active non-terminated candidates (circles) becoming finalized (stars) after following the bottom-most branch, corresponding to the end token w eos . in the end, beam search yields k finalized candidates ȳi1 . . . ȳik compared to a single ȳi in greedy search. variable-width beam search reduces the computational cost of the above fixed-width beam search by pruning the full beam b lt i using heuristics h, for example at @xmath5. the width of the resulting pruned beam h(b lt i ) is no longer always exactly equal to k, and may vary over time. ## streaming variable-length decoding as an example, consider translating a batch of n german sentences into english via a traditionallybatched variable-width beam search (e.g., @xcite ) on a gpu. henceforth we refer to this baseline as var-batch. an inefficiency results from decoding being inherently variable in length: after t steps, we may have completed @xmath6, but the last n -m beams may take several more timesteps. for example, in figure 1 , our initial batch consists of the dark blue, brown, and purple beams. after the dark blue and brown beams terminate, we would still be stuck decoding the purple beam by itself. the resulting gpu underutilization motivates our streaming approach var-stream applied to variable-width beam search (figure 1 ). for batch size n, var-stream initially proceeds identically to var-batch. but when the number of remaining beams drops below n for some @xmath7, 1), var-stream encodes a new batch of inputs x to "refill" its batch to size n. 2 this occurs at @xmath8, where we refill the batch using the green and light blue beams. note the active beams are no longer of equal length l @xmath9. at each subsequent t, var-stream only expands beams h(b lt i ) with minimal l t ; in particular, the purple beam in figure 1 pauses computation at @xmath10. 3 when decoding with state-of-the-art transformer architectures for mt, it is advantageous to expand only beams with minimal l t at each step, because self-attention causes steps at higher l t to be more expensive; see appendix a.2. (for rnn-based architectures, it may be faster to expand all active beams at each step.) we emphasize that var-stream is an implementation optimization, exactly matching the output of var-batch. full details in algorithm 1. when the memory bottleneck is partially the decoding process itself rather than caching the input encodings e(x i ) or beams h(b lt i ), var-stream can cache additional encodings and beams on gpu. at each t, var-stream then selects beams up to some limit on total beam width, filling gpu capacity even in the case of variable-width beams. this batching constraint addresses a second inefficiency in gpu utilization: the widths of the pruned beams h(b lt i ) may vary over time. we exploit this in semantic (sec. 4.2) and syntactic parsing (sec. 4.3). ## experiments we apply var-stream to variable-width beam search in machine translation, semantic parsing, 2 our method is relatively insensitive to (appendix a.3). 3 as very long translations could get "stuck" in the batch, one can periodically finish computation on all remaining beams in the batch if latency is a concern in addition to throughput. algorithm 1 var-stream input: inputs @xmath11, . . . x n }, model (e, d), batch size n, refill threshold , beam width k, pruning heuristics h 1: procedure decode(x, e, d, n, , k, h) 2: # initialize encodings, beams, final outputs # refill batch by @xmath12 with min l t 14: for (e i , h(b lt i )) ∈ (e s , b s ) do 15: for y l t @xmath13: compute expansion d(e i , y l t ij ) 17: update h(b lt i ) to h(b lt+1 i ) 18: add finalized candidates to y [i] 19: remove terminated beams from e, b 20: return y and syntactic parsing. we use the absolute threshold and max candidates heuristics of freitag and al-onaizan (2017) as h, modifying only the heuristic hyperparameters for each domain based on a development set. the absolute threshold heuristic prunes candidates y lt ij whose log-probabilities fall short of the best candidate y lt i1 's by some threshold δ, i.e. log p (y lt ij ) < log p (y lt i1 ) -δ. the max candidates heuristic prevents the search from selecting more than @xmath14. in each domain we compare four methods: 1. greedy, a greedy search, 2. fixed, a fixed-width beam search, 3. var-batch, a batched variable-width beam search, and 4. var-stream, our streaming method. we sort and bucket inputs by length for batching. ## discussion in this work, we have proposed a streaming method for variable-length decoding to improve gpu utilization, resulting in cheaper inference. applied to a state-of-the-art machine translation model, our method yields substantial speed improvements compared to traditionally-batched variable-width beam search. we also apply our method to both semantic and syntactic parsing, demonstrating our method's broader applicability to tasks that process variable-output-length data in a sequential manner.