Discovering Non-monotonic Autoregressive Orderings with Variational Inference

1 INTRODUCTION . Autoregressive models have a rich history . Early papers that studied autoregressive models , such as ( Uria et al. , 2016 ) and ( Germain et al. , 2015 ) , showed an interest in designing algorithms that did not require a gold-standard autoregressive order to be known upfront by researchers . However , these papers were overshadowed by developments in natural language processing that demonstrated the power of the left-to-right autoregressive order ( Cho et al. , 2014 ; Sutskever et al. , 2014a ) . Since then , the left-to-right autoregressive order has been essential for application domains such as image captioning ( Vinyals et al. , 2015b ; Xu et al. , 2015 ) , machine translation ( Luong et al. , 2015 ; Bahdanau et al. , 2015 ) and distant fields like image synthesis ( van den Oord et al. , 2016 ) . However , interest in non left-to-right autoregressive orders is resurfacing ( Welleck et al. , 2019b ; Stern et al. , 2019 ) , and evidence ( Vinyals et al. , 2016 ; Gū et al. , 2018 ; Alvarez-Melis & Jaakkola , 2017 ) suggests adaptive orders may produce more accurate autoregressive models . These positive results make designing algorithms that can leverage adaptive orders an important research domain . Inferring autoregressive orderings in a data-driven manner is challenging . Modern benchmarks for machine translation ( Stahlberg , 2019 ) and other tasks ( Oda et al. , 2015 ) are not labelled with gold-standard orders , and left-to-right seems to be the default . This could be explained if domainindependent methodology for identifying high-quality orders is an open question . Certain approaches ( Stern et al. , 2019 ; Welleck et al. , 2019b ; Ruis et al. , 2020 ) use hand-designed loss functions to promote a genre of orders—such as balanced binary trees . These loss functions incorporate ∗Authors contributed equally . certain domain-assumptions : for example , they assume the balanced binary tree order will not disrupt learning . Learning disruption is an important consideration , because prior work shows that poor orders may prohibitively slow learning ( Chen et al. , 2018 ) . Future approaches to inferring autoregressive orders should withhold domain knowledge , to promote their generalization . To our best knowledge , we propose the first domain-independent unsupervised learner that discovers high-quality autoregressive orders through fully-parallelizable end-to-end training without domainspecific tuning . We provide three main contributions that stabilize this learner . First , we propose an encoder architecture that conditions on training examples to output autoregressive orders using techniques in combinatorical optimization . Second , we propose Variational Order Inference that learns an approximate posterior over autoregressive orders . Finally , we develop a practical algorithm for solving the resulting non-differentiable ELBO end-to-end with policy gradients . Empirical results with our solution on image captioning , code generation , text summarization , and machine translation tasks suggest that with similar hyperparameters , our algorithm is capable of recovering autoregressive orders that are even better than fixed orders . Case studies suggest that our learned orders depend adaptively on content , and resemble a type of best-first generation order , which first decodes focal objects and names . Our experimental framework is available at this link . 2 RELATED WORKS . Autoregressive Models Autoregressive models decompose the generation of a high dimensional probability distribution by generating one dimension at a time , with a predefined order . Combined with high capacity neural networks , this approach to modeling complex distributions has been very successful ( Sutskever et al. , 2011 ; Mikolov et al. , 2012 ) . Recent works have achieved great improvements with autoregressive models in many applications , including language modeling ( Radford et al. , 2018 ; 2019 ; Brown et al. , 2020 ) , machine translation ( Sutskever et al. , 2014b ) and image captioning ( Karpathy & Fei-Fei , 2015 ) . Most previous works on autoregressive models use a fixed ordering pre-defined by the designer with left-to-right emerging as the primary choice . In contrast , our method is capable of learning arbitrary orderings conditioned on data and is more flexible . Non-Monotonic Autoregressive Orderings Ford et al . ( 2018b ) shows that a sub-optimal ordering can severely limit the viability of a language model and propose to first generate a partially filled sentence template and then fill in missing tokens . Previous works have also studied bidirectional decoding ( Sun et al. , 2017 ; Zhou et al. , 2019 ; Mehri & Sigal , 2018 ) and syntax trees based decoding ( Yamada & Knight , 2001 ; Charniak et al. , 2003 ; Dyer et al. , 2016 ; Aharoni & Goldberg , 2017 ; Wang et al. , 2018 ) in the natural language setting . However , all of the works mentioned above do not learn the orderings and instead opt to use heuristics to define them . Chan et al . ( 2019 ) performs language modeling according to a known prior , such as balanced binary tree , and does not allow arbitrary sequence generation orders . Welleck et al . ( 2019a ) proposes to use a tree-based recursive generation method to learn arbitrary generation orders . However , their performance lags behind that of left-toright . Gu et al . ( 2019a ) proposes Transformer-InDIGO to allow non-monotonic sequence generation by first pretraining with pre-defined orderings , such as left-to-right , then fine-tuning use Searched Adaptive Order ( SAO ) to find alternative orderings . They report that without pretraining , the learned orders degenerate . In addition , they perform beam search when decoding each token during training , which can not be efficiently parallelized at the sequence length dimension . Emelianenko et al . ( 2019 ) proposes an alternative to SAO , but suffers from similar poor time complexity . In contrast , our method learns high-quality autoregressive orderings directly from data under fully-parallelizable end-to-end training . Variational Methods Our method optimizes the evidence lower bound , or ELBO in short . ELBO is a quantity that is widely used as an optimization proxy in the machine learning literature , where the exact quantity is hard to compute or optimize . Variational methods have achieved great success in machine learning , such as VAE ( Kingma & Welling , 2013 ) and β-VAE ( Higgins et al. , 2017 ) . Combinatorial Optimization Recent works have studied gradient-based optimization in the combinatorial space of permutations ( Mena et al. , 2018 ; Grover et al. , 2019 ; Linderman et al. , 2018 ) . These works have been applied in tasks such as number sorting , jigsaw puzzle solving , and neural signal identification in worms . To our best knowledge , we are the first to build on these techniques to automatically discover autoregressive orderings in vision and language datasets . 3 PRELIMINARIES . The goal of autoregressive sequence modelling is to model an ordered sequence of target values y = ( y1 , y2 . . . , yn ) : yi ∈ R , possibly conditioned on an ordered sequence of source values x = ( x1 , x2 . . . , xm ) : xi ∈ R , where ( x , y ) is sampled from the dataset D. Inspired by Vinyals et al . ( 2015a ) and Gu et al . ( 2019a ) , we formulate the generation process of y as a 2n step process , where at time step 2t − 1 we generate a value , and at timestep 2t we select a not-yet-chosen position in { 1 , 2 , · · · , n } to insert the value . Thus , we introduce the latent sequence variable z = ( z1 , z2 . . . , zn ) : z ∈ Sn , where Sn is the set of one-dimensional permutations of { 1 , 2 , · · · , n } , and zt is defined as the absolute position of the value generated at time step 2t− 1 in the naturally ordered y . Then p ( y , z|x ) denotes the probability of generating y in the ordering of z given the source sequence x . We can thus factorize p ( y , z|x ) using the chain rule : p ( y , z|x ) = p ( yz1 |x ) p ( z1|yz1 , x ) n∏ i=2 p ( yzi |z < i , yz < i , x ) p ( zi|z < i , yz < =i , x ) ( 1 ) For example , p ( y1 , y2 , z1 = 2 , z2 = 1|x ) = p ( y2|x ) p ( z1|y2 , x ) p ( y1|z1 , y2 , x ) p ( z2|y1 , z1 , y2 , x ) is defined as the probability of generating y2 in the first step , then inserting y2 into absolute position 2 , then generating y1 , and finally inserting y1 into absolute position 1 . Note that in practice , the length of y is usually varied . Therefore , we do not first create a fixed-length sequence of blanks and then replace the blanks with actual values . Instead , we dynamically insert a new value at a position relative to the previous values . One common approach to predict such relative position is Pointer Network ( Vinyals et al. , 2015a ) . In other words , at timestep t , we insert the value at position rt relative to the previous generated values . Here , for any z ∈ Sn , r = ( r1 , r2 , . . . , rn ) is constructed such that there is a bijection between Sn and the set of all constructed r. Due to such bijection , we can use z and r interchangeably . We will use z throughout the paper . 4 VARIATIONAL ORDER INFERENCE ( VOI ) Starting from just the original data y in natural order , we can use variational inference to create an objective ( 2 ) that allows us to recover latent order z , parametrized by two neural networks θ and φ . The encoder network φ samples autoregressive orders given the ground truth data , which the decoder network θ uses to recover y . More specifically , φ is a non-autoregressive network ( permutation generator in Fig . 1 ) that takes in the source sequence x and the entire ground truth target sequence y and outputs latent order z in a single forward pass . θ is an autoregressive network ( autoregressive decoder in Fig . 1 ) that takes in x and predicts both the target sequence y and the ordering z through the factorization in Equation ( 1 ) . We name this process Variational Order Inference ( VOI ) . E ( x , y ) ∼D [ log pθ ( y|x ) ] = E ( x , y ) ∼D [ logEz∼qφ ( z|y , x ) [ pθ ( y , z|x ) qφ ( z|y , x ) ] ] ≥ E ( x , y ) ∼D [ Ez∼qφ ( z|y , x ) [ log pθ ( y , z|x ) ] +Hqφ ( ·|y , x ) ] ( 2 ) Here , Hqφ is the entropy term . During training , we train φ and θ jointly to maximize the ELBO in ( 2 ) . During testing , we only keep the decoder θ . To optimize the decoder network θ in ( 2 ) , for each y , we first sample K latents z1 , z2 , . . . , zK from qφ ( ·|y , x ) . We then update θ using the Monte-Carlo gradient estimate Ey∼D [ 1 K ∑K i=1∇θ log pθ ( y , zi|x ) ] . Algorithm 1 Variational Order Inference 1 : Given : encoder network φ with learning rate αφ , decoder network θ with learning rate αθ , entropy coefficient β , batch of training data ( X , Y ) = { ( xb , yb ) } Nb=1 sampled from dataset D 2 : Set gradient accumulators gφ = 0 , gθ = 0 3 : for ( x , y ) ∈ ( X , Y ) do . In practice , this is done through parallel tensor operations 4 : X = φ ( y , x ) 5 : Sample K doubly stochastic matrices B1 , B2 , . . . , BK ∈ Bn×n from G.S . ( X , τ ) 6 : Obtain P1 , P2 , . . . , PK ∈ Pn×n from B1 , B2 , . . . , BK using Hungarian Algorithm 7 : Obtain latents z1 , z2 , . . . , zK = f−1len ( y ) ( P1 ) , f −1 len ( y ) ( P2 ) , . . . , f −1 len ( y ) ( PK ) 8 : gθ = gθ + 1 N ·K ∑K i=1∇θ log pθ ( y , zi|x ) 9 : Calculate log qφ ( zi|y , x ) = 〈X , Pi〉F − log ( perm ( exp ( X ) ) ) ≈ 〈X , Pi〉F − log ( permB ( exp ( X ) ) ) 10 : Calculate b ( y , x ) = 1K ∑K i=1 log pθ ( y , zi|x ) 11 : gφ = gφ + 1 N ·K ∑K i=1∇φ log qφ ( zi|y , x ) ( log pθ ( y , zi|x ) − b ( y , x ) ) + β · ∇φHqφ ( ·|y , x ) 12 : end for 13 : φ = φ+ αφ · gφ 14 : θ = θ + αθ · gθ Optimizing the encoder network φ is tricky . Since z is a discrete latent variable , the gradient from log pθ ( y , z ) does not flow through z . Thus , we formulate ( 2 ) in a reinforcement learning setting with a one-step Markov Decision Process ( S , A , R ) . Under our setting , the state space S = D ; for each state ( x , y ) ∈ D , the action space A ( x , y ) = Slength ( y ) with entropy term Hqφ ( ·|y , x ) ; the reward function R ( ( x , y ) , z ∈ Slength ( y ) ) = log pθ ( y , z|x ) . We can then set the optimization objective L ( φ ) to be ( 2 ) . In practice , we find that adding an entropy coefficient β and gradually annealing it can speed up the convergence of decoder while still obtaining good autoregressive orders . To compute ∇φL ( φ ) , we derive the policy gradient with baseline formulation ( Sutton et al. , 2000 ) : ∇φL ( φ ) = E ( x , y ) ∼D [ Ez∼qφ [ ∇φ log qφ ( z|y , x ) ( log pθ ( y , z|x ) − b ( y , x ) ) ] + β∇φHqφ ] ( 3 ) where b ( y , x ) is the baseline function independent of action z . The reason we use a state-dependent baseline b ( y , x ) instead of a global baseline b is that the the length of y can have a wide range , causing significant reward scale difference . In particular , we set b ( y , x ) = Ez∼qφ [ log pθ ( y , zi|x ) ] . If we sampleK ≥ 2 latents for each y , then we can use its Monte-Carlo estimate 1K ∑K i=1 log pθ ( y , zi|x ) . Since we use policy gradient to optimize φ , we still need a closed form for the distribution qφ ( z|y , x ) . Before we proceed , we define Pn×n as the set of n × n permutation matrices , where exactly one entry in each row and column is 1 and all other entries are 0 ; Bn×n as the set of n × n doubly stochastic matrices , i.e . non-negative matrices whose sum of entries in each row and in each column equals 1 ; R+n×n as the set of non-negative n×nmatrices . Note that we have the relationship Pn×n ⊂ Bn×n ⊂ R+n×n . To obtain qφ ( z|y , x ) , we first write z in two-dimensional form . For each z ∈ Sn , let fn ( z ) ∈ Pn×n be constructed such that fn ( z ) i = one hot ( zi ) , where fn ( z ) i is the i-th row of fn ( z ) . Thus fn is a natural bijection from Sn to Pn×n , and we can rewrite qφ as a distribution over Pn×n such that qφ ( fn ( z ) |y , x ) = qφ ( z|y , x ) . Next , we need to model the distribution of qφ ( ·|y , x ) . Inspired by ( Mena et al. , 2018 ) , we model qφ ( ·|y , x ) as a Gumbel-Matching distribution G.M . ( X ) over Pn×n , where X = φ ( y , x ) ∈ Rn×n is the output of φ . Then for P ∈ Pn×n , qφ ( z|y , x ) = qφ ( f−1n ( P ) |y , x ) = qφ ( P |y , x ) ∝ exp 〈X , P 〉F ( 4 ) where 〈X , P 〉F = trace ( XTP ) is the Frobenius inner product of X and P . To obtain samples in Pn×n from the Gumbel-Matching distribution , ( Mena et al. , 2018 ) relaxes Pn×n to Bn×n by defining the Gumbel-Sinkhorn distribution G.S . ( X , τ ) : τ > 0 over Bn×n , and proves that G.S . ( X , τ ) converges almost surely to G.M . ( X ) as τ → 0+ . Therefore , to approximately sample from G.M . ( X ) , we first sample from G.S . ( X , τ ) , then apply Hungarian algorithm ( Munkres , 1957 ) to obtain P ∈ G.M. ( X ) . Further details are presented in Appendix A . The Gumbel-Matching distribution allows us to obtain the numerator for the closed form of qφ ( z|y , x ) = qφ ( f−1n ( P ) |y , x ) , which equals exp 〈X , P 〉F . However , the denominator is intractable to compute and equals ∑ P∈Pn×n exp 〈X , P 〉F . Upon further examination , we can express it as perm ( exp ( X ) ) , the matrix permanent of exp ( X ) , and approximate it using permB ( exp ( X ) ) , its Bethe permanent . We present details about matrix permanent and Bethe permanent along with the proof that the denominator of qφ ( ·|y , x ) equals perm ( exp ( X ) ) in Appendix B . After we approximate qφ , we can now optimize φ using the policy gradient in ( 3 ) . We present a diagram of our architecture in Figure 1 , and a pseudocode of our algorithm in Algorithm 1 . Note that even though latent space Sn is very large and contains n ! permutations , in practice , if pθ ( y , z ∗|x ) ≥ pθ ( y , z|x ) ∀z ∈ Sn , then pθ ( y , z|x ) tends to increase as the edit distance between z and z∗ decreases . Therefore , φ does not need to search over the entire latent to obtain good permutations , making variational inference over Sn feasible .

Summary: This paper proposes to model the generation order as latent variables for sequence generation tasks, by optimizing the ELBO involving a proposed process of Variational Order Inference (VOI). To alleviate the difficulty of optimizing discrete latent variables, the authors propose to cast it as a one-step Markov Decision problem and optimize it using the policy gradient. The authors also introduce the recent developed Gumbel-matching techniques to derive the close-form of the posterior distribution.