Iterative Decoding for Compositional Generalization in Transformers

1 INTRODUCTION . Deep learning architectures achieve state-of-the-art results in a wide array of machine learning problems , where their impressive performance is attributed to their ability to generalize ( Goodfellow et al. , 2016 ; LeCun et al. , 2015 ) . However , this ability is typically limited to generalization under the statistical learning paradigm , i.e. , in-distribution generalization , and does not encompass generalizing compositionally . Compositional generalization is the ability of a model to combine a set of learned primitives to execute more complex tasks . For instance , for a ground robot whose motion planner has learned to execute the instructions “ walk ” , “ jump ” , and “ jump right ” , generalizing compositionally would be to be able to execute the instruction “ walk right ” ( Lake & Baroni , 2018 ) . In machine learning models , compositional generalization is desirable for two reasons . First , because it is a crucial aspect of intelligence observed in both humans and classical artificial intelligence techniques . In humans , a prevailing example is the way children can solve complicated mathematical expressions after being taught basic arithmetics . Second , because it can increase a model ’ s data efficiency . By endowing models with the ability to extrapolate to unseen examples that are more complex , or complex in different ways than seen at training , compositionality acts as an implicit mechanism for data augmentation . In this paper , our goal is to increase compositional generalization in transformers , with particular focus on natural language-like tasks—where compositionality is key . Understanding that for models to be able to execute composite tasks they need to be taught how to compose , we introduce iterative decoding , an alternative to sequence-to-sequence ( seq2seq ) learning that decomposes the process of mapping the inputs to the outputs of each example into a sequence of intermediate steps which the transformer learns to perform iteratively . During training , each input-output pair is converted into a sequence of “ intermediate input-intermediate output ” pairs , which are task specific . During prediction , the intermediate outputs predicted by the transformer are adapted into the subsequent intermediate inputs , which are fed back to the transformer until an end-of-iteration token is produced . Our main contributions are ( i ) showing that iterative decoding , especially when combined with other architectural modifications such as relative attention ( Shaw et al. , 2018 ) and copy decoders ( Ontañón et al. , 2021 ) , largely improves compositional generalization in transformers , and ( ii ) evidencing that , in general , seq2seq transformers can not learn iterations unless they are unrolled . We show this through numerical experiments on two compositionally hard splits of PCFG ( Hupkes et al. , 2020 ) , a string editting dataset ; and on a cartesian product dataset , where the goal is to generalize to longer input vectors than those on which the model is trained . We also present numerical results on CFQ ( Keysers et al. , 2019 ) , a semantic parsing dataset consisting of natural language questions paired with SPARQL queries . Importantly , the results obtained in this dataset evidence a limitation of iterative decoding , which is that it can be sensitive to the ordering of the intermediate steps . The rest of this paper is organized as follows . Section 2 discusses related work on compositional generalization , transformers and on the PCFG , cartesian product , and CFQ datasets . Section 3 introduces iterative decoding . Iterative decoding strategies specific to each dataset and numerical results are presented and discussed in Section 4 . Section 5 presents future research directions and concluding remarks . 2 BACKGROUND . In this section , we introduce some background and related work on compositional generalization , transformer architectures and on the PCFG , Cartesian product and CFQ datasets , used in our experiments . 2.1 COMPOSITIONAL GENERALIZATION . Compositional generalization ( or compositionality ) refers to the ability of a model that has learned to perform a set of basic operations—primitives—to generalize to more complex operations , i.e. , operations consisting of compositions of the learned primitives ( Lake & Baroni , 2018 ) . Examples of operations requiring compositionality are shown in Figure 1 for three datasets . For instance , the top-left corner shows an example from the PCFG ( Hupkes et al. , 2020 ) dataset . In some versions of this dataset , the model is trained to solve several atomic string editing operations ( such as copy and swap first last ) , and how to compose them . During testing , the model is tested by asking it to compose a longer number of operations than seen during training . Hence , the model has to be able to generalize compositionally . This string editing example can be seen as an instance of productivity , one of the five types of compositional generalization identified by Hupkes et al . ( 2020 ) which involves generalizing to longer examples than those seen during training . Another type of compositional generalization is systematicity , the ability to recombine known parts and rules in ways different than those seen during training . Early works on compositionality have explored the limitations of different machine learning models in generalizing compositionally . Liška et al . ( 2018 ) showed that , while it is theoretically possible for a recurrent neural network ( RNN ) to generalize in this way , only a small fraction of the models they trained behaved compositionally . Lake & Baroni ( 2018 ) proposed SCAN , a dataset consisting of navigation commands to be mapped to action sequences , and observed that while RNNs trained on it only generalized well when the differences between the training and test sets where small , they failed when more systematic compositional skills were required . Other datasets created with the purpose of measuring compositionality include PCFG ( Hupkes et al. , 2020 ) and CFQ ( Keysers et al. , 2019 ) , where both long-short term memory ( LSTM ) and transformer-based architectures have been shown to perform poorly . More recently , a popular research direction is to try to endow these machine learning models with a “ compositional generalization bias ” . Kim et al . ( 2021 ) saw benefits in converting CFQ into a classification task and using structural annotations ( e.g. , entity links ) as attention masks in transformers . Ontañón et al . ( 2021 ) were able to improve transformer compositional generalization on a variety of compositionally hard datasets by making architectural modifications such as relative attention , copy decoders , and weight sharing . Taking a similar approach , Csordás et al . ( 2021 ) observed performance improvements from relative position encodings and scaled embeddings . Other strategies to improve compositional generalization include increased pretraining ( Furrer et al. , 2020 ) , data augmentation ( Andreas , 2019 ) and differentiable neurocomputers ( Graves et al. , 2016 ) . Closely related to our work , PonderNet ( Banino et al. , 2021 ) trains a model that iterates internally to achieve a better compromise between training accuracy and generalization . It predicts both an output and a halting probability at each step , operating recurrently . Iterative decoding also operates recurrently , but two important differences are that ( i ) in iterative decoding the intermediate steps are supervised and interpretable , and ( ii ) rather than predicting a halting probability at each step , iterative decoding trains the model to produce a special token to indicate end of iteration . 2.2 TRANSFORMER MODEL . In this paper we focus on transformer models . Despite struggling to generalize compositionally , transformer-based architectures such as BERT ( Devlin et al. , 2018 ) and T5 ( Raffel et al. , 2019 ) were popularized by their remarkable performance in machine translation ( Zhu et al. , 2020 ) , question answering ( Ainslie et al. , 2020 ) , summarization ( Zhang et al. , 2019 ) and other natural language processing ( NLP ) tasks . Introduced by Vaswani et al . ( 2017 ) , the basic transformer model is composed of an encoder and a decoder . The encoder is made up of layers consisting of a self-attention sublayer and a feedforward sublayer . The decoder has the same structure , but with an additional attention sublayer to compute the decoder-to-encoder attention . The input to the transformer is a sequence of token embeddings . Since these embeddings do not carry information about the position of each token in the sequence , a position encoding is typically added to the input embeddings . These are then fed to the encoder , which encodes all tokens at once and forwards the result to the decoder . From the encoded input and the decoded output tokens generated so far , the decoder generates the distribution of the next output token , one token at a time . We experiment with two extensions to the original transformer architecture that were previously shown to improve compositional generalization ( Ontañón et al. , 2021 ) : relative position encodings and copy decoding . To each pair of tokens in the input , relative position encodings assign a label that is equal to the minimum between their relative distance and a fixed relative attention radius . An important characteristic of relative position encodings is that they are position invariant , which means that two tokens that are k positions apart will attend to each other in the same way regardless of their absolute positions in the sequence . Copy decoding involves adding a learnable parameter that allows to switch between the decoder and a copy decoder which produces an independent embedding that can be interpreted as a “ copy ” from the input sequence . 2.3 DATASETS . We consider three datasets—PCFG , cartesian product , and CFQ—which are discussed in detail in the following sections . 2.3.1 PCFG . PCFG is an artificial translation dataset proposed by Hupkes et al . ( 2020 ) and generated by a probabilistic context free grammar . The input data consists of string editing instructions with four types of tokens : unary operation tokens ( e.g. , reverse ) , binary operation tokens ( e.g. , append ) , a string separation token “ , ” ( to separate arguments of binary operations ) , and string elements ( e.g. , B10 , D2 ) . The output data consists of the strings resulting from the application of the operations ; see the top left corner of Figure 1 for an example . There are six training-test splits of the PCFG dataset . The first is a random split which we use as a baseline . The other five are compositionally hard splits used to measure the five different types compositional generalization . We focus on two of them : productivity and systematicity . In the productivity split , the training samples have up to 8 string operations , while the test samples have 9 or more . In the systematicity split , there is no restriction on the number of operations in either set , but in the test set they are combined in different ways than in the training set . The main challenge of the PCFG dataset is that it requires learning ten string editing operations , some of which are very similar . The unary operation echo , for instance , only differs from copy by repeating the last element of the string . While transformers generally achieve good performance in the random split of the PCFG dataset , the productivity and systematicity splits are harder because transformers tend to learn shortcuts ( instead of learning the mechanics of each operation ) . In particular , a key difficulty of the productivity split is that the model needs to learn to do ” recursion ” and apply an arbitrary number of operations when input examples grow in length . In some cases , the strings to modify can also be very long , which places an additional capacity burden on transformers by requiring them to learn how to copy strings .

Summary: This paper studies compositional generalization in the context of Transformer models. Noting that "one-shot" sequence-to-sequence learning with Transformer models often results in models that do not generalize compositionally, the authors propose an iterative scheme whereby intermediate input-output pairs are constructed for each task (how this is constructed varies depends on the task), and the decoding proceeds as a series of seq2seq tasks. The approach is found to perform well on synthetic tasks (PCFG, Cartesian product) where there is a "natural" way to construct the intermediate input-output pairs, but fails to perform well on more real-world tasks such as CFQ.