Quality-Aware Translation Models: Efficient Generation and Quality Estimation in a Single Model (2024)

Abstract: Maximum-a-posteriori (MAP) decoding is the most widely used decoding strategy for neural machine translation (NMT) models. The underlying assumption is that model probability correlates well with human judgment, with better translations getting assigned a higher score by the model. However, research has shown that this assumption does not always hold, and generation quality can be improved by decoding to optimize a utility function backed by a metric or quality-estimation signal, as is done by Minimum Bayes Risk (MBR) or Quality-Aware decoding. The main disadvantage of these approaches is that they require an additional model to calculate the utility function during decoding, significantly increasing the computational cost. In this paper, we propose to make the NMT models themselves quality-aware by training them to estimate the quality of their own output. Using this approach for MBR decoding we can drastically reduce the size of the candidate list, resulting in a speed-up of two-orders of magnitude. When applying our method to MAP decoding we obtain quality gains similar or even superior to quality reranking approaches, but with the efficiency of single pass decoding.

## introduction most state-of-the-art models for natural language processing tasks are probabilistic, with the most frequent parameterization being based on neural networks. once these models are trained, the prevailing decoding strategy for natural language generation is map decoding, i.e. select the hypothesis that maximizes the conditional probability given an input. as an exact maximization is computationally intractable, typically beam search or greedy decoding are used to approximate the search for the best * work carried out while interning at google. correspondence to: christian tomani (christian.tomani@tum.de), david vilar (vilar@google.com). hypothesis. neural machine translation is a prominent example of these types of models, where the system is trained to generate a sentence in a target language given a source sentence in another language. nonetheless, @xcite have demonstrated that map decoding methods may be suboptimal due to the presence of misaligned probability distributions. moreover, nmt models often assign human translations lower probabilities than their own beam search outputs due to calibration issues @xcite . @xcite applied mbr decoding for nmt models as an alternative generation approach. mbr decoding follows a selfconsistency approach by sampling from the model distribution and giving preference to hypotheses that exhibit greater similarity to all other hypotheses. in contrast to map decoding, mbr decoding's objective is not centered on generating the translation with the highest estimated model probability, instead it selects the translation that exhibits the highest quality based on a utility metric. subsequent research conducted by @xcite showed that mbr decoding with neural utility metrics leads to significant improvements over beam search decoding. however, mbr is computationally intensive, with a time complexity of o(m 2 ) for a candidate list containing m samples, ideally @xmath0. note than when using neural metrics, each "computation step" in the quadratic complexity is itself computationally expensive, requiring a forward pass through a large neural network. as an alternative to mbr decoding, we can use a quality-aware decoding strategy, generating a list of candidate translations and reranking them using a neural quality estimation (qe) metric that computes a quality score for the translation conditioned only on the source sentence. this method offers the advantage of being more efficient than mbr decoding, as its inference speed scales linearly with the number of candidate translations. a study conducted by @xcite showed that employing neural metrics for qe reranking exhibits comparable advantages to those seen with mbr decoding. however, this approach still demands the use of a separate, computationally expensive qe model to evaluate the quality of each candidate. in our work, we propose a novel method that moves quality awareness inside the model itself, enabling us to guide the generation process towards higher-quality translations, and eliminating the need for an external qe model during decoding. specifically, we investigate two key strategies: quality-aware prompting, where we use quality prompts that explicitly encourage the generation of high-quality translations, and quality-aware prediction, where we enable an nmt model to judge the quality of its own translations. both strategies add special quality tokens to each nmt training example. the strategies differ only in whether the token is included in the source or the target sentence. our main scientific contributions are the introduction of quality-aware translation models, and their application for improved, more efficient decoding strategies. we analyze two use cases: ## related work machine translation (mt) metrics can be divided into two high-level categories: reference-based and reference-free, also known as quality estimation (qe) metrics. qe metrics compute a quality score for the mt output conditioned only on the source text. reference-based metrics, on the other hand, require a human-generated reference translation to compare the mt output with. a multitude of reference-based metrics are available to evaluate the quality of translated content. some metrics rely on lexical overlap, such as bleu @xcite , meteor (banerjee and lavie, 2005) or chrf @xcite . the wmt metrics task @xcite demonstrated that the new generation of metrics -neural fine-tuned metrics like bleurt @xcite and comet @xcite ) -have significantly higher correlation with human judgement than traditional word overlap metrics. consequently, we focus on neural fine-tuned metrics in this study. quality estimation for mt began as confidence estimation @xcite , but has recently shifted to embrace close kinship with reference-based metrics, with recent neural examples including openkiwi @xcite , transquest @xcite , and comet-qe @xcite . a prominent example of a reference-based neural metric is bleurt @xcite , and its extension to metricx, which was the winning entry in the wmt22 metrics task @xcite . building on this foundation, @xcite recently introduced metricx-qe, which we use as our primary qe metric in this work. reranking has a long history in translation, starting with @xcite , where a discriminative model is learned to rank a candidate list to maximize a reference-based metric, with recent examples including @xcite and @xcite , who both train using bleu. our approach is closest to that of @xcite , who rerank using various translation quality-estimation metrics, as opposed to training a special-purpose discriminative reranker. we differ from these works in that we do not need any external quality signal, which is instead provided by the nmt system itself. while conventional mt research often relies on map decoding or generating k-best lists through beam search for mbr decoding, eikema and aziz (2020) proposed an approximation of mbr decoding via unbiased sampling. their method aims to address the limitations of map decoding @xcite @xcite by demonstrating that samples drawn from the nmt model align more faithfully with training data statistics when compared to beam search. @xcite showed that using neural metrics instead of overlap metrics results in significant improvements in translation quality. as a follow up, @xcite reported that the choice of sampling approach is important and epsilon sampling @xcite is ideal for mbr decoding and reranking. @xcite introduced an orthogonal method to our proposed approach. they speed up mbr decoding by gradually increasing the number of samples used to estimate the utility while pruning hypotheses. our quality-aware prompting approach extends a long line of methods where tagged training data has been used to control nmt output for different properties, including target language @xcite , formality level @xcite , politeness @xcite , domain @xcite , gender @xcite , syntactic structure @xcite , complexity @xcite and reading level @xcite . the approaches closest to ours identify attributes related to translation quality, such as tagging back-translated examples to control away from synthetic data @xcite , or tagging target-original examples to control toward natural-sounding output @xcite . to the best of our knowledge, we are the first to use quality estimation to tag training data, allowing nmt to discriminate between different translation qualities, and allowing us to prompt the model to generate high quality translations. ## preliminaries we are given a nmt model p θ (y|x) which serves to estimate the probability of a hypothesis segment @xmath1, given a source segment x. here, θ denotes the learned parameters of the neural network and y the set of all possible hypotheses. there are two widely used approaches for generating the translations of a given sentence. ## experimental setup model: our model is a transformer consisting of 6 encoder and 6 decoder layers, 16 attention heads with a dimension of 128, a hidden dimension of 8192, and a model dimension of 1024, resulting in 551m parameters. we employ a shared vocabulary of 32k tokens and impose a maximum sentence length of 128 tokens. we utilize gelus with gated activation functions. the baseline system is trained on the entire available dataset. all models are trained on tpus (v3) until they reach convergence. the metricx-qe model used for quality estimation is a transformer based model with a total of 2b parameters, as described in @xcite . to assess the applicability of our approach to llms, we also fine-tune and evaluate a qualityaware llm, with results in appendix c. ## conclusion this paper introduces a novel approach to enhance nmt by making the models quality-aware. our approach addresses the issue of misalignment be-tween outputs generated via map decoding and human judgment. we achieve this by training nmt models to assess the quality of their own translations, effectively circumventing the limitations of conventional decoding methods. as a result this new approach yields improvements similar or superior to qe-reranking approaches, but with the efficiency of map-decoding, i.e. with single-pass decoding (for qa prompting, the best performing method). qe-reranking in contrast needs a sampling step followed by a reranking step using an external, computationally expensive model. by leveraging the model's quality signal internally during mbr decoding, not only does translation quality further improve, but computational efficiency is also dramatically enhanced, reducing inference time by two orders of magnitude. this improvement comes from a drastic reduction in the necessary size of the candidate list needed for producing good quality translations. overall this research opens up exciting possibilities for advancing the field of nmt, offering both improved translation quality and faster processing speeds without the need for additional, computationally intensive models.