Minimum Bayes risk (MBR) decoding (Bickel and Doksum, 1977; Goel and Byrne, 2000) 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 Eikema and Aziz (2020) propose an sampling-based approximation algorithm.

MBR has been shown to outperform MAP beam search in both automatic and qualitative evaluation in a diverse range of tasks (Suzgun et al., 2023), including NMT (Freitag et al., 2022a) and code generation (Shi et al., 2022). MBR also generalizes other previously proposed decoding methods and explains their success (Bertsch et al., 2023).

The accuracy improvement from MBR comes at a heavy cost: the number of samples used can reach thousands (Freitag et al., 2023), 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++ (Popović, 2015), and COMET (Rei et al., 2020) 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.

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 $y_{<t}$ with a neural network $p_{\theta}$. This model can be used to assign probabilities to full sequences $p_{\theta}(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^{MAP}=\operatorname*{arg\,max}_{y}{p_{\theta}(y|x)}.$ The exact solution is generally intractable, and beam search is used to find an approximation.

In contrast, minimum Bayes risk decoding (MBR) (Goel and Byrne, 2000) outputs:

for some utility function $u$, a measure of similarity between sequences, and $U(y,\mathcal{Y})=\operatorname{\mathop{{}\mathbb{E}}}_{\hat{y}\sim\mathcal{Y}}[u(y,\hat{y})]$, where $\mathcal{Y}$ is either a probability distribution or an array of samples. We call $U$ the expected utility function. Eikema and Aziz (2020) propose a sampling method for neural language models where hypotheses $\mathcal{H}$ and pseudo-references $\mathcal{R}$ are generated with unbiased sampling, and $y^{MBR}$ is estimated as:

This method, which we refer to as "standard" MBR, requires $|\mathcal{H}|+|\mathcal{R}|$ samples (assuming $\mathcal{H}\neq\mathcal{R}$) and $|\mathcal{H}||\mathcal{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 (Eikema and Aziz, 2022; Freitag et al., 2023; Yan et al., 2022), and pruning $\mathcal{H}$ to a smaller size with a faster method prior to running standard MBR (Eikema and Aziz, 2022; Fernandes et al., 2022).

Freitag et al. (2022a) and Eikema and Aziz (2022) show that MBR is more effective relative to MAP when the utility metric has high segment-level accuracy as measured by Freitag et al. (2021). So, we use COMET (Rei et al., 2020), one of the best available metrics for NMT, and chrF++ (Popović, 2015), as a simpler and faster yet reasonably good lexical metric. We heed the call of Freitag et al. (2022b) and do not use BLEU, which is obsoleted by newer metrics as both an evaluation and utility metric (Freitag et al., 2022a).

Sampling-based MBR returns the highest-utility hypothesis from a set measured over pseudo-references 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 pseudo-references. 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 $\mathcal{H}_{1}$, and at each time step $t$, a pruning function uses a pseudo-reference list $\mathcal{R}_{t}$ of size $r_{t}$ to select $\mathcal{H}_{t+1}\subseteq\mathcal{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 $\mathcal{R}_{t}$ grows according to a pre-defined "schedule" $r_{1},...,r_{T}$.

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 $\operatorname*{arg\,max}_{y\in\mathcal{H}_{t}}U(y,p_{\theta}(\cdot|x))$, the true "winner".

We propose to prune hypotheses in $\mathcal{H}_{t}$ with low probability of being the true winner and to estimate this probability using nonparametric bootstrap resampling (Efron, 1979; Koehn, 2004); given an initial collection of i.i.d. samples $\mathcal{S}$ from an unknown distribution $\mathcal{X}$, our beliefs about the true value of any statistic $T(\mathcal{X})$ are represented by the distribution $p(T(\hat{\mathcal{S}}))$, where $\hat{\mathcal{S}}\sim\operatorname*{boot}(\mathcal{S})$ , and $\operatorname*{boot}(\mathcal{S})$ returns a with-replacement size-$|\mathcal{S}|$ resample of $\mathcal{S}$.

In our case, we want to estimate the probability that $y$ is the true winner in $\mathcal{H}_{t}$:

which we estimate as the chance that $y$ wins in a bootstrap resample. Let $\hat{\mathcal{R}_{t}}\sim\operatorname*{boot}(\mathcal{R}_{t})$. Then the bootstrap estimator is:

This estimator can be high variance because the probability $y$ winning in a bootstrap sample is very small when $\mathcal{H}_{t}$ is large, so instead we use the probability that $y$ outranks a particular $\bar{\bar{y}}\in\mathcal{H}_{t}$:

This statistic is invariant to the size of $\mathcal{H}_{t}$ because it only considers utility estimates of $y$ and $\bar{\bar{y}}$. It is an upper bound of Equation 4 because the probability of $y$ winning against all $\bar{y}\in\mathcal{H}_{t}$ cannot be higher than the probability of winning against a particular $\bar{\bar{y}}\in\mathcal{H}_{t}$. $\bar{\bar{y}}$ can be any element in $\mathcal{H}_{t}$, but we set it to the winner under $\mathcal{R}_{t}$, i.e. $\bar{\bar{y}}=\operatorname*{arg\,max}_{\bar{y}\in\mathcal{H}_{t}}U(\bar{y},\mathcal{R}_{t})$, to achieve a tighter upper bound.

We propose to prune $y$ if its estimated probability of beating $\bar{\bar{y}}$ is less than $1-\alpha$, where $\alpha$ is a confidence threshold $0\leq\alpha\leq 1$. In summary, this procedure prunes some but not necessarily all $y\in\mathcal{H}_{t}$ which are estimated to have less than $1-\alpha$ chance of being the true winner. 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 $\mathcal{H}_{t},\mathcal{R}_{t}$. Also note that the bootstrap estimator is always biased because $\mathcal{R}_{t}$ is never equal to $p_{\theta}(\cdot|x)$, and the bias is worse when $\mathcal{R}_{t}$ is small. Nonetheless, we show empirically that bootstrapping is effective in our pruning algorithm for modest sizes of $\mathcal{R}_{t}$.

Another benefit of our pruning method compared to standard MBR is that it can terminate early if $\mathcal{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 $y\in\mathcal{H}_{t}$ by $U(y,\mathcal{R}_{t})$ and prune the bottom-$\beta$ proportion. At $\beta=0$, no pruning occurs and standard MBR decoding is recovered. We refer to this baseline as $\operatorname*{prune}_{\beta}$ and our confidence-based method as $\operatorname*{prune}_{\alpha}$.

We perform our experiments on NMT models which we train on the German-English (de-en), English-Estonian (en-et), and Turkish-English (tr-en) 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 Eikema and Aziz (2022). 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 Vaswani et al. (2017) and are trained without label smoothing until convergence.

For all language pairs and validation/test datasets, we generate $\mathcal{H}_{1}$ with beam top-$k$ with $k=256$¹¹1Different sequences may have the same detokenization result, so we have fewer than 256 hypotheses on average.. We generate 1024 pseudo-references $\mathcal{R}^{\ast}$ with epsilon sampling (Hewitt et al., 2022) with $\epsilon=0.02$, following Freitag et al. (2023). In order to run multiple random trials efficiently, we simulate sampling pseudo-references by sampling from $\mathcal{R}^{\ast}$ without replacement. The experiments in Sections 4.1 and 4.2 show averaged results from 10 trials. We use chrF++ and COMET as utility functions and always match the evaluation metric to the utility. chrF++ is computed using SacreBLEU (Post, 2018) with default settings. COMET is computed with the COMET-22 model from Rei et al. (2022). We use 500 bootstrap samples for $\operatorname*{prune}_{\alpha}$.

For the pseudo-reference sample size schedule $r_{1},...,r_{T}$, the choice of $r_{1}$ is pivotal in the speed-accuracy trade-off; $|\mathcal{H}_{1}||\mathcal{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, $\operatorname*{arg\,max}_{\bar{y}\in\mathcal{H}}U(\bar{y},\mathcal{R}^{\ast})$ is pruned under different choices of $\alpha$ and $|\mathcal{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²²2https://github.com/juliusc/pruning_mbr.

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 sampling-based MBR without sacrificing accuracy. Our method is a drop-in substitute for standard MBR that requires no knowledge about the model $p_{\theta}$, how $\mathcal{H}_{1}$ is generated, or the utility function.

Even with our pruning algorithm, MBR is many times more costly to run than beam search.

An important hyperparameter in our method is the sample size schedule. We show why it is important to carefully choose the size of the first sample, but not how the remaining schedule should be set, opting to simply double the size at each step. We leave this issue to future work.

Methods such as MBR and reranking that directly optimize a metric may exploit noise in the metric to improve the score without actually improving quality (Fernandes et al., 2022). In these settings, automatic evaluation is less trustworthy and should ideally be combined with human evaluation. However, human evaluation is difficult and expensive to obtain.

Julius Cheng is supported by a scholarship from Huawei. The authors would like to thank the bilingual readers who helped with the human evaluation and the anonymous reviewers for their helpful suggestions.