Enabling Arbitrary Translation Objectives with Adaptive Tree Search

1 INTRODUCTION . Conditional text generation tasks , such as machine translation , consist of two parts : a model that assigns scores to candidate outputs , and a search component that interacts with the model in order to find an output that maximizes the score assigned by the model . This search problem is a hard combinatorial optimization problem , and as a result , constraints are frequently imposed on the structure of the model to make solving or approximating the search problem easier . In neural machine translation , an autoregressive factorization of the output probability distribution is widely used ( Kalchbrenner & Blunsom , 2013 ; Sutskever et al. , 2014 ; Vaswani et al. , 2017 ) , and a variety of conditional independence assumptions are made in other model classes from statistical translation models ( Brown et al. , 1993 ; Koehn et al. , 2003 ) to non-autoregressive neural models ( Lee et al. , 2018 ) . Although these assumptions enable fast and accurate approximations to the search problem with simple and efficient algorithms ( e.g. , beam search ) , which can be crucial for efficient production applications , they limit the form of the models and thereby restricting the kinds of architectures that can be used to address observed model failures . ∗Authors contributed equally . †Work carried out while Wang Ling was at DeepMind . Despite the algorithmic benefits that beam search provides , we argue that it is a poor foundation for long-term scientific progress toward an accurate and reliable translation model whose scores adequately predict translation quality . First , it can only be applied to autoregressive models , which have well-known length calibration problems due to a tendency to drop or “ hallucinate ” content , and this tendency has been remarkably resistant to remedy across a variety of different autoregressive architectures ( Koehn & Knowles , 2017 ; Lin et al. , 2020 ) . The existing heuristic solutions—e.g. , nonlocal corrections to the search objective at decoding time , or global statistics about the population of per-word probabilities ( Wu et al. , 2016 ; Meister et al. , 2020 ) —point to non-autoregressive components in the search objective as necessary parts of the solution.1 We therefore would like to directly work with model classes that contain these solutions , rather than being dependent on limited heuristics that are imposed after the fact . Second , beam search is strongly biased towards translations that yield high initial scores due to the heuristic used to score partial translations ( Stahlberg & Byrne , 2019 ; Meister et al. , 2020 ) . While this feature provides short term translation quality gains as it addresses many of the calibration issues in autoregressive models , it is undesirable in the longer term as it masks many modeling issues that need addressing . More importantly , as scientific progress drives better correlations between model scores and translation quality , search errors caused by this bias will inevitably have a negative impact on both model scores and translation quality . To address these shortcomings , we introduce the beam adaptive tree search ( BATS ) algorithm , which is based on Monte Carlo tree search ( Coulom , 2006 ; Browne et al. , 2012 ) . Like MCTS , BATS estimates the value of internal nodes ( i.e. , partial translations ) in the search tree with estimates from an expected-outcome model based on playouts from an auxiliary model , and these estimates are refined as the search progresses . Because BATS is guided from the start by the true objective—whether autoregressive or not , and due to the refinement of initial score estimates , the BATS decoder exhibits fewer biases as the search budget increases than beam search or reranking algorithms . Our experimental section aims to characterise the impact of decoding mechanisms on both nonautoregressive and autoregressive models . For autoregressive models , we show that the calibration issues in autoregressive models can be addressed by adding a non-autoregressive component that augments the autoregressive sequence score with the lowest scoring produced token in the autoregressive decomposition of the sequence , which we name max rank . We further show that existing non-autoregressive approaches , such as the “ noisy channel ” model ( Yu et al. , 2017 ; Yee et al. , 2019 ; Ng et al. , 2019 ; Yu et al. , 2020b ; a ; Liu et al. , 2021 ) , which factorize the translation probability according to Bayes ’ rule and have been argued to be better calibrated , are ( 1 ) poorly optimized by standard reranking-based approaches , and ( 2 ) ultimately have similar calibration failures as neural autoregressive models have . Incorporating the max rank component to existing objectives benefits both beam search and BATS , but the latter yields both the best translation and model scores . Crucially , we show that our decoder can search substantially longer and achieve higher model scores before BLEU starts to deteriorate , which suggests a negative impact of the search bias in beam search . Finally , we show that once autoregressive models become robust enough to address the calibration issues , this bias has an equally negative impact on translation quality . By fine-tuning an autoregressive model to better correlate with BLEU using minimum risk training ( Shen et al. , 2016 ) , we show that BATS can achieve higher translation quality and a better model score compared to beam search . 2 BACKGROUND . 2.1 BEAM SEARCH . Translation tasks operate on the space of possible translations Y = Σ∗ ◦ { EOS } where Σ is a finite vocabulary , and EOS is a symbol represents the end of a sequence . Elements in Y correspond to full sentences y = y1 , . . . , yn with n− 1 tokens , where yn = EOS ( end of sentence ) . Each translation y is conditioned on a source sentence x and is assigned a score s ( x , y ) , which measures how well x translates to y. Decoding algorithms aim to find the best hypothesis y∗ ∈ Y under the search objective s ( x , y ) . While autoregressive models may assign scores to prefixes of translations , we 1In this paper , we will use the term non-autoregressive to refer to any model whose scores do not decompose additively with the words in the output sequence . These include models that make conditional independence assumptions and generate each word independent of the others ( Lee et al. , 2018 ) , but also energy based models that require a complete translation hypothesis to compute a score , and models that make a Bayes ’ rule decomposition of the translation probability . only assume that s is well defined for full translations , y ∈ Y . In the particular case of machine translation under standard autoregressive models , the search objective is defined as the conditional log probability : sAR ( x , y ) = log p ( y | x ) = |y|∑ i=1 log p ( yi | y1 , . . . , yi−1 , x ) . As autoregressive models have probability emissions at the token level , the search problem can be cast into a shortest path problem for weighted graphs , where each node p is identified by a sequence y ( p ) . Nodes define an additional space Y∗ including Y as well as partial translations in Σ∗ that do not terminate with a EOS token . Each node contains |Σ|+ 1 edges , each appending a different word yi ∈ Σ ∪ { EOS } to y ( p ) generating a new node p ◦ yi . Edge weights are given by the emission log probability log p ( yi | y ( p ) , x ) according to the autoregressive model . Search starts from the root node ε with an empty sequence and nodes that generate the EOS token have a single edged with weight 0 that lead to a single shared terminal node . One solution to this problem is A∗ search ( Hart et al. , 1968 ) , which is a best-first algorithm that iteratively searches nodes p with the highest value v̂ ( p ) = c ( p ) + f ( p ) , where c ( p ) is the sum of weights of all edges from ε to p , and f ( p ) is a heuristic function that attempts to estimate the sum of the edges on the highest scoring path from p to the terminal node . Although theoretically appealing , good A∗ heuristics are difficult to obtain , and search can be extremely expensive with poor heuristics . To remedy this , beam search introduces approximations . Rather than a best-first traversal , beam search proceeds iteratively along a frontier at a certain depth . Exploration is limited to b nodes at each depth , where b is selected as a hyperparameter that trades search accuracy ( in terms of model score ) for speed . From depth 0 composed of just root node ε until a maximum depth limit Y is reached , or another stopping criteria is reached ( Klein et al. , 2017 ) , beam search progressively scores all children of the nodes at the current depth and prunes all generated nodes except for the top scoring b nodes . Beam search scores each node with only the current cost v̂ ( p ) = c ( p ) , and , in the case of neural machine translation , setting f ( p ) = 0 . Discarding the future cost biases search towards nodes with high scores without regard to whether they lead to a good path to the terminal state . Thus , a large space of potentially good translations with low initial scores is never explored . Examples include re-orderings that place high-entropy words at the start of the sentence or shorter sentence constructions ( e.g . “ Help me ” vs. “ Lend me a hand ” ) . Interestingly , autoregressive models tend to overestimate the probability of short and ungrammatical translations that do not translate the entirety of the source sentence , which are pruned by this scoring heuristic ( Stahlberg & Byrne , 2019 ; Holtzman et al. , 2019 ) . Thus , while b may be set low to increase speed , it is often set low to improve translation quality . However , we believe that model changes evaluated with beam search ’ s biases may be obscured . We seek to propose modeling improvements to mitigate degenerate solutions , such that search quality and model quality are aligned . 2.2 MONTE CARLO TREE SEARCH . In Monte Carlo tree search ( Coulom , 2006 , MCTS ) , the Monte Carlo method replaces the heuristically driven measure of value v̂ ( p ) for a node p with an expected-outcome model based on random game playouts . For instance , given a node p representing a state of a game , one can assign a value to p by randomly playing from that state a certain number of times and computing the average score obtained from the playouts . Thus , no burden is placed on the form of the objective function , allowing the definition of arbitrary complex objectives ( e.g . winner of a chess game ) . Additionally , as each playout yields a possible terminal state , both current c ( p ) and future costs f ( p ) are naturally embedded within the obtained estimate . Unlike beam search and A∗ search , where the search direction is determined by value v̂ ( p ) , MCTS diversifies the search space by allocating budget to less explored areas in the search space ( Kocsis & Szepesvári , 2006 ) , and continually refines value estimates .

Summary: This paper presents a novel algorithm called beam adaptive tree search (BATS) to enable the incorporation of search objectives that cannot be easily factorized. Unlike the regular Monte Carlo tree search algorithms that relies on a large number of playouts to update the value function, BATS relies on "informed playout", which is guided by the greedy decoding of the autoregressive models. To further constrain the search space, the paper also proposed a constrained node expansion criterion that gradually increases the lower limit of the node depth $d_{min}$, and allows for expanding $k$ nodes for each depth limit $d_{min}$.  Aside from the search algorithm, the paper also explored a couple of model changes to counteract the calibration issues, including two modified search objectives and using MRT-trained autoregressive models. Experiments show that while beam search performs better for the weaker models, BATS does improve the translation quality of the hypotheses when stronger MT models are used. Analysis on different beam size also show that unlike beam search, BATS can operate at large search budget without performance degradation under stronger models.