We thank the reviewer for their suggestion, encouragement, and comments. We are happy that you appreciated the clean idea behind our work: that sparsity naturally suggests a set of “core” tasks, which can then be trained on robustly, and this can ensure length generalization. We hope that this conceptual simplicity/message will be useful more broadly, beyond the specific problems in this paper.

Thank you also for the detailed and valuable suggestion: we will rephrase to present the proofs more transparently, and will also unify the presentation styles. We completely agree that it will make the ideas of the paper more accessible, and will help the reader. We will also add more connections and comparisons to previous work.

Please find responses and clarifications to the other questions and comments below.

The Role of Encodings: Thanks for the question. We view the role of position encodings as complementary (and indeed, additive) to our proposed recipe. Typically in order to train a model with positional encoding for length generalization standard techniques such as a random position offset is used during training. If positional encodings are present, the same techniques will apply to our framework too.

Choice of Auxiliary Tasks: It is a great question to design auxiliary tasks automatically for more real world settings such as natural language tasks and math tasks. One way would be to have a “library” of “basis” functions for the “local” tasks, and then use these as subroutines in a general training setup where they can hopefully enable length generalization for a variety of downstream applications.

Comparison to other Models: Regarding the use of scratchpads, we would like to note that in all our experiments the baseline models use scratchpads and still do not length-generalize; this is consistent with what other literature on this topic has observed. Index hinting techniques such as position coupling are quite powerful for length generalization but rely on the right hints to be fed into the model at inference time as well. These are very strong hints; our focus was to make the model naturally learn to attend to the right tokens via robust training without such direct hints, since such hints will typically not be available at inference time. We are happy to include these methods for comparison in the next version of the paper.

I thank the authors for their response. I am not quite sure, I am fully following the responses made by the authors and would like more details from the authors.

Regarding the choice of encodings for example. My question about them was motivated by the fact that I don't quite believe that the integration of standard techniques (like random offset for APE) for example would then retain the computational efficiency that the authors are hoping to achieve with their proposed approach. It would be good if the authors can provide a bit more intuition about why they think, that the same techniques as and when applied to standard training would also be applicable with their methods, and not loose out on the efficiency of their suggested approach. With APE and the positional offset trick for example, it takes quite a lot more samples at various positions (such that most positions are seed during training), to be able to achieve length generalization. Since that overhead is already quite high, would the approach suggested by authors still have a significant increase in efficiency while trying to achieve length generalization. The answer to this question highly impacts the overall contribution of the authors work. This is precisely why I would like to request the authors to engage a bit more in trying to explain their reasoning.

For the choice of Auxiliary tasks for realistic settings, could the authors provide a simple example of how they foresee having a library of basis functions, this currently seems very abstract to me, and I am not sure I quite follow what the authors mean here.

For the comparison to other models, I gave the paper a quick skim through again, to me, it was not clear where this was mentioned, that the baseline method already uses scratchpads. I would request the authors to provide a bit more details about how that baseline training was done (or kindly point me to a section in the main paper / appendix where more details are already provided), as that is quite relevant to understand the experimental setup and the significance of the results. I also do not quite understand why index hinting would not be available at inference time. It can always be fed. I understand that index hinting is costly (which is precisely why a comparison would be useful), but I don't quite understand why you cannot design the training in a way where models learn to use index hints during training and then during inference also (conditioned on the input being fed with the index hints), then use the same to solve the task. For example in Zhou et al [1], they were able to deploy index hints quite well, but maybe I am missing some details that make it hard to apply in general.

Additionally I would like to acknowledge that I was not aware of the paper Awasthi and Gupta [2], which I only have become aware of because of reviewer hBbS, and the authors response to the same. I would like to thank both the reviewer and the authors for the same. While I acknowledge that there are certain additional improvements in the current work from that paper, the originality of the work is severely impacted by the fact there is a significant overlap with [AG23]. In my head, the strength of the paper is slightly weakened by the extent of the overlap, despite there being improvements.

[1] Zhou et al, What algorithms can transformers learn ? [2] Awasthi and Gupta, Improving Length-Generalization in Transformers via Task Hinting

We thank the reviewer for their suggestions and comments. We are happy that you concur with the fact that this is a key problem on which progress is essential, and that you also see the promise of our approaches to getting better length generalization.

Please find responses and clarifications to your specific questions and comments below.

Link between Theoretical Results and Training Methodology: Recall that we focus on sparse functions. For this we consider $(N,n)$-perturbations (which are sequences of length $n$ produced by suitably sub-sampling $N$-length sequences), and train for both the auxiliary and main tasks on these sequences. Given a problem, we do the following

a. Identifying Auxiliary Tasks: For sparse problems with multi-token outputs, each token at timestep $t$ is obtained by applying some “local” function $f_t$ to a small number of previous tokens. Given this sparsity, there can only be a bounded number of such functions; moreover, for many common tasks, there are just a handful of functions being repeatedly applied, e.g., finding the minimum and successor in sorting, or doing single-digit addition in the increment task, or summing two tokens in the SLiM task, etc. These commonly-used functions for a particular problem are natural auxiliary tasks for that problem.

b. Robust Training for Auxiliary Tasks : Having identified these auxiliary tasks, we train the model to be $(N,n)$-robust on auxiliary tasks (in addition to training on the main task). Recall that this training is on $n$-length sequences.

c. Robustness implies Length Generalization: We can now use our theoretical guarantees, which say that if we are $(N,n)$-robust with good parameters, then length generalization happens (i.e., low error on $N$-length sequences).

We hope this clarifies the setup. We will make this more explicit in the next version.

Clarification of Sorting Task: For sorting, the auxiliary task is as follows: given an input unordered sequence $S$ of tokens (which in our case are numbers modulo 100), and a specific “query” token $q$, the auxiliary task is to lookup the smallest number in $S$ that is greater than $q$ (“successor”). Note that this “lookup” task is precisely what sorting repeatedly performs, since it outputs the minimum element, and then repeatedly outputs the successor of the last-output token. We will clarify this in the revision.

Motivation for the SLiM Task: The motivation for the SLiM task was to serve as a simple example of contextual recall, where each number pays attention to the closest previous number of the same “type” and sums the two. This gives us a precisely defined task on which we can objectively measure length generalization, while capturing some of the essential properties of language (i.e., contextual lookup and then some transformation of the tokens looked up this way). One could consider more sophisticated probabilistic versions, e.g., each token $a_t$ pays attention to multiple previous tokens having some types (that depends on $a_t$’s type), and then outputs the next token chosen from a probability distribution. This is another step closer to natural language and we leave this for future work.

We thank the reviewer for their encouragement and comments. We are glad that you liked the theoretical justification of the training methodology we propose, and that the experiments showed significant gains when baseline training methods failed.

Please find responses and clarifications to your specific questions and comments below.

Relationship to Chugtai et al: We were unable to locate a 2023 paper by "Chugtai et al." on this topic and suspect the reviewer may be referring to Awasthi and Gupta [AG23], given the similar problem space. We sincerely apologize for this oversight in our initial submission and will, of course, add a detailed comparison to our Related Work section.

Our work presents three strict improvements over [AG23]:

1. Theoretical Justification: In [AG23], there was no formalization of why task hinting would lead to length generalization. We, on the other hand, provide a theoretical framework to show why auxiliary tasks help with length generalization.

2. Theory suggests new Training Recipe: Our theoretical analysis leads to a robust training recipe, which was not present in [AG23].

3. Improved Performance via Robust Training: The new robust training recipe leads to better empirical performance across a variety of tasks. For example, for the increment task [AG23] length-generalized (i.e., obtain more than 90% test accuracy) only by a factor of at most 1.2. In our work, however, we can length generalize to a factor of more than 1.5 for the same task.

Use of Synthetic Tasks: Synthetic tasks are standard in the literature on length generalization as it is easy to clearly define the task for various input lengths ($n$). Moreover, we wanted to focus the investigation on tasks with precise answers, like sorting, increment, etc., where the effectiveness of length generalization is concretely and objectively measurable. In fact, the SLIM task is an attempt at a concrete problem capturing “contextual recall”, where each number looks for the previous number of the same “type” and sums them.

Choice of Auxiliary Functions: This is a great question. While the choice of auxiliary tasks currently does depend on the specific problem, we believe this is natural for problems with precise answers. Indeed, such problems require particular tasks to be done correctly and repeatedly, and hence it makes sense to robustly train for them. If we want to consider multiple problems simultaneously, one can go beyond the naive approach of training on the auxiliary tasks for each of these problems by use a Boosting/Hedge-based approach, where we adaptively train the model on a small set of auxiliary tasks on which the model has the worst performance.

It is a very interesting question about how to isolate these auxiliary tasks automatically: one way would be to have a “library” of “basis” functions for the “local” tasks, and then use these as subroutines in a general training setup where they can hopefully enable length generalization for a variety of downstream applications. Thank you for raising this point; we will add it to the discussion part of the revision.

Comparisons to Other Approaches: In the revision, we will add the comparisons to the mentioned approaches. Please note that we did not add those because the approaches are complementary to ours: curriculum training increases the complexity of the training data, whereas architectural constraints try to enforce additional structure to achieve their results. In our framework, our focus is quite different.

Extensions to “Fuzzy” Problems: This is a great question. Our focus has been on “algorithmic” problems for which the auxiliary tasks can be naturally identified, and where length generalization can be precisely and objectively measured. Language modeling is a more amorphous task, and we feel that your question will require further research. E.g., one can imagine that for languages for which we do not have huge data corpuses, we should be able to use its grammar to define auxiliary tasks, and train the models to be better at handling intricate language constructs. (Our SLiM task tries to capture the kind of contextual recall typical in natural language.)

Further Experimental Comparisons: We currently do not have experimental comparisons to curriculum training, since we consider it orthogonal to our focus. We will mention this in the “Limitations” section.

We thank the reviewer for their encouragement and comments. We are glad that you liked the general framework, the fact that the experiments achieve the significant gains as suggested by the theoretical model. Please find responses to your specific questions and comments below.

Clarification of the Setup: Your statement regarding the problem setup is correct: we focus on sparse functions, consider $(N,n)$-perturbations (which are sequences of length $n$ produced by sub-sampling $N$-length sequences), and train for both the auxiliary and main tasks on these sequences. To complete the picture, the steps are as follows:

a. Identifying Auxiliary Tasks: For a sparse function, we produce multi-token outputs, where each token at timestep $t$ is obtained by applying some “local” function $f_t$ to a small number of previous tokens. Given the sparsity, there can only be a bounded number of such functions; moreover, for many common tasks, there are just a handful of functions being repeatedly applied, e.g., finding the minimum and successor in sorting, or doing single-digit addition in the increment task, or summing two tokens in the SLiM task, etc. These commonly-used functions for a particular problem are natural auxiliary tasks for that problem.

b. Robust Training for Auxiliary Tasks : Having identified these auxiliary tasks, we train the model to be $(N,n)$-robust on auxiliary tasks (in addition to training on the main task). Recall that this training is on $n$-length sequences.

c. Robustness implies Length Generalization: We can now use our theoretical guarantees, which say that if we are $(N,n)$-robust with good parameters, then length generalization happens (i.e., low error on $N$-length sequences).

Relationship to other Approaches: We feel the ideas in our approach and in curriculum training are orthogonal. The latter focuses on training on progressively more complex instances (e.g., the approach of [1] buckets the sequences by their lengths, and then trains on some carefully chosen mixture of these sequences). Our approach is to identify useful building blocks (i.e., the auxiliary tasks) that can improve the training performance, and then to perform robust training for them (in addition to the main task). We feel the two approaches could be fruitfully combined in the future. We will add a remark to this effect in the revision.

Data Usage Identical: Yes, the amount of data is the same for the two approaches. So the baseline model is in total using the same amount of data as the proposed model (with data for main task + data for auxiliary task. On the point of unique data, it is an interesting comment: in fact, that is to some extent what our theory proposes—if we show more robust trajectories (which will be unique) to the model then it will length-generalize.

Choice of Auxiliary Functions: This is a great question. While the choice of auxiliary tasks currently does depend on the specific problem, we believe this is natural for problems with precise answers. Indeed, such problems require particular tasks to be done correctly and repeatedly, and hence it makes sense to robustly train for them. If we want to consider multiple problems simultaneously, we can definitely go beyond the naive approach of training on the auxiliary tasks for each of these problems—we can use a Boosting/Hedge-based approach, where we adaptively train the model on a small set of auxiliary tasks on which the model has the worst performance.

It is a very interesting question about how to isolate these auxiliary tasks automatically: one way would be to have a “library” of “basis” functions for the “local” tasks, and then use these as subroutines in a general training setup where they can hopefully enable length generalization for a variety of downstream applications. Thank you for raising this point; we will add it to the discussion part of the revision.

Clarification of Auxiliary Task for SLiM: For SLiM the auxiliary task is as follows: given a valid sequence $S$ of tokens (produced as in the original problem), and a specific “query” token $q$, the auxiliary task is to lookup the last token $r$ in the input sequence $S$ of the same type as $q$, and to output $(q+r)$ modulo $100$. Note that this contextual lookup/recall is precisely the task the SLiM has to repeatedly perform using the last output token as its query, which is why it was chosen as an auxiliary task. Moreover, since contextual lookup/recall is an essential part of language, we think of this as a first yet concrete step towards understanding length generalization in broader contexts.