Faster Minimum B ayes Risk Decoding with Confidence-based Pruning (2023)

Abstract: Minimum Bayes risk (MBR) decoding outputs the hypothesis with the highest expected utility over the model distribution for some utility function. It has been shown to improve accuracy over beam search in conditional language generation problems and especially neural machine translation, in both human and automatic evaluations. However, the standard sampling-based algorithm for MBR is substantially more computationally expensive than beam search, requiring a large number of samples as well as a quadratic number of calls to the utility function, limiting its applicability. We describe an algorithm for MBR which gradually grows the number of samples used to estimate the utility while pruning hypotheses that are unlikely to have the highest utility according to confidence estimates obtained with bootstrap sampling. Our method requires fewer samples and drastically reduces the number of calls to the utility function compared to standard MBR while being statistically indistinguishable in terms of accuracy. We demonstrate the effectiveness of our approach in experiments on three language pairs, using chrF++ and COMET as utility/evaluation metrics.

## introduction minimum bayes risk (mbr) decoding @xcite has recently gained renewed attention as a decision rule for conditional sequence generation tasks, especially neural machine translation (nmt). in mbr, the sequence with the highest expected utility with respect to thez model distribution is chosen as the output, where the utility is usually some measure of text similarity. this contrasts with the more commonly used maximum a posteriori (map) decision rule, which returns the sequence with the highest probability under the model. map is generally intractable, and beam search is typically used to find an approximation. mbr is likewise intractable, and @xcite propose an samplingbased approximation algorithm. mbr has been shown to outperform map beam search in both automatic and qualitative evaluation in a diverse range of tasks @xcite , including nmt @xcite and code generation @xcite . mbr also generalizes other previously proposed decoding methods and explains their success @xcite . the accuracy improvement from mbr comes at a heavy cost: the number of samples used can reach thousands @xcite , and the number of calls to the utility function required is quadratic in the number of samples. often, the utility function itself is a deep neural model, rendering mbr prohibitively expensive for many use cases. in this work, we address the computational efficiency of mbr with an iterative pruning algorithm where low-performing hypotheses are removed while the number of samples used to estimate utilities grows. hypotheses are pruned based on their estimated probability of being the true best hypothesis under the mbr objective, thus avoiding making expensive fine-grained utility estimates for hypotheses which are unlikely to be the final prediction. in nmt experiments on three language pairs using chrf++ @xcite , and comet @xcite as mbr utility and evaluation metrics, we show that our method maintains the same level of accuracy as standard mbr while reducing the number of utility calls by a factor of at least 7 for chrf++ and 15 for comet. our algorithm can also use fewer samples to reach a prediction by terminating early, unlike standard mbr. ## minimum bayes risk decoding conditional sequence generation problems such as neural machine translation (nmt) model the probability of the next token y t given a source sequence x and prefix @xmath0. this model can be used to assign probabilities to full sequences p θ (y|x) via the chain rule. at test time, a decision rule is employed to select a single "best" sequence. the most common decision rule is to return the highest probability sequence y m @xmath1. the exact solution is generally intractable, and beam search is used to find an approximation. in contrast, minimum bayes risk decoding (mbr) @xcite outputs: for some utility function u, a measure of similarity between sequences, and u (y, y) = e ŷ∼y [u(y, ŷ)], where y is either a probability distribution or an array of samples. we call u the expected utility function. @xcite propose a sampling method for neural language models where hypotheses h and pseudo-references r are generated with unbiased sampling, and y m br is estimated as: this method, which we refer to as "standard" mbr, requires |h| + |r| samples (assuming h ̸ = r) and |h||r| calls to u. the latter is the main computational bottleneck which we address in this work. recent works on mbr focus on identifying accurate and efficient generation methods @xcite @xcite , and pruning h to a smaller size with a faster method prior to running standard mbr @xcite . @xcite and @xcite show that mbr is more effective relative to map when the utility metric has high segmentlevel accuracy as measured by @xcite . so, we use comet @xcite , one of the best available metrics for nmt, and chrf++ @xcite , as a simpler and faster yet reasonably good lexical metric. we heed the call of @xcite and do not use bleu, which is obsoleted by newer metrics as both an evaluation and utility metric @xcite . ## confidence-based hypothesis pruning sampling-based mbr returns the highest-utility hypothesis from a set measured over pseudoreferences sampled from the model. speedups can be achieved if low-performing hypotheses are removed from consideration based on coarse utility estimates obtained from a subset of the pseudoreferences. in other words, we can save time by not computing precise utility estimates for hypotheses which are unlikely to be chosen in the end. we propose an iterative algorithm for mbr where the hypothesis set is gradually shrunk while the pseudo-reference list grows. the procedure is shown in algorithm 1. we start with an initial hypothesis set h 1 , and at each time step t, a pruning function uses a pseudo-reference list r t of size r t to select h t+1 ⊆ h t . after the maximum time step is reached or when the current hypothesis set contains one element, terminate and return the highest utility hypothesis under all available pseudo-references. the size of r t grows according to a pre-defined "schedule" r 1 , ..., r t . algorithm 1 pruning mbr input: source sentence x. constants: sample size schedule @xmath2, ..., r t , expected utility function u , model parameters θ, pruning function prune, hypothesis generation function gen(x, θ). output: an mbr prediction. 1: r 0 ← [ ] 2: t ← 1 3: h 1 ← gen(x, θ) 4: while @xmath3: r t ← r t-1 6: while |r t | < r t do 7: append ŷ ∼ p θ (•|x) to r t 8: end while 9: the goal of the pruning function is to exclude as many hypotheses as possible to reduce the number of utility calls made in the future without excluding the true top-ranking mbr hypothesis arg max @xmath4, p θ (•|x)), the true "winner". we propose to prune hypotheses in h t with low probability of being the true winner and to estimate this probability using nonparametric bootstrap resampling @xcite ; given an initial collection of i.i.d. samples s from an unknown distribution x , our beliefs about the true value of any statistic t (x ) are represented by the distribu-tion p(t ( ŝ)), where ŝ ∼ boot(s) , and boot(s) returns a with-replacement size-|s| resample of s. in our case, we want to estimate the probability that y is the true winner in h t : this estimator can be high variance because the probability y winning in a bootstrap sample is very small when h t is large, so instead we use the probability that y outranks a particular ȳ ∈ h t : this statistic is invariant to the size of h t because it only considers utility estimates of y and ȳ. it is an upper bound of equation 4 because the probability of y winning against all ȳ ∈ h t cannot be higher than the probability of winning against a particular ȳ ∈ h t . ȳ can be any element in h t , but we set it to the winner under r t , i.e. ȳ = arg max ȳ∈ht u (ȳ, r t ), to achieve a tighter upper bound. we propose to prune y if its estimated probability of beating ȳ is less than 1α, where α is a confidence threshold 0 ≤ α ≤ 1. in summary, this procedure prunes some but not necessarily all @xmath5. algorithm 2 shows the procedure in detail. note that bootstrapped utility estimates do not require additional utility function calls because they reuse utility values already computed over h t , r t . also note that the bootstrap estimator is always biased because r t is never equal to p θ (•|x), and the bias is worse when r t is small. nonetheless, we show empirically that bootstrapping is effective in our pruning algorithm for modest sizes of r t . another benefit of our pruning method compared to standard mbr is that it can terminate early if h t has only one remaining hypothesis, reducing the total number of pseudo-references needed. as a baseline pruning function for comparison, we rank each @xmath6, r t ) and prune the algorithm 2 confidence-based pruning function input: hypothesis set h, pseudo-reference list r. constants: expected utility function u , confidence threshold α, number of bootstrap samples n. output: a subset of h. ri ← boot(r) 6: end for 7: for @xmath7: ## experiments we perform our experiments on nmt models which we train on the german-english (de-en), english-estonian (en-et), and turkish-english (tren) news datasets from wmt18. we use the data preprocessing steps provided by the wmt18 organizers, except that we exclude paracrawl from the de-en dataset following @xcite . the final training sets have 5.8 million, 1.9 million, and 207 thousand sentence pairs respectively. all models are transformers of the base model size from @xcite and are trained without label smoothing until convergence. for all language pairs and validation/test datasets, we generate h 1 with beam top-k with @xmath8. we generate 1024 pseudo-references r * with epsilon sampling @xcite with ϵ = 0.02, following @xcite . in order to run multiple random trials efficiently, we simulate sampling pseudo-references by sampling from r * without replacement. the experiments in sections 4.1 and 4.2 show averaged results from 10 trials. we use chrf++ and comet as utility func-tions and always match the evaluation metric to the utility. chrf++ is computed using sacrebleu @xcite with default settings. comet is computed with the comet-22 model from @xcite . we use 500 bootstrap samples for prune α . for the pseudo-reference sample size schedule r 1 , ..., r t , the choice of r 1 is pivotal in the speedaccuracy trade-off; |h 1 ||r 1 | is a lower bound on the number of utility calls needed, but the bootstrap estimate is more biased when the sample size is small. in a preliminary experiment on the validation set, we measure the "false pruning rate", the rate that the estimated true winner, arg max ȳ∈h u (ȳ, r * ) is pruned under different choices of α and |r|. based on results shown in figure 1 , we set r 1 to 8 for comet and 16 for chrf++ for all experiments. r 2 , ..., r t are set by doubling at each time step until reaching 256. more experimental details and figures for language pairs not shown here are in the appendix. our code is publicly available foot_2 . ## conclusion we propose an iterative pruning algorithm for mbr along with a pruning criterion based on confidence estimates derived from bootstrap resampling. in experiments across diverse language pairs and metrics, we show that our method consistently outperforms our proposed baseline and achieves significant computational savings over standard samplingbased mbr without sacrificing accuracy. our method is a drop-in substitute for standard mbr that requires no knowledge about the model p θ , how h 1 is generated, or the utility function.