PaperlayerView exact revision on arXiv

Source-linked AI summary

Machine Translation Decoding beyond Beam Search

Rémi Leblond, Jean-Baptiste Alayrac, Laurent Sifre, Miruna Pislar, Jean-Baptiste Lespiau, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals

arXiv:2104.05336v1cs.CLcs.LG

TL;DR

Beam search seeks high-likelihood outputs but does not optimize the end metric practitioners care about. This paper compares metric-driven decoding algorithms, including value-based methods and MCTS, and finds that the best choice depends on the metric: beam search leads for privileged metrics, while alternatives including MCTS outperform it for unprivileged scores.

  • Problem

    Beam search optimizes model likelihood rather than the metric practitioners care about, motivating tests of whether metric-driven decoding can replace it.

    Source: §1 · p. 1

  • Method

    The paper compares decoding algorithms that optimize metrics directly through a learned value function or, when available, the metric itself, including MCTS variants.

    Source: §1 · p. 1

  • Results

    Beam search is best only for privileged metrics; for unprivileged scores, competitors including MCTS outperform it, while value-guided methods perform significantly better than regular beam size.

    Source: §1 · p. 2, §6.2 · p. 7

  • Takeaways & Limitations

    The best decoding algorithm depends heavily on the metric, so machine-translation pipelines should choose a decoding mechanism rather than defaulting to beam search.

    Source: §7 · p. 8

    Value networks are relatively low quality in the experiments, especially for privileged metrics, where value-based decoding yields comparatively little improvement, particularly on small ENDE.

    Source: §7 · p. 9

Abstract

from arXiv · show

Beam search is the go-to method for decoding auto-regressive machine translation models. While it yields consistent improvements in terms of BLEU, it is only concerned with finding outputs with high model likelihood, and is thus agnostic to whatever end metric or score practitioners care about. Our aim is to establish whether beam search can be replaced by a more powerful metric-driven search technique. To this end, we explore numerous decoding algorithms, including some which rely on a value function parameterised by a neural network, and report results on a variety of metrics. Notably, we introduce a Monte-Carlo Tree Search (MCTS) based method and showcase its competitiveness. We provide a blueprint for how to use MCTS fruitfully in language applications, which opens promising future directions. We find that which algorithm is best heavily depends on the characteristics of the goal metric; we believe that our extensive experiments and analysis will inform further research in this area.

1. Introduction

Likelihood-based decoding is tractable only through heuristics, and beam search optimizes model likelihood rather than practitioners’ target metrics. The paper compares metric-driven alternatives and finds that the best algorithm depends on the metric, with beam search strongest for privileged metrics and alternatives including MCTS outperforming it for unprivileged scores.

Motivation

  • Beam search maintains parallel hypotheses and improves on greedy decoding, but remains agnostic to the metrics practitioners want to optimize.Exact search is generally infeasible for token-by-token models, so practitioners use heuristic mechanisms.Source: §1 · p. 1
  • Prior approaches focused on training or differentiable search but still decoded by likelihood; reported RL gains were limited and poorly understood.The paper instead compares algorithms that target metrics directly through a value function or the metric itself when available.Source: §1 · p. 1

Contributions

  • The study tests whether metric optimization at inference improves performance and whether it scales with computation, unlike beam search’s reported degradation at large beam sizes.The experiments focus on machine translation, a well-studied sequence-to-sequence task with available data and benchmarks.Source: §1 · p. 2
  • Beam search is best only for privileged metrics; for unprivileged scores, competitors including MCTS outperform it.For privileged metrics, value-based alternatives falter because the value problem ultimately relies on reconstructing hidden information.Source: §1 · p. 2
  • The paper distinguishes privileged metrics, which use ground-truth translations, from unprivileged scores, and introduces Multilingual BERTScore as an illustrative unprivileged metric.It also presents a blueprint for applying MCTS in NLP.Source: §1 · p. 2

2. Related Work

Prior work documents limitations of likelihood-based decoding and explores alternatives including minimum Bayes risk, diverse sampling, and value-guided search. The paper’s MCTS approach is related to value-guided beam search but searches further into the future.

Beam search variants

  • Beam search variants include out-of-order decoding, look-ahead, and methods that speed up search for monotonous scores, but the cited MT results do not surpass beam search.Look-ahead improves low-data tasks, while out-of-order decoding benefits other task types but not MT.Source: §2 · p. 2

Beam search and model likelihood

  • Beam search performance can decline beyond a fairly small beam, often misses the model’s most likely sequence, and exact search finds the empty sequence in MT more than half the time.Minimum Bayes risk decoding uses the distribution rather than only its mode and can outperform vanilla beam search in low-resource scenarios.Source: §2 · p. 2

Sampling diversity

  • Diversity methods include top-k and nucleus sampling, Gumbel-Max unique sampling, and MCTS that uses expensive rollouts to add diverse samples to a larger NMT system.The cited diversity work is motivated mainly by applications other than machine translation.Source: §2 · pp. 2–3

Metric-driven decoding

  • Value-guided beam search combines policy logits with a network’s estimate of the eventual score, while MCTS can explore further into the future.Value-guided search can improve on vanilla beam search and is less sensitive to beam size, but scoring every token can be prohibitively expensive with MT vocabularies above 30,000.Source: §2 · p. 3

3. Machine Translation metrics

The paper distinguishes metrics that require reference translations from metrics that assess outputs against source sentences. It introduces Multilingual BERTScore as an illustrative, imperfect unprivileged metric and reports that beam outputs score above greedy outputs on it.

Metric types

  • Privileged metrics require reference translations, whereas unprivileged metrics assess translations without ground-truth sequences and can be used on specific unseen inputs.Privileged metrics are suited to average-performance comparisons but depend on test-set translation quality.Source: §3 · p. 3

Privileged metrics

  • BLEU is a corpus-level n-gram precision metric with a brevity penalty, so it is unavailable for decoding-time assessment without references.BERTScore compares candidate and reference sentences using contextual token embeddings and is sentence-level before corpus averaging.Source: §3 · p. 3

Unprivileged metrics

  • Multilingual BERTScore compares a candidate translation with its source using multilingual BERT embeddings, serving as an illustrative unprivileged metric.The authors do not advocate widespread adoption; its alignment design may favor language pairs with syntactically similar, one-to-one token correspondence.Source: §3 · pp. 3–4
  • Table 1 reports that beam-search outputs consistently outscore greedy outputs on Multilingual BERTScore.The paper also notes that this score is lower than corresponding monolingual BERTScore values and that trained model outputs score higher than ground-truth outputs.Source: Table 1 · p. 4, §3 · p. 4

4. Decoding algorithms

The paper groups decoding methods into likelihood-based, value-based, and ranking-based approaches, which differ in whether they use model likelihood, an estimated score, or directly computed scores. Value guidance supports metrics unavailable at test time, while reranking can directly select high-scoring finished candidates when scores are tractable.

4.1. Likelihood-based decoding

  • Likelihood-based beam search expands each of k prefixes with up to k likely tokens, then retains the k highest-likelihood prefixes, making it metric-agnostic.Greedy decoding is the baseline that selects the highest-likelihood token at each step.Source: §4.1 · p. 4

4.2. Value-based decoding

  • Value-based decoding treats translation as an MDP and uses a value function to estimate the final score from an unfinished state, providing forward-looking guidance.Value-guided beam search combines length-normalized log-likelihood with the network value; V-MCTS allocates simulations to explore future trajectories using a policy and value network.Source: §4.2 · p. 4
  • In V-MCTS, selection follows pUCT using action values, policy probabilities, temperature, and visit counts; expansion evaluates policy and value, and backup updates search statistics.The next token is chosen from root-child statistics, such as visit counts or aggregated values.Source: §4.2 · p. 5

4.3. Reranking-based decoding

  • For tractable unprivileged metrics, reranking scores finished samples directly and selects the highest-scoring candidate; S+RV instead ranks by estimated value.MCTS+Roll replaces a node’s value estimate with a more expensive score from a greedy rollout, serving as a proof of concept for measuring value-function loss.Source: §4.3 · p. 5

5. Training a value network

The section describes training a value network on policy-generated samples and explains why privileged metrics are difficult to predict. A hybrid training pathway uses reference information to improve the regular value representation while keeping inference non-cheating.

5. Training a value network

  • Training data are augmented by replacing each reference with a greedy sample from the supervised policy and scoring that sample against the reference for privileged metrics or the source for unprivileged metrics.The method then trains a dual-headed network with a shared transformer encoder-decoder torso and policy and value heads.Source: §5 · pp. 5–6
  • Using policy-generated samples is mandatory for privileged metrics because ground-truth-versus-ground-truth scores are perfect and therefore uninformative for value-network training.Samples also give values associated with the trained policy used during decoding.Source: §5 · p. 6
  • The policy is trained to stay close to the initial supervised policy with KL divergence, while value regression discretizes scores into buckets and adds loss at every decoder step.The decoder causality mask is reused to train on unfinished samples.Source: §5 · p. 6
  • For privileged metrics, the hardest value-network subtask is assessing translation quality across languages, while score approximation is easy and end-of-trajectory prediction improves with target access.The ablation finds that assessing finished sentence pairs has roughly the same error as the full setup.Source: §5 · p. 6
  • The proposed explanation is that predicting a privileged score requires inferring the reference from the source, making value prediction empirically harder than policy prediction on this dataset.The value network receives the source and a sample prefix, whereas the privileged score compares a sample with a ground-truth reference.Source: §5 · p. 6
  • The hybrid training regimen significantly reduces the performance gap between the non-cheating model and the cheating value network, while inference uses only the regular pathway.The training setup applies value loss to the reference-conditioned pathway and aligns the pathways’ final layers with an ℓ2 loss.Source: §5 · p. 6

6. Experiments

The experiments compare likelihood-, value-, and score-based decoding across privileged and unprivileged metrics, then tune the algorithms and test scaling with search budget. Which approach performs best depends on the metric setting and the quality or availability of scores.

6.1. Experimental setup

  • Experiments use WMT’14 ENDE with roughly 4.5 million training pairs and ENFR with just under 41 million, allowing dataset-scale comparisons; sampled solutions receive 50 inferences per token.Greedy decoding uses 1 inference per token and beam search uses 4.Source: §6.1 · pp. 6–7

6.2. Main results analysis

  • Beam search is strong on privileged metrics, while value-guided methods prevail on unprivileged scores; direct score-based methods outperform their value-guided counterparts, though V-MCTS remains competitive.Table 4 summarizes the comparison on ENDE and ENFR.Source: Table 4 · p. 7, §6.2 · p. 7

6.1. Experimental setup

  • Shared keys and values across attention heads yield an almost linear speedup with head count at a small accuracy cost, while 128-dimensional keys and values remove TPUv3 padding overhead.The adapted models still match or outperform the original Transformer models.Source: §6.1 · p. 7

6.2. Main results analysis

  • On ENFR, MCTS optimizing MLBERTScore reaches 40.31 BLEU, more than a full BLEU point above beam search, while larger budgets can favor less natural word-for-word translations.The authors report that MLBERTScore is not perfectly aligned with BLEU.Source: §6.2 · p. 7

6.3. Algorithms tuning and ablations

  • Beam search mostly performs best with default hyperparameters; value-guided beam works best with plain values, length-normalized likelihoods, and α = 0.5 for privileged versus 0.9 ENDE and 0.95 ENFR for unprivileged metrics.The larger unprivileged weights are consistent with the authors’ observation that value functions are higher quality in that setting.Source: §6.3 · p. 8
  • For unprivileged MCTS, max backups and selecting the root child with the highest aggregate value work best; MCTS+rollouts also favors aggressive value following when estimates are good but expensive.For privileged metrics, the reverse MCTS choices work better because value networks are lower quality.Source: §6.3 · p. 8
  • Sampling and reranking is sensitive to temperature: high values (τ ≥1) yield nonsensical outputs, while low values (τ ≤0.5) produce almost no diversity.The best dev-set temperatures are τ = 0.75 for MLBERTScore and τ = 0.25 for privileged metrics using S+RV.Source: §6.3 · p. 8

6.4. Scaling search with computational budget

  • With more computation, performance quickly plateaus or degrades for imperfect privileged-metric value networks, rises more steadily for higher-quality unprivileged-metric values, and keeps increasing when scores are accessed directly.Table 5 compares search budgets from 1 to 100 inferences per token on MLBERTScore.Source: §6.4 · p. 8, Table 5 · p. 8

7. Discussion

The best decoding algorithm depends heavily on the metric, so practitioners should choose a method rather than defaulting to beam search. The experiments also expose limits of value-based decoding for privileged metrics and identify settings where MCTS or S+R performs well.

7. Discussion

  • The best decoding algorithm depends heavily on the metric, making metric-specific selection preferable to defaulting to beam search.Source: §7 · p. 8
  • Value-based decoding for privileged metrics yields comparatively little improvement, especially on the small ENDE dataset, while unprivileged metrics appear more promising.Distinguishing large quality gaps is easier than modelling the policy, but discriminating between good candidates is similarly hard because estimating the ground-truth sentence is required.Source: §7 · pp. 8–9
  • MCTS is a valid MT decoding method and the best option in some use cases, provided necessary adjustments are made.Its performance is closely linked to how easily a good value function can be learned.Source: §7 · p. 9
  • When direct access to the score is possible, the simple S+R method performs well, though further experimentation is needed to understand why.Source: §7 · p. 9
  • A future direction is to train scores directly from human assessments of translation pairs and then optimize them with MCTS.Source: §7 · p. 9
  • Another proposed direction is iterative reinforcement learning that improves policies through value-guided search and trains value functions on search-improved policies.Source: §7 · p. 9

Outline

The paper moves from prior work and metric distinctions to decoding methods, value-function training, and experimental evaluation, before closing with discussion and future directions.

Outline

  • Appendix A details network architectures and optimisation hyper-parameters, including a hybrid architecture to improve value-network training in the hard case of privileged metrics; Appendix B details tuning of decoding mechanisms used throughout the paper.Source: Outline · p. 12
  • Appendix A covers architectures, hyper-parameters and a hybrid value-network architecture for privileged metrics; B decoding tuning; C computational-budget scaling; D policy/value training trade-offs; E MCTS trees; F batched NumPy MCTS.Source: Outline · p. 12
  • Appendix A details network architectures and optimization hyperparameters and describes a hybrid architecture for improving value-network training in the hard case of privileged metrics, while Appendices B–F cover decoding, scaling, trade-offs, trees, and implementation.Source: Outline · p. 12

A. Network architectures and training

The paper uses a dual-headed transformer for policy and value prediction, and adds a shared-weight hybrid pathway to improve value learning for privileged metrics. The hybrid approach yields more substantial MCTS gains than the normal architecture, though its cheating pathway cannot be used at decoding time.

Dual-head transformer architecture

  • The transformer uses six encoder and decoder layers, hidden size 512, and inference-oriented multi-query attention that reduces key-value memory use by a factor of 16 and yields near-linear speedup with attention-head count.The architecture also uses 16 attention heads and a feedforward hidden size of 3072.Source: A. Network architectures and training · p. 12
  • The dual-headed transformer adds a policy head that predicts vocabulary distributions and a value head that predicts distributions over 500 value buckets.The value estimate is computed as the sum of bucket probabilities multiplied by their average values.Source: A. Network architectures and training · p. 12

Hybrid architecture for privileged metrics

  • With the hybrid architecture, MCTS improves more substantially over greedy decoding than with the normal architecture on ENDE BLEU and BERTScore; the cheating-model results are from a prohibited setting.Table 6 reports the comparisons for 50 simulations.Source: Table 6 · p. 13, A. Network architectures and training · p. 13
  • The hybrid model shares transformer weights across regular and cheating pathways: the regular encoder receives the source sentence, while the cheating encoder receives the ground-truth target unavailable at test time.Source: A. Network architectures and training · pp. 12–13
  • Training applies policy and value losses to the regular pathway, a value loss to the cheating pathway, and an ℓ2 distillation loss between their final layers with a stop gradient on the cheating pathway.No policy loss is applied to the cheating pathway because it would encourage reproducing its input.Source: A. Network architectures and training · p. 13
  • Best performance combines shared weights, distillation, and cheating-value loss, with weights 1.0 for Lπ and Lvc and 0.1 for Lv and D.Source: A. Network architectures and training · p. 13

B. Experimental details and ablations

The ablations show that decoding settings and value reliance should vary with the metric: privileged metrics favor more conservative value use, while Multilingual BERTScore supports stronger value guidance. Search variants also require distinct settings for temperature, aggregation, and action selection.

Beam Search

Value-guided beam search

  • For value-guided decoding, logits temperature 1.0 performs best by a significant margin, while tuned log-value normalization matches but does not beat the plain-value formulation.Source: B. Experimental details and ablations · p. 14
  • For the linear policy-value weight α, privileged metrics perform best at 0.5, while Multilingual BERTScore performs best at 0.9 on ENDE and 0.95 on ENFR.The authors hypothesize that the latter metric's value function is more reliable.Source: B. Experimental details and ablations · p. 14, Table 8 · p. 14

MCTS variants

  • For privileged metrics, MCTS performs best with running-average backups and choosing the root child with the most visits; Multilingual BERTScore favors maximum backups and choosing the child with the highest aggregated value.The latter results indicate stronger reliance on the value function for the unprivileged metric.Source: Table 9 · p. 15, B. Experimental details and ablations · p. 15
  • For privileged metrics, the best UCT settings are τ = 0.9 and cpuct = 3.0; for Multilingual BERTScore, they are τ = 1.1 and cpuct = 8.0.Higher temperature and cpuct encourage exploration and increase the relative weight of value guidance.Source: B. Experimental details and ablations · p. 15

Sampling + Ranking variants

  • S+R uses temperature τ = 0.75 to balance sample diversity and correctness; S+RV uses τ = 0.25 for privileged metrics but τ = 0.75 for Multilingual BERTScore.The lower privileged-metric temperature makes S+RV rely more on the policy than on the lower-quality value function.Source: B. Experimental details and ablations · p. 15

C. Scaling search with computational budget

Decoding methods scale differently with computational budget depending on the quality and type of metric signal. Score-based search can overtake MCTS at higher simulation budgets, while beam search plateaus quickly.

C. Scaling search with computational budget

D. Supervised policy vs Distilled policy

Distilled-policy models perform slightly below plainly supervised models, with a larger gap on ENFR. Pairing a supervised policy with a distinct value model lets MCTS exceed other decoding methods on the larger dataset, though the comparison uses twice as many parameters.

D. Supervised policy vs Distilled policy

  • For privileged metrics, ground-truth targets would give every training item a value of 1, so the authors create a distillation dataset using greedy policy samples instead.Source: D. Supervised policy vs Distilled policy · p. 16
  • Distilled-policy models perform slightly worse than plainly supervised counterparts, with a larger performance effect on ENFR.Source: D. Supervised policy vs Distilled policy · p. 16, Table 12 · p. 16
  • With a supervised policy and a distinct Multilingual BERTScore value model, MCTS improves BLEU by more than 1 point over the supervised baseline on the larger dataset using 25 simulations.The result is reported as outperforming any other decoding type on that dataset; adding more computation does not help because decoding targets a different metric than BLEU.Source: Table 13 · p. 16, D. Supervised policy vs Distilled policy · p. 16
  • The supervised-policy and distinct-value-model comparison is not parameter-matched because the value network doubles the parameter count.The authors suggest a lightweight value head as a possible alternative and leave it for future work.Source: D. Supervised policy vs Distilled policy · p. 16

E. MCTS examples

F. Batched Numpy-friendly MCTS

The Numpy-compatible MCTS implementation batches search computations to better use accelerator hardware. It stores search state by batch item and tree node, limits action storage to policy-ranked candidates, and repeats selection, expansion, and value backup for each simulation.

F. Batched Numpy-friendly MCTS

  • A Numpy-compatible MCTS implementation runs on the accelerator device to better use batched neural-network computation and avoid queuing inefficiencies.Source: F. Batched Numpy-friendly MCTS · p. 20
  • Search storage is indexed by batch item and node; sparse tensors retain only the top A policy-ranked actions to reduce costly operations.The implementation also stores each node’s transformer state for incremental inference.Source: F. Batched Numpy-friendly MCTS · p. 20
  • The stored tree state tracks visit counts, raw and aggregated values, parent links, actions, and node depth.Source: F. Batched Numpy-friendly MCTS · p. 20
  • Each simulation selects an unexplored node with the UCT formula, expands it by computing policy and value, then backs the new value up the tree.Selection, expansion, and backup are described as batched operations.Source: F. Batched Numpy-friendly MCTS · pp. 21–23
Loading 2104.05336v1…