Source-linked AI summary
SEARNN: Training RNNs with Global-Local Losses
Rémi Leblond, Jean-Baptiste Alayrac, Anton Osokin, Simon Lacoste-Julien
TL;DR
MLE training for RNNs may be a poor surrogate for structured test errors and conditions predictions on ground-truth histories unlike test-time prediction. SEARNN uses learning-to-search exploration to create global-local losses, and the authors report outperforming MLE on three challenging structured-prediction tasks.
Problem
MLE is not a close surrogate for structured test errors and trains next-token predictions on ground-truth histories rather than the model’s own test-time predictions.
Source: §1 · p. 1Method
SEARNN uses learning-to-search roll-outs to score completed sequences with the task’s test error and derive global-local losses attached to RNN steps.
Source: §4 · p. 5Results
SEARNN outperforms MLE on three challenging structured-prediction tasks.
Source: §7 · p. 11Takeaways & Limitations
The authors conclude SEARNN can be trained from scratch and applied to large-vocabulary structured tasks such as neural machine translation.
Source: §7 · p. 11Computing the expectation over an exponential number of possibilities is computationally intractable.
Source: §7 · p. 11
Abstract
from arXiv · showhide
We propose SEARNN, a novel training algorithm for recurrent neural networks (RNNs) inspired by the "learning to search" (L2S) approach to structured prediction. RNNs have been widely successful in structured prediction applications such as machine translation or parsing, and are commonly trained using maximum likelihood estimation (MLE). Unfortunately, this training loss is not always an appropriate surrogate for the test error: by only maximizing the ground truth probability, it fails to exploit the wealth of information offered by structured losses. Further, it introduces discrepancies between training and predicting (such as exposure bias) that may hurt test performance. Instead, SEARNN leverages test-alike search space exploration to introduce global-local losses that are closer to the test error. We first demonstrate improved performance over MLE on two different tasks: OCR and spelling correction. Then, we propose a subsampling strategy to enable SEARNN to scale to large vocabulary sizes. This allows us to validate the benefits of our approach on a machine translation task.
1 INTRODUCTION
SEARNN adapts learning to search for RNN training, replacing MLE with a global-local loss intended to better match test error. The paper reports improvements on OCR, spelling correction, and machine translation, and investigates subsampling to reduce training time while maintaining performance over MLE.
Background
- RNNs generate tokens sequentially, conditioning each next-token prediction on earlier outputs, and are commonly trained by maximizing ground-truth probability with MLE.The paper notes their use in machine translation (Sutskever et al., 2014), parsing (Ballesteros et al., 2016), and caption generation (Vinyals et al., 2015).Source: §1 · p. 1
Limitations of MLE
- MLE is a poor surrogate for structured test errors because it treats near and distant candidates similarly and trains on ground-truth histories unlike test-time prediction.At test time, the RNN feeds its own previous predictions into the next step; this train-test mismatch is called exposure or exploration bias (Ranzato et al., 2016).Source: §1 · p. 1
Contributions
- SEARNN uses learning-to-search ideas to train RNNs with a global-local loss that is closer to test error than MLE.The approach draws on the L2S framework introduced by Daumé et al. (2009) and refined by Ross & Bagnell (2014) and Chang et al. (2015).Source: §1 · p. 2
- The paper reports significant improvements over MLE on OCR and spelling correction, and on machine translation after introducing a subsampling strategy.Subsampling is explored to reduce training time while maintaining improved performance compared to MLE.Source: §1 · p. 2
- The paper reviews MLE limitations, clarifies related claims, connects RNNs to learning to search, and contrasts SEARNN with related L2S- and RL-inspired methods.It also situates its approach among RL-derived training losses, including REINFORCE and ACTOR-CRITIC (Ranzato et al., 2016; Bahdanau et al., 2017).Source: §1 · pp. 1–2
2 TRADITIONAL RNN TRAINING AND ITS LIMITATIONS
Teacher-forced maximum likelihood training can mismatch test-time prediction and structured task losses, motivating exploration of RNN training through structured prediction, particularly learning to search.
Traditional RNN training
- RNNs produce token distributions from hidden states, and teacher-forced MLE trains their joint sequence probability via the chain rule.The model conditions each next-token distribution on preceding tokens and the input; the training inputs are ground-truth tokens.Source: §2 · p. 2
The limitations of MLE training
- Teacher-forced MLE causes exposure bias because test predictions condition on the model’s own prior outputs, so early mistakes can compound along unseen paths.During training, next-token probabilities are conditioned on ground-truth tokens; at test time, they are conditioned on the model’s previous predictions.Source: §2 · p. 2
- MLE maximizes ground-truth sequence probability but ignores structured-loss information, failing to distinguish predictions close to the target from those far away.A single surrogate may not approximate task-specific test errors such as edit distance or F1 score.Source: §2 · p. 3
Local vs. sequence-level
- Although MLE decomposes into local token losses, it is sequence-level under full-history conditioning; conditioning only on recent predictions makes it local and Markovian.The latter assumption reflects imperfect propagation of earlier information and can prevent accurate representation of long-term dependencies.Source: §2 · p. 3
Motivation
- The paper explores structured-prediction approaches to training RNNs, particularly the learning-to-search approach, rather than adapting reinforcement-learning ideas.Source: §2 · p. 3
3 LINKS BETWEEN RNNS AND LEARNING TO SEARCH
L2S reduces structured prediction to sequential cost-sensitive classification, and its roll-in/roll-out search offers RNN training a way to use structured-loss information beyond MLE.
L2S training
- L2S reduces structured prediction to multiclass classification, using a shared local classifier to predict each token conditioned on the input and past decisions (Beygelzimer et al., 2016).Source: §3 · p. 3
- In each round, roll-in creates trajectories and roll-outs score alternative tokens, producing per-time-step cost vectors for training the shared classifier before the policy is updated.Roll-in controls the explored search space, while roll-out determines token costs; policies may be reference, learned, or mixed.Source: §3 · pp. 3–4
Links to RNNs
- RNNs share a cell across time steps and predict recursively from past decisions, matching L2S's shared local classifier; teacher forcing is reference roll-in and model predictions are learned roll-in.Source: §3 · p. 4
- Standard RNN training has no roll-outs, so the section considers whether L2S search can address MLE's limitations by leveraging structured-loss information.Source: §3 · p. 4
4 IMPROVING RNN TRAINING WITH L2S
SEARNN uses roll-outs to derive sequence-level costs for local RNN training losses, bringing structured test-error information into training. The section reports gains over MLE on OCR and spelling correction and identifies effective loss and policy choices.
The SEARNN Algorithm
- At each step of a roll-in trajectory, SEARNN forces each candidate token in turn, rolls out the rest of the sequence, and scores it with the test error.The resulting cost vectors define local losses containing sequence-level information, which are summed across the T steps.Source: §4 · pp. 4–5
Choosing a multi-class classifier
- SEARNN uses the RNN itself as a deep shared classifier, rather than a shallow classifier on cell features, and avoids reducing training to |A|^2 binary classifiers.During roll-outs, a forced token is fed back when the policy requires it and included in the output sequence used to compute cost.Source: §4 · p. 5
Choosing a cost-sensitive loss
- SEARNN's log-loss targets the minimum-cost token, while its KL loss matches the model distribution to a cost-derived target distribution.The log-loss is structurally similar to MLE and becomes equivalent to it with reference roll-in and roll-out when copying the ground truth is optimal; KL uses the costs more directly, with α controlling target-distribution peakedness.Source: §4 · p. 6
Optimization
- SEARNN takes one gradient step on a random mini-batch per round, fixing roll-in and roll-out costs during loss minimization.The resulting gradient is unbiased under round-specific losses; differentiating through the roll-in and roll-out instead remains an open problem.Source: §4 · p. 6
Expected benefits
- SEARNN's test-error-based costs may improve on MLE surrogates, capture differences among candidate predictions, mitigate teacher-forcing exploration bias, and provide full-sequence signals at each cell.The authors suggest that richer structured errors may yield greater improvements and that sequence-level signals may help with vanishing gradients.Source: §4 · p. 6
Key takeaways
- SEARNN significantly outperforms MLE on the OCR and spelling tasks.The experiments use greedy decoding; OCR is evaluated with normalized total Hamming error, while spelling uses edit distance normalized by ground-truth length. The authors report that LL and KL perform best, learned roll-in with mixed roll-out works best, and gains also occur when copying the ground truth is optimal.Source: Table 1 · p. 7, §4 · p. 7
5 SCALING UP SEARNN
SEARNN’s rollout cost grows with sequence length and vocabulary size, limiting its use on large-vocabulary tasks. Subsampling cells and tokens reduces that cost while retaining reported improvements and yielding a 5× running time speedup.
Scaling up SEARNN
- The original |A|T rollouts scale with sequence length and possible tokens, creating a computational bottleneck for large-vocabulary tasks such as machine translation.Forward passes can be parallelized more readily than backward passes, but their asymptotic cost remains O(dT).Source: §5 · p. 7
Key takeaways
- 5× running time speedup accompanied recovery of all full-SEARNN improvements while sampling only a fraction of possible tokens; the best loss was task-dependent.Table 2 compares SEARNN with MLE across datasets using sLL and sKL.Source: Table 2 · p. 8, §5 · p. 8
Scaling up SEARNN
- Subsampling reduces rollout costs by computing cost vectors for only a subset of cells and possible tokens.The expected speedup comes from reducing both quantities that determine the number of rollouts.Source: §5 · p. 8
Sampling strategies
- Step sampling is uniform; token sampling may be uniform, policy-based, or top-k, and every strategy includes the ground-truth action.Top-k is described as a simplified variant of targeted sampling.Source: §5 · p. 8
Adapting our losses to sampling
- sLL and sKL apply a softmax only to sampled tokens, leaving unsampled-token scores alone in the first-order update and avoiding sKL’s need for a default score.The sampled-token layer follows Jean et al. (2015).Source: §5 · p. 8
Experiments
- Experiments sampled five tokens per cell on the two prior datasets, sampling tokens but all cells, with learned roll-in and mixed roll-out.For sLL, top-k performed best; for sKL, biased current-policy sampling performed best.Source: §5 · p. 8
6 NEURAL MACHINE TRANSLATION.
On IWSLT 2014 German-English translation, SEARNN improves over MLE by 2 BLEU points without dropout and performs on par with or better than related methods with comparable baselines. A learned roll-in performed poorly, so the experiments used a reference roll-in.
6 NEURAL MACHINE TRANSLATION.
- The evaluation used IWSLT 2014 German-English data and comparable encoder-decoder architectures, including a bidirectional GRU encoder for BSO and ACTOR-CRITIC comparisons.The training set contained roughly 153k sentence pairs; the MIXER comparison used a convolutional encoder.Source: §6 · p. 9
- SEARNN used a reference roll-in, mixed roll-out, 25 sampled tokens per cell, and validation-selected KL loss scaled by 200; test results were corpus-wide BLEU.Smoothed BLEU supplied sentence-level costs, with a custom reference policy.Source: §6 · p. 9
- SEARNN improves over MLE by 2 BLEU points without dropout, showing benefits on a large-scale, challenging structured prediction task at reasonable computational cost.Source: §6 · p. 9
- Against comparable baselines, SEARNN is similar to MIXER with a convolutional encoder, exceeds BSO in absolute performance and gains, and trails ACTOR-CRITIC in absolute performance despite similar gains.SEARNN requires half as many parameters during training as ACTOR-CRITIC.Source: §6 · p. 9
- The learned roll-in performed poorly, and scheduled sampling yielded no improvement; possible explanations include a weak reference policy or harder optimization with a learned roll-in.Source: §6 · p. 9
7 DISCUSSION
SEARNN uses structured cost information and global-local losses to address limitations of MLE training for RNNs. The paper reports training from scratch, outperforming MLE on three structured-prediction tasks, and scaling to large-vocabulary tasks.
Traditional L2S approaches
- Like LOLS, SEARNN permits customizable roll-in and roll-out strategies; the authors favor learned roll-in with mixed roll-out, except reference roll-in when the reduced problem is too hard.Each intermediate dataset takes one gradient step, so SEARNN’s policy interpolation differs from LOLS and SEARN.Source: §7 · p. 10
- SEARNN integrates learning-to-search into RNN training with adapted intermediate datasets, an RNN-suitable classifier, and cost-informed surrogate losses that are easy to optimize.Its feature extraction is integrated into the model and learnable; roll-out costs for arbitrary tokens require feeding those tokens back into the RNN.Source: §7 · p. 10
L2S-inspired approaches
- Unlike BSO, SEARNN handles general structured losses without requiring meaningful partial-sequence losses; unlike neural AGGREVATE variants, it needs no free cost-to-go oracle.BSO’s sparse sequence-level signal requires warm-starting, whereas SEARNN does not.Source: §7 · pp. 10–11
RL-inspired approaches
- SEARNN’s local losses carry full-sequence costs, avoiding sparse sequence-level feedback and warm-start dependence; deterministic roll-outs also match test decoding.Unlike RL-inspired methods that train baselines or critics, SEARNN needs no additional models; fixing parameter-dependent costs approximates the gradient. With KL loss, it can be viewed as learned label smoothing, deriving targets from intermediate costs rather than adding uniform noise.Source: §7 · p. 11
Conclusion and future work
- SEARNN trains RNNs from scratch, outperforms MLE on three challenging structured-prediction tasks, and scales to tasks with large output vocabularies such as neural machine translation.Its efficient scaling techniques enable its application to tasks with very large output vocabularies.Source: §7 · p. 11
- Future directions include bandit L2S for single observed trajectories, focused costing to lift quadratic sequence-length dependence, and targeted sampling for more efficient large-scale exploration.Source: §7 · p. 11
A ALGORITHMS · A.1 SEARN (ADAPTED FROM DAUMÉ ET AL. (2009), FIGURE 1.) · A.2 SEARNN: REFERENCE ROLL-IN WITH AN RNN.
SEARN builds cost-sensitive training examples by rolling in under a policy, evaluating possible actions with expert roll-outs, and updating the policy iteratively. SEARNN’s reference roll-in uses teacher forcing to obtain the state and ground-truth prefix needed for roll-outs and cost computation.
A.1 SEARN (ADAPTED FROM DAUMÉ ET AL. (2009), FIGURE 1.)
- SEARN initializes its policy from the reference policy, then creates a cost-sensitive dataset in each round over ground-truth input/output pairs.The algorithm repeats this process for N rounds.Source: §A.1 · p. 14
- For each example, SEARN rolls in under the current policy, computes context features at each step, and initializes an empty cost vector.The roll-in predictions are computed once for the example.Source: §A.1 · p. 14
- At each step, SEARN evaluates every possible token by applying an expert policy from the current prefix and token, then comparing the resulting sequence with the ground truth.The comparison supplies the token’s cost in the cost vector.Source: §A.1 · p. 14
- After collecting examples, SEARN learns a classifier and interpolates it with the previous policy before continuing to the next round.The algorithm returns the resulting policy after N rounds.Source: §A.1 · p. 14
A.2 SEARNN: REFERENCE ROLL-IN WITH AN RNN.
- SEARNN treats teacher forcing as the RNN’s reference-policy roll-in.The section develops this analogy for obtaining the state and predictions needed for roll-outs and costs.Source: §A.2 · p. 14
- For a roll-out from cell t, SEARNN obtains the hidden state by running the RNN with teacher-forced ground-truth outputs, and uses the copied ground-truth prefix as prior predictions.It discards the RNN outputs before cell t because the reference roll-in’s predictions up to that point are copied from the ground truth.Source: §A.2 · p. 14
B DESIGN DECISIONS
SEARNN uses the RNN as its classifier, which determines how gradients are handled, and relies on a reference policy whose rollout heuristic can affect performance. The section contrasts this design with a cell-level classifier and discusses when the basic ground-truth heuristic may be suboptimal.
B DESIGN DECISIONS
- SEARNN uses the full RNN as its classifier, with the input and previous decisions forming the intermediate state alongside the cost vector.In an encoder-decoder architecture, the classifier includes both encoder and decoder RNNs.Source: B DESIGN DECISIONS · p. 14
- Using the full RNN requires fully backpropagating each loss gradient, whereas a cell-level classifier would not backpropagate gradient updates.The cell-level view is less aligned with L2S because feature extraction is learned by shared instances of the same RNN cell rather than fixed.Source: B DESIGN DECISIONS · p. 15
- The reference policy selects the action minimizing cost under optimal future decisions, and its roll-in phase always follows the ground truth.The policy is attributed to Daumé et al. (2009).Source: B DESIGN DECISIONS · p. 15
- Computing the optimal reference policy is intractable in general, so rollout requires heuristics that can perform relatively poorly.The paper notes that better reference policies yield better performance, making heuristic choice important.Source: B DESIGN DECISIONS · p. 15
- The paper mostly uses ground truth as its rollout heuristic, but tailored alternatives can improve results on tasks such as machine translation where it is suboptimal.For example, after a skipped token, skipping the next token in rollout may be better than repeating it.Source: B DESIGN DECISIONS · p. 15
C ADDITIONAL EXPERIMENTAL DETAILS · C.1 LOSSES.
The appendix compares alternative losses with the main-text methods: cost-augmented log-loss did not significantly improve over log-loss, while structured hinge and consistent loss also failed to improve over MLE. It also notes equivalences between KL, label smoothing, and MLE under specified setups.
C.1 LOSSES.
- Cost-augmented softmax log-loss incorporates full token costs in the exponential, with a tunable scaling parameter α that balances model scores and costs.The resulting gradient discriminates between tokens according to their costs.Source: §C.1 · p. 15
- Cost-augmented log-loss showed no significant improvement over log-loss; structured hinge was slightly worse than MLE, and consistent loss yielded no significant improvement over MLE.The authors attribute the hinge result possibly to RNNs having more difficulty optimizing its objective than objectives closer to the MLE objective.Source: §C.1 · pp. 15–16
- Structured hinge loss is considered as a cost-sensitive structured SVM loss and can be viewed as a smooth version of cost-augmented softmax log-loss.The structured SVM connection is attributed to Tsochantaridis et al. (2005).Source: §C.1 · p. 16
- The authors suggest RNNs may optimize structured hinge less easily than objectives closer to MLE, which they had been tuned to train well on.Despite enabling RNN learning, structured hinge performed slightly worse overall than MLE.Source: §C.1 · p. 16
- The consistent loss, inspired by traditional structured prediction and Lee et al. (2004), encountered optimization issues that prevented significant gains over MLE.The authors report that they could not obtain significant improvements over the MLE baseline.Source: §C.1 · p. 16
- With Hamming loss, log-loss using reference roll-in and roll-out is equivalent to MLE, while KL produces the same target distribution as label smoothing.The KL equivalence follows because the cost vector becomes equal coordinates minus a one-hot ground-truth vector before softmax.Source: §C.1 · p. 16
C.2 NMT
For NMT, SEARNN used custom sampling and reference policies and switched from learned to reference roll-in. The paper suggests learned roll-in may create a harder reduced problem and lead to worse overall performance.
C.2 NMT
- For NMT, sampling used 15 top-k tokens per cell plus 10 neighboring ground-truth labels to account for common skipping or repetition errors.The vocabulary was large, making low-probability tokens unattractive to sample.Source: §C.2 · p. 16
- The NMT reference policy selected the suffix with the highest BLEU-1 score among ground-truth suffixes added to current predictions.BLEU-4 was avoided because it could favor adding the entire sequence, producing uninformative costs.Source: §C.2 · p. 16
- The NMT experiment switched from learned to reference roll-in, citing the weak reference policy's greater impact on learned roll-in and a harder optimization problem.The authors also proposed a possible gap in standard L2S reduction theory as an explanation.Source: §C.2 · p. 16
- The authors suggest learned roll-in may have produced a harder reduced problem and ultimately worse overall performance than reference roll-in.Standard reduction analysis (Chang et al., 2015) guarantees that reduced-problem classifier performance translates to the initial problem, but does not account for how roll-in/roll-out choices affect reduced-problem difficulty.Source: §C.2 · p. 16
Related papers
This paper cites
- Sequence to Sequence Learning with Neural Networks2014
- Rethinking the Inception Architecture for Computer Vision2015
- Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation2014
- Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks2015
- On Using Very Large Target Vocabulary for Neural Machine Translation2014
- Reinforcement and Imitation Learning via Interactive No-Regret Learning2014