Task Diversity Shortens the ICL Plateau

## Abstract

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.

## 1 Introduction

In-context learning (ICL), first reported by Brown et al. () with GPT-3, describes a language model’s ability to generate outputs based on a set of input demonstrations and a subsequent query. In ICL, the model discerns the task implied by the context of these demonstrations without explicit descriptions, indicating that the model may internally implement an *algorithm* or engage in a *reasoning process*. To understand this remarkable capability that emerges in language models trained on complex real-world language data, researchers such as Garg et al. () have studied simplified, stylized models. In these studies, transformers are trained from scratch to learn simple functions in context, such as linear regression. We thoroughly review this prior work in Section 1.1.

An intriguing phenomenon observed in these works is the *long loss plateau* in training for ICL. Throughout these plateaus, models display minimal performance improvement, followed by a sudden, rapid surge of learning—deviating from the typical smooth reduction in training loss. In this work, we reveal that training on multiple ICL tasks simultaneously shortens the loss plateaus as illustrated in Figure 1. 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.

In the language of multi-task learning, our findings present an instance where multi-task learning is *easier* than single-task learning in the sense of training dynamics. While previous research has primarily focused on the *statistical* benefits of task diversity in multi-task learning, our findings reveal *optimization* benefits. This insight 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.

##### Organization.

Section 2 describes the experimental setup for training for multiple ICL tasks. Section 3 articulates our main claim that task diversity shortens the ICL plateau and presents experimental evidence with transformers and state-space models trained on synthetic and natural language ICL tasks. Section 4 investigates the underlying reasons for this phenomenon. Section 4.1 characterizes the model’s behavior during the plateau, which we refer to as the no-context learning regime. Section 4.2 shows that there is a common structure shared across the ICL tasks and that task diversity accelerates the learning of this shared structure. Section 5 presents further experimental details. Section 6 concludes the paper.

### 1.1 Related works

##### In-context learning.

In-context learning (ICL) abilities of pretrained Large Language Models have gained significant attention since first investigated by GPT-3 (Brown et al., ).
Following this, a large body of empirical studies have explored ICL in Large Language Models (Min et al., ; ; Liu et al., ; Nie et al., ; Rubin et al., ; Wei et al., ). Given the complexity of real-world data, researchers have explored ICL in more stylized and simplified setups. Garg et al. () formalized an approach to studying transformer’s performance to learn simple function class in context. Building on this work, several works have investigated the ability of models to learn stylized function classes in various settings, including boolean functions (Bhattamishra et al., ),
regular language (Akyürek et al., ), as well as task mixtures (Tripuraneni et al., ), and the ability of Mamba (Park et al., ; Grazzi et al., ; Li et al., ). For more comprehensive overview, please refer to the survey by Dong et al. ().

While these studies have scrutinized ICL abilities across various setups, most have not delved into the in-context algorithms that are implemented by models. Xie et al. () suggested that the ICL process can be interpreted as Bayesian inference. A number of works have argued, both theoretically and empirically, that transformers implement gradient descent to learn linear regression in-context (Akyürek et al., ; von Oswald et al., ; Mahankali et al., ; Zhang et al., ; Ahn et al., ). Beyond these works, many have sought to uncover the internal ICL procedure of transformers, in more complex algorithms (Fu et al., ; Giannou et al., ; Cheng et al., ; von Oswald et al., ; Lin & Lee, ); when handling more intricate tasks (Wang et al., ; Guo et al., ; Bai et al., ; Lin et al., ; Wang et al., ).

##### Abrupt phase transition in in-context learning.

Many studies (Srivastava et al., ; Wei et al., ; He et al., ; Chan et al., ; Raventós et al., ; Raventos et al., ) have shown that a model’s ability to perform ICL emerges abruptly, with respect to dataset and model size. While these studies do not focus on the training process, abrupt performance gains during training have also been reported in various works. This transition is typically associated with escaping *loss plateaus*. During these plateaus, no performance gains are observed, but once the plateau is escaped, the model begins to learn in-context abruptly. This phenomenon has been observed in various setups (Garg et al., ; Bhattamishra et al., ; Park et al., ; Li et al., ; Kirsch et al., ), although these works did not primarily focus on it.

Beyond these, several studies have explored loss plateaus themselves. Fu et al. () suggested that models learn the features of dataset during loss plateaus. Chen et al. () theoretically demonstrated that a plateau occurs during training when a one-layer transformer is trained by linear regression ICL task. Olsson et al. () proposed that the plateau escape and the formation of ‘induction head’ simultaneously happen. Further work by Reddy (); Singh et al. (); Song et al. (), focused on two-layer transformers, explicitly defining the induction head with transformer parameters and characterizing the internal mechanisms behind its formation. In contrast, Section 4.1 offers a new perspective by providing the explicit form of the model’s output during plateaus. Furthermore, in Section 4.2.3, we provide a discrepancy between our finding and induction head.

##### Task diversity.

Multi-task training (Caruana, ; Baxter, ) is a widely used approach for model pretraining. In particular, researchers have identified that *task diversity* is crucial for pre-trained models to outperform in downstream tasks, alongside a large body of research across various domains: supervised learning (Tripuraneni et al., ; ; Du et al., ; Maurer et al., ; Crawshaw, ; Ruder, ), reinforcement learning (Zhang et al., ; Jin et al., ; Hu et al., ; Yang et al., ; Collins et al., ; Lu et al., ; Cheng et al., ), and natural language processing (Zhang et al., ; Zhao et al., ; Zhang et al., ; Hu et al., ; Song et al., ; Zhou et al., ; Gunasekar et al., ; Sharma et al., ). These works mostly focused on statistical benefits, whereas our finding emphasize on the optimization benefits.

## 2 Experimental setup

Figure: Figure 2: To generate a training sequence, we sample the function $f\in\mathcal{F}_{1}\cup\ldots\cup\mathcal{F}_{k}$, where $\mathcal{F}_{1},\dots,\mathcal{F}_{k}$ are different in-context function classes, sample $x_{1},\dots,x_{n}$ IID, and form $(x_{1},f(x_{1}),\cdots,x_{n-1},f(x_{n-1}),x_{n},f(x_{n}))$. We refer to the case $k=1$ as single-task ICL and $k>1$ as multi-task ICL. The sequence model is trained with autoregressive next-token prediction, i.e., the model predicts $f(x_{i})$ conditioned on $(x_{1},f(x_{1}),x_{i-1},f(x_{i-1}),\cdots,x_{i})$ for $i=1,\dots,n$.
Refer to caption: extracted/5896652/figure/figure2_a.png

Our experimental setup follows Garg et al. () and Bhattamishra et al. (). Consider a function class $\mathcal{F}$ with domain $\mathcal{X}$. A sequence model $M_{\theta}$ (transformer or state-space model) is trained to identify $f\in\mathcal{F}$ in context and make a prediction on the subsequent query. The training data consists of sequences of the form $P=(x_{1},f(x_{1}),x_{2},f(x_{2}),\dots,x_{n-1},f(x_{n-1}),x_{n},f(x_{n}))$, where $f$ is sampled from a distribution $\mathcal{D}_{\mathcal{F}}$ and $x_{1},\dots,x_{n}$ are independently sampled from a distribution $\mathcal{D}_{\mathcal{X}}$. We train $M_{\theta}$ for the next token-prediction task: The loss over the sequence $P$ is given by $\frac{1}{n}\sum_{i=1}^{n}\ell(M_{\theta}(P_{i}),f(x_{i}))$, where $\ell(\cdot,\cdot)$ is an appropriate loss function and $P_{i}:=(x_{1},f(x_{1}),x_{2},f(x_{2}),\dots,x_{i-1},f(x_{i-1}),x_{i})$ is the $i$-th prefix for $i=1,\dots,n$. In our experiments, $n=120$ is a predetermined number shared across all tasks. This procedure is illustrated on Figure 2.

ICL tasks.
We consider *continuous ICL tasks* and *boolean ICL tasks*. For each ICL task, an $f\in\mathcal{F}$ is chosen, where $\mathcal{F}$ is a function class. For continuous ICL tasks, $\mathcal{F}$ consists of $f\colon\mathbb{R}^{d}\to\mathbb{R}$ for $d\in\mathbb{N}$ and the probability distribution on the domain is assumed to be $\mathcal{D}_{\mathcal{X}}=\mathcal{N}(0,I_{d})$. For boolean ICL tasks, $\mathcal{F}$ consists of $f\colon\{\pm 1\}^{d}\to\{\pm 1\}$ for $d\in\mathbb{N}$ and the probability distribution on the domain is assumed to be $\mathcal{D}_{\mathcal{X}}=\mathrm{Unif\,}(\{\pm 1\}^{d})$. We consider $10$ different ICL tasks: Linear Regression, Quadratic Regression, Sparse Linear Regression, ReLU Regression, Decision Tree, Sparse Parity($2$), Sparse Parity($3$), Parity, Conjunction, and Disjunction. The function class $\mathcal{F}$ defining each of these $10$ ICL tasks is precisely stated in Section 5.

Multi-task ICL training.
Most prior work on ICL, such as Garg et al. (); von Oswald et al. (); Bhattamishra et al. (); Park et al. (), considers ICL training with a single function class $\mathcal{F}$.
In this work, we train models to learn functions from the union of function classes $\bigcup_{m=1}^{k}\mathcal{F}_{m}$ in context, where $\mathcal{F}_{1},\dots,\mathcal{F}_{k}$ are distinct ICL task among the $10$ that we list in Section 5. For instance, if the model is trained on the union of linear and quadratic regression tasks, then $\mathcal{F}_{1}\cup\mathcal{F}_{2}=\{f\,|\,f(x)=w^{\intercal}x\}\cup\{f\,|\,f(
x)=x^{\intercal}Wx\}$.
For each $m=1,\dots,k$, we sample $f\sim\mathcal{D}_{\mathcal{F}_{m}}$
and $x_{1},\dots,x_{n}\stackrel{{\scriptstyle\text{IID}}}{{\sim}}\mathcal{D}_{
\mathcal{X}_{m}}$ and form the sequence $(x_{1},f(x_{1}),x_{2},f(x_{2}),\dots,x_{n},f(x_{n}))$. This sampling process is repeated $B$ times for each $m=1,\dots,k$, making $kB$ the total batch size. To balance the loss scales across different tasks, we normalize each task’s loss with constant factors $c_{1},\dots,c_{k}$ (further discussed in Section 4.1) so that the training loss stabilizes around $1$ during the plateaus. Appendix A provides further experimental details. In expectation, we minimize the loss

$$ $L(\theta)=\sum_{m=1}^{k}c_{m}\!\!\!\!\!\mathop{\mathbb{E}}_{\begin{subarray}{c }f\sim\mathcal{D}_{\mathcal{F}_{m}}\\ x_{1},\dots,x_{n}\stackrel{{\scriptstyle\text{IID}}}{{\sim}}\mathcal{D}_{ \mathcal{X}_{m}}\end{subarray}}\!\!\!\!\!\bigg{[}\frac{1}{n}\sum_{i=1}^{n}\ell \big{(}M_{\theta}(P_{i}),f(x_{i})\big{)}\bigg{]}.$ $$

Test loss. To evaluate the model’s performance on ICL tasks, we measure the error of the model’s prediction on the last ($n$-th) output of the prompt. For continuous ICL tasks, the test error is $(M_{\theta}(P_{n})-f(x_{n}))^{2}$. For boolean ICL tasks, the test accuracy is measured by $\mathbf{1}[\mathrm{sign}(M_{\theta}(P_{n}))=f(x_{n})]$, where $\mathbf{1}$ is the indicator function.

## 3 Task diversity shortens ICL plateaus

Long loss plateaus have been commonly reported in the various setups for training sequence models from scratch to perform ICL, including simple in-context functions (Garg et al., ; Chen et al., ; Li et al., ; Bhattamishra et al., ; Park et al., ), image datasets (Fu et al., ; Kirsch et al., ; Singh et al., ; Reddy, ), and language datasets (Akyürek et al., ; Olsson et al., ).
In this section, we present the following claim:

Here, task diversity refers to learning multi-task ICL on a mixture of distinct function classes. The claim 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, not shorten it.

We provide comprehensive experimental evidence to support this claim. Table 1 summarizes our experimental results on transformers. Figure 3 shows the results of state space models, specifically Mamba (Gu & Dao, ) and Hyena (Poli et al., ). Across the hundreds of task combinations in these setups, we consistently observed that task diversity enables training to escape plateaus more quickly. Figure 7 shows results in natural language ICL tasks reinforcing our claim. Refer to Appendix A for experimental details and Appendix B for additional tables.

However, not all ICL tasks mutually reduce the duration of plateaus. For instance, we consider Regbench task (Akyürek et al., ), which is generated by a random automata. Our experiments in Appendix C.2 show that combining Regbench with the ICL tasks we consider does not reduce the plateau of Regbench, but does reduce the plateau of the other ICL tasks.

Therefore, the claim should be understood as a description of a general tendency rather than a universal law. Nevertheless, our finding is broadly and robustly observed, as borne out by our extensive experiments.

Figure: Figure 3: Mamba (First row) and Hyena (Second row). Red lines and Blue lines respectively represent the loss dynamics of single-task training and multi-task training. (Left): Quadratic Regression$+$Linear Regression compared to Quadratic Regression. (Middle): Sparse Parity(2)$+$Linear Regression compared to Sparse Parity(2). (Right): Sparse Parity(3)$+$Conjunction compared to Sparse Parity(3).
Refer to caption: extracted/5896652/figure/mamba-QR.png

**Table 1: Task diversity shortens the ICL plateau. We train a transformer with various combinations of $9$ different tasks ($d=10$). For each run, we report two metrics: the time to escape the plateau and the time to complete the learning of the task (written in parentheses). The rows correspond to the number of tasks trained together, and each entry in the table corresponds to the average time across the training runs that include the given task. For example, the entry at (Number of tasks = 4, Conjunction) shows the average of $\binom{8}{3}$ results. We find that multi-task training shortens the ICL plateau. Precise details are provided in Appendix A.**
|  | Boolean Tasks | Continuous Tasks |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Number of<br>tasks | Sparse Parity(2) | Sparse Parity(3) | Conjunction | Disjunction | Linear<br>Regression | Quadratic<br>Regression | ReLU<br>Regression | Sparse Linear<br>Regression | Decision<br>Tree |
| 1 | $>1000$k | $>1000$k | $2.4$k ($2.7$k) | $2.5$k ($2.7$k) | $9.0$k ($12.1$k) | $13.1$k ($14.2$k) | $3.0$k ($3.7$k) | $16.5$k ($18.4$k) | $4.3$k ($16.0$k) |
| 2 | $12.2$k ($17.2$k) | $22.4$k ($26.7$k) | $1.6$k ($1.9$k) | $1.3$k ($1.6$k) | $2.7$k ($3.4$k) | $5.5$k ($7.5$k) | $2.2$k ($2.8$k) | $3.1$k ($3.8$k) | $5.9$k ($19.3$k) |
| 4 | $2.1$k ($2.4$k) | $2.9$k ($3.9$k) | $1.5$k ($3.0$k) | $1.2$k ($1.5$k) | $1.4$k ($2.1$k) | $2.1$k ($4.5$k) | $1.8$k ($2.5$k) | $1.4$k ($2.1$k) | $2.5$k ($18.0$k) |
| 8 | $1.5$k ($2.0$k) | $2.1$k ($2.4$k) | $0.9$k ($1.5$k) | $1.0$k ($1.5$k) | $0.9$k ($2.0$k) | $1.7$k ($5.4$k) | $0.9$k ($2.0$k) | $0.9$k ($1.9$k) | $2.0$k ($20.0$k) |
| 9 | $1.6$k ($1.8$k) | $2.2$k ($2.3$k) | $0.9$k ($1.5$k) | $0.9$k ($1.4$k) | $0.8$k ($1.8$k) | $1.6$k ($6.6$k) | $0.9$k ($1.9$k) | $0.8$k ($1.8$k) | $1.8$k ($22.0$k) |

Figure: Figure 4: Complexity model of multi-task training
Refer to caption: extracted/5896652/figure/illustration_3.png

Model of task complexity.
We quickly describe our mental model of the aggregate complexity of multi-task ICL training and the speedup due to task diversity. Let the ‘complexity’ of an ICL task be the time it takes to escape from the plateau. The complexity of a single-task ICL is the complexity observed when training with just the single task. The complexity of a multi-task ICL is the sum of the complexities of the constituent tasks, but the task diversity reduces it.
This reduction makes the aggregate complexity of the multi-task ICL less than the complexity of the individual single-task ICLs. Section 4.2 discusses when and why task diversity could reduce individual complexities. Figure 4 illustrates this notion.

##### Escaping plateau ≈ \approx ≈ training completion.

In the ICL setups we consider, we observe that (i) there is only one plateau, and (ii) learning is very rapid once this plateau is escaped from. This implies that the time at which training is completed, defined as the moment when the model reaches near-perfect training accuracy, is very close to the time of escaping from the (first) plateau. The results of Table 1 confirm this.

It should be noted that previous studies have demonstrated that multi-stage learning does occur in supervised learning, both theoretically (Ghosh et al., ; Bietti et al., ; Jin et al., ; Wang & Ma, ; Berthier et al., ) and empirically (Nakkiran et al., ; Refinetti et al., ; Rubruck et al., ). Therefore, we expect that ICL tasks involving complex hierarchical structures may exhibit multiple plateaus.

## 4 Why is task diversity helpful?

In Section 3, we demonstrated that task diversity shortens plateaus but did not explore *why* this effect occurs. In this section, we provide partial answers and hypotheses toward understanding this phenomenon.

### 4.1 Plateau is task-wise no-context learning

Figure: Figure 5: During the plateau, the model output very closely matches the task-wise optimal no-context function. Solid lines ( ) denote the raining loss while the dotted lines ( ) denote the squared distance between model output and task-wise optimal no-context function. (Left): Linear Regression($\mu=-0.5$) + Quadratic Regression($\mu=0.5$). (Right): Linear Regression($\mu=1$) + Sparse Parity(2).
Refer to caption: extracted/5896652/figure/ncl_lr+qr.png

At first glance, plateaus might appear to be a failure mode where no meaningful learning occurs.
However, in the Sparse Parity(2) task, for example, both test and train accuracies hover around $0.55$ during plateaus, as illustrated in Figure 3.
If the model learned nothing and were making random predictions, the accuracy should be $0.5$.
This deviation implies the model is learning something.

For continuous ICL tasks, define the optimal *no-context function* as

$$ $g_{\mathcal{F}}^{\star}:=\mathrm{argmin\,}_{g}\mathop{\mathbb{E}}_{f\sim \mathcal{D}_{\mathcal{F}},x\sim\mathcal{D}_{\mathcal{X}}}\left[(g(x)-f(x))^{2} \right],$ $$

i.e., $g_{\mathcal{F}}^{\star}$ is the context-independent function that minimizes the test error. For boolean ICL tasks, $g_{\mathcal{F}}^{\star}$ is analogously defined with the argmax of the test accuracy. In many cases, the optimal no-context function has a simple closed-form expression. For instance, if $\mathcal{F}=\{f\,|\,f(x)=w^{\intercal}x\}$ and $\mathcal{D}_{\mathcal{F}}$ is given by $w\sim\mathcal{N}(\mu,I_{d})$, then $g_{\mathcal{F}}^{\star}(x)=\mu^{\intercal}x$. As another example, if $\mathcal{F}=\{f\,|\,f(x)=x^{\intercal}Wx\}$ and $\mathcal{D}_{\mathcal{F}}$ is given by $W\sim\mathcal{N}(U,I_{d\times d})$, then $g_{\mathcal{F}}^{\star}(x)=x^{\intercal}Ux$. The ICL plateau corresponds to *task-wise no-context learning*, which we describe in the following. When we train a model $M_{\theta}$ for ICL with function classes $\mathcal{F}_{1},\dots,\mathcal{F}_{k}$, the model’s output during its plateau corresponds to

$$ $\displaystyle M_{\theta}(P_{n})=g^{\star}_{\mathcal{F}_{m}}(x_{n}),\quad\text{ $P_{n}$ is sampled from $\mathcal{F}_{m}$}.$ $$

In other words, the model identifies the function class $\mathcal{F}_{m}\in\{\mathcal{F}_{1},\dots,\mathcal{F}_{k}\}$
and then applies the optimal no-context function corresponding to $\mathcal{F}_{m}$. The in-context demonstrations $(x_{1},f(x_{1}),\dots,x_{n-1},f(x_{n-1}))$ are used to determine the function class $\mathcal{F}_{m}\in\{\mathcal{F}_{1},\dots,\mathcal{F}_{k}\}$ but not to determine the specific $f\in\mathcal{F}_{m}$.
This claim can be verified by measuring the error between $M_{\theta}(P_{n})$ and $g^{\star}_{\mathcal{F}}$, which we plot in Figure 5. Revisiting the Sparse Parity(2) task, the accuracy during plateau is indeed attributed to $0.55=\mathbb{E}_{f\sim\mathcal{D}_{\mathcal{F}},x\sim\mathcal{D}_{\mathcal{X}}
}[\mathbf{1}(g_{\mathcal{F}}^{\star}(x)=f(x))]$. Appendix F provides further details on no-context learning regime.

In Section 2, we used the scaling factors $c_{1},\dots,c_{k}$ to normalize each task’s empirical loss. Specifically, we set $c_{m}=1/\mathbb{E}_{f\sim\mathcal{D}_{\mathcal{F}},x\sim\mathcal{D}_{\mathcal{
X}}}\left[\ell(g_{\mathcal{F}_{m}}^{\star}(x),f(x))\right]$. Our findings on no-context learning imply that the normalized loss will have a plateau of height $1$.

### 4.2 Common structure across ICL tasks

Figure: Figure 6: (Left) Illustration. We pre-train on Task A and extract a checkpoint as training escapes the plateau. We then train on Task B starting from the checkpoint. (Right) Plateau escape time comparison. Each cell represents the ratio of plateau escape time, with lower ratios ( blue color) indicating that the model pre-trained on Task A significantly aids the learning of Task B.
Refer to caption: extracted/5896652/figure/ckpt_illust.png

Consider a multi-task ICL setup, where a model is trained on a set of tasks $\bigcup_{m=1}^{k}\mathcal{T}_{m}$. Suppose there exists a “common structure” $\mathcal{C}$ shared across tasks. Denoting the remaining part of each task as $\mathcal{I}_{m}$, we can decompose each task as $\mathcal{T}_{m}=\mathcal{C}+\mathcal{I}_{m}$ for $m=1,\dots,k$. Thus, multi-task ICL training is decomposed into two sub-problems: [learning $\mathcal{C}$] and [learning $\mathcal{I}_{1},\dots,\mathcal{I}_{k}$].
We argue that the main claim of Section 3 can be explained by the following hypotheses:

- (1)
There exists a common structure shared across the multiple ICL tasks.
- (2)
The ICL plateau arises from the difficulty of learning this common structure.
- (3)
Training multiple tasks jointly with a shared structure makes it easier to learn that structure.

In the following, we present evidence supporting these hypotheses. Section 4.2.1 demonstrates
(1) and (2). Section 4.2.2 elucidates (3) with a toy experiment on feature learning.

#### 4.2.1 Checkpoint experiment

Consider the following checkpoint experiment. For each single-task training with task A, we save the model as it escapes the plateau as illustrated in Figure 6 (left). Using this checkpoint model as the initialization, we train it on task B.

The findings, summarized in Figure 6, indicate that the checkpoint model transferred from task A quickly learns task B with a shortened plateau. This implies that (1) task A and task B share a common structure and that (2) the plateau arises from the difficulty of learning the common structure. We conduct an analogous experiment on natural language ICL tasks and observe similar results, as shown in Figure 8 of the appendix. Further details are presented in Appendix D.

#### 4.2.2 Generality of our hypothesis

Consider the following intuition. When $\bigcup_{m=1}^{k}\mathcal{T}_{m}$ are trained concurrently, the model receives multiple “views” of the common structure $\mathcal{C}$ through the different compositions $\mathcal{C}+\mathcal{I}_{m}$ for $m=1,\dots,k$. We hypothesize (3): These multiple “views” of $\mathcal{C}$ make $\mathcal{C}$ easier to learn in the sense of a more favorable optimization landscape. This may be the key mechanism allowing task diversity to shorten the ICL plateau, and this phenomenon may extend beyond the ICL setup.

The following feature learning experiment makes this intuition more concrete. Consider the 2-layer feature learning setup, a setup with a large body of prior work (Damian et al., ; Ba et al., ; Dandi et al., ; Wang et al., ). For input $x\in\mathbb{R}^{d}$, the true function to learn is $f_{\star}(x)=U\sigma(A_{\star}x)$, where $U\in\mathbb{R}^{k\times h}$ and a true feature matrix $A_{\star}\in\mathbb{R}^{h\times d}$. The goal is to learn $f\approx f_{\star}$ with $f(x)=V\sigma(Wx)$, where $h^{\prime}\gg h$, $V\in\mathbb{R}^{k\times h^{\prime}}$, and $W\in\mathbb{R}^{h^{\prime}\times d}$. The loss function is:

$$ $\displaystyle L(W,V)=\sum_{m=1}^{k}\mathbb{E}_{x\sim\mathcal{D}_{x}}\left[ \left(v_{i}^{\intercal}\sigma(Wx)-u_{i}^{\intercal}\sigma(A_{\star}x)\right)^{ 2}\right],\quad U=\left[\begin{matrix}u_{1}&\cdots&u_{k}\end{matrix}\right],\, V=\left[\begin{matrix}v_{1}&\cdots&v_{k}\end{matrix}\right].$ $$

The idea is that $u_{1},\dots,u_{k}$ represent $k$ sub-tasks that share the common feature matrix $A_{\star}$. If we sample $u_{1},\dots,u_{k}$ from $k$ different distributions, this corresponds to multi-task training of $k$ distinct tasks. Conversely, if we sample $u_{1},\dots,u_{k}$ from a single distribution, this corresponds to single-task training, as the $k$ sub-tasks become identical. Interestingly, we find that multi-task training exhibits significantly shorter plateau compared to single-task training, as shown in Figure 7 (right).
Figure 12 of the appendix shows that additional results with different hyperparameter configurations provide qualitatively similar results.

Although this toy model is a simple supervised learning setup, without sequence models or in-context learning, it reproduces the shortened plateau and makes our intuition more concrete through analogy. The results of this toy model, shown in Figure 7, lead us to make the general hypothesis (3): Training multiple tasks jointly with a shared structure makes it easier to learn the common structure.

Figure: Figure 7: (Left) Language ICL task. Previous work Fu et al. () identified the difficulty of learning the WordSelection$(4)$ task. We found that mixing it with WordLength or WordSelection$(2)$ reduces the plateau. Refer to Appendix C.1 for further details. (Right) Feature learning setup. For the toy model described in Section 4.2.2, multi-task feature learning significantly shortens the loss plateau. Refer to Appendix E for further details.
Refer to caption: extracted/5896652/figure/lang_mix_loss_comb.png

#### 4.2.3 The common structure is not just an induction head

So then, what specifically is this common structure? We believe it must involve some algorithmic component, as all of the ICL tasks necessitate an internal algorithm to identify the specific function being demonstrated by the in-context demonstrations.

A plausible candidate is the *induction head*, a circuit that searches over the sequence for previous instances of a current token and predicts the same completion again. Indeed, Olsson et al. () argued that the development of an induction head coincides with the escape from the training plateau. To test this idea, we designed the following Retrieval ICL task, inspired by the prior ICL tasks from Park et al. (); Singh et al. (); Reddy ().

Retrieval. Sample $1024$ 5-tuples of one key and four values: $\{(k_{i},v_{i}^{1},v_{i}^{2},v_{i}^{3},v_{i}^{4})\}_{i=1}^{1024}$. The $k_{i}$s and $v_{i}^{j}$s are independently sampled from $\mathcal{D}_{\mathcal{K}}$ and $\mathcal{D}_{\mathcal{V}}$, respectively. To generate prompts, we uniformly sample $(n-1)$ 5-tuples without replacement and uniformly choose one $v_{i}$ per 5-tuple, resulting in $\{(\mathbf{k}_{i},\mathbf{v}_{i})\}_{i=1}^{n-1}$. Next, we sample $p\sim\mathrm{Unif\,}(\{1,\dots,n-1\})$ and set $\mathbf{q}=\mathbf{k}_{p}$. Finally, given $P_{n}=(\mathbf{k}_{1},\mathbf{v}_{1},\dots,\mathbf{k}_{n-1},\mathbf{v}_{n-1},
\mathbf{q})$, the task is to predict $\mathbf{v}_{p}$. We consider two Retrieval tasks: Gaussian Retrieval with $(\mathcal{D}_{\mathcal{K}},\mathcal{D}_{\mathcal{V}})=(\mathcal{N}(0,I_{d}),
\mathcal{N}(0,1))$
and Boolean Retrieval with $(\mathcal{D}_{\mathcal{K}},\mathcal{D}_{\mathcal{V}})=(\mathrm{Unif\,}\{\pm 1
\}^{d},\mathrm{Unif\,}\{\pm 1\})$.

Note that the induction head is precisely the mechanism for solving this Retrieval task. We conduct a checkpoint experiment for the Retrieval tasks, allowing the models to learn the induction head. To train continuous ICL tasks, we use a checkpoint model pre-trained with Gaussian Retrieval. To train Boolean ICL tasks, we use a checkpoint model pre-trained with Boolean Retrieval. We find that the checkpoint does *not* significantly shorten the ICL tasks’ plateaus. This result demonstrates that the common structure shared by other ICL tasks is not just an induction head. (It is unclear whether an induction head is useful for our ICL tasks at all.) For further details, please refer to Appendix D.2.

Presently, the problem of characterizing the common structure with any specificity remains unresolved. We defer further investigation of this matter to future work.

## 5 In-context learning tasks

In this section, we quickly list and define the ICL tasks that we consider, which are primarily adapted from
Garg et al. (); Bhattamishra et al. ().
Each ICL task is specified by $\mathcal{F}$ the function class, $\mathcal{D}_{\mathcal{F}}$ a probability distribution over the function class, and $\mathcal{D}_{\mathcal{X}}$ a probability distribution over the inputs. For continuous ICL tasks, $\mathcal{D}_{\mathcal{X}}=\mathcal{N}(0,I_{d})$. For boolean ICL tasks, $\mathcal{D}_{\mathcal{X}}=\mathrm{Unif\,}\{\pm 1\}^{d}$.

Continuous ICL tasks:

- •
Linear Regression. $\mathcal{F}=\{f\,|\,f(x)=w^{\intercal}x\}$. $\mathcal{D}_{\mathcal{F}}$: Each element of $w\in\mathbb{R}^{d}$ is independently sampled from $\mathcal{N}(\mu,1)$.
- •
Quadratic Regression. $\mathcal{F}=\{f\,|\,f(x)=x^{\intercal}Wx\}$. $\mathcal{D}_{\mathcal{F}}$: Each element of $W\in\mathbb{R}^{d\times d}$ is independently sampled from $\frac{1}{\sqrt{d}}\mathcal{N}(\mu,1)$.
- •
Sparse Linear Regression. $\mathcal{F}=\{f\,|\,f(x)=w_{\text{sparse}}^{\intercal}x\}$. $\mathcal{D}_{\mathcal{F}}$: Each element of $w\in\mathbb{R}^{d}$ is independently sampled from $\mathcal{N}(\mu,1)$. To sample $w_{\text{sparse}}$, we uniformly choose $k=3$ coordinates and retain the corresponding coordinates of $w$.
- •
ReLU Regression. $\mathcal{F}=\{f\,|\,f(x)=\mathrm{ReLU\,}(w^{\intercal}x)\}$. $\mathcal{D}_{\mathcal{F}}$: Each elements of $w\in\mathbb{R}^{d}$ is independently sampled from $\mathcal{N}(\mu,1)$.
- •
Decision Tree. Consider a full binary tree of fixed depth$=4$. The values of leaf nodes and branch nodes are independently sampled from $\mathcal{N}(\mu,1)$ and $\mathrm{Unif\,}(1,\dots,d)$, respectively. When $x$ traverses through the tree, at each branch node with index $i$, we move right if $x[i]>0$ and move left otherwise. $f(x)$ corresponds to the value of the leaf node reached at the end of the traversal.

Boolean ICL tasks:

- •
Sparse Parity$(k)$. $\mathcal{F}=\{f\,|\,f(x)=\prod_{i\in A}x[i]\}$.
$\mathcal{D}_{\mathcal{F}}$:
$A\subseteq\{1,\dots,d\}$ is a uniformly sampled subset of size $k$.
- •
Parity. $\mathcal{F}=\{f\,|\,f(x)=\prod_{i\in A}x[i]\}$. $\mathcal{D}_{\mathcal{F}}$: $A\subseteq\{1,\dots,d\}$ is a uniformly sampled subset, regardless of the size.
- •
Conjunction. $\mathcal{F}=\{f\,|\,f(x)=\left(\wedge_{i\in A}x[i]\right)\wedge\left(\wedge_{i
\in B}\bar{x}[i]\right)\}$. $\mathcal{D}_{\mathcal{F}}$: $A$ is an uniformly sampled subset of $\{1,\dots,d\}$. Thereafter, we uniformly sample a subset $B\subseteq\{1,\dots,d\}\backslash A$. .
- •
Disjunction. $\mathcal{F}=\{f\,|\,f(x)=\left(\vee_{i\in A}x[i]\right)\wedge\left(\wedge_{i
\in B}\bar{x}[i]\right)\}$. $\mathcal{D}_{\mathcal{F}}$: $A$ is an uniformly sampled subset of $\{1,\dots,d\}$. Thereafter, we uniformly sample a subset $B\subseteq\{1,\dots,d\}\backslash A$.

## 6 Conclusion

In this work, we identify that training on a diverse set of multiple ICL tasks is surprisingly easier than training for a single ICL task in the sense of a more favorable optimization landscape. This observation aligns with the “blessing of dimensionality/scale” seen in the modern era of deep learning. Indeed, LLM training via next-token prediction can be thought of effectively as a highly diverse multi-task learning, requiring a wide range of reasoning skills for a wide range of text types, and the success of LLM training 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.

This insight opens new avenues for future work. It may be that explaining and understanding the success of large-scale deep learning requires considering not just the large data, large network, and large compute but also the large (effective) task diversity.

## Acknowledgement

We thank Wonnjong Rhee for prompting this project as part of a course at Seoul National University. We are also grateful to the participants of the KAIST-SNU workshop for their discussions at the early stages of this work. Special thanks to Chanwoo Park for sharing his empirical insights. We thank Minkyu Kim for providing feedback about figures. We also thank Hyunsik Chae, Junsu Kim, and Jaesung Park for reviewing the manuscript and providing valuable feedback.

## Appendix A Experimental details

##### Model Architectures.

We focus on transformer (Vaswani et al., ), Mamba (Gu & Dao, ), and Hyena (Poli et al., ). To handle model architectures that process inputs and outputs as vectors in an embedding space, we first pad each $f(x_{i})$ with $(d-1)$ zeros to match the dimension of $x_{i}$s. We then add a learnable linear layer to map these vectors into the embedding space. A second learnable linear layer maps the model’s output to a scalar. Each model’s hyperparameters are dictated as follows. For transformer, we use GPT2 (Radford et al., ), with 12 decoder layers of embedding size 256 with 8 heads and relative positional encoding. Hyena follows the exact configuration of transformer’s setup for the relevant hyperparameters. On the other hand, Mamba employs 24 layers, doubling the number of layers to match the model parameter size as is standard for Mamba. All models contain input and output projection heads that map input to the hidden state and hidden state to output. To train a model from scratch, we randomly initialize the model parameters.

##### Configuration.

We use the Adam optimizer (Kingma & Ba, ) with a learning rate of $0.0001$. The batch size $B=64$ corresponds to each task. Regarding the choice of $\ell(\cdot,\cdot)$, we use mean-squared error loss for continuous ICL tasks and cross-entropy loss for boolean ICL tasks. The input dimensions considered are $d=10$ and $15$.

##### Plateau and Training Exit Conditions.

The plateau escape time $t_{\text{plateau}}$ is defined as $t_{\text{plateau}}=\min_{t>100}\left[\frac{1}{100}\sum_{t^{\prime}=t-99}^{t}
\text{Training loss}(t^{\prime})<0.8\right]$. This effectively captures the plateau escape as we have normalized each task’s loss around $1$. The training exit conditions are defined as follows $t_{\text{exit}}=\min_{t>100}\left[\frac{1}{100}\sum_{t^{\prime}=t-99}^{t}\text
{Test error}(t^{\prime})<0.2\right]$ for continuous ICL tasks and $t_{\text{exit}}=\min_{t>100}\left[\frac{1}{100}\sum_{t^{\prime}=t-99}^{t}\text
{Test accuracy}(t^{\prime})>0.95\right]$ for boolean ICL tasks. $t_{\text{plateau}}$ and $t_{\text{exit}}$ are measured for each individual task.

##### Batch generation process

The following pseudo algorithm summarizes the batch generation process of (multi)-task ICL training.

Figure: Batch Prompt Generation Algorithm

After constructing a batch, we minimize the following empirical loss. Denote $f^{(j,m)}$ and $x_{i}^{(j,m)}$ as the corresponding in-context function and $i$-th input for each $P^{(j,m)}$.

$$ $\displaystyle\hat{L}(\theta)\colon=\sum_{m=1}^{k}c_{m}\left[\sum_{j=1}^{B} \left[\frac{1}{n}\sum_{i=1}^{n}\ell\big{(}M_{\theta}(P_{i}^{(j,m)}),f^{(j,m)}( x_{i}^{(j,m)})\big{)}\right]\right]$ $$

## Appendix B Additional tables

To demonstrate the generality of our claim, we conducted extensive experiments beyond those shown in Table 1. The following tables summarize results for transformer (input dimension $d=15$), Mamba (input dimension $d=10$), and Hyena (input dimension $d=10$). We report two metrics for each run: time to escape the plateau and the time to complete the learning (written in parentheses).

**Table 2: Task diversity shortens the ICL plateau. We train a transformer with various combinations of $9$ different tasks ($d=15$).**
|  | Boolean Tasks | Continuous Tasks |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Number of<br>tasks | Sparse Parity(2) | Sparse Parity(3) | Conjunction | Disjunction | Linear<br>Regression | Quadratic<br>Regression | ReLU<br>Regression | Sparse Linear<br>Regression | Decision<br>Tree |
| 1 | $>1000$k | $>1000$k | $46.9$k ($66.6$k) | $22.6$k ($50.7$k) | $58.8$k ($61.3$k) | $25.0$k ($25.1$k) | $4.1$k ($5.4$k) | $31.6$k ($35.3$k) | $32.4$k ($86.2$k) |
| 2 | $38.9$k ($52.6$k) | $104.7$k ($148.2$k) | $7.5$k ($16.5$k) | $5.0$k ($11.6$k) | $12.4$k ($13.9$k) | $14.8$k ($17.2$k) | $4.1$k ($5.1$k) | $4.8$k ($5.4$k) | $26.6$k ($70.4$k) |
| 8 | $4.0$k ($4.8$k) | $10.1$k ($13.1$k) | $2.8$k ($8.5$k) | $2.7$k ($8.7$k) | $2.3$k ($3.7$k) | $4.0$k ($4.3$k) | $2.4$k ($3.9$k) | $2.2$k ($3.6$k) | $13.1$k ($44.4$k) |
| 9 | $3.6$k ($4.1$k) | $10.9$k ($11.6$k) | $2.5$k ($7.9$k) | $2.5$k ($7.8$k) | $2.0$k ($3.2$k) | $3.5$k ($3.6$k) | $2.0$k ($3.4$k) | $1.9$k ($3.2$k) | $13.8$k ($50.0$k) |

**Table 3: Task diversity shortens the ICL plateau (Mamba). We train a Mamba with various combinations of $9$ different tasks ($d=10$).**
|  | Boolean Tasks | Continuous Tasks |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Number of<br>tasks | Sparse Parity(2) | Sparse Parity(3) | Conjunction | Disjunction | Linear<br>Regression | Quadratic<br>Regression | ReLU<br>Regression | Sparse Linear<br>Regression |
| 1 | $7.5$k ($7.7$k) | $194.3$k ($196.5$k) | $0.4$k ($0.4$k) | $0.4$k ($0.4$k) | $0.2$k ($0.4$k) | $17.7$k ($37.7$k) | $0.2$k ($1.0$k) | $0.2$k ($0.4$k) |
| 2 | $2.9$k ($3.2$k) | $14.4$k ($15.6$k) | $0.4$k ($0.5$k) | $0.3$k ($0.4$k) | $0.2$k ($0.5$k) | $7.7$k ($22.8$k) | $0.3$k ($0.9$k) | $0.2$k ($0.5$k) |
| 4 | $1.5$k ($1.7$k) | $6.0$k ($6.7$k) | $0.3$k ($0.4$k) | $0.3$k ($0.4$k) | $0.2$k ($0.6$k) | $3.4$k ($17.1$k) | $0.3$k ($0.9$k) | $0.2$k ($0.6$k) |
| 8 | $2.4$k ($2.4$k) | $3.7$k ($4.3$k) | $0.3$k ($0.4$k) | $0.3$k ($0.4$k) | $0.3$k ($0.7$k) | $2.8$k ($20.5$k) | $0.3$k ($0.9$k) | $0.3$k ($0.6$k) |

**Table 4: Task diversity shortens the ICL plateau (Hyena). We train a Hyena with various combinations of $9$ different tasks ($d=10$).**
|  | Boolean Tasks | Continuous Tasks |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Number of<br>tasks | Sparse Parity(2) | Sparse Parity(3) | Conjunction | Disjunction | Linear<br>Regression | Quadratic<br>Regression | ReLU<br>Regression | Sparse Linear<br>Regression |
| 1 | $1.1$k ($1.3$k) | $53.9$k ($58.6$k) | $0.2$k ($0.3$k) | $0.2$k ($0.3$k) | $0.1$k ($0.4$k) | $11.1$k ($35.4$k) | $0.2$k ($0.7$k) | $0.1$k ($0.4$k) |
| 2 | $2.0$k ($2.4$k) | $23.7$k ($27.2$k) | $0.2$k ($0.3$k) | $0.2$k ($0.3$k) | $0.1$k ($0.4$k) | $7.6$k ($33.6$k) | $0.2$k ($0.7$k) | $0.1$k ($0.4$k) |
| 4 | $2.4$k ($2.5$k) | $16.8$k ($17.5$k) | $0.2$k ($0.3$k) | $0.2$k ($0.3$k) | $0.1$k ($0.5$k) | $6.3$k ($41.7$k) | $0.2$k ($0.8$k) | $0.1$k ($0.5$k) |
| 8 | $3.8$k ($4.0$k) | $13.6$k ($14.2$k) | $0.2$k ($0.3$k) | $0.2$k ($0.3$k) | $0.2$k ($0.5$k) | $7.1$k ($55.0$k) | $0.2$k ($0.9$k) | $0.2$k ($0.5$k) |

For Mamba and Hyena, loss plateaus are less consistently observed. Specifically, training on Conjunciton, Disjunction, Linear regression, ReLU Regression, and Sparse Linear Regression rarely exhibit a loss plateau. This tendency, however, does not violate our claim, as we focus on how task diversity shortens the plateau when it *does occur*.

## Appendix C Language ICL tasks

### C.1 Natural Language ICL tasks

Figure: Figure 8: Checkpoint experiments on language ICL tasks.
Refer to caption: extracted/5896652/figure/lang_ckpt_loss_comb.png

To further verify our claim on tasks with real-world language data, we examine ICL in natural language processing (NLP) tasks.

##### Model.

We use GPT2 (Radford et al., ) model with $12$ layers, $12$ attention heads, $768$ embedding size, and $50257$ vocabularies. We train GPT2 from scratch
to learn in-context tasks.

##### Dataset.

We use the word data from Nguyen et al. (). Following the approach in Todd et al. (), we choose $5245$ words that can be tokenized to a single token. Words are then randomly chosen and arranged into sequences according to the task. Each sequence consists of six examples where the preceding five serve as context examples, and the last one as the query example.

##### Prompting.

We convert a given sequence of words to a prompt by following Todd et al. (), where words are separated by a space, and examples are separated by a line:

$<$bos$>$Q. word${}^{1}_{1}$ word${}^{1}_{2}$ $\dots$ word${}^{1}_{d}$
A. answer^1
⋮
Q. word${}^{5}_{1}$ word${}^{5}_{2}$ $\dots$ word${}^{5}_{d}$
A. answer^5
Q. word${}^{6}_{1}$ word${}^{6}_{2}$ $\dots$ word${}^{6}_{d}$
A.

##### WordSelection ( d ) 𝑑 (d) ( italic_d ) .

We devise WordSelection$(d)$ by generalizing the WordSelection Task (Fu et al., ). The goal of this task is to select a single word from given $d$ words. The model’s goal is to learn in-context which word to select.

##### WordLength.

We also devise WordLength Task. Let $\text{len}(w)$ be the length of word $w$ and $\%$ be the remainder operator. Given four words $w_{1},\dots w_{4}$, the goal of this task is to learn in-context either $\sum_{i=1}^{4}\text{len}(w_{i})$ or $\sum_{i=1}^{4}\text{len}(w_{i})\%10$ from given examples.

##### Result.

Prior work Fu et al. () reported the difficulty of learning WordSelection$(4)$, by showing that it cannot learn within about $78$k iterations. Indeed, as shown in (Figure 7, Right), single-task training of WordSelection$(4)$ exhibits a plateau for over $180$k iterations. However, mixing WordSelection$(4)$ and WordSelection$(2)$ significantly shortens the plateau to $3.2$k iterations, accelerating the learning. The same phenomenon is observed when WordSelection$(4)$ and WordLength are trained simultaneously. These results reveal that our claim generally holds in NLP tasks.

We also conduct the checkpoint experiment, where we use *checkpoint* models on WordSelection$(2)$ and WordLength as baseline models and train them respectively on WordSelection$(4)$. As shown in Figure 8, using *checkpoint* models shortens the plateau, demonstrating that NLP tasks also share a certain common structure.

### C.2 Formal Language ICL task: Regbench

We follow the configuration of Akyürek et al. (). To jointly train Quadratic Regression and DFA tasks, since their input formats differ, we use two different input embedding layers. Figure 9 illustrates the result.

Figure: Figure 9: Regbench. DFA task shortens the Quadratic Regression’s plateau, while Quadratic Regression does not shorten the DFA’s plateau. Red line shows the multi-task training dynamics and blue and green lines show the single-task training dynamics.
Refer to caption: extracted/5896652/figure/regbench_dfa.png

## Appendix D Checkpoint experiment

### D.1 Experimental details

In this section, we provide the configurations for the checkpoint experiments in Section 4.2.1. The input dimension is set to $d=10$. For each single-task training, we save the model once it escapes the plateau. The exact iteration numbers where the models are saved are provided in Table 5. Since Sparse Parity(2) and Sparse Parity(3) do not escape the plateau within observable training steps, no checkpoint models could be obtained.

**Table 5: Checkpoint models.**
|  | Boolean Task | Continuous Task |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Conjunction | Disjunction | Boolean<br>Retrieval | Linear<br>Regression | Quadratic<br>Regression | ReLU<br>Regression | Sparse Linear<br>Regression | Gaussian<br>Retrieval |
| Iteration number | $2.2$k | $2.6$k | $105$k | $19$k | $11.7$k | $2.8$k | $7$k | $50$k |

After saving the baseline models, we train each task (Conjunction, Disjunction, Linear Regression, Quadratic Regression, ReLU Regression, Sparse Linear Regression, Sparse Parity(2), and Sparse Parity(3)) according to the configuration in Appendix A. To calculate the ratio, we divide the result by the plateau escape time summarized in the first row of Table 1.

##### Learning time comparison of checkpoint models.

In Figure 6, we demonstrated that checkpoint models significantly shorten the plateaus. However, since the time required to escape the plateau and to learn the task may differ, it is necessary to verify whether the same phenomenon holds for learning time. Figure 10 confirms that this is indeed the case, although the effect is less robust compared to Figure 6.

Figure: Figure 10: Learning time ratios of checkpoint models.
Refer to caption: extracted/5896652/figure/after_plateau_complete.png

### D.2 Experiment on Retrieval task

Recall the description of Retrieval task.

Retrieval. Sample $1024$ 5-tuples of one key and four values: $\{(k_{i},v_{i}^{1},v_{i}^{2},v_{i}^{3},v_{i}^{4})\}_{i=1}^{1024}$. The $k_{i}$s and $v_{i}^{j}$s are independently sampled from $\mathcal{D}_{\mathcal{K}}$ and $\mathcal{D}_{\mathcal{V}}$, respectively. To generate prompts, we uniformly sample $(n-1)$ 5-tuples without replacement and uniformly choose one $v_{i}$ per 5-tuple, resulting in $\{(\mathbf{k}_{i},\mathbf{v}_{i})\}_{i=1}^{n-1}$. Next, we sample $p\sim\mathrm{Unif\,}(\{1,\dots,n-1\})$ and set $\mathbf{q}=\mathbf{k}_{p}$. Finally, given $P_{n}=(\mathbf{k}_{1},\mathbf{v}_{1},\dots,\mathbf{k}_{n-1},\mathbf{v}_{n-1},
\mathbf{q})$, the task is to predict $\mathbf{v}_{p}$. We consider two Retrieval tasks: Gaussian Retrieval and Boolean Retrieval, with $(\mathcal{D}_{\mathcal{K}},\mathcal{D}_{\mathcal{V}})=(\mathcal{N}(0,I_{d}),
\mathcal{N}(0,1))$ and $(\mathcal{D}_{\mathcal{K}},\mathcal{D}_{\mathcal{V}})=(\mathrm{Unif\,}\{\pm 1
\}^{d},\mathrm{Unif\,}\{\pm 1\})$, respectively.

To train continuous ICL tasks, we use a checkpoint model saved from Gaussian Retrieval, and for boolean ICL tasks, we use a checkpoint model saved from Boolean Retrieval. We follow the configuration in Appendix A. For each experiment, we measure the time spent to learn each ICL task. The Results are illustrated in Figure 11.

Figure: Figure 11: Checkpoint experiment on Retrieval.
Refer to caption: extracted/5896652/figure/retrieval_after_plateau_complete.png

## Appendix E Feature learning experiment

The true function is $f_{\star}(x)=U\sigma(A_{\star}x)$, where $U\in\mathbb{R}^{k\times h}$ and $A_{\star}\in\mathbb{R}^{h\times d}$. Columns of $U$, denoted as $u_{m}$s, are independently sampled from $\mathcal{N}(\mu_{m},I_{d})$, with $(\mu_{1},\ldots,\mu_{k})$ being pre-determined vectors prior to training.

To design multi-task learning, we independently sample $k$ different vectors $\mu_{1},\ldots,\mu_{k}$ from $\sqrt{h}\times\mathbb{S}^{h-1}$, making $\mathcal{N}(\mu_{m},I_{d})$s differ. In contrast, for single-task training, all $\mu_{1},\ldots,\mu_{k}$ are set to the same vector $\mu$, which is also sampled from $\sqrt{h}\times\mathbb{S}^{h-1}$.

The model that learns the true function is $f(x)=V\sigma(Wx)$ where $V\in\mathbb{R}^{m\times h^{\prime}}$ and $W\in\mathbb{R}^{h^{\prime}\times d}$. We set $d=150$, $h\in\{10,20\}$, $h^{\prime}\in\{100,300\}$, and $k\in\{15,30\}$.
The Adam optimizer is used with a learning rate of $\text{lr}=0.001$ and we apply $0.1\times\text{lr}$ for the model’s heads, following the approach in Berthier et al. (). We use the sigmoid activation function
$\sigma(x)=\frac{1}{1+e^{-x}}$.

##### Additional Figures.

Figure 7 represents the result when $(d,h,h^{\prime},k)=(150,10,100,15)$. To enhance the generality, we present figures for other hyperparameters as well. For each figure, multi-task training ( blue lines) exhibits a much shorter plateau compared to single-task training ( red lines).

Figure: (a) $(h^{\prime},k)=(100,15)$
Refer to caption: extracted/5896652/figure/d150_m10_n100_div15_2nn_from_csv.png

## Appendix F No-context learning

### F.1 Optimal no-context functions

In this section, we provide a detailed explanation of the optimal no-context function. We have defined the optimal no-context function as follows. For continuous ICL tasks,
$g_{\mathcal{F}}^{\star}\colon=\mathrm{argmin}_{g}\mathbb{E}_{f\sim\mathcal{D}_
{\mathcal{F}},x\sim\mathcal{D}_{\mathcal{X}}}\left[(g(x)-f(x)^{2})\right]$, where the argmin is taken over $\{g\,|\,g\colon\mathbb{R}^{d}\to\mathbb{R}\}$. For boolean ICL tasks, $g_{\mathcal{F}}^{\star}\colon=\mathrm{argmax}_{g}\mathbb{E}_{f\sim\mathcal{D}_
{\mathcal{F}},x\sim\mathcal{D}_{\mathcal{X}}}\left[\mathbf{1}(g(x)=f(x))\right]$, where the argmin is taken over $\{g\,|\,g\colon\{\pm 1\}^{d}\to\{\pm 1\}\}$.

For boolean ICL tasks, we can calculate $g_{\mathcal{F}}^{\star}(x)$ for each $x\in\{\pm 1\}^{d}$. The closed form of $g_{\mathcal{F}}^{\star}(x)$ can be expressed using $G_{x}\colon=\{f(x)\,|\,f\in\mathcal{F}\}$ as follows.

$$ $\displaystyle g_{\mathcal{F}}^{\star}(x)=\begin{cases}1\quad&\#(1\in G_{x})>\# (-1\in G_{x})\\ -1\quad&(\text{otherwise})\end{cases}.$ $$

For continuous ICL tasks, we cannot apply the same approach since the output space is not discrete. However, we can still derive the closed form for $g_{\mathcal{F}}^{\star}(x)$ for most tasks: $g_{\mathcal{F}}^{\star}(x)\colon=\mathbb{E}_{f\sim\mathcal{D}_{\mathcal{F}}}
\left[f(x)\right]$. This is because

$$ $\displaystyle g_{\mathcal{F}}^{\star}(x)\colon=\mathrm{argmin}_{a\in\mathbb{R} }\mathbb{E}_{f\sim\mathcal{D}_{\mathcal{F}}}\left[(f(x)-a)^{2}\right]=\mathrm{ argmin}_{a\in\mathbb{R}}\left(a-\mathbb{E}_{f\sim\mathcal{D}_{\mathcal{F}}}[f( x)]\right)^{2}.$ $$

For tasks other than ReLU Regression, we can obtain the closed forms of $g_{\mathcal{F}}^{\star}(x)$ using the above property, as listed follows.

- •
Linear Regression:
$g_{\mathcal{F}}^{\star}(x)=(\mu,\ldots,\mu)^{\intercal}x$.
- •
Quadratic Regression: $g_{\mathcal{F}}^{\star}(x)=x^{\intercal}Ux,(U)_{ij}=\mu.$
- •
Sparse Linear Regression: $g_{\mathcal{F}}^{\star}(x)=\frac{k}{d}(\mu,\ldots,\mu)^{\intercal}x$.
- •
Decision Tree: $g_{\mathcal{F}}^{\star}(x)=\mu$.

For ReLU Regression, we can estimate $g^{\star}_{\mathcal{F}}(x)$ with small error through a sufficient number of samples from $\mathcal{D}_{\mathcal{F}}$.

### F.2 Experimental details and additional figures

Let tasks with in-context function classes $\bigcup_{m=1}^{k}\mathcal{F}_{m}$ be trained simultaneously. Note that we do not discern the single-task and multi-task training: $k$ can be $1$ as well. During training, we measure

$$ $\displaystyle\mathbb{E}_{f\sim\mathcal{D}_{\mathcal{F}_{m}},x_{1},\dots,x_{n} \stackrel{{\scriptstyle\text{IID}}}{{\sim}}\mathcal{D}_{\mathcal{X}_{m}}} \lvert M_{\theta}(P_{n})-g^{\star}_{\mathcal{F}_{m}}(x_{n})\rvert^{2},\quad P_ {n}=(x_{1},f(x_{1}),\ldots,x_{n-1},f(x_{n-1}),x_{n})$ $$

for continuous ICL tasks and

$$ $\displaystyle\mathbb{E}_{f\sim\mathcal{D}_{\mathcal{F}_{m}},x_{1},\dots,x_{n} \stackrel{{\scriptstyle\text{IID}}}{{\sim}}\mathcal{D}_{\mathcal{X}_{m}}}\left [\mathbf{1}(\mathrm{sign}(M_{\theta}(P_{n}))\neq g^{\star}_{\mathcal{F}_{m}}(x _{n}))\right],\quad P_{n}=(x_{1},f(x_{1}),\ldots,x_{n-1},f(x_{n-1}),x_{n}).$ $$

for boolean ICL tasks. As shown by the following figures, these measurements are close to 0 0, which indicates that the model’s output corresponds to *task-wise optimal no-context function*. We follow Appendix A for other configurations. Moreover, to prevent the optimal no-context be merely zero function, we use $\mu\neq 0$ to parameterize the in-context function classes.

Figure: Figure 13: No-context learning regime. During the plateau, the model output matches the task-wise optimal no-context function. (Top) Sparse Parity(2). (Middle) Quadratic Regression ($\mu=-1$). (Bottom) Quadratic Regression ($\mu=-1$) + Sparse Parity(2).
Refer to caption: extracted/5896652/figure/ncl_sp2.png

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 4
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?
soundness: 1
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 4
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.
soundness: 2
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.
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.
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.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Please refer to my weaknesses part.
soundness: 2
strengths: This paper is clearly written. The figures and captions are easy to follow, and the motivation is clear.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: N/A
soundness: 3
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).
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).
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.

Meta-review: ## 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.

Decision: Reject