A Continuous Relaxation of Beam Search for End-to-end Training of Neural Sequence Models	Beam search is a desirable choice of test-time decoding algorithm for neural sequence models because it potentially avoids search errors made by simpler greedy methods. However, typical cross entropy training procedures for these models do not directly consider the behaviour of the final decoding method. As a result, for cross-entropy trained models, beam decoding can sometimes yield reduced test performance when compared with greedy decoding. In order to train models that can more effectively make use of beam search, we propose a new training procedure that focuses on the final loss metric (e.g. Hamming loss) evaluated on the output of beam search. While well-defined, this"direct loss"objective is itself discontinuous and thus difficult to optimize. Hence, in our approach, we form a sub-differentiable surrogate objective by introducing a novel continuous approximation of the beam search decoding procedure. In experiments, we show that optimizing this new training objective yields substantially better results on two sequence tasks (Named Entity Recognition and CCG Supertagging) when compared with both cross entropy trained greedy decoding and cross entropy trained beam decoding baselines.

This objective does not directly reason about the behaviour of the final decoding method. As a result, for cross-entropy trained models, beam decoding can sometimes yield reduced test performance when compared with greedy decoding  ,  ,  . These negative results are not unexpected. The training procedure was not search-aware: it was not able to consider the effect that changing the model's scores might have on the ease of search while using a beam decoding, greedy decoding, or otherwise. We hypothesize that the under-performance of beam search in certain scenarios can be resolved by using a better designed training objective. Because beam search potentially offers more accurate search when compared to greedy decoding, we hope that appropriately trained models should be able to leverage beam search to improve performance. In order to train models that can more effectively make use of beam search, we propose a new training procedure that focuses on the final loss metric (e.g. Hamming loss) evaluated on the output of beam search. While well-defined and a valid training criterion, this “direct loss” objective is discontinuous and thus difficult to optimize. Hence, in our approach, we form a sub-differentiable surrogate objective by introducing a novel continuous approximation of the beam search decoding procedure.

In order to accomplish this, we propose a continuous relaxation to the composition of our final loss metric, INLINEFORM0 , and our decoder function, INLINEFORM1 : INLINEFORM2  Specifically, we form a continuous function softLB that seeks to approximate the result of running our decoder on input INLINEFORM0 and then evaluating the result against INLINEFORM1 using INLINEFORM2 . By introducing this new module, we are now able to construct our surrogate training objective: DISPLAYFORM0  Specified in more detail in Section SECREF9 , our surrogate objective in Equation 2 will additionally take a hyperparameter INLINEFORM0 that trades approximation quality for smoothness of the objective. Under certain conditions, Equation 2 converges to the objective in Equation 1 as INLINEFORM1 is increased. We first describe the standard discontinuous beam search procedure and then our training approach (Equation 2) involving a continuous relaxation of beam search.

Results	As shown in Tables 1 and 2, our approach INLINEFORM0 shows significant improvements over the locally normalized CE baseline with greedy decoding for both the tasks (+5.5 accuracy points gain for supertagging and +1.5 F1 points for NER). The improvement is more pronounced on the supertagging task, which is not surprising because: (i) the evaluation metric is tag-level accuracy which is congruent with the Hamming loss that INLINEFORM1 directly optimizes and (ii) the supertagging task itself is very sensitive to the search procedure because tags across time-steps tend to exhibit long range dependencies as they encode specialized syntactic information about word usage in the sentence. Another common trend to observe is that annealing INLINEFORM0 always results in better performance than training with a constant INLINEFORM1 for both INLINEFORM2 (Section SECREF9 ) and INLINEFORM3 (Section SECREF12 ). This shows that a stable training scheme that smoothly approaches minimizing the actual direct loss is important for our proposed approach.

Comparison with Max-Margin Objectives 	Max-margin based objectives are typically motivated as another kind of surrogate training objective which avoid the discontinuities associated with direct loss optimization. Hinge loss for structured prediction typically takes the form: INLINEFORM0   where INLINEFORM0 is the input sequence, INLINEFORM1 is the gold target sequence, INLINEFORM2 is the output search space and INLINEFORM3 is the discontinuous cost function which we assume is decomposable across the time-steps of a sequence. Finding the cost augmented maximum score is generally difficult in large structured models and often involves searching over the output space and computing the approximate cost augmented maximal output sequence and the score associated with it via beam search. This procedure introduces discontinuities in the training procedure of structured max-margin objectives and renders it non amenable to training via backpropagation. Related work  on incorporating beam search into the training of neural sequence models does involve cost-augmented max-margin loss but it relies on discontinuous beam search forward passes and an explicit mechanism to ensure that the gold sequence stays in the beam during training, and hence does not involve back propagation through the beam search procedure itself.

Model	We denote the seq2seq model parameterized by INLINEFORM0 as INLINEFORM1 . We denote the input sequence as INLINEFORM2 , the gold output sequence as INLINEFORM3 and the result of beam search over INLINEFORM4 as INLINEFORM5 . Ideally, we would like to directly minimize a final evaluation loss, INLINEFORM6 , evaluated on the result of running beam search with input INLINEFORM7 and model INLINEFORM8 . Throughout this paper we assume that the evaluation loss decomposes over time steps INLINEFORM9 as: INLINEFORM10 . We refer to this idealized training objective that directly evaluates prediction loss as the “direct loss” objective and define it as: DISPLAYFORM0   Unfortunately, optimizing this objective using gradient methods is difficult because the objective is discontinuous. The two sources of discontinuity are: We introduce a surrogate training objective that avoids these problems and as a result is fully continuous.

In experiments, we show that optimizing this new training objective yields substantially better results on two sequence tasks (Named Entity Recognition and CCG Supertagging) when compared with both cross-entropy trained greedy decoding and cross-entropy trained beam decoding baselines. Several related methods, including reinforcement learning  ,  , imitation learning  ,  ,  , and discrete search based methods  ,  ,  ,  , have also been proposed to make training search-aware. These methods include approaches that forgo direct optimization of a global training objective, instead incorporating credit assignment for search errors by using methods like early updates  that explicitly track the reachability of the gold target sequence during the search procedure. While addressing a related problem – credit assignment for search errors during training – in this paper, we propose an approach with a novel property: we directly optimize a continuous and global training objective using backpropagation. As a result, in our approach, credit assignment is handled directly via gradient optimization in an end-to-end computation graph. The most closely related work to our own approach was proposed by Goyal et al.  . They do not consider beam search, but develop a continuous approximation of greedy decoding for scheduled sampling objectives.

CCG Supertagging	We used the standard splits of CCG bank  for training, development, and testing. The label space of supertags is 1,284 which is much larger than NER. The distribution of supertags in the training data exhibits a long tail because these supertags encode specific syntactic information about the words' usage. The supertag labels are correlated with each other and many tags encode similar information about the syntax. Moreover, this task is sensitive to the long range sequential decisions and search effects because of how it holistically encodes the syntax of the entire sentence. We perform minor preprocessing on the data similar to the preprocessing in  . For this task, we used hidden layers of size 512 and the supertag label embeddings were also of size 512. The standard evaluation metric for this task is the word level label accuracy which directly corresponds to Hamming loss.

We observed in our experiments that hard beam search resulted in predicting more `O's which also hurt the prediction of tags at future time steps and hurt precision as well as recall. Encouragingly, INLINEFORM1 optimization, even though warm started with a CE trained model that performs worse with beam search, led to the NER model becoming more search aware, which resulted in superior performance. However, we also observe that the approximate max-margin approach ( INLINEFORM2 ) performs poorly here. We attribute this to a deficiency in the max-margin objective when coupled with approximate search methods like beam search that do not provide guarantees on finding the supremum: one way to drive this objective down is to learn model scores such that the search for the best hypothesis is difficult, so that the value of the loss augmented decode is low, while the gold sequence maintains higher model score. Because we also warm started with a pre-trained model that results in a worse performance with beam search decode than with greedy decode, we observe the adverse effect of this deficiency. The result is a model that scores the gold hypothesis highly, but yields poor decoding outputs.

Named Entity Recognition	For named entity recognition, we use the CONLL 2003 shared task data  for German language and use the provided data splits. We perform no preprocessing on the data. The output vocabulary length (label space) is 10. A peculiar characteristic of this problem is that the training data is naturally skewed toward one default label (`O') because sentences typically do not contain many named entities and the evaluation focuses on the performance recognizing entities. Therefore, we modify the Hamming cost such that incorrect prediction of `O' is doubly penalized compared to other incorrect predictions. We use the hidden layers of size 64 and label embeddings of size 8. As mentioned earlier, seq2seq models are not an ideal choice for NER (tag-level correlations are short-ranged in NER – the unnecessary expressivity of full seq2seq models over simple encoder-classifier neural models makes training harder). However, we wanted to evaluate the effectiveness of our approach on different instantiations of seq2seq models.

Decoding 	The motivation behind our approach is to make the optimization aware of beam search decoding while maintaining the continuity of the objective. However, since our approach doesn't introduce any new model parameters and optimization is agnostic to the architecture of the seq2seq model, we were able to experiment with various decoding schemes like locally normalized greedy decoding, and hard beam search, once the model has been trained. However, to reduce the gap between the training procedure and test procedure, we also experimented with soft beam search decoding. This decoding approach closely follows Algorithm SECREF7 , but along with soft back pointers, we also compute hard back pointers at each time step. After computing all the relevant quantities like model score, loss etc., we follow the hard backpointers to obtain the best sequence INLINEFORM0 . This is very different from hard beam decoding because at each time step, the selection decisions are made via our soft continuous relaxation which influences the scores, LSTM hidden states and input embeddings at subsequent time-steps. The hard backpointers are essentially the MAP estimate of the soft backpointers at each step. With small, finite INLINEFORM1 , we observe differences between soft beam search and hard beam search decoding in our experiments.

Training with Continuous Relaxation of Beam Search 	We describe our approach in detail in Algorithm 3 and illustrate the soft beam recurrence step in Figure 1. For composing the loss function and the beam search function for our optimization as proposed in Equation 2, we make use of decomposability of the loss function across time-steps. Thus for a sequence y, the total loss is: INLINEFORM0 . In our experiments, INLINEFORM1 is the Hamming loss which can be easily computed at each time-step by simply comparing gold INLINEFORM2 with INLINEFORM3 . While exact computation of INLINEFORM4 will vary according to the loss, our proposed procedure will be applicable as long as the total loss is decomposable across time-steps. While decomposability of loss is a strong assumption, existing literature on structured prediction  ,  has made due with this assumption, often using decomposable losses as surrogates for non-decomposable ones.

Discontinuity in Beam Search	[t] continuous-top-k-argmax [1] INLINEFORM0 INLINEFORM1 , s.t. INLINEFORM2 INLINEFORM3 INLINEFORM4 = 1 to k peaked-softmax will be dominated by scores closer to INLINEFORM5 INLINEFORM6 The square operation is element-wise Formally, beam search is a procedure with hyperparameter INLINEFORM7 that maintains a beam of INLINEFORM8 elements at each time step and expands each of the INLINEFORM9 elements to find the INLINEFORM10 -best candidates for the next time step. The procedure finds an approximate argmax of a scoring function defined on output sequences. We describe beam search in the context of seq2seq models in Algorithm SECREF1 – more specifically, for an encoder-decoder  model with a nonlinear auto-regressive decoder (e.g. an LSTM  ). We define the global model score of a sequence INLINEFORM0 with length INLINEFORM1 to be the sum of local output scores at each time step of the seq2seq model: INLINEFORM2 .

Comparison	We report performance on validation and test sets for both the tasks in Tables 1 and 2. The baseline model is a cross entropy trained seq2seq model (Baseline CE) which is also used to warm start the the proposed optimization procedures in this paper. This baseline has been compared against the approximate direct loss training objective (Section SECREF9 ), referred to as INLINEFORM0 in the tables, and the approximate max-margin training objective (Section SECREF12 ), referred to as INLINEFORM1 in the tables. Results are reported for models when trained with annealing INLINEFORM2 , and also with a constant setting of INLINEFORM3 which is a very smooth but inaccurate approximation of the original direct loss that we aim to optimize. Comparisons have been made on the basis of performance of the models under different decoding paradigms (represented as different column in the tables): locally normalized decoding (CE greedy), hard beam search decoding and soft beam search decoding described in Section SECREF11 .

Conclusion	While beam search is a method of choice for performing search in neural sequence models, as our experiments confirm, it is not necessarily guaranteed to improve accuracy when applied to cross-entropy-trained models. In this paper, we propose a novel method for optimizing model parameters that directly takes into account the process of beam search itself through a continuous, end-to-end sub-differentiable relaxation of beam search composed with the final evaluation loss. Experiments demonstrate that our method is able to improve overall test-time results for models using beam search as a test-time inference method, leading to substantial improvements in accuracy.

Experimental Setup	Since our goal is to investigate the efficacy of our approach for training generic seq2seq models, we perform experiments on two NLP tagging tasks with very different characteristics and output search spaces: Named Entity Recognition (NER) and CCG supertagging. While seq2seq models are appropriate for CCG supertagging task because of the long-range correlations between the sequential output elements and a large search space, they are not ideal for NER which has a considerably smaller search space and weaker correlations between predictions at subsequent time steps. In our experiments, we observe improvements from our approach on both of the tasks. We use a seq2seq model with a bi-directional LSTM encoder (1 layer with tanh activation function) for the input sequence INLINEFORM0 , and an LSTM decoder (1 layer with tanh activation function) with a fixed attention mechanism that deterministically attends to the INLINEFORM1 -th input token when decoding the INLINEFORM2 -th output, and hence does not involve learning of any attention parameters. Since, computational complexity of our approach for optimization scales linearly with beam size for each instance, it is impractical to use very large beam sizes for training.

We detail the continuous relaxation to beam search in Algorithm SECREF7 with INLINEFORM5 being the cumulative loss of beam element INLINEFORM6 at time step INLINEFORM7 and INLINEFORM8 being the embedding matrix of the target vocabulary which is of size INLINEFORM9 where INLINEFORM10 is the size of the embedding vector. In Algorithm SECREF7 , all the discrete selection functions have been replaced by their soft, continuous counterparts which can be backpropagated through. This results in all the operations being matrix and vector operations which is ideal for a GPU implementation. An important aspect of this algorithm is that we no longer rely on exactly identifying a discrete search prediction INLINEFORM0 since we are only interested in a continuous approximation to the direct loss INLINEFORM1 (line 18 of Algorithm SECREF7 ), and all the computation is expressed via the soft beam search formulation which eliminates all the sources of discontinuities associated with the training objective in Equation 1. The computational complexity of our approach for training scales linearly with the beam size and hence is roughly INLINEFORM2 times slower than standard CE training for beam size INLINEFORM3 .

Hyperparameter tuning	For tuning all the hyperparameters related to optimization we trained our models for 50 epochs and picked the models with the best performance on the development set. We also ran multiple random restarts for all the systems evaluated to account for performance variance across randomly started runs. We pretrained all our models with standard cross entropy training which was important for stable optimization of the non convex neural objective with a large parameter search space. This warm starting is a common practice in prior work on complex neural models  ,  ,  .

Finally, backpointers at each time step to the beam elements at the previous time step are also stored for identifying the best sequence INLINEFORM19 , at the conclusion of the search procedure. A backpointer at time step INLINEFORM20 for a beam element INLINEFORM21 is denoted by INLINEFORM22 which points to one of the INLINEFORM23 elements at the previous beam. We denote a vector of backpointers for all the beam elements by INLINEFORM24 . The INLINEFORM25 operation takes as input backpointers ( INLINEFORM26 ) and candidates ( INLINEFORM27 ) for all the beam elements at each time step and traverses the sequence in reverse (from time-step INLINEFORM28 through 1) following backpointers at each time step and identifying candidate words associated with each backpointer that results in a sequence INLINEFORM29 , of length INLINEFORM30 . The procedure described in Algorithm SECREF1 is discontinuous because of the top-k-argmax procedure that returns a pair of vectors corresponding to the INLINEFORM0 highest-scoring indices for backpointers and vocabulary items from the score matrix INLINEFORM1 . This index selection results in hard backpointers at each time step which restrict the gradient flow during backpropagation.

In neural models, the function INLINEFORM3 is implemented as a differentiable mapping, INLINEFORM4 , which yields scores for vocabulary elements using the recurrent hidden states at corresponding time steps. In our notation, INLINEFORM5 is the hidden state of the decoder at time step INLINEFORM6 for beam element INLINEFORM7 , INLINEFORM8 is the embedding of the output symbol at time-step INLINEFORM9 for beam element INLINEFORM10 , and INLINEFORM11 is the cumulative model score at step INLINEFORM12 for beam element INLINEFORM13 . In Algorithm SECREF1 , we denote by INLINEFORM14 the cumulative candidate score matrix which represents the model score of each successor candidate in the vocabulary for each beam element. This score is obtained by adding the local output score (computed as INLINEFORM15 ) to the running total of the score for the candidate. The function INLINEFORM16 in Algorithms SECREF1 and SECREF7 yields successive hidden states in recurrent neural models like RNNs, LSTMs etc. The INLINEFORM17 operation maps a word in the vocabulary INLINEFORM18 , to a continuous embedding vector.

Hence, beam size for all the beam search based experiments was set to 3 which resulted in improvements on both the tasks as discussed in the results. For both tasks, the direct loss function was the Hamming distance cost which aims to maximize word level accuracy.

Additionally, we did not observe a large difference when our soft approximation is used for decoding (Section SECREF11 ) compared to hard beam search decoding, which suggests that our approximation to the hard beam search is as effective as its discrete counterpart. For supertagging, we observe that the baseline cross entropy trained model improves its predictions with beam search decoding compared to greedy decoding by 2 accuracy points, which suggests that beam search is already helpful for this task, even without search-aware training. Both the optimization schemes proposed in this paper improve upon the baseline with soft direct loss optimization ( INLINEFORM0 ), performing better than the approximate max-margin approach.  For NER, we observe that optimizing INLINEFORM0 outperforms all the other approaches but we also observe interesting behaviour of beam search decoding and the approximate max-margin objective for this task. The pretrained CE baseline model yields worse performance when beam search is done instead of greedy locally normalized decoding. This is because the training data is heavily skewed toward the `O' label and hence the absolute score resolution between different tags at each time-step during decoding isn't enough to avoid leading beam search toward a wrong hypothesis path.

The argmax operation can be represented as: INLINEFORM2   which can be relaxed by replacing the indicator function with a peaked-softmax operation with hyperparameter INLINEFORM0 : INLINEFORM1   As INLINEFORM0 , INLINEFORM1 so long as there is only one maximum value in the vector INLINEFORM2 . This peaked-softmax operation has been shown to be effective in recent work  ,  ,  involving continuous relaxation to the argmax operation, although to our knowledge, this is the first work to apply it to approximate the beam search procedure. Using this peaked-softmax operation, we propose an iterative algorithm for computing a continuous relaxation to the top-k-argmax procedure in Algorithm SECREF6 which takes as input a score matrix of size INLINEFORM0 and returns INLINEFORM1 peaked matrices INLINEFORM2 of size INLINEFORM3 . Each matrix INLINEFORM4 represents the index of INLINEFORM5 -th max. For example, INLINEFORM6 will have most of its mass concentrated on the index in the matrix that corresponds to the argmax, while INLINEFORM7 will have most of its mass concentrated on the index of the 2nd-highest scoring element.

Continuous Approximation to top-k-argmax 	[t] Continuous relaxation to beam search [1] INLINEFORM0 , INLINEFORM1 , INLINEFORM2 , INLINEFORM3 , INLINEFORM4 t = 0 to T INLINEFORM5 i=1 to k INLINEFORM6 INLINEFORM7 is a local output scoring function INLINEFORM8 INLINEFORM9 is used to compute INLINEFORM10 INLINEFORM11 Call Algorithm 2 i = 1 to k INLINEFORM12 Soft back pointer computation INLINEFORM13 Contribution from vocabulary items INLINEFORM14 Peaked distribution over the candidates to compute INLINEFORM15 INLINEFORM16 INLINEFORM17 INLINEFORM18 j = 1 to k Get contributions from soft backpointers for each beam element INLINEFORM19 INLINEFORM20 INLINEFORM21 INLINEFORM22 is a nonlinear recurrent function that returns state at next step INLINEFORM23 Pick the loss for the sequence with highest model score on the beam in a soft manner. The key property that we use in our approximation is that for a real valued vector INLINEFORM0 , the argmax with respect to a vector of scores, INLINEFORM1 , can be approximated by a temperature controlled softmax operation.

Introduction	[t] Standard Beam Search [1] INLINEFORM0 , INLINEFORM1 , INLINEFORM2 , INLINEFORM3 t = 0 to T i = 1 to k INLINEFORM4 INLINEFORM5 INLINEFORM6 is the local output scoring function INLINEFORM7 top-k-max INLINEFORM8 Top k values of the input matrix INLINEFORM9 top-k-argmax INLINEFORM10 Top INLINEFORM11 argmax index pairs of the input matrix i = 1 to k INLINEFORM12 embedding( INLINEFORM13 ) INLINEFORM14 INLINEFORM15 is a nonlinear recurrent function that returns state at next step INLINEFORM16 INLINEFORM17 follow-backpointer( INLINEFORM18 ) INLINEFORM19 Sequence-to-sequence (seq2seq) models have been successfully used for many sequential decision tasks such as machine translation  ,  , parsing  ,  , summarization  , dialog generation  , and image captioning  . Beam search is a desirable choice of test-time decoding algorithm for such models because it potentially avoids search errors made by simpler greedy methods. However, the typical approach to training neural sequence models is to use a locally normalized maximum likelihood objective (cross-entropy training)  .

In the next section, we describe a continuous relaxation to the top-k-argmax procedure which forms the crux of our approach.

This observation indicates that using max-margin based objectives with beam search during training actually may achieve the opposite of our original intent: the objective can be driven down by introducing search errors. The observation that our optimization method led to improvements on both the tasks–even on NER for which hard beam search during decoding on a CE trained model hurt the performance–by making the optimization more search aware, indicates the effectiveness of our approach for training seq2seq models.

Our continuous approximation to beam search can very easily be modified to compute an approximation to the structured hinge loss so that it can be trained via backpropagation if the cost function is decomposable across time-steps. In Algorithm SECREF7 , we only need to modify line 5 as: INLINEFORM0   and instead of computing INLINEFORM0 in Algorithm SECREF7 , we first compute the cost augmented maximum score as: INLINEFORM1   and also compute the target score INLINEFORM0 by simply running the forward pass of the LSTM decoder over the gold target sequence. The continuous approximation to the hinge loss to be optimized is then: INLINEFORM1 . We empirically compare this approach with the proposed approach to optimize direct loss in experiments.

Since we have established the pointwise convergence of peaked-softmax to argmax as INLINEFORM4 for all vectors that have a unique maximum value, we can establish pointwise convergence of objective in Equation 2 to objective in Equation 1 as INLINEFORM5 , as long as there are no ties among the top-k scores of the beam expansion candidates at any time step. We posit that absolute ties are unlikely due to random initialization of weights and the domain of the scores being INLINEFORM6 . Empirically, we did not observe any noticeable impact of potential ties on the training procedure and our approach performed well on the tasks as discussed in Section SECREF4 . DISPLAYFORM0   We experimented with different annealing schedules for INLINEFORM0 starting with non-peaked softmax moving toward peaked-softmax across epochs so that learning is stable with informative gradients. This is important because cost functions like Hamming distance with very high INLINEFORM1 tend to be non-smooth and are generally flat in regions far away from changepoints and have a very large gradient near the changepoints which makes optimization difficult.

Other related work involves training a generator with a Gumbel reparamterized sampling module to more reliably find the MAP sequences at decode-time  , and constructing surrogate loss functions  that are close to task losses.

Specifically, we obtain matrix INLINEFORM8 by computing the squared difference between the INLINEFORM9 -highest score and all the scores in the matrix and then using the peaked-softmax operation over the negative squared differences. This results in scores closer to the INLINEFORM10 -highest score to have a higher mass than scores far away from the INLINEFORM11 -highest score. Hence, the continuous relaxation to top-k-argmax operation can be simply implemented by iteratively using the max operation which is continuous and allows for gradient flow during backpropagation. As INLINEFORM0 , each INLINEFORM1 vector converges to hard index pairs representing hard backpointers and successor candidates described in Algorithm SECREF1 . For finite INLINEFORM2 , we introduce a notion of a soft backpointer, represented as a vector INLINEFORM3 in the INLINEFORM4 -probability simplex, which represents the contribution of each beam element from the previous time step to a beam element at current time step. This is obtained by a row-wise sum over INLINEFORM5 to get INLINEFORM6 values representing soft backpointers.

Question: Which loss metrics do they try in their new training procedure evaluated on the output of beam search?
Answer: continuous relaxation to top-k-argmax