Compositional Generalization without Trees using Multiset Tagging and Latent Permutations
Year: 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.1

heading: 1 Introduction
text: 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 (Lake and Baroni, 2018; Finegan-Dollak et al., 2018; Kim and Linzen, 2020; Hupkes et al., 2020), 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.
1https://github.com/namednil/multiset-perm
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 (Wang et al., 2021; Lindemann et al., 2023), 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.
14488
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 (Kim and Linzen, 2020). 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: • a flexible seq2seq model that performs well on
structural generalization in semantic parsing without assuming that input and output are related to each other via a tree structure.
• a differentiable algorithm to parameterize and predict permutations without a priori restrictions on what permutations are possible.

heading: 2 Overview and Motivation
text: Our approach consists of two stages. In the first stage (multiset tagging), we predict a multiset zi of tokens for every input token xi from the contextualized representation of xi. 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 x1 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 (Koehn et al., 2007) and methods from dependency parsing, where directly modeling bi-lexical relationships on hidden states has proven successful (Kiperwasser and Goldberg, 2016). Note also that any permutation can be represented with jumps. In contrast, prior work (Wang et al., 2021; Lindemann et al., 2023) 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.
Setup. We assume we are given a dataset D = {(x1, y1), . . .} of input utterances x and target sequences y. If we had gold alignments, it would be straightforward to train our model. Since we do not have this supervision, we have to discover during training which tokens of the output y belong into which multiset zi. We describe the model and the training objective of the multiset tagging in Section 3. After the model is trained, we can annotate our training set with the most likely z, and then train the permutation model.
For predicting a permutation, we associate a score with each possible jump and search for the highest-scoring sequence of jumps. We ensure that the jumps correspond to a permutation by means of constraints, which results in a combinatorial optimization problem. The flexibility of our model and its parametrization come with the challenge that this problem is NP-hard. We approximate it with a regularized linear program which also ensures differentiability. Our permutation model and its training are described in Section 4. In Section 5, we discuss how to solve the regularized linear program and how to backpropagate through the solution.

heading: 3 Learning Multisets
text: For the multiset tagging, we need to train a model that predicts the multisets z1, . . . , zn of tokens by conditioning on the input. We represent a multiset zi as an integer-valued vector that contains for every vocabulary item v the multiplicity of v in zi, i.e. zi,v = k means that input token i contributes k occurrences of output type v. If v is not present in multiset zi, then zi,v = 0. For example, in Fig. 1, z3,sleep = 1 and z2,x1 = 2. As discussed in Sec-
tion 2, we do not have supervision for z1, . . . , zn and treat them as latent variables. To allow for efficient exact training, we assume that all zi,v are independent of each other conditioned on the input:
P (z | x) = ∏
i,v
P (zi,v | x) (1)
where v ranges over the entire vocabulary. Parametrization. We parameterize P (zi,v | x) as follows. We first pass the input x through a pretrained RoBERTa encoder (Liu et al., 2019), where ENCODER(x) is the output of the final layer. We then add RoBERTa’s static word embeddings from the first, non-contextualized, layer to that:
hi = ENCODER(x)i + EMBED(xi) (2)
We then pass hi through a feedforward network obtaining h̃i = FF(hi) and define a distribution over the multiplicity of v in the multiset zi:
P (zi,v = k|xi) = exp
( h̃Ti wv,k + bv,k )
∑ l exp ( h̃Ti wv,l + bv,l ) (3)
where the weights w and biases b are specific to v and the multiplicity k. In contrast to standard sequence-to-sequence models, this softmax is not normalized over the vocabulary but over the multiplicity, and we have distinct distributions for every vocabulary item v. Despite the independence assumptions in Eq. (1), the model can still be strong because hi takes the entire input x into account.
Training. The probability of generating the overall multiset m as the union of all zi is the probability that for every vocabulary item v, the total number of occurrences of v across all input positions i sums to mv:
P (m|x) = ∏
v
P (z1,v + . . .+ zn,v = mv|x)
This can be computed recursively:
P (z1,v + . . .+ zn,v = mv | x) =∑
k
P (z1,v = k|x)P (z2,v + . . . zn,v = mv − k|x)
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):
∑
(x,y)∈D logP (m(y) | x) (4)
Like Lindemann et al. (2023), we found it helpful to initialize the training of our model with high-confidence alignments from an IBM-1 model (Brown et al., 1993) (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 zi. We handle this by imposing a canonical order ORDER(zi) on the elements of zi by sorting the multiset tokens by their id in the vocabulary. They are then concatenated to form the input z′ = ORDER(z1) . . . ORDER(zn) to the permutation model.

heading: 4 Relaxed Permutations
text: 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 Vi,j = 1 if position i in z′ is mapped to position j in the output y. 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. Vij = 1). We associate this feature with the score si 7→j and use these scores only to model what the first and the last token in the output should be. That is, si 7→1 models the preference to map position i in the input to the first position in the output, and analogously si 7→n models the preference to put i into the last position in the output. For all output positions j that are neither the first nor the last position, we simply set si 7→j = 0.
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 sk↷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 (Vij = 1) and it must also map input position k to output position j − 1 (Vk,j−1 = 1). Hence, the product Vk,j−1Vij is 1 iff there is a jump from k to i at output position
j. Based on this, the sum ∑
j Vk,j−1Vij 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:
arg max V∈P
∑
i,j
Vijsi 7→j +
∑
i,k
sk↷i
 ∑
j
Vk,j−1Vij
  (5)
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 Vk,j−1Vij with an auxiliary variable Wijk. The variable Wijk is designed to take the value 1 if and only if Ui,j = 1 and Uk,j−1 = 1. We achieve this by coupling W and U using constraints. Then, the optimization problem becomes:
arg max U,W
∑
i,j
Uijsi 7→j + ∑
i,j,k
Wijksk↷i (6)
subject to ∑
i Uij = 1 ∀j (6a) ∑
j Uij = 1 ∀i (6b) ∑
k Wijk = Uij ∀j > 1, i (6c) ∑
i
Wijk = Uk(j−1) ∀j > 1, k (6d)
U,W ≥ 0 (6e)
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 Uij(logUij−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.
Predicting permutations. At test time, we want to find the highest scoring permutation, i.e. we want to solve Eq. (5). We approximate this by using U∗(s) instead, the solution to the entropy regularized version of Eq. (6). Despite using a low temperature τ , there is no guarantee that U∗(s) can be trivially rounded to a permutation matrix. Therefore, we solve the linear assignment problem with U∗(s) as scores using the Hungarian Algorithm (Kuhn, 1955). The linear assignment problem asks for the permutation matrix V that maximizes ∑ ij VijU ∗(s)ij .

heading: 4.1 Parametrization
text: We now describe how we parameterize the scores s to permute the tokens into the right order. We first encode the original input utterance x like in Eq. (2) to obtain a hidden representation hi for input token xi. Let a be the function that maps a(i) 7→ j if the token in position i in z′ came from the multiset that was generated by token xj . For example, in Fig. 1, a(6) = 3 since sleep was predicted from input token slept. We then define the hidden representation h′i as the concatenation of ha(i) and an embedding of z′i:
h′i = [ ha(i); EMBED(z′i) ] (7)
We parameterize the scores for starting the output with token i as
si 7→1 = wTstartFFstart(h ′ i)
and analogously for ending it with token i:
si 7→n = wTendFFend(h ′ i)
We set si 7→j = 0 for all other i, j. We parameterize the jump scores sk↷i using Geometric Attention (Csordás et al., 2022) from h′k to h ′ i. Intuitively, Geometric Attention favours selecting the ‘matching’ element h′i that is closest to h′k in terms of distance |i− k| in the string. We refer to Csordás et al. (2022) for details.

heading: 4.2 Learning Permutations
text: We now turn to training the permutation model. At training time, we have access to the gold output y and a sequence z′ from the output of the multiset tagging (see the end of Section 3). We note that whenever y (or z′) contains one vocabulary item at least twice, there are multiple permutations that can be applied to z′ to yield y. Many of these permutations will give the right result for the wrong reasons and the permutation that is desirable for generalization is latent. For example, consider Fig. 2. The token agent is followed by the entity who performs the action, whereas theme is followed by the one affected by the action. The permutation indicated by dashed arrows generalizes poorly to a sentence like Emma knew that James slid since slid will introduce a theme role rather than an agent role (as the sliding is happening to James). Thus, this permutation would then lead to the incorrect output know theme Emma ... slide agent James, in which Emma is affected by the knowing event and James is the one who slides something.
In order to train the permutation model in this setting, we use a method that is similar to EM for structured prediction.2 During training, the model output U∗(s) and W∗(s) often represents a soft permutation that does not permute z′ into y. Our goal is to push the model output into the space of (soft) permutations that lead to y. More formally, let Q ∈ Q(y, z′) be the set of bistochastic matrices such that Qi,j = 0 iff z′i ̸= yj . That is, any permutation included in Q(y, z′) leads to the gold output y when applied to z′.
First, we project the current prediction U∗(s) and W∗(s) into Q(y, z′) using the KL divergence
2It can also be derived as a lower bound to the marginal loglikelihood logPθ(y|x, z′) (see Appendix A.4), and be seen as a form of posterior regularization (Ganchev et al., 2010)
as a measure of distance (E-step):
Û, Ŵ = arg min U∈Q(y,z′),W
KL(U || U∗(s))+ (8)
KL(W || W∗(s))
subject to Eq. (6a) to Eq. (6e). Similar to the M-step of EM, we then treat Û and Ŵ as soft gold labels and train our model to minimize the KL divergence between labels and model:
KL(Û || U∗(s)) + KL(Ŵ || W∗(s))
Eq. (8) can be solved in the same way as the entropy regularized version of Eq. (6) because expanding the definition of KL-divergence leads to a regularized linear program with a similar feasible region (see Appendix A.3 for details).

heading: 5 Inference for Relaxed Permutations
text: 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.

heading: 5.1 Bregman’s Method
text: Bregman’s method (Bregman, 1967) is a method for constrained convex optimization. In particular, it can be used to solve problems of the form
x∗ = arg max x∈C0∩C1...∩Cn,x≥0 sT x + τH(x)︸ ︷︷ ︸ regularizer (9)
where C0, . . . , Cn are linear equality constraints, H(x) = −∑i xi(log xi − 1) is a form of entropy regularization, and τ determines the strength of the regularization. Note that our parameterization of permutations (Eq. (6)) has this form.
Bregman’s method is a simple iterative process. We start with the scores s and then cyclically iterate over the constraints and project the current estimate xi onto the chosen constraint until convergence:
x0 = exp s τ
xi+1 = arg min x∈Ci mod (n−1)
KL(x || xi) (10)
where KL(x | y) = ∑i xi log xiyi − xi + yi is the generalized KL divergence. We call
Algorithm 1 Bregman’s method for Eq. (6) with entropic regularization
function BREGMAN(s, τ ) Uij = exp(τ−1si 7→j) Wijk = exp(τ−1sk↷i) while within budget and not converged do
U,W = KL project(U,W; 6a,6c) with Prop. 2 U,W = KL project(U,W; 6a,6d) with Prop. 2 U = KL project (U, 6b) with Prop. 1
end while return U,W
end function
arg minx∈C KL(x | xi) a KL projection. In order to apply Bregman’s method, we need to be able to compute the KL projection arg minx∈C KL(x | xi) for all C0, . . . , Cn in closed-form. We discuss how to do this for Eq. (6) in the next section.
As an example, consider a problem of the form Eq. (9) with a single linear constraint C∆ = {x |∑
i xi = 1}. In this case, Bregman’s method coincides with the softmax function. This is because the KL projection x∗ = arg minx∈C∆ KL(x || y) for y > 0 has the closed-form solution x∗i = yi∑ i yi
. If we have a closed-form expression for a KL projection (such as normalizing a vector), we can use automatic differentiation to backpropagate through it. To backpropagate through the entire solver, we apply automatic differentiation to the composition of all projection steps.

heading: 5.2 Bregman’s Method for Eq. (6)
text: In order to apply Bregmans’ method to solve the entropy regularized version of Eq. (6), we need to decompose the constraints into sets which we can efficiently project onto. We choose the following three sets here: (i), containing Eq. (6a) and Eq. (6c), and (ii), containing Eq. (6a) and Eq. (6d), and finally (iii), containing only Eq. (6b). We now need to establish what the KL projections are for our chosen sets. For (iii), the projection is simple:
Proposition 1. (Benamou et al. (2015), Prop. 1) For A,m > 0, the KL projection arg minU KL(U || A) subject to ∑ j Uij = mi is given by U∗ij = mi Aij∑ j′ Aij′ .
Let us now turn to (i) and (ii). The constraints Eq. (6d) and Eq. (6c) are structurally essentially the same, meaning that we can project onto (ii) in basically the same manner as onto (i). We project onto (i), with the following proposition:
Proposition 2. For A,B > 0, the KL projection
arg min U,W
KL(U || A) + KL(W || B)
subject to ∑
i Uij = 1 ∀j ∑
k
Wijk = Uij ∀j, i
(11)
is given by:
U∗ij = Tij∑ i′ Ti′j
W∗ijk = U ∗ ij Bijk∑ k′ Bijk′
where Tij = √ Aij · ∑
k Bijk. The proof can be found in Appendix A.2. We present the overall algorithm in Algorithm 1, and note that it is easy to implement for GPUs. In practice, we implement all projections in log space for numerical stability.

heading: 6 Evaluation

heading: 6.1 Doubling Task
text: Our permutation model is very expressive and is not limited to synchronous context-free languages. This is in contrast to the formalisms that other approaches rely on (Wang et al., 2021; Lindemann et al., 2023). To evaluate if our model can structurally generalize beyond the synchronous contextfree languages in practice, we consider the function F = {(w,ww) | w ∈ Σ∗}. This function is related to processing challenging natural language phenomena such as reduplication and cross-serial dependencies. We compare our model with an LSTM-based seq2seq model with attention and a Transformer in the style of Csordás et al. (2021) that uses a relative positional encoding. Since the input is a sequence of symbols rather than English, we replace RoBERTa with a bidirectional LSTM and use randomly initialized embeddings. The models are trained on inputs of lengths 5 to 10 and evaluated on longer examples. The results can be found in Fig. 3. All models get perfect or close to perfect accuracy on inputs of length 11 but accuracy quickly deteriorates for the LSTM and the Transformer. In contrast, our model extrapolates very well to longer sequences.

heading: 6.2 COGS
text: COGS is a synthetic benchmark for compositional generalization introduced by Kim and Linzen
(2020). Models are tested for 21 different cases of generalization, 18 of which focus on using a lexical item in new contexts (Lex). There are 1000 instances per generalization case. Seq2seq models struggle in particular with the structural generalization tasks (Yao and Koller, 2022), and we focus on those: (i) generalization to deeper PP recursion than seen during training ("Emma saw a hedgehog on a chair in the garden beside the tree..."), (ii) deeper CP recursion ("Olivia mentioned that James saw that Emma admired that the dog slept"), and (iii) PPs modifying subjects when PPs modified only objects in the training data (OS).
We follow previous work and use a lexicon (Akyurek and Andreas, 2021) to map some input tokens to output tokens (see Appendix C.2 for details). We also use this mechanism to handle the variable symbols in COGS.
We report the means and standard deviations for 10 random seeds in Table 1. Our approach obtains high accuracy on CP and PP recursion but exact match accuracy is low for OS. This is in part because our model sometimes predicts semantic representations for OS that are equivalent to the gold standard but use a different order for the conjuncts. Therefore, we report accuracy that accounts for this in Table 2. In both tables, we also report the impact of using a simple copy mechanism instead of the more complex lexicon induction mechanism (-Lex). Our model outperforms all other non-treebased models by a considerable margin.
Structural generalization without trees. All previous methods that obtain high accuracy on recursion generalization on COGS use trees. Some approaches directly predict a tree over the input (Liu et al., 2021; Weißenhorn et al., 2022), while others use derivations from a grammar for data augmentation (Qiu et al., 2022) or decompose the input along a task-specific parse tree (Drozdov et al., 2022). Our results show that trees are not as important for compositional generalization as their success in the literature may suggest, and that weaker
structural assumptions already reap some of the benefits.
Logical forms with variables. COGS uses logical forms with variables, which were removed in conversion to variable-free formats for evaluation of some approaches (Zheng and Lapata, 2021; Qiu et al., 2022; Drozdov et al., 2022). Recently, Wu et al. (2023) have argued for keeping the variable symbols because they are important for some semantic distinctions; we keep the variables.

heading: 6.3 ATIS
text: While COGS is a good benchmark for compositional generalization, the data is synthetic and does not contain many phenomena that are frequent in semantic parsing on real data, such as paraphrases that map to the same logical form. ATIS (Dahl et al., 1994) is a realistic English semantic parsing dataset with executable logical forms. We follow the setup of our previous work (Lindemann et al., 2023) (L’23) and use the variable-free FunQL rep-
resentation (Guo et al., 2020). Apart from the usual iid split, we evaluate on a length split, where a model is trained on examples with few conjuncts and has to generalize to longer logical forms with more conjuncts. For a fair comparison with previous work, we do not use the lexicon/copying. We also evaluate a version of our model without RoBERTa that uses a bidirectional LSTM and GloVe embeddings instead. This mirrors the model of L’23.
Table 3 shows mean accuracy and standard deviations over 5 runs. Our model is competitive with non-pretrained models in-distribution, and outperforms all other models on the length generalization. The high standard deviation on the length split stems from an outlier run with 18% accuracy – the second worst-performing run achieved an accuracy of 44%. Even without pretraining, our model performs very well. In particular, without grammarbased decoding our model performs on par or outperforms L’23 with grammar-based decoding.
The runtime of the model in L’23 is dominated by the permutation model and it takes up to 12 hours to train on ATIS. Training the model presented here only takes around 2 hours for both stages.
Performance breakdown. In order for our approach to be accurate, both the multiset tagging model and the permutation model have to be accurate. Table 4 explores which model acts as the bottleneck in terms of accuracy on ATIS and COGS. The answer depends on the dataset: for the synthetic COGS dataset, predicting the multisets correctly is easy except for OS, and the model struggles more with getting the permutation right. In contrast, for ATIS, the vast majority of errors can be attributed to the first stage.

heading: 6.4 Okapi
text: Finally, we consider the recent Okapi (Hosseini et al., 2021) semantic parsing dataset, in which an English utterance from one of three domains (Calendar, Document, Email) has to be mapped to an API request. We again follow the setup of L’23 and evaluate on their length split, where a model has to generalize to longer logical forms. In contrast to all other datasets we consider, Okapi is quite noisy because it was collected with crowd workers. This presents a realistic additional challenge on top of the challenge of structural generalization.
The results of 5 runs can be found in Table 5. Our model outperforms both BART (Lewis et al., 2020) and the model of L’23. In the comparison without pretraining, our model also consistently achieves higher accuracy than the comparable model of L’23 without grammar-based decoding.

heading: 7 Related Work
text: Predicting permutations. Mena et al. (2018) and Lyu and Titov (2018) use variational autoencoders based on the Sinkhorn algorithm to learn latent permutations. The Sinkhorn algorithm (Sinkhorn, 1964) is also an instance of Bregman’s method and solves the entropy regularized version of Eq. (6) without the W-term. This parameterization is con-
siderably weaker than ours since it cannot capture our notion of ‘jumps’.
Wang et al. (2021) compute soft permutations as an expected value by marginalizing over the permutations representable by ITGs (Wu, 1997). This approach is exact but excludes some permutations. In particular, it excludes permutations needed for COGS.3 In addition, the algorithm they describe takes a lot of resources as it is both O(n5) in memory and compute. Devatine et al. (2022) investigate sentence reordering methods. They use bigram scores, which results in a similar computational problem to ours. However, they deal with it by restricting what permutations are possible to enable tractable dynamic programs. Eisner and Tromble (2006) propose local search methods for decoding permutations for machine translation.
Outside of NLP, Kushinsky et al. (2019) have applied Bregman’s method to the quadratic assignment problem, which Eq. (5) is a special case of. Since they solve a more general problem, using their approach for Eq. (6) would require O(n4) rather than O(n3) variables in the linear program.
Compositional generalization. Much research on compositional generalization has focused on lexical generalization with notable success (Andreas, 2020; Akyurek and Andreas, 2021; Conklin et al., 2021; Csordás et al., 2021). Structural generalization remains more challenging for seq2seq models (Yao and Koller, 2022).
Zheng and Lapata (2022) modify the transformer architecture and re-encode the input and partially generated output for every decoding step to disentangle the information in the representations. Structure has also been introduced in models by means of grammars: Qiu et al. (2022) heuristically induce a quasi-synchronous grammar (QCFG, Smith and Eisner (2006)) and use it for data augmentation for a seq2seq model. Kim (2021) introduces neural QCFGs which perform well on compositional generalization tasks but are very compute-intensive. Other works directly parse into trees or graphs inspired by methods from syntactic parsing (Liu et al., 2021; Herzig and Berant, 2021; Weißenhorn et al., 2022; Jambor and Bahdanau, 2022; Petit and Corro, 2023).
Several approaches, including ours, have decoupled the presence or absence of output tokens from their order: Wang et al. (2021) train a model end-
3This is mostly due to the treatment of definite descriptions, which appear clustered together at the start of the logical form.
to-to-end to permute the input (as discussed above) and then monotonically translate it into an output sequence. Lindemann et al. (2023) also present an end-to-end differentiable model that first applies a ‘fertility step’ which predicts for every word how many copies to make of its representation, and then uses the permutation method of Wang et al. (2021) to reorder the representation before translating them. Cazzaro et al. (2023) first translate the input monotonically and feed it into a second model. They use alignments from an external aligner to train the first model. The second model is a tagger or a pretrained seq2seq model and predicts the output as a permutation of its input. We compare against such a baseline for permutations in Appendix B, finding that it does not work as well as ours in the compositional generalization setups we consider.

heading: 8 Conclusion
text: 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.
Limitations
The conditional independence assumptions are a limitation for the applicability of our multiset tagging model. For example, the independence assumptions are too strong to apply it to natural language generation tasks such as summarization. From a technical point of view, the independence assumptions are important to be able to induce the latent assignment of output tokens to multisets efficiently. Future work may design multiset tagging methods that make fewer independence assumptions.
While our method for predicting permutations is comparatively fast and only has a memory requirement of O(n3), inference on long sequences, e.g. with more than 100 tokens, remains somewhat
slow. In future work, we plan to investigate other approximate inference techniques like local search and dual decomposition.
Regarding the importance of trees for compositional generalization, our model has no explicit structural inductive bias towards trees. However, we do not exclude that the pretrained RoBERTa model that we use as a component implicitly captures trees or tree-like structures to a certain degree.

heading: Acknowledgements
text: We thank Bailin Wang and Jonas Groschwitz for technical discussions; we thank Hao Zheng for discussions and for providing system outputs for further analysis. We also say thank you to Christine Schäfer and Agostina Calabrese for their comments on this paper.
ML is supported by the UKRI Centre for Doctoral Training in Natural Language Processing, funded by the UKRI (grant EP/S022481/1), the University of Edinburgh, School of Informatics and School of Philosophy, Psychology & Language Sciences, and a grant from Huawei Technologies. IT is supported by the Dutch National Science Foundation (NWO Vici VI.C.212.053).

heading: A Math Details
text: A.1 NP-hardness
We show that Eq. (5) can be used to decide the Hamiltonian Path problem. Let G = (V,E) be a graph with nodes V = {1, 2, . . . , n}. A Hamiltonian path P = v1, v2, . . . , vn is a path in G (i.e. (vi, vi+1) ∈ E for all i) such that each node of G appears exactly once. Deciding if a graph has a Hamiltonian path is NP-complete.
Reduction of Hamiltonian path to Eq. (5). Note that a necessary but not sufficient condition for P to be a Hamiltonian path is that P is a permutation of V . This will be ensured by the constraints on the solution in Eq. (5).
We construct a score function
sk↷i = { 1 if (k, i) ∈ E 0 else
(12)
and let si 7→j = 0 for all i, j. If we find the solution of Eq. (5) for the score function Eq. (12), we obtain a permutation P of V , which may or may not be a path in G. In a path of n nodes, there are n − 1 edges that are crossed. If the score of the solution is n − 1, then all node pairs (vi, vi+1) that are adjacent in P must have had a score of 1, indicating an edge (vi, vi+1) ∈ E. Therefore, P must be a Hamiltonian path. If the score of the solution is less than n− 1, then there is no permutation of V that is also a path, and hence G has no Hamiltonian path.
A.2 Proof of Proposition 2
We now prove Prop. 2 using a very similar technique as Kushinsky et al. (2019). As the constraints
in Prop. 2 for any value of j do not interact with constraints for other values of j, we can assume w.l.o.g. that j takes a single value only and drop it in the notation. We want to solve:
arg min x,Y
KL(x || z) + KL(Y || W)
subject to ∑
i
xi = 1
∑
k
Yik = xi ∀i
(13)
We can find Y∗ = arg minY KL(Y || W) subject to ∑ j Yi,j = xi based on Prop. 1: Y ∗ i,j = xiWi,j∑ j Wi,j . That is, we can express Y∗ as a function of x (which we write Y∗(x)), and therefore our overall problem is now a problem in one variable (x):
arg min x
KL(x || z) + KL(Y∗(x) || W)
subject to ∑
i
xi = 1 (14)
We can now rewrite the objective function as
arg min x
KL(x || q)
subject to ∑
i
xi = 1 (15)
where qi = √ zi · ∑
j′ Wi,j′ . This step is justified in detail in Fig. 4.
The rewritten optimization problem has the right form to apply Prop. 1 a second time. We obtain:
x∗i = qi∑ i′ qi′
By plugging this into Y∗(x), we obtain the solution to the overall optimization problem.
A.3 Reduction of Eq. (8) to Eq. (6) In this section, we show how Eq. (8) can be reduced to a problem of the entropy regularized version of Eq. (6). This is useful because it means we can use Algorithm 1 to solve Eq. (8).
First, we show that computing a KL projection is equivalent to solving an entropy-regularized linear program. Let C be the feasible region of the linear constraints.
arg max x∈C
sT x − τ ∑
i
xi(log xi − 1)
= arg min x∈C
τ ∑
i
xi(log xi − si τ − 1)
= arg min x∈C KL(x || exp( s τ ))
Due to this, Eq. (8) is equivalent to a linear program that has the same feasible region as Eq. (6) except for the additional constraint U ∈ Q(y, z′). Note that U ∈ Q(y, z′) essentially rules out certain correspondences. Therefore we can approximately enforce U ∈ Q(y, z′) by masking U∗(s) such that any forbidden correspondence receives a very low score.
A.4 Derivation of loss function as ELBO We now show how the training procedure we use to train our permutation model can be derived from a form of evidence lower bound (ELBO).
Ideally, our permutation model would be a distribution Pθ(R|x, z′) over permutation matrices R and we would maximize the marginal likelihood, i.e. marginalizing over all permutations:
P (y|x, z′) = ∑
R∈P Pθ(R|x, z′)P (y|z′,R) (16)
where P (y|z′,R) = ∏j ∑
iRij · 1(yj = zi) with 1 being the indicator function. P (y|z′,R) returns 1 iff applying the permutation R to z′ results in y. Unfortunately, computing Eq. (16) exactly is intractable in general due to the sum over permutation matrices. We instead use techniques from variational inference and consider the following evidence lower bound (ELBO):
logP (y|x, z′) ≥max Q ER∼Q(R|x,z′,y) logP (y|z′,R)
− KL(Q(R|x, z′, y) || Pθ(R|x, z′)) (17)
where Q(R|x, z′, y) is an approximate variational posterior. We now relax the restriction that P (R|x, z′) places non-zero mass only on permutation matrices and use the following definition of Pθ(R|x, z′):
Pθ(Rij = 1|x, z′) = U∗(s)ij
where U∗(s) is the solution to Eq. (6) with added entropy regularization.
It turns out, in our case, we can easily construct a variational posterior Q that has zero reconstruction loss (the first term on the right side in Eq. (17)): we can choose any Q(R|x, z′, y) ∈ Q(y, z′) where Q(y, z′) is the set of bistochastic matrices such that Q(R|x, z′, y)i,j = 0 iff z′i ̸= yj . To see that this gives zero reconstruction error, consider position j in the output: The probability mass is distributed across precisely those positions i in z′ where the
right kind of token lives. In other words, any alignment with non-zero probability will reconstruct the output token at position j.
Therefore we can use the following lower bound to the log-likelihood:
log P (y|x, z′) ≥ (18) − min
Q∈Q(y,z) KL(Q(R|x, z′, y) || Pθ(R|x, z′))
During training, we need to compute the gradient of Eq. (18). By Danskin’s theorem (Danskin, 1967), this is:
−∇θKL(Q∗ || Pθ(R|x, z′)) (19)
where Q∗ ∈ Q(y, z′) is the minimizer of Eq. (18). Note that Q∗ can equivalently be characterised as Û (Eq. (8)).
In practice, we also add −KL(Ŵ|W∗(s)) to our objective in Eq. (18) to speed up convergence; this does not change the fact that we use a lower bound.

heading: B Additional results and analysis
text: Okapi. In Fig. 5 we show the accuracy of our model on the document domain in comparison with previous work by number of conjuncts in the logical form.
Permutation baseline. A simpler approach for predicting a permutation of the output z′ from
the multiset tagging is to use a seq2seq model. In order to compare our approach to such a baseline, we concatenate the original input x with a separator token and z′. We then feed this as input to a BART-base model which is trained to predict the output sequence y. At inference time, we use beam search and enforce the output to be a permutation of the input. As detailed in Table 6, this approach works well in-distribution and it also shows a small improvement over finetuning BART directly on the length split of ATIS. However, it does not perform as well as our approach. On COGS, our model outperforms the permutation baseline by an even bigger margin. Unseen variable symbols could be a challenge for BART on COGS which might explain part of the gap in performance.
This approach towards predicting permutations is similar to that of Cazzaro et al. (2023) except that they do not constrain the beam search to permutations. We found that not constraining the output to be a permutation worked worse in the compositional generalization setups.

heading: C Further model details
text: C.1 Parametrization of permutation model
We do not share parameters between the multiset tagging model and the permutation model.
Tokens that appear more than once in the same multiset have the same representation h′i in Eq. (7). In order to distinguish them, we concatenate another embedding to h′i: if the token z′i is the k-th instance of its type in its multiset, we concatenate an embedding for k to h′i. For example, in Fig. 1, z′5 = ‘x1’ and it is the second instance of ‘x1’ in its multiset, so we use the embedding for 2.
We found it helpful to make the temperature τ of the scores for Algorithm 1 dependent on the number of elements in the permutation, setting τ = (log n)−1, so that longer sequences have slightly sharper distributions.
Since the permutation model is designed to model exactly permutations, during training, z and y must have the same elements. This is not guaranteed because z is the prediction of the multiset model which may not have perfect accuracy on the training data. For simplicity, we disregard instances where z and y do not have the same elements. In practice, this leads to a very small loss in training data for the permutation model.
C.2 Lexicon mechanism
The lexicon L is a lookup table that deterministically maps an input token xi to an output token L(xi), and we modify the distribution for multiset tagging as follows: P ′(zi,v = k|x) = { P (zi,L = k|xi) if v = L(xi) P (zi,v = k|x) else
where P (zi,v = k|x) is as defined in Eq. (3) and L is a special lexicon symbol in the vocabulary. P (zi,L |xi) is a distribution over the multiplicity of L(xi), independent of the identity of L(xi). We use the ‘simple’ lexicon induction method by Akyurek and Andreas (2021). Unless otherwise specified during learning, L(xi) = xi like in a copy mechanism.
Handling of variables in COGS. For the COGS dataset, a model has to predict variable symbols. The variables are numbered (0-based) by the input token that introduced it (e.g. in Fig. 1, slept, the third token, introduces a variable symbol x2). In order to robustly predict variable symbols for sentences with unseen length, we use a similar mechanism as the lexicon look up table: we introduce another special symbol in the vocabulary, Var. If Var is predicted with a multiplicity of k at i-th input token, it adds the token xi−1 to its multiset k times.
C.3 Initialization of Multiset Tagging model
If there are l alignments with a posterior probability of at least χ that an input token i produces token v, we add the term λ logP (zi,v ≥ l | x) to Eq. (4). λ is the hyperparameter determining the strength. This additional loss is only used during the first g epochs.

heading: D Datasets and Preprocessing
text: We show basic statistics about the data we use in Table 7. Except for the doubling task, all our datasets are in English. COGS uses a small fragment of English generated by a grammar, see Kim and Linzen (2020) for details.
Doubling task. For the doubling task, we use an alphabet of size |Σ| = 11. To generate inputs with a specific range of lengths, we first draw a length from the range uniformly at random. The symbols in the input are also drawn uniformly at random and then concatenated into a sequence. Examples of lengths 5 - 10 are used as training, examples of
length 11 are used as development data (e.g. for hyperparameter selection), and examples of length 11 - 20 are used as test data.
D.1 Preprocessing
COGS. Unlike Zheng and Lapata (2022); Qiu et al. (2022); Drozdov et al. (2022) we do not apply structural preprocessing to the original COGS meaning representation and keep the variable symbols: all our preprocessing is local and aimed at reducing the length of the logical form (to keep runtimes low). We delete any token in {",","_","(",")","x",".",";","AND"} as these do not contribute to the semantics and can be reconstructed easily in post-processing. The tokens {"agent", "theme", "recipient", "ccomp", "xcomp", "nmod", "in", "on", "beside"} are always preceded by a "." and we merge "." and any of those tokens into a single token.
Example: * cookie ( x _ 3 ) ; * table ( x _ 6 ) ; lend . agent ( x _ 1 , Dylan ) AND lend . theme ( x _ 1 , x _ 3 ) AND lend . recipient ( x _ 1 , x _ 9 ) AND cookie . nmod . beside ( x _ 3 , x _ 6 ) AND girl ( x _ 9 )
Becomes * cookie 3 * table 6 lend .agent 1 Dylan lend .theme 1 3 lend .recipient 1 9 cookie .nmod .beside 3 6 girl 9
ATIS. We follow the pre-procressing by Lindemann et al. (2023) and use the variable-free FunQL representation as annotated by Guo et al. (2020). We use spacy 3.0.5 (model en_core_web_sm) to tokenize the input.
Okapi. Again, we follow the preprocessing of Lindemann et al. (2023). We use spacy 3.0.5 (model en_core_web_sm) to tokenize both the input utterances and the output logical forms.

heading: E Details on evaluation metrics
text: We provide code for all evaluation metrics in our repository.
Doubling. We use exact match accuracy on the string.
COGS. For COGS we use exact match accuracy on the sequence in one evaluation setup. The other evaluation setup disregards the order of conjuncts: we first remove the ‘preamble’ (which contains all the definite descriptions) from the conjunctions. We count a prediction as correct if the set of definite descriptions in the preamble matches the set of definite descriptions in the gold logical form and the set of clauses in the prediction match the set of clauses in the gold logical form.
ATIS. We allow for different order of conjuncts between system output and gold parse in computing accuracy. We do this by sorting conjuncts before comparing two trees node by node. This is the same evaluation metric as used by Lindemann et al. (2023).
Okapi. We follow Hosseini et al. (2021); Lindemann et al. (2023) and disregard the order of the parameters for computing accuracy. We use a case-insensitive string comparison.

heading: F Hyperparameters
text: We use the same hyperparameters for all splits of a dataset. For our model, we only tune the hyperparameters of the multiset tagging model; the permutation model is fixed, and we use the same configuration for all tasks where we use RoBERTa. For model ablations where we use an LSTM instead of RoBERTa, we use the same hyperparameters for Okapi and ATIS, and a smaller model for the doubling task. These configurations were determined by hand without tuning. For BART, we use the same hyperparameter as Lindemann et al. (2023).
We follow the random hyperparameter search procedure of Lindemann et al. (2023) for the multiset tagging models and the LSTM/transformer we train from scratch: we sample 20 configurations and evaluate them on the development set. We run the two best-performing configurations again with a different random seed and pick the one with the highest accuracy (comparing the union of the predicted multisets with the gold multiset). We then train and evaluate our model with entirely different random seeds.
The chosen hyperparameters along with the search space are provided in the github repository.

heading: G Number of parameters, computing infrastructure and runtime
text: We show the number of parameters in the models we train in Table 9.
All experiments were run on GeForce GTX 1080 Ti or GeForce GTX 2080 Ti with 12GB RAM and Intel Xeon Silver or Xeon E5 CPUs.
The runtime of one run contains the time for training, evaluation on the devset after each epoch and running the model on the test set. We show runtimes of the model we train in Table 8. Since we evaluate on 5 random seeds (10 for COGS due to high variance of results), our experiments overall took around 64 hours of compute time on our computing infrastructure.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:

heading: 3 A1. Did you describe the limitations of your work?
text: Limitations, at the end of the paper
7 A2. Did you discuss any potential risks of your work? No apparent societal risks.

heading: 3 A3. Do the abstract and introduction summarize the paper’s main claims?
text: Abstract and 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts?

heading: 6, Appendix D
text: 3 B1. Did you cite the creators of artifacts you used? 6
7 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? The data we create programmatically is likely too simple to be protected by copyright.

heading: 3 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided
text: that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? 6
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Not applicable. Left blank.

heading: 3 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and
text: linguistic phenomena, demographic groups represented, etc.? 6, Appendix D
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. 6, Appendix D
C 3 Did you run computational experiments? 6

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? Appendix H
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? 6, Appendix G, code submission.

heading: 3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? 6, Appendix B

heading: 3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did
text: you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Appendix D
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: Future work may design multiset tagging methods that make fewer independence assumptions. While our method for predicting permutations is comparatively fast and only has a memory requirement of O(n3), inference on long sequences, e. g. with more than 100 tokens, remains somewhat
slow. In future work, we plan to investigate other approximate inference techniques like local search and dual decomposition. Regarding the importance of trees for compositional generalization, our model has no explicit structural inductive bias towards trees. However, we do not exclude that the pretrained RoBERTa model that we use as a component implicitly captures trees or tree-like structures to a certain degree