Task Diversity Shortens the ICL Plateau

Authors: Jaeyeon Kim, Sehyun Kwon, Joo Young Choi, Jongho Park, Jaewoong Cho, Jason D. Lee, Ernest K. Ryu
Primary area: foundation or frontier models, including LLMs
Keywords: In-context learning, Transformers, Training dynamics, Task diversity, Multi-task

Abstract:
In-context learning (ICL) describes a language model's ability to generate outputs based on a set of input demonstrations and a subsequent query. To understand this remarkable capability, researchers have studied simplified, stylized models. These studies have consistently observed long loss plateaus, during which models exhibit minimal improvement, followed by a sudden, rapid surge of learning. In this work, we reveal that training on multiple diverse ICL tasks simultaneously shortens the loss plateaus, making each task easier to learn. This finding is surprising as it contradicts the natural intuition that the combined complexity of multiple ICL tasks would lengthen the learning process, not shorten it. Our result suggests that the recent success in large-scale training of language models may be attributed not only to the richness of the data at scale but also to the easier optimization (training) induced by the diversity of natural language training data.

Decision: reject
Reviewer scores: 3, 6, 3, 6

Reviews and discussion:
<thread>
creation date:
1729996983797

thread_type:
official_review

<thread_object>
creator: reviewer_soyu

<content>
summary:
This paper studies how multi-task learning improves in-context learning on individual tasks. Through experiments on synthetic tasks, it shows multi-task training with in-context examples shortens the loss plateaus of single in-context learning tasks, making the learning easier. The phenomenon also apears for other model architectures (Mamba) and tasks (simple natural language tasks).

strengths:
1. The paper studies the learning dynamics of the in-context learning ability of language models, which is an important direction for understanding how language models obtain their capability.

2. "Task Diversity Shortens the ICL Plateau" is an interesting and valuable insight for designing new methods to accelerate the learning of LMs.

3. The experiments are sufficient to verfy the claims in the paper (at least in the toy settings).

weaknesses:
1. An important property of ICL training is its generalization to unseen tasks [1,2]. This paper only shows that increasing the task diversity improves the learning tasks also in the training set. It would be better to examine how the diversity of ICL training tasks affects the ICL on unseen tasks.

2. The training tasks seems to have equal weights in the experiments (by sampling the same number of instances and balancing the losses). However, the tasks usually follow a long-tail distribution in the real-world pre-training corpus. Discussion on the setting where the tasks are unbalanced whould make the conclusions more close to the practical scenarios.

[1] MetaICL: Learning to Learn In Context. 2022. In NAACL.

[2] Pre-Training to Learn in Context. 2023. In ACL.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730035000169

thread_type:
official_review

<thread_object>
creator: reviewer_6hyd

<content>
summary:
This paper examines the phenomenon of *in-context learning* (ICL) in language models, where a model learns from example inputs and outputs within a context to respond accurately to subsequent queries. Researchers have found that during ICL, models often undergo periods of minimal learning progress (“loss plateaus”) before rapidly improving. In this study, the authors show that training on a variety of diverse ICL tasks concurrently reduces these loss plateaus, making each task quicker and easier to learn. This result challenges the expectation that increased task complexity would slow down the learning process. The authors suggest that the enhanced performance of large-scale language models could be due to both the richness of large datasets and the facilitation of learning through task diversity, which helps streamline optimization.

strengths:
This paper is clearly written. The figures and captions are easy to follow, and the motivation is clear.

weaknesses:
The primary conclusion of this paper, “Task diversity shortens the ICL plateau,” may appear somewhat trivial. The study demonstrates that training models with multiple tasks simultaneously involves using significantly more data (i.e., training with n tasks simultaneously implies utilizing n times the amount of training data). Though these data sets pertain to different in-context learning (ICL) tasks, the observed improvements in learning efficiency could stem from the increased data volume rather than solely from “task diversity.” In this sense, the results align with well-established concepts of “knowledge transfer” in multitask learning, where it is well-documented that similar tasks can support each other’s learning. This paper essentially reinforces that ICL tasks share underlying similarities, which is trivial.

Additionally, it’s unclear how this finding can be practically applied to large language model (LLM) fine-tuning. If task diversity is indeed advantageous, does this imply a need to incorporate task-irrelevant data when fine-tuning a task-specific LLM? Clarifying this connection could greatly enhance the practical relevance of the paper’s conclusions.

questions:
Please refer to my weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730347703887

thread_type:
official_review

<thread_object>
creator: reviewer_uzef

<content>
summary:
In this work, the authors explores how training models with ICL examples from multiple tasks simultaneously can improve the rate at which models learn the different individual tasks, even faster compared to when learning on one individual task. In particular, the authors characterize through a phenomena of the loss plateau, which is a scenario where the model's loss plateaus for a significant number of steps before observing sudden large decreases, by showing that multi-task ICL learning significantly shortens this plateau duration. They use this to conclude that this task diversity appears to aid optimization, making ICL training more efficient.

strengths:
The manuscript very well written; the division of the paper makes it very intuitive to follow and the plots/tables are informative. The main findings are also quite interesting in how they present a setting where (if all holds) then learning various tasks at one can potentially improve the efficiency of learning a marginally or completely un-related task. The notion of measuring how the loss plateau contracts is a rather novel way of investigation and I see avenues where it could be useful for exploring the idea of other types of emergent properties in language models.

weaknesses:
Some points made by the authors do not have completely solid grounding in the results or lack some support.

For example, on L16 and L199-202, the authors state that "multi-task ICL is easier to learn than single-task ICL is surprising as it contradicts the natural intuition that the combined complexity of multiple ICL tasks would lengthen the learning process". I'm not particularly convinced of this notion, at least in the regime within which the authors present their results. The sample is broken up into tokens, therefore the explicit rules of the functions that are being sampled for the in-context examples are likely not being directly captured in the sequence being presented. Additionally, since samples are consist of examples from a fixed function (that differs between samples), learning multiple tasks together can serve to learn more robust representations that won't as quickly over-fit to individual samples [1]. 

I also believe that to arrive at the author conclusions, further baselines need to be included. More explicit control over the exact sampling of tasks from $\mathcal{F}$, as well as the sampling of the parameters associated with the specific instance of $f$ should be necessary. For example, tasks can be sampled with uneven probability; if the shortening is consistent among different task classes but with the same sampling probabilities, then this can serve as a more explicit signal of it directly being due to the use of multiple tasks rather than a set of very few pairs of tasks that observe greater transfer compared to others.

Another particular concern of mine is that the paper is lacking in any specific formal analysis of the optimization landscape between the two settings. From my perspective, this causes a significant portion of the results, while interesting, to appear as speculative at best, which combined with the lack in explicit diversity among the different ICL tasks (either regression or boolean) to suggest a possible level of difficulty of extrapolating these observations to more general scenarios. I appreciate the examples provided in Appendix C on other natural language tasks, but given the fact that the results are somewhat inconclusive with regards to what settings multi-task ICL can help me I believe it's quite essential to investigate the differences from a optimization perspective.

In summary: I did enjoy quite a few points about this paper, but the authors need to reformulate some critical points (or provide details to demonstrate that those are not critical). If the necessary clarifications can be provided, I am open to changing my opinion of the work but as of the moment I believe there still remains some work to be done to solidify the details in it.

-----

[1] Sebastian Ruder. An Overview of Multi-Task Learning in Deep Neural Networks. arXiv preprint, 2017.

questions:
- Perhaps I am overthinking, but I want to just make sure my understanding of the training mechanism is correct. I understand that in the multi-task setting, $\mathcal{F} = \bigcup_{i=1}^k \mathcal{F}_i$ from which $f$ is sampled. Then $x_i,\dots,x_n$ is sampled from the domain of the sampled $f$. Therefore the training (for any given sample) is on any arbitrary task in the set of tasks that were selected. If this understanding is correct, please ignore the next two comments.
- Is $n$ (the number of samples from each function class) fixed at the same value for all values of $k$, or do they scale together; I apologize but I couldn't find this exact detail either in the main text or in the Appendix.
- I believe it would be useful to attempt an experiment where you can sample different functions from the same function class for constructing the tasks and see how this affects learning.
- L157: Why this complexity measure was chosen isn't particularly justified. Importantly, although you measure based on the number of learning steps, there are a number of different 
- For Section 4.2, I believe that the authors should be able to use synthetically generated data where an explicit common structure exists to better demonstrate the claims.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730650604617

thread_type:
official_review

<thread_object>
creator: reviewer_aaqm

<content>
summary:
The paper explores the impact of multi-task learning ("task diversity" in the title) on the speed of training on individual tasks in the context of In-Context Learning (ICL) of various language-modeling architectures (Transformer, Mamba, Hyena). Quite surprisingly, the results suggest that multi-task setup leads to faster training/convergence of individual tasks, compared to training on those individual tasks.

The experiments show primarily synthetic tasks, but language tasks are also included as validation. The paper also investigates different explanations for its findings, e.g., finding that the creation of Induction Heads on its own doesn't explain the results.

strengths:
The paper showcases a good number of experiments with solid results. I also appreciate the explanation of the experimental setup, and detailed experiments with multiple tasks and models. The authors also present a good analysis of some possible explanations for the results achieved. The presentation is great, and so is the writing.

weaknesses:
The most significant weakness of the work is the construction of the single-task baselines. In my opinion, this issue alone should justify the rejection of the paper. In short - **the experiments compare single-task models trained with batch size B with multi-task models (with k tasks) with batch size kB - k times larger batch size, and k times larger the training set/compute cost!** Then, the authors use the total number of iterations (batches) to compare models. While this is mentioned in the paper (line 177 and Appendix A), I think results should be shown with a more standard x-axis showing the number of tokens and not the number of iterations.

While the results, as initially presented, sound surprising, the fact that **multi-task models train on multiple times as much data compared to the baseline** makes the results fairly trivial or at least extremely easy to achieve positive results. Imagine this experiment: instead of training multi-task model on different tasks F1 and F2, let's set F1 to be equal to F2 (the same task twice). Such a "multi-task" model is then trained on simply a 2x bigger batch size than the baseline, and so, obviously, it converges faster in terms of iterations (but probably not faster in terms of tokens!) Of course, authors set different F1 and F2, but they are fairly similar, so I expect the results the authors get by default (if the model is not undersized wrt. problem complexity).

The paper shows model trained on 6 tasks jointly (using 6x bigger batch size). Then, convergence between 1.4x and 9.2x faster in terms of iterations (see Figure 1) is no longer surprising! In terms of tokens, that would mean a relative speed between 0.23x and 1.53x, which means multi-task learning is between **4x slower** and 1.5x faster. (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.)

I'd expect either comparison on the same number of tokens - preferably with the same total batch sizes for both single-task and multi-task settings, as training with different batch sizes and number of batches makes the setup prone to mistuning the baseline on accident.

While there is a possibility that I don't understand something fundamental here, I'm inclined to reject the paper for the above reason. With proper comparison metrics, I believe the results are fairly trivial. What is worse, the paper in its current state is a bit misleading.

questions:
As stated in the weakness section.
1. Why was a comparison on iterations, not total tokens, chosen?
2. Were baselines tuned for their smaller batch sizes?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732253975857

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We sincerely thank all the reviewers for their invaluable feedback. We were happy to see that the reviewers overall agreed that our presented main observation is interesting. We have addressed the individual concerns in the individual responses.

In this common response, we reiterate the main value of our paper. Our discovery shows that training on diverse ICL tasks is more effective than training on a single ICL task, from an **optimization perspective**. This contrasts with prior notions such as "knowledge transfer” in multi-task learning which is a **statistical perspective**. In other words, it is known that multi-task training leads to better/robust features and therefore better generalization **after training** has been completed. However, our results concern the dynamics **during training**. We believe that the observation that multi-task training reduces training time is novel.


In relation to LLMs, pre-training via next-token prediction on a diverse corpus can be regarded as a highly diverse multi-task training. In light of our observations, we argue that the success of LLM training may be attributed not only to the **statistical** richness of the multi-task training data but also to the improved **optimization** dynamics induced by the diversity of the pre-training corpus.


We ask the reviewers and the area chair to consider whether our results are novel, valuable, and convincing. We argue that our presented results do have such qualities.
</content>
</thread_object>

</thread>


Metareview:
metareview:
## Summary: 
The paper investigates the effects of multi-task learning on the training speed of individual tasks when training models with In-Context Learning (ICL) tasks (synthetic tasks in the paper). Surprisingly, the results indicate that training on multiple diverse ICL tasks concurrently can accelerate the learning process for individual tasks compared to training on them individually. This approach shortens the loss plateaus typically observed during training, contrary to the expectation that increased task complexity would prolong learning. The study suggests that task diversity in training not only aids optimization but also contributes to the efficiency of training Transformer models by facilitating faster convergence on individual tasks. However, the experiments in this paper are not conducted on language model tasks but are limited to a few categories of synthetic tasks.  

## Strengths:
1. The study provides valuable insights into the learning dynamics of language models when trained with single ICL task vs. multiple ICL tasks.
1. The paper presents extensive experiments on toy and synthetic tasks, whose detailed setup and analysis of multiple tasks and models enhance the paper's credibility.
1. The presentation and writing are of high quality, making the content easy to follow.

## Weaknesses:
1. The original draft applied a much larger (K times, for K tasks) batch size for multi-task learning than single-task learning, while the loss curves merely show how the loss changes with the number of batches (iterations). This makes the shortened loss plateau a questionable and misleading claim since multi-task learning consumes much more training tokens than single-task learning at the same number of iterations. The authors corrected this issue in the rebuttal by providing updated curves of loss vs. number of samples, which shows less speedup on fewer categories of tasks by multi-task learning (4-5x speedup on 3/6 task categories and infinity on Parity tasks). While the authors' quick action of correcting the issue in a limited time frame is greatly appreciated, the new result weakens the generalizability of the major claim to some extent. 
1. A formal analysis of the loss landscape from the optimization perspective is not provided, though a brief discussion of the common structure is provided. The paper lacks a theoretical explanation of the observed phenomenon. 
1. The analysis in this paper is limited to a few specific categories of synthetic or toy tasks. It is not clear whether the observation can be generalized to other categories of tasks such as the ones on LLMs. 
1. Solely increasing the task diversity may introduce negative transfer in multi-task learning. It is necessary to discuss the transferability between selected tasks and why/when a negative transfer is not harmful in the proposed setting. 
1. The paper only studies the shortened loss plateau on the training tasks rather than unseen tasks, which is more related and better aligned with the interests of the foundation model or LLM researchers. 
1. The experiments applied even weights to multiple tasks. The authors' rebuttal provided experiments using uneven weights on which the conclusion still holds. More comprehensive experiments of different uneven weights can strengthen the conclusion further.  

## Decision:
The authors provided detailed clarifications and various additional experimental results in the rebuttal. Three out of the four reviewers responded to the rebuttal and confirmed that some concerns had been addressed promisingly. However, they also have some remaining concerns that cannot be fully addressed by the authors' responses. One reviewer raised the rating but also lowered the confidence, leading to two 3s and two 6s. The meta-reviewer carefully read the paper and all the discussions. While the observation that multi-task learning speeds up learning is interesting and relevant to foundation models such as LLMs, the limitation of experimental settings and tasks in this paper still leads to gaps in achieving a general conclusion. Moreover, a more rigorous and in-depth theory behind the observation is preferred, especially in explaining the difference between existing theories of multi-task learning and transferability. The experiments and analysis provided by the authors in the rebuttal and discussion can greatly strengthen the paper if a more complete version can be included. Although the current draft is still not ready for publication due to the above reasons, the meta-reviewer believes an improved version can present an important contribution to the community.