Structured Reordering for Modeling Latent Alignments in Sequence Transduction

1 Introduction Recent advances in deep learning have led to major progress in many domains , with neural models sometimes achieving or even surpassing human performance [ 49 ] . However , these methods often struggle in out-of-distribution ( ood ) settings where train and test examples are drawn from different distributions . In particular , unlike humans , conventional sequence-to-sequence ( seq2seq ) models , widely used in natural language processing ( NLP ) , fail to generalize systematically [ 4 , 27 , 28 ] , i.e. , correctly interpret sentences representing novel combinations of concepts seen in training . Our goal is to provide a mechanism for encouraging systematic generalization in seq2seq models . To get an intuition about our method , consider the semantic parsing task shown in Figure 1 . A learner needs to map a natural language ( NL ) utterance to a program which can then be executed on a knowledge base . To process the test utterance , the learner needs to first decompose it into two segments previously observed in training ( shown in green and blue ) , and then combine their corresponding program fragments to create a new 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . program . Current seq2seq models fail in this systematic generalization setting [ 12 , 24 ] . In contrast , traditional grammar formalisms decompose correspondences between utterances and programs into compositional mappings of substructures [ 46 ] , enabling grammar-based parsers to recombine rules acquired during training , as needed for systematic generalization . Grammars have proven essential in statistical semantic parsing in the pre-neural era [ 51 , 57 ] , and have gained renewed interest now as a means of achieving systematic generalization [ 18 , 41 ] . However , grammars are hard to create and maintain ( e.g. , requiring grammar engineering or grammar induction stages ) and do not scale well to NLP problems beyond semantic parsing ( e.g. , machine translation ) . In this work , we argue that the key property of grammar-based models , giving rise to their improved ood performance , is that a grammar implicitly encodes alignments between input and output segments . For example , in Figure 1 , the expected segment-level alignments are ‘ the length → len ’ and ‘ the longest river → longest ( river ( all ) ) ’ . The encoded alignments allow for explicit decomposition of input and output into segments , and consistent mapping between input and output segments . In contrast , decision rules employed by conventional seq2seq models do not exhibit such properties . For example , recent work [ 15 ] shows that primitive units such as words are usually inconsistently mapped across different contexts , preventing these models from generalizing primitive units to new contexts . Instead of developing a full-fledged grammar-based method , we directly model segment-level alignments as structured latent variables . The resulting alignment-driven seq2seq model remains end-to-end differentiable , and , in principle , applicable to any sequence transduction problem . Modeling segment-level alignments requires simultaneously inducing a segmentation of input and output sequences and discovering correspondences between the input and output segments . While segment-level alignments have been previously incorporated in neural models [ 50 , 55 ] , to maintain tractability , these approaches support only monotonic alignments . The monotonicity assumption is reasonable for certain tasks ( e.g. , summarization ) , but it is generally overly restrictive ( e.g. , consider semantic parsing and machine translation ) . To relax this assumption , we complement monotonic alignments with an extra reordering step . That is , we first permute the source sequence so that segments within the reordered sequence can be aligned monotonically to segments of the target sequence . Coupling latent permutations with monotonic alignments dramatically increases the space of admissible segment alignments . The space of general permutations is exceedingly large , so , to allow for efficient training , we restrict ourselves to separable permutations [ 5 ] . We model separable permutations as hierarchical reordering of segments using permutation trees . This hierarchical way of modeling permutations reflects the hierarchical nature of language and hence is arguably more appropriate than ‘ flat ’ alternatives [ 33 ] . Interestingly , recent studies [ 45 , 47 ] demonstrated that separable permutations are sufficient for capturing the variability of permutations in linguistic constructions across natural languages , providing further motivation for our modeling choice . Simply marginalizing over all possible separable permutations remains intractable . Instead , inspired by recent work on modeling latent discrete structures [ 9 , 14 ] , we introduce a continuous relaxation of the reordering problem . The key ingredients of the relaxation are two inference strategies : marginal inference , which yields the expected permutation under a distribution ; MAP inference , which returns the most probable permutation . In this work , we propose efficient dynamic programming algorithms to perform exact marginal and MAP inference with separable permutations , resulting in effective differentiable neural modules producing relaxed separable permutations . By plugging these modules into an existing module supporting monotonic segment alignments [ 55 ] , we obtain end-to-end differentiable seq2seq models , supporting non-monotonic segment-level alignments . In summary , our contributions are : • A general seq2seq model for NLP tasks that accounts for latent non-monotonic segment-level alignments . • Novel and efficient algorithms for exact marginal and MAP inference with separable permutations , allowing for end-to-end training using a continuous relaxation.1 • Experiments on synthetic problems and NLP tasks ( semantic parsing and machine translation ) showing that modeling segment alignments is beneficial for systematic generalization . 1Our code and data are available at https : //github.com/berlino/tensor2struct-public . 2 Background and Related Work . 2.1 Systematic Generalization . Human learners exhibit systematic generalization , which refers to their ability to generalize from training data to novel situations . This is possible due to the compositionality of natural languages - to a large degree , sentences are built using an inventory of primitive concepts and finite structurebuilding mechanisms [ 8 ] . For example , if one understands ‘ John loves the girl ’ , they should also understand ‘ The girl loves John ’ [ 13 ] . This is done by ‘ knowing ’ the meaning of individual words and the grammatical principle of subject-verb-object composition . As pointed out by Goodwin et al . [ 15 ] , systematicity entails that primitive units have consistent meaning across different contexts . In contrast , in seq2seq models , the representations of a word are highly influenced by context ( see experiments in Lake and Baroni [ 27 ] ) . This is also consistent with the observation that seq2seq models tend to memorize large chunks rather than discover underlying compositional principles [ 20 ] . The memorization of large sequences lets the model fit the training distribution but harms out-ofdistribution generalization . 2.2 Discrete Alignments as Conditional Computation Graphs . Latent discrete structures enable the incorporation of inductive biases into neural models and have been beneficial for a range of problems . For example , input-dependent module layouts [ 2 ] or graphs [ 37 ] have been explored in visual question answering . There is also a large body of work on inducing task-specific discrete representations ( usually trees ) for NL sentences [ 9 , 17 , 36 , 54 ] . The trees are induced simultaneously with learning a model performing a computation relying on the tree ( typically a recursive neural network [ 43 ] ) , while optimizing a task-specific loss . Given the role the structures play in these approaches – i.e. , defining the computation flow – we can think of the structures as conditional computation graphs . In this work , we induce discrete alignments as conditional computation graphs to guide seq2seq models . Given a source sequence x with n tokens and a target sequence y withm tokens , we optimize the following objective : X = Encodeθ ( x ) Lθ , φ ( x , y ) = − logEpφ ( M |X ) pθ ( y|X , M ) ( 1 ) where Encode is a function that embeds x into X ∈ Rn×h with h being the hidden size , M ∈ { 0 , 1 } n×m is the alignment matrix between input and output tokens . In this framework , alignments M are separately predicted by pφ ( M |X ) to guide the computation pθ ( y|X , M ) that maps x to y . The parameters of both model components ( φ and θ ) are disjoint . Relation to Attention Standard encoder-decoder models [ 3 ] rely on continuous attention weights i.e. , M [ : , i ] ∈ 4n−1 for each target token 1 ≤ i ≤ m. Discrete versions of attention ( aka hard attention ) have been studied in previous work [ 11 , 53 ] and show superior performance in certain tasks . In the discrete case M is a sequence of m categorical random variables . Though discrete , the hard attention only considers word-level alignments , i.e. , assumes that each target token is aligned with a single source token . This is a limiting assumption ; for example , in traditional statistical machine translation , word-based models ( e.g. , [ 6 ] ) are known to achieve dramatically weaker results than phrase-based models ( e.g. , [ 26 ] ) . In this work , we aim to bring the power of phrase-level ( aka segment-level ) alignments to neural seq2seq models . 2 3 Latent Segment Alignments via Separable Permutations . Our method integrates a layer of segment-level alignments with a seq2seq model . The architecture of our model is shown in Figure 2 . Central to this model is the alignment network , which decomposes the alignment problem into two stages : ( i ) input reordering and ( ii ) monotonic alignment between the reordered sequence and the output . Conceptually , we decompose the alignment matrix from Eq 1 into two parts : M = MpeMmo ( 2 ) 2One of our models ( see Section 3.2 ) still has a flavor of standard continuous attention in that it approximates discrete alignments with continuous expectation . where Mpe ∈ Rn×n is a permutation matrix , and Mmo ∈ Rn×m represents monotonic alignments . With this conceptual decomposition , we can rewrite the objective in Eq 1 as follows : Lθ , φ ( x , y ) = − logEpφ ( Mpe|x ) Epφ′ ( Mmo|MpeX ) pθ ( y|MpeX , Mmo ) ( 3 ) where MpeX denotes the reordered representation . With a slight abuse of notation , φ now denotes the parameters of the model generating permutations , and φ′ denotes the parameters used to produce monotonic alignments . Given the permutation matrix Mpe , the second expectation Epφ ( Mmo|MpeX ) pθ ( y|MpeX , Mmo ) , which we denote as pθ , φ′ ( y|MpeX ) , can be handled by existing methods , such as SSNT [ 55 ] and SWAN [ 50 ] . In the rest of the paper , we choose SSNT as the module for handling monotonic alignment.3 We can rewrite the objective we optimize in the following compact form : Lθ , φ , φ′ ( x , y ) = − logEpφ ( Mpe|x ) pθ , φ′ ( y|MpeX ) ( 4 ) 3.1 Structured Latent Reordering by Binary Permutation Trees . Inspired by Steedman [ 47 ] , we restrict word reorderings to separable permutations . Formally , separable permutations are defined in terms of binary permutation trees ( aka separating trees [ 5 ] ) , i.e. , if a permutation can be represented by a permutation tree , it is separable . A binary permutation tree over a permutation of a sequence 1 . . . n is a binary tree in which each node represents the ordering of a segment i . . . j ; the children exhaustively split their parent into sub-segments i . . . k and k+1 . . . j . Each node has a binary label that decides whether the segment of the left child precedes that of the right child . saw the girl the hedgehog re-ordered sentence Bracketing transduction grammar [ BTG , 52 ] , which is proposed in the context of machine translation , is the corresponding context-free grammar to represent binary permutation trees . Specifically , BTG has one non-terminal ( X ) and three anchored rules : Si , j , k : Xki Straight−−−−−→ XjiXkj Ii , j , k : Xki Inverted−−−−−−→ XjiXkj Ti : Xi+1i → xi where Xki is the anchored non-terminal covering the segment from i to k ( excluding k ) . The first two rules decide whether to keep or invert two segments when construct- ing a larger segment ; the last rule states that every word xi in an utterance is associated with a non-terminal Xi+1i . An example is shown in Figure 3 . Through this example , we note that the first two rules only signify which segments to inverse ; an additional process of interpreting the tree ( i.e. , performing actual actions of keeping or inverting segments ) is needed to obtain the permutated sequence . This hierarchical approach to generating separable permutations reflects the compositional nature of language , and , thus , appears more appealing than using ‘ flat ’ alternatives [ 10 , 16 , 33 ] . Moreover , with BTGs , we can incorporate segment-level features to model separable permutations , and design tractable algorithms for learning and inference . 3In our initial experiments , we found that SWAN works as well as SSNT but is considerably slower . By assigning a score to each anchored rule using segment-level features , we obtain a distribution over all possible derivations , and use it to compute the objective in Eq 4. pφ ( D|x ) = ∏ R∈D fφ ( R ) Z ( x , φ ) , Lθ , φ , φ′ ( x , y ) = − logEpφ ( D|x ) pθ , φ′ ( y|M D pe X ) ( 5 ) where fφ is a score function assigning a ( non-negative ) weight to an anchored rule R ∈ { S , I , T } , Z ( x , φ ) = ∑ D′ ∏ R∈D′ fφ ( R ) is the partition function , which can be computed using the inside algorithm , MDpe is the permutation matrix corresponding to the derivation D. BTG , along with the weight assigned for each rule , is a weighted context-free grammar ( WCFG ) . In this WCFG , the weight is only normalized at the derivation level . As we will see in Algorithm 1 , we are interested in normalizing the weight of production rules and converting the WCFG to an equivalent PCFG following Smith and Johnson [ 42 ] , so that the probability of a derivation can be computed as follows : pφ ( D|x ) = ∏ R∈D Gφ ( R ) ( 6 ) where Gφ ( R ) is the weight of the production rule R under the transformed PCFG . The details of the conversion are provided in the Appendix . The challenge with optimizing the objective in Eq 5 is that the search space of possible derivations is exponential , making the estimation of the gradients with respect to parameters of the reordering component ( φ ) non-trivial . We now present two differentiable surrogates we use .

Summary: The authors propose a model for sequence transduction tasks that allows for source tokens to be reordered and then monotonically aligned with the target tokens. The reordering and the alignment are latent, and since all monotonic alignments can be marginalized over using the approach introduced in the SSNT paper (Yu et al., 2016), the authors focus on approximating the marginalization over the latent source permutation. The authors consider in particular separable permutations, whose derivations are induced by a PCFG. The authors propose two approximation approaches, one which uses a dynamic program to obtain an expected permutation matrix (which is used instead of the expected likelihood under the permutation-matrix distribution), and one which samples a hard permutation matrix during learning and uses a ST-Gumbel-Softmax gradient estimator. The authors experiment on a number of synthetic and small-scale real-world tasks and generally find that their approach outperforms baselines, and is especially performant in settings that require generalization to longer inputs than those seen in training.