Centroid-Based Efficient Minimum Bayes Risk Decoding (2024)

Abstract: Minimum Bayes risk (MBR) decoding has achieved state-of-the-art translation performance using COMET, which is a neural metric that has a high correlation with human evaluation. However, MBR decoding requires quadratic time because it computes the expected score between a translation hypothesis and all reference translations. We propose centroid-based MBR (CBMBR) decoding to improve the speed of MBR decoding. Our method clusters reference translations in the feature space and then calculates the score using the centroids of each cluster. The experimental results demonstrate that our CBMBR not only improved the decoding speed of the expected score calculation by 5.7 times but also outperformed vanilla MBR decoding in terms of translation quality by up to 0.5 COMET% in the WMT’22 En↔Ja, En↔De, En↔Zh, and WMT’23 En↔Ja translation tasks.1

heading: 1 Introduction
text: Minimum Bayes risk (MBR) decoding achieves robust and high-quality translation by selecting the output sentence that maximizes the expected metric score computed from the set of translation hypotheses (Kumar and Byrne, 2004; Eikema and Aziz, 2020; Müller and Sennrich, 2021). Recently, neural evaluation metrics that have a high correlation with human evaluation have been proposed (Rei et al., 2020, 2022a; Sellam et al., 2020; Zhang et al., 2020), and MBR decoding using such neural metrics has achieved state-of-the-art translation performance in human evaluation compared with conventional maximum-a-posteriori (MAP) decoding using beam search (Fernandes et al., 2022).
However, because of its formulation, typical MBR decoding that regards the hypothesis set as a pseudo-reference set requires the computational time of O(N2) when N translation hypotheses are
1https://github.com/naist-nlp/mbrs
given. In recent work, the number of hypotheses N exceeded 1,000 candidates (Freitag et al., 2023), which makes the quadratic order of computational time a challenge for MBR decoding, particularly when expensive neural metrics are used. Several pruning methods have been proposed (Eikema and Aziz, 2022; Cheng and Vlachos, 2023) to improve the decoding speed. These approaches require the careful selection of a proxy metric (Eikema and Aziz, 2022), or it is difficult to take advantage of computational parallelism because hypotheses are pruned iteratively (Cheng and Vlachos, 2023).
Given that expensive neural metrics, e.g., COMET or BLEURT, are trained to output high scores when a hypothesis sentence and reference translation are semantically similar, we hypothesize that the distance between sentence vectors of similar sentences in the feature space of their models is close. We leverage sentence similarity to improve the decoding speed of COMET-MBR by clustering sentence vectors of the translation candidates into k ≪ N clusters as shown in Figure 1. Then, we calculate the COMET scores using k centroid vectors of their clusters, instead of using N sentence vectors.
In experiments, our proposed method not only achieved a speed-up of 5.7 times in the calculation of the expected score but also an improvement in the COMET score of up to 0.5% compared with naive MBR decoding in the WMT’22 En↔Ja, En↔De, En↔Zh, and WMT’23 En↔Ja transla-
11009
tion tasks.

heading: 2 Background
text: MBR decoding MBR decoding has been demonstrated to be effective in fields such as statistical automatic speech recognition (Goel and Byrne, 2000) and statistical machine translation (Kumar and Byrne, 2004). In recent years, it has been applied to neural machine translation (Eikema and Aziz, 2020; Müller and Sennrich, 2021). Furthermore, it is more suitable for multiple translation systems than ensemble models (Ito et al., 2023).
Let X and Y be the spaces of possible source sentences and target sentences, respectively. MAP decoding generates the target sentence y∗MAP ∈ Y using y∗MAP = argmaxy∈Y pθ(y|x), where θ denotes the parameter of the translation model that calculates the likelihood of an output sentence y given an input sentence x ∈ X . Because it is infeasible to calculate probabilities for all possible y ∈ Y , beam search is typically used to obtain the solution.
By contrast, MBR decoding determines the output sentence y∗MBR ∈ Y by maximizing the expected utility as follows:
y∗MBR = argmaxh∈H Eŷ∼P (y|x) [u(h, ŷ)] , (1) ≈ argmaxh∈H Eŷ∈Ŷ [u(h, ŷ)] , (2)
where u : Y × Y → R denotes the utility function, which represents the preference relation and H = {hi}|H|i=1 ⊂ Y denotes the set of translation hypotheses. P (y|x) is the true probability of being translated from a given input sentence x ∈ X and is approximated using the sampled reference translations Ŷ = {ŷi}|Ŷ|i=1 ⊂ Y , as shown in Equation 2, because the true probability is unknown. Typical MBR decoding considers the hypothesis set itself as the pseudo-reference set, i.e., Ŷ := H. Note that the time complexity is O(N2), where N := |H|, which is time-consuming.
COMET-MBR COMET is an evaluation metric of translation quality that achieves a high correlation with human evaluation. The COMET model consists of the XLM-RoBERTa (XLM-R)-based sentence encoder (Conneau et al., 2020) and the output layer, and is trained to predict direct assessment scores (Rei et al., 2020, 2022a). It first encodes the source sentence x ∈ X , hypothesis sentence h ∈ Y , and reference sentence ŷ ∈ Y
into their D-dimensional sentence vectors independently. Then the COMET score is computed from the triplet of sentence vectors in the output layer. Let f : X ∪ Y → RD be the function of sentence encoding and s : RD × RD × RD → R be the output layer. The COMET score is computed using s(f(x), f(h), f(ŷ)). MBR decoding with COMET (COMET-MBR) replaces the utility u in Equation 2 with the COMET score:
y∗COMET-MBR = argmaxh∈H Eŷ∈Ŷ [s(f(x), f(h), f(ŷ))] . (3)

heading: 3 Proposed Method
text: Our proposed centroid-based MBR (CBMBR) approximates the expected utility using the centroids of similar sentence vectors. CBMBR decodes by computing the expected utility according to the following procedures: sentence encoding, clustering, and calculating the expected utility.
Encoding First, we compute the sentence vector of the source f(x) ∈ RD, the hypotheses {f(hi)}|H|i=1 ⊂ RD, and the pseudo-references {f(ŷi)}|Ŷ|i=1 ⊂ RD. Clustering Next, we cluster the sentence vectors of the pseudo-references into k ≪ N clusters and obtain the centroid vectors of each cluster C = {ci}ki=1 ⊂ RD. Here, we employ kmeans++ (Arthur and Vassilvitskii, 2007) to prevent the centroids from being biased. kmeans++ selects the initial centroids so that the distances between each pair of centroids are farther according to the weights calculated from the distances between vectors. We describe the details of the algorithm in Appendix C.1. Then, we cluster the vectors {f(ŷi)}|Ŷ|i=1 using the standard kmeans algorithm. Specifically, we calculate the following steps iteratively: 1) assign a vector to its nearest neighbor centroid and 2) update the centroid using the vectors assigned to its cluster.
Expected utility Finally, we calculate the expected utility by replacing pseudo-reference vectors f(ŷ) ∈ RD with centroids c ∈ RD in Equation 3:
y∗CBMBR = argmaxh∈H Ec∈C [s(f(x), f(h), c)] . (4)
The conventional method requires O(N2) of computational time to compute the expected utility for all hypotheses, whereas our CBMBR computes it in O(Nk). Note that k (1 ≤ k ≤
N ) is a hyperparameter that balances the tradeoff between the decoding speed and approximation accuracy. Especially, when k = 1, i.e., C = {c1}, the centroid c1 can be calculated as the average of all pseudo-reference vectors, i.e.,
c1 = 1 |Ŷ| ∑|Ŷ| i=1 f(ŷi), and the time complexity of CBMBR is O(N), which is equivalent to DeNero et al. (2009) and Vamvas and Sennrich (2024). Our proposed method approximates the expected utility using centroid representations, which implicitly assumes that the score function s is approximately linear. Specifically, when k = 1, we approximate the expected utility Eŷ∈Ŷ [s(f(x), f(h), f(ŷ))] as s(f(x), f(h),Eŷ∈Ŷ [f(ŷ)]) in our method by assuming that the score function s is roughly linear.

heading: 4 Experiments
text: Setup We conducted translation experiments with two settings; one used diversified translation candidates and the other simulated a more realistic scenario, multi-system translation. We evaluated the translation quality using the COMET score, which is the same as the utility function of MBR decoding used in our experiments. For comparison, we also performed translation candidate reranking using a quality estimation model COMETKIWI (Rei et al., 2022b) (QE)2, MBR decoding with confidence-based pruning (PruneMBR) (Cheng and Vlachos, 2023), and evaluated the quality upper bound (Oracle), which selects the hypothesis with the best score according to COMET using reference translations. We also compared CBMBR without kmeans++, where we randomly selected the initial centroids from the sample set (w/o kmeans++). We used COMET22 (Rei et al., 2022a) for the evaluation metric and utility function. In MBR decoding, we treat the hypothesis set as the pseudo-reference set, i.e., Ŷ := H. We set the number of centroids to k = 64. The details of our setup are shown in Appendix D.
Diverse translation candidates In this setting, we evaluated translation quality in six language directions: En↔Ja, En↔De, and En↔Zh
2Unlike MBR decoding, which requires calculating multiple pairwise scores from a single sentence representation, the QE model computes a single score from a single representation. Additionally, COMETKIWI does not compute explicit sentence vectors that are independent between the source and translation sentences but directly estimates the score from the concatenated two sentences. Therefore, the QE model does not cache the sentence vectors of both the source and hypotheses, like MBR decoding.
Decoding en-ja ja-en en-de de-en en-zh zh-en avg.
in the WMT’22 translation task (Kocmi et al., 2022). We generated translation candidates using the pre-trained multilingual translation model, M2M100 (Fan et al., 2021). We used beam search with a beam size of 256 for MAP decoding, and generated 1,024 translations using epsilon sampling with ϵ = 0.02 (Freitag et al., 2023) for MBR decoding.
Table 1 shows the translation quality of each decoding method. From the average scores in the table, compared with MAP decoding, both MBR decoding and the proposed CBMBR decoding improved the COMET score by +6.8 and +6.7%, respectively, and the gap between Oracle and both MBR and CBMBR narrowed. The results also show that the difference between CBMBR and MBR was within 0.1% using kmeans++ initialization.
Next, we compared the decoding time of each method as shown in Table 2. From the table, the total time of COMET-MBR (E2E) indicates that CBMBR was 1.8 times faster than vanilla MBR. Specifically, in the computation of the expected utility, which required quadratic time, the speed
increased by 5.7 times when including kmeans++, and by 16.0 times when comparing only the utility computation. We confirmed that, compared with PruneMBR, the speed of the expected utility calculation improved by 1.4 times. One reason for this improvement is that, unlike PruneMBR, CBMBR computes the expected utility in a single transaction, which makes it easier to leverage the parallel computation capabilities of the GPU.
To summarize, CBMBR maintains translation quality comparable with naive MBR decoding while accelerating the computational time of the expected utility calculation by 5.7 times, including clustering.
Multi-system translation We also evaluated the effectiveness of our CBMBR in the setting where translation candidates are generated from multiple translation systems. In particular, we followed the practice in Deguchi et al. (2023), in which 18 candidate sets with each set comprising the 50-best translations were generated from nine models and two decoding methods: beam search and top-p sampling (p = 0.7) with a beam size of 50. We evaluated the translation quality in two language directions: En↔Ja in the WMT’22 and WMT’23 translation tasks (Kocmi et al., 2022, 2023).
Table 3 shows the results. Unlike the diverse translation candidates setting, CBMBR improved the translation quality by up to 0.5% compared with naive MBR. Naive MBR calculates the expected utility using all samples equally, which is prone to translation bias when candidates have a multimodal distribution. By contrast, CBMBR estimates the expected utility using only centroids; therefore, it decodes robustly, even if the distribution is multimodal. A detailed analysis is shown in Appendix F.
To summarize, we found that CBMBR not only improved the decoding speed but also improved
translation quality compared with vanilla MBR when the translation was determined from the candidate sets generated from multiple translation systems.

heading: 5 Discussion

heading: 5.1 Number of centroids k
text: We evaluated the COMET scores of various k ∈ {2i}8i=0 in the multi-system translation setting. Figure 2 shows the results. When k = 1, while the time complexity was linear time O(N), the COMET score of CBMBR was degraded by 0.5% compared with vanilla MBR. The figure shows that translation quality improved as k increased and CBMBR outperformed vanilla MBR when k ≥ 16. Additionally, translation quality was better when we used kmeans++ instead of the standard kmeans.

heading: 5.2 Distance between similar sentence vectors
text: As mentioned in Section 3, CBMBR implicitly assumes that the score function s has linearity. Because s in the COMET model is a nonlinear multi-layer perceptron, it is not clear whether it is appropriate to use the averaged vector of hypothesis representations as representative points of the clusters. To verify the assumption, we investigated the distances between sentence vectors of similar sentences using the semantic textual similarity benchmark (STS-B) task (Cer et al., 2017). We evaluated the Pearson correlation coefficient
r using the ground truth similarity score. Table 4 shows the experimental results. Despite sentence vectors not being trained explicitly like contrastive learning, COMET demonstrated a strong correlation of 73.6. Moreover, the result demonstrates that it implicitly learned sentence similarity through the training of score prediction, as evidenced by its significantly better correlation of 73.6 compared with the pre-trained XLM-R score of 31.6. Furthermore, we confirmed that COMET outperformed LaBSE (Feng et al., 2022) trained using contrastive learning.
To summarize, the sentence vectors of COMET demonstrated a strong correlation with gold scores in the STS-B task although the sentence representations were not explicitly trained. Additionally, we confirmed that the encoder of COMET implicitly learned sentence similarity through score prediction. From the results, we verified that approximating the expected utility by aggregating sentence vectors that are close to each other is reasonable because the distance between sentence vectors represents semantic similarity.

heading: 6 Conclusion
text: In this paper, we proposed CBMBR, which improved the speed of MBR decoding by clustering the sentence vectors of similar sentences and computing the score using the centroid representations of each cluster. Our CBMBR achieved a 5.7 times speed-up in the expected score calculation and an improvement in COMET of up to 0.5% compared with vanilla MBR decoding in the WMT’22 En↔Ja, En↔De, En↔Zh, and WMT’23 En↔Ja translation tasks. For future work, we will apply our method to other evaluation metrics, including both neural and non-neural metrics.
Limitations
In this study focused only on improving the speed of MBR decoding, particularly the neural evaluation metric, COMET. For non-neural metrics, it is necessary to apply the appropriate clustering method for each metric.
In COMET-MBR, there are two bottlenecks for the computational time: the calculation of the expected utility and sentence encoding. However, we only improved the computation speed of the expected utility, which took quadratic time. Although sentence encoding can be computed in linear time, the sentences are encoded using the expen-
sive XLM-R encoder, which is time-consuming. Our method can only be applied to metrics for which we can compute the representation independently for each sentence. This limitation is the same as that of the method of DeNero et al. (2009) and is also explained in their paper.
We measured the decoding times reported in this paper on a single computer and only in a single run; the amount of speed improvement may differ when different computer architectures are used.
Ethical Consideration
Both vanilla MBR decoding and CBMBR decoding select output sentences from a set of translation candidates generated by translation systems; hence, if the systems generate toxic text, it may be selected.

heading: Acknowledgements
text: This work was partially supported by JSPS KAKENHI Grant Number JP21H05054.

heading: A Licenses
text: In our experiments, we used the COMET-22 model licensed under the Apache-2.0 license and
COMETKIWI model licensed under the CC BYNC-SA 4.0 license. We evaluated our method on the test sets of WMT’21, WMT’22, and WMT’23 translation tasks under the following policy: “The data released for the WMT General MT task can be freely used for research purposes”.

heading: B Details of the Datasets
text: Table 5 shows the number of sentences for each dataset we used in our experiments.
Dataset en-ja ja-en en-de de-en en-zh zh-en
WMT’21 1,000 1,005 1,002 1,000 1,002 1,948 WMT’22 2,037 2,008 2,037 1,984 2,037 1,875 WMT’23 2,074 1,992 – – – –
Table 5: Number of sentences for each dataset we used.

heading: C Details of the Algorithms and Models
text: C.1 kmeans++ We describe the algorithm for the initial centroid selection of kmeans++:
1. Pick up the first centroid from the set Ŷ and add it to C.
2. Calculate the squared Euclidean distance between vector f(ŷi) and its nearest centroid d2(ŷi) = minc∈C∥f(ŷi)− c∥22.
3. Sample vector f(ŷi) from the multinominal distribution according to the weights
d2(ŷi)∑|Ŷ| j=1 d 2(ŷj) and add it to the set C.
4. Repeat steps 2 and 3 until k centroids are selected.
C.2 COMET model Figure 3 shows the overview of the COMET model. A triplet of sentences are independently encoded into their sentence vectors, and then the COMET score is calculated from the vectors.

heading: D Details of the Experimental Setup
text: Table 6 shows the details of our experimental setup. We implemented vanilla MBR, PruneMBR, and CBMBR using PyTorch.
3https://huggingface.co/Unbabel/ wmt22-comet-da
4https://huggingface.co/Unbabel/ wmt22-cometkiwi-da
5https://huggingface.co/facebook/m2m100_418M 6https://huggingface.co/naist-nlp/wmt23

heading: E Other Experimental Results
text: E.1 Translation quality on the development set in the diverse translation candidates setting
Table 7 shows the experimental results of the diverse translation candidates setting on the development set. In the table, “niter” denotes the number of iterations of kmeans clustering. We chose niter = 1 from the results.
E.2 Decoding speed in the multi-system translation setting
Table 8 shows the decoding speed in the multisystem translation setting measure in the WMT’22 and WMT’23 En↔Ja translation tasks. As shown in the table, our CBMBR improved the speed of the expected score calculation by 5.0 times compared with vanilla MBR and 1.5 times compared with PruneMBR in the multi-system setting.
E.3 Relationship between the number of centroids and translation quality in the diverse translation setting
Figure 4 shows the translation quality of various k in the diverse translation candidates setting. We observed that the COMET score increased as k increased, except for an outlier, i.e., k = 2.
E.4 Evaluation of non-target metrics
We also evaluated translation quality using metrics that were not used as the utility function. Table 9 and 10 show the results of the diverse translation candidates setting, and Table 11 and 12 show the results of the multi-system translation setting, respectively.

heading: F Multimodality of Translation Candidates
text: In the multi-system translation setting, CBMBR also outperformed vanilla MBR in terms of transla-
Decoding en-ja ja-en en-de de-en en-zh zh-en avg.
tion quality, as shown in Table 3 and Figure 2. In this section, we discuss the multimodal nature of translation, two approximations of MBR decoding, and why our CBMBR outperformed vanilla MBR in the multi-system translation setting.
The n-best translations generated by beam search are often similar to each other (Vijayakumar et al., 2018). To diversify the candidates while maximizing translation quality, Deguchi et al. (2023) generated the 50-best translation sets from each translation system, which resulted in candidates that exhibited multimodality. Vanilla MBR decoding calculates the expected score by treating all samples equally, which means that it is prone to being affected by the number of similar transla-
Decoding en-ja ja-en en-de de-en en-zh zh-en avg.
Decoding en-ja ja-en en-de de-en en-zh zh-en avg.
tion samples for the candidates that have such a multimodal distribution.
Now, there are two approximation variants of MBR decoding in our CBMBR. One is our proposed method, which calculates the expected score using centroid representations:
y∗CBMBR = argmaxh∈H Ec∈C [s(f(x), f(h), c)] . (5) The other, CBMBRcnt, multiplies each centroidbased score by the weight according to the number of samples in each cluster:
y∗CBMBRcnt = argmaxh∈H Ec∈C [s(f(x), f(h), c)× w(c)] , (6)
where w : RD → [0, 1] returns the weight of the given centroid as follows:
w(c) = count(c)
∑k i=1 count(ci) , (7)
count(c) = ∣∣∣ { ŷ ∈ Ŷ : c = NN(f(ŷ), C) }∣∣∣ , (8)
NN(q, C) = argminc∈C∥q − c∥2, (9)
where NN: RD×C → RD finds the nearest neighbor centroid of a vector q ∈ RD from the given set
of centroids C and count : RD → N ∪ {0} counts the number of samples in the cluster of the given centroid. CBMBRcnt, which uses the number of samples in a cluster, can be regarded as more accurately approximating vanilla MBR compared with our CBMBR, which ignores the number of samples in a cluster.
We compared the translation quality of our CBMBR and CBMBRcnt with that in the multisystem translation setting. Table 13 shows the results. From the results, the difference between vanilla MBR and CBMBRcnt narrowed to 0.1% and degraded by 0.6% compared with CBMBR; that is, CBMBRcnt, which more accurately approximates vanilla MBR, was worse than our proposed CBMBR. We attribute this observation to the biased distribution caused by beam search or sampling. With CBMBR, we can robustly decode against bias.

heading: Future Work
text: For future work, we will apply our method to other evaluation metrics, including both neural and non-neural metrics. Limitations
In this study focused only on improving the speed of MBR decoding, particularly the neural evaluation metric, COMET. For non-neural metrics, it is necessary to apply the appropriate clustering method for each metric. In COMET-MBR, there are two bottlenecks for the computational time: the calculation of the expected utility and sentence encoding. However, we only improved the computation speed of the expected utility, which took quadratic time. Although sentence encoding can be computed in linear time, the sentences are encoded using the expen-
sive XLM-R encoder, which is time-consuming. Our method can only be applied to metrics for which we can compute the representation independently for each sentence. This limitation is the same as that of the method of DeNero et al. (2009) and is also explained in their paper. We measured the decoding times reported in this paper on a single computer and only in a single run; the amount of speed improvement may differ when different computer architectures are used. Ethical Consideration
Both vanilla MBR decoding and CBMBR decoding select output sentences from a set of translation candidates generated by translation systems; hence, if the systems generate toxic text, it may be selected