Source-linked AI summary
Reward Augmented Maximum Likelihood for Neural Structured Prediction
Mohammad Norouzi, Samy Bengio, Zhifeng Chen, Navdeep Jaitly, Mike Schuster, Yonghui Wu, Dale Schuurmans
TL;DR
Structured prediction needs objectives that optimize nondifferentiable task rewards rather than relying only on conditional likelihood. RAML addresses this by sampling reward-weighted auxiliary outputs and training their likelihood, improving maximum-likelihood baselines in speech recognition and translation.
Problem
Task rewards used for evaluation are often nondifferentiable, while maximum likelihood does not explicitly distinguish alternative outputs by their task rewards.
Method
RAML samples auxiliary outputs proportionally to exponentiated scaled rewards and optimizes their conditional log-likelihood, using a temperature to control distribution smoothness.
Results
RAML consistently improves maximum-likelihood baselines on speech recognition and machine translation, with best reported gains of 2.3% PER on TIMIT test and 0.4 average BLEU on WMT’14.
Takeaways & Limitations
Reward-weighted output augmentation provides a computationally efficient way to incorporate task rewards into supervised likelihood training.
Abstract
from arXiv · showhide
A key problem in structured output prediction is direct optimization of the task reward function that matters for test evaluation. This paper presents a simple and computationally efficient approach to incorporate task reward into a maximum likelihood framework. By establishing a link between the log-likelihood and expected reward objectives, we show that an optimal regularized expected reward is achieved when the conditional distribution of the outputs given the inputs is proportional to their exponentiated scaled rewards. Accordingly, we present a framework to smooth the predictive probability of the outputs using their corresponding rewards. We optimize the conditional log-probability of augmented outputs that are sampled proportionally to their exponentiated scaled rewards. Experiments on neural sequence to sequence models for speech recognition and machine translation show notable improvements over a maximum likelihood baseline by using reward augmented maximum likelihood (RAML), where the rewards are defined as the negative edit distance between the outputs and the ground truth labels.
1 Introduction
Structured prediction needs training objectives that reflect nondifferentiable task rewards, while retaining maximum likelihood’s computational efficiency. RAML connects these goals by training on reward-weighted auxiliary outputs and improves results on speech recognition and translation.
- Task rewards such as BLEU and word error rate are difficult to optimize directly because they are typically nondifferentiable.
- Maximum likelihood trains on ground-truth conditional log-probabilities without distinguishing among alternative incorrect outputs.
- Maximum likelihood remains computationally efficient because locally normalized RNN objectives factorize and require neither inference nor sampling during training.
- RAML samples outputs proportionally to exponentiated scaled rewards, then maximizes their conditional log-likelihood instead of training only on original targets.
- RAML and regularized expected reward optimize KL divergences in opposite directions, while reward-weighted sampling can estimate entropy-regularized reward without sampling from the model.
- RAML consistently improves over maximum-likelihood RNNs on TIMIT speech recognition and WMT’14 machine translation using edit-distance-based output sampling.
2 Reward augmented maximum likelihood
The paper relates maximum likelihood and reward optimization through exponentiated payoff distributions, then uses reward-weighted target sampling to obtain a practical training objective. RAML preserves likelihood-style optimization while incorporating graded preferences among outputs.
- Structured output models score possible outputs, while empirical reward evaluates predictions against ground-truth outputs but is not directly amenable to numerical optimization.
- Maximum likelihood increases target probability and decreases incorrect-output probability, treating all negative outputs as equally wrong.
- Regularized expected reward uses a reward function and entropy term, but stochastic-gradient optimization is difficult because its gradients have high variance.
- The exponentiated payoff distribution weights outputs according to scaled rewards, and the optimal regularized expected-reward model matches this distribution.
- RAML generalizes maximum likelihood by using a nonzero temperature while optimizing the KL divergence from the exponentiated payoff distribution to the model.
- RAML gradients are estimated by sampling auxiliary outputs from the reward distribution and applying mean log-likelihood gradients; at zero temperature, this becomes ordinary maximum likelihood.
- Unlike reinforcement learning, RAML samples from a stationary reward distribution rather than the evolving model distribution, reducing reliance on variance-reduction tricks.
- Edit-distance sampling uses stratification by distance and approximate counts, while importance sampling can extend the approach to BLEU and other sequence metrics.
3 RAML analysis
RAML analyzes the opposite KL directions between the exponentiated payoff distribution and the model distribution, relating their gap to variance differences.
- RAML minimizes D_KL(q ∥ p_θ) rather than the regularized expected-reward objective’s D_KL(p_θ ∥ q).Here q is the exponentiated payoff distribution and p_θ is the model distribution.
- The KL divergence is treated as a Bregman divergence generated by the negative entropy potential, yielding τ-scaled KL divergence for probability vectors.
- The gap between regularized expected reward and the τ-scaled RAML criterion equals a difference of variances on interpolating distributions.Its magnitude decreases as regularization increases.
- The analysis motivates sampling from exponentiated scaled rewards because suitable temperatures can produce lower-variance interpolating distributions than the model distribution.
4 Related Work
Related structured-prediction methods either ignore task reward, use reward without supervision, or combine both; RAML belongs to the hybrid category and emphasizes efficient fixed-distribution sampling.
- Method categories: Structured-output learning is organized into supervised methods, reinforcement-learning methods, and hybrid methods combining supervision with task reward.The paper places RAML in the hybrid category.
- Supervised methods: Conditional log-likelihood methods improve maximum likelihood training but fundamentally do not incorporate task reward.
- Reinforcement learning: Reinforcement-learning approaches optimize task reward without using other supervision, including policy-gradient, actor-critic, Q-learning, and SEARN methods.
- Hybrid methods: Large-margin structured prediction uses supervision and a surrogate task-loss bound but requires loss-augmented inference that is generally inefficient for arbitrary task losses.
- RAML efficiency: RAML samples from a fixed exponentiated payoff distribution, avoiding model sampling and enabling augmentation preprocessing to run in parallel with training.The paper contrasts this with reinforcement learning, scheduled sampling, and structural SVM training costs.
5 Experiments
Experiments compare RAML with maximum likelihood on TIMIT speech recognition and WMT’14 English–French translation, varying temperature and output augmentation.
- Overall results: RAML considerably outperforms maximum likelihood on both speech recognition and machine translation sequence-prediction tasks.
- Speech recognition: The TIMIT system changes only output targets for gradient computation, using negative edit distance to sample augmented outputs with insertions, deletions, and substitutions.
- Speech recognition: At τ = 0.6, more than 60% of TIMIT outputs remain intact, while at τ = 0.9, 5 to 9 edits each have sampling probability larger than 0.1.Augmentation becomes more severe as output length increases.
- Speech recognition: 2.9% PER on the dev set and 2.3% PER on the test set are achieved over the ML baseline at τ = 0.9 on TIMIT.Performance improves through τ = 0.9 and worsens beyond it.
- Machine translation: WMT’14 augmentation uses Hamming distance and one sampled output per sentence at each step for efficiency on the 36M-pair training corpus.Insertions, deletions, and sentence-BLEU sampling are left for future work.
- Machine translation: 0.4 average BLEU points and 0.35 best BLEU points are gained over the ML baseline on WMT’14 at τ = 0.85.BLEU is tokenized and evaluated on the newstest-2014 test set.
6 Conclusion
The paper presents RAML as an efficient maximum-likelihood generalization that directly optimizes task reward through augmented output targets.
- RAML generalizes maximum likelihood by enabling direct optimization of a task reward metric for structured output prediction.The method requires augmenting output targets within a log-likelihood objective.
- Augmented outputs sampled according to edit distance improve maximum-likelihood baselines on both machine translation and speech recognition.
- The framework is intended to apply to probabilistic models with arbitrary reward functions, while other models and more complicated metrics remain future work.
A Proofs
The proofs develop Bregman-divergence identities through convex conjugates and gradient-map transfers, then specialize these relationships to entropy-related potentials and KL divergences.
- Definitions: A Bregman divergence is defined from a strictly convex differentiable closed potential, whose conjugate induces inverse gradient-map transfers.The conjugate is defined by a supremum, with f = ∇F, f* = ∇F*, and f* = f^-1.
- Bregman identities: The divergence identities relate D_F(p ∥ q) and D_F*(r ∥ s) after mapping p and q through the corresponding gradients.The reverse-direction relation uses s = f(p) and r = f(q), and the two expressions are generally unequal.
- KL specialization: Two special potential-function choices yield KL divergences, connecting the general Bregman framework to entropy-based distributions.The passage identifies F_τ(p) = −τH(p) as one of the relevant special cases; the displayed conjugate expressions use log-sum-exp structure.
- Proposition 2: Proposition 2 expresses the KL divergence in both directions and analyzes their difference using Taylor expansions around p and q.The proof introduces an intermediate point from Taylor’s theorem and derives the stated relations through subsequent algebraic identities.
- Proof steps: The proof further rewrites the divergence difference using Hessians and a variance term before applying earlier bounds to obtain the final equality and inequality.The intermediate expressions involve H_F*τ and Var_y∼f*τ(b)[s(y) − r(y)].