Generalizing Reasoning Problems to Longer Lengths

$**Dear Reviewer AyZa,**$

Thank you very much for your positive comments. We hope this response can address your question.

$**Q1**$: How robust is the approach to errors in the steps of the CoT? Can you conduct some experiments to study this and report in the final paper?

$**Response**$: We have experimented on parity-[2] by choosing a token at random in each CoT step and replacing it with another token chosen at random from the vocabulary, and found that its accuracy is still 1.0 on the tested lengths. We'll continue adding more noise on parity and other problems and report the results in the final revised paper.

$**M1**$: L 133: “Note that we do not define CoT formally” - this is a weakness for a theory paper and should ideally be fixed.

$**Response**$: We have defined CoT in lines 122-126.

$**M2**$: L 136 (Problem statement): The text says “performs well“ which is very informal ...... being precise is important.

$**Response**$: Fixed. Changed to ``perform perfectly.''

$**M3**$: L 170 says that “with only the continuity bias […] it is almost impossible to predict” ...... as the gap between $N$ and $N'$ grows”.

$**Response**$: We rewrite Theorem 3.1.

$**M4**$: L536-539: I would have liked to see that much earlier in the paper (e.g., around L176).

$**Response**$: Thanks for the suggestion. It has been done. See footnote 4 in the revised paper.

$**M5**$: L 251: $s_{j,l}$ and $s_{j,r}$ swapped.

$**Response**$: We believe that it is correct.

$**M6**$: L 347: “We use 4 classes” - should be 5.

$**Response**$: Fixed.

$**I3**$: Strengthen the point of how the current theoretical understanding can help design learnable CoT schemes that generalize. The empirical examples in the paper are good, but how much did the theoretical understanding contribute? Could it help to come up with a less manual process or procedure to design CoT schemes? How? Currently, the manual decomposition / rewriting of the original problem does all the heavy lifting.

$**Response**$: Yes, we have a method capable of automatically designing CoT schemes to $(n,r)$-consistency. However, its computational complexity is prohibitively high, making it impractical for real-world use. This method is outlined in Appendix B.

$**I4**$: The notation in Sec. 3.3 is quite tedious to parse - a figure would really help (probably all of Def. 3.2 could easily be shown graphically).

$**Response**$: Thanks for the suggestion. We thought about it and found it quite hard to draw a figure. But we will keep thinking. If we find a good way, we will add it to the final revised paper.

$**I5**$: Theorem 3.6 is fine (there always exists a transformer that can solve the CoT problem), but what about learnability (via SGD)? How do we find that transformer in practice? Is that always easy or only sometimes, and if so, when?

$**Response**$: As we stated in footnote 8 in the revised paper (footnote 4 in the original submission), this paper does not study learnability. Our experiments show that Transformer-based models can learn to achieve LG for challenging mathematical reasoning tasks. We find the Transformer by training with Adam. In experiments, we can always find the Transformer to achieve LG on the tested lengths given the same hyperparameters (e.g. learning rate, train/test length, etc).

$**Q1**$: What can be said theoretically about the complexity of (n,r) consistent CoT schemes? How does the number of CoT steps scale with problem complexity and problem length (e.g. the CoT scheme for multiplication seems to not scale very well).

$**Response**$: We still do not have a mature understanding of the complexity of (n, r)-consistent.

Regarding the number of CoT steps, we think it is not a function of $(n, r)$ by the following example. For a problem, let $S_0, S_1, \dots, S_T$ to be the CoT process. We could simply define a new problem by letting $S_0, S_1$ to be the CoT process and $S_1$ to be the final answer. This new problem shares the same (n, r)-consistent property but only has one CoT step for all problem instances.

Regarding the complexity, we conjecture that the complexity is a problem-specific function of the length, which positively correlates to $(n, r)$ and the number of steps, e.g. $C = C(N) \propto n \cdot \log T$, where $C$ is complexity, $N$ is length, $T$ is the number of steps. Further research is still needed.

$**Q2**$: How would one use the insights from the paper to design reasoning systems that can reason in many settings, or even at LLM / Frontier model scale? Currently the CoT format for each task is very different, and it is unclear that there is any synergy between learning different tasks simultaneously with a single model. The very different CoT schemes may even hurt performance when learning many tasks simultaneously.

$**Response**$: Yes, for different problems their CoT schemes are different because the calculation operations are quite different for different problems as we explained in the response to the first weakness and in Section 3.5.

We don't expect different tasks to be learned simultaneously. Our vision is that only the basic reasoning functions need to be learned in the proposed way individually. Those more complex reasoning tasks are simply combinations of the basic functions, for which we can make the LLMs/Frontier models to call the basic functions when needed. The current approach of using ever large amount of data to train LLMs may not be the best solution for solving reasoning tasks.

$**Q3**$: L 186: Why is it theoretically important that the lengths are the same?

$**Response**$: This is due to the requirement in (iv) of Definition 3.2 for (n, r)-consistency. See lines 259-262.

$**Q4**$: How can one in general determine if a problem is (n,r) consistent? Only by finding a valid decomposition (CoT scheme)? How does one find n and r in practice?

$**Response**$: We discussed this issue in Section 3.5. We still do not have a general procedure for these. The general idea is to use tags to let the system know what elements are involved in each reasoning or calculation step such that there is no ambiguity. Also, see our responses to your first point of Weaknesses and the third point of Improvements.

$**Dear Reviewer sCY1,**$

We would like to express our sincere thanks for your constructive and in-depth review and advice, which we have followed in revising our paper. Specifically, we have reoriented the writing to focus on the empirical contributions. We have strong results that have never been achieved before and we also provide proofs for the solution, which again has not been done before. If you have additional questions or concerns, we will be very happy to address them. Below, we address your comments and answer your questions.

$**W1**$: The CoT decomposition does most of the heavy lifting ...... relates to general function and complexity classes (e.g., regular languages)

$**Response**$: We do not have a general and practical way (see our response to question 3) that can be used for all problems. We believe that it is probably unlikely there is a general bias that can be applied to solve all reasoning problems because they are so different from each other, e.g., the calculations involved in addition and multiplication are entirely different. LG is an extrapolation problem and require the learning of underlying rules of calculation of each problem rather than short-cut regularities in the training data, which are only sufficient for interpolation.

For humans, when we teach children to do addition or multiplication, we teach them the computation rules rather than letting them guess based on the input and output alone. The tags in the CoT scheme of each problem basically inform the learner what elements or positions are involved in the calculation in each step of the problem. This enables the learner to learn the rules of calculation to achieve LG.

Regarding complexity classes, we thought about them but did not have a mature idea. Thus, we wanted this paper to focus on the LG problem alone. Thanks for your advice given in your first point of the ``Improvements'' section below. We followed it and reoriented the writing of the paper to make it focus on the empirical side of the contribution. That is, the paper defines a problem class for which we proposed a condition for its CoT schemes to achieve LG. We believe that our work is still a strong contribution compared to the existing approaches as none of them can solve the multiplication problem perfectly on the tested problem lengths and none have even attempted the division problem. Additionally, none of them has any proof of their solutions.

The changes made to the paper can be found in the abstract and introduction and in lines 137-139 and footnote 3, which are all highlighted in blue. Minor edits are also made in numerous other places to align with the new focus, but not highlighted.

$**W2**$: Writing throughout most of the manuscript makes claims ...... I will be more specific in the Questions section.

$**Response**$: There may be some misunderstanding here. We have not claimed that $(n, r)$-consistency is necessary. We only claimed that it is a sufficient condition. As described in the previous response, we have reoriented our paper to focus on the empirical contributions.

$**W3**$: Learnability and generalization with transformers in practice remain unclear ...... reaching the scale of modern frontier models?

$**Response**$: As we discussed in footnote 9, this work has not studied learnability. Our theory and proof show the existence of a solution, and our empirical results running on a transformer give perfect accuracy on the tested lengths. Regarding the scale of CoT steps, please refer to our response to your question 1. Regarding learning all tasks simultaneously, please refer to our response to your question 2.

$**I1**$: A significant overhaul / rewrite to emphasize that the main contribution is not a theory of length generalization and identification of a novel root cause (the root cause is well known in machine learning, mathematics, and philosophy), but a definition of a problem class that has favorable properties w.r.t. length generalization built in, and these properties can be exploited by transformers.

$**Response**$: Thanks for the valuable advice. We have reoriented the writing to shift the focus of our paper to the empirical contributions. See our response to the first weakness.

Based on the strong empirical contribution, we hope it is sufficient for you to raise your score, as none of the existing approaches in the literature can solve multiplication on the tested problem lengths and none has attempted division. Furthermore, none of the existing papers has any proof of their approaches.

$**I2**$: If the focus of a major revision lies on the theory ...... what are counterexamples?

$**Response**$: See our response to the last point above.

$**Q3**$: I'm not sure if I completely understand how the padding works. When you pad, you still have blank tokens at the start and end of the sequence. However, they are not the middle of any interval. So how are they predicted? Also are the padding tokens provided for all the examples in Appendix D?

$**Response**$: The padding tokens (or elements) at the start and the end are not predicted. They are added to ensure that every predicted token can be the central or middle token of a r-length interval. Since the first token has no token on the left, then it cannot be the central/middle token of an interval. To resolve it, we add about $r/2$ padding tokens before the first token. For the same reason, we add some padding tokens after the end token. They are not provide in examples.

$**Q4**$: In experiments, the length generalization ability is tried on different lengths and for all of them we see 100% performance. So I wonder what's the point that we see some reduction in the performance? Why haven't you tested the length generalization ability on longer sequences? The experiments are done with specialized Transformers. I wonder what would be the length generalization performance if CoT steps were learned by a typical model, e.g., GPT2 or llama?

$**Response**$: We have run experiments on Parity-[2] with length 1000 and still got 100% accuracy. This is not surprising as Theorem 3.5 shows the existence of a 'perfect' Transformer. We also ran Addition-[2], which gives 100% accuracy on length 20 (used in the paper), 30, 40, but 73% on length 50. We think the length amplified the mismatch between the 'learned' Transformer and the 'perfect' Transformer. We will conduct comprehensive experiments on all problems, including multiplication and division, to find the longest lengths that can achieve 100% accuracy and report in the final revised paper.

Regarding using a typical model, our Transformer adds a relative encoder with mask, which depends on $(n, r)$-consistency property of the reasoning problem (Sec.4.3, 1st paragraph). This mask ensures that every r-length interval is 'clean' i.e. without mixing with noisy tokens beyond the r-length interval. The proof of Theorem 3.5 is a process of matching a sequence of 'clean' r-length intervals. If replacing our Transformer with a typical one, with high probability, the LG would decay a lot as there is no 'clean' r-length interval and the matching process can be chaotic.

$**M1**$: From what I understand, it's not important where the intervals other than the anchor interval are located, is that right? I think it would be nice if this is further clarified in the paper.

$**Response**$: Yes, that is right. We clarified this in footnote 6 (line 268).

$**M2**$: I think Theorem 3.1 is doing more harm than good ...... I would really suggest to use a simpler theorem on discrete token space.

$**Response**$: Following your suggestion, we have simplified the theorem using a discrete token space. See lines 149-157. The proof is also updated in Appendix C.

$**M3**$: The abstract and intro imply that authors have found a necessity condition for length generalization and thus I think they should be revised.

$**Response**$: We have revised it. See lines 18-20 and the first contribution starting from line 63.

$**M4**$: I find the proof sketch part very interesting, and I'd suggest authors to further elaborate on it. (In that case, you may want to move some of the experiment details to the appendix).

$**Response**$: We tried but found that there is not much to add without referring to the full proof in the appendix.

$**Dear Reviewer F4Li**$,

Thank you very much for your valuable and constructive comments and questions. Below, we address them point by point. We have updated the paper to reflect the changes.

$**W1**$: One of the main weaknesses is the modeling assumption ...... Further, this problem becomes more significant in two ways:

$**Response**$: Thanks for the suggestion. We have acknowledged the difference from the original CoT modeling in footnote 2 in line 160.

$**W2**$: The model doesn't learn the termination condition ...... but this is manually controlled by the user.

$**Response**$: Thanks for your suggestion. We have added one special 'EOS' token at the end of the CoT process for each training instance in the dataset. In testing, the output generation process terminates iff 'EOS' is predicted. We ran the experiments again and the results remain the same. We reflected this change in line 451 and line 502 in the revised paper.

$**W3**$: It seems some further processing are done on the inputs and outputs. In particular, it seems input[k+1] is not exactly output[k].

$**Response**$: Yes. In each step, the input and output elements need to be aligned as required by $(n,r)$-consistency (see (iv) in Definition 3.2 in line 259). For the arithmetic problem, after a step, empty elements are removed and for addition, an empty element is added for padding (see more details about padding below).

$**W4**$: All of these processings and manual handlings seem like ...... in particular to problems that we have less control over the data.)

$**Response**$: We have a different opinion here. We would not regard the proposed method as an engineering hack as we have a theory and proof to support the method. The proposed $(n, r)$-consistency is a general condition. If the CoT scheme of any problem meets the condition, it can achieve LG. $(n, r)$-consistency is not specific to the problems used in the paper because our proofs are not specific to these problems. We believe that existing methods are mainly heuristics as they have no proof of the generality of their methods. Their methods are also much weaker than ours and no paper has attempted the division problem.

$**Q1**$: From what I understand from Appendix D, we see that output[k] is not always exactly equal to input[k+1] (e.g., in the addition examples). I think this is some sort of processing that is not necessarily generalizable to other tasks because it uses the knowledge of the task too much (I personally see it an engineering hack). Is there a way to avoid that and be more agnostic towards the task?

$**Response**$: See our response to ''the processing'' above. Regarding the use of the knowledge of the task, we think it is unavoidable because to achieve LG, the learner has to learn the rules of computation for each problem in order to extrapolate rather than short-cut regularities in the training data, which may be enough for interpolation but not enough for extrapolation or LG. That is why we require the CoT scheme of the reasoning problem to satisfy $(n,r)$-consistency. Designing a CoT scheme to meet $(n,r)$-consistency needs the calculation knowledge of the problem. Intuitive, Theorem 3.1 shows that a higher dimensional function is not predictable with only partial observation of the function projected onto a lower dimension. It is thus necessary to have a bias that limits the hypothesis space of the higher dimensional function. We think different problems may need different biases. But it may be possible to have a more general condition than $(n, r)$-consistency.

As humans, when we teach children to do addition or multiplication, we teach them the rules of computation rather than letting them guess based on the input and output alone. The tags in the CoT scheme serve to inform the learner what elements or positions are involved in the computation in each calculating step (also see Section 3.5 in the revised paper.).

That being said, we do have a method capable of automatically designing CoT schemes to satisfy $(n,r)$-consistency. However, its computational complexity is prohibitively high, making it impractical for real-world use. This method is outlined in Appendix B.

$**Q2**$: Similar to the concern above, In Appendix D, we see some orange lines that are not predicted by the model and are manually inserted. I think this again reduces the learning part of the paper and is an engineering trick. Can it be avoided?

$**Response**$: These orange lines are only for easy reading or understanding. They are static and pure input and do not change during the CoT process. See examples for multiplication.

$**Dear Reviewer F4Li**$,

Thank you very much for your new comments and questions. We are happy to address them.

$**Q1**$: How does this approach generalize and how do you think your approach will be useful in practice in the future?

$**Response**$: Yes, it can achieve generalization, but use a different approach than the familiar methods. We believe that only a finite set of basic reasoning functions needs to be learned individually using the proposed method. More complex reasoning tasks can then be addressed as combinations of these basic functions, with the complex tasks invoking the appropriate basic functions as needed. We suspect this mirrors how humans learn and execute reasoning tasks. Relying on ever-growing amounts of sequential data to train large models in a brute-force manner may not be the most effective solution.

$**Q2**$: Can you train a model by just using some sequence of tokens? (without extra interventions in the process such as adding/removing paddings manually at test time)

$**Response**$: We cannot, as certain problems (e.g., arithmetic, addition) cannot guarantee $(n,r)$-consistency without them. It remains unclear whether these problems can achieve $(n,r)$-consistency through some transformations that avoid adding or removing padding. We leave this question for future work.

$**Q3**$: Another concern of mine is that this works completely ignores the positional embeddings ...... Do you think your circumventions would generalize in practice?

$**Response**$: Manipulating positional embeddings as input cannot guarantee length generalization (LG), even though existing heuristic methods show some improvement with carefully designed positional embeddings. This limitation arises because such approaches fall within the scope of Theorem 3.1, where the learner does not encounter positional embeddings of longer sequences in training and, therefore, cannot ensure LG in testing. The positional embedding is thus insufficient. In contrast, our $(n, r)$-consistency-based method only requires the learner to see all $n$ $r$-length intervals to guarantee LG, irrespective of the length of the training or testing sequences. We have theoretically proven the existence of a solution that achieves perfect LG under this framework.

$**Minor typo**$: there seems to be a typo in the orange line of parity output[3].

$**Response**$: We have fixed it. Thanks.

$**In summary**$, our work introduces a theoretical framework for understanding reasoning that extends beyond specific problems or architectures. It also proposes a theoretically grounded solution that utilizes the proven effectiveness of Transformers. While we do not claim that our approach is the optimal solution, we hope it serves to inspire further investigation guided by principled methodologies.

We understand that you may not fully share our perspective. However, we trust you value the importance of pursuing new directions to advance the field. With the proposed theoretical framework and strong empirical results, we hope you agree that our paper represents a significant contribution.

Thank you very much, and we look forward to hearing your feedback.

Authors

$**Dear Reviewer F4Li**$,

With the discussion period nearing its end, we would greatly appreciate your feedback on our responses to your recent questions. If you have any additional questions or comments, we would be more than happy to address them.

Thank you in advance,

Authors.