Linear-time Minimum Bayes Risk Decoding with Reference Aggregation (2024)

Abstract: Minimum Bayes Risk (MBR) decoding is a text generation technique that has been shown to improve the quality of machine translations, but is expensive, even if a sampling-based approximation is used. Besides requiring a large number of sampled sequences, it requires the pairwise calculation of a utility metric, which has quadratic complexity. In this paper, we propose to approximate pairwise metric scores with scores calculated against aggregated reference representations. This changes the complexity of utility estimation from O(n) to O(n), while empirically preserving most of the quality gains of MBR decoding. We release our source code.1

heading: 1 Introduction
text: The idea of generating translations by maximizing a metric of translation quality (Kumar and Byrne, 2004) has recently been revived in the context of neural machine translation. In sampling-based MBR decoding (Eikema and Aziz, 2020), many hypotheses are sampled from the model distribution, and their expected utility is estimated using Monte Carlo (MC) sampling. This approach has been shown to improve translation quality compared to beam search, especially when neural metrics are used for utility estimation (Freitag et al., 2022).
Estimating utility through MC sampling has quadratic complexity in the number of samples, which limits practical application. Previous work suggested pruning the number of samples based on a cheaper metric or a smaller number of references (Eikema and Aziz, 2022; Cheng and Vlachos, 2023). In this paper, we propose reference aggregation, an alternative efficiency technique that exploits the fact that most common metrics represent text sequences in averageable form, e.g., as n-gram statistics or as embeddings. Specifically,
1https://github.com/ZurichNLP/mbr
we combine representations of the references into an aggregate reference representation, which we then use for utility estimation. Our proposed approximation still relies on MC sampling, but on a lower level: Rather than computing an MC estimate of the expected utility, we compute an MC estimate of the “true” reference representation in the feature space of the given utility metric. Since this estimate only needs to be computed once, our approach has linear complexity in the number of sampled hypotheses and references.
We report empirical results for four translation directions and two utility metrics: CHRF (Popović, 2015), which is based on character n-gram overlap, and COMET (Rei et al., 2020), a neural network trained with examples of human translation quality judgments. For CHRF, we find that reference aggregation reduces the time needed for computing the utility of 1024 samples by 99.5%, without affecting translation quality. For COMET, metric accuracy does decrease with aggregation, but to a lesser extent than with simply reducing the number of references. Depending on the COMET model, computation time is reduced by 95–99%, which makes reference aggregation an efficient method for hypothesis pruning with COMET.

heading: 2 Background and Related Work
text: Sampling-based MBR (Eikema and Aziz, 2020) selects a translation hyp∗ out of a set of translation hypotheses hyp1, . . . , hypn ∈ hyps by maximizing (expected) utility:
hyp∗ = argmax hyp∈ hyps utility(hyp). (1)
The set of hypotheses is sampled from the model distribution p(hyp|src). Eikema and Aziz (2020) propose to approximate the utility using MC sampling: sample a set of pseudo-references refs = {ref1, . . . , refm} ∼ p(ref|src) from the model and
790
calculate a metric against each sampled reference:
utility(hyp) ≈ 1 m
∑
ref∈ refs metric(hyp, ref). (2)
For machine translation, typical such metrics are CHRF (Popović, 2015) and BLEU (Papineni et al., 2002), which are based on n-gram statistics, or neural metrics such as COMET (Rei et al., 2020) and BLEURT (Sellam et al., 2020).
A line of research has focused on improving the efficiency of sampling-based MBR. Eikema and Aziz (2022) propose coarse-to-fine MBR, which prunes the hypotheses based on a cheaper metric, and N-by-S MBR, which uses fewer references than hypotheses. Cheng and Vlachos (2023) propose confidence-based pruning, where the number of hypotheses is iteratively reduced based on an increasing number of references. Jinnai and Ariu (2024) interpret sampling-based MBR as an instance of medoid identification and apply an established approximation algorithm to this problem. A line of work uses MBR outputs as a training reward, avoiding the inefficiency of MBR during deployment (Finkelstein et al., 2023; Yang et al., 2023). Finally, alternative reranking approaches that do not require pairwise comparisons have been proposed (Fernandes et al., 2022).
Several other works investigate the aggregation of reference representations to develop a faster variant of MBR decoding. DeNero et al. (2009) perform reference aggregation in the context of statistical machine translation (SMT). Since SMT does not afford random sampling of pseudo-references, they aggregate references from translation forests or k-best lists. Our study shows the effectiveness of reference aggregation from sampled pseudoreferences, and for neural metrics such as COMET. Furthermore, concurrent to our work, Deguchi et al. (2024) propose to aggregate the sentence embeddings of COMET, and use k-means to group the references into multiple clusters.

heading: 3 Reference Aggregation
text: Our approach is based on the observation that most metrics that are commonly used for MBR make use of feature representations that can be aggregated. For example, the n-gram statistics used by CHRF can be aggregated by averaging the counts of the n-grams across all references; and the sentence embeddings used by COMET can be aggregated by calculating an average sentence embedding.
For simplicity, we re-use the above notation, where hyp is a hypothesis and ref is a reference, but we now assume that they are represented in an averageable form. We then combine the set of references refs into an aggregate representation ref:
ref = 1
m
∑
ref∈ refs ref. (3)
We approximate the expected the utility of a sampled hypothesis by calculating a single metric score against this aggregate representation:
utility(hyp) ≈ metric(hyp, ref). (4)
Like with standard sampling-based MBR, it is possible to interpret this approximation as MC sampling: By averaging over representations of sampled references, we estimate a representation of the “true” reference, which we then use for approximating the expected utility of each sampled hypothesis. Importantly, the computational complexity of our approach is in O(|hyps| + |refs|) rather than O(|hyps| · |refs|); see Appendix D for a discussion.

heading: 3.1 Application to chrF Metric
text: CHRF (Popović, 2015) is defined as an F-score over character n-grams:
CHRFβ = (1 + β2) · CHRP · CHRR β2 · CHRP + CHRR , (5)
where
CHRP = |hyp ∩ ref| |hyp| and CHRR = |hyp ∩ ref| |ref| ,
and the parameter β controls the relative importance of precision and recall. The representations hyp and ref are bags of n-grams, i.e., objects that map each n-gram to its count in the string.
We apply reference aggregation to CHRF by averaging the counts of n-grams across all references:
ref = 1
m
⊎
ref∈ refs ref, (6)
where ⊎
is an operation that sums up the counts of each n-gram. We then approximate the expected utility of a hypothesis by calculating CHRFβ(hyp, ref). Appendix A provides a more formal definition of reference aggregation for CHRF.

heading: 3.2 Application to COMET Metric
text: COMET (Rei et al., 2020) is a pre-trained Transformer model (Vaswani et al., 2017) that has been fine-tuned to predict human judgments of translation quality. In this paper, we focus on the Estimator model architecture, which directly estimates a quality score given a hypothesis, a reference and the source sequence. COMET separately encodes these three inputs into fixed-size embeddings:
hyp, ref , src = emb(hyp), emb(ref), emb(src).
The three embeddings are then fed into a feedforward module, which outputs a scalar score:
comet(hyp) = score(hyp, ref , src). (7)
We apply reference aggregation to COMET by averaging the reference embeddings:
ref = 1
m
∑
ref∈ refs emb(ref), (8)
calculating a single score per hypothesis:
comet(hyp) ≈ score(hyp, ref , src). (9)

heading: 3.3 Partial Aggregation
text: To better understand the loss of accuracy incurred by aggregation, we experiment with partial aggregation, where we vary the number of references
that are combined into an average. Given m references and a desired number of references s that should effectively be used for pairwise utility estimation, we partition the set of references into s subsets and create an aggregate reference for each subset. Appendix B presents a formal description of partial aggregation.

heading: 3.4 Aggregate-to-fine MBR
text: Analogously to coarse-to-fine MBR (Eikema and Aziz, 2022), we evaluate an aggregate-to-fine MBR approach. Specifically, we use the aggregate reference to prune the number of hypotheses to 20 in a first step. In a second step, we use standard MBR to select the best hypothesis from the pruned set. A formal description is provided in Appendix C.

heading: 4 Experimental Setup
text: Data We use newstest21 (Akhbardeh et al., 2021) as validation data and newstest22 (Kocmi et al., 2022) as test data.
Generation Parameters As baselines, we evaluate beam search with a beam size of 5 and epsilon sampling (Hewitt et al., 2022) with ϵ = 0.02. For MBR, we generate 1024 samples per segment using epsilon sampling and re-use the same samples as references. While this approach does not guarantee
that the estimation of the expected utility is unbiased (Eikema and Aziz, 2022), it has empirically been found to work well (Freitag et al., 2023).
Models We use open-source NMT models trained for the EN–DE, DE–EN, EN–RU and RU– EN translation directions (Ng et al., 2019).2 The authors provide an ensemble of four models per direction, but we restrict our experiments to one single model per direction. We use the Fairseq codebase (Ott et al., 2019) for model inference.
Metrics For estimating the utilities with CHRF, we use a custom implementation of CHRF3 that is equivalent to SacreBLEU (Post, 2018) with default settings4. As COMET model, we use COMET-22 (Rei et al., 2022a); because this model was not trained on annotations of newstest21 or newstest22, a train–test overlap can be ruled out. We estimate wall-clock time based on a part of the segments, using a system equipped with an NVIDIA GeForce RTX 3090 and an AMD EPYC 7742 64-core processor.
2The models were trained with a label smoothing of ϵ = 0.1 (Szegedy et al., 2016), which is a common choice in NMT. Some previous studies of MBR trained custom models without label smoothing (e.g., Eikema and Aziz, 2020). We argue that this is only necessary if unbiased utility estimates are sought through ancestral sampling, and should be less of a concern with epsilon sampling.
3https://github.com/jvamvas/fastChrF 4chrF2|#:1|case:mixed|eff:yes|nc:6|nw:0|space:no|v:2.0.0

heading: 5 Results

heading: 5.1 Validation results
text: Figure 1 evaluates how accurately MBR efficiency methods approximate standard MBR. We report top-20 accuracy, motivated by the idea of coarse-to-fine MBR: any method with perfect top-20 accuracy could be used for pruning the hypothesis set to 20 without affecting quality. Results for top-1 accuracy are reported in Appendix I.5
For CHRF, we observe that reference aggregation is Pareto superior to N-by-S, maintaining nearperfect top-20 accuracy even if a single aggregate reference is used. For COMET, reference aggregation causes some loss of accuracy, but outperforms N-by-S if the number of effective references is ≤ 16, where efficiency is highest. In addition, we find that reference aggregation approximates standard (pairwise) COMET much better than using CHRF as a coarse metric does, providing a clear motivation for aggregate-to-fine MBR as an alternative to coarse-to-fine MBR.

heading: 5.2 Test results
text: In Table 1, we report test results for newstest22, focusing on a comparison between fast baseline algorithms (beam search and sampling) and MBR (with or without reference aggregation). We perform an automatic evaluation using BLEURT-20 (Sellam et al., 2020), chosen because it is unrelated to the utility metrics we use for MBR. CHRF and
5Accuracy was proposed by Cheng and Vlachos (2023) as an evaluation metric for MBR efficiency methods.
COMET scores are reported in Appendix F. The results show that reference aggregation narrows the efficiency gap between MBR and beam search while preserving most of the quality gain of standard MBR. Reference aggregation speeds up utility estimation by 99.5% for CHRF and 95.1% for COMET-22, reducing the total time needed for translation by 75.5% and 78.8%, respectively. Using an aggregate-to-fine approach has a lower loss of quality and still reduces the total translation time by 73.6–77.4%.
Reference aggregation is thus a successful strategy to overcome the quadratic complexity of MBR. However, it is still slower than beam search, as the cost of sampling is now the dominant factor. Future work could focus on sampling efficiency, e.g., by using fewer hypotheses, improved caching, or speculative sampling approaches (Leviathan et al., 2023; Chen et al., 2023).

heading: 6 Conclusion
text: We proposed reference aggregation, a technique that boosts the efficiency of MBR decoding by shifting the MC sampling from the utility estimation to the reference representation. Experiments on machine translation showed that reference aggregation speeds up utility estimation by up to 99.5% while minimally affecting translation quality. This reduces the gap to beam search and makes MBR more practical for large-scale applications.
Limitations
This work has two main limitations:
1. Reference aggregation requires a utility metric based on averageable representations.
2. For trained metrics, the effectiveness of aggregation needs to be evaluated empirically.
We have demonstrated that reference aggregation is a viable technique for MBR with CHRF and COMET, leading to a considerable speed-up with minor quality losses. In the case of CHRF, reference aggregation entails a slight modification of the metric definition, but is otherwise exact and not an approximation. We thus expect that reference aggregation could be applied in a straightforward manner to other lexical overlap metrics such as CHRF++ (Popović, 2017) and BLEU (Papineni et al., 2002).
For COMET, which is a trained metric, reference aggregation involves the averaging of fixedsize sentence embeddings. We empirically studied the loss of accuracy incurred by this averaging and found that there is a favorable trade-off between speed and accuracy for the COMET models we evaluated. We recommend that future work validates the effectiveness of reference aggregation for other trained metrics.
While CHRF and COMET are among the most commonly used metrics for MBR, previous work has also proposed metrics that are not based on averageable reference representations. For example, BLEURT (Sellam et al., 2020), a trained metric that was shown to be effective for MBR (Freitag et al., 2022), is based on a cross-encoder architecture that creates a joint representation for each hypothesis–reference pair. Future work could investigate in what form, if at all, reference aggregation can be applied to cross-encoder architectures.
Finally, this work studies MBR decoding with a classical sequence-to-sequence NMT model and in the context of sentence-level MT. While MBR decoding has also been successfully applied to MT with large language models (Farinhas et al., 2023), more research is needed on MBR decoding with large language models, especially on the document level.

heading: Acknowledgments
text: We thank Clara Meister and Bryan Eikema for helpful discussions and feedback. This work was funded by the Swiss National Science Foundation (project MUTAMUR; no. 213976).

heading: A Formal Definition of Reference Aggregation for ChrF
text: The CHRF metric (Popović, 2015) is a harmonic mean of precision and recall scores:
CHRFβ = (1 + β2) · CHRP · CHRR
β2 · CHRP + CHRR . (10)
Internally, CHRF converts hypotheses and references into bags of character n-grams. Such bags can be represented as multisets (Knuth, 1997, Section 4.6.3) or as (sparse) vectors. We will use vector notation in this formal definition, which allows us to define reference aggregation with standard vector operations.
Let hyp ∈ R|V| and ref ∈ R|V| be bags representing a hypothesis and a reference, where V is the vocabulary of all character n-grams up to maximum order n, and the entries hypj and ref j are the counts of n-gram j ∈ V in the hypothesis and reference, respectively.
For a given n-gram order i ∈ {1, . . . , n}, precision and recall are defined as:
CHRPi(hyp, ref) = ∑
j∈Vi min(hypj , ref j)∑ j∈Vi hypj ,
(11)
CHRRi(hyp, ref) = ∑
j∈Vi min(hypj , ref j)∑ j∈Vi ref j ,
(12)
where Vi is the set of all character n-grams of order i. Overall precision and recall are calculated as the arithmetic mean of the precision and recall scores for each n-gram order:
CHRP(hyp, ref) = 1
n
n∑
i=1
CHRPi(hyp, ref),
(13)
CHRR(hyp, ref) = 1
n
n∑
i=1
CHRRi(hyp, ref).
(14)
When CHRF is used as a utility metric in a standard MBR setting, the expected utility of a hypothesis is estimated based on a set {ref (1), . . . , ref (m)} of m references:
utility(hyp) = 1
m
m∑
k=1
CHRFβ(hyp, ref (k)).
(15)
In contrast, reference aggregation first calculates the arithmetic mean of the reference bags:
ref = [ 1
m
m∑
k=1
ref (k)1 , . . . , 1
m
m∑
k=1
ref (k)|V| ], (16)
and estimates the utility as:
utilityagg(hyp) = CHRFβ(hyp, ref). (17)
Note that the only mathematical difference between pairwise calculation of chrF and using the aggregate reference is that the F-score is averaged across sentences in the pairwise calculation, and computed over the global precision and recall with reference aggregation.

heading: B Formal Definition of Partial Aggregation
text: We conceptualize partial aggregation as follows:
1. The set of individual references contains m references. 2. We randomly partition the set of references into s groups of equal size. 3. Each group is combined into an average reference representation, resulting in s aggregate references ref(i), . . . , ref(s).
The expected utility of each sampled hypothesis is then approximated as the average metric score over all aggregate references:
utility(hyp) ≈ 1 s
s∑
i=1
metric(hyp, ref(i)). (18)
Like with N-by-S MBR, the parameter s can be seen as the number of effective references that determines the computational complexity of the utility estimation. The case s = m corresponds to standard MBR, where each sampled hypothesis is compared to each reference in a pairwise fashion. The case s = 1 corresponds to the full aggregation approach, where a single aggregate reference is created from all references.

heading: C Formal Definition of Aggregate-to-fine MBR
text: Aggregate-to-fine MBR is a special case of coarseto-fine MBR (Eikema and Aziz, 2022), which uses a cheap proxy utility function to prune the number of hypotheses. In the case of aggregate-to-fine MBR, the proxy utility function is based on an aggregate reference representation.
The general definition of coarse-to-fine MBR is as follows: Given the original set of sampled hypotheses H̄(x) and a proxy utility function uproxy, coarse-to-fine MBR selects a subset of T hypotheses:
H̄T (x) := top-T hyp∈H̄(x) uproxy(hyp). (19)
In the second step, the utility of each hypothesis in the pruned set is estimated using the fine-grained utility function utarget:
yC2F := argmax hyp∈H̄T (x) utarget(hyp). (20)
When experimenting with aggregate-to-fine MBR, we re-use the same utility metric for both steps, but first with an aggregate reference and then with the full set of references:
uproxy(hyp) = metric(hyp, ref), (21)
utarget(hyp) = 1
m
∑
ref∈refs metric(hyp, ref). (22)
Note that using the same metric in both steps is not strictly necessary, but has the advantage that the features (e.g., embeddings) only need to be computed once.

heading: D Complexity Analysis
text: Generally, reference aggregation reduces the complexity of utility estimation from O(nm) to O(n+ m), where n is the number of hypotheses and m is the number of references. The exact complexity depends on the specifics of the utility metric. Here, we provide a more detailed analysis for CHRF and COMET.
Above, we stated that utility estimation with these metrics usually has two stages: feature extraction and scoring. The feature extraction stage is not affected by reference aggregation, and previous work has already remarked that reference features
can be extracted once and re-used for all hypotheses (Amrhein and Sennrich, 2022). If the reference set is identical to the set of hypotheses, the feature extraction stage is in O(n), otherwise O(n+m).
The scoring stage of CHRF is dominated by the element-wise minimum function in Eqs. 11 and 12 (or, if the bags of n-grams are represented as multisets, by the intersection operation hyp ∩ ref). Because this operation is performed separately for each hypothesis–reference pair, the complexity is in O(nm). Reference aggregation reduces the complexity to O(n+m), given that the aggregate reference can be computed once and then re-used for all hypotheses.6
The same analysis applies to COMET. With standard MBR, Eq. 7 is evaluated for each hypothesis–reference pair; with reference aggregation, it is only evaluated once for each hypothesis. The aggregate reference embeddings can be computed once and re-used for all hypotheses.
In practice, the runtime of utility estimation is affected by additional factors. There may be duplicates among the samples, so the number of scores that effectively need to be computed can vary. In addition, most aspects of utility estimation can be computed in parallel, which makes the effective runtime highly implementation-dependent.
6For CHRF, reference aggregation can result in an aggregate bag of n-grams that is larger that the bags of the individual references; in the theoretical worst case, where all the references are disjoint, even in an aggregate bag that is m times larger. However, this is a highly unlikely scenario in practice, since different translations of the same source will have substantial overlap, and even if |ref| ≫ |ref|, the cost of intersection only depends on |hyp|, assuming that a constanttime hash table is used to check whether each item in hyp is contained in ref.
E Data Statistics
F Extended Test Results

heading: H Effect of Larger Beam Size

heading: G Test Results with 256 Samples

heading: I Top-1 Accuracy of Efficiency Methods
text: Utility metric: CHRF
J Validation Results for Cometinho
Top-20 accuracy
n 1 n 2 n 4 n 8 . . . 1
0%
50%
100%
Number of eff. references (log scale)
To p-
20 ac
cu ra
cy EN–DE
Reference aggregation N-by-S CHRF
n 1 n 2 n 4 n 8 . . . 1
DE–EN
n 1 n 2 n 4 n 8 . . . 1
EN–RU
n 1 n 2 n 4 n 8 . . . 1
RU–EN
Top-1 accuracy

heading: Future Work
text: We recommend that future work validates the effectiveness of reference aggregation for other trained metrics. While CHRF and COMET are among the most commonly used metrics for MBR, previous work has also proposed metrics that are not based on averageable reference representations. For example, BLEURT (Sellam et al. , 2020), a trained metric that was shown to be effective for MBR (Freitag et al. , 2022), is based on a cross-encoder architecture that creates a joint representation for each hypothesis–reference pair. Future work could investigate in what form, if at all, reference aggregation can be applied to cross-encoder architectures. Finally, this work studies MBR decoding with a classical sequence-to-sequence NMT model and in the context of sentence-level MT. While MBR decoding has also been successfully applied to MT with large language models (Farinhas et al. , 2023), more research is needed on MBR decoding with large language models, especially on the document level