Universal Length Generalization with Turing Programs

We thank the reviewer for their feedback. We respond to the main points raised by the reviewer below.

Theoretical analysis may be missed: If the reviewer can further explain what is lacking in terms of the theoretical analysis, we would love to further explain and add to the revision. From our perspective, theorem 4.2 is the main theoretical result: we show a Transformer of some fixed size can simulate Turing Programs up to length that is exponential in their size, which suggests that “small” Transformers can execute very long programs. This expressivity result is connected to length generalization because it is a necessary (but not sufficient) condition for length generalization to happen with Turing Programs (not sufficient because gradient descent might not learn this constructed solution). Therefore, we feel it provides theoretical backing for the empirical success we see in algorithmic tasks.

Why special data format works: We should have made it clearer why the Turing Program data format is useful for length generalization. When we write out a task in this format, the task gets broken down to two subtasks: 1. modifying the tape content at a single position and 2. copying the tape content. The modification only requires the token at the head position and the positions where the Turing machine state is located, so is relatively independent to the overall length of the input. Copying is length dependent, but a past work (https://arxiv.org/abs/2402.01032) already showed that Hard-ALiBi could achieve length-generalization on copying. Thus, we expect Turing Programs combined with Hard-ALiBi to deliver length generalization results.

We explained this understanding in Section 2.2 of the paper (quoted below for ease of reading), but can add further explanation in the revision.

“We begin with the observation that the scratchpad technique can be realized as an iterative sequence of copying operations, where at each iteration the input is slightly modified. Building on previous works showing that with the right positional encoding, transformers can achieve length generalization on the copying operation, we hypothesize that combining the scratchpad technique with a favorable positional encoding can unlock length generalization capabilities.”

Could be used for language modeling?: Our method currently is most suited for algorithmic tasks, i.e. problems that can be decomposed into step-by-step solutions. Our work focuses on algorithmic problems that have closed and known solutions (e.g. arithmetic problems), similarly to many existing works on length generalization (e.g. https://arxiv.org/abs/2402.09371). However, we believe that this method can be adapted to a broader set of natural language problems that are of algorithmic nature, such as mathematical reasoning problems, and leave this study to future work.

We thank you again for the helpful comments on our paper. We are happy to provide further clarification, and would appreciate it if you would raise your score if you believe that your concerns were addressed.

We thank the reviewer for their feedback. We respond to the main points raised by the reviewer below.

Key claims of the paper: we want to emphasize that the key claim of our paper is that Transformers can achieve length generalization (generalization to problems longer than the ones observed in the training data) for a large class of problems. That is, we do not argue that we establish novel expressivity results—as you pointed out, the fact that language models can express a large class of functions using chain-of-thought has already been shown in various prior works. Rather, we argue, using a combination of empirical experiments and theoretical observations, that Transformers can extrapolate to longer sequence lengths when provided with chain-of-thought/scratchpad data of a particular format which tracks the step-by-step operation of a “Turing Machine” (the Turing Programs). This connection between the chain-of-thought format and length generalization goes far beyond what has been shown in prior works (including the works that you mentioned), which focus on a narrow set of algorithmic tasks. Instead, our results demonstrate the potential of “universal length generalization” — length generalization for any algorithmic tasks. While we do not establish this formally (due to limitations of our theoretical analysis), we believe that the combination of our extensive experimental results and theoretical insights suggests that Transformers can length-generalize on a far larger class of problems than previously acknowledged, a result that we see as the main novel contribution of our work. Therefore, we believe that our theoretical result should be viewed in the broader scope of the paper, coupled with the experiments on length generalization, and not as an independent expressivity result. Thank you for pointing this out, and we will clarify this in the final version of the manuscript.

Comparison to other baselines: throughout the paper, we compare our method to training without chain-of-thought, which displays poor length generalization performance. However, we agree that adding a baseline where we train with another chain-of-thought format can help establish our claims, and we plan to run additional baseline experiments and add them to the final version of the paper. Specifically, we will compare our scratchpad technique to more minimal chain-of-thought, for example one that tracks only the current number and the carry digit in the case of multi-digit addition.

The choice of n in non-repeated n-grams: we would like to clarify that Theorem 4.2 holds for any choice of n. I.e., for any $n \in \mathbb{N}$ there exists a RASP program of size (number of lines) which grows linearly with the chosen n, that satisfies the conditions of the theorem. We will clarify this in the final version of the paper.

Citations: we will fix the citations format in the paper.

We thank you again for the helpful comments on our paper. We believe that we answered the main drawbacks raised in the review (in particular, about the key claim of the paper), and would appreciate it if you would raise your score if you believe that your concerns were indeed addressed.

We thank the reviewer for their positive feedback. We respond to the main points raised by the reviewer below.

Universality: We understand that the word “universal” may not accurately capture the nature of our results, and we are open to removing it from the title if the reviewer thinks this will be adequate.

N-gram repetition: We agree that relying on n-gram repetition may be a significant restriction of our theoretical construction, but note that for “random enough” inputs, repeated n-grams become very unlikely for large enough n. Additionally, we suspect that transformers in practice may be in fact utilizing an n-gram matching mechanism (as observed in prior works, e.g. https://arxiv.org/pdf/2402.01032), which means that this limitation reflects a true limitation of transformers and not just a problem in our theoretical construction. In that sense, we believe that our theoretical construction truly captures the nature of the solutions learned by the transformers, including the potential limitation of these learned solutions.

Gap with real-world tasks and data: This is a legitimate concern. We want to make two points here:

• For this work, our goal is to do a deep evaluation of how length generalization can arise in language models trained on next-token prediction. Following a large body of prior work (see the many papers on addition, such as https://arxiv.org/abs/2402.09371), we conduct these experiments on synthetic tasks. Therefore, we hope this won’t be considered as a fatal weakness for the paper.
• There may be applications to math problems that have good algorithmic solutions. Consider the game of 24 (analyzed in https://arxiv.org/abs/2404.03683): you win by manipulating 4 numbers to reach 24. It can be solved by DFS, which can be encoded into the Turing Program format. It would be interesting to see if including Turing Programs of various math problems into the training data mix can improve length generalization performance. We leave this study to future work.

Role of Positional Encoding in Length Generalization: This is a good point. We agree that many factors contribute to length generalization performance, but we chose to focus on the choice of positional encoding as it was pointed out to be a key factor in prior works studying length generalization. Controlled experiments on how other variables affect length generalization should be done in future research.

We thank you again for the helpful comments on our paper. We believe that we answered the main drawbacks raised in the review, and would appreciate it if you would raise your score if you believe that your concerns were indeed addressed.

We thank the reviewer for their positive feedback. We respond to the main points raised by the reviewer below.

Only algorithmic problems: The reviewer is right to point out that currently the result is not immediately practical. For this work, our goal is to do a deep evaluation of how length generalization can arise in language models trained on next-token prediction. Following a large body of prior work (see the many papers on addition, such as https://arxiv.org/abs/2402.09371), we conduct these experiments on synthetic tasks. Therefore, we hope this won’t be considered as a fatal weakness for the paper.

Larger LMs: We expect no reason why the same generalization won’t hold when our technique is applied to larger models. We want to make two points:

• As model size grows, we expect the model to perform complex Turing programs better. (e.g. harder arithmetic tasks). It would be an interesting direction of future research to quantify and see how model size and task difficulty (maybe captured by the RASP program size as observed in https://arxiv.org/pdf/2310.16028) affect length generalization results.
• Turing program may be a way to construct CoT for certain math problems (see the next section), which can be used to train large models.

Improvement in math reasoning: Although it is unclear how general QA can be improved by the current iteration of Turing Program, there is certainly application to math problems that have good algorithmic solutions. Consider the game of 24 (analyzed in https://arxiv.org/abs/2404.03683): you win by manipulating 4 numbers to reach 24. It can be solved by DFS, which can be encoded into the Turing Program format. It would be interesting to see if including Turing Programs of various math problems into the training data mix can improve length generalization performance. We leave the study of how to use Turing Programs for more realistic math problems to future work.

We thank you again for the helpful comments on our paper. We believe that we answered the main drawbacks raised in the review, and would appreciate it if you would raise your score if you believe that your concerns were indeed addressed.