Centroid-Based Efficient Minimum Bayes Risk 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 $\mathcal{X}$ and $\mathcal{Y}$ be the spaces of possible source sentences and target sentences, respectively. MAP decoding generates the target sentence $y^{\ast}_{\text{MAP}}\in\mathcal{Y}$ using $y^{\ast}_{\text{MAP}}=\operatorname{argmax}_{y\in\mathcal{Y}}p_{\theta}(y|x)$, where $\theta$ denotes the parameter of the translation model that calculates the likelihood of an output sentence $y$ given an input sentence $x\in\mathcal{X}$. Because it is infeasible to calculate probabilities for all possible $y\in\mathcal{Y}$, beam search is typically used to obtain the solution.

By contrast, MBR decoding determines the output sentence $y^{\ast}_{\text{MBR}}\in\mathcal{Y}$ by maximizing the expected utility as follows:

where $u\colon\mathcal{Y}\times\mathcal{Y}\to\mathbb{R}$ denotes the utility function, which represents the preference relation and $\mathcal{H}=\{h_{i}\}_{i=1}^{|\mathcal{H}|}\subset\mathcal{Y}$ denotes the set of translation hypotheses. $P(y|x)$ is the true probability of being translated from a given input sentence $x\in\mathcal{X}$ and is approximated using the sampled reference translations $\hat{\mathcal{Y}}=\{\hat{y}_{i}\}_{i=1}^{|\hat{\mathcal{Y}}|}\subset\mathcal{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., $\hat{\mathcal{Y}}\coloneqq\mathcal{H}$. Note that the time complexity is $\mathcal{O}(N^{2})$, where $N\coloneqq|\mathcal{H}|$, which is time-consuming.

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\in\mathcal{X}$, hypothesis sentence $h\in\mathcal{Y}$, and reference sentence $\hat{y}\in\mathcal{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\colon\mathcal{X}\cup\mathcal{Y}\to\mathbb{R}^{D}$ be the function of sentence encoding and $s\colon\mathbb{R}^{D}\times\mathbb{R}^{D}\times\mathbb{R}^{D}\to\mathbb{R}$ be the output layer. The COMET score is computed using $s(f(x),f(h),f(\hat{y}))$. MBR decoding with COMET (COMET-MBR) replaces the utility $u$ in Equation 2 with the COMET score:

First, we compute the sentence vector of the source $f(x)\in\mathbb{R}^{D}$, the hypotheses $\{f(h_{i})\}_{i=1}^{|\mathcal{H}|}\subset\mathbb{R}^{D}$, and the pseudo-references $\{f(\hat{y}_{i})\}_{i=1}^{|\hat{\mathcal{Y}}|}\subset\mathbb{R}^{D}$.

Next, we cluster the sentence vectors of the pseudo-references into $k\ll N$ clusters and obtain the centroid vectors of each cluster $\mathcal{C}=\{\bm{c}_{i}\}_{i=1}^{k}\subset\mathbb{R}^{D}$. Here, we employ $k$means++ Arthur and Vassilvitskii (2007) to prevent the centroids from being biased. $k$means++ 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(\hat{y}_{i})\}_{i=1}^{|\hat{\mathcal{Y}}|}$ using the standard $k$means 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.

Finally, we calculate the expected utility by replacing pseudo-reference vectors $f(\hat{y})\in\mathbb{R}^{D}$ with centroids $\bm{c}\in\mathbb{R}^{D}$ in Equation 3:

The conventional method requires $\mathcal{O}(N^{2})$ of computational time to compute the expected utility for all hypotheses, whereas our CBMBR computes it in $\mathcal{O}(Nk)$. Note that $k$ ($1\leq k\leq N$) is a hyperparameter that balances the trade-off between the decoding speed and approximation accuracy. Especially, when $k=1$, i.e., $\mathcal{C}=\{\bm{c}_{1}\}$, the centroid $\bm{c}_{1}$ can be calculated as the average of all pseudo-reference vectors, i.e., $\bm{c}_{1}=\frac{1}{|\hat{\mathcal{Y}}|}\sum_{i=1}^{|\hat{\mathcal{Y}}|}f(\hat{y}_{i})$, and the time complexity of CBMBR is $\mathcal{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 $\operatorname{\mathbb{E}}_{\hat{y}\in\hat{\mathcal{Y}}}\left[s(f(x),f(h),f(\hat{y}))\right]$ as $s(f(x),f(h),\operatorname{\mathbb{E}}_{\hat{y}\in\hat{\mathcal{Y}}}\left[f(\hat{y})\right])$ in our method by assuming that the score function $s$ is roughly linear.

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 Unlike 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. , 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 $k$means++, where we randomly selected the initial centroids from the sample set (w/o $k$means++). We used Comet-22 (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., $\hat{\mathcal{Y}}\coloneqq\mathcal{H}$. We set the number of centroids to $k=64$. The details of our setup are shown in Appendix D.

In this setting, we evaluated translation quality in six language directions: En$\leftrightarrow$Ja, En$\leftrightarrow$De, and En$\leftrightarrow$Zh 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 $\epsilon=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 $k$means++ 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 $k$means++, 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.

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$\leftrightarrow$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.