What Algorithms can Transformers Learn? A Study in Length Generalization
Authors: Hattie Zhou
Arwen Bradley
Etai Littwin
Noam Razin
Omid Saremi
Joshua M. Susskind
Samy Bengio
Preetum Nakkiran
Venue: ICLR 2024 poster (2024)
Keywords: length generalization
systematic generalization
understanding
transformer
scratchpad
LLM
algorithmic reasoning
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We show that length generalization of Transformer models trained from scratch strongly correlates with the simplicity of the true RASP-L program for the task.

Abstract: Large language models exhibit surprising emergent generalization properties, yet also struggle on many simple reasoning tasks such as arithmetic and parity. In this work, we focus on length generalization, and we propose a unifying framework to understand when and how Transformers can be expected to length generalize on a given task. First, we show that there exist algorithmic tasks for which standard
decoder-only Transformers trained from scratch naturally exhibit strong length generalization. For these tasks, we leverage the RASP programming language (Weiss et al., 2021) to show that the correct algorithmic solution which solves the task can be represented by a simple Transformer. We thus propose the RASP-Generalization Conjecture: Transformers tend to learn a length-generalizing solution if there exists a short RASP-L program that works for all input lengths. We present empirical evidence to support the correlation between RASP-simplicity and generalization. We leverage our insights to give new scratchpad formats which yield strong length generalization on traditionally hard tasks (such as parity and addition), and we illustrate how scratchpad can hinder generalization when it increases the complexity of the corresponding RASP-L program. Overall, our work provides a novel perspective on the mechanisms of length generalization and the algorithmic capabilities of Transformers.

Published as a conference paper at ICLR 2024
WHAT ALGORITHMS CAN TRANSFORMERS LEARN? A STUDY IN LENGTH GENERALIZATION
Hattie Zhou St, Arwen Bradley', Etai Littwin', Noam Razin Omid Saremi', Josh Susskind', Samy Bengio', Prectum Nakkiran
i Apple Mila, Université de Montréal 'Tel Aviv University
ABSTRACT
Large language models exhibit surprising emergent generalization properties, yet also struggle on many simple reasoning tasks such as arithmetic and parity. In this work, we focus on length generalization, and we propose unifying framework to understand when and how Transformers can be expected to length generalize on a given task. First, we show that there exist algorithmic tasks for which standard decoder-only Transformers trained from scratch naturally exhibit strong length generalization. For these tasks, we leverage the RASP programming language (Weiss et al., 2021) to show that the correct algorithmic solution which solves the task can be represented by simple Transformer. We thus propose the RASP- Generalization Conjecture: Transformers tend to learn a -engtt-generalizing solu- tion if there exists short RASP-L program that works for all input lengths. We present empirical evidence to support the correlation between RASP-simplicity and generalization. We leverage our insights give new scratchpad formats which yield strong length generalization on traditionally hard tasks (such as parity and addition), and we illustrate how scratchpad can hinder generalization when in- creases the complexity of the corresponding RASP-L program. Overall, our work provides a novel perspective on the mechanisms of length generalization and the algorithmic capabilities of Transformers.
INTRODUCTION
Large language models (LLMs) have shown impressive abilities in natural language generation, reading comprehension, code-synthesis, instruction-folloiingg commonsense reasoning, and on (Brown et al., 2020; Chen et al., 2021; Chowdhery et al., 2022; Lewkowycz et al., 2022; Gunasekar al., 2023; Touvron al., 2023). However, when evaluated in controlled studies, Transformers often struggle with systematic generalization (Nogueira et al., 2021; Ontañón et al., 2022; Dziri al., 2023; Wu al., 2023; Saparov et al., 2023). It is thus not clear how to reconcile Transformers' performance in some settings with their fragility in others.
In this work, we aim to understand the factors that determine standard decoder-only Transformer's ability to generalize systematically. Recent studies have focused on length generalization on algo- rithmic tasks as a measure of how well language models can learn to reason (Nogueira et al., 2021; Kim et al., 2021; Anil et al., 2022; Lee et al., 2023; Dziri et al., 2023; Welleck et al., 2022: Liu et al., 2023). Length generalization evaluates the model on problems that are longer (and harder) than seen in the training set, and is used as proxy for whether the model has learned the cor- rect problem-solving strategy for the given task. There is currently scattered evidence regarding the length generalization capabilities of Transformers. Standard Transformers trained from scratch on addition and other arithmetic tasks exhibit little to no length generalization (Nye et al., 2021; Nogueira al., 2021; Lee al., 2023), and even models finetuned from pretrained LLMS struggle on simple algorithmic tasks (Anil et al., 2022). Dziri et al. (2023) propose that Transformers only solve tasks via "analogical pattern matching" instead of learning the true algorithm, and thus will not generalize robustly. On the other hand, length generalization can occur for particular architectural choices and scratchpad formats (Jelassi et al., 2023; Kazemnejad et al., 2023).
Work done while interning Apple.

---

Published as a conference paper at ICLR 2024
1.0
RASP-L Program Exists
No RASP-L Program
0.8
Task Name
Test EM
Task Name
Test EM (+10
0.6
(+10
length)
length)
Count
100%
0.2
Mode
100%
Mode (hard)
0%
0.0
Unique Copy
96%
Repeat Copy
0%
Addition (new)
100%
Addition
0%
Length Training Length
Parity (new)
100%
Parity
50%
Len=20
Len=40
Len=60
Len=80
(a) Transformer length generalization
(b) Length generalization on count
Figure 1: (a) A selection of tasks studied in this paper partitioned by whether they can be solved by programs the RASP-L programming language (discussed in Section 4). Test EM denotes the exact match accuracy (EM) for test inputs of length 10 greater than train. We show that all tasks which admit a short solution in the RASP-L programming language also exhibit strong length generalization performance, and vice which admit RASP-L solutions, by carefully modifying the input and scratchpad format, and these versions length generalize. We also show how poor scratchpad formats make tasks harder, by giving an example for the Mode task. (b) Length generalization for the counting task (described below). Transformers are trained on sequences of varying length, and tested at different levels of out-of-distee distriiibbttin over the maximum training length. Models trained on sequences of length 60 or more exhibit near perfect length generalization up to length 150 (max evaluation length).
For certain hard tasks construct "new" versions
As a starting point of our work, show that there exist algorithmic tasks where Transformers trained from scratch generalize naturally outside of the training distribution. This observation suggests that length generalization is not inherently problematic for the Transformer architecture, though clearly does not occur for all tasks. Why then do Transformers exhibit strong length generalization certain tasks and others, and what are the mechanisms behind generalization when occurs? In the following sections, we will propose unifying framework to predict cases of successful length generalization and describe the possible underlying mechanisms. We discuss additional related works in Appendix A.
Preview of Length Generalization. We begin by introducing simple task that exhibits strong length generalization. The task is "counting": given a prompt SoS a b for numbers a, the model must count from a b inclusive, and terminate with "Eos". An example is: Ko3 We train Transformer with learned positional embeddings on count sequences lengths up to 50, with random and points in [0, 155]. This trained model then length-generalizes perfectly when prompted to count sequences of length 100 (see Figure lb).
SoSS
Possible Mechanisms. is helpful to first consider: why should length generalization be possible all? The crucial observation is that the Transformer architecture already equipped with a natural notion of length-extension. If we omit positional encodings for simplicity, then a fixed setting of Transformer weights defines a sequence-to-seeuuncee function on sequences of arbitrary length. If this function applies the correct transformation for inputs of any length in the context, then we can expect length generalize.
For the count task, length generalization is possible if the model somehow learns a correct algorithm solve the count task. One such algorithm is follows. To predict the next token:
1. Search for the most-recent SoS token, and read the following two numbers a, b.
2. Read the previous token as If (x==!>'), output If (x==b), output EoS.
3. Otherwise, output (xr 1).
This program applies to sequences of all lengths. Thus, if the model ends up learning this program from short sequences, then it will automatically
to long sequences. The discussion so far could apply any auto-regressive model, not just Transformers. What is special about the
2

---

Published as a conference paper at ICLR 2024
Transformer architecture and the count task, though, is that Transformer can easily represent the above program, uniformly for all input lengths. This is not trivial: we claim that the same exact Transformer weights which solve the task at length 20, can also solve the task at length 50, and all greater lengths (Figure la). In fact, we provide evidence in Appendix C that the trained models are actually implementing this algorithm.
The main message our work, inspired by the above example, is that is actually possible for Transformers to approximately' learn length-generalizing algorithm, if the correct algorithm is both possible and "simple" to represent with a Transformer. To reason about what is "simple" for Transformers, we leverage the RASP programming language (Weiss al., 2021; Lindner et al., 2023), which is essentially an "assembly language for Transformers." (We describe RASP, and our extension of it RASP-L, Section 4). We then propose the following toy model of learning.
Toy Model (RASP Learning). For symbolic tasks, Transformers tend to learn the shortest RASP-L program which fis the training set, if one exists. Thus, if this minimal program exists and correctly length-generalicn generalizes, then will the Transformer
We find this toy model to be a useful conceptual tool for predicting when Transformers will length- generalize, and is consistent with our experiments on many tasks. Although many prior works have conjectured similar "simplicity bias" in Transformers (Abbe et al., 2023; Bhattamishra et al., 2023), the notion of simplicity we use here is mulor-made for the Transformer architecture: by construction, each line of RASP can be compiled into most layer of Transformer. Moreover, this is a notion of simplicity over programs, rather than over functions with fixed input dimension. Our main phenomenological conjecture, described in Section 2, is more precise version of this toy model. On the theoretical side, we also give an example where the imin-degree-inteepplaooo model of learning from Abbe et al. (2023) does not produce correct predictions for Transformers, but our conjecture does (Appendix H).
The RASP perspective thus helps unify our understanding of length generalization: instead of de- veloping specialized tools to study generalization of cach individual task (such as addition, multipli- cation, parity, with and without scratchpads, etc), we can now apply a generic hammer to all tasks. To predict whether a task is likely to generalize, the RASP Conjecture tells us we should first see can be solved by a RASP-L program for input distributions. While we focus on studying length generalization in this work, we expect the intuitions developed here can apply to systematic generalization more broadly.
MAIN CONJECTURE
We describe main conjecture. In this conjecture, and throughout the paper, consider Transformers "trained to completion," " meaning trained to near-optimal performance on their training distribution. That is, assume that in in-distribution generalization is achieved nearly-optimally, and focus our attention on the induced out-of-distribution generalization. The exact training procedure consider is given in Section 3.
A key ingredient we develop is a restriction of RASP which we call RASP-L, that we describe in more detail in Section 4. For now, suffices to say RASP is human-readable programming lan- guage which defines sequence-to-sequenee programs that can be compiled into Transformer weights, such that each line RASP roughly corresponds to one Transformer layer. Our conjecture follows.
Conjecture. A decoder-only autoregressive Transformer is likely to length- generalize when trained completion on an algorithmic task if the following conditions hold.
1. Realizability. The true next-token function for the task can be represented by a single decoder-only Transformer which works on all input lengths.
2. Simplicity. This representation is "simple". ". meaning it can be written in RASP-L (a leamn able subset of RASP defined in Section 4).
We do not claim the model will be exactly equivalent to the correct algorithu on all possible inputs: this is unlikeley for technical reasons- numerical precision, noise in the training and sampling procedures, finite optimization time, However, we hope for strong approximations.
3

---

Published as a conference paper at ICLR 2024
3. Diversity. The training data is sufficiently diverse, such that there does not exist any shorter RASP-L program which agrees with the task in-distribution but not out-of out-of-distrinee distriibtbut..n.
The first condition of realizabilitty is actually quite stringent, because it requires a single Transformer to be able to solve the task all lengths² Decoder-only Transformers (also known as causal Trans- formers) define a particular computational model, and not all sequence-to-sequeee solved within this model. For example, next-token functions which require (2(n")) computation time on inputs of length n provably cannot be represented by a Transformer, however large- a conse- quence of the Time Hierarchy Theorem (e.g. Arora & Barak (2009)). The realizability condition may seem stronger than required, because in practice, we do not actually need length generalization for arbitrary unbounded lengths- only lengths up to some maximum context size. Nevertheless, we find that considering representability in the unbounded length setting is a good heuristic for learn- ability in bounded length settings. Intuitively, input length, then it may be an "unnatural" task for Transformers, and unlikely to generalize well.
tasks can be
task requires different Transformer for each
We emphasize that our conjecture is primarily phenomenological, as opposed to mechanistic. That is, we do not make strong claims about which exact function the Transformer learns, and whether or not it is equivalent to a RASP-L program. Our main conjecture characterizes which tasks Trans- formers are likely to sength-generalle
on, and not why how they do so. Although we believe the toy model is a plausible mechanism that implies our conjecture, we leave investigating this more fully important question for future work.
Limitations of Scope and Strength. We acknowledge that our Main Conjecture not fully for- mal, because there are aspects do not fully understand. For example, we cannot precisely predict the extent of length generalization for different tasks. Moreover, since it is likely intractable to de- termine the minimum RASP-L program that fits given training set, cannot predict priori what forms of "data diversity" are required to ensure strong length generalization, even if our conjecture holds true. Nonetheless, we view our conjecture of Transformers, as it has more predictive power than many prior theories. For example, in Ap- pendix H we give simple theoretical setting where the popular "miin-degre--interpolatoo" model of learning from Abbe et al. (2023) does not correctly predict Transformers' out-of cittof--ditributtinn be- havior. but our conjecture does. Developing more formal and precise conjectures an important question for future work.
step forward in understanding the implicit bias
EXPERIMENTAL SETUP
A Transformer (Vaswani et al., 2017) refers to a decoder-only causal Transformer architecture with constant depth. width, and fixed setting of weights, along with any computable positional embedding scheme?. As technical point, we allow the transformer weights to take values in the extended real line R U {+: to allow saturating the softmax at arbitrary context lengths". We consider only greedy sampling throughout, since our tasks are deterministic.
We train all of our models to convergence on the train distribution where possible. For all tasks, the length of training examples is sampled uniformly from length 1 up to the max training length. We train Transformer models from scratch and use learned positional embedding on all tasks. At train time, we "pack the context", filling the Transformer's 's context window with multiple indepen- dent samples of the task, and we randomly shift the Transformer along its context window. This procedure of packing and shifting the context mirrors standard practice in LLM training real data (Karpathy, 2023; Brown et al., 2020), but is typically not done in prior works using synthetic tasks", It is an important detail: packing and shifting the context allows all positional embeddings to be trained, and encourages the transformer to treat all positions symmetrically. At test time, we
Thii corresponds to uniform model of computation, in the terminology of computational complexity. See Merrill & Sabharwal (2023) for discussion of uniformity in the context of Transformers.
"This is a technical detail: we consider position encoding schemes which can he uniformly generated, i.e. there exists a Turing machine which on input (i, m), produces the positional embedding vector for index out of total.
*This bypasses the limitations presented in Hahn (2020), which exploit non-saturating softmaxes.
'One exception this is Liu al. (2023), which performed similar type of shifting.
4

---

Published as a conference paper at ICLR 2024
evaluate examples without packing and shifting. We measure the exact match (EM) all outputs, which is 1 the entire output sequence is correct, and otherwise. Full details in Appendix B.
RASP-L: WHAT ALGORITHMS CAN TRANSFORMERS LEARN?
We will now define version of the RASP programming language (Weiss al., 2021), which we call RASP-L, to heuristically capture the set of algorithms which Transformers can both represent and learn. To apply our conjecture, we need determine if a given task can be solved in RASP-L. Since RASP-L programming is fairly non-intuitive,, we present a "standard library" of useful RASP- functions, which one can compose to solve more complex tasks. We also discuss which types of operations are hard or impossible in RASP-L, to give intuition about which programs cannot be representable or learnable by Transformers.
4.1 RASP: A PRIMER
The original RASP language can be thought of as domainn-peecifi--laggaaa for specifying Trans- former weights, in human-readable form (Weiss et al., 2021). Importantly, RASP was designed for the computational model of Transformers, short RASP programs define functions which are "easy to represent" for Transformers. Although RASP was conceived as a separate language with its own syntax, we also realize RASP restricted subset of Python where only a few operations allowed. We show how to do this explicitly in Appendix F.1, and just include a few examples here. Every RASP program accepts an imput sequence of length n, for all n E N, and returns an output sequence of the exact same length just like a Transformer. The core operations allowed in RASP are: arbitrary elementwise operations over sequences (map and seq_map), and very particular type of non-elementwise operation kqv, which simulates a causal Attention layer. Moreover, no control flow is allowed: all programs must be straight-line programs, with no branching or loops.
For example, suppose we want to write a causal RASP program which always outputs the second-to-last token of the input sequence, e.g- in Python: def f (x) return x [--]] To represent this function in causal sequence-to-sequence manner, we need to take in the sequence x and output the same sequence but shifted by 2. In pure Python, the function is lambda to maintain the same dimension. This lambda function can be implemented in RASP as: lambda x: kqv findices (x) +2, indices (x), x, equals). This RASP program uses the fact that positional embeddings can be specially constructed such that shift-by-two is an elementwise operation on these embeddings. Then, kqv function simulates an attention layer where the current position indices (as query) are matched to the shifted-by-2 indices (as key), and applied to input sequence x (as values).
[0]*2 x2 + x[2:1, where we pad the shifted with 2 tokens in the front
Crucially, since we study autoregressive decoder-only Transformers, we must use the causal version of RASP, where all seq-to-seq operations are executed causally. Moreover, while RASP programs define sequence-tor to-seqeence functions, we interpret them as requence-to-tokes functions, by taking the last token of output sequence as the next-token prediction (in the standard autoregressive man- ner). This setting differs from most prior literature on RASP, which typically consider non-causal models, and these differences significantly change the nature of RASP programming.
Intuition. The intuition to takeaway from this example is that RASP only allows parallelizable operations, because Transformers are an inherently parallel model of computation. This makes performing inherently-sequentia computation, such iterating through each input symbol and up- dating an internal state, tricky not impossible to write in RASP. This is why loops are not allowed in RASP: because Transformer has only constant depth, and cannot directly simulate an arbitrary number of loop iterations. The one way of bypassing this limitation is to exploit the autoregressive inference procedure: since the model is called iteratively at inference time, this effectively provides an "outer-loop" that can enable sequential computation which uses the input context as the state. elaborate in Section 5.
This is exactly what scratchpads enable,
4.2 RASP-L: LEARNABLE RASP
The original RASP technically allows for certain operations which are possible to represent, but not "easy" represent or learn. For example, arbirarily-complex tokenwise operations R - IR
5

---

Published as a conference paper at ICLR 2024
1.0
Unique Tokens
Test Accuracy att Langm 50
1.0
0.8
Test EM at Length 50
0.8
Repeat Tolonns
0.0
Le
Test DM
0.6
0.6
Scratch
Lene40
W/ Scretch (final)
Len-50
0.4
0.4
W/ Scratch (EM)
0.2
Len-70 70
0.2
0.2
Len-B0
40
10
Additional Length over Max Trin Langth
Max Train Length
Max Length
(a) Length generalization mode
(b) Copy task
(c) Mode task (w/ scratchpad)
Figure 2: (a) Length generalization performance for the mode task. All models generalize perfectly to 10 or more additional tokens at test time. (b) Performance on copy tasks for test sequences of length 50, and varying train lengths. Copying unique tokens makes length generalization much easier, since can be solved by an induction head. (c) Performance on mode task for test sequences of length 50, and varying training lengths. The scratchpad hurts in this setting (both final answer accuracy and exact match), since not computable by a short RASP-L program.
are allowed. To disallow such pathologies, we define a "learnable" subset of RASP which we call RASP-L. The technical details of RASP-L are in Appendix F.2, but the primary restrictions are: all variables are bounded integers 8) to avoid and token indices are treated specially. Roughly, token indices in RASP-L can only be operated on in simple ways (order comparison, predecessor, successor)- arbitrary arithmetic involving indices is not allowed. We find that empirically, such operations involving index-arithmetic are often not learned robustly, as we elaborate on Appendix F.2.
and numerical stability issues,
A RASP-L Standard Library. In Appendix F.3.2 we provide a small library of useful functional- ity built the RASP-L core. This also serves as representative sample of functions that are "easy" to implement RASP-L, to build intuition about programming causal Transformers. The library in- cludes where, which is functionally equivalent branching (though inefficiently, since all branches are always evaluated). It also includes the fairly self-explanatory maxima maximum, minimum argmin n and argmax. Recall that all functions are seq- to-seq and causal, so for instance maximum returns the array of running maxima. The functions num_prev, firsts, and induct all take an input sequence x and a query sequence q, and re- turn, respectively, the number of previous elements of equal to each q [. the index of the first occurrence of each [ii] in x, and the token following the first occurrence of o [i] in x.
numpy. where, and allows for rudimentary
4.3
CASE STUDIES
In this section, we experimentally evaluate 3 tasks that have simple RASP-L programs and 3 tasks that do not. We show that RASP-L representability correlated with length generalization per- formance. The three easy tasks we consider are: count, mode, and copy with unique tokens. We provide the RASP program for these tasks in Appendix F.3 (Listings 3, 4, and respectively). Detailed training procedures and hyperparameters are provided in Appendix B.
Count. We described the count task in the Introduction and showed results in Figure 1b. This task can be solved by a RASP-L program that essentially translates the pseudocode in the Introduction (defails in Listing 3). We find that models trained count can generalize near perfectly to double the training lengths. It is crucial that our training distribution contain samples ranging from length to maximum training length, which adds diversity necessary in preventing shortcut programs from being learned: no generalization is observed if we train sequences of all the same length.
scale up training length. These factors are
Mode. The mode task identifies the most frequent element in sequence. We constrain the quences such that the answer is unique. An example is: a b shows the results on mode, when training and testing random sequences from alphabet of 52
c b a c b
>|b]. Figure 2a

---

Published as a conference paper at ICLR 2024
symbols. We find that models trained on mode generalizes strongly to sequence lengths much longer than the maximum training lengths. Interestingly, increasing training set complexity here does not result in huge improvements in length generalization: Even models trained on max sequence length of only 10 can achieve a median test accuracy of 50% of sequences of length 60.
Copy with unique tokens. The copy task repeats the prompt sequence in the output. We constrain the sequences to have all unique tokens in the prompt. An example is: |
Figure 2b shows the results on copy with unique tokens. For models trained on sequence length up to 40, find that they can generalize perfectly to length of 50. Intuitively, this task is easy because we can leverage something called an "induction head" (Olsson et al., 2022). Induction heads work by identifying a previous instance of the current token, find the token that came after it, and predict the same completion to the current token. Olsson et al. (2022) found that induction heads are reliably learned even by simple Transformers, and conjectured them to be component of what enables in- context learning in LLMs. Induction heads are simple to implement in RASP-L, as the induct function. Thus, the next token can be generated by simply using an induction head on the current token, since all tokens are unique. This is exactly what the RASP-L
does, in Listing 5.
Next, we identify three tasks that do not admit simple RASP-L solutions: addition, parity, and copy with repeating tokens. We discuss reasons why Transformer models struggle to generalize on these tasks by highlighting the operations these algorithms require, but which are unnatural for a Transformer to represent.
Addition & Parity. Both these tasks have been studied as difficult tasks for Transformers: mod- els trained from scratch show little et al., 2023) and parity (Bhattamishra et al., 2020; Chiang & Cholak, 2022; Ruoss et al., 2023; Delé- tang et al., 2023), and even pretrained LLMs cannot solve these tasks robustly (Brown et al., 2020; Chowdhery et al., 2022; Anil et al., 2022) without careful prompting (Zhou et al., 2022b).
length generalization on addition (Nye et al., 2021; Lee
Indeed, addition is also difficult to write in RASP-L. To see why, consider the standard addition program shown in Appendix F.3.6. This algorithm requires the carry value to be propagated in reverse order from least- to most-significant digit, but this is difficult simulate due to causal mask- ing. Moreover, the most prohibitive aspect is the index-related operations. The standard addition algorithm requires -rit-eei (e.g. finding the middle of the prompt sequence) and precise indexing operations (e.g. look up the corresponding summand digits for the current output digit). Such operations are forbidden in RASP-L, as they require index- -aith-arithmetic which are difficult to rep- resent in a global, genggh--gneralzingggg way (see Appendix F22 for more discussion). Similarly, parity without any scratchpad requires operations that are forbidden under RASP-L. Since a Transformer cannot update its state sequentially, it must solve parity with parallel operations only. Intuitively, this requires taking the sum of the entire sequence, then determining the parity of the sum. This cannot naturally be computed in numerically stable way for arbitrarily large sums, and we cannot expect to learn a 'sum" operation which generalizes to numbers larger than the training sequences. Indeed, many works have shown that Transformer cannot even fit the training set of parity sequences over some minimal length (Bhattamishra et al.. 2020; Chiang & Cholak, 2022). Under our experimental settings, we find that no length generalization is observed for both addition and parity tasks. We evaluate different test lengths in increments of 5, and all 20 runs on all training lengths up to 45 showed a test EM of approximately 0% on addition and parity.
Copy with repeating tokens. For this task, constrain the sequences to consist only of 2 pos- sible tokens. An example is:
>|a a a Since the tokens are no longer unique, the induction- head is no longer helpful. Instead, the model must perform precise index-arithmetic, which is difficult for the same reason that indexing is difficult in addition. We show in Figure 2b that models completely fail to generalize to longer lengths this task.
APPLICATION: IMPROVING LENGTH GENERALIZZAION
In this section, we demonstrate how our RASP conjecture can go beyond post-hoc explanations, by constructing interventions that predictably change length generalization performance. We study how reformatting tasks to allow shorter RASP-L programs can improve generalization performance, and how inereasing diversity in the training data allows the model to perform well on tasks that require more complex RASP-L programs.
7

---

Published as a conference paper at ICLR 2024
1.0
1.0
Forward EM al Length
0.8
Reverse EM at Lengghh
0.8
Text EM at Length 50
0.0
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
30
Max Train Length
Max Train Length
Max Train Length
(a) Forward-order addition with
(b) Reverse-order addition with
(c) Parity with scratchpad
index hints
index hints
Figure 3: Length generalization on addition and parity. Plot shows 20 individual trials per length, as well as their median (solid line). (a) Shows generalization performance for forward addition with index hints on hard carry examples of length 50. No length generalization is observed. (b) Shows the generalization of reverse addition with index hints on hard carry examples of length 50. Most runs start length generalize perfectly 50 digit addition once the training length is greater than 40. (c) Shows generalization performance for parity with scratchpad, on length 50 inputs. Most runs start to generalize perfectly on 50 digit addition once the training length is greater than 35.
5.1 DEEP DIVE ON ADDITION
Reducing the RASP-L complexity for addition. In the previous section, we noted two aspects of naive addition algorithm that pose problems for RASP-L: anie--aith (to query the summand digits for the current output digit), and non-causality (for the carry). To address the difficulty with indexing operations, we can leverage induction heads to simplify the addition algorithm for Trans- former by adding "index hints" to the prompt and answer: For example, 5 > 9 becomes P This enables us to get the corresponding digits for each sum step by calling induct on its index hint (a or b), thus sidestepping the need to precisely access and manip- ulate positional indices. To address non-causality of the carry operation, we can format the output in reverse-order. For example, a/ This enables simple and causal propagation of the carry where each step can reference the previous output to determine the current carry, similar to the standard addition algorithm. A RASP-L pro- gram for reverse-order addition, with index- is provided answer digits greatly simplifies the carrying procedure, it is still possible implement algorithm for addition in forward order (Listing 8). This algorithm is nontrivial, because we essentially need propagate a carry through all input digits in order determine the first digit of output, the most significant digit (see Figure 12). We show how to construct a RASP-L program for forward-order addition in Listing 8. Comparing Listing 7 and Listing 8 reveals how much more complicated the forward algorithm is is- it results in a much longer RASP-L program.
a 5
becomes
l
Listing 7. Although reversing the
Index hints enables generalization on addition. We evaluate addition in two settings: "easy" carry and "hard" carry. In easy carry, the two numbers are sampled randomly and independently this is what is typically done in the literature. However, uniformly random summands will only produce addition instances with short carry-chains (in expectation)- and for such instances, each output digit only depends on a small number of input digits. We thus also test "hard" carry instances, where constrain the examples have the longest possible carry chain for the given length. For example, hard carry instance of length 3 is 381+619 = 1000, which requires the model to compute the carry over chain of 3 digit positions. The performance on "easy" carry is shown in Figure 8, and the performance on "hard" carry in Figure 3. We find that index hints allow both forward and reverse addition to length generalize on "casy" carry. However, on "hard" carry questions that involve carry- chains longer than seen at training, reverse addition maintains strong length generalization while forward addition exhibits no generalization.
Diversity enables generalization forward addition. Another lever for performance improve- ment suggested by the RASP conjecture is to increase training data diversity, such that shortcut programs can longer fit the training set. Since forward addition does admit RASP-L program,
8

---

Published as a conference paper at ICLR 2024
albeit a more complex one, we would expect it is possible to learn if we "try harder," e.g. use more careful and diverse train distribution. We explore this by training with balanced carry sampling instead of sampling the two numbers independently, we first sample the length of the carry chain uniformly between O and question length, then sample a random question that contains the given earry chain length. This ensures that the model sees a significant percentage of questions containing long carry chains, thus increasing the diversity and difficulty of the training data. The results of the balanced carry training approach for both forward and reverse addition are shown in Figure 9. We see that this more careful training unlocks the model's ability to length generalize on forward addition, even under the hard carry evaluation. To our knowledge, these results demonstrate the first instance of strong length generalization on decimal addition for Transformers trained from scratch.
5.2
WHY DO SCRATCHPADS HRLP?
The RASP conjecture provides a natural way to understand why scratchpads (Nye et al., 2021; Wei et al., 2022) can be helpful: seratchpads can simplify the next-token prediction task, making it amenable to a short RASP-L program. One especially common type of simplification is when a scratchpad is used to "unroll" loop, turning next-token problem that requires n sequential steps into n next-token problems that are each only one step. In the following examples, we construct "good" scratchpads which simplifies the target RASP-L program, and "bad" scratchpads which seem natural to humans but require more complex RASP-L program than the original task. We show that these interventions lead to predictable changes length generalization.
Seratchpad enables generalization on parity. We leverage the "unrolling" intuition to design a scratchpad for parity. Similar to addition, add index hints to the prompt to simplify the indexing operation. In the scratchpad output, we locate index hints that precede each 1 in the prompt, and keep track of the running parity with symbols + (even) and (odd). The last output token corresponds to the final answer. For example: || performance of the proposed parity scratchpad. We see that some of the runs trained with sequences up to 30 in length can generalize perfectly on sequences of length 50. When training length reaches 45, all models achieve perfect length generalization on length 50. These results demonstrate the first instance of strong length generalization on parity for Transformer models trained from scratch.
| + Figure 3c shows the exact match
Scratchpad hurts generalization on mode. Now we consider the mode task and look at how scratchpad might affect a task that a Transformer is naturally amenable to. A natural algorithm one might come up with to calculate the counts of each unique token in the sequence, then out- put the token with the maximum count. To encourage the model to learn this algorithm, we might utilize the following scratchpad, where we output the frequency of each token in ascending order: bbb. The last token in the scratchpad is then the correct answer. How- ever, although this scratchpad provides more supervision for what algorithm the model should learn, is more difficult task when considered in terms of RASP-L. Finding and comparing the frequency is simple to do internally, but converting this implicit representation into an integer token adds ad- ditional complexity. We show in Figure 2c that the seratchpad performs significantly worse than no scratchpad, both when measured on exact match and also on the accuracy of the final answer.
abblcb c
CONCLUSION
We proposed a model for understanding and predicting when Transformers are likely to exhibit strong length generalization on algorithmic tasks. We do so by reasoning about the unique information- constraints of the Transformer architecture through a variant of the RASP lan- guage. We conjecture that algorithms which are simple to represent by a Transformer are also more likely to be learned, and use this model to predict and improve length generalization on set of algorithmic tasks. Our toy model demystifies certain observations of strong reasoning and out -ut-of-distribution abilities of Transformers, and shows that they can occur for potentially simple reasons. Although our studies focused on length generalization, the intuitions are not specific to it, and they lay out path towards a possible mechanism for systematic generalization. Studying these mechanisms beyond length generalization, and their interaction in multi-task settings, are important directions for future work.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We thank (alphabetically) Samira Abnar, Madhu Advani, Jaroslaw Blasiok, Stefano Cosentino, Lau- rent Dinh, Fartash Faghri, Spencer Frei, Yejin Huh, Vaishaal Shankar, Vimal Thilak, Russ Webb, Jason Yosinski, and Fred Zhang for feedback on early drafts and discussions throughout the project.
REFERENCES
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization on the unseen, logic reasoning and degree curriculum, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15660 2022.
Cem Anil, Yuhuai Wu, Anders Johan Andreassen, Aitor Lewkowyez, Vedant Misra, Vinay Venkatesh Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. In Advances Neural formation Processing Systems, 2022. URL https:/ http:///peeeeeeeeeeeeeeeee wttpsvviwwwe/fforrrn 2SkYVeX7bC4.
Sanjeev Arora and Boaz Barak. Computational complexity: a modern approach. Cambridge Uni- versity Press, 2009.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009. 11264, 2020.
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901,
François Charton. Linear algebra with transformers. arXiv preprint 2 2021.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2 03374, 2021.
Ta-Chung Chi, Ting-Har Han Fan, Alexander Rudnicky, and Peter Ramadge. Transformer working memory enables regular language reasoning and natural language length extrapolation. arXiv preprint Stti:2305.03799 2023.
David Chiang and Peter Cholak. Overcoming a theoretical limitation of self-attention, 2022.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
Antonia Creswell and Murray Shanahan. Faithful reasoning using large language models, 2022.
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, and Pedro A. Ortega. Neural networks and the chomsky hierarchy, 2023.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, al. Faith and fate: Limits of transformers on compositionality. arXiv preprint arXivv2305. 18654, 2023.
Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. How self-attention networks recognize dyck-n languages? In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4301-4306, 2020.
10

---

Published as a conference paper at ICLR 2024
Daniel Furrer, Mare Zee, Nathan Scales, and Nathanael Schärli. Compositional generalization in semantic parsing: Pre-traimin training vs. specialized architectures, 2021.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
Suriya Gunasekar, Yi Zhang. Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah, Harkirat Singh Behl, Xin Wang. Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need, 2023.
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
Samy Jelassi, Stéphane d' Ascoli, Carles Domingo- Enrich, Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers, 2023.
Lukasz Kaiser and Iyy Sutskever. Neural gpus learn algorithms. arXiv preprint arXiv:1511.0822 2015.
Andrej Karpathy. nanogpt. h http://cw...///e/ptttt//aannnnnn
2023.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers, 2023.
Jeonghwan Kim, Giwon Hong, Kyung-min Kim, Junmo Kang, and Sung-Hyon Myaeng. Have you seen that number? investigating extrapolation in question answering models. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2021.
Andrei N Kolmogorov. On tables of random numbers. Sankhyā: The Indian Journal of Statistics, Series A, pp. 369-376, 1963.
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers, 2023.
Aitor Lewkowycz. Anders Andreassen, David Martin Dohan, Ethan S Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with language models. 2022. URL https: http//rxvvvvvv//ee//.aa abs 2200 14858.
Yuxuan Li and James McClelland. Representations and computations transformers that support generalization on structured tasks. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https:: / /pwprevvee.weetteeeeeeeeeeeeeee.
David Lindner, János Kramár, Matthew Rahtz, Thomas McGrath, and Vladimir Mikulik. Tracr: Compiled transformers a laboratory for interpretability. arXiv preprint arXiv:2301.00062, 2023.
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata, 2023.
Aman Madaan and Amir Yazdanbakhsh. Text and patterns: For effective chain of thought, it takes two tango, 2022.
Sadegh Mahdavi, Kevin Swersky, Thomas Kipf, Milad Hashemi, Christos Thrampoulidis, and Ren- jie Liao. Towards better ditto--distibbttion generalization of neural algorithmic reasoning tasks. Trans. Mach. Learn. Res., 2023, 2022. URL https: /appi semant iectttcccchlllcc /org/ Corpus 253255355.
William Merrill and Ashish Sabharwal. The Parallelism Tradeoff: Limitations of Log-Piccision Transformers. Transactions of the Association for Computational Linguistics, 11:531-545, 06 2023. ISSN 2807-387X. doi: 10.1 162/tcll a 00562. URL https: http:://///ooo/////e/.///oo 1162/ tacl 000662.

---

Published as a conference paper at ICLR 2024
William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-8555 856, 2022.
Rodrigo Nogueira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks, 2021.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint 1rXiv:2122.001144 2021.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfiel-Doddss Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kermion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. http://tannsffmmmeecrrruitttatt wttp://p////-nn-ottttteeeeeeeeee context-learning wawphng--nnnidndud munductiion-eads//ndee
Santiago Ontañón, Joshua Ainslie, Vaclav Cvicek, and Zachary Fisher. Making transformers solve compositional tasks, 2022.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, (1):))4333349777 2021.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL
ht https: wpwnreview.eeeeeoo net
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Ben- nani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. Toronto, Canada, 2023. Association for Computational Linguistics.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. arXiv preprint 2rXiiv2200.0220 2023.
Abulhair Saparov, Richard Yuanzhe Pang, Vishakh Padmakumar, Nitish Joshi, Seyed Mchran Kazemi, Najoung Kim, and He He. Testing the general deductive reasoning capacity large language models using ood examples, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.
Josef Valvoda, Naomi Saphra, Jonathan Rawski, Adina Williams, and Ryan Cotterell. Benchmark- ing compositionality with formal languages. In Proceedings of the 29th International Conference on Computational Linguistics, pp. 6007-6018, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz. Kaiser, and Illia Polosukhin. Attention you need. Advances in neural informa- tion processing systems, 30, 2017.
Petar Veličković and Charles Blundell. Neural algorithmic reasoning. Patterns, 2(7):100273, 2021.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers, 2021.
Sean Welleck, Peter West. Jize Cao, and Yejin Choi. Symbolic brittleness in sequence models: systematic generalization in symbolic mathematics. In AAAI, 2022. URL https: / : /arxiv.
brg//df/21099 13986.90 pdf.
12

---

Published as a conference paper at ICLR 2024
Zhaofeng Wu, Linlu Qiu, Alexis Ross, Ekin Akyürek, Boyuan Chen, Bailin Wang, Najoung Kim, Jacob Andreas, and Yoon Kim. Reasoning or reciting? exploring the capabilities and limitations of language models through counterfactual tasks, 2023.
Shizhuo Dylan Zhang, Curt Tigges, Stella Biderman, Maxim Raginsky, and Talia Ringer. Can transformers learn to solve problems recursively?, 2023.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schu- urmans, Olivier Bousquet, Quoe Le, and Ed Chi. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205. 10625, 2022a.
Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning, 2022b.
13

---

Published as a conference paper at ICLR 2024
A ADDITIONAL RELATED WORKS
Our paper is related to the line of work that seeks to understand the capabilities and limitations of Transformer models when it comes to algorithmic reasoning (Kaiser & Sutskever, 2015; Veličković & Blundell, 2021). Specifically, we focus on simple tasks like arithmetic and study length general- ization on the standard Transformer architecture. Related this, Lee et al. (2023) study how well transformers trained from scratch can learn simple arithmetic tasks, and finds that no length gen- eralization is observed. Nogueira et al. (2021) find that partial length generalization on addition is observed only when models reach 3B parameters and when the addition questions are presented in reverse order. Jelassi et al. (2023) study models trained on addition and find strong generalization performance when using few examples of longer sequences. However, they required non-standard architectures (non-causal Transformers) and training procedures (artificial padding) to observe this, and still find that the model does not generalize unseen length that is in between the minimum and maximum lengths seen during training.
Moreover, our work contributes to the study of what makes for effective scratchpads. Other pa- pers have also found that using positional tokens in the prompt can help with length generalization (Nogueira al., 2021; Li & McClelland, 2023). However, these works do not provide a framework for understanding why these tricks are helpful. A number of papers also study how chain-of-thought style prompting helps with reasoning performance (Wei al., 2022; Zhou et al., 2022b;a; Creswell & Shanahan, 2022; Madaan & Yazdanbakhsh, 2022), but these focus on in-context learning and do not study the effect of training models on these formats.
Other papers also aim to understand the limits of what Transformers can learn and represent. Bhat- tamishra et al. (2020) and Delétang et al. (2023) study the ability of Transformers to represent and generalize on families of formal languages. Delétang et al. (2023) evaluated Transformers on tasks at all levels of the Chomsky hierarchy. and found that "Transformers and LSTMs are not well- aligned with the Chomsky hierarchy." This is consistent with our results, since the Transformer computational model does not directly correspond to a natural level in the hierarchy. Valvoda et al. mandomly-generatee
(2022) benchmarked Transformers state transducers, and evaluated an even stronger metric than maggth-generalization (compositional generalization). They found that Transformers only generalize on certain tasks, and only when the training data has sufficient "coverage" of the task. As an aside, we chose not to use randomly- generated tasks in our setting. because our target program space much larger than Valvoda et al. (2022), so our experimental results would be depend very strongly on the exact sampling procedure for random RASP-L programs. For example, a naive "random RASP-L program" is likely to al- ways produce the constant O output, due to degeneracies. Zhang et al. (2023) evaluate the ability of transformer models to emulate the behavior structurally recursive functions from input-output examples. Liu et al. (2023) study how shallow Transformers can simulate recurrent dynamics rep- resentable by finite-state automata. Both works identify shortcut solutions that become brittle on out-of-distributie samples. To address these shortcomings, various works have proposed modifica- tions or alternatives to the Transformer architecture in order to improve generalization, such as Press et al. (2022); Chi al. (2023); Mahdavi et al. (2022).
tasks produced by deterministic finite-
Bhattamishra et al. (2023) suggest that Transformer models have an inductive bias towards learning functions with low sensitivity, such as sparse boolean functions, but focus on the in-distribution setting. Abbe et al. (2023) also propose a simplicity bias in Transformers, but use "minimum- degree" as their notion of function simplicity. However, they only consider functions with fixed input dimension rather than programs on arbitrary input lengths.
Lastly, there have been many other approaches to improving length generalization in Transformers. These include studying how various training hyperparameters and design choices influence com- positional generalization (Furrer et al., 2021; Ontañón al., 2022), and designing better positional embeddings (Press et al., 2022; Ontañón et al., 2022; Kazemnejad et al., 2023; Ruoss et al., 2023).
B ADDITIONAL EXPERIMENTAN DETAILS
For all experiments, we tokenize every character individually, including digits of number. We train in the online setting, where each batch is sampled iid from the train distribution instead of from finite train set this avoids overfitting issues, and is closer to the training of modern LLMss.
14

---

Published as a conference paper at ICLR 2024
Unless otherwise specified, evaluate test performance on 5x the batch size number of samples. Unless otherwise specified, we run each experiment on 20 random seeds and report the median of the runs. We select hyperparameters for each task based what is required to fit the training set. Hyperparameter details can be found in Table 1.
Count. For the count task, we train with an alphabet size of 155 and evaluate on test sequences up to 150 in length. Given the nature of the task, we enumerate all possible sequences of each test length at evaluation time. The alphabet is ordered such that the sequence between any two tokens in the alphabet is clearly defined. At train time, the length of each example is sampled uniformly between 1 and the maximum training length.
Mode. For the count task, we train on an alphahet of 52 tokens. Each example consist of 5 unique tokens sampled randomly from the alphabet. The length of each example is sampled uniformly between 1 and the maximum training length, and the sequence is randomly sampled from the 5 selected tokens. If there is a tie for the most frequent element, we randomly select a token from one set and changes it to a token from the other set, thus ensuring that there is one unique answer.
Copy. For the copy task with unique tokens, we train on an alphabet size of 100. The length of each example is sampled uniformly between 1 and the maximum training length, and the sequence randomly sampled from the alphabet without replacement. For the copy task with repeat tokens, we use the same sampling procedure, but now on an alphabet size of 2.
Addition. For the addition task, we sample the length of each of the two numbers independently, from up to the maximum training length. We then pad the two numbers with 0 in the front such that they have the same length. We pad the numbers with an extra 0 to allow for the potential of an extra digit in the answer due to carry. For addition with index hints, we sample the index hints as random slice from longer contiguous block of tokens, to encourage learning all hints and their linear ordering. This is similar to our training strategy for the count task.
Parity. For the parity task, we sample the length of each parity sequence from up to the maxi- mum training length. We then sample randomly from {1,0) sequence of the given length. We note that the definition of length we use is based on the sequence length and not based on the number of Is in the sequence.
Experimental hyperparameters. All experiments use Adam W optimizer and cosine learning rate schedule. Count and Copy use weight decay of 0.1 and grad clip of O. Parity and Mode use weight decay of 0.1 and grad clip of I. Addition uses weight decay of 0 and grad clip of 1.
Table I:
Task
Model Size
Train Iter
Context Len
Batch Size
Learning Rate
Count
6 layer; head; 64 cmb
10000
256
128
lc-3 to lc-5
Mode
6 layer; head; 512 emb
10000
256
128
le-3 to le-6
Copy
6 layer; head; 512 emb
100000
512
128
le-4 to le-6
Addition
6 layer; head; 512 emb
30000
512
61
le-4 to
Parity
6 layer; 8 head; 512 emb
10000
512
256
le-3 to 1c-6
C APPENDIX: COUNTERFACTUAL ANALYSIS ON COUNT
In this section, we probe whether models trained on count actually learn the count algorithm that we intuitively want. To reiterate, one simple algorithm that solves the count task is as follows. To predict the next token:
1. Search for the most-recent SoS token, and read the following two numbers as a,
2. Read the previous token as If (x==1>> >!!, ") output a. If (x==b), output Eos.
3. Otherwise, output (xz + 1).
15

---

Published as a conference paper at ICLR 2024
1.0
1.0
Accuracy of Start-on-FFig on--Fist
Accuracy End-on-Last
0.6
0.6
0.4
0.4
length
length
OOD Length
0OD Length
0.0
0.0
120
160
40
80
120
160
Max Train Length
Max Train Length
(a) Counterfactual test of starting on first token (b) Counterfactual test ending on last token
We measure performance of models trained on the count task on counterfactual tests designed to evaluate whether the model simulates the correct generalizing algorithm on random sequences far aut-of-distribution. We see that models always output the start token the first output, but do not always output the EoS token once the ending token has been outputted.
Figure 4:
Since there is no easy way to formally check if the Transformer model learns this exact algorithm internalyy, we employ the simple heuristic of running the model on counterfactual examples. The allows us to stress test the models behavior and see if it performs the expected algorithmic steps in an input-agnostic way. To do so, we performance inference on randomly generated input sequences that are designed to test the model in four ways:
1. The model should always output the start token in the prompt (a) as the first output. (Figure 4a)
2. The model should always output Eos following token that matches the ending token in the prompt (b). (Figure 4b)
3. In all other settings, the model should increment the previous token by 1. (Figure 5a)
4. The model should not output Eos prematurely. (Figure 5b)
We create the counterfactual dataset by sampling start and end tokens of varying distances, then generate a sequence of random tokens of the length specified by the distance between the start and end token. We then pass this sequence through a trained model and look at its predictions at each token position. The goal of the four proposed tests random sequences is probe whether the model learned the expected algorithmic behavior rather than learning something that would strongly depend statistics of the training distribution. We sample examples for in-distribution lengths and oottof-distribution lengths based on the training distribution of each model. For simplicity, we choose 1 model with strong length generalization performance from each maximum training length setting. The performance on each test is shown in Figure 4 and Figure 5.
We that for the start-on-first test and the increment-by-1 test, all models exhibit near perfect per- formance both in- and out-of-distribution. For the end-on-last test, we see that models trained with shorter lengths do not learn robustly output Eos on long test sequences once the ending condi- tion is met. However, models trained on longer sequences (and has better length generalization), this behavior more robust. Lastly, when we measure the percentage of Fos which are correct, we see that models that do not have strong generalization also fails to output EoS only at the ap- propriate time. This failing is observed on both in-distribution and out-of-distribution lengths. This suggests that the failures of length generalization can be attributed to prematurely outputting Eos token before the full length of the sequence outputted. Overall, we observe strong correspondence between the model's behavior and what would expect from the correct algorithm. This lends credence to the intuition that the model learns the correct RASP-L program and generalizes because of it.
16

---

Published as a conference paper at ICLR 2024
1.0
1.0
Eos Outputs
Accuracy Increment
0.6
0.6
0.4
0.4
length
IID length
0OD Length
0OD Length
0.0
0.0
40
120
160
80
120
160
Max Train Length
Max Train Length
(a) Counterfactual increment-by-1
(b) Counterfactual test percentage EoS outputs which are correct
Figure 5: We measure performance of models trained the count task counterfactual tests designed to evaluate whether the model simulates the correct generalizing algorithm on random se- quences far out-of-distribution. We see that models almost always increments the previous token by 1, no matter what the previous sequence is. However, it sometimes output the EoS token pre- maturely, especially on lengths longer than seen in training. This likely explains failures of length generalization observed in Figure 1b.
D APPENDIX: TRAINING SPEED
The RASP Generalization Conjecture suggests that simple-to-represen programs are also more likely to be learned. Intuitively, we should expect that simple-to-represent programs also exhibits faster optimization during training. In this section, evaluate different strategies of presenting the addition and the parity task, and see whether the training speed of these variants correspond to the simplicity of their corresponding RASP-L programs (and by extension their length generalization performance).
As discussed in Section 5, reverse addition simplifies the function that the model needs learn when compared to forward addition. Figure 6a shows the training curves for each of these settings. Consistent with intuition,
that reverse addition converges more quickly than forward addition.
For the parity task, we introduce an additional scratchpad format for comparison. This scratchpad outputs the sum-mod- 10 of the parity sequence before outputting the final parity. An example is 3 2 0 This scratchpad does not simplify the problem as much as the main scratchpad presented Section 5 because it does not leverage the autoregressive inference pro- cess the task sequentially. However, it is still simpler than parity without any scratchpad because it helps to simplify the final operation getting the parity the sum. Instead of doing this internally, the model can now reference the output of the sum-mod-le 10 and learn a simple mapping between that and the corresponding parity. Figure 6a shows the training curves that the main scratchpad ("Easy Scratchpad") optimizes much more quickly than the sum-mod- 10 scratchpad ("Hard Scratchpad"). We also observe that Easy Scratchpad exhibits significantly stronger length generalization than Hard Scratchpad, shown in Figure 7b. Both seratchpads opti- mizes much better than parity with no scratchpad, which is unable to even fit the training set and demonstrates no length generalization.
each of these settings. We
E APPENDIX: ADDITIONAL ABLATIONS
In this section, we include additional experiments to support the results in the main paper.
In Section we introduced a scratchpad for mode, which orders the intermediate counts in order of frequency. This may overly demanding, it requires the model to know the order of the frequencies before outputting them. Another variant of this could output the scratchpad in order of appearance in the original sequence. Moreover,
output the token first before outputting
17

---

Published as a conference paper at ICLR 2024
1.0
1.0
Train EM (Max Len-45) 451
0.8
0.8
0.4
0.4
EM
Easy Scratchpad
0.2
Reverse Addition
0.2
Hard Scratchpad
Forward Addition
No Scratchpad
0.0
0.0
5000
10000 15000 20000 25000 30000
2000
4000
6000
8000
10000
Train Iteration
Train Iteration
(a) Training speed comparison for addition with index (h) Training speed comparison for parity with different scratchpads
hints
Figure 6: We compare the training speed as measured by exact match on maximum training length of 45. (a) compares the convergence speed of models trained on forward addition vs reverse addition with index hints. (b) compares the convergence speed of models trained on different scratchpads on parity.
Test Accuracy Length 50
1.0
Scratch
1.0
Scratch (Final)
0.8
0.8
Scratch (EM)
EM Length
0.4
0.2
0.2
0.0
0.0
10
20
30
30
35
Length
Length
(a) Mode using scratchpad ordered by appearance
(b) Parity with sum-mod-10 scratchpad
Figure 7: (a) compares performance of mode with without scratchpad. In this case, we the scratchpad presented in order of appearance. We see that no scratchpad significantly outperforms scratchpad, whether measured on the final answer accuracy or the exact match of the entire scratch- pad output. (b) illustrates the generalization performance for parity with scratchpad on length 50. We see that no runs show significant length generalization in this setting.
their count, which may help the model reference the relevant token for this step. example b
b
The performance of this scratchpad is shown in Figure We see that utilizing this scratchpad still results in much worse length generalization performance than using no scratchpad.
Lastly, we showease the spread of different training runs on the count and mode tasks in Figure and the copy tasks in Figure 11.
F APPENDIX: RASP DETAILS
F.1 RASP SPECIFICATION
Here we describe how to realize RASP as a restricted subset of Python (with numpy). First, every RASP program accepts input sequence of length for arbitrary E N, and returns output sequence of the exact same length- just like transformer. The restrictions on Python are: All variables either numpy arrays of length ("sequences"), binary matrices in {0, 1 jaxn ("se-
18

---

Published as a conference paper at ICLR 2024
1.0
1.0
Forward Length
0.8
Reverse Length
0.8
0.6
0.6
0.2
0.2
0.0
0.0
35
40
35
40
Length
Length
(a) Forward addition on easy carry examples
(b) Reverse addition on easy carry examples
Figure 8: Length generalization addition with index hints. Each diamond shows the performance of one of 20 runs, illustrating the spread of different training runs. (a) illustrates the generalization performance for forward addition with index hints easy carry examples of length 50. (b) illus- trates the generalization performance for reverse addition with index hints easy carry examples of length 50. Easy earry examples consist of addition questions where the two numbers are sampled randomly and independently, which is the setting considered in prior works studying addition. We see that both settings demonstrate strong length generalization, thus demonstrating the usefulness of the index hints.
Easy Carry
Easy Carry
Hard Carry
Hard Carry
0.8
0.8
EM for Length
EM for Length
0.6
0.6
0.4
0.4
0.2
0.2
10
20
30
40
10
20
30
40
Max Train Length
Max Train Length
(a) Forward addition trained with balanced carry
(b) Reverse addition trained with balanced carry
Figure 9: Length generalization on addition with index hints, trained with balanced carry distribu- tion. Each point shows the median test performance over 20 runs. (a) illustrates the generalization performance for forward addition with index hints on both easy and hard carry examples of length 50. (b) illustrates the generalization performance for reverse addition with index hints both easy and hard carry examples of length 50. With balanced carry training, strong length generalization is observed for both forward and reverse addition on the hard carry evaluation setting. This demon- strates that with increased data diversity, we also increase the likelihood that the model learns solution.
lectors").. No control flow allowed: all programs are straight-line programs, with no branching or loops. Finally, every line of the program must be call to either one of the core functions defined in Listing I, or to another RASP program.
it easy confirm that this is equivalent to the original presentation of RASP in Weiss et al. (2021).
F.2 RASP-L: LEARNABLE RASP
The original RASP technically allows for certain operations which are possible to represent, but not "easy" represent learn. For example, any arbitrarity-complex tokenwise operation IR - IR
19

---

Published as a conference paper at ICLR 2024
1.0
1.0
0.8
0.8
Test |Length
EM Leegth
0.2
0.2
0.0
0.0
20
40
60
80
100
120
140
40
80
100
120
140
Length
Length
(a) Test EM on Count for Length 150
(b) Test EM on Count for Length 150
Figure 10: (a) illustrates the generalization performance for count with varying training length on test examples of length 150. (b) illustrates the generalization performance for mode with varying training length test examples of length 150.
1.0
1.0
Test Length
EM Length
0.4
0.4
0.2
0.2
0.0
0.0
Max Train Length
Max Train Length
(a) Test EM on Unique Copy Length
(b) Test EM Repeat Copy for Length 50
Figure 11: (a) illustrates the generalization performance for unique copy with varying training length on test examples of length 50. (b) illustrates the generalization performance for repeat copy with varying training length test examples of length 50.
is allowed. To disallow such pathologies, define "learnable" subset of RASP which call RASP-L. RASP-L enforces the following additional restrictions on RASP.
First, all non-index values, including the input and intermediate variables, must be of type i This constraint handles issues with both infinite-precision and with learnability of the token- wise operations (since all tokenwise functions now have small finite domains and co-domains, int8 + int 8, and so can be easily memorized). Although disallowing floating-point opera- tions and unbounded numbers may seem like strong restriction, find these are necessary for most symbolic programs.
Moreover, token indices are treated specially. Standard RASP allows arbitrary arithmetic operations on indices, such as division-by-two: + [i/2]. However, transformers must decode index infor- mation from their positional embeddings. so any RASP-L operation on indices must be "easy" to perform on the appropriate positional embeddings. We find that empirically, such operations in- -engtthgeneralizing manner. This potentially
volving -ari--aiithmetic are often not learned because it is difficult to learn globally valid arithmetic structures (such division by two) from sentially local examples, iee. short contexts. To reflect this RASP-L, we only allow the following operations indices: order comparisons with other indices, and computing successor/predeccsssor Formally, the RASP-L core function indi ces (x) returns a special type IndexInt. which can take values in N, but only allows these restricted operations. That is, allow adding to
20

---

Published as a conference paper at ICLR 2024
def nap 1x, zunel:
(tunc for xi
def (x Vr cuncl:
return mp.artay( (tune yi for xia yi 13
sclect ik, 3. pred):
5 (e)
for app.zeoo((Ca s), deygerbool)
Linddx anidexx co causazity
for 3 in (1-11
M[i, grediklt), qUAII
def sel :AIt
zeturn dot (A, Bp. cp.cnes (Len
aggs v):
ouz dop (A, V)
norn sci_vidth(N)
divide
out-ap. aeo-p Kheze-
def indicer (K)
zeturn arange (ken
(x,l oonal))
return aryars aasarryleennss. [Iconat] 00)
Convenivnce lonz
kge 4, V. pred):
return agge isclect (s, 9, pred), vI
Listing RASP-L core functions.
IndexInt, but do not allow adding two IndexInts, casting between IndexInt and int8.
There is one additional restriction RASP, involving the "selector width" operation. Selector- width in standard RASP returns the number of prior elements that are selected by a binary Atten- tion matrix, for each token position. The return type of Selector-width in RASP-L inherits from IndexInts thus it can represent unbounded numbers of selected elements, but can only operate on them in restricted ways. Moreover, every call to sel_ sel_width inherits from IndexInt, and these types cannot be compared to each other. That is, the sequences returned by two different calls sel ww dth are incomparable. The reason for these restrictions, which may otherwise seem contrived, is that sel_widte can be used to simulate indi ces, by calling it on the all-ones selector matrix. Thus, we must restrict the output sel_width suf- ficiently not allow bypassing the intended restrictions index-arithmetic. There may also be more mechanistic motivations for such restrictions, since the Transformer implementation of selec- tor width requires weights which grow linearly with sequence length (Lindner et al., 2023).
RASP-L returns a new type which
F.3 RASP PROGRAMS
In this section, we provide the RASP-L core (1) and library (2); RASP-L programs for the tasks discussed in the paper, namely counting (3), mode (4), copy-with-unique-toke (5), addition with reverse-order and index-hints (7); and naive non non-RASI RASP addition algorithm (6).
21

---

Published as a conference paper at ICLR 2024
reverse-order
forward-order
3
6
8
1
8
+
6
3
4
5
+
2
6
3
4
5
9
9
12
8
3
?
Intermaciaale Veriableo
Intermediate Variables
output
2
?
.aWe.** output
current
compute
current carry depends on subsequeet inputs output known
previous output:
diract carry prav digit sum
indirect canyy i prev digit sum = prev
output
Figure 12: Intuition for why reverse order addition is simpler implement in RASP than forward order. In reverse order, the carry calculation is simple with the help of the most recent output digit, which acts like built-in scratchpad. In forward order, the carry requires more complex calculation involving all remaining input digits. This intuition reflected in the RASP-L program for forward addition, which is much longer than that of reverse addition (see Listings 7 and 8).
22

---

Published as a conference paper at ICLR 2024
def vhere (candicion, X_if, y_Hlse):
equivalent ap.ahice cconiiiinn
y_else]
x mested seq nap if, y_meakedd seg.neply_et (ylse,,,,, condition fan, lanbda nt y 4f not wlse 0)
condition, Lambda x, nt else
return seo (K_nasked, y_nasked, lambda
LE =* C else
makk maakk, nask_val=();
it equivaient to ssosoL_ass oetauitr/-oool ma56) return where (boo)
fill mast_val)1
def cursuu arrayy)
returns nimber of previacy True eienents in bool array
return sel widhh fseleot (bucl_atsa)) boal_atray, lambda k, k)
shift rhightright Is, default=0)s
3 fiw seguence to the right by postl lons return (indlcees)))
indices(x),
setault-detaulit)
def maximun_xY v)1
returns ost w/argnexrx, (rrppertx____eeee_ Jie 2127 wh which.st select num_at cei_wisthhwhhihhhpp
highest (num gt
idx_hale hold curaun (uraun (higgheat) zero order hold of the number of the highnss Incides so Far. tox_hoiddtt highttt eequnap hold, highest, lambda 1, else kqp return 2
highest, hold, equalas
def meximm(())
return moximun_xY (X,
def ()))
-maximun(-x)
def crgnax 181
return meximun_xV IX, insiccsix)
argnia
return argmax(-X)
def
querics):
output [d] number presious elements equai to queries(2j, inciusive return sel_widch (K) quaries, equals))
def tirstsix,
decouit-01:
find index i out 1n ap. fistnonzero (xt:1+1] [ulrios[1]).nit .iia() VAL -- special coktn, cannct appear in x bas_prev kevishlttrrigt rirst_oce 1-has_prev
first occurrence of quozyji7
fillix,
equals) xf1] cccured provicusiy
fesst_oce_only where ifisst_een, 8, fi1l1s, VULL_VAL) z ouc xf3] i2 fiest_oooff] else VAL zetura koy (cleat__oce_nnn,, queries, indices (6,, equals, defaaltideraoit)
flad_las:__ok (x, tok):
1 finds the indox of the jast occurronce of tok in sequence x (couselly, iaclusive) toke f:1l18, tok)
apсey [e1_xiithh(sseee (N, toks, aquais)]
nprov_ms:cher_onte onlly mask mask inprev, matches) zero-sut noa-toks idxs kay (npcov_natcies_onll, oprav, indices(Mt) equals)
def solex (s, ids, defaclt-D1:
i indoxes into sequense x, via inder scquence iox stide) if ids[]] else default zeturn kay (indlcess))) (0), idx, & equals, dafaiit-dafault)
def first defaull=-1)s
first ssenctrue kgvis, fill D), fillis, 1), equals, defauls=0)
Loeet_ooe seggnaptaeen shift right (sven 1), lanbda ousr, prev cuse and not pewv) zeturn (firattooo fill (6, 1), indicesia), equals, dofaciz
def induct ik, q. offset, defauit=D, vai-
i velow of index of first cocurrence of Gif Zound) plus offeet special cokea (ehift
aull
aspeaz offset, defaul vall,, & default-nal) vall)
provent accidental natohes
indices tol nogy I capy valuee of k at indiowe La cupy fase requested defauli for invalid Indices) cspind_values indox_select ik, indices_to_oopy, deault-defacl1) copied_
Listing 2: RASP-L library functions.
23

---

Published as a conference paper at ICLR 2024
ooant (beg):
First, find tnd index of nost-recent START_IOK (begninng of curzent sequence)
start idx find last toklseg, START_ TOR)
It Then, corpute the start/end nunbers of the current sequence
start rune index select lseg, start start_idxe))
indox_oclocttsse
Rool arrayer whecher we're predicting ing the first / last cakens the current sequence
cissttost seq_nap (Lndices (C2g), stast_idx+2, equals) pred_tinei_pos i-pred_tirst.
sec_naplseg,
noxt_iok vtore [p:oo_
if predicting the first token:
next tox starting nun
Where tpere
predictine final
(so), END_ TOR),
acxt__ok
Beg 11)
eiser nexc_tok tok prey tok
zeturn next_tok
Listing 3: RASP-L program for Count.
def node (K)
sun_pprvv_nachiing (scleet oquais))
ids gagnaa/munn.ieynaaahhaoo intox_peloct 1X, idx)
def binary_noda sal fillix, equals))
pun_prov_zaree
pre_ppee_poees sil (select be, fillix, 1h, eguelal) aode_val aeq aeg_nappnur pnun_prov__nes,, pan_proy__ezos,, gu1+1 zeturn moda_val
Listing 4: RASP-L program for Mode.
def copy ar irue (seg):
return induct (seg, offset=l)
copy or ique ar(x):
START, END 1,
((
seq ()
while eegi-1] 1- ESD:
tox copy (negl [-1] seg np. ap..conatonnte (1sog, [
return seg
511)
Listing 5: RASP-L program for Copy with unique tokens.
def and illlgallinng
zero-padoed digits xO,
sur_dicits int (1an(1pp)/22 z np.zeros (nem_aigite) carsy
index cypes
raversed_range
inan_digital [1s-1] IsasCal: reversal
seversed__angee
JELECALI Joops
*0, x1 inplil, inp Laglnun_digits111 IILBOAL; variabies canact bo used as Indicea dici-_sun xo xi xii] tsigit_sun cazzy digit_aun
return 2
Listing addition program that is illegal in RASP-L for several reasons.
24

---

Published as a conference paper at ICLR 2024
#2 Constanta and holper Fuactions
2
STA87_28ONET -1
FLUS
EQUALS_SIGN
END END_RESE RESPONSE ---
XONE -88
7
mask_potveen_tokmm
tekt, tokl]:
seer_iokō has_sean (208, full (seq, taxC]1
10
szor_iok1 tas_sean(sent, full (seq, taxl)1
ind_between reg_nsp(econ__oRO szen_toki, lanbda and not ind(tok3) ind(to
int_betwson
def sx, y1! 14
indox-hinta
17
* -okent function
18
def wtw___oorrroddoooooooooooooo
prompi_neck 1-has_seen (oeq, (ull, EQUALS_SIGN)
vecond_ mak. Iseq, PLUS, BQUALS_AIGN) prompt mazk (seq, prompt_nask]
22
fatta first other_rurmns__doos
the avanand with second. induct pr-ooop,, riahiftt_rittttptooo 1), offset-1)
pairsans seq (seg, at auher_unmeni :lgght, _atki_aae; # chis wligns paicouve with the 2nd sommen pairsans nask (pairaun, secenn__urman__nna HCNE)
pairsuns_sh (psi(paissnn,,
HONE)
hints: only dieits
precooctptt$iii sight_right (seg, 1)
30
curr_pzirsun induet (pairsuns, abift_right (seq, 2), offset=l) f pairsun that goacrated juzi_outpuc_diqie pest_pairsun induel (pairsuns, seg, offsel=))
START CRANCES
34
direci_casty curs_pairmu
previoue sum givea lauzt_pupput_digit
Indirect_carzy (surr_palzaur text_tok_.get__arr direct_cerry incireot__arry
earlier cazzy
2 (sinple) indox-hint conputacions: 38
Final_Hint (swg, 1003 finei first_iist induct fawg, (seg, BQUNL8_AIGN), offsei- 2) fixat hint plases before naxt_hinz shift_sight (sea, 11 was [next_ hint hintl_hint) CHANGES
alweys
42
next_iok mest_pairrsmm next_iok 1- toxt__ok cout__ook_gets_aerry
Finally, handle the case of outout Ing index-hinte hints -100) 8 al) index-hints -100
next_.tok in__iisiidees hint
next_tok_iia_nnde
next tok wheret next..to i8. index hint, next hint, next tok) nexi_lok wheret fail (awy, END (RD__NR) nexi_too) next_tok return next_.ook
(seg EQUALS, SIGN), tirst_hint, next_zok!
Listing 7: RASP-L program for addition, with output in reverse order, and index-hints. See Sec- tion 5.1 for details on prompt format. For addition in forward order, the highlighted codeblock is replaced with Listing 8.
25

---

Published as a conference paper at ICLR 2024
CHANGES
gives_carry tot_mapipairremm_. nh, lambda _Xt it else 0)
cumsunCipai (ipai 1-
nh NONE)
macklz, girce_corty, mack_val SONE:
v tok_napits it,, lambda _X 10
chain, end idd_ Firsta(s, deFAult-HOMR)
(left) anding indices of cerry-ohain
corr_tot_gor_cene next_.oo tok inside darry, chain (next pei raum 91 cirr toc to...gtt corrrotttrr.........
!- cure wotpottttttpppppeeeeeeee _iggtt
10
in the middle of cazry-chain? (NOTB: assumes the paizsume has rirst clenert 0)
noxt_iok_idx karipairsums, seg, Andicesiseg) equals;
which ensver-position are не st? (indices sligned to pairsunsj
14
next next_tok_chela shd kqyl easin_end_ idxe tokt_tok.__ tok_idx tull (eeq, 31, equale, detsult=0)
carsy-shain?
next_cot_eess_coe tesf_tot_tncdddccerry..o
next_rok_chetn_tng
18
t (simpic) index-bint computations:
Hinsi_Hint insuct_prer fiest_hint (uld, -100) next_ainz shict_ (scg, 11 eos Inext_ tintl_hiat)
tull EQUALE_SIGN) ottset--21 rinsi nint is piaces botore for
22
CHANGES
Listing 8: The required patch to Listing 7, to produce program for addition in forward order. This code replaces the highlighted block in Listing 7. See Section 5.1 for details on prompt format.
26

---

Published as a conference paper at ICLR 2024
G LIMITATIONS AND OPEN QUESTIONS
Limitations of RASP-L. The RASP and RASP-L programming languages were remarkably use- ful in our work, for reasoning about which algorithms are easy to represent and learn. However, they have clear limitations. For one, RASP and RASP-L are not complete: not all functions which can be efficiently represented by Transformers have efficient RASP implementations. notable class of Transformer algorithms which are ill-suited for RASP are numerical algorithms, such as in-context gradient-descent and linear regression, which involve high-dimensional matrix and vector floating- point operations (Akyürek ct al., 2022; Garg et al., 2022: Charton, 2021). Morcover, RASP supports only deterministic outputs and binary attention masks. In addition, some the RASP-L restrictions may seem somewhat contrived (such a more natural way to capture the set of rasseyblepreeentable algorithms for standard Transformers. Nonetheless, we view RASP and RASP-L as important steps towards reasoning about Transformer algorithms, and we hope to see future work in this direction.
floating-points and index restrictions) there may be
Limitations of Complexity Measures. We considered for simplicity a basic notion of function complexity, which the minimum RASP-L program length. This intuitive notion captures many empirical behaviors, as we demonstrated. However, depending on the operations, it can be possible represent multiple lines of RASP-L in single layer of a Transformer, and so RASP program length does not perfectly correspond to Transformer cor.ffomeeccomlexittt.... There are likely more refined notions of complexity, such as the "parallel depth" of RASP-L programs, or lower-level measures like the minimum- woight-nn among all weight-noise-robune function. Many of these notions, including ours, have the drawback of being likely intractable to compute- intuitivly, it may be difficult to find the minimum RASP-L program for task for similar reasons that Kolmogorov complexity is uncomputable (Kolmogorov, 1963). We leave investigating more refined notions of complexity to future works.
Transformers which implement the
On Formal Language Characterizations. One important question that remains open is whether there exists a natural complexity-theoreii definition of the class of tasks which are "simple" for autoregressive Transformers to represent. For example, it is well-known that RNNs tend to length generalize on tasks equivalent to regular languages, like Parity (e.g. Delétang et al. (2023)). This intuitively because regular languages nite automata) which are "simple" to represent by RNNs, and thus plausibly easy to learn. We would like analogous characterization of which tasks admit algorithms with simple and natural Transformer representation. Recent works have characterized which functions are possible to rep- resent by Transformer", but this representation is not always "simple" enough be learnable, and not always uniform (Hahn, 2020; Merrill et al., 2022; Chiang & Cholak, 2022: Pérez et al., 2021; Bhattamishra et al., 2020; Ebrahimi et al., 2020; Merrill & Sabharwal, 2023; Sanford et al., 2023). Our presentation of RASP-L is meant to be one way of defining algorithms which are "simple" to represent- those expressable as short RASP-L programs- but this definition is not explicit (and likely not complete).
be decided by a class of algorithms (deterministic fi-
H COMPARISON TO MIN-DEGREE-INTERPPLAAN
An essential aspect of our work is our Transformer-specifie notion of function complexity: the minimum RASP-L program length. Here we show why this choice is important, by contrasting it with another popular notion of complexity: minimum polynomial degree. Concretely, Abbe et al. (2023) recently proposed model of learning in which Transformers learn the minimun-degree function which interpolates their train set. We will give a simple example where our RASP toy model correctly predicts a Transformer's dustof-distribution generalization behavior, but the min- ingree-ineerooator model does not. We emphasize that these results are not inconsistent with Abbe et al. (2023): neither Abbe et al. (2023) nor our current work claim to apply in all settings. Rather, this section illustrates how a Transformer-specite measure of complexity can be more predictive than arcitecttue-agnostie measures of complexity, in certain settings.
'Note that the exact set of which functions are representable depends on certain detinitional details of Transformers such as finite infinite precision, bounded vs. unbounded weights, etc, which is why some of these references arrive different conclusions.
27

---

Published as a conference paper at ICLR 2024
H.1 THE SETTING: BOOLEAN CONJUNCTION
setting of Abbe et al. (2023), for direct compar- ison. Our target function is simply boolean conjunction. Given n 20 input bits 24 E {0, 1), the ground truth function is the boolean AND of all bits: (21,,
We consider the "Gennerlizatiiononnitt-----me
Fn. = Aeets) 24 That is, f (z) 1 iff all bits 2; = 1, and f (a: 0 otherwise. Now, the train distribution supported on inputs x where the last k = bits of x are always 1. Specifically: with probability 1/2, x is the all-ones vector, otherwise is all ones with single in a random location among the first 15 bits. That is.
x x~ (~
i~ Unif[1, 15].
where e, E {0, 1)" the ith standard basis vector. Note the ground-truth label y = fx) for this distribution is balanced. The unseen test distribution is identical to the train distribution, except the '0' is only among the last 5 bits. That is,
x. Unseen) (1 B(1/2)ea); i~ 11i 20]-
We now ask: when a Transformer is trained on the above train distribution, what does it predict on unseen test distribution?
Experimental Result. We train standard decoder-only Transformer autoregressively in the above setting, with sequence distribution [z, (x)) for sampled from the train distribution. (Ex- perimental details are given below.) We find that the trained Transformer reaches 100% test accuracy on the unseen test set. That is, the Transformer correctly computes the boolean AND, even on bits which were irrelevant at train time.
Experimental detials: For this task, we train on sequences of length 20 and evaluate on test sequences of length 20. We train 2 layer, 4 head autoregressive Transformer, with embedding dimension 64. We train with Adam for 10000 iterations, using context size of 128, batch size 128, and learning rate le-3. To compare to the setting in Abbe et al. (2023), we do not pack the context here and train on single examples in the context window. The training distribution has 50% chance of being sequence of all 1s, and a 50% chance of having one 0 element in the sequence. The position of the 0 element is sampled uniformly between positions 0 and 15. The last 5 elements in the training sequence are always 1s. At test time, there is a 50% chance of being a sequence all 1s, and a 50% chance of having
element in the last 5 elements in the sequence.
H.2 THE MIN-DEGREF INTERPOLATOR
We now claim that the minimum-degree-inn inteereeeiinerrolatoor of the train set does not behave like the Trans- will not
former in the above experiment. To see this, observe that the minimum-dee ceanuuu--ddgre-inerppllaa depend on the last k 5 bits of the input, since these bits are constant on the train set. This can be formalized via the following simple lemma, which uses the same notion of "degree profile" (DegP) Abbe et al. (2023)².
Lemma 1. For all subsets S {0, 1)" and all boolean functions 10, R. the following holds. Let g² : {0, 13" + R be the boolean function of minimum degree-profile which agrees with on S. That is,
g"
arg min DegP(g).
9=00,11--8 s. gls=fls
Let C [n] be the subset of indices (if any) on which S is constant. That is. TI (S), the projection of S to coordinates I, is singleton. Then, the minimun-degee degiiuum-dgrre--ntrppaaaat g" also does not depend on indices I. That is,
x = y.vi 1 9")
(y)-
This lemma follows from the fact that an interpolator depends on bits which are always constant on the train set S, then its degree- could be reduced without changing its output on S, and thus cannot be a con--derre-interpolle
For completeness, we state this formally below.
Briefly, the degree profile of boolean function is the tuple of its Fourier weights each level, with the natural total ordering which refines the standard polynomial degree.
28

---

Published as a conference paper at ICLR 2024
RASP-Length. On the other hand, there is one-line RASP-L program which computes the boolean AND all input bits:
def output (x) kgv (x, full (x, 0), full (x, equals, default=1)
This program can be represented by a I-layer Transformer, where the attention layer simply "searches for 0" in the prior context. Thus, experimental result in this setting. The intuition here is that it is actually "easier" for a Transformer represent the boolean conjection of all input bits, and learns to do so. Learning ignore the last bits would have been possible, but is an additional complication, and does not help in fitting the train set in any case. Notably, this intuition (and indeed, the experimental result) does hold for MLPs; it is specific to the Transformer architecture.
RASP toy model of learning predicts the correct
H.2.1 PROOF OF LEMMA
Lemma 1 follows directly from the following fact. We state this with the boolean hypercube iden- tified with {=1}" as standard in boolean function analysis, but these statements can trivially be translated to {0, 1)"
Lemma 2. Let f [+1|- - IR be a boolean function. Suppose f depends on its first coordinate;
that is suppose 32 E {+1]--1: f(1oz) # f(-1o Then, restricting the first coordinate
to be strictly reduces its degree profile:
DegP (f ==11) < (f)
Proof. (Sketch). Consider the multilinear representation of f, and factor out terms containing æ1
f P() (T)
where E denotes (æ2, Tz. Fn.. When restricted with 1, we have
f(11 = P(z) + Q(x)
(2)
Now, because f depends on 21 by assumption, we must have P * 0. Thus, comparing the above two, at least one monomial containing has reduced degree with the restriction. The conclusion follows.
29


Reviews:
[{"review_id": "LajyCQnoQG", "reviewer": "Reviewer_Ho88", "summary": "This paper proposes to predict the feasibility of length generalization by checking whether the task can be implemented by RASP-L (L for Learnable).\n\n- RASP-L is a constrained version of the RASP language, where all variables are constrained to int8. This excludes certain operations such as index arithmetic (but still allow order comparison, predecessor and successor).\n- The paper shows that the difficulty in length generalization is correlated with the difficulty in writing the RASP-L solution.\n    - Tasks that have a simple RASP-L solution: count, mode, copy.\n    - Tasks that do not have a simple RASP-L solution: addition, parity, copy with repeated tokens.\n- The paper then shows that whether scratchpad is helpful can be explained by RASP-L as well: a properly chosen scratchpad can help with performance (e.g. by making the task solvable with induction heads), but improperly chosen scratchpad, which increases the difficulty of writing in RASP-L, can hurt the performance (e.g. on the mode task).\n\n_Post rebuttal update_: I've increased my score based on the discussions and changes in the revised paper.", "questions": "- Page 2, \"Possible Mechanisms\": it's not true that omitting positional encodings may lead to \"arbitrary length\" generalization, due to failures in both MLP and attention (please see the 4 citations mentioned above).\n- It would be better to include discussions on prior work that relate algorithmic reasoning tasks to formal languages.\n  - For example, for the discussion towards the end Section 2, please consider relating to e.g. Merrill and Sabharwal 2022 and the notion of uniformity in computational complexity (which footnote 2 has mentioned).\n- Fig 3: How many test sequences are there for each trial?\n- Fig 6(b): this is not really about length generalization, but more a failure of optimization?\n\n\nMerrill and Sabharwal 2022: The Parallelism Tradeoff: Limitations of Log-Precision Transformers", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "- The paper analyses different factors affecting length generalization.\n- The paper provides a variety of empirical evidence.", "weaknesses": "- The paper studies length generalization, though it's unclear what is considered as successful length generalization in the paper. [_Update_: the author has added a clarification.]\n    - Fig 1(b) is considered as demonstrating successful length generalization, even though the performance from all curves are dropping, some even to 0.\n    - In Fig 2(a) (i.e. the mode task), the paper seems to consider a test accuracy of 50% as reasonable generalization performance, per the comment in Sec 4.3.\n- The paper proposes to use the difficulty of RASP-L programs as an indicator for length generalization performance. However, I'm not sure what this view can teach us, either theoretically or empirically. [_Update_: the authors modified Sec 5.1 significantly on implications of the RASP-L formulation.]\n    - It's unclear how to empirically verify or act upon this conjecture. It's also not clear how to convert a task into RASP-L except for applying the definition directly, i.e. checking for whether certain criteria are satisfied.\n        - For example, two criteria provided in the paper are 1) whether the operation is non-causal and 2) whether precise index arithmetic is required. Hence one could also say that these criteria are indicative of length generalization performance, without resorting to the RASP-L language.\n    - I'm not sure whether these criteria provide new insights beyond findings in the existing literature.\n        - It is well known that auto-regressive generation cannot uncover non-causal information. [_Update_: previously I misunderstood what \"non-causal\" mean; the authors have clarified this in the rebuttal.]\n        - The concern about index arithmetic is essentially a problem with approximation (e.g. precision and weight norms), which has been raised in prior work. In particular, Hahn 2020, Chiang & Cholak 2022 and Liu et al. 2023(b) discussed limitations in the attention module, and Liu et al. 2023(a) discussed limitations in MLP.\n- Some experiments have been covered in prior work.\n    - The shifted position experiments are also included in Liu et al. 2023(a) (Fig 6 & 7).\n    - The scratchpad experiments are similar to those in Liu et al. 2023(a). This paper uses a different form of scratchpad, but the generalization performance is not very good (e.g. generalizing from length 25 to length 50 gets around 0.3 accuracy), which is related to my earlier comment about what is considered as successful length generalization.\n\n*References*\n\nHahn 2020: Theoretical Limitations of Self-Attention in Neural Sequence Models\n\nChiang & Cholak 2022: Overcoming a Theoretical Limitation of Self-Attention\n\nLiu et al. 2023(a): Transformers Learn Shortcuts to Automata\n\nLiu et al. 2023(b): Exposing Attention Glitches with Flip-Flop Language Modeling\n\n_Note: the paper has cited Hahn 2020, Chiang & Cholak 2022 and Liu et al. 2023(a)._", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cJxaLFHVsq", "reviewer": "Reviewer_tXQ4", "summary": "This paper proposed a conjecture about when a decoder-only autoregressive Transformer is likely to have strong length generalization ability on symbolic algorithmic tasks. There are three conditions: 1) the true next token function for the task is realizable with a *causal Transformer*; 2) the next token function can be represented by RASP-L (a subset of the RASP language) programs; 3) the training data is sufficiently diverse, which makes sure that the shortest RASP-L program can length-generalize. The conjecture states that Transformer is likely to generalize when the three conditions are satisfied. It is shown that the conjecture is supported by some case studies. On 3 tasks which have simple RASP-L programs, Transformer empirically generalize well with a moderate max train length (e.g., >30 on the *copy with unique tokens* task). On 3 tasks that \"do not admit\" simple RASP-L solutions, Transformers struggle to generalize. The conjecture is then applied to interpret why scratchpads can improve generalization. Specifically, on a case where the scratchpad increases the complexity of the RASP-L program, Transformers show decreased generalization performance.", "questions": "- Please refer to the weakness part.\n- Minor question: Empirical results show that the max train length has a great impact on the generalization ability of Transformers. Is the proposed method helpful in analyzing this phenomenon?\n- Minor suggestion: Some important definitions, e.g., that of RASP-L, would better to be included in the main paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- A novel approach for understanding the reasoning and out-of-distribution generalization ability of Transformers on the symbolic tasks.\n- The analysis in the case studies is well supported by empirical results.\n- Experiment details are well introduced.", "weaknesses": "- Some concepts are not clearly defined, making their meanings obscure and the statements less rigorous. For example, there is no reference or definition for the \"causal Transformer\" and \"causal masking\"; how to identify whether a RASP-L program is \"simple\" or \"difficult\".\n- The correlation between RASP-L representable and the generalization of Transformer need further discussions. Are the selected cases representative enough? Is there any task which can be represented by \nRASP-L but Transformers hard to generalize, or vise-versa? These answers are needed to make a stronger statement on the correlation between the two.\n- No discussions on the limitations of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "634UeSHVkO", "reviewer": "Reviewer_q9T5", "summary": "This paper discusses an important problem: in which algorithmic task transformers are more likely to length generalize in which format. \n\nThey propose a conjecture: transformers are likely to length generalize for a task if there exists a \"simple\" \"RASP-L\" program to solve that task. (The conjecture includes other conditions to ensure almost perfect in-distribution performances.) (The \"RASP\" programs are easier to represent and learn by transformers in the sense that their program parameterizations encourage parallel operators and discourage sequential, for-loop, and branch operations. The \"RASP-L\" programs further rule out numerically unstable operations such as they forbid float values and complex token indices processing.)\n\nThe authors also provide some experimental results that empirically align with the conjecture.", "questions": "* How did you choose the \"learnable\" operators? Will the representation power of \"RASP\" programs be affected by this additional constraint? Can all \"generalizable\" tasks' solutions be represented by the \"RASP-L\" programs? \n* Are there results on the difference between the IID and OOD performances for most of the tasks?\n* Are there constraints as almost perfect IID performances before evaluating their length generalization performance? Are there results with metrics such as maximum among all runs?\n* Did you achieve almost perfect IID performances for the addition task with 40 digits with the naive forward order? \n* Can you explain why the length generalization performance of transformers is still limited (e.g., from length<=40 to length=50) even when the conjecture conditions are fulfilled?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "This paper discusses an important topic: the length generalization of transformers. It is easy to understand with many illustrative visualizations. They test the conjecture in six simple tasks whose positive/negative results all align with the conjecture. They also show improved/degraded length generalization performances as guided by the conjecture after reformulating some tasks. \n\nRegarding the concrete conjecture, at least some parts of it align with the reviewer's intuition and experiences. For example,\n* If human beings can implement a simple RASP-L program to solve a task, we do expect better length generalization of learned transformers.\n* Transformers do suffer from numerical issues in practice, either for counting tokens or modeling complex indexing mechanisms. The constraints on \"learnable\" in addition to \"RASP\" are interesting. \n\nThe reviewer also appreciates the simple trick on position embeddings for better length generalization: they concatenate all samples in a long sequence with random shifting while training transformers. This is different from the common choice, which only puts one sample in one sequence, for synthetic tasks. They find that \"packing and shifting the context allows all positional embeddings to be trained, and encourages the transformer to treat all positions symmetrically,\" therefore, towards better length generalization of transformers.", "weaknesses": "The main weakness is lacking the support and validation of the conjecture. The definition of RASP-L programs is also vague and unclear.\n\nRegarding the support and validation of the conjecture: \n* There is no proof or reasoning for the conjecture. There is little evidence saying that transformers actually learned or have any relationships with the \"RASP-L\" program in general. According to the experimental results where \"transformers mostly generalize to length at most 50 when trained with length<=40\", it is also hard to believe that the true \"RASP-L\" programs have been learned. \n* The experimental results and comparison may not be solid with confusing metrics:\n   - The metrics should characterize the differences between OOD and IID, instead of the absolute performances on OOD samples. As shown in Figure 1(a), generalization is about the difference in performances between in-distribution and out-of-the-distribution samples. This is important for, e.g., the addition task, because addition with 40 digits may already be out of the representation power of a 6-layer transformer in the naive forward-order. The IID performances could be as bad as the OOD ones. \n  - There is also a lack of constraints on achieving almost perfect in-distribution performances before evaluating their length generalization. \"Almost perfect IID performance\" is assumed in the conjecture. It is also strange to evaluate the OOD performances of random runs with imperfect IID performances. If this constraint is not enforced, the current analysis may confuse the optimization difficulties with the OOD generalization performances. \n   - Similarly, the metric as \"median among 20 runs\" is not enough to fully characterize the transformers' performance, given their non-convex optimization nature. Other metrics, such as \"maximum\" or \"ratio of almost perfect,\" could be more informative and practical. \n  - The tasks are simple: from count to addition/parity. \n* Ablation studies on \"without the index hints\" are needed. \n\nRegarding the format of the conjecture as \"RASP-L\" programs,\n* The learnable constraints are interesting, but there are no reasoning/proof/experiments to support their choices, such as which operators are \"learnable\" in which sense, how to find and choose them, and if the current set of operators can represent all \"learnable\" operations.\n* Similarly, there is no characterization of the function space represented by the \"RASP-L\" programs. \n* The optimization efficiency conjecture in the appendix is also poorly supported. I would suggest the authors remove it if there are no supporting results from more tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z2BMAvkHTU", "reviewer": "Reviewer_LyN1", "summary": "This paper focuses on the standard decoder-only transformers model for algorithmic tasks. Prior work has shown that standard transformers exhibit little length generalization on certain arithmetic tasks. Accordingly, the author first propose the RASP Conjecture to describe the kind of problems on which the transformers can perform length generalizations. The author also present empirical evidence to support the RASP Conjecture and length generalization. Finally, the author analyzes the hard tasks for transformers, such as parity and decimal addition, and proposes particular scratchpad formats to enhance the length generalization on these tasks.\n\n**Post rebuttal update**\n\nHaving reviewed the author's response, I am impressed by the meticulous effort the authors have done in answering questions and refining the manuscript. The work has effectively addressed my initial concerns. As a result, I have raised my point.", "questions": "See the weakness section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper is well-written: discussions and presentations are clear. \n\n* The RASP Conjecture is well-motivated.\n\n* The case study clearly explains why certain tasks are solvable (i.e., length generalizable) and why others aren't. For each solvable task, there is a corresponding RASP format.\n\n* The author designs particular scratchpad formats for parity and decimal addition and experimentally verifies the effectiveness of these designs.", "weaknesses": "* The scratchpad designs for parity and decimal addition, though effective, are a bit different from the RASP format. For example, it seems that the proposed scratchpad for addition is motivated by correcting the failure modes discussed in section 4.3. It would be nice if the author could clarify the connection between the scratchpad designs and the RASP. \n\n* This paper classifies the algorithmic tasks into easy and hard ones. The easy ones are solvable (i.e., length generalizable) by the standard transformers. As a general question, can the author recommend the possible directions to solve the hard ones? The proposed scratchpad solution requires a detailed understanding of the failure modes for each hard task, which doesn't seem quite scalable.\n\n* Missing citations. Transformers for algorithmic tasks is an interesting topic in the recent literature. It would be nice if the authors could discuss a few more papers. From the modeling perspective, [1] introduces certain twists on the transformer architecture to enable the length generalization on algorithmic/regular language tasks such as Even Pairs, Modular Arithmetic, Parity Check, and Cycle Navigation (i.e., the tasks in [2]). From the task perspective, [3] evaluates the model performance over randomly generated finite-state transduction tasks. The idea of evaluating over randomly generated tasks is complementary to the usual conduct of evaluating on particular tasks.\n\n\n[1] Chi, T.C., Fan, T.H., Rudnicky, A.I. and Ramadge, P.J., 2023. Transformer Working Memory Enables Regular Language Reasoning and Natural Language Length Extrapolation. arXiv preprint arXiv:2305.03796.\n\n[2] Delétang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wenliang, L.K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J. and Ortega, P.A., 2022. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.02098.\n\n[3] Valvoda, J., Saphra, N., Rawski, J., Williams, A. and Cotterell, R., 2022, October. Benchmarking compositionality with formal languages. In Proceedings of the 29th International Conference on Computational Linguistics (pp. 6007-6018).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)