A UTOMATICALLY C OMPOSING R EPRESENTATION T RANSFORMATIONS AS A M EANS FOR G ENERALIZATION
A BSTRACT
A generally intelligent learner should generalize to more complex tasks than it has previously encountered, but the two common paradigms in machine learning – either training a separate learner per task or training a single learner for all tasks – both have difﬁculty with such generalization because they do not leverage the compositional structure of the task distribution.
This paper introduces the compositional problem graph as a broadly applicable formalism to relate tasks of different complexity in terms of problems with shared subproblems.
We propose the compositional generalization problem for measuring how readily old knowledge can be reused and hence built upon.
As a ﬁrst step for tackling compositional generalization, we introduce the compositional recursive learner , a domaingeneral framework for learning algorithmic procedures for composing representation transformations, producing a learner that reasons about what computation to execute by making analogies to previously seen problems.
We show on a symbolic and a high-dimensional domain that our compositional approach can generalize to more complex problems than the learner has previously encountered, whereas baselines that are not explicitly compositional do not
I NTRODUCTION
This paper seeks to tackle the question of how to build machines that leverage prior experience to solve more complex problems than they have previously encountered.
How does a learner represent prior experience?
How does a learner apply what it has learned to solve new problems?
Motivated by these questions, this paper aims to formalize the idea of, as well as to develop an understanding of the machinery for, compositional generalization in problems that exhibit compositional structure.
The solutions for such problems can be found by composing in sequence a small set of reusable partial solutions, each of which tackles a subproblem of a larger problem.
The central contributions of this paper are to frame the shared structure across multiple tasks in terms of a compositional problem graph , propose compositional generalization as an evaluation scheme to test the degree a learner can apply previously learned knowledge to solve new problems, and introduce the compositional recursive learner , a domain-general framework 1 for sequentially composing representation transformations that each solve a subproblem of a larger problem.
The key to our approach is recasting the problem of generalization as a problem of learning algorithmic procedures over representation transformations.
A solution to a (sub)problem is a transformation between its input and output representations, and a solution to a larger problem composes these subsolutions together.
Therefore, representing and leveraging prior problem-solving experience amounts to learning a set of reusable primitive transformations and their means of composition that reﬂect the structural properties of the problem distribution.
This paper introduces the compositional recursive learner (CRL), a framework for learning both these transformations and their composition together with sparse supervision, taking a step beyond other approaches that have assumed either pre-speciﬁed transformation or composition rules (Sec. 5).
CRL learns a modular recursive program that iteratively re-represents the input representation into more familiar representations it knows how to compute with.
In this framework, a transformation between representations is encapsulated into a computational module , and the overall program is the sequential combination of the inputs and outputs of these modules, whose application are decided by a controller .
What sort of training scheme would encourage the spontaneous specialization of the modules around the compositional structure of the problem distribution?
First, exposing the learner to a diverse distribution of compositional problems would help it pattern-match across problems to distill out common functionality that it can capture in its modules for future use.
Second, constraining the model class of these modules guides them to capture the structure of the problem distribution; without such constraints the learner can easily overﬁt to the empirical training distribution and fail to generalize to harder problems in a combinatorially large problem space.
Third, training the learner with a curriculum encourages the learner to build off old solutions to solve new problems by re-representing the new problem into one it knows how to solve, rather than learning from scratch.
How should the learner learn to use these modules to exploit the compositional structure of the problem distribution?
The application of a sequence of modules can be likened to the execution trace of the program that CRL automatically constructs, where a computation is the application of a module to the output of a previous computation.
The automatic construction of the program can be formulated as the solution to a sequential decision-making problem in a meta-level Markov decision process (MDP)
(Hay et al., 2014), where the state space is the learner’s internal states of computation and the action space is the set of modules.
Under this formulation, the controller can implement loops and recursion, and it can also decide on which part of the current state of computation to apply a module.
Our experiments on solving multilingual arithmetic problems and recognizing spatially transformed
MNIST digits show that the above proposed training scheme prescribes a type of reformulation : re-representing a new problem in terms of other problems by implicitly making an analogy between their solutions.
We also show that our meta-reasoning approach for deciding what modules to execute achieves better generalization to more complex problems than monolithic learners that are not explicitly compositional.
C OMPOSITIONAL G ENERALIZATION
Solving a problem simply means representing it so as to make the solution transparent.
(S IMON , 1988)
Humans navigate foreign cities and understand novel conversations despite only observing a tiny fraction of the true distribution of the world.
Perhaps they can extrapolate in this way because the world contains compositional structure, such that solving a novel problem is possible by composing previously learned partial solutions in a novel way to ﬁt the context.
With this perspective, we propose the concept of compositional generalization , adding to the standard learning theory paradigm additional assumptions that are tailored for studying generalization in the compositional context.
The key assumption of compositional generalization is that harder problems are made of easier problems.
The problems from the training and test sets share the same primitive subproblems, but differ in the manner and complexity with which these subproblems are combined.
Therefore, problems in the test set can be solved by combining solutions learned from the training set in novel ways.
Deﬁnition.
Let a problem P be a pair ( X in , X out ) , where X in and X out are random variables that respectively correspond to the input and output representations of the problem.
Let the distribution of X in be r in and the distribution of X out be r out .
To solve a particular problem P = p is to transform X in = x in into X out = x out .
A composite problem p a = p b ◦ p c is that for which it is possible to solve by ﬁrst solving p c and then solving p b with the output of p c as input.
p b and p c are subproblems with respect to p a .
The space of compositional problems form a compositional problem graph , whose nodes are the representation distributions r .
A problem is described as pair of nodes between which the learner must learn to construct an edge or a path to transform between the two representations.
Characteristics.
First, there are many ways in which a problem can be solved.
For example, translating an English expression to a Spanish one can be solved directly by learning such a transformation, or a learner could make an analogy with other problems by ﬁrst translating English to French, and then French to Spanish as intermediate subproblems.
Second, sometimes a useful (although not only) way to solve a problem is indicated by the recursive structure of the problem itself: solving the arithmetic expression 3 + 4 × 7 modulo 10 can be decomposed by ﬁrst solving the subproblem 4 × 7 = 8 and then 3+8 = 1 .
Third, because a problem is just an (input, output) pair, standard problems in machine learning ﬁt into this broadly applicable framework.
For example, for a supervised classiﬁcation problem, the input representation can be an image and the output representation a label, and intermediate subproblems can be transforming some intermediate representations to other intermediate representations.
Sec. 4 demonstrates CRL on all three of the above examples.
Broad Applicability.
Problems in supervised, unsupervised, and reinforcement learning can all be viewed under the framework of transformations between representations.
What we gain from the compositional problem graph perspective is a methodological way to relate together different problems of various forms and complexity, which is especially useful in a lifelong learning setting: the knowledge required to solve one problem is composed of the knowledge required to solve subproblems seen in the past in the context of different problems.
For example, we can view latent variable reinforcement learning architectures such as (Ha & Schmidhuber, 2018; Nair et al., 2018) as simultaneously solving an image reconstruction problem and an action prediction problem, both of which share the same subproblem of transforming a visual observation into a latent representation.
Lifelong learning, then, can be formulated as continuing to make more connections between nodes in the compositional problem graph, gradually expanding the frontier of nodes explored.
Sec. 4 describes how CRL takes advantage of this compositional formulation in a multi-task zero-shot generalization setup to solve new problems by re-using computations learned from solving past problems.
Evaluation.
To evaluate a learner’s capacity for compositional generalization, we introduce two challenges.
The ﬁrst is to generalize to problems with different subproblem combinations from what the learner has seen.
The second is to generalize to problems with longer subproblems combinations than the learner has seen.
Evaluating a learner’s capability for compositional generalization is one way to measure how readily old knowledge can be reused and hence built upon.
A L EARNER T HAT P ROGRAMS I TSELF
This paper departs from the popular representation-centric view of knowledge (Bengio et al., 2013) and instead adopts a computation-centric view of knowledge: our goal is to encapsulate useful functionality shared across tasks into specialized computational modules – atomic function operators that perform transformations between representations.
This section introduces the compositional recursive learner (CRL), a framework for training modules to capture primitive subproblems and for composing together these modules as subproblem solutions to form a path between nodes of the compositional problem graph.
3. C OMPOSITIONAL R ECURSIVE L EARNER
The CRL framework consists of a controller π , a set of modules m ∈ M , and an evaluator E .
Training CRL on a diverse compositional problem distribution produces a modular recursive program that is trained to transform the input X in into its output X out , the corresponding samples of which are drawn from pairs of nodes in the compositional problem graph.
In this program, the controller looks at the current state x
i of the program and chooses a module m to apply to the state.
The evaluator executes the module on that state to produce the next state x i +1 of the program.
X in is the initial state of the program, ˆ X out is the last, and the intermediate states
X i of the execution trace correspond to the other representations produced and consumed by the modules.
The controller can choose to re-use modules across different program executions to solve different problems, making it straightforward to re-use computation learned from solving other problems to solve the current one.
The controller can also choose to reuse modules several times within the same program execution, which produces recursive behavior.
3. D ECIDING W HICH C OMPUTATIONS T O E XECUTE
The sequential decision problem that the controller solves can be formalized as a meta-level Markov decision process (meta-MDP) (Hay et al., 2014), whose state space corresponds to the intermediate states of computation X , whose action space corresponds to the modules M , and whose transition model corresponds to the evaluator E .
The symbiotic relationship among these components is shown in Fig.
In the bounded version of CRL (Sec. 4.2), the meta-MDP has a ﬁnite horizon whose length is determined by the complexity of the current problem.
In the unbounded version of CRL (Sec. 4.1), the program itself determines when to halt when the controller selects the the HALT signal.
When the program halts, in both versions the current state of computation is produced as output ˆ x out , and CRL receives a terminal reward that reﬂects how ˆ x out matches the desired output x out .
The unbounded CRL also incurs a cost for every computation it executes to encourage it to customize its complexity to the problem.
Note the following key characteristics of CRL.
First, unlike standard reinforcement learning setups, the state space and action space can vary in dimensionality across and within episodes because
CRL trains on problems of different complexity, reducing more complex problems to simpler ones (Sec. 4.1).
Second, because the meta-MDP is internal to CRL, the controller shapes the meta-MDP by choosing which modules get trained and the meta-MDP in turn shapes the controller through its non-stationary state-distribution, action-distribution, and transition function.
Thus CRL simultaneously designs and solves reinforcement learning problems “in its own mind,” whose dynamics depend just as much on the intrinsic complexity of the problem as well as the current problem-solving capabilities of CRL.
3. M AKING A NALOGIES IN THE C OMPOSITIONAL P ROBLEM G RAPH
The solution that we want CRL to discover lies between two extremes, both of which have their own drawbacks.
One extreme is where CRL learns a module specialized for every pair of nodes in the compositional problem graph, and the other is where CRL only learns one module for all pairs of nodes.
Both extremes yield a horizon-one meta-MDP and are undesirable for compositional generalization in a lifelong learner: the former does not re-use past knowledge and the latter cannot ﬂexibly continuously learn without suffering from negative transfer.
We want to encourage the modules to capture the most primitive subproblems, such that they can be composed as atomic computations for other problems.
In our experiments, we use domain knowledge to restrict the representation vocabulary and the function class of the modules.
Constraining each module to learn only a simple transformation gives them a local view on the global problem, thereby forcing the controller to compose the modules together to generalize.
Because the the space of subproblem compositions is combinatorially large, we use a curriculum to encourage solutions for the simpler subproblems to converge somewhat before introducing more complex problems, for which CRL can learn to solve by composing together the modules that had been trained on simpler problems.
Lastly, to encourage the controller to generalize to new node combinations it has not seen, we train on a diverse distribution of compositional problems, such that the controller does not overﬁt to any one problem.
This encourages controller to make analogies between problems during training by re-using partial solutions learned while solving other problems.
Our experiments show that this analogy-making ability helps with compositional generalization because the controller solves new or more complex subproblem combinations by re-using modules that it learned during training.
E XPERIMENTS
To evaluate compositional generalization, we select disjoint subsets of node pairs for training and evaluating the learner.
Evaluating on problems distinct from those in training tests the learner’s ability to apply what it has learned, rather than directly performing a simple nearest-neighbor lookup.
To demonstrate the broad applicability of the compositional graph, we consider the structured symbolic domain of multilingual arithmetic and the underconstrained and high-dimensional domain of transformed-MNIST classiﬁcation.
We compare CRL with “monolithic” architectures that directly map input to output without explicit modularity.
In our experiments, the controller and modules begin as randomly initialized neural networks.
After the program halts, a loss is computed by comparing the current state of computation with the desired output.
The loss is backpropagated through the modules, which are trained with Adam (Kingma &
Ba, 2014).
The controller receives a sparse reward derived from the loss at the end of the computation, and a small cost for each computational step.
The model is trained with proximal policy optimization (Schulman et al., 2017).
4. M ULTILINGUAL A RITHMETIC
This experiment evaluates the unbounded CRL in a multi-objective, variable-length input, symbolic reasoning multi-task setting.
A task is to simplify an arithmetic expression expressed in a source language, encoded as variable-length sequences of one-hot tokens, and produce the answer modulo 10 in a given target language.
To evaluate compositional generalization, we test whether, after having trained on 46200 examples of 2, 3, 4, 5-length expressions ( 2 .
76 · 10 − 4 of the training distribution) involving 20 of the 5 × 5 = 25 pairs of ﬁve languages, the learner can generalize to 5-length and 10-length expressions involving the other ﬁve held-out language pairs (problem space: 4 . 92 · 10 15 problems).
To handle the multiple target languages, the CRL controller receives a one-hot token for the target language at every computational step additional to the arithmetic expression.
The
CRL modules consist of reducers and translators, which do not know the target language and so can only make local progress on the global problem.
Reducers transform a consecutive window of three tokens into one token, and translators transform all tokens in a sequence by the same transformation.
The CRL controller also selects where in the arithmetic expression to apply a reducer.
We trained by gradually increasing the complexity of arithmetic expressions from length two to length ﬁve.
Quantitive results in Fig.
2 show that CRL achieves signiﬁcantly better compositional generalization than an RNN baseline (Cho et al., 2014) trained to directly map the expression to its answer, even
[Figure]
when the RNN has been pretrained or receives 10x more data.
Fig.
8 shows that CRL achieves about 60% accuracy for extrapolating to 100-term problems (problem space: 4 . 29 · 10 148 ).
The curriculum-based training scheme encourages CRL to designs its own edges and paths to connect nodes in the compositional problem graph, solving harder problems with the solutions from simpler ones.
It also encourages its internal representations to mirror the external representations it observes in the problem distribution, even though it has no direct supervision to do so.
However, while this is often the case, qualitative results in Fig.
3 show that CRL also comes up with its own internal language – hybrid representations that mix different external representations together – to construct compositional solutions during for novel problems.
Rather than learn translators and reducers that are speciﬁc to single input and single output language pair as we had expected, the modules, possibly due to their nonlinear nature, tended to learn operations speciﬁc to the output language only.
4. I MAGE T RANSFORMATIONS
This experiment demonstrates the bounded CRL in a single-objective, latent-structured, highdimensional multi-task setting.
A task is to classify an MNIST digit, where the MNIST digit has been randomly translated (left, right, up, down), rotated (left, right) and scaled (small, big).
Suppose CRL has knowledge of what untransformed MNIST digits look like; is it possible that CRL can learn to compose appropriate spatial afﬁne transformations in sequence to convert the transformed
MNIST digit into a “canonical” one, such that it can use a pre-trained classiﬁer to classify it?
To reformulate a scenario to one that is more familar is characteristic of compositional generalization humans: humans view an object at different angles yet understand it is the same object; they may have an accustomed route to work, but can adapt to a detour if the route is blocked.
To evaluate compositional generalization, we test whether, having trained on images produced by combinations of two spatial transformations, CRL can can generalize to different length-2 combinations as well as length-3 combinations.
A challenge in this domain is that the compositional structure is latent, rather than apparent in the input for the learner to exploit.
CRL is initialized with four types of modules: a Spatial Transformer Network (STN) (Jaderberg et al., 2015) parametrized to only rotate, an STN that only scales, an STN that only translates, and an identity function.
All modules are initialized to perform the identity transformation, such that symmetry breaking (and their eventual specialization) is due to the stochasticity of the controller.
Quantitative results in Fig.
4 show that CRL achieves signiﬁcantly better compositional generalization than both ﬁnetuning the convolutional neural network (Springenberg et al., 2014) pretrained classiﬁer and training an afﬁne-STN as a pre-processor to the classiﬁer.
Both baselines perform better than CRL on the training set.
Note that in Fig.
5 the sequence of transformations CRL performs are not necessarily the reverse of those that generated the original input.
R ELATED W ORK
Our work is one of a small growing body of work bridging reformulation and metareasoning with deep learning; our emphasis on learning algorithmic procedures as a means for generalization departs from work that seeks to learn shared representations; we show the beneﬁt of hierarchical learning in low data regimes; we demonstrate a fully disentangled self-organizing neural network where both the transformations and the means of their combination are learned with sparse supervision, a step beyond prior work.
We hope this paper provides a point of uniﬁcation among these areas through which further connections can be strengthened
Reformulation and cognitive development: Reformulation , or re-representing a problem to make it easier to solve, dates back at least to the 1960s (Holte & Choueiry, 2003; Simon, 1969; Anderson, 1990), but is relatively unexplored in deep learning.
CRL’s ability to dynamically re-represent a new problem in the language of previous problems is an initial step towards bridging the idea of re-representation with deep learning.
How the primitive transformations for such re-representation are acquired relates to the study of representational change in cognitive science.
The simultaneous optimization over the the continuous function parameters and their discrete compositional structure in CRL can be seen as an attempt to build a learner that emulates in a limited sense the interplay between abstract and concrete knowledge that is hypothesized to characterize cognitive development: abstract structural priors serve as a scaffolding within which concrete, domain-speciﬁc learning takes place (Spelke, 1990; Pinker, 1994) but domain-speciﬁc learning about the continuous semantics of the world can also provide feedback to update the more discrete structural priors (Gopnik &
Wellman, 2012; Carey, 2015).
Combinatorial generalization through metareasoning:
To solve out-of-distribution problems, both metalearning and metareasoning approaches assume that the new problem shares regularities with previous problems but take different perspectives on how to exploit those regularities.
The metalearning framework (Thrun & Pratt, 2012; Schmidhuber, 1987) seeks to learn a learner that can learn quickly on the new problem.
CRL departs from this paradigm by taking a metareasoning approach (Russell & Wefald, 1991; Hay et al., 2014), where the learner reasons about its own knowledge and computations to solve the new problem.
Much recent metalearning work (Finn et al., 2017;
Nichol et al., 2018; Ravi & Larochelle, 2016; Andrychowicz et al., 2016; Grant et al., 2018;
Mishra et al., 2018;
Lake et al., 2015; Frans et al., 2017; Gupta et al., 2018b;a; Srinivas et al., 2018) focuses on domains where achieving zero-shot generalization is unclear and follow a common paradigm which assumes that the shared representation of monolithic architectures (with the exception of concurrent work (Alet et al., 2018)) can be shaped by the diversity of tasks in the training distribution as good initializations for future learning.
But for combinatorial generalization, monolithic architectures may itself be a barrier to zero-shot generalization: our experiments suggest that such architectures are signiﬁcantly less data efﬁcient than compositional modular approaches.
In these cases it may be possible to extract enough structure during training to generalize out-of-distribution without learning at all.
Indeed, the monolithic/shared-representation paradigm can be improved in two complementary ways.
The ﬁrst, advocated in contemporaneous work (Battaglia et al., 2018), focuses on learning more structured representations of entities and relations with graph networks.
The second, the focus of this paper, focuses on learning structured algorithmic procedures that can implement recursion and composition of partial solutions that operate on state representations.
Instead of learning an algorithm to learn, the learner can learn an algorithm to reason: casting the generalization problem as an algorithm design problem is a growing hypothesis for the mechanism with which humans can generalize to new scenarios (Grifﬁths et al., 2015; Callaway et al., 2017;
Lieder et al., 2017).
Our work is one of a small but growing community (Hamrick et al., 2017;
Graves, 2016;
Oh et al., 2017) linking metareasoning with deep learning.
Compositional neural program induction: Compositional approaches (as opposed to memoryaugmented (Graves et al., 2014; Sukhbaatar et al., 2015; Joulin & Mikolov, 2015; Grefenstette et al., 2015; Kurach et al., 2015; Andrychowicz et al., 2016; Graves et al., 2016) or monolithic (Zaremba & Sutskever, 2014; Kaiser & Sutskever, 2015) approaches for learning programs) to the challenge of discovering reusable primitive transformations and their means of composition generally fall into two categories.
The ﬁrst assumes pre-speciﬁed transformations and learns the structure (from dense supervision on execution traces to sparse-rewards)
(Reed & De Freitas, 2015; Cai et al., 2017; Xu et al., 2017; Chen et al., 2017; Ganin et al., 2018; Bunel et al., 2018; Feser et al., 2016; D ˇ zeroski et al., 2001; Zaremba et al., 2016).
The second learns the transformations but pre-speciﬁes the structure (Devin et al., 2017;
Andreas et al., 2016;
Riedel et al., 2016).
These approaches are respectively analogous to our hardcoded-functions and hardcoded-controller ablations in Fig.
(Zaremba &
Sutskever, 2014) noted the challenge of learning the individual operators and how to combine them; our paper can be considered a candidate to address this challenge.
The closest work in this respect to ours is (Gaunt et al., 2016) which assumes a differentiable interpreter grounded in source code; we take a different perspective of learning algorithms through re-representations, which enables our method to adapt on-line to new problem classes.
Our perspective comes with its own challenges; without grounding in source code, the concept of a for -loop is implicit in the controller, rather than an explicit concept that can exist independently of the data, which is desirable for constructing more complex programs.
Self-organizing learners, disentanglement, hierarchical learning: Research on self-organizing learners (Schmidhuber, 1987;
Dechter et al., 2013; Schmidhuber, 2009; 2012) has so far been relatively theoeretical; CRL is a synthesis of recent work in disentanglement and hierarchical learning that attempts to bring this old idea closer to reality.
Just as the motivation behind disentangled representations (Whitney et al., 2016; Kulkarni et al., 2015; Chen et al., 2016; Thomas et al., 2017;
Bengio et al., 2013) is to uncover the latent factors of variation, the motivation behind disentangled programs is to uncover the latent organization of a task.
Finding such structure is important for generalization because it shields various parts of the representation or program from changing when one component changes.
For example, (Hinton et al., 2018; Rosenbaum et al., 2018; Fernando et al., 2017) have begun to examine how to route through an ﬁxed architecture to encapsulate computation, but the hard limitation that prevents ﬁxed architectures from scaling to problems with complexity beyond that which they were trained is their limited number of computation steps.
A parallel line of work runs through the hierarchical reinforcement learning (HRL) literature (Barto &
Mahadevan, 2003), with recent work (Bacon et al., 2017; Kulkarni et al., 2016; Frans et al., 2017;
Vezhnevets et al., 2017; Nachum et al., 2018) attempting to learn both lower-level policies as well as a higher-level policy that calls them.
Several works have investigated the conditions in which hierarchy is useful for humans (Botvinick et al., 2009;
Solway et al., 2014; Sanborn et al., 2018); our experiments show that the hierarchical structure of CRL is more useful than the ﬂat structure of the RNN for low-data compositional problems.
The fast-weights (FW) literature (Schmidhuber, 1992; Ba et al., 2016) is also related to the idea of selecting different weights at different steps of computation, but with a different motivation for learning context-dependent associative memory (Hopﬁeld, 1982;
Willshaw et al., 1969; Kohonen, 1972; Anderson & Hinton, 2014;
Ha et al., 2016), with the exception of (Schlag & Schmidhuber, 2017).
By combining adaptive computation time (Graves, 2016) with routing, CRL synthesizes all of these above ideas by dynamically composing learned reusable subnetworks together as fully disentangled programs, taking a step towards making the topology of the computation graph self-organizing as well.
D ISCUSSION
This paper sought to tackle the question of how to build machines that leverage prior experience to solve more complex problems than they have seen.
This paper makes three contributions towards the solution.
First, we formalized the compositional problem graph as a language for studying compositionally-structured problems of different complexity that can be applied on various problems in machine learning.
Second, we introduced the compositional generalization evaluation scheme for measuring how readily old knowledge can be reused and hence built upon.
Third, we presented the compositional recursive learner, a domain-general framework for learning a set of reusable primitive transformations and their means of composition that reﬂect the structural properties of the problem distribution.
Grifﬁths et al.
(2019) argued that the efﬁcient use cognitive resources in humans may also explain their ability to generalize, and this paper provides evidence that reasoning about what computation to execute by making analogies to previously seen problems achieves signiﬁcantly higher compositional generalization than non-compositional monolithic learners.
Because problems in supervised, unsupervised, and reinforcement learning can all be expressed under the framework of transformations between representations in the compositional problem graph, we hope that our work motivates further research for tackling the compositional generalization problem in many other domains to accelerate the long-range generalization capabilities that are characteristic of general-purpose learning machines.
A CKNOWLEDGMENTS
The authors would like to thank Alyosha Efros, Dinesh Jayaraman, Pulkit Agrawal, Jason Peng,
Erin Grant, Rachit Dubey, Thanard Kurutach, Parsa Mahmoudieh, Aravind Srinivas, Fred Callaway, Justin Fu, Ashvin Nair, Marvin Zhang, Shubham Tulsiani, Peter Battaglia, Jessica Hamrick,
Rishabh Singh, Feras Saad, Kai-I Shan, and others in the Berkeley Artiﬁcial Intelligence Research
Lab for helpful feedback, discussions, and support.
The authors are grateful for computing support from Amazon, NVIDIA, and Google.
This work was supported in part by the Berkeley EECS
Department Fellowship for ﬁrst-year Ph.D. students, travel funding from Bloomsbury AI, contract number FA8650-18-2-7832 from the Defence Advanced Research Projects Agency (DARPA) under the Lifelong Learning Machines program, contract number FA9550-18-1-0077 from the Air Force
Ofﬁce of Scientiﬁc Research (AFOSR), and the National Science Foundation (NSF) Graduate Research Fellowship Program.
Any opinions, ﬁndings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reﬂect the views of DARPA, AFOSR, or the NSF.
A D ATA
Numerical arithmetic (Sec.
D.1):
The dataset contains arithmetic expressions of k terms where the terms are integers ∈ [0 , 9] and the operators are ∈ { + , × , −} .
The number of possible problems is (10 k )(3 k − 1 ) .
The learner sees 5810 / (2 . 04 · 10 14 ) = 2 . 85 · 10 − 11 of the training distribution.
The number of possible problems in the extrapolation set is (10 20 )(3 19 )
= 1 .
An input expression is a sequence of one-hot vectors of size 13.
[Table]
Multilingual arithmetic (Sec. 4.1):
The dataset contains arithmetic expressions of k terms where the terms are integers ∈ [0 , 9] and the operators are ∈ { + , · , −} , expressed in ﬁve different languages.
With 5 choices for the source language and target language, the number of possible problems is (10 k )(3 k − 1 )(5 2 ) .
In training, each source language is seen with 4 target languages and each target language is seen with 4 source languages: 20 pairs are seen in training and 5 pairs are held out for testing.
The learner sees 46200 / (1 . 68 · 10 8 ) = 2 .
76 · 10 − 4 of the training distribution.
The entire space of possible problems in the extrapolation set is (10 10 )(3 9 )(5 2 ) = 4 .
92 · 10 15 out of which we draw samples from the 5 held-out language pairs (cid:0) (10 10 )(3 9 )(5)
= 9 .
84 · 10 14 possible (cid:1) .
An input expression is a sequence of one-hot vectors of size 13 × 5 + 1 = 66 where the single additional element is a STOP token (for training the RNN).
[Figure]
[Table]
Spatially transformed MNIST (Sec. 4.2): The generative process for transforming the standard
MNIST dataset to the input the learner observes is described as follows.
We ﬁrst center the 28x MNIST image in a 42x42 black background. We have three types of transformations to apply to the image: scale, rotate, and translate. We can scale big or small.
We can rotate left or right.
We can translate left, right, up, and down, but the degree to which we translate depends on the size of the object: we translate the digit to the edge of the image, so smaller digits get translated more than large digits.
In total there are 2 + 2 + 4 × 3 = 16 individual transformation operations used in the generative process. Because some transformation combinations are commutative, we deﬁned an ordering with which we will apply the generative transformations: scale then rotate then translate. For length-2 compositions of generative transformations, there are scale-small-then-translate ( 1 × 4 ), scale-big-then-translate ( 1 × 4 ), rotate-then-translate ( 2 × 4 ), and scale-then-rotate ( 2 × 2 ). We randomly choose 12 of these 16 for training, 2 for validation, 2 for test, as shown in Figure 4.
For length-3 compositions of generative transformations, there are scale-small-then-rotate-then-translate ( 1 × 2 × 4 ) and scale-big-then-rotate-then-translate ( 1 × 2 × 4 ).
All 16 were held out for evaluation.
B L EARNER D ETAILS
All learners are implemented in PyTorch (Paszke et al., 2017) and we will make the code and data available online.
B. B ASELINES
The RNN is implemented as a sequence-to-sequence (Sutskever et al., 2014)
GRU (Cho et al., 2014). The CNN is a variant of an all-convolutional network (Springenberg et al., 2014) fully convolutional network. The afﬁne-STN predicts all 6 learnable afﬁne parameters as in Jaderberg et al.
B. CRL
Controller: The controller consists of a policy network and a value function, each implemented as
GRUs that read in the input expression.
The value function outputs a value estimate for the current expression.
For the numerical arithmetic task, the policy network ﬁrst selects a reducer and then conditioned on that choice selects the location in the input expression to apply the reducer.
For the multilingual arithmetic task, the policy
ﬁrst samples whether to halt, reduce, or translate, and then conditioned on that choice (if it doesn’t halt) it samples the reducer (along with an index to apply it) or translator.
Modules: The reducers are initialized as a two-layer feedforward network with ReLU non-linearities (Nair & Hinton, 2010).
The translators are a linear weight matrices.
C E XPERIMENT D ETAILS
Training procedure: The training procedure for the controller follows the standard Proximal Policy
Optimization training procedure, where the learner samples a set of episodes, pushes them to a replay buffer, and every k episodes updates the controller based on the episodes collected.
Independently, every k (cid:48) episodes we consolidate those k (cid:48) episodes into a batch and use it to train the functions.
We found via a grid search k = 1024 and k (cid:48)
= 256 .
Through an informal search whose heuristic was performance on the training set, we settled on updating the curriculum of CRL every 10 5 iterations and updating the curriculum of the RNN every 5 · 10 4 iterations.
Domain-speciﬁc nuances
: In the case that HALT is called to early, CRL treats it as a no-op.
Similarly, if a reduction operator is called when there is only one token in the expression, the learner also treats it as a no-op.
There are other ways around this domain-speciﬁc nuance, such as to always halt whenever HALT is called but only do backpropagation from the loss if the expression has been fully reduced (otherwise it wouldn’t make sense to compute a loss on an expression that has not been fully reduced).
The way we interpret these “invalid actions” is analogous to a standard practice in reinforcement learning of keeping an agent in the same state if it walks into a wall of a maze.
Symmetry breaking: We believe that the random initialization of the functions and the controller breaks the symmetry between the functions.
For episodes 0 through k the controller still has the same random initial weights, and for episodes 0 through k (cid:48)
the functions still have the same random initial weights.
Because of the initial randomness, the initial controller will select certain functions more than others for certain inputs; similarly initially certain functions will perform better than others for certain inputs.
Therefore, after k iterations, the controller’s parameters will update in a direction that will make choosing the functions that luckily performed better for certain inputs more likely; similarly, after k (cid:48) iterations, the functions’ parameters will update in a direction that will make them better for the inputs they have been given.
So gradually, functions that initially were slightly better at certain inputs will become more specialized towards those inputs and they will also get selected more for those inputs.
Training objective: The objective of the composition of functions is to minimize the negative log likelihood of the correct answer to the arithmetic problem.
The objective of the controller is to maximize reward.
It receives a reward of 1 if the token with maximum log likelihood is that of the correct answer, 0 if not, and − 0 .
for every computation step it takes.
The step penalty was found by a scale search over {− 1 , − 0 .
1 , − 0 .
01 , − 0 .
001 } and − 0 .
01 was a penalty that we found balanced accuracy and computation time to a reasonable degree during training.
There is no explicit feedback on what the transformations should be and on how they are composed. [Figure]
D A DDITIONAL E XPERIMENTS D. N UMERICAL M ATH
The input is a numerical arithmetic expression (e g. 3 + 4 × 7 ) and the desired output (e.g. 1 ) is the evaluation of the expression modulo 10 .
The input and output have the same type.
In our experiments we train on a curriculum of length-2 expressions to length-10 expressions, adding new expressions to an expanding dataset over the course of training in a lifelong learning fashion.
The ﬁrst challenge is to learn from this limited data (only 6510 training expressions) to generalize well to unseen length-10 expressions in the test set ( ≈ 2 14 possible).
The second challenge is to extrapolate from this limited data to length-20 expressions ( ≈ 10 29 possible).
Expressions are represented as sequences of one-hot vectors.
The functions (“reducers”) used in
CRL reduce a length k expression to a length k − 3 one: the controller chooses a window of 3 terms in the expression, and the reducer transforms that into a softmax distribution over a single term in the vocabulary.
This process is repeated until the expression is reduced to one term, at which point the expression persists until the controller HALT s. Choosing to attend to a subset of the input is not domain-speciﬁc; this strategy is applicable in many settings (Mnih et al., 2014; Eslami et al., 2016;
Xu et al., 2015).
We compare with an RNN architecture (Chung et al., 2014) directly trained to map input to output.
Results: Interestingly, though the RNN eventually generalizes to different 10-length expressions and extrapolates to 20-length expressions (yellow in Fig. 6) with 10 times more data as CRL, it completely overﬁts when given the same amount of data (gray).
In contrast, CRL (red) does not overﬁt, generalizing signiﬁcantly better to both the 10-length and 20-length test sets.
We believe that the modular disentangled structure in CRL biases it to cleave the problem distribution at its joints, yielding this 10-fold reduction in sample complexity relative to the RNN.
Indeed, we found that the controller naturally learned windows centered around operators (e.g. 2+ rather than × 4 − ), suggesting that it has discovered semantic role of these primitive two-term expressions by pattern-matching common structure across arithmetic expressions of different lengths.
Note that CRL’s extrapolation accuracy here is not perfect compared to (Cai et al., 2017): but what we believe is noteworthy is that CRL achieves such high extrapolation accuracy with only sparse supervision, without the step-by-step supervision on execution traces, the stack-based model of execution, and hardcoded transformations. [Figure]
D. V ARIATIONS
Here we study the effect of varying the number of functions available to our learner.
Fig.
7a, 7b highlights a particular pathological choice of functions that causes CRL to overﬁt.
If CRL uses four reducers and zero translators (red), it is not surprising that it fails to generalize to the test set: recall that each source language is only seen with four target languages during training with one held out; each reducer can just learn to reduce to one of the four target languages.
What is interesting though is that when we add ﬁve translators to the four reducers (blue), we see certain runs achieve 100% generalization, even though CRL need not use the translators at all in order to ﬁt the training set.
That the blue training curve is slightly faster than the red offers a possible explanation: it may be harder to ﬁnd a program where each reducer can reduce any source language to their specialized target language, and easier to ﬁnd programs that involve steps of re-representation (through these translators), where the solution to a new problem is found merely by re-representing that problem into a problem that learner is more familiar with.
The four-reducers-ﬁve-translators could have overﬁtted completely like the four-reducers-zero-translators case, but it consistently does not.
Indeed, we ﬁnd that when we vary the number of reducers in { 1 , 3 } and the number of translators in { 5 , 8 } in Fig.
7c, the extrapolation performance is consistent across the choices of different numbers of functions, suggesting that CRL is quite robust to the number of functions in non-pathological cases.
D. H OW FAR CAN WE PUSH EXTRAPOLATION ?
The extrapolation accuracy from 6 to 100 terms after training on a curriculum from 2 to 5 terms (46200 examples) on the multilingual arithmetic task (Sec. 4.1).
The number of possible 100-term problems is (10 100 )(3 99 )(5 2 ) = 4 .
29 · 10 148 and CRL achieves about 60% accuracy on these problems; a random guess would be 10% .
D. E XECUTION T RACES : F UNCTION S ELECTION
Fig.
9 compares the execution traces of CRL on the validation, test, and extrapolation sets after training.
The language pairs in the test and extrapolation sets are the same, and the language pairs among training, validation, and test are disjoint
. D. E XECUTION T RACES : E XAMPLES
Here are two randomly selected execution traces from the numerical arithmetic extrapolation task (train on 10 terms, extrapolate to 20 terms), where CRL’s accuracy hovers around 80%. These
[Figure] expressions are derived from the internal representations of CRL, which are softmax distributions over the vocabulary (except for the ﬁrst expression, which is one-hot because it is the input).
The expressions here show the maximum value for each internal representation.
This is a successful execution.
The input is 6*1*3-4+6*0*0+1-7-3+3+3*4+1+1+3+3+6+2+
and the correct answer is 3 .
Notice that the order in which controller applies its functions does not strictly follow the order of operations but respects the rules of order of operations: for example, it may decide to perform addition (A) before multiplication (B) if it doesn’t affect the ﬁnal answer [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
This is an unsuccessful execution trace.
The input is 5+6-4+5*7*3*3*8*0*1-4+6-3*5*3+6-0+0-4-6 and the correct answer is 0 .
Notice that it tends to follow of order of operations by doing multiplication ﬁrst, although it does make mistakes (D), which in this case was the reason for its incorrect answer.
Note that CRL never receives explicit feedback about its mistakes on what its functions learn to do or the order in which it applies them; it only receives a sparse reward signal at the very end.
Although (C) was a calculation mistake, it turns out that it does not matter because the subexpression would be multiplied byanyways.
[Table]
[Table]
[Table]