Guiding Transformers to Process in Steps

1 INTRODUCTION . Daniel Kahneman has pointed out that there is a fundamental difference in how humans solve the following two tasks ( Kahneman , 2011 ) : ( 1 ) complete the phrase “ bread and . . . ” ; ( 2 ) complete the equation “ 17× 24 = . . . ” . The answer to ( 1 ) comes to mind instantly , with no mental effort , and is a result of a computation that one is not consciously aware of and could not explain . The answer to ( 2 ) requires time and effort to come up with , and is a result of a sequence of computations that one consciously carries out . The former mode of cognition is what Kahneman calls System 1 , the latter System 2 . We currently have separate tools for solving each of these two kinds of problems on a computer , but it is not yet clear how to build a single system capable of both modes of cognition simultaneously . Many tasks that exercise System 2 in humans involve executing some fully-specified algorithm and thus are quite straightforward to solve using conventional computer programs . However , classical programming is not applicable to System 1 tasks as these typically correspond to functions that we do not know how to implement . Instead , System 1 tasks are usually tackled by approximating the target functions from observed input-output pairs . Whether such a learning-based approach could be successfully extended to System 2 tasks is an open question , and one that we aim to explore in this paper . We argue that the classical System 1 paradigm of training neural networks to approximate functions from inputs and outputs corresponding to a given task is not a scalable approach for tackling tasks from the System 2 domain . We claim that , given only inputs and outputs , neural networks will not receive enough supervision to converge to a solution when trained to perform algorithmic tasks that humans solve via System 2 . To circumvent this , we propose to supplement the training data with intermediate results that would be helpful to compute before arriving at the final output . By modifying the training data in this way , we are effectively changing the training objective — instead of having to learn any arbitrary way of mapping inputs to outputs , the neural networks are now trained to compute the outputs from the inputs through a particular sequence of intermediate steps . We hypothesize that such additional guidance might be necessary in order for neural networks to learn algorithmic tasks , and we empirically evaluate this hypothesis . Much of existing work in applying neural networks for System 2 tasks has been focused on making neural architectures more aligned with the execution of algorithms . While fine-tuning the model architectures may ultimately lead to superior System 2 capabilities ( after all , the architecture of a calculator allows it to perform arithmetic at a super-human level ) , we instead explore the challenge of extending the capabilities of existing state-of-the-art System 1 models — namely , the Transformer — towards the System 2 domain . In particular , we take inspiration from the human ability to solve algorithmic tasks by writing down symbols on paper and propose to formulate such tasks as sequence-to-sequence problems where the target sequence includes both the outputs and all the intermediate symbols that a human would write ( or more ) . That is , instead of modifying the models , we experiment with modifying the data . Our main contributions can be summarized as follows : • We show that it is possible to hand-code a 1-layer 1-head Transformer to compute any finite function if enough intermediate steps are used ; • We show that a 1-layer 1-head Transformer can be trained to perform binary addition if the target sequences include intermediate results , whereas it fails to learn the task when predicting the outputs directly from the inputs ; and • We show that a Frozen Pretrained Transformer can be trained to perform binary addition if the target sequences include intermediate results , whereas it fails to learn the task when predicting the outputs directly from the inputs . 2 RELATED WORK . The need of bridging the gap between System 1 and System 2 capabilities in current deep learning models has recently been emphasized by Yoshua Bengio ( Bengio , 2019 ; Bengio et al. , 2021 ) . The work of Bengio and colleagues addresses the higher-level cognition in its broadest sense , seeking to model and incorporate into current systems concepts such as consciousness ( Bengio , 2017 ) , causality ( Bengio et al. , 2020 ) , agency ( Thomas et al. , 2018 ) , and global workspace ( Goyal et al. , 2021 ) . In Goyal & Bengio ( 2020 ) , they question the paradigm of classical statistical learning and propose to shift from training models on curated datasets towards training agents in complex non-stationary environments . In contrast to that , our approach lies within a conventional learning framework ( namely , sequence-to-sequence modeling with Transformers ) , and addresses System 2 in a slightly narrower sense ( namely , as algorithmic execution ) . A considerable fraction of recent work in applying neural networks for algorithmic tasks has been concerned with modifying or augmenting neural architectures to make them more suitable for algorithmic execution . One very common approach involves trying to bring components from classical computer architecture into the neural setting : Neural Turing Machines give neural networks access to dynamic external memory ( Graves et al. , 2014 ; 2016 ) , Stack-augmented RNNs augment the networks with an infinite pushdown stack ( Joulin & Mikolov , 2015 ) , Neural Random Access Machines use registers and introduce the ability to manipulate and dereference pointers ( Kurach et al. , 2016 ) , while Neural Arithmetic Logic Modules represent neural versions of the ALU ( Trask et al. , 2018 ; Madsen & Johansen , 2020 ) . A parallel but closely related line of research frames these new kinds of augmented architectures as neural controllers endowed with access to external interfaces and applies reinforcement learning techniques to train them to solve algorithmic tasks by interacting with the interfaces ( Zaremba & Sutskever , 2015 ; Zaremba et al. , 2016 ) . The approach of using variable amount of computation per input , known as Adaptive Computation Time ( ACT ) , has been introduced by Graves ( 2016 ) . Universal Transformers ( Dehghani et al. , 2019 ) integrate ACT into the Transformer architecture , allowing each output symbol to be a result of a variable number of applications of a single Transformer layer . While these works are , like ours , motivated by the apparent necessity of intermediate processing in certain tasks , they use learned ponder time and learned ponder content whereas our work uses given ponder time and given ponder content . The main advantage of learned ponder time and content is that the model is free to discover and learn to perform the intermediate computations that it finds most useful . This also means that the researcher does not need to know the correct intermediate steps to train the model and can thus tackle a potentially larger class of problems . The main disadvantage of learned ponder time and content is that the amount of supervision per forward pass decreases as the model uses more intermediate steps , and the signal can quickly become too weak for the model to train at all . It is also worth noting that intermediate steps are typically given when humans are taught to perform algorithmic tasks ( rather than being asked to infer intermediate computations by just looking at input-output pairs ) , which might be an important argument for given ponder content . The idea of providing supervision beyond input-output examples has already appeared in several projects , although , in our view , it still remains largely underexplored . Reed & de Freitas ( 2016 ) train Neural Programmer-Interpreters by providing supervision on the correct action sequences ( execution traces ) of the recurrent controller , Mirman et al . ( 2018 ) train differential Neural Computational Machines with extra supervision on the movements of the read-write heads , and Mirman et al . ( 2018 ) train Neural Execution Engines with extra supervision on the attention masks . While the added training signal does lead to better sample efficiency and improved generalization , a major limitation of all of these approaches , as acknowledged by Mirman et al . ( 2018 ) , is that the extra supervision needs to be highly specialized as it applies to very specific components of each architecture . Our approach of supplementing the target sequences with intermediate results is completely architectureindependent and thus does not share this limitation . Veličković et al . ( 2020 ) use extra supervision at the data level as well , though their work is focused specifically on tasks involving graph-structured inputs and is thus based on neural network architectures that are tailored for processing graphs . In contrast to that , we adopt a general sequence-to-sequence modeling framework , with an intention to imitate a human executing an algorithm on a piece of paper . Our main goal is to explore how and whether powerful existing System 1 models could be used in the System 2 domain , rather than to find a neural architecture that would be most suitable for executing algorithms . We therefore use vanilla decoder-only Transformers in our experiments . 3 MOTIVATION . The distinguishing characteristic of System 1 tasks is that they can be solved instantly and unconsciously , in something akin to a single “ forward pass ” through a human neural network . It thus seems plausible that all input-output mappings corresponding to System 1 tasks are in principle computable via a single forward pass through a sufficiently large artificial neural network . Recent success in deep learning shows that training bigger models on more data indeed makes it possible to solve more and more System 1 tasks , and it is not unreasonable to expect this trend to continue . When it comes to solving System 2 tasks , however , scaling up appears to be a dead-end . Consider the problem of completing the sentence “ The first digit of the n-th Fibonacci number is . . . ” , where n is replaced by some positive integer . Since a single forward pass through a neural network ( no matter how large ) involves a constant amount of computation , there will be some number N ∈ N such that for all n > N the correct output is not computable . It follows that the only viable approach for solving these kinds of algorithmic tasks is to arrive at the output through intermediate steps . The main value of introducing intermediate steps is that it gives control over the level of model expressivity required to implement a solution to a particular task . No matter how many atomic operations separate the outputs from the inputs in a given algorithmic problem , performing those exact operations in sequence would make each execution step only as complex as a single atomic operation . This is arguably the main reason why a human brain can “ implement ” something like integer multiplication even though it does not contain a circuit for multiplying arbitrarily large numbers directly — the key is that a complex problem can be decomposed into simpler ones until each sub-problem becomes solvable “ atomically ” via System 1 . By the same token , a neural network that is not expressive enough to compute some output through 0 intermediate steps may be expressive enough to compute the same output through l intermediate steps for some l > 0 , since each of the l steps would involve computing a simpler function . Since many System 2 tasks are algorithmic , their outputs can be made arbitrarily distant from the inputs ( in terms of the number of atomic operations separating the two ) , which implies that step-by-step processing is the only scalable approach for solving them . The necessity of intermediate processing does not by itself imply the amount of guidance that the neural networks should receive during training . The exact intermediate computations can either remain unspecified and be inferred by the model , or be fixed and provided as part of the training data . Although there are pros and cons to both alternatives ( as we have outlined in section 2 ) , we see the latter as more promising and choose to include the intermediate results in the target sequences used during training . Our position is based on the following three observations : first , it is natural to provide complete supervision over the intermediate computations when teaching humans to perform algorithmic tasks ; second , learned intermediates create the vanishing supervision problem ( the longer the ponder time , the less training signal is received per forward pass ) ; and third , for all algorithmic tasks the exact intermediate results are fully known anyway . While manually increasing the amount of supervision goes against the tradition of minimizing the reliance on expert knowledge that arguably lead to most of the recent achievements in the System 1 domain , System 2 tasks represent a structurally different class of problems for which a different set of constraints may apply . The arguments laid out above have the following important implications : • Expressiveness : a neural network implementing only some very basic atomic operations should in principle be capable of expressing arbitrarily complex functions as long as enough intermediate steps are used ; • Training : a neural network that fails to learn some algorithmic task when trained without intermediate processing may be able to learn the same task if enough intermediate steps are used ; and • Composability : a neural network that is pretrained on some general-purpose task and does not contain a circuit implementing some algorithmic computation directly may nevertheless be able to solve the same task by computing the output through simpler operations which it does contain circuits for . We establish the first of these implications in Section 4 and we empirically validate the latter two in Sections 5 and 6 . Even though demonstrating a single example where a statement holds does not prove it in general , our empirical results confirm that the hypothesized phenomena do occur in the particular contexts we considered in this work and thus serve as evidence in favor of the latter two implications .

Summary: The paper presents the authors’ hypothesis: intermediate steps are necessary for Transformers to perform well on algorithmic tasks. They analyse Transformers on a simple binary addition task with and without intermediate steps. They show that only the one with immediate steps and supervision over these steps can generalise. In contrast to most of the prior work, the authors argue for supervision in the output domain. This makes the approach architecture-agnostic.