[Paper Review] Reward Augmented Maximum Likelihood for Neural Structured Prediction

Mohammad Norouzi, Samy Bengio+5 more|arXiv (Cornell University)|Sep 1, 2016

Topic Modeling49 references89 citations

TL;DR

This paper proposes Reward Augmented Maximum Likelihood (RAML), a method that improves neural structured prediction by incorporating task-specific rewards into maximum likelihood training. By sampling outputs proportionally to their exponentiated rewards (e.g., inverse edit distance to ground truth), RAML optimizes log-likelihood on these augmented outputs, achieving significant gains over standard maximum likelihood in speech recognition (TIMIT) and machine translation (WMT’14), with improvements of up to +0.8 BLEU and -0.6 PER.

ABSTRACT

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.

Motivation & Objective

• Address the limitation of maximum likelihood training in structured prediction, which treats all incorrect outputs equally regardless of their proximity to the ground truth.
• Overcome the non-differentiability and sparsity of common evaluation metrics like BLEU and word error rate, which are difficult to optimize directly.
• Develop a computationally efficient alternative to reinforcement learning-based methods that require high-variance policy gradients and complex sampling during training.
• Bridge the gap between maximum likelihood and expected reward optimization by introducing a regularized objective that favors high-reward outputs.
• Enable direct optimization of task-specific rewards while preserving the computational efficiency and ease of implementation of standard maximum likelihood training.

Proposed method

• Introduce a reward-augmented objective where, for each input, outputs are sampled with probability proportional to exp(β × R(y|x)), where R(y|x) is the reward of output y given input x and β is a temperature hyperparameter.
• Optimize the conditional log-likelihood of these sampled outputs instead of the ground truth, effectively regularizing the model to assign higher probability to high-reward outputs.
• Theoretical analysis shows that RAML minimizes a KL divergence between the exponentiated reward distribution and the model’s predicted distribution, with the optimal solution balancing expected reward and entropy.
• Use a temperature parameter τ (inverted β) to control the concentration of the sampling distribution: lower τ increases focus on high-reward outputs.
• Apply the method to sequence-to-sequence models using RNNs with attention, where rewards are defined as negative edit distance (or Hamming distance) between generated and reference sequences.
• Train using standard stochastic gradient descent with asynchronous updates, maintaining the same architecture and hyperparameters across baseline and RAML models.

Experimental results

Research questions

• RQ1 Can we improve maximum likelihood training in structured prediction by incorporating task-specific rewards without introducing high-variance policy gradient updates?
• RQ2 How does the performance of RAML compare to standard maximum likelihood training when optimizing for non-differentiable metrics like BLEU or word error rate?
• RQ3 What is the optimal level of reward-based augmentation (i.e., temperature τ) for achieving the best generalization on structured prediction tasks?
• RQ4 Does sampling from a reward-proportional distribution lead to better alignment with test-time evaluation metrics than standard likelihood training?
• RQ5 Can RAML be applied effectively to both sequence-to-sequence models in machine translation and speech recognition with minimal architectural changes?

Key findings

• On the TIMIT speech recognition dataset, RAML reduced phone error rate (PER) by up to 0.6 points on the dev set and 0.8 points on the test set compared to the maximum likelihood baseline.
• On the WMT’14 English-to-French machine translation benchmark, RAML with τ=0.85 achieved an average BLEU score of 36.91 and a best BLEU of 37.23, outperforming the ML baseline by 0.41 and 0.36 points respectively.
• The optimal temperature τ was found to be 0.85 for machine translation and 1.00 for speech recognition, indicating that the degree of reward emphasis depends on the task and reward function.
• Performance degraded when augmentation was too strong (e.g., τ < 0.75), suggesting that overemphasizing high-reward outputs can hurt generalization.
• The improvement from RAML was consistent across multiple random seeds and evaluation runs, indicating robustness and reproducibility.
• Despite its simplicity, RAML significantly outperformed the strong maximum likelihood baseline, demonstrating that reward-aware training can enhance model performance without complex training procedures.