Task Diversity Shortens the ICL Plateau

We thank the reviewer for bringing in the prior views on multi-task learning into the discussion.

We firmly disagree that the conclusion of our work is trivial. As we discuss in the common response, we maintain that our contribution, an optimization perspective, differs from the prior statistical perspective of multi-task learning.

We kindly ask the reviewer for further clarity on the following points.

1. The Reviewer 6hyd references the concerns of Reviewer aaqM. Can the reviewer comment on whether they find our rebuttal to Reviewer aaqM convincing or not? If convincing, perhaps the initial criticism is no longer relevant. If not, we would like to understand why.

We also remark that looking to other reviewers to find additional points of criticism that the reviewer did not originally raise may not be informative, especially when our rebuttal to the said criticism is not addressed. If the reviewer has new questions or new points of criticism, we will happily address them. (Also, as stated in our response to Reviewer aaqM, the main experimental results have been updated accordingly.)

2. We apologize for not addressing this question. We initially misinterpreted the remark as a rhetorical point. We now know that the reviewer considered it a crucial question in need of a response. Our answer is yes. One of the key messages of our paper is: Even if you are only interested in ICL task A, it is beneficial to jointly train ICL tasks B and C because doing so speeds up training. Within the context of ICL, our answer is a very direct and unequivocal yes, and we provide the empirical evidence to back up the claim.

3. First of all, we point out that the reviewer's claim that "The primary conclusion of this paper, ... appear somewhat trivial" is a very strong claim that should be substantiated with evidence or an argument. The reviewer does provide an argument, but we point out that the reviewer's argument is statistical, while our main claim is in the sense of optimization. As we elaborate in the common response, this distinction between the statistical and optimization claims is crucial. We wonder if the reviewer agrees that this distinction is a meaningful one.

The reviewer acknowledges using "GPT-4o to identify relevant references" due to being "not an expert in this domain". While leveraging search tools—whether based on traditional IR methods or LLMs—is undoubtedly useful for identifying relevant prior work, I hope the reviewer understands that this comment makes us feel a little bit disheartened. Ideally, one expects reviews to reflect expert opinions (informed by search results).

Substantively, the claim that the prior conclusions of [1] (in particular, Figure 1 in Appendix C of [1]) overlap with our main findings seems to be a wholly hallucinated one.

Specifically, [1] does not consider ICL, but does consider multi-track training, and starts from the claim that naive multi-task training has the problem of overfitting. The goal of the paper is to remedy this overfitting, but, in any case, the work does not make any comparisons with single-task training, so the findings have no bearing on whether multi-task training is easier or harder than single-task training.

We are very happy to hear that the reviewer found our work interesting and valuable. The individual comments are addressed in the following.

Weakness 1. As our paper concentrates on the optimization itself rather than the out-of-distribution (OOD) generalization capability, we did not perform evaluations on unseen ICL tasks. Nevertheless, although this is beyond the scope of our work, we strongly anticipate that the multi-task trained model would do better on the OOD ICL tasks.

Weakness 2. The reviewer makes an interesting point about mixing tasks with uneven probability, so we conducted experiments on this. We selected five ICL tasks and mixed them using probabilities of $\left(\frac{1}{2},\frac{1}{8},\frac{1}{8},\frac{1}{8},\frac{1}{8}\right)$, resulting in five different combinations. Under these uneven mixing conditions, we consistently observed a shortened plateau, further indicating that our findings hold in such cases.

Plateau reduction with uneven sampling

We sincerely thank the reviewer for their insightful comments. We address the concerns in the following.

Regarding the point on the difference between single-task and multi-task training. As we understand, the reviewer commented that our argument is not convincing as the example sequence is broken into different tokens. In our setup, however, since the in-context demonstrations are already real numbers, we use a single linear embedding layer without a tokenizer, following the experimental setup of prior works [1,2]. Therefore, the explicit function rule is accessible to the model through the sequence.

Further baselines. The reviewer made an excellent point that a further baseline comparison would be valuable, as the plateau reduction might be due to very few pairs of tasks that observe greater transfer compared to others. We also thought about this question, which led to Table 1 (line 247). For Table 1, we evaluated all possible combinations of different tasks and averaged the iteration numbers to escape the plateau. Table 1 reveals that all ICL tasks, not just a selected few, tend to benefit from each other.

Additionally, as suggested by the reviewer, we conducted experiments on multi-task training with uneven sampling probabilities. To construct a multi-task ICL, we selected five ICL tasks. These tasks were then mixed using probabilities of $\left(\frac{1}{2},\frac{1}{8},\frac{1}{8},\frac{1}{8},\frac{1}{8}\right)$, resulting in five different combinations. Under these uneven mixing conditions, we consistently observed a shortened plateau, further indicating that our findings hold in such cases.

Plateau reduction with uneven sampling.

Optimization landscape. While we do not provide a formal/rigorous analysis of the optimization landscape, Section 4.2 provides plausible insights into this phenomenon. Section 4.2.2 demonstrates a common structure across ICL tasks and hypothesizes that task diversity plays a critical role in learning this common structure, particularly from an optimization perspective. Section 4.2.3 verifies the generality of this hypothesis through experiments in the supervised learning setup.

Question 1. Yes! To generate the prompt, for each $i=1,\ldots,m$, we sample $f \in \mathcal{F}_i$ from a designated distribution $\mathcal{D}_F$s. (These are described in Section 5).

Therefore, as noted, $f$ changes across the task, and also across the prompt, thus a model needs to `in-context learn' $f$ through demonstrations $(x_1,f(x_1),x_2,f(x_2),\ldots,x_{n-1},f(x_{n-1}))$.

Question 4: The choice of complexity measures We chose the iteration number as the notion of complexity because it is a directly observable and practically relevant quantity. There are certainly other options, but we would like the reviewer to consider this part of the paper as a very rough intuitive discourse surrounding the newly observed phenomenon.

Question 5: Common structure experiment As the reviewer suggested, we ran the same checkpoint experiments (Section 4.2.2.) for the data with an explicit common structure. We consider a linear in-context regression, $f(x)=w^\intercal x$, but $w$ is sampled from three different distributions: $w \sim \mathcal{N}(0,I_d)$, $w \sim \mathcal{N}(1,I_d)$, and $w \sim \mathcal{N}(-1,I_d)$. We can still observe the shortened plateau, compared to single-task baselines.

Common Structure

We believe these answers clarify the reviewer's concerns. If the reviewer feels that our results are valuable and novel, we kindly ask the reviewer to increase the score.

[1] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. [2] Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and llms by learning to learn discrete functions.

We thank the reviewer for raising very precise points regarding (1) batch size and (2) hyperparameter tuning. We recognize that these are central issues to the evaluation of our work. So, we would like to quickly address Reviewer aaqM’s concerns while we prepare a more comprehensive response for the full rebuttal.

Batch size: We would like to clarify that our intention behind the original batch size choice was to track the sample complexity for the single task in consideration. However, the reviewer raises an excellent point that this comparison would be misleading when two tasks are very similar. We therefore re-ran the experiments with a uniform batch size of $64$ instead of $k \times 64$ for the multi-task ICL. The results are shown in Figure (https://drive.google.com/file/d/1ABDNowiZuhXSEz-jCdyvmiEGamyxWO6u/view?usp=sharing), and we see that the plateau-shortening effect persists, except from the Conjunction setup.

(We also note that running $N$ iterations with batch size $kB$ is not necessarily the same as running $kN$ iterations with batch size $B$, especially when training is stuck in a plateau. So taking the plot of our original submission and elongating the multi-task ICL curve by a factor of k to match the batch size is not necessarily a fair comparison. The only fair comparison is a direct experiment, and it shows that the plateau-shortening effect persists.)

Hyperparameter tuning: In our original experiments, we did no parameter tuning and blindly followed the setup from [1]. We thought that this would be fair, or perhaps even favorable to the baseline, since the single-task batch size of 64 is inherited from [1], and the learning rates of [1] were tuned with the batch size of $64$ and not the batch size of $k × 64$ of the multi-task ICL setups.

In any case, we now see that it would be better to tune all setups, especially to double-check the validity of our new results. The results are shown in Figure (https://drive.google.com/file/d/1GRuMlpSEf7cHqucbv-CJy-367Fvp2eh_/view?usp=sharing), and they are qualitatively the same. The absence of tuning did not lead to confounding or erroneous conclusions.

Parity task and local minimum: The reviewer made the comment: "Although parity tasks break this pattern, with single-task, the model ends up at a bad local minima. I’d expect the issue to be with training stability and insufficient baseline tuning, however." We clarify that the difficulty of training ICL parity tasks has also been reported in prior works [2,3], and this provides us with, in a sense, "social proof" that the training failure is not due to poor implementation.

In fact, while we did not emphasize this point in the writing, we considered to have provided a solution to the open problem of training ICL parity. Some may have thought ICL parity to be too combinatorial and challenging for transformers to learn, but we show the positive result that transformers can learn to solve parity, and our multi-task training scheme elegantly resolves the training difficulty.

Conclusion: In about a week, we will provide a full rebuttal addressing all comments of Reviewer aaqM and the other reviewers. In particular, we will re-run all experiments with equitable batch sizes. We wanted to quickly respond to the main concerns of Reviewer aaqM, as it is central to the evaluation of our work. We maintain that the main claim of our paper is, and we will make our complete case with our updated full set of experiments. In the meantime, we welcome Reviewer aaqM to provide any comments or questions if anything is unclear about the experiments of this intermediate rebuttal.

[1] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. NeurIPS, 2022.

[2] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. ICML, 2024.

[3] Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and llms by learning to learn discrete functions. ICLR, 2024.