How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad

We thank the reviewer’s feedback on the writing of the paper, we will revise the paper accordingly to enhance its readability. In particular, we will use different colors alongside annotations for the inductive scratchpad examples to make reading and parsing them easier.

Further, note that we include the details of the implementation of the inductive scratchpad in Appendix C.2; does the reviewer suggest that we should move this to the main text? This may be challenging with the available space.

Please read the answers to the questions below.

Q. How would the inductive scratchpad scale? What’s the impact on the compute demands?

The inductive scratchpad approach can scale easily to larger models, datasets, and context lengths as it reduces the attention computations and the effective context size. Assume, we have states s $0$ , …, s $k$ in the inductive scratchpad. The inductive scratchpad enforces s $i$ to only use the information of s $i-1$ . This behavior can be either implemented by masking the attention between s $i$ and all the previous states except s $i-1$ or by trimming the states prior to s $i-1$ manually. As a result, the inductive scratchpad actually reduces the ‘effective’ context size. Thus, this method can easily scale. Note that due to the reduced effective context length, one can also use less memory during both training and inference. We have already discussed the implementation of the inductive scratchpad and its efficiency in Appendix C.2 and we will further extend it focusing on its scalability.

Q. Does this approach generalize to other architectures and language models?

Yes, this approach is easily implementable for other variants of Transformers and language models. First, note that the inductive behavior is controlled by two special tokens: the <START> and <STATE> (denoted #) tokens. Thus, for inputs that do not require induction (such as normal text generation) the model would keep its default behavior (as it would not generate <START>/<STATE> tokens). Regarding the implementation, assume that we have states s $0$ , …, s $k$ . As discussed in Appendix C.2, there are two methods for implementing the inductive scratchpad. The first approach is based on masking the attention between s $i$ and all previous states except s $i-1$ and reindexing the positions of the tokens of s $i-1$ and s $i$ as if s $i-1$ was generated right after the question. This implementation would only require attention masking and reindexing of positions of tokens which is easily doable in all variants of Transformers (e.g., whether they use absolute or relative positional embedding). The second approach is also implemented by manually removing the states prior to s $i-1$ for generating s $i$ before giving the text to the model during inference and splitting the input into pairs of states (s $i-1$ #s $i$ ) during training. The second implementation method only modifies the input/training data of the Transformer prior to feeding it to the model (for generation this is in the outer loop that calls the model for generating one token at a time) and hence would work with any Transformer model.

Finally, note that our current implementation is for the decoder-only Transformers as they are more common nowadays. Nevertheless, encoder-decoder Transformers can be handled in the same fashion. The only difference is that in the encoder-decoder models, the question (Q) is already separated from the rest of the tokens as it is in the encoder part of the architecture, this would make the implementation slightly easier as we can consider removing the <START> token.

Q. Can you provide a step by step method for implementing this in other kinds of problems?

As a large class, here we consider the algorithms that have a loop in which they update some variables/arrays until the answer is computed. To design an inductive scratchpad for such algorithms, we first put the input data of the question (e.g., a graph) and all other constant and permanent information before the <START> token so that the Transformer attends to this information in all of the generation steps. Afterward, we define the states of the inductive scratchpad to be all of the variables that are being updated in the loop of the algorithm (including the loop counters). Now, generating the next step/state in the scratchpad becomes equivalent to one iteration in the loop. For example for the cycles task, we keep the graph input as the permanent info in the question part before the <START> token. After that in each state, we keep track of the node that we are visiting in our search, this is equivalent to having a variable that keeps track of the current node of the search in a for loop. Also, the Transformer learns the termination condition of this loop which is reaching the source/query vertex. Similarly, for the parity/addition task (random space method), we keep the input as the permanent part before the <START> token. In the states of the inductive scratchpad, we keep track of the location of the corresponding digit(s) and their values along with other variables like the current value of the parity, carry, and the addition result. This is again, similar to a for loop that updates these variables in each iteration. Note that the inductive scratchpads for addition use some minor adjustment in the inductive scratchpad format, i.e., the initial generation of a random answer. This part is simply added for the length generalization of our experiments and to prevent scenarios in which the Transformer has to generate tokens for unseen positions during training (note that we train our Transformers from scratch with absolute positional embedding). However, we expect one can skip such adjustment when using a pre-trained Transformer with proper relative positional embedding.

Q. Broader application and comparison of the inductive scratchpad?

Please see the global response.

Q. How does the inductive scratchpad compare to SOTA methods?

The length generalization for tasks such as parity and addition has been studied rather extensively recently and overall the reported improvements remain fairly modest. We note that all of the proposed schemes have some additional conditions for their inputs and/or solution designs that are specific for the tasks (e.g., tailored positional embeddings). Our work trains the Transformers from scratch and uses absolute positional embeddings, however, we also ask for a mild condition on the input: random spaces for parity and additions. Since each work has its own modifications and solution design, it is not trivial to compare just the performance. However if one wants to compare the performance with a rough description of the approach, it would look like this:

Addition

work | performance | method and assumptions
[1] | from 8 digits to 9 digits | Using NoPE (no positional embedding)
[2] | from 10 digits to 12 digits | Scratchpad with recursive format
[3] | from 40 digits to 50 digits | Reverse order of the output + ‘index hints’ (special tokens that go before each digit). An example looks like a5b4 + a3b7 = b1a9.
[4] | from 5 digits to 15 digits | encoder-only model (no autoregressive generation) + relative positional embedding + padding inputs (input looks like 1 2 <PAD> + 3 9 <PAD>)
[5] | form 40 digits to 60 digits | Fire relative positional embedding + randomized position encodings + reversed output + index hints (similar to the above input looks like a5b4 + a3b7)
Our random space method | from 10 to 18 digits | Using random space in the input + inductive scratchpad. An input looks like 94_+_3__1= (we use _ instead of space for better readability).
Our shift method | from 4 to 26 digits | Using a random text before each number + inductive scratchpad. An input looks like fs\$46+ih\\$98.

Parity

work | performance | method and assumptions
[1] | from 8 bits to 12 bits | Using NoPE (no positional embedding)
[3] | from 30 bits to 50 bits | Using scratchpad + ‘index hints’ (special tokens that go before each bit in the input). An example looks like a0b0c1d1e0 > +c-d+
[6] | from 8 bits to 20 bits | Using pretrained large models (128B) + prompting + fine tuning + scratchpad
Our method | from 30 bits to 55 bits | Using random spaces in the input + inductive scratchpad. An input looks like _01_10_0__1_

(reported values are for median of the seeds, some seeds do better than others)

Therefore our solution requires one of the least stringent modifications of the inputs and provides a significant improvement of the length generalization compared to prior works.

[1] The Impact of Positional Encoding on Length Generalization in Transformers, Kazemnejad et al. 2023
[2] Positional description matters for transformers arithmetic, Shen et al., 2023
[3] What algorithms can transformers learn? a study in length generalization, Zhou et al., 2023
[4] Length generalization in arithmetic transformers, Jelassi et al., 2023
[5] Transformers Can Achieve Length Generalization But Not Robustly, Zhou et al., 2024
[6] Exploring Length Generalization in Large Language Models, Anil et al., 2022

We thank the reviewer for the constructive feedback. We address the rest of the remarks and questions below.

Q. Is the graph classification task possible to be solved by a Transformer regardless of training? This is quite important since the work is focused on the learnability of Transformers instead of expressiveness but if the Transformers cannot express the task, then the learnability impossibility seems trivial.

Yes it can express the task. As mentioned in the general comment, this is not a case where the failure comes from expressivity limitations. In the specific version of Theorem 1, the problem is equivalent to checking if a product of permutations in $S\_3$ is the identity, which is solvable by a constant depth transformer. The versions with a less restrictive formatting of the input are probably not expressible by a transformer with log-precision, constant depth and polynomial size. However, we do not restrict ourselves to this specific regime, and our definition of T-regular allows an arbitrary polynomial depth, and a logarithmic depth transformer can solve the graph classification task by finding the vertex $2$ edges from each vertex, then finding the vertex $4$ edges from each vertex, and then finding the vertex $8$ edges from each vertex and so on and then combining those to find the vertex $n$ edges from some designated starting vertex and checking if it is the same vertex.

Q. Appendix B repeatedly refers to $d\_{emb}$ and seems to use $d\_{emb}$ to denote the maximum length of the input string. However, $d\_{emb}$ typically refers to the embedding dimension of every token. Is there a relationship between the two?

Here we used embedding dimension $d\_{emb}$ for the input space as in “a two-digit number embedded in a dimension $d\_{emb}=5$ such as _6_1_”. In other words, we used embedding dimension as the number of tokens that the input uses given the random spacing that we use. This is completely independent of the embedding dimension of the Transformer architecture. You’re right that this is not an ideal notation. We will update this.

Q. In Appendix B.2 “Length generalization for addition”, why is generating an initial random answer ans $0$ (starting with $ token) useful?

We generate a random sequence of tokens with size+2 elements starting with \$as ans$ 0 $. In each of the steps afterward we shift ans$ i $one unit to the right (losing the rightmost token) and concatenate the computed digit of the correct answer from the left, i.e., ans$ i+1 \= \<computed digit of the answer\> ShiftRight(ans i $). Note that initializing ans$ 0 $with a random placeholder text ensures that the length of ans$ i $does not change during different generation steps. If we had initialized ans$ 0 $with an empty string, then ans$ i $would have had$i$tokens. As a result, during the length generalization, the Transformer would have had to write tokens on the positions it had never seen during the training which is impossible as we train from scratch and use absolute positional embedding. Hence, the initial random text for ans$ 0 $ensures that the Transformer will not see unseen positions during generation. We expect one can avoid this placeholder by using an appropriate relative positional embedding (potentially in addition to pre-training). We simply use \\$ as the first token of ans $0$ so that the final answer can always be easily recognized (it’s the number left to the $).

Q. Does “pointer” denote a separate pointer token for each possible input position? Are “ $00$ ” one single token or 3/4 separate tokens? What is the meaning of "the value of the ith bit can be retrieved using the pointer", is this specially processed in the generation process? It seems that for different pointers the model needs to copy different positions in the input (the pointer position), doesn’t that result in a high locality?

In this example, $00$ has 4 separate tokens. In the examples of lines 594 and 629, each character is an individual token (other than special tokens <START> and <EOS>). We emphasize that we do not do any processing for the pointers during the generation and all is learned by the Transformer model itself. By retrieving, we simply meant that the Transformer can learn to generate the token that is appearing at a specific location indicated by the pointer. We will further clarify this part and the tokenization process in our revision.

Also, note that the pointer retrieval is a low locality operation in the setting that we consider and hence easily learnable. For simplicity, assume that we have 2 pointer tokens $pq$ and we have $N$ potential tokens (t_0,..., t_{N-1}) for retrieving. Denote the output by $Y$. For simplicity also assume that the probability for retrieving each of these tokens is the same. In this case, the locality is 1. Indeed it’s enough for the model to attend to t_0 to get a significant correlation with the output. More precisely, with probability $1/N$ (if the pointer is pointing to the first token), the output $Y$ will be equal to t_0. Hence, the mutual information $I(Y;t\_0) = poly(1/N)$ (as we can determine $Y$ in $1/N$ of the cases based on t_0). Therefore, retrieving based on the pointer is indeed a low locality task.

Q. How does the inductive scratchpad work for more complex reasoning tasks such as mathematical reasoning?

Please see the global response on the broader applicability with a focus on math.

Q. The theoretical result does not rigorously prove the “negative side” of the conjecture.

Indeed, we do not prove the “negative side of Conjecture 1” and are careful to never make that claim. We prove only a specific case falling under the negative side of Conjecture 1 that relies on this variant of the cycle task. We believe that this variant is interesting in view of prior lower-bound work as it does not follow from any previous papers providing negative results such as $12, 13, 28$ (refs in our paper).

Q. The presentation can be improved. In particular, in section 2, the authors should consider using shortened/less formal definitions and remarks and putting extended versions in separate sections or the appendix. For example: Conjecture 1: Can be shortened using “if and only if” instead of having 2 sentences for “if” and “only if”?

Conjecture 1 is if and only if for constant size alphabets; we will update this; please see the global comment at the beginning in that regard.

Q. Remark 2 seems loosely relevant to the main flow of the paper. It is recommended to have a separate discussion section for these extensions and/or put them in the supplemental materials

Thanks for the suggestion, we will revise the structure.

Q. Conjecture 2 is hard to understand on its own and is not explained from a higher-level perspective.

We will break Theorem 1 and Conjecture 2 into a definition and then a conjecture/theorem to enhance readability. In the agnostic scratchpad version of the problem, we are giving the transformer access to a scratchpad but we have no prior knowledge on what the transformer should write in it. So, in order to train it to use the scratchpad well, we consider an entry it writes in the scratchpad to be “correct” if it leads to the transformer giving the right output at the end and “incorrect” if it does not and we define the loss of the scratchpad to be equal to the loss of the output. We can add more explanation to the paper.

Q. In Theorem 1, it’s unclear why are the label names a_i, b_i, c_i necessary? It seems that the label names do not impact the correctness of the theorem and thus their descriptions can be removed for clarity of the theorem.

The justification for the vertex names in Theorem 1 is not trivial. In Theorem 1 the input is formatted in such a way that it can be divided into blocks that each specify how a_i, b_i, and c_i are connected to a_{i+1}, b_{i+1}, and c_{i+1} for some i (each corresponding to a permutation in $S_3$). Each block has 6 possible values and these blocks are independent except for the fact that their overall product is always an even permutation. If we were to change the way we assign labels to vertices in such a way as to make it so that the set of vertex names mentioned in a given block was not always the same then that would no longer be the case and the proof would require additional arguments.

We thank the reviewer for the constructive feedback. We address the rest of the remarks and questions below.

Q. What is the connection to the “Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective“ paper which considers dynamic programming tasks solved by scratchpad?

Thank you for introducing this paper, we will add it to our literature review. In comparing it with the inductive scratchpad part of our paper, the inductive scratchpad is indeed suitable for learning DP tasks. DP tasks can be implemented using a loop, and the inductive scratchpad can easily learn algorithms with loops. In order to learn such algorithms with loops one has to put all the variables/arrays that are updated in the loop in the state of the inductive scratchpad (including the counters of the loop). In that case, each state generated by the inductive scratchpad becomes similar to one iteration of the loop. The inductive scratchpads that we have provided for the parity, addition, and the cycles task can also be viewed as algorithms with loops.

Nevertheless, we clarify that the mentioned paper does not use induction or any similar idea. Note that in our paper, we make Transformers behave inductively on such reasoning tasks with the introduction of two special <START> and <STATE> (also denoted as #) tokens and dynamic attention masking according to these new tokens. As a result, we obtain strong length generalization performances (e.g., almost double length for parity and addition tasks) while without the inductive scratchpad idea length generalization property is weak or absent (e.g., in the DP paper the length generalization is from 15 to 17/18). Further, note that the inductive scratchpad allows Transformers to work with reasoning tasks with longer scratchpads/contexts thanks to the dynamic attention masking. For example, the tasks considered in the DP paper have remarkably shorter scratchpads than the parity and addition tasks we have in our work and hence could be considered simpler for the Transformers.

On the theoretical side, the suggested paper focuses only on the representation power of constant-depth transformers. It shows that there are tasks that cannot be represented by constant-depth Transformers, however, they can be expressed by a constant-depth Transformer with a scratchpad. We note that our paper considers learnability hardness (beyond expressivity). For example, the cycles task can be represented by polynomial size transformers with logarithmic depth however cannot be learned by them as shown in Theorem 1.

Q. Missing reference on line 229.

Thanks for noticing the typo. The intended reference was $12$ (Emmanuel Abbe and Colin Sandon. Polynomial-time universality and limitations of deep learning. Communications on Pure and Applied Mathematics, 76(11):3493–3549, 2023).

Q. Line 231: What does k\wedge (n-k) \= \\omega(1) mean? Elaborate on the whole sentence.

Here, by $k\wedge (n-k)$ we meant $\\min\\{k, n-k\\}$ operation and by the little omega $\\omega(1)$ we meant any quantity that grows at any rate beyond constants (as $n$ scales). In other words, a parity function of degree $k$ is learnable in polynomial time only if either $k$ or $n-k$ is constant. Note that the locality of the parity task is also equal to $\\min\\{k,n-k\\}$ (because of the histogram). Therefore, the constant locality requirement is consistent with prior works on the hardness of parities. Note that in the locality definition, we have access to the histogram of all tokens, and as a result, the full parity $x\_1x\_2 \cdots x\_n$ has locality $0$. Thus, thanks to the histogram, the hardness of learning for $x\_1x\_2 \cdots x\_k$ (degree $k$ parity) and $x\_{k+1} \cdots x\_{n}$ (degree $n-k$ parity) is similar, hence we use $\\min\\{k, n-k\\}$.

We thank the reviewer for the constructive feedback. We address the rest of the remarks and questions below.

Q. Intuitively, why is global reasoning hard for transformers?

Please see first the global response. For more intuition: we train transformers using a gradient descent algorithm, which means we need to have a significant gradient in order to make progress. I.e., we need the target function to have a nontrivial correlation to the derivative of the function the transformer is currently computing with respect to at least one of its weights. If there is a constant-sized set of inputs that can be used to predict the target function with nontrivial accuracy then that means there is a low-degree polynomial in the inputs that is correlated with the target function, and we would expect at least some of the derivatives to be significantly correlated with it. However, if there is no such set of inputs then that means that no low-degree polynomial is significantly correlated with the target function. There are a superpolynomial number of uncorrelated high-degree polynomials, so the derivatives of the transformer with respect to its weights cannot be nonnegligibly correlated with a nonnegligible fraction of them. So, we will tend to be unable to learn the target function in that case.

Q. Terms like inverse polynomial edge in Conjecture 1 can be further clarified. Theorem 1 and Conjecture 2 are written in a verbose way.

Please see the general comments. We can split Theorem 1 and Conjecture 2 into a definition part first and then the theorem/conjecture statement to enhance the readability.

Q. The paper can use some more intuition. E.g., why is the mutual information targeted to be $n^{-O(1)}$? What's the role of the histogram. Lemma 1 can be further explained. A sketch of the proof from the appendix would be valuable.

We can certainly add more intuition, differing other parts to the appendix. Please see the general comment and the previous question’s answer as a starting point. Regarding Lemma 1’s intuition: if you see only $n-1$ edges in this task, then those $n-1$ edges do not tell you anything useful because you miss one edge to conclude that you have in fact a short n-long cycle closing, or instead a longer path in the (2n)-long cycle not closing. Because of the uniform distribution on how the vertices are labeled, the two scenarios are equally likely and you can thus not infer any useful information. This means that you need at least $n$ edges to get at best a non-zero correlation, so the locality is at least $n$. Further see Appendix E.

The role of histogram is also natural for the Transformer’s architecture. To see this, consider removing positional embeddings; the remaining architecture would see the input as a set, using the number of times each token is appearing. This effect can also be observed if for example positional embeddings are initialized small enough. As a result, Transformers can easily access the token count (histogram) information of the input. Here, the histogram is similar to the bag of words in the natural language processing domain. We will add more details on this.

Also, note that there is a sketch of proof at the beginning of Appendix D.

Q. While the notion of an inductive scratchpad is nicely general, the specifics of the state are highly dependent on the problem at hand. How can other inductive tasks be learned, in particular using "in-context learning".

Please see our general response on the broader applicability of the inductive scratchpad. In-context learning of new inductive tasks is indeed an interesting future research direction. We can imagine two settings for in-context learning of inductive tasks. (1) Having a normal pre-trained model (without the inductive scratchpad implementations or tokens). With the general advancement of LLMs, we think in-context learning of some inductive tasks may be possible. However, the family of such tasks will probably be limited as the model needs to see similar tasks to the target in question during pre-training. Also, the (OOD) performance of the induction would probably be much more limited than the architecture that enforces the inductive behavior. As seen in our paper, even training on standard scratchpad fails at OOD generalization. (2) The second setting is to have a model with inductive scratchpad (with the special tokens) pre-trained on other inductive tasks. In this setting, in-context learning of new inductive tasks seems more likely; nevertheless, this direction requires new investigations for future works.

Q. Line 85: Did you mean to cite the paper titled The Parallelism Tradeoff: … by the same authors? That work showed TC^0 as the bound, not TC^1, though the generalization you mention in line 171 is accurate.

You are right we will cite [The Parallelism tradeoff…] when there is no CoT, and you are also right it should say TC^0 instead of TC^1 there.

Q. In line 193, shouldn't the vertex be labeled a_i (and not b_i) if it's at distance i from a_0?

The 3 nodes that are at distance i from a_0, b_0, c_0 are randomly called a_i, b_i, c_i. If we always called the vertex i edges from a_0 a_i then solving the problem would reduce to checking if the next vertex after a_{n-1} was a_0, which would be easy to learn. Randomizing the names is necessary to make it harder. Naming the vertices in a manner that does not specify how far they are from a_0, b_0, or c_0 would probably make it even harder to learn, but it would not work with our current proof technique, so we are sticking with these vertex labels for now. And again, the original cycle task is also expected to be hard enough, but this would likely require even more proof technicalities.

We also thank the reviewer for the typos; we will fix them.