Paper: What Algorithms can Transformers Learn? A Study in Length Generalization (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
This is indeed a challenging aspect of the framework. Currently predicting or improving length generalization entails finding a RASP-L program that solves the task. We provided a few examples that can be generally more applicable, such as adding index hints for tasks that require complex indexing operations, and scratchpad for unrolling loops in the algorithm. Nonetheless, one promising future direction is to identify “unnatural operations” as suggested by RASP-L, but which are common subroutines for a large class of algorithms, and increase the training data coverage for those operations. This may help the model to perform these operations through memorization, and leverage this to learn other algorithms in a more generalizing way.

Reviewer comment (Theory):
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.

Author response:
>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.
>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?

We agree that our work has no theoretical justification for the conjecture; it is currently a purely empirical claim. Nonetheless, we view such empirical work as an important first step towards understanding new phenomena -- a step upon which theory can be built. We hope that our conjectures can eventually be formalized and proven, but we leave this (significant) endeavor for future work.

(As an aside, the reviewer may be interested in our new Appendix H, where we added a theoretical example from boolean function analysis that illustrates our conjecture in a simple case).

We wish to clarify the scope of our main claims. Our Main Conjecture (Section 2) only characterizes which tasks Transformers are likely to length-generalize on, and not why or how they do so. Specifically, our Main Conjecture only claims that RASP-L is a useful predictive tool: empirically, if a task can be solved by a simple RASP-L program, then it is likely to exhibit strong length generalization, and vice versa. This is a phenomenological claim, as opposed to a mechanistic one. Nonetheless, we believe that the toy model in our introduction suggests a plausible mechanism for Transformer learning, and could guide new lines of mechanistic investigations in future work.

**Re. degradation at longer lengths:**

There are a number of factors that could influence how robustly the correct solution is learned, if indeed it is. We do not expect (and indeed do not observe) perfect length generalization over arbitrary lengths. Some level of degradation is likely fundamental due to issues trying to approximate a discrete program with a continuous space (e.g. noisy optimization, finite precision, etc). 

One other possible explanation, compatible with the RASP conjecture, is that there exists other “shortcut” solutions which do not generalize, but are simple enough to compete with the true RASP-L program during training. Since it is likely intractable to determine the minimum RASP-L program that fits a given training set, we cannot predict a priori what forms of “data diversity'' are required to rule out such ”shortcuts“, even if our conjecture holds true. Nonetheless, our study focuses on the relevant characteristics of tasks that influence generalization performance, thus even good albeit imperfect generalization lends strong signal to the usefulness of this perspective. We have added several new paragraphs in the revised pdf (in Section 2, and Appendix G), which highlights these limitations for the reader.

Reviewer comment (Evaluation):
There is 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.

Author response:
**Regarding OOD vs IID performance:** 

In all of our experiments except for parity, we only evaluate models that achieve ~100% accuracy on in-distribution test samples. This includes addition on all training lengths. Parity without scratchpad is a special case as it is well known that Transformers cannot even fit the training set for long parity sequences (e.g. Chiang & Cholak and Bhattamishra et al). All parity with scratchpad experiments are trained to convergence. We hope this fully addresses your concern.

Reviewer comment (Evaluation):
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.

Author response:
>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.

In Fig 3, 7b, and 8, we show the individual test performance of all 20 runs (plotted as points). For brevity we default to reporting the median, but most experiments follow the same pattern as those in Fig 3. We hope these plots can give you a fuller picture of Transformer’s performance. If there is a specific additional metric you would like to see, we are happy to compute it on the 20 runs shown.

Reviewer comment (Experiments):
The tasks are simple: from count to addition/parity.

Author response:
>The tasks are simple: from count to addition/parity.

On the positive side (length-generalizing tasks), we agree that it would be interesting to demonstrate more complex tasks which length-generalize. We focused our paper on simple tasks for exposition purposes, and to better connect with the existing literature on length generalization (where such simple tasks are standard). However, our RASP-L framework is more general, and can indeed express more complex tasks than we presented here due to space constraints (such as sorting and finite state automata). We can consider adding such tasks into the full version of this paper.

On the negative side, we could have included a large number of complex tasks, because most tasks will not length-generalize and will also not have a RASP-L program. A number of more complex tasks have been tested in the literature, where they were shown to fail at length-generalization. These tasks include: dynamic programming, multiplication, and certain kinds of logic puzzles (as tested in Dziri et al. 2023). We thought about these tasks, and could not devise RASP-L programs for them (even for e.g. multiplication), suggesting that they are unlikely to have simple RASP-L solutions. In general, Transformers only seem to length-generalize on very particular types of tasks, and so we focused on identifying positive instances of length generalization, since negative instances abound.

Reviewer comment (Experiments):
Ablation studies on "without the index hints" are needed.

Author response:
>Ablation studies on "without the index hints" are needed.

We mention briefly in Section 4.3 that no generalization is observed on standard addition or parity (no index hints) under our experimental setting. We evaluate different lengths in increments of 5, and all 20 runs for all length showed a test EM of ~0% on these tasks. For brevity we do not include a plot, but we have added this detail to the updated pdf for better clarity.

Reviewer comment (Theory):
There is no characterization of the function space represented by the "RASP-L" programs.

Author response:
>There is no characterization of the function space represented by the "RASP-L" programs.

We agree that it is an important open question, so we have added this discussion explicitly in Appendix G, paragraph “On Formal Language Characterizations.” This is also an open question even for the original RASP language of Weiss et al. (2021).

Reviewer comment (Evaluation):
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 RASP-L but Transformers hard to generalize, or vise-versa? These answers are needed to make a stronger statement on the correlation between the two.

Author response:
This is a fair question. Although we cannot prove conclusively that RASP-L covers all cases of generalization success and failures, and in fact there are likely exceptions due to the complex nature of learning and generalization, we have considered most empirical analyses of Transformer length generalization on algorithmic tasks in the literature to the best of our knowledge. We have not observed exceptions in these cases. See Figure 1a for a clearer summary.

Given that most studies in the literature focus on highlighting length generalization failures, we introduced a number of new tasks to showcase that Transformers do in fact have the ability to show strong length generalization. We have identified more of these positive examples than we had room for in the manuscript, such as sorting and finite state automata. We can consider adding such tasks into the full version of this paper. On the negative side, it is easy to find tasks which do not admit RASP-L programs. Dziri et al. 2023 studied a number of tasks (e.g. dynamic programming, multiplication, and certain kinds of logic puzzles), and showed failures in length generalization. We could not devise RASP-L programs for these algorithms, suggesting that they are unlikely to have simple RASP-L solutions. Nonetheless, RASP-L only covers a subset of possible programs that could be represented by a Transformer. In particular, RASP-L focuses on symbolic tasks. A large class of learnable but non-RASPable tasks are numerical algorithms, such as linear regression, which involve high-dimensional matrix and floating-point operations. We have added a discussion of the limitations in Appendix G.

Reviewer comment (Evaluation):
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?

Author response:
In this work, we used max train length as a measure of training data diversity. As proposed by our Main Conjecture, one condition on length generalization is diversity--- the training data should be sufficiently diverse, such that there does not exist any shorter RASP-L program which agrees with the task in-distribution but not out-of-distribution. Intuitively, a dataset containing only 5 different lengths is more likely to have shortcut solutions than a dataset that contain 50 possible lengths. We have expanded Section 5 to include a deep dive on addition, which includes a section on how increasing data diversity by changing the sampling procedure also leads to improved performance. The Main Conjecture can be used to understand both phenomena (max train length and improved sampling procedures).

Reviewer comment (Evaluation):
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.]

Author response:
>The paper studies length generalization, though it's unclear what is considered as successful length generalization in the paper.

Thank you for this question. In this paper, we consider successful length generalization to mean near perfect test EM on length that is at least 10 longer than the maximum training length, and we consider failures of length generalization to mean near 0 test performance on +10 length. We have added this clarification in a new figure (see Fig 1a) that summarizes the different tasks.

The specific choice of 10 is somewhat arbitrary. It reflects the level that all tasks which demonstrates non-trivial generalization can readily achieve without significant changes to model / data size and training strategy. Certain tasks such as count also exhibit significantly stronger length generalization than +10. However, it is also a reasonable proxy for capturing the differences between the set of tasks that are amenable to RASP-L solutions and the set that is not. Existing literature on length generalization reports many failure cases, including addition and parity, and in these cases performance starts to degrade as soon as the length goes out-of-distribution (see e.g. Nye et al Figure 3 and Anil et al). Being able to generalize perfectly for length of +10 or more suggests a qualitative difference from the failures cases reported in the literature. Nevertheless, how to predict the exact level of length generalization for a given task is still an open question.

Reviewer comment (Theory):
The paper proposes to use the difficulty of RASP-L programs as an indicator for length generalization performance. However, it's unclear 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.]

Author response:
> 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. [...] I'm not sure whether these criteria provide new insights beyond findings in the existing literature.

We emphasize that our work focuses on the empirical learnability of tasks, rather than their theoretical representability.
While there is a long line of work on the latter question (we have included more references in Appendix G), the former question is very much still open --- it is essentially about understanding the inductive bias of Transformers.
Our work does not resolve this question, but we believe it takes an important and significant step forward.

To be concrete, we have added an explicit theoretical example (Appendix H) where our conjecture’s prediction differs significantly from prior work. We consider the popular “minimum-degree-interpolator” conjecture proposed by Abbe et al (ICML 2023 Outstanding Paper: https://openreview.net/forum?id=3dqwXb1te4).
We give a simple theoretical setting (learning boolean conjunction) where our RASP-L conjecture correctly predicts successful length-generalization, but the min-degree conjecture does not. This simple example is emblematic of the value in our perspective.

Moreover, as we discuss in the Intro (paragraph 2), the prior literature did not have a consistent answer to whether Transformers length-generalize. Most works found that standard Transformers typically fail catastrophically on even mild length-generalization, e.g. Delétang et al. 2023,   Abbe et al 2023, and Liu et al. 2023(a).  However, it was also known that Transformers can sometimes length-generalize on particular tasks. Our contribution sheds light on which tasks: what is special about these length-generalizing tasks, that separates them from non-length-generalizing ones? We claim RASP-L representation is an empirically good heuristic for answering this question.

We hope these remarks help clarify the contribution of this work in the context of prior literature and address your concern regarding contribution. Please let us know if there are further concerns regarding this point.

Reviewer comment (Evaluation):
It's unclear how to empirically verify or act upon the conjecture that the difficulty of RASP-L programs indicates length generalization performance. 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.

Author response:
>It's unclear how to convert a task into RASP-L except for applying the definition directly, i.e. checking for whether certain criteria are satisfied. 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.

These two criteria you mentioned are implications of our RASP-L conjecture, but our conjecture is significantly stronger. 
To give a simple example, the Parity task satisfies the two criteria you mentioned, but cannot be written in RASP-L.

We acknowledge that “It's unclear how to convert a task into RASP-L” without actually trying to write the RASP-L program. This is however generally true of many computational complexity classes — for example, it's unclear whether how to determine if a task is in Polynomial time, without exhibiting a poly-time algorithm for the task. (Of course, certain well-studied complexity classes have structural features which help us determine membership. We do not yet have such a structural understanding of RASP-L, but we view the definition as an important first step). We have added a discussion of these limitations in Appendix H in the revised pdf.

Reviewer comment (Theory):
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.

Author response:
>It's unclear how to convert a task into RASP-L except for applying the definition directly, i.e. checking for whether certain criteria are satisfied. 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.

These two criteria you mentioned are implications of our RASP-L conjecture, but our conjecture is significantly stronger. 
To give a simple example, the Parity task satisfies the two criteria you mentioned, but cannot be written in RASP-L.

We acknowledge that “It's unclear how to convert a task into RASP-L” without actually trying to write the RASP-L program. This is however generally true of many computational complexity classes — for example, it's unclear whether how to determine if a task is in Polynomial time, without exhibiting a poly-time algorithm for the task. (Of course, certain well-studied complexity classes have structural features which help us determine membership. We do not yet have such a structural understanding of RASP-L, but we view the definition as an important first step). We have added a discussion of these limitations in Appendix H in the revised pdf.

Reviewer comment (Theory):
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.]

Author response:
>It is well known that auto-regressive generation cannot uncover non-causal information.

We are unaware of a formalization of this claim in the literature. This claim has been made in informal/imprecise ways, but note that it is tricky to formally define what it means for a task to “require non-causal operations” without a model of computation such as RASP-L.

For example, prior works postulated that forward-order addition would not length-generalize because the naive algorithm, on a von Neumann computer, involves non-causal operations. However, Transformers empirically actually can length generalize on forward-addition, and we show there exists a RASP-L algorithm which solves this task. This algorithm is not the naive one; it is specific to the Transformer architecture (not the von Neumann architecture), and it is thus able to “work around” the causality restrictions by taking clever advantage of parallelism.

Reviewer comment (Theory):
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).

Author response:
>it’s not true that omitting positional encodings may lead to "arbitrary length" generalization, due to failures in both MLP and attention

Thank you for flagging this; we will clarify the statement in the paper. Note that the issue is subtle: there exist certain tasks where arbitrary-length generalization is provably possible, and such failures of MLP and attention do not manifest. For example, consider the simple Transformer that solves the boolean conjunction task in Appendix H. (The limitations of Hahn 2020 do not apply in our setting, since we allow weights to take values in the extended real line, as described in Footnote 3).

Reviewer comment (Novelty):
It would be better to include discussions on prior work that relate algorithmic reasoning tasks to formal languages.

Author response:
>It would be better to include discussions on prior work that relate algorithmic reasoning tasks to formal languages.

We have added more discussion of such works in Appendix G, and we also have added a footnote about uniformity to Section 2 as suggested.

Reviewer comment (Novelty):
For the discussion towards the end of 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).

Author response:
>It would be better to include discussions on prior work that relate algorithmic reasoning tasks to formal languages.

We have added more discussion of such works in Appendix G, and we also have added a footnote about uniformity to Section 2 as suggested.

Reviewer comment (Experiments):
Fig 3: How many test sequences are there for each trial?

Author response:
>Fig 3: How many test sequences are there for each trial?

The test data size is 5 * batch size for each task, which is typically 5 * 128 = 640 examples.

Reviewer comment (Evaluation):
Fig 6(b): this is not really about length generalization, but more a failure of optimization?

Author response:
>Fig 6(b): this is not really about length generalization, but more a failure of optimization?

Indeed, for parity without scratchpad we see the difficulty of optimization, which has also been noted in prior work (e.g. Chiang & Cholak and Bhattamishra et al). We focus on comparing parity with different scratchpads, based on how simple the scratchpads are according to RASP-L. We observe that the tasks with simpler RASP-L solutions also show faster optimization, which lends some support to the intuition that simple RASP-L programs are “easier-to-learn”.