Compositional Generalization without Trees using Multiset Tagging and Latent Permutations (2023)

Abstract: Seq2seq models have been shown to struggle with compositional generalization in semantic parsing, i.e. generalizing to unseen compositions of phenomena that the model handles correctly in isolation. We phrase semantic parsing as a two-step process: we first tag each input token with a multiset of output tokens. Then we arrange the tokens into an output sequence using a new way of parameterizing and predicting permutations. We formulate predicting a permutation as solving a regularized linear program and we backpropagate through the solver. In contrast to prior work, our approach does not place a priori restrictions on possible permutations, making it very expressive. Our model outperforms pretrained seq2seq models and prior work on realistic semantic parsing tasks that require generalization to longer examples. We also outperform non-tree-based models on structural generalization on the COGS benchmark. For the first time, we show that a model without an inductive bias provided by trees achieves high accuracy on generalization to deeper recursion depth.

## introduction sequence-to-sequence models have been very successfully applied to many structural tasks in nlp such as semantic parsing. however, they have also been shown to struggle with compositional generalization @xcite @xcite , i.e. the model fails on examples that contain unseen compositions or deeper recursion of phenomena that it handles correctly in isolation. for example, a model which correctly parses 'mary knew that jim slept' should also be able to parse sentences with deeper recursion than it has seen during training such as 'paul said that mary knew that jim slept'. this sort of generalization is easy for humans but challenging for neural models. in order for a model to generalize compositionally in semantic parsing, it has to identify reusable 'fragments' and be able to systematically combine them in novel ways. one way to make a model sensitive to fragments is to make it rely on a tree that makes the compositional structure explicit. however, this complicates the training because these trees need to be obtained or induced. this is computationally demanding or at least requires structural preprocessing informed by domain knowledge. in this paper, we propose a two-step sequencebased approach with a structural inductive bias that does not rely on trees: the 'fragments' are multisets of output tokens that we predict for every input token in a first step. a second step then arranges the tokens we predicted in the previous step into a single sequence using a permutation model. in contrast to prior work @xcite , there are no hard constraints on the permutations that our model can predict. we show that this enables structural generalization on tasks that go beyond the class of synchronous context-free languages. we overcome two key technical challenges in this work: firstly, we do not have supervision for the correspondence between input tokens and output tokens. therefore, we induce the correspondence during training. secondly, predicting permutations without restrictions is computationally challenging. for this, we develop a differentiable gpu-friendly algorithm. on realistic semantic parsing tasks our approach outperforms previous work on generalization to longer examples than seen at training. we also outperform all other non-tree models on the structural generalization tasks in semantic parsing on cogs @xcite . for the first time, we also show that a model without an inductive bias towards trees can obtain high accuracy on generalization to deeper recursion on cogs. to summarize, our main contributions are: ## overview and motivation our approach consists of two stages. in the first stage (multiset tagging), we predict a multiset z i of tokens for every input token x i from the contextualized representation of x i . this is motivated by the observation that input tokens often systematically correspond to a fragment of the output (like slept corresponding to sleep and agent and a variable in fig. 1 ). importantly, we expect this systematic relationship to be largely invariant to phrases being used in new contexts or deeper recursion. we refer to the elements of the multisets as multiset tokens. in the second stage (permutation), we order the multiset tokens to arrive at a sequential output. conceptually, we do this by going from left to right over the output and determining which multiset token to put in every position. consider the example in fig. 1 . for the first output position, we simply select a multiset token (* in the example). all subsequent tokens are put into position by 'jumping' from the token that was last placed into the output to a new multiset token. in fig. 1 , we jump from * to girl (shown by the outgoing red edge from *). this indicates that girl is the successor of * in the output and hence the second output token. from girl we jump to one of the x 1 tokens to determine what the third output token is and so on. since we predict a permutation, we must visit every multiset token exactly once in this process. the jumps are inspired by reordering in phrasebased machine translation @xcite and methods from dependency parsing, where directly modeling bi-lexical relationships on hidden states has proven successful @xcite . note also that any permutation can be represented with jumps. in contrast, prior work @xcite has put strong restrictions on the possible permutations that can be predicted. our approach is more flexible and empirically it also scales better to longer inputs, which opens up new applications and datasets. ## learning multisets for the multiset tagging, we need to train a model that predicts the multisets z 1 , . . . , z n of tokens by conditioning on the input. we represent a multiset z i as an integer-valued vector that contains for every vocabulary item v the multiplicity of v in z i , i.e. z i,@xmath0. if v is not present in multiset z i , then z i,@xmath1. for example, in fig. 1 , z 3,@xmath2,x 1 = 2. as discussed in sec-tion 2, we do not have supervision for z 1 , . . . , z n and treat them as latent variables. to allow for efficient exact training, we assume that all z i,v are independent of each other conditioned on the input: parametrization. we parameterize p (z i,v | x) as follows. we first pass the input x through a pretrained roberta encoder @xcite , where @xcite is the output of the final layer. we then add roberta's static word embeddings from the first, non-contextualized, layer to that: we then pass h i through a feedforward network obtaining @xmath3: training. the probability of generating the overall multiset m as the union of all z i is the probability that for every vocabulary item v, the total number of occurrences of v across all input positions i sums to m v : this can be computed recursively: let m(y) be the multiset of tokens in the gold sequence y. we train our model with gradient ascent to maximize the marginal log-likelihood of m(y): like @xcite , we found it helpful to initialize the training of our model with high-confidence alignments from an ibm-1 model @xcite (see appendix c.3 for details). preparation for permutation. the scoring function of the permutation model expects a sequence as input. there is no a priori obvious order for the elements within the individual multisets z i . we handle this by imposing a canonical order order(z i ) on the elements of z i by sorting the multiset tokens by their id in the vocabulary. they are then concatenated to form the input z ′ = order(z 1 ) . . . order(z n ) to the permutation model. ## relaxed permutations after predicting a multiset for every input token and arranging the elements within each multiset to form a sequence z ′ , we predict a permutation of z ′ . we represent a permutation as a matrix v that contains exactly one 1 in every row and column and zeros otherwise. we write v i,@xmath4. let p be the set of all permutation matrices. now we formalize the parametrization of permutations as discussed in section 2. we associate a score predicted by our neural network with each permutation v and search for the permutation with the highest score. the score of a permutation decomposes into a sum of scores for binary 'features' of v. we use two types of features. the first type of feature is active if the permutation v maps input position i to output position j (i.e. v @xmath5. we associate this feature with the score s @xmath6. that is, s @xmath7, and analogously s @xmath8. for all output positions j that are neither the first nor the last position, we simply set s @xmath9. the second type of feature models the jumps we introduced in section 2. we introduce a feature that is 1 iff v contains a jump from k to i, and associate this with a score s k↷i . in order for there to be a jump from k to i, the permutation v must map input i to some output position j (v @xmath10,j-1 = 1). hence, the product v k,j-1 v ij is 1 iff there is a jump from k to i at output position j. based on this, the sum j v k,j-1 v ij equals 1 if there is any output position j at which there is a jump from k to i and 0 otherwise. this constitutes the second type of feature. multiplying the features with their respective scores, we want to find the highest-scoring permutation under the following overall scoring function: let v * (s) be the solution to eq. ( 5 ) as a function of the scores. unfortunately, v * (s) does not have sensible derivatives because p is discrete. this makes v * (s) unsuitable as a neural network layer. in addition, eq. ( 5 ) is np-hard (see appendix a.1). we now formulate an optimization problem that approximates eq. ( 5 ) and which has useful derivatives. firstly, we relax the permutation matrix v to a bistochastic matrix u, i.e. u has non-negative elements and every row and every column each sum to 1. secondly, note that eq. ( 5 ) contains quadratic terms. as we will discuss in the next section, our solver assumes a linear objective, so we replace v k,j-1 v ij with an auxiliary variable w ijk . the variable w ijk is designed to take the value 1 if and only if u i,@xmath11,j-1 = 1. we achieve this by coupling w and u using constraints. then, the optimization problem becomes: finally, in combination with the linear objective, the argmax operation still causes the solution u * (s) of eq. ( 6 ) as a function of s to have no useful derivatives. this is because an infinitesimal change in s has no effect on the solution u * (s) for almost all s. to address this, we add an entropy regularization term τ (h(u) + h(w)) to the objective eq. ( 6 ), where h(u) =ij u ij (log u ij -1), and τ > 0 determines the strength of the regularization. the entropy regularization 'smooths' the solution u * (s) in an analogous way to softmax being a smoothed version of argmax. the parameter τ behaves analogously to the softmax temperature: the smaller τ , the sharper u * (s) will be. we discuss how to solve the regularized linear program in section 5. ## inference for relaxed permutations now we describe how to solve the entropy regularized form of eq. ( 6 ) and how to backpropagate through it. this section may be skipped on the first reading as it is not required to understand the experiments; we note that the resulting algorithm (algorithm 1) is conceptually relatively simple. before describing our method, we explain the general principle. ## conclusion in this paper, we have presented a flexible new seq2seq model for semantic parsing. our approach consists of two steps: we first tag each input token with a multiset of output tokens. then we arrange those tokens into a sequence using a permutation model. we introduce a new method to predict and learn permutations based on a regularized linear program that does not restrict what permutations can be learned. the model we present has a strong ability to generalize compositionally on synthetic and natural semantic parsing datasets. our results also show that trees are not necessarily required to generalize well to deeper recursion than seen at training time.