Paper: Improving length generalization in transformers via task hinting (ICLR-2024)

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

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

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

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

| n->    | 7 |  8 |  9 |  10 |  11 |  12 |
| -------- | ------- |  ------- |  ------- |  ------- |  ------- |  ------- |
| Standard  | 99.9    |  99.4    |  94.8    |  84.5    |  67.6    |  41.3   |
| Hinting |    99.9     |  99.9    |  98.4    |  91.9    |  79.3    |  56.5    |

Reviewer comment (Experiments):
It's also unclear if the same performance gain could be achieved via curriculum learning.

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

| n->    | 50 | 100 |
| -------- | ------- |------- |
| Standard  | 98    |0    |
| Curriculum  | 98.4    |1.2    |
| Scratchpad | 98.2     |0     |
| Hinting   | 99.8    |92.6    |

Reviewer comment (Experiments):
The comparison between this approach and scratchpad training is not made. What if we provide auxiliary hints as a part of the scratchpad?

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

| n->    | 50 | 100 |
| -------- | ------- |------- |
| Standard  | 98    |0    |
| Curriculum  | 98.4    |1.2    |
| Scratchpad | 98.2     |0     |
| Hinting   | 99.8    |92.6    |

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

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

| n->    | 50 | 100 |
| -------- | ------- |------- |
| Standard  | 98    |0    |
| Curriculum  | 98.4    |1.2    |
| Scratchpad | 98.2     |0     |
| Hinting   | 99.8    |92.6    |

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

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

Reviewer comment (Writing):
Could you provide more details on the fill hint task? Its current explanation is unclear.

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

Reviewer comment (Experiments):
Does this idea work for tasks other than sorting? Could you apply the same to many other tasks?

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

| n->    | 7 |  8 |  9 |  10 |  11 |  12 |
| -------- | ------- |  ------- |  ------- |  ------- |  ------- |  ------- |
| Standard  | 99.9    |  99.4    |  94.8    |  84.5    |  67.6    |  41.3   |
| Hinting |    99.9     |  99.9    |  98.4    |  91.9    |  79.3    |  56.5    |

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

[1] https://arxiv.org/abs/2310.04861

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

Reviewer comment (Novelty):
Multi-task is not novel. How / why it can help length generalization is still unclear in this paper.

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

Reviewer comment (Writing):
The writing is poor. 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.

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

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

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

Reviewer comment (Presentation):
This paper is not completed within the required number of pages by ICLR.

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

Reviewer comment (Experiments):
One major weakness of the paper is that its scope is a bit limited. While it shows the efficacy of using auxiliary data for length generalization, its results are limited largely to the problem of sorting.

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

| n->    | 7 |  8 |  9 |  10 |  11 |  12 |
| -------- | ------- |  ------- |  ------- |  ------- |  ------- |  ------- |
| Standard  | 99.9    |  99.4    |  94.8    |  84.5    |  67.6    |  41.3   |
| Hinting |    99.9     |  99.9    |  98.4    |  91.9    |  79.3    |  56.5    |

Reviewer comment (Experiments):
In this paper, no positional embeddings seem to have been used, does the efficacy of the proposed auxiliary task change when using positional embeddings?

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