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 presents a type of “task re-formulation” that re-arranges the labels in task such that they are easier for transformer-based encoder-decoder models. This technique is specifically meant to address compositional generalization, which roughly means to generalize to instances longer than seen during training, with combinations of tokens not seen during training, and so on (further clarified in the paper). The paper presents interesting results on three tasks (PCFG, Cartesian Product, and CFQ). I find the problem of compositional generalization with transformers to be very important for the ML community, and that the results are mostly useful. Nonetheless, there is a substantial amount of missing related work, and the writing could greatly be improved.