Decoding As Dynamic Programming For Recurrent Autoregressive Models

1 INTRODUCTION . Neural autoregressive models ( ARMs ) have shown remarkable performance on various natural language processing tasks such as question answering , machine translation , summarization and reading comprehension ( Anderson et al. , 2018 ; Bahdanau et al. , 2014 ; Wan et al. , 2019 ) . These models are usually trained in an end-to-end manner by optimizing the training objective to learn model parameters . Once the model has been trained , the output is generated by searching for a high scoring sequence given the context and the trained model . This is referred to as the decoding problem . ARMs create a sequence by repeatedly generating the next symbol conditioned on all previous symbols generated . Symbols play dual duty : first they are generated , and next they are incorporated into the conditioning of subsequent decisions . As subsequent decisions can be arbitrarily distant in the sequence , these generation models are non-Markovian , i.e . they do not have a bounded Markov order . As prominent examples , ARMs include Elman ’ s recurrent neural networks ( RNNs ) ( Elman , 1990 ) , conditional text generation models with RNN-based decoder ( Sutskever et al. , 2014 ) , and transformers ( Vaswani et al. , 2017 ) . Exact decoding in ARMs is computationally hard , as the output search space is exponentially large and does not lend itself to efficient algorithms . This is due to non-decomposable long-range interdependencies among the output variables , i.e . an output token directly depends on all of the previously generated tokens . Standard uni-directional decoding algorithms , e.g . greedy and beam search , are ineffective in producing high-scoring output sequences , as errors in the decoding history can adversely affect the future . These algorithms make local decisions to extend an incomplete sequence ( hypothesis ) by selecting the token with the maximum likelihood at each time step , hoping to get a globally optimal complete sequence ( Bahdanau et al. , 2014 ; Sutskever et al. , 2014 ; Mikolov et al. , 2010 ) . In this paper , we present a novel decoding method , based on the method of auxiliary coordinate ( MAC ) , which has been mainly investigated for training deep neural networks ( Carreira-Perpinan & Wang , 2014 ) . Our approach introduces discrete variables for output tokens , and auxiliary continuous variables representing the states of the underlying ARM . The auxiliary variables lead to a factor graph approximation of the ARM with a bounded Markov order ( see Figure 1 ) . We then alternate between optimizing over the output variables and the state variables . The state variables are updated to respect the state dynamics of the underlying ARM and the currently fixed output tokens . The output variables are updated by dynamic programming to exactly optimise a global scoring function , decomposed over local factors determined by the currently fixed state variables . We then extend our MAC-based decoding approach to decode under product of ARM experts ( Hinton , 1999 ) , i.e . an ensemble of ARMs combined additively in log-space , with each using a different generation order . To validate our approach , we evaluate on the text infilling task , which consists of filling missing parts of a sentence or a paragraph ( Horvat & Byrne , 2014 ; Tromble & Eisner , 2009 ; Schmaltz et al. , 2016 ) . Text infilling is challenging as it requires the global structure of the sentence to fill a blank properly . As shown in the figure to right , both greedy and beam search fail to fill the blanks correctly . We conduct experiments on two datasets : SWAG ( Zellers et al. , 2018 ) and Daily dialogue ( Li et al. , 2017 ) with various mask rates . We show that our decoding approach achieve remarkable improvements against the greedy and beam search algorithms as well as TIGS ( Liu et al. , 2019 ) , a recently introduced strong inference method for this task . TIGS decodes by iteratively optimizing the output variables in a continuous relaxation of the discrete output space , and then projecting back the fractional solution to the discrete space . In contrast , the outputs in our approach stay in the discrete space , and we iteratively optimize over both the outputs and the continuous state variables . 2 RELATED WORK . Inference in Relaxed Output Spaces . One approach to inference in deep neural networks is to relax the discrete output space to a continuous space , and then use gradient descent optimisation over the relaxed output variables , including structured prediction energy networks ( SPENs ) ( Belanger & McCallum , 2016 ) and deep value networks ( DVNs ) ( Gygli et al. , 2017 ) . The gradient descent optimisation for inference in SPENs is replaced by an inference network in Tu & Gimpel ( 2018 ) , in order to speed up the prediction process . In all of these works , the relaxed output should be projected back to the discrete space at the end of the inference in the relaxed space . This is usually done by rounding , similar to rounding the fractional solution of a linear program relaxation of an integer linear program . Performing the best projection , in terms of the loss function of interest , is usually a hard problem in itself ; therefore , the quality of the solution is often degraded in the projection step . Our inference approach is similar to that of SPENs and DVNs in that it iteratively optimises and refines the predicted output . However , the output variables always stay in the discrete space in our approach , so we do not need a projection step of the inference process for rounding . Our approach , instead , relaxes the space of continuous state variables to those which may not strictly satisfy the constraints imposed by the dynamics of the underlying ARMs . Cho ( 2016 ) proposed to inject random noise to the hidden states of a conditional recurrent neural language model during greedy or beam search , and execute multiple parallel decoding runs . Our approach , instead , makes structured changes to the state variables in the direction of optimising the decoding objective . Method of Auxiliary Coordinates . Carreira-Perpinan & Wang ( 2014 ) introduced the method of auxiliary variables ( MAC ) to train a neural network . Their idea was to break the interdependence in layers of the neural network by introducing additional variables per data point and per hidden unit . This resulted in smaller problems that could be solved easily given a careful design of the architecture . Taylor et al . ( 2016 ) extends MAC by introducing different auxiliary variables for the linear operations and non-operations in a neural network . This results in three sets of variables involving two auxiliary variables and the weights , which is then optrimized by alternating direction method of multipliers ( ADMM ) . Similarly , ( Wang et al. , 2019 ) introduces new variables and train the neural network using ADMM to globally optimize the training objective , given a careful design of the architecture . 3 OUR DECODING FRAMEWORK . Notations . We denote scalars , vectors and matrices using lower-case , bold lower-case and bold upper-case letters , e.g . y , y and Y . Individual elements of y are denoted as yi . Problem Formulation . Consider a recurrent neural network model for text generation . The probability of generating a sequence y1 , .. , yn under the RNN is decomposed as P ( y1 , .. , yn ) = n∏ i=1 P ( yi|y < i ) , ( 1 ) where P ( .|y < i ) : = softmax ( Whi + b ) , and the state dynamics is hi = f ( hi−1 , yi−1 ) . Decoding then refers to the following optimization problem , arg max y1 , .. , yn logP ( y1 , .. , yn ) = n∑ i=1 logP ( yi|y < i ) ( 2 ) The above optimization problem is computationally hard as it decomposes to conditional probabilities with unbounded length for the conditioning contexts , i.e . a non-Markovian model . Method of Auxiliary Coordinate ( MAC ) . Our goal is to decompose the optimisation problem in eqn ( 2 ) into smaller optimisation problems , which can be solved jointly and coupled via the state variables . This would make decoding more resilient compared t uni-directional decoding , where early errors can adversely affect the future . Let us start by considering the following decomposition of the decoding problem , arg min yn1 , g n 1 − ∑ i logP ( yi|gi ) s.t . gi = f ( gi−1 , yi−1 ) ∀i ∈ [ 2 , .. , n ] . ( 3 ) Notice the appearance of the new explicit variables g in this optimisation formulation of the decoding problem , which mirror the role of h in the underlying RNN . Eliminating the g variables would make this optimisation problem exactly equivalent to the original decoding problem in eqn ( 2 ) . The g serve as continuous auxiliary coordinates , following the MAC technique ( Carreira-Perpinan & Wang , 2014 ) . The objective function in eqn ( 3 ) links the state variable gi to token variable yi generated at the time step i . This is a zero-order decomposition of the log-likelihood function in eqn ( 2 ) , as the terms in the objective do not condition on the previous tokens . Interestingly , we can generalise to a kth order decomposition by linking the state variable gi−k from k steps in the past to the generation of the current token yi . This is due to the fact that the current hidden state , responsible for generating the current token in RNNs , is a deterministic function of the past hidden state gi−k and all of the k generated tokens between that time step and the current time step yi−1i−k . Hence , our k th order decoding optimisation problem is written as follows : arg min yn1 , g n 1 − ∑ i logP ( yi|f ( k ) ( gi−k , yi−1i−k ) ) s.t . gi = f ( gi−1 , yi−1 ) ∀i ∈ [ 2 , .. , n ] , ( 4 ) where f ( k ) ( gi−k , yi−1i−k ) denotes k repeated applications of the RNN ’ s state transition function f ( . ) to compute the current hidden state from the past state gi−k and the tokens observed since then yi−1i−k . For example with k = 2 , f ( 2 ) ( gi−2 , y i−1 i−2 ) = f ( Wf ( Wgi−2 , yi−2 ) , yi−1 ) . The variables gi and yi for indices i ≤ 0 are assigned null values , i.e. , 0 , or a sentence start sentinel . Assuming the constraints are satisfied , this constrained optimisation problem is equivalent to the decoding problem in eqn ( 2 ) . Otherwise , when the constraints are not met , the kth order objective results in a more accurate approximation to the original decoding problem compared to the first-order formulation in eqn ( 3 ) . As k is made larger , the reliance on g is reduced , as f ( k ) uses the RNN ’ s state transition function directly . This will be important in early stages of optimisation , where g does not accurately reflect the RNN state dynamic . The cost of using a higher order , is a higher complexity of inference , which we discuss in Section 4 . In summary , our decoding approach results in a constrained optimisation problem , incorporating two types of factors corresponding to the likelihood terms in the objective function and the constraint violations , respectively , which are denoted by blue and red , in the factor graph of Figure 1 .

Summary: The paper proposes a decoding algorithm for auto-regressive models (ARMs) to improve the decoding accuracy. The key idea is to introduce auxiliary continuous states variables, and alternatively optimize the discrete output variables (i.e., the tokens) and the continuous states. Given the hidden states, the decoding can be efficiently done by veterbi-like algorithms. Then the constraint of the auxiliary variables can be imposed by penalty based continuous based optimization. The paper also extends the idea to the ensemble of ARMs.  On two NLP tasks, the paper shows improvement upon the existing greedy or beach search based approaches.