Disentangled Sequence to Sequence Learning for Compositional Generalization (2022)

Abstract: There is mounting evidence that existing neural network models, in particular the very popular sequence-to-sequence architecture, struggle to systematically generalize to unseen compositions of seen components. We demonstrate that one of the reasons hindering compositional generalization relates to representations being entangled. We propose an extension to sequence-to-sequence models which encourage disentanglement by adaptively re-encoding (at each time step) the source input. Specifically, we condition the source representations on the newly decoded target context which makes it easier for the encoder to exploit specialized information for each prediction rather than capturing it all in a single forward pass. Experimental results on semantic parsing and machine translation empirically show that our proposal delivers more disentangled representations and better generalization.

## introduction when humans use language, they exhibit compositional generalization; they are able to produce and understand a potentially infinite number of novel linguistic expressions by systematically combining known atomic components @xcite . for example, if a person knows the meaning of the utterance "a boy ate the cake on the table in a house" and the verb "like", it is natural for them to understand the utterance "a boy likes the cake on the table in a house" when they encounter it for the first time (see there has been a long standing debate whether this systematicity can be captured by connectionist architectures @xcite @xcite and recent years have witnessed a resurgence of interest thanks to the tremendous success of neural networks at various natural language understanding and generation tasks @xcite @xcite . mounting evidence, however, suggests that existing models, in particular the very popular sequence-tosequence architecture, struggle with compositional generalization @xcite @xcite . this failure may be due to spurious correlations which hinder out-of-distribution generalization @xcite @xcite or limited robustness to perturbations in the input @xcite . in this paper, we identify an entanglement problem with how different semantic factors (e.g., lexical meaning and semantic relations) are represented in neural sequence models that hurts generalization. in theory, neural networks should represent semantic factors in a disentangled way by virtue of the principle of compositionality @xcite which implies that semantic properties of syntactic constituents are to a certain extent context invariant and the semantic primitives they express are conditionally independent. disentangled meaning representations ought to preserve this conditional independence, and neural units modeling a particular semantic factor should be relatively invariant to changes in other factors @xcite . for example, the relation between "table" and "house" in motivated by this finding, we propose an extension to sequence-to-sequence (seq2seq) models which allows us to learn disentangled representations for compositional generalization. specifically, at each time step of the decoding, we adaptively re-encode the source input by conditioning the source representations on the newly decoded target context. we therefore build specialized representations which make it easier for the encoder to exploit relevant-only information for each prediction. experiments on three benchmarks, namely cogs @xcite , cfq @xcite , and cognition @xcite , empirically verify that our proposal leads to better generalization, outperforming competitive baselines and more specialized techniques. ## disentanglement in a toy experiment we first shed light on the problem of entangled representations with a toy experiment and then move on to describe our modeling solution. for simplicity, we only focus on relations as the kind of semantic factors a model aims to represent, but the entanglement issue could also exist in representations of other factors, such as lexical meaning. data creation let @xmath0, r 1 , e c , r 2 , e 2 ] denote a sequence of symbols. we want to predict the relation between e 1 and e c , and e c and e 2 , which we denote by @xmath1, y 2 ), with y 1 ∈ l 1 and y 2 ∈ l 2 where l 1 are a set of relation labels for y 1 and l 2 are a set of relation labels for y 2 . for simplicity, we set e 1 , e c , and e 2 to the same symbol e (i.e., e 1 , e c , e 2 ∈ {e}) whereas r 1 ∈ r 1 and r 2 ∈ r 2 denote different relation symbols, and r 1 and r 2 are the corresponding sets of relation candidates. in this toy setting, we will further assume that different relation symbols determine different relation labels (e.g., for the phrases "cat in house" and "cat with house", "in" and "with" represent two distinct relations between "cat" and "house"). in reality, relations between words could be dependent on broader context or not verbalized at all. we also assume that there is a one-to-one mapping between relation symbols and relation labels (i.e., between l 1 and r 1 and l 2 and r 2 ). we construct a training set by including examples [e 1 , r 1 , e c , r 2 , e 2 ] where r 1 is the same relation symbol throughout while r 2 can be any relation symbol in r 2 (r 1 ∈ {r train }, r 2 ∈ r 2 ). we also include examples [e 1 , r 1 , e c ] with all relation symbols from r 1 occurring in isolation (r 1 ∈ r 1 ). this way, the training set covers all primitive relations, but contains only a particular type of relation composition (i.e., {r train } × r 2 ). in contrast, the test set contains all unseen compositions [e 1 , r 1 , e c , r 2 , e 2 ] (i.e., r 1 ∈ r 1 \{r train }, r 2 ∈ r 2 ) which will allow us to evaluate a model's ability to generalize. we set each relation set to include 10 relation symbols ( finally, we simplistically only consider the relations of target word e c with its left and right words e 1 and e 2 . in reality, a model would be expected to capture sentence-level semantics, i.e., a word's relation to all context words in a sentence (including no relation). modeling for each input symbol, we sample a vector from a gaussian distribution n (0, 0.2 2 i) and freeze it during training. we then embed each example x into a sequence of vectors [w 1 , w 2 , ..., w n ] (where @xmath2, h 2 , .., h n ] using a transformer encoder @xcite . to predict the relation between two symbols, we concatenate their corresponding representations and feed the resulting vector to an mlp for classification. to study how changes in relation y 1 affect the prediction of y 2 at test time, we explore two training methods. one is joint training where a model learns to predict both y 1 and y 2 (i.e., h 1 and h 3 are concatenated to predict y 1 or h 3 and h 5 are concatenated to predict y 2 ). the other method is separate training where a model is trained to only predict y 2 (i.e., only h 3 and h 5 are concatenated to predict y 2 ). for separate training, we basically ignore examples [e 1 , r 1 , e c ] which only include r 1 , as they have no bearing on the prediction of y 2 . observation with separate training, the model learns to ignore r 1 , the accuracy of predicting y 2 on the test set is 100%, regardless of which value r 1 takes. this indicates that random perturbation of r 1 alone does not lead to generalization failure. it also follows that there is no spurious correlation between r 1 and y 2 . however, when the model is trained to predict both relations (which is what happens in realistic settings since we need to capture all possible relations) r 1 has a huge impact on the prediction of y 2 whose accuracy drops to approximately 55%. taken together, these results suggest that the model fails to generalize to new relation compositions due to its internal representations being entangled and as a result changes in one relation affect the representation of others. why is there a wide performance gap between joint and separate training? at test time the model processes the same utterance (no matter whether it is trained jointly or separately), and could in theory be susceptible to both r 1 and r 2 . however, the induced representations show fundamentally different behaviors, and remain invariant to r 1 with separate training. a possible explanation is that modern neural networks trained with sgd have a learning bias towards simple functions @xcite . when r 1 is not predictive of y 2 , relying only on r 2 whilst remaining invariant to r 1 constitutes a simpler function than making use of both r 1 and r 2 . as a result, in separate training the model learns to ignore extraneous information, focusing exclusively on r 2 . on the contrary, in joint training the target of predicting both y 1 and y 2 forces the hidden states (e.g., h 3 ) to capture information about both relations, leading to the entanglement problem discussed above. a simple solution although separate training presents a solution to entanglement, it is unrealistic for real-wold data as it would be extremely inefficient to train separate models for each relation (the number of relations is quadratic with respect to sentence length). instead, we explore a simple but effective approach where a single model takes as input an utterance enriched with different indicator features for different targets. specifically, given utterance [e 1 , r 1 , e c , r 2 , e 2 ], and assuming we wish to predict relation y 1 , we add indicator feature 1 for symbols e 1 , r 1 , and e c (marking the relation and its immediate context), and 0 for all other symbols. the model then takes as input the utterance and relation indicators, i.e., [1, 1, 1, 0, 0] for y 1 and [0, 0, 1, 1, 1] for y 2 , and learns embeddings for indicators during training. it thus learns specialized representations for each prediction rather than shared representations for all predictions. based on the simplicity bias, the two representations will guide the model towards exclusively relying on r 1 and r 2 , naturally disentangling different relations by encoding them separately. such a model predicts y 1 with 100% test accuracy and y 2 with 97%. @xcite have argued that failure to capture systematicity is a major deficiency of neural architectures, contrasting human learners who can readily apply known grammatical rules to arbitrary novel word combinations to individually memorizing an exponential number of sentences. however, our toy experiment shows that neural networks are not just memorizing sentences but implicitly capturing structure. with separate training or joint training enhanced with the marking strategy, the neural model manages to remain robust to interference from r 1 and properly represent r 2 even for unseen examples, i.e., new compositions of r 1 and r 2 . this generalization ability implies that neural models do not need to see all exponential compositions in order to produce plausible representations of them. instead, with appropriate training and model design, they could uncover and represent the structure underlying systematically related sentences. ## learning to disentangle while the marking strategy offers substantial benefits in learning disentangled relation representations, we typically do not have access to explicit labels indicating which words are helpful for predicting a specific relation. nevertheless, the idea of learning representations specialized for different predictions (albeit with shared parameters) is general and could potentially alleviate the entanglement problem for compositional generalization. let [x 1 , x 2 , ..., x n ] denote a source sequence. canonical seq2seq models like the transformer @xcite first encode it into a sequence of contextualized representations which are then used to decode target symbols [y 1 , y 2 , ..., y m ] one by one. the same source encodings are used to predict all target symbols, and are therefore expected to capture all semantic factors in the input. however, these could be entangled as demonstrated in our analysis above. to alleviate this issue, we propose to learn specialized source representations for different predictions by adaptively re-encoding the source input at every step of the decoding. specifically, at the t-th time step, we concatenate the source input with the previously decoded target and obtain the context for the current prediction c @xmath3, x 2 , ..., x n , y 1 , ..., y t-1 , [ph]] where [ph] is a placeholder (e.g., a mask token when using a pretrained encoder). c t is then fed to a standard encoder (e.g., the transformer encoder) to obtain the contextualized representations h @xmath4,1 , h t,2 , ..., h t,n , h t,n+1 , ..., h t,n+t ]: the key difference from the encoder in standard seq2seq models is that at each time step we adaptively re-compute source encodings h t,@xmath5,1 , ..., h t,n ] that condition on the newly decoded target [y 1 , ..., y t-1 ]. this way, target context informs the encoder of predictions of interest at each time step. this simple modification unburdens the model from capturing all source information through a forward pass of encoding. instead, based on the simplicity bias, the model tends to zero in on information relevant for the current prediction, remaining invariant to irrelevant details, thereby improving disentanglement. one might argue that the decoder in standard seq2seq models could also extract specialized information for each prediction (through the cross attention mechanism). however, it would fail to do so when working with an entangled encoder that produces problematic representations for out-of-distribution examples and breaks down the decoding process. we propose two strategies for exploiting the target-informed encoder. firstly, we use a multilayer perceptron (mlp) to predict y t based on the encoder's output, i.e., the last hidden states h t,n+t : secondly, we incorporate the proposed encoder into the standard encoder-decoder architecture: we take source encodings h t,n and feed them together with the previous target [y 1 , ..., y t-1 ] to a standard decoder (e.g., transformer-based) to predict y t : for complex tasks like machine translation, preserving the encoder-decoder architecture is essential to achieving good performance. we adopt the transformer architecture to instantiate the encoder and decoder, however, the proposed method is generally applicable to any seq2seq model. we maintain separate position encodings for source and target symbols (e.g., x 1 and y 1 correspond to the same position). to differentiate between source and target content, we also add a source(target) type embedding to all source(target) token embeddings. compared to the classical transformer, our proposal increases running time from o(n 2 + m 2 ) to o(m(n 2 + m 2 )) where n is input length and m is output length. improving the efficiency of our approach is deferred to future work. ## experiments: semantic parsing in this section, we present our experiments for evaluating the proposed disentangled seq2seq model which we call dangle. we refer to the two variants of dangle as dangle-enc and dangle-encdec. we first focus on semantic parsing benchmarks which target compositional generalization. our second suite of experiments reports results on compositional generalization for machine translation. ## related work the realization that neural sequence models struggle in settings requiring compositional generalization has led to numerous research efforts aiming to understand why this happens and how to prevent it. one line of research tries to improve compositional generalization by adopting a more conventional grammar-based approach @xcite , incorporating a lexicon or lexiconstyle alignments into sequence models @xcite , and augmenting the standard training objective with attention supervision losses @xcite . other work resorts to data augmentation strategies as a way of injecting a compositional inductive bias into neural models @xcite @xcite and meta-learning to directly optimize for out-ofdistribution generalization @xcite . there are also several approaches which explore the benefits of large-scale pre-trained language models @xcite . in this work we identify the learning of representations which are not disentangled as one of the reasons why neural sequence models fail to generalize compositionally. disentanglement, i.e., the ability to uncover explanatory factors from data, is often cited as a key property of good representations @xcite . for example, a model trained on 3d objects might learn factors such as object identity, position, scale, lighting, or colour. several types of variational autoencoders (kingma and welling, 2014) have been proposed for the unsupervised learning of disentangled representations in images @xcite @xcite . however, some of the underlying assumptions of these models have come under scrutiny recently @xcite . disentanglement for linguistic representations remains under-explored, and has mostly focused on separating the style of text from its content @xcite . in the context of sentence-level semantics, disentangled representations should be able to discriminate among lexical meanings and semantic relations between words. we highlight the entanglement problem in neural sequence models when trained with explicit factor supervision which, however, does not cover the entire exponential space of compositions for different factors. instead of encouraging disentanglement with some regularization @xcite , we propose a modification to sequence-to-sequence models which achieves this by re-encoding the source based on newly decoded target context. it may be counterintuitive that we are disentangling by conditioning on more information, but it is feasible thanks to the inherent simplicity bias in neural models. ## conclusions in this paper we proposed an extension to sequenceto-sequence models which allows us to learn disentangled representations for compositional generalization. we have argued that taking into account the target context makes it easier for the encoder to exploit specialized information for improving its predictions. experiments on semantic parsing and machine translation have shown that our proposal improves compositional generalization without any model, dataset, or task specific modification.