A UTOMATICALLY C OMPOSING R EPRESENTATION T RANSFORMATIONS AS A M EANS FOR G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
How can we build a learner that can capture the essence of what makes a hard problem more complex than a simple one, break the hard problem along characteristic lines into smaller problems it knows how to solve, and sequentially solve the smaller problems until the larger one is solved?
To work towards this goal, we focus on learning to generalize in a particular family of problems that exhibit compositional and recursive structure: their solutions can be found by composing in sequence a set of reusable partial solutions.
Our key idea is to recast the problem of generalization as a problem of learning algorithmic procedures: we can formulate a solution to this family as a sequential decision-making process over transformations between representations.
Our formulation enables the learner to learn the structure and parameters of its own computation graph with sparse supervision, make analogies between problems by transforming one problem representation to another, and exploit modularity and reuse to scale to problems of varying complexity.
Experiments on solving a variety of multilingual arithmetic problems demonstrate that our method discovers the hierarchical decomposition of a problem into its subproblems, generalizes out of distribution to unseen problem classes, and extrapolates to harder versions of the same problem, yielding a 10-fold reduction in sample complexity compared to a monolithic recurrent neural network.
We further show that our method can compose learned spatial transformations to recover canonical MNIST digits from transformed ones.
I NTRODUCTION
Learning to generalize to harder problems from solving simpler ones is a difﬁcult challenge in machine learning because the learner must discover regularities that cover a combinatorially large space of problems from only limited data.
To make progress towards this goal, this paper focuses on a particular class of problems that exhibit compositional and recursive structure: their solutions can be found by composing in sequence a small set of reusable partial solutions.
The key perspective of this paper is to recast the problem of generalization to a problem of learning algorithmic procedures over representation transformations: discovering the structure of a family of problems amounts to learning a set of reusable primitive transformations and their means of composition.
Approaches to this challenge generally have either assumed pre-speciﬁed transformations but learned the structure, often with dense supervision, or learned the transformations but pre-speciﬁed the structure (see Sec. 5).
Pre-specifying one or the other is a major limitation to learning programs that can adapt to different types of data.
The contribution of this paper is to propose a direction towards learning both the structure and transformations together with sparse supervision.
We propose the compositional recursive learner (CRL), which learns both the structure and transformations of a modular recursive program that iteratively re-represents the input representation into more familiar representations it knows how to compute with.
Our approach is based on the following core ideas.
The ﬁrst core idea is that training on a distribution of composite problems can encourage the spontaneous specialization of reusable computational units that perform transformations between representations.
Rather than learning to solve a new problem from scratch, a learner can make an analogy between the new problem and previously en- countered problems and re-represent the new problem into one it knows how to solve.
The second core idea is to treat the construction of a program as a sequential decision-making problem over these transformations.
This formulation not only implements loops, recursion, composition, and selective attention, but enables metareasoning , where the learner dynamically customizes its program complexity to the problem instance.
The third core idea is to operationalize the learner in differentiable programs by using reinforcement learning algorithms to solve the discrete optimization problem of selecting which transformation to apply, while simultaneously learning the transformations with backpropagation.
Multilingual arithmetic problems are a challenging but minimal family of tasks that exhibit compositional and recursive properties (see Sec. 2).
Understanding how to learn to decompose and exploit the structure of arithmetic is a ﬁrst step towards solving any more complex domain, which would at least have the structure we study here in arithmetic.
Our experiments on generalization to a completely unseen distribution of harder problems from only observing less than 3 × 10 − 4 of the training distribution show that CRL outperforms baselines without recursive or compositional reusability in their structure, even when the baseline is pretrained on all of the primitive transformations of the dataset.
We further show that our method can compose learned spatial transformations to recover canonical MNIST digits from transformed ones.
Summary of contributions: (1) To consider a learner whose primitive transformations and means of their composition match the structure of the problem distribution, we narrow the scope of problems to those that exhibit compositional and recursive structure, which we formalize in Sec. 2.
In Sec.
3 we propose a hypothesis that the emergence of primitive transformations can be encouraged by training on a distribution of composite problems.
(3) We frame the problem of learning the structure of a program as a sequential decision-making problem over these transformations.
These two pieces comprise the CRL framework.
(4) To implement CRL in differentiable programs, we present a formalism for expressing the construction of a neural network in terms of an Markov decision process that generates computation graphs.
(5) Our experimental ﬁndings in Sec.
4 highlight the limitations of monolithic static network architectures in generalizing to structurally similar but more complex problems than they have been trained on, providing evidence in favor of learners that dynamically re-program themselves, from a repertoire of learned reusable computational units, to the current problem they face.
G ENERAL S ETUP
It would be impossible to generalize if the universe exhibited no regularities.
Discovering and exploiting the structure in their world helps humans generalize make predictions and solve problems in new situations.
We formalize the notion of compositional structure in terms of problems , describe the characteristics of these problems, and deﬁne the challenge of how to learn programs that reﬂect the structure of problems.
Problems: To solve a problem P i is to transform its input representation x i into its output representation y i .
x i and y i are associated with their respective types t x and t y .
In a recursive problem , the vocabulary with which x i and y i are represented is the same.
For example, arithmetic is recursive problem, because different arithmetic expressions are represented using the same vocabulary.
Language translation can be thought of as a recursive problem if all words in the language are drawn from the same large vocabulary that includes all words in all languages. However, because different languages occupy different sections of the vocabulary, a word in one language is a different type than a word in another language.
Composite problems: Problems can be compositions of other problems, which themselves can be composite.
A composite problem P a = P b ◦ P c means that the solution of P a can be found by ﬁrst solving P c , then solving P b with the output of P c as input.
In this case, P b and P c are subproblems with respect to P a .
The goal: Let P be a set of composite recursive problems that share a set of regularities. We want the learner to extrapolate to solve tasks whose pairs of types were not observed with supervisory signal during training.
For example, if the learner has received supervision for translating English to
French (EF), and French to Spanish (FS), but not English to Spanish (ES), we would like the learner to still be able to translate English to Spanish by composing together EF and FS transformations that it had learned during training.
We also want the learner to extrapolate to solve recursive problems that have higher complexity than it has seen (e.g., the learner was trained to solve up to 5-length arithmetic expressions but is tested on 10-length expressions).
The challenge: For a particular problem P i , the learner only is given an input and the output type ( x i , t y i ) and is expected to produce output y i .
The only supervision signal it receives is its prediction error.
The learner is trained on a set of problems P train ∈ P and tested on a disjoint set of problems P test ∈ P .
Different problems may require different compositions of partial solutions to subproblems, so ﬂat learners that directly attempt to map input to output may quickly overﬁt to P train and do not generalize to P test .
However, although in general the speciﬁc composition of subproblems of a particular problem instance is unknown, it may be possible to leverage domain-speciﬁc knowledge of how problems from a task distribution are decomposed in general to learn how to compose partial solutions to the subproblems of a particular problem instance. The next section describes our approach for doing so.
A L EARNER T HAT P ROGRAMS I TSELF
The Compositional Recursive Learner (CRL) explicitly attempts to learn transformations that it can compose together to extrapolate to new problems. CRL begins by incorporating domain-speciﬁc knowledge of the compositional structure of the problem distribution as an inductive bias that restricts the representation vocabulary and the functional form of its modules.
Given such constraints, learning to compose these modules can be formulated as a sequential decision-making problem in a meta-level Markov decision process, in which CRL can leverage deep reinforcement learning algorithms to optimize the composition of the modules while learning their parameters with backpropagation.
3. P ATTERN -M ATCHING OVER P ROBLEMS VIA A NALOGY
If CRL is trained on a diverse distribution of composite problems that share enough structure, the hope is that it can distill out the individual structural components into specialized, modular computational units : atomic function operators that perform transformations between representations.
These modules, each one of which is learned from the empirical distribution, can then be composed to iteratively re-represent an out-of-distribution problem into a in-distribution one, thereby providing a possible solution to the extrapolation problem.
Indeed, often it is easier to solve a problem P a by making an analogy to a problem P b that one already knows how to solve, rather than retraining to solve P b from scratch (Schmidhuber, 2015; Simon, 1969).
Transforming the type of input representation from t x a to t x b is one way of making such analogies.
Therefore, if a learner hasn’t ever translated French to Spanish but knows how to translate French to English and English to Spanish, it can reduce the French to Spanish problem to a English to Spanish problem by transforming the French input to English input, which it knows how to compute with.
If a learner hasn’t ever seen a 10 -length arithmetic expression before, but it knows how to reduce an n -length expression to an ( n − 1) -length one, then it can iteratively reduce the 10 -length expression down to a length it knows.
To encourage such compositionality and specialization to emerge requires constraints on the internal states and modules. This paper contains experiments where the internal states live in the same vocabulary as the external representations, the modules are parametrized the same way as the generative transformations, and the learner is trained with a curriculum of tasks to encourage the learner to reuse previously learned modules for more complex tasks.
3. L EARNING THE S TRUCTURE OF A C OMPUTATION AS A S EQUENTIAL D ECISION -M AKING P ROCESS
A transformation between representations can be generalized as any computation which changes the state of a ﬁnite state machine to another.
Learning how to apply computational units in sequence can then be formulated as a sequential decision-making problem, where the state space is the intermediate results produced by a program and the action space is the set of computations.
This kind of sequential decision problem can be formalized as a meta-level Markov decision process (MDP) (Hay et al., 2014), deﬁned by a tuple ( X , F , P meta , r, γ ) .
X is the set of information states (intermediate results of computation), C is the set of computations, P meta ( x j , f j , x j +1 ) is the transition model that expresses the probability that at step j the computation f j will change the information state from x j to x j +1 , γ is a discount factor.
The goal of CRL is to select a series of computations f to iteratively transform the input x into its predicted output ˆ y .
When CRL selects a special computation, the HALT signal, the current result is produced as output.
The learner incurs a cost for every computation it executes and receives a terminal reward that reﬂects how its output ˆ y matches the desired output y .
The computation cost and the HALT signal differentiate the meta-level
MDP from a generic MDP, requiring CRL to balance program complexity and performance.
The result is a program composed of a sequence of computations that customizes its complexity to the problem.
3. L EARNING THE T RANSFORMATIONS AND THEIR M EANS OF C OMPOSITION FROM S PARSE R EWARDS
The meta-level MDP describes how to choose transformations, but not how to learn the form of the transformations themselves.
We solve this problem by learning differentiable programs composed of neural networks, in which both the transformations and their composition are learned.
The implementation consists of a controller π ( f | x, t y ) , a set of functions f k ∈ F , and an evaluator (Fig. 1).
At step j , the controller observes the intermediate state of computation x j and the target output type t y .
It selects a function f k and a portion of x j to operate on.
The evaluator applies f k to transform x j into a new state of computation x j +1 .
When π selects HALT , a loss is computed by comparing the current state of computation with the desired output.
The loss is backpropagated through the functions, which are trained with Adam (Kingma & Ba, 2014).
The controller receives a sparse reward derived from the loss, incurs a cost for each computation, and is trained with proximal policy optimization (Schulman et al., 2017).
The state space (excluding the initial state, which is CRL’s input) consists of the outputs of the functions, the action space consists of the functions themselves, and the transition model is the evaluator.
Unlike standard reinforcement learning problems, here the state space and action space can vary in dimensionality, as we show in our experiments.
Moreover, as CRL trains, the function parameters change, so transition dynamics are non-stationary.
To overcome these optimization difﬁculties, we draw from (Dechter et al., 2013) and gradually grow the complexity of the problems during training such that learning more complex compositions are bootstrapped from transformations learned in the context of simple compositions earlier on (see Appx.
C for intuition).
Because the outputs of a function are the internal representations of the larger neural network, CRL simultaneously designs its own internal representation language and the transformations that convert between them.
CRL’s states of computation are its internal representations, which may or may not correspond with the external representation given by the problem.
The meta-MDP is the the learner’s own internal representation of the problem graph that it builds from the many tasks it encounters.
Ideally, eventually its internal states and modules will converge to respectively mirror the representations and generative transformations in the external problem graph, which in general is unknown, by pattern-matching across tasks to extract shared subproblems whose solutions it can encapsulate in its modules.
3. D ISCUSSION OF D ESIGN C HOICES
The dynamic nature of the CRL computation graph highlights a connection with other work in the recent literature on dynamic neural network architectures.
One important approach is Routing Networks (Rosenbaum et al., 2018), which learn to route through different choices of a ﬁxed set of fully-connected layers of a convolutional neural network. However, in contrast to CRL, they do not incorporate constraints on the intermediate representations and modules and do not use a variablelength computation horizon. This limits Routing Networks from re-representing new problems in terms of subproblems it has seen before, which is at the core of how CRL approaches the extrapolation problem.
Our framework provides a language for describing a learner that designs itself: making the connection between programs, neural networks, and MDPs establishes the equivalence between computational units, functions, and representation transformations and opens opportunities to borrow tools from reinforcement learning and deep learning as a practical means for solving both the continuous and discrete optimization problems of respectively learning the transformations and their means composition.
By explicitly attempting to discover and exploit the compositional structure of composite problems, we transform the generalization problem into a problem of learning algorithmic procedures over transformations between representations.
E XPERIMENTS
The following experiments explore the compositional learning capabilities of CRL.
In the highly structured domain of multilingual arithmetic, CRL discovers primitive modules that capture primitive operations and composes these modules to extrapolate to a combinatorially larger space of problems it has never seen.
In underconstrained domains such as compositional image transformations, CRL learns to compose imperfect knowledge of primitive transformations to iteratively re-represent transformed MNIST digits into canonical ones.
4. M ULTILINGUAL A RITHMETIC
Multilingual arithmetic problems is a domain that exactly tests a learner’s capability to exploit compositionality to solve problems beyond its training distribuition.
Arithmetic expressions are modulo10 come in m different types , each corresponding to a language that expresses it.
An example input is three plus four times seven (the source language is English) and an example corresponding output is uno (the target language is Spanish).
We arbitrarily chose m = 5 languages:
English, Numerals, PigLatin, Reversed-English, Spanish.
During training, each source language is seen with four target languages (and one held out for testing) and each target language is seen with four source languages (and one held out for testing).
If a learner successfuly captures the compositional structure of these problems, it should automatically (1) solve problems of signﬁcantly longer length by recursively reducing these problems to smaller instances(2) solve problems with unseen different language pairs by translating between languages.
The input is a tuple ( x ( s ) , t y ) , where x is the arithmetic expression expressed in source language s , and t y is the language with which the learner should output the answer.
The learner trains on a curriculum of a limited set of 2, 3, 4, 5-length expressions ( 2 . 76 · 10 4 of the training distribution), and is asked to generalize to 5-length expressions (test set) extrapolate to 10-length expressions (extrapolation set) with unseen language pairs (problem space: 4 . 92 · 10 15 problems; see Appx.
A).
CRL begins with several randomly initialized reducers and translators , subject to the constraints that a reducer transforms a window of three terms (selected by the conroller) into a softmax distri- [Figure] bution over the vocabulary, and a translator applies the same function to every element of a length- k sequence to produce a length- k sequence of softmax distributions.
The controller has access to t y at every step of computation, but the functions do not, forcing the functions to learn behaviors agnostic of the language pair and putting the burden of choosing how to generalizing to different input-output types on the controller.
CRL is compared with an RNN architecture (Chung et al., 2014) directly trained to map input to output, where t y is concatenated with every element of x ( s ) .
Results: With 10 times more data the RNN (yellow) generalizes to the 5-terms test set but fails to extrapolate to the 10-terms dataset.
This is an intriguing result that suggests that with enough data the RNN’s distributed representation captures some compositionality in the language pairs but cannot extrapolate to a larger problem space than which it has been trained.
In contrast, CRL (red in Fig. 2) can extrapolate consistently better because, by design, it can solve the longer expressions by composing its learned primitives more times.
It is not just domain-speciﬁc knowledge: Although they are all learned, the existence of CRL’s attention mechanism and the topology of its reducers and translators may give it a domain-speciﬁc advantage over the RNN.
However, an RNN that was pre-trained to 100% accuracy on domainspeciﬁc auxiliary tasks (translating between numerical math to all other languages and of solving 3-term numerical math problems) still cannot extrapolate with 10 times more data, even though these auxiliary tasks give the RNN a nontrivial boost on the test set.
This experiment highlights a fundamental gap in knowledge representation between a learner endowed with an inductive bias for compositionality and traditional monolithic function approximators, echoed in Fodor & Pylyshyn (1988); Lake & Baroni (2017).
Qualitative Analysis
Though the randomly initialized functions initially have no semantic meaning, as training progresses the functions gradually specialize such that after convergence the controller chooses a particular ordering of the functions given a new input expression.
The controller gradually learned windows centered around operators (e.g. 2 + 3 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.
Fig.
3 elucidates how CRL makes analogies to re-represent an unseen problem into one that is more unfamiliar.
Because the language pair is one it has not seen, the controller learns to route through expressions it does know how to work with to solve a new problem that it has never seen before.
CRL learns reducers that reduce to a particular language from any language, which result in intermediate hybrid-language expressions.
Notably, after it has reduced the entire expression, it takes an additional step to translate the ﬁnal term to a term in the target language, even though the term before translation is not in either the source or the target language.
Indeed, this shows that the controller reasons about its own computation, taking additional steps that it deems necessary to compose the primitives (the translators and reducers)
it had learned to solve unseen problems (although it sometimes tries to HALT unnecessarily: see Appx.
D.5).
Rather than learn translators and reducers that are speciﬁc to single input and single output language pair as we had expected, the functions tended to learn operations speciﬁc to the output language only, possibly due to their nonlinear nature
4. I MAGE T RANSFORMATIONS
This section explores how CRL may be applied to signiﬁcantly less structured domains.
The task is to classify an MNIST digit, where the MNIST digit has been randomly translated, rotated, and resized.
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 common in the real world: 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 be able to re-represent unfamiliar situations into familiar ones by composing primitive operations that capture the axes of variation in a domain is a beneﬁt of a learner endowed with compositionality.
In Fig.
4, CRL is trained on a distribution of MNIST digits that were ﬁrst modiﬁed by either one transformation or two transformations.
CRL is initalized with four modules: a Spatial Transformer
Network (STN) (Jaderberg et al., 2015) parametrized to only rotate (action 0), an STN that only scales (action 1), an STN that only translates (action 2), and an identity function (action 3).
The bottom row of Fig.
4 is on an unseen dataset of MNIST images generated by composing three transformations (one more than CRL has ever seen).
In this transfer setting, we tested whether it is possible to freeze the modules learned from seeing only one or two transformations, and only re-learning the controller on the dataset with three transformations.
Quantitatively, CRL averages about 70% accuracy, oscillating between 60% and 80% accuracy (as opposed to about 85%, oscillating between 80% and 90% in the one-two transformation set).
One possible reason for a poorer performance on the transfer set can be because freezing the modules limits the ability of CRL to improve: as a result, if the modules are imperfect and do not capture precisely the transformations in the underlying generative process of the data, the controller is limited in what it can do.
Furthermore, this domain is signﬁcantly more challenging than the mulilingual arithmetic domain because the underlying compositional structure of the data is latent: the learner only receives pixels as input, so to discover that there was a sequence of transformations that generated its input is nontrivial.
Lastly, this domain is high dimensional, which poses several challenges for preventing degeneracies.
Unless more constraints are applied to the funcitons and the intermediate representations, CRL could easily fall into a trap of only choosing one module (since all modules take in pixels and produce pixels) and get stuck in suboptimal local optima.
We plan to investigate better ways of stabilizing and regularizating the training for future work.
R ELATED W ORK
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); CRL is an initial step towards bridging this idea with deep learning.
Indeed for generalizing to multiple different tasks, CRL departs from the metalearning framework (Thrun & Pratt, 2012; Schmidhuber, 1987), which seeks to learn a learner that can learn quickly on the new problem and takes 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.
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 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 growing community (Hamrick et al., 2017; Graves, 2016;
Oh et al., 2017) linking metareasoning with deep learning.
Compositionality is a major motivation for neural program induction, whose work generally falls 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).
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.
Work that has explored compositionality in the input space include (Schlag & Schmidhuber, 2017; Battaglia et al., 2018).
Research on self-organizing learners (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.
Parascandolo et al.
(2017) focuses on a complementary direction to ours; whereas the focus on learning causal mechanisms for a single step, we focus on learning how to compose the modules.
D ISCUSSION
We have considered learning to generalize and extrapolate with limited data to harder compositional problems than a learner has previously seen.
We have taken steps toward this challenge by presenting a characterization, algorithm, and implementation of a learner that programs itself automatically to reﬂect the structure of the problem it faces.
Our key ideas are (1) transforming representations with modular units of computation is a potential solution for decomposing problems in a way that reﬂects their subproblem structure; (2) learning the structure of a computation can be formulated as a sequential decision-making problem; (3) tools from reinforcement learning and deep learning can serve as a practical means for solving both the continuous and discrete optimization problems of respectively learning the transformations and their means of composition with sparse rewards.
Our experiments support the core perspective of this paper: for compositional problems under low data regimes, casting the generalization problem as a problem of learning algorithmic procedures over representation transformations offers signiﬁcant advantages over monolithic architectures that only rely on the diversity in the data to shape a shared representation.
To truly learn the subproblem decomposition of a task distribution involves discovering its compositional structural form and exploiting this structural form as an inductive bias to extrapolate to new problems that exhibit the same structural form but require different compositions of representation transformations.
This paper takes a step toward this goal by assuming the structural form as given, thereby focusing on the problem of learning to compose learnable domain-restricted representation transformations that can be reused to extrapolate to new problems.
This work brings together ideas from various areas of research – reformulation, metareasoning, program induction, modularity, self-organizing neural networks – to provide a proof-of-concept for learning compositional and recursive learners that design themselves and reason about their own computations
A D ATA
Numerical math:
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 math:
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 theheld-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).
[Table]
[Figure]
[Table]
[Figure]
[Table]
B L EARNER D ETAILS
All models 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.
The routing networks (experimental comparison in Appx. D.2) implementation used the same hyperparameters as those for the best CRL model.
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.
Functions:
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.
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
[Figure]
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. 5) 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.
D. A BLATION S TUDY
[Figure] [Figure]
We compare our approach to a baseline that ﬁxes the number of computation steps, uses a different controller per target language and per computation step, and uses a different set of modules per computation step.
This baseline can be seen as a multi-step mixtures of experts, where there is a hard choice of the expert at every step.
This experiment shows the beneﬁt of allowing a ﬂexible computation horizon and of reusing modules across computation steps.
Because the baseline uses a different set of functions per computation step, it is unable to exploit compositional structure to reuse its learned functions for simpler versions of the same problem.
A ﬁxed computation horizon also limits how the baseline can adapt to problems of variable length.
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.
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 : 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 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]
This is an unsuccessful execution trace.
The input is
[Table]
[Table]
[Table]
[Table]
{+[Table] [Table] 8*0*1-4+6-3*5*3+6-0+0-4-
+} 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 by 0 anyways.
[Table]
[Table] [Table]
D. E XECUTION T RACES : F UNCTION S ELECTION
[Figure]
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.
The functions beginning with R are reducers, and the functions beginning with T are translators.
CRL takes an additional step to translate during the test and extrapolation execution traces.
The traces that are signiﬁcantly longer are because CRL tries to HALT too early and wastes computation steps.