Improving length generalization in transformers via task hinting

Thanks for the review. We address your points below.

It is not a surprising phenomenon that simultaneously training the model with some related task can help to improve the generalization capability of transformers.

We argue that the situation is more nuanced. While it does seem natural that related tasks should help, we found it surprising that some tasks that a-priori seem quite natural do not help length generalization, whereas others do! In particular, for the sorting problem the successor task helped the most, whereas the counting and fill tasks (which are also natural subtasks) did not help almost at all. Our investigations in Sections 3 and 4 were directly motivated by this surprising difference in performance: we aimed to provide a more detailed analysis of why some particular tasks were so useful.

Another thing that surprised us was the extent to which task hinting helped. For instance, the accuracy jumped from ~0% to 92% when considering length 100 sequences. In fact, even for sequences with length 500, our models gave outputs with very small edit distances, having been trained on sequences of length only up to 20! These are the first instances we know of where such extreme length generalization has been observed.

The proposed method to solve length generalization is not generic. For any specific task, "task hinting" requires researchers to search for a good task that could improve length generalization. Searching for such a task may be complicated and highly task-dependent.

We agree with you that in general, the problem of finding appropriate tasks may itself be quite challenging. (E.g., we discuss this in the limitations section in Appendix A.) Our hope in doing a comprehensive study for the sorting problem (and particularly for doing the interpretability analysis) is to better understand when and how task-hinting is helpful. Based on our analysis in Section 4 we posit that tasks that align well with the internal learning biases of the model are the ones that help the most. Understanding automated/algorithmic ways to find such tasks for an arbitrary setting is a fascinating direction for future research.

The authors only studied the sorting task (and another task hidden in the appendix). It is unclear if "task hinting" could be easily developed as a generic method for solving length generalization.

To demonstrate the effectiveness of our proposed task-hinting framework beyond sorting, we have added results for the LEGO task, which has been studied in recent works in the context of generalization. These results are given in Appendix C.2, and again demonstrate the benefits of task hinting. Briefly, the LEGO task consists of a variable resolution problem over 12 variables where during training the model is only provided supervision for the first $n=6$ variables. We then measure the accuracy of the model in predicting the correct values for variables 7 to 12. For easier readability we include the results in the table below as well. Please see Appendix C.2 for more details.

Thanks for the review. We address your points below.

It's unclear if the same performance gain could be achieved via curriculum learning. In some sense, the auxiliary task can be viewed as a partial sorting, i.e., so it's a simpler task than the original target task.

Thanks for your suggestions for other methods such as curriculum learning and scratchpad training. Typically in curriculum learning, one has to progressively include sequences of higher and higher lengths. Our goal instead is to understand whether we can generalize from training on short sequences alone (length up to 20). Nonetheless, we have included curriculum learning as a baseline in the updated draft.

Regarding scratchpad training, the work of Anil et al.’22 observed that fine-tuning on scratchpad data only helps length generalization to a limited extent. Is there a particular scratchpad that you had in mind? We had trouble coming up with meaningful scratchpads for the sorting problem that were length-efficient. (E.g., we considered adding all the steps of some sorting algorithm like merge-sort or insertion-sort as a scratchpad, but these blow up the input length considerably, making it super inefficient for both training and inference. Our approach of training via subtasks does not suffer from this issue.) In Appendix E we have added a comparison to a length-efficient scratchpad—however, this scratchpad is not particularly different from the natural input/output representation. We would be interested in hearing what scratchpads you think might be useful to compare to.

To summarize, we followed your suggestion and added comparisons to the above two methods. The results are in Appendix E of the updated draft. These experiments show that task hinting is the only method in this collection which achieves non-trivial performance for sequence lengths beyond 50. For easier readability the results are also included in the table below.

Could you clarify the role of the mask values in Figure 3.1? Why are there mask values of 0 after some 1's? Not sure why you need to pad when training decoder models.

A mask value of 1 means that we penalize the predictions made at the corresponding position during the loss function computation. Since we use padding, we use a mask value of 0 to avoid adding terms corresponding to these padding positions into our loss function. (The use of padding is simply for ease of training: the padding makes each training example be the same length.) We will clarify this further in the final version.

Could you provide more details on the fill hint task? Its current explanation is unclear.

Thanks for the question. The motivation for the fill hint task is the following: to sort data with repeated elements, a sorting algorithm must be able to keep track of the number of copies of an element that have been output already, and the number still to be output. Instances of the fill task consist of a number, say 17, repeated some number k (say k=5) times, followed by a prefix consisting of the same number 17 repeated r < k times. The goal of the task is to teach the model that k-r copies of the number 17 still need to be output. We will make this discussion clearer in the final version.

Does this idea work for tasks other than sorting? Could you apply the same to many other tasks?

To demonstrate the effectiveness of our proposed task-hinting framework beyond sorting, we have added results for the LEGO task, which has been studied in recent works in the context of generalization. These results are given in Appendix C.2, and again demonstrate the benefits of task hinting. Briefly, the LEGO task consists of a variable resolution problem over 12 variables where during training the model is only provided supervision for the first $n=6$ variables. We then measure the accuracy of the model in predicting the correct values for variables 7 to 12. For easier readability we include the results in the table below as well. Please see Appendix C.2 for more details.

Can you use the factor analysis (e.g., those introduced in the recent paper [1]) for the interpretation part? Especially see the induction head analysis of this aforementioned paper.

Thanks for the reference! This work seems to be subsequent to the submission of our paper. We will definitely look into it and investigate whether this analysis can be applied to our setting to further demystify the mechanisms of the learned models.

Thanks for the review. We address your points below.

Multi-task is not novel. How / why it can help length generalization is still unclear in this paper

We do not claim that we are proposing multitask learning as a new paradigm. Our main contribution is the idea of using auxiliary task-hinting to improve length generalization. In order to incorporate auxiliary tasks into the training setup, we use multitask learning via fairly standard setups. This is already discussed in Appendix A.

As to the question of how/why it helps length-generalization: this is exactly the goal of Section 4. That section is devoted to understanding why task hinting helps, and why certain tasks are better than others. We use our interpretability analysis to posit a concrete conjecture, which says that tasks which align well with the internal learning biases of the model are the ones that help the most for length generalization. We show some corroborating evidence for the case of the sorting problem.

This paper is not organized well. For example, at the beginning, authors try to categorize length generalization into two categories. However, the taxonomy does not make sense and it cannot fit the task-hinting well

Perhaps there is some confusion here. Our goal is to characterize existing approaches to length generalization into two categories: (a) approaches that finetune an existing pre-trained LLM and (b) those that consider specific problems and train models for these specific problems from scratch. In this sense, our work belongs to the second category: it studies specific problems and trains models from scratch. We can reword this in the final version to further clarify this.

The results of tempered softmax are not surprising. With longer sequences, self-attention requires scaling is well-known. (The hugging face transformer library has even involved this feature). To be more specific, the scaling factor should be "log_n scale'.

Tempered softmax is a natural operation that comes out of our theoretical analysis: we then make this a trainable parameter for each layer in our experiments. However, this investigation of tempered softmax is secondary to the main results of the paper. If an existing library indeed incorporates it, we view this as a positive for our theoretical work, which provides further evidence for it being a good idea. On that note, we searched the HuggingFace library to find uses of “log_n scale” which you mentioned, but we were unsuccessful. If you could please point us to the right online references, we would be happy to connect our work to it, and also provide further clarifications on any similarities and differences.

This paper is not completed within the required number of pages by ICLR.

We respectfully disagree. The main paper is within the first 9 pages followed by the citations and the appendix. The call for papers clearly permits the use of an arbitrary number of pages for the supplementary material in the appendices.

Thank you for your comments and for your appreciation of the work. To demonstrate the effectiveness of our proposed task-hinting framework beyond sorting, we have added results for the LEGO task, which has been studied in recent works in the context of generalization. These results are given in Appendix C.2, and again demonstrate the benefits of task hinting. Briefly, the LEGO task consists of a variable resolution problem over 12 variables where during training the model is only provided supervision for the first $n=6$ variables. We then measure the accuracy of the model in predicting the correct values for variables 7 to 12. For easier readability we include the results in the table below as well. Please see Appendix C.2 for more details.

We agree with you that in general, the problem of finding appropriate tasks may itself be quite challenging. (E.g., we discuss this in the limitations section in Appendix A.) Our hope in doing a comprehensive study for the sorting problem (and particularly for doing the interpretability analysis) is to better understand when and how task-hinting is helpful. Based on our analysis in Section 4 we posit that tasks that align well with the internal learning biases of the model are the ones that help the most. Understanding automated/algorithmic ways to find such tasks for an arbitrary setting is a fascinating direction for future research.

Regarding positional embeddings, we did not find any changes in the efficacy of the tasks when positional information is present vs. when it is not. Indeed, for the sorting problem, the successor task continued to be the best among the three types of tasks considered.