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 designed a new generative model by capturing the auto-regressive order as latent variables for sequence generation task. Based on combinatorical optimization techniques, the authors derived an policy gradient algorithm to optimize the variational lower bound. Empirical results on image caption and code generation showed that this method is superior than both fixed-order generation and previous adaptive-order method transformer-InDIGO. The authors further analyzed the learned orders on global and local level on COCO2017 dataset, demonstrating that the arrangement tend to follows the best-first strategy.