Learning to Recombine and Resample Data For Compositional Generalization

1 INTRODUCTION How can we build machine learning models with the ability to learn new concepts in context from little data ? Human language learners acquire new word meanings from a single exposure ( Carey & Bartlett , 1978 ) , and immediately incorporate words and their meanings productively and compositionally into larger linguistic and conceptual systems ( Berko , 1958 ; Piantadosi & Aslin , 2016 ) . Despite the remarkable success of neural network models on many learning problems in recent years—including one-shot learning of classifiers and policies ( Santoro et al. , 2016 ; Wang et al. , 2016 ) —this kind of few-shot learning of composable concepts remains beyond the reach of standard neural models in both diagnostic and naturalistic settings ( Lake & Baroni , 2018 ; Bahdanau et al. , 2019a ) . Consider the few-shot morphology learning problem shown in Fig . 1 , in which a learner must predict various linguistic features ( e.g . 3rd person , SinGular , PRESent tense ) from word forms , with only a small number of examples of the PAST tense in the training set . Neural sequenceto-sequence models ( e.g . Bahdanau et al. , 2015 ) trained on this kind of imbalanced data fail to predict past-tense tags on held-out inputs of any kind ( Section 5 ) . Previous attempts to address this and related shortcomings in neural models have focused on explicitly encouraging rule-like behavior by e.g . modeling data with symbolic grammars ( Jia & Liang , 2016 ; Xiao et al. , 2016 ; Cai et al. , 2017 ) or applying rule-based data augmentation ( Andreas , 2020 ) . These procedures involve highly task-specific models or generative assumptions , preventing them from generalizing effectively to less structured problems that combine rule-like and exceptional behavior . More fundamentally , they fail to answer the question of whether explicit rules are necessary for compositional inductive bias , and whether it is possible to obtain “ rule-like ” inductive bias without appeal to an underlying symbolic generative process . This paper describes a procedure for improving few-shot compositional generalization in neural sequence models without symbolic scaffolding . Our key insight is that even fixed , imbalanced training datasets provide a rich source of supervision for few-shot learning of concepts and composition rules . In particular , we propose a new class of prototype-based neural sequence models ( c.f . Gu et al. , 2018 ) that can be directly trained to perform the kinds of generalization exhibited in Fig . 1 by explicitly recombining fragments of training examples to reconstruct other examples . Even when these prototype-based models are not effective as general-purpose predictors , we can resample their outputs to select high-quality synthetic examples of rare phenomena . Ordinary neural sequence models may then be trained on datasets augmented with these synthetic examples , distilling the learned regularities into more flexible predictors . This procedure , which we abbreviate R & R , promotes efficient generalization in both challenging synthetic sequence modeling tasks ( Lake & Baroni , 2018 ) and morphological analysis in multiple natural languages ( Cotterell et al. , 2018 ) . By directly optimizing for the kinds of generalization that symbolic representations are supposed to support , we can bypass the need for symbolic representations themselves : R & R gives performance comparable to or better than state-of-the-art neuro-symbolic approaches on tests of compositional generalization . Our results suggest that some failures of systematicity in neural models can be explained by simpler structural constraints on data distributions and corrected with weaker inductive bias than previously described.1 2 BACKGROUND AND RELATED WORK . Compositional generalization Systematic compositionality—the capacity to identify rule-like regularities from limited data and generalize these rules to novel situations—is an essential feature of human reasoning ( Fodor et al. , 1988 ) . While details vary , a common feature of existing attempts to formalize systematicity in sequence modeling problems ( e.g . Gordon et al. , 2020 ) is the intuition that learners should make accurate predictions in situations featuring novel combinations of previously observed input or output subsequences . For example , learners should generalize from actions seen in isolation to more complex commands involving those actions ( Lake et al. , 2019 ) , and from relations of the form r ( a , b ) to r ( b , a ) ( Keysers et al. , 2020 ; Bahdanau et al. , 2019b ) . In machine learning , previous studies have found that standard neural architectures fail to generalize systematically even when they achieve high in-distribution accuracy in a variety of settings ( Lake & Baroni , 2018 ; Bastings et al. , 2018 ; Johnson et al. , 2017 ) . Data augmentation and resampling Learning to predict sequential outputs with rare or novel subsequences is related to the widely studied problem of class imbalance in classification problems . There , undersampling of the majority class or oversampling of the minority class has been found to improve the quality of predictions for rare phenomena ( Japkowicz et al. , 2000 ) . This can be combined with targeted data augmentation with synthetic examples of the minority class ( Chawla et al. , 2002 ) . Generically , given a training dataset D , learning with class resampling and data augmentation involves defining an augmentation distribution p̃ ( x , y | D ) and sample weighting function u ( x , y ) and maximizing a training objective of the form : L ( θ ) = 1|D| ∑ x∈D log pθ ( y | x ) ︸ ︷︷ ︸ Original training data + E ( x , y ) ∼p̃ u ( x , y ) log pθ ( y | x ) ︸ ︷︷ ︸ Augmented data . ( 1 ) In addition to task-specific model architectures ( Andreas et al. , 2016 ; Russin et al. , 2019 ) , recent years have seen a renewed interest in data augmentation as a flexible and model-agnostic tool for encouraging controlled generalization ( Ratner et al. , 2017 ) . Existing proposals for sequence models are mainly rule-based—in sequence modeling problems , specifying a synchronous context-free grammar ( Jia & Liang , 2016 ) or string rewriting system ( Andreas , 2020 ) to generate new examples . Rule-based data augmentation schemes that recombine multiple training examples have been proposed 1Code for all experiments in this paper is available at https : //github.com/ekinakyurek/compgen . We implemented our experiments in Knet ( Yuret , 2016 ) using Julia ( Bezanson et al. , 2017 ) . for image classification ( Inoue , 2018 ) and machine translation ( Fadaee et al. , 2017 ) . While rulebased data augmentation is highly effective in structured problems featuring crisp correspondences between inputs and outputs , the effectiveness of such approaches involving more complicated , context-dependent relationships between inputs and outputs has not been well-studied . Learned data augmentation What might compositional data augmentation look like without rules as a source of inductive bias ? As Fig . 1 suggests , an ideal data augmentation procedure ( p̃ in Eq . 1 ) should automatically identify valid ways of transforming and combining examples , without pre-committing to a fixed set of transformations.2 A promising starting point is provided by prototype-based models , a number of which ( Gu et al. , 2018 ; Guu et al. , 2018 ; Khandelwal et al. , 2020 ) have been recently proposed for sequence modeling . Such models generate data according to : d ∼ prewrite ( · | d′ ; θ ) where d′ ∼ Unif ( D ) ; ( 2 ) for a datasetD and a learned sequence rewriting model prewrite ( d | d′ ; θ ) . ( To avoid confusion , we will use the symbol d to denote a datum . Because a data augmentation procedure must produce complete input–output examples , each d is an ( x , y ) pair for the conditional tasks evaluated in this paper . ) While recent variants implement prewrite with neural networks , these models are closely related to classical kernel density estimators ( Rosenblatt , 1956 ) . But additionally—building on the motivation in Section 1—they may be viewed as one-shot learners trained to generate new data d from a single example . Existing work uses prototype-based models as replacements for standard sequence models . We will show here that they are even better suited to use as data augmentation procedures : they can produce high-precision examples in the neighborhood of existing training data , then be used to bootstrap simpler predictors that extrapolate more effectively . But our experiments will also show that existing prototype-based models give mixed results on challenging generalizations of the kind depicted in Fig . 1 when used for either direct prediction or data augmentation—performing well in some settings but barely above baseline in others . Accordingly , R & R is built on two model components that transform prototype-based language models into an effective learned data augmentation scheme . Section 3 describes an implementation of prewrite that encourages greater sample diversity and well-formedness via a multi-prototype copying mechanism ( a two-shot learner ) . Section 4 describes heuristics for sampling prototypes d′ and model outputs d to focus data augmentation on the most informative examples . Section 5 investigates the empirical performance of both components of the approach , finding that they together provide they a simple but surprisingly effective tool for enabling compositional generalization . 3 PROTOTYPE-BASED SEQUENCE MODELS FOR DATA RECOMBINATION . We begin with a brief review of existing prototype-based sequence models . Our presentation mostly follows the retrieve-and-edit approach of Guu et al . ( 2018 ) , but versions of the approach in this paper could also be built on retrieval-based models implemented with memory networks ( Miller et al. , 2016 ; Gu et al. , 2018 ) or transformers ( Khandelwal et al. , 2020 ; Guu et al. , 2020 ) . The generative process described in Eq . 2 implies a marginal sequence probability : p ( d ) = 1 |D| ∑ d′∈D prewrite ( d | d′ ; θ ) ( 3 ) Maximizing this quantity over the training set with respect to θ will encourage prewrite to act as a model of valid data transformations : To be assigned high probability , every training example must be explained by at least one other example and a parametric rewriting operation . ( The trivial solution where pθ is the identity function , with pθ ( d | d′ = d ) = 1 , can be ruled out manually in 2As a concrete example of the potential advantage of learned data augmentation , consider applying the GECA procedure of Andreas ( 2020 ) to the language of strings anbn . GECA produces a training set that is substitutable in the sense of Clark & Eyraud ( 2007 ) ; as noted there , anbn is not substitutable . GECA will infer that a can be replaced with aab based on their common context in ( aabb , aaabbb ) , then generate the malformed example aaababbb by replacing an a in the wrong position . In contrast , recurrent neural networks can accurately model anbn ( Weiss et al. , 2018 ; Gers & Schmidhuber , 2001 ) . Of course , this language can also be generated using even more constrained procedures than GECA , but in general learned sequence models can capture a broader set of both formal regularities and exceptions compared to rule-based procedures . the design of pθ . ) When D is large , the sum in Eq . 3 is too large to enumerate exhaustively when computing the marginal likelihood . Instead , we can optimize a lower bound by restricting the sum to a neighborhood N ( d ) ⊂ D of training examples around each d : p ( d ) ≥ 1 |D| ∑ d′∈N ( d ) prewrite ( d | d′ ; θ ) . ( 4 ) The choice of N is discussed in more detail in Section 4 . Now observe that : log p ( d ) ≥ log ( |N ( d ) | ∑ d′∈N ( d ) 1 |N ( d ) | prewrite ( d | d′ ; θ ) ) − log |D| ( 5 ) ≥ 1 |N ( d ) | ∑ d′∈N ( d ) log prewrite ( d | d′ ; θ ) + log ( |N ( d ) | |D| ) ( 6 ) where the second step uses Jensen ’ s inequality . If all |N ( d ) | are the same size , maximizing this lower bound on log-likelihood is equivalent to simply maximizing∑ d′∈N ( d ) log prewrite ( d | d′ ; θ ) ( 7 ) over D—this is the ordinary conditional likelihood for a string transducer ( Ristad & Yianilos , 1998 ) or sequence-to-sequence model ( Sutskever et al. , 2014 ) with examples d , d′ ∈ N ( d ) .3 We have motivated prototype-based models by arguing that prewrite learns a model of transformations licensed by the training data . However , when generalization involves complex compositions , we will show that neither a basic RNN implementation of prewrite or a single prototype is enough ; we must provide the learned rewriting model with a larger inventory of parts and encourage reuse of those parts as faithfully as possible . This motivates the two improvements on the prototype-based modeling framework described in the remainder of this section : generalization to multiple prototypes ( Section 3.1 ) and a new rewriting model ( Section 3.2 ) . 3.1 n-PROTOTYPE MODELS To improve compositionality in prototype-based models , we equip them with the ability to condition on multiple examples simultaneously . We extend the basic prototype-based language model to n prototypes , which we now refer to as a recombination model precomb : d ∼ precomb ( · | d′1 : n ; θ ) where d′1 : n def = ( d′1 , d ′ 2 , . . . , d ′ n ) ∼ pΩ ( · ) ( 9 ) A multi-protype model may be viewed as a meta-learner ( Thrun & Pratt , 1998 ; Santoro et al. , 2016 ) : it maps from a small number of examples ( the prototypes ) to a distribution over new datapoints consistent with those examples . By choosing the neighborhood and implementation of precomb appropriately , we can train this meta-learner to specialize in one-shot concept learning ( by reusing a fragment exhibited in a single prototype ) or compositional generalization ( by assembling fragments of prototypes into a novel configuration ) . To enable this behavior , we define a set of compatible prototypes Ω ⊂ Dn ( Section 4 ) and let pΩ def = Unif ( Ω ) . We update Eq . 6 to feature a corresponding multi-prototype neighborhood N : D → Ω . The only terms that have changed are the conditioning variable and the constant term , and it is again sufficient to choose θ to optimize ∑ d′1 : n∈N ( d ) log precomb ( d | d′1 : n ) over D , implementing precomb as described next .

Summary: To tackle situations where compositionality is mostly required at inference time, the paper proposes a novel data augmentation method with an RNN based generator (recombination); to make the generator generate highly compositional patterns, the paper proposes a resampling method. The methods have been tested on two benchmarks focusing on the issue, SCAN and morphological analysis. The system performs on par with recently proposed GECA for SCAN and favorably to GECA on morphological analysis.