A Streaming Approach For Efficient Batched Beam Search

Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, pages 4526–4535, November 16–20, 2020. c⃝2020 Association for Computational Linguistics 4526 A Streaming Approach For Efficient Batched Beam Search Kevin Yang, Violet Yao, John DeNero, Dan Klein UC Berkeley {yangk,violetyao,denero,klein}@berkeley.edu Abstract We propose an efficient batching strategy for variable-length decoding on GPU architec- tures. During decoding, when candidates ter- minate 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, experi- ments show that our method can speed up de- coding in other domains, such as semantic and syntactic parsing. 1 Introduction While inference is often cheap compared to training in modern neural models, one may need to run in- ference 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 (Turovsky, 2016). Large- scale inference is also required for methods such as iterative backtranslation and knowledge distilla- tion to generate training data (Hoang et al., 2018; Kim and Rush, 2016). 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 nev- ertheless often preferred in MT. Not only is beam search usually more accurate than greedy search, but it also outputs a diverse set of decodings, en- abling reranking approaches to further improve ac- curacy (Yee et al., 2019; Ng et al., 2019; Charniak and Johnson, 2005; Ge and Mooney, 2006). However, it is challenging to optimize the per- formance of beam search for modern neural ar- chitectures. 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 (Freitag and Al-Onaizan, 2017; Wu et al., 2016). When decoding with Fairseq’s state-of- the-art WMT’19 model (Ng et al., 2019), our method further improves over the speed of base- line variable-width beam search: up to 16.5% on a 32GB V100 GPU, without changing BLEU (Pa- pineni et al., 2002). Our approach also improves decoding efficiency in lightweight models for se- mantic and syntactic parsing.1 In principle, our method can be applied to any task which …

Content truncated to 500 words