Task Diversity Shortens the In-Context Learning Plateau
Authors: Jaeyeon Kim
Sehyun Kwon
Joo Young Choi
Jongho Park
Jaewoong Cho
Jason D. Lee
Ernest K. Ryu
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (08/2025)
Task Diversity Shortens the Learning Plateau
Jacycon Kim Horvard University
jaeycon jaeyoon_ kiwggnkaagg@aaaaa.......
Sehyun Kwon' Sumsung Research
shyan.kuon@ssmmsunggcom
Joo Young Choi KRAFTON
jychoit@kraftonn.co com
Jongho Park
ijkpark@Werkaleey.ooo edu
KRAFTON
UC Berkeley
Jaewoong Cho KRAPTON
jucho
Jason D. Lee UC Berkeley
josondlee@berkeeey.com edu
Ernest K. Ryu UCLA
eryu@math, Cayyu@math.ulla.coo edu
Reviewed on OpenReview: https: // openreview. form?? id- 7d5
Abstract
In-context learning (ICL) describes language model's ability generate outputs based on set of input demonstrations and a subsequent query. To understand this remarkable capability, rescarchers have studied simplified, stylized models. These studics have consistently observed long loss plateaus, during which models exhibit minimal improvement, followed by sudden, rapid surge of learning. In this work, we reveal that training on multiple diverse ICL tasks simultancously shortens the loss platcaus, making cach task casier to learn. This finding is surprising as contradicts the natural intuition that the combined complexity 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 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. (2020) with GPT-3, describes a language model's ability to generate oulputs 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 reasoning process. To understand this remarkable capability that emerges in language models trained complex real-world language data, researchers such Garg et al. (2022) have studied simplified, stylized models. Tn these studies, transformers are trained from scratch to learn simple functions in context, such linear regression. We thoroughly review this prior work in Section 1.1.
*These authors contributed equally co-first authors. Code available https://gitta https://tthhoo.cm///symkkkooo

---

Published in Transactions on Machine Learning Research (08/2025)
Sparse Linear Regression
Linear Regression
6.0 X
1.4
6.1 X
1.0
0.5
0.0
0
1.7k
10.1k lterations
20k
0 1.7k
10.4k Iterations
20k
Quadratic Regression
LeakyReLU Regression
Parity
1.4
4.8 X
2.8 X
1.4
00 X
1.0
(%%
50
25
2.7k
12.9k
20k 1.7k
4.7k
8k 0
10k
36k 990k
1000k
Iterations
Iterations
Iterations
Figure 1: We train a transformer from scratch on in-context learning tasks. Single-task ICT: Training loss (-) and test error arror/accura /accuracy (---) when each task is trained individually, The Parity task cannot be learned even after 1000k training steps. Multi-task ICL: Training loss (-) and test crror/accurae (---) when all five tasks are trained simultaneously. Green lines mark the plateau escape points. Surprisingly, multi-tast ICL training significantly shortens these plateaus, making training easier.
An intriguing phenomenon observed in these works is the long loss plateau training for ICL, Throughout these platcaus, models display minimal performance improvement, followed by sudden, rapid surge of learning deviating from the typical smooth reduction in training loss. In this work, we reveal that training 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 lcarning process, not shorten it.
In the language multi-task learning, our findings present an instance where multi-task learning easier than single-task learning in the sense 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 describes the experimental setup for training for multiple ICL tasks. Section 3 articulates our main claim that lask diversity shortens the ICL plateau and presents experimental evidence with transformers and state-space models trained 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 pretrained Large Language Models have gained significant attention since first investigated by GPT-3 (Brown et al., 2020). Following this, large body of empirical studies have explored ICL in Large Language Models (Min et al., 2022a;b; Liu al., 2021: Nie et al., 2022; Rubin et al., 2022; Wei et al., 2023). Given the complexity of real-world data, researchers have
2

---

Published in Transactions on Machine Learning Research (08/2025)
explored ICL in more stylized and simplified setups. Garg et al. (2022) formalized 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., 2024), regular language (Akyürek et al., 2024), as well as task mixtures (Tripuraneni et al., 2023), and the ability of Mamba (Park et al., 2024; Grazzi et al., 2024; Li et al., 2024). For more comprchensive overview, please refer to the survey by Dong et al. (2024).
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. (2022) suggested that the ICL process can be interpreted as Bayesian inference. A number of works have argued, both thcoretically and empirically, that transformers implement gradient clescent to learn linear regression in-context (Akyürek et al., 2023; von Oswald et al., 2023; Mahankali et al., 2024; Zhang et al., 2024; Ahn et al., 2023). Beyond these works, many have sought to uncover the internal ICL procedure of transformers, in more complex algorithms (Fu et al., 2023; Giannou et al., 2024; Cheng et al., 2024; von Oswald ct al., 2024; Lin & Lee, 2024); when handling more intricate tasks (Wang et al., 2024; Guo et. al., 2024; Bai et al., 2023; Lin et al., 2024; Wang et al., 2024).
Abrupt phase transition in in-context lcarning. Many studies (Srivastava et al., 2023; Wei et al., 2022; He et al., 2024 Chan et al., 2022; Raventós et al., 2023; Raventos et al., 2023) have shown that a model's ability to perform ICT, 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 plateans, 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 varions sctups (Garg al., 2022; Bhattamishra et al., 2024; Park et al., 2024; Li et al., 2024; Kirsch et al., 2024), although these works did not primarily focus on it.
Beyond these, several studies have explored loss plateaus themselves. Fu et al. (2024) suggested that models learn the features of dataset during loss platcaus. Chen et al. (2024) theoretically demonstrated that a plateau occurs during training when a one-layer transformer is trained by linear regression ICL lask. Olsson et al. (2022) proposed that the plateau escape and the formation of 'induction head' simultaneously happen. Further work by Reddy (2024); Singh et al. (2024); Song et al. (2024), 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, 1997; Baxter, 2000) is a widely used approach for model pretraining, In particular, researchers have identified that task diversily 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., 2020; 2021; Du et al., 2021; Maurer et al., 2016; Crawshaw, 2020; Ruder, 2017), reinforcement learning (Zhang et al., 2023a; Jin et al., 2020; Hu ct al., 2021; Yang et al., 2021; Collins et al., 2021; Lu el al., 2022; Cheng et al., 2022), and natural language processing (Zhang et al., 2023a; Zhao el al., 2023; Zhang et al., 2023b; Hu et al., 2020; Song et al., 2020; Zhou et al., 2019; Gunasekar et al., 2023; Sharma et al., 2023). These works mostly focused on statistical benefits, whereas our finding emphasize on the optimization benefits.
2 Experimental setup
Our experimental setup follows Garg et al. (2022) and Bhattamishra et al. (2024). Consider function class FF with domain X. A sequence model Me (transformer or state-space model) is trained to identify f e F in context and make a prodiction on the subsequent query. The training data consists of sequences of the form P (I1, I [xx1, x2, J(x2)...,,, ,
1), (xx)), where I is sampled from a distribution Dz and task: The loss over the sequence P is given by 1 DT (MM ( (a:)), where E(.,.) is an appropriate loss function N
Ti...., In are independently sampled from a distribution Dx. We train Mg for the next coken-predictee
and P.- = ((1), I2, fx2). Z;-1, (I,,1), 21) is the i-th prefix for 1....,n. In our experiments, n- 120 is a predetermined number shared across all tasks. This procedure is illustrated in Figure 2.
3

---

Published in Transactions on Machine Learning Research (08/2025)
f
fa 1)
Single Task
Multi Tasks
fe
JEFU
Sequence Model (TF, Mamba,
Generate prompt with given
1
(I1, f(z1), ..... 1- (In-1), Uns fzx))
211
f(z1)
Fn-1 f( un
To generate a training sequence, we sample the function FE..U....U where Fi,..., Fu are different in-context function classes, sample 21...., Fn IID, and form (I1, f(z1).(...... Fn fIn- Fn. f f(x))) We refer to the case is trained with autoregressive next-token prediction, ie., the model predicts f(x;) conditioned on (I), fz.)....., I(2,,....1 (2; for i ===,.... I..
Figure 2:
single-task ICL and k 1 as multi-task ICL. The sequence model
ICL tasks. We consider conlinuous ICL lasks and boolean ICL lasks. For each ICL task, an E F is chosen, where F is function class. For continuous ICL tasks, F consists of f: R - IK for d E N and the probability distribution
the domain is assumed to be Dx (00 Ia). For boolean ICL tasks, F consists of f: {+1]" + {=1} for d E N and the probability distribution on the domain is assumed to be Dx = Unif ((+1))). We consider 7 different TCL. tasks: Linear Regression, Quadratic Regression, Sparse Linear Regression, LeakyReLU Regression, Sparse Parity(2), Sparse Parity(3), Parity, The function class F defining each of these 7 ICL lasks precisely stated in Section 5.
Multi-task ICL training. Most prior work on ICL, such as Garg et al. (2022); von Oswald al. (2023); Bhattamishra et al. (2024); Park et al. (2024), considers ICL training with single function class F. In this work, we train models learn functions from the union of function classes -m= Fm in context, where FI....,FN are distinct ICL task among the 10 that we list in Section For instance, if the model is trained the union of linear and quadratic regression tasks, then Fi FF [flf(x)) w U (f xWax.. For each
I,..., ki sample f D and Z1...., IID Dx and form the sequence fZ11, fx2)).... (xx)). This sampling process is repeated B/k times for each m = 1,...,k. making the total batch size. This choice of batch size ensures a fair comparison between single- and multi-task training, as the number of tokens during training exactly proportional to the iteration count. Additionally, we consider a multi-task mixture with uneven probabilities further support. the generality of our claim. To balance the loss scales across different tasks, we normalize each task's loss with constant factors cI....,C Cs (further discussed in Section 4.1) that the training loss stabilizes around during the plateaus. Appendix A provides further experimental details. In expectation, minimize the loss
L(0) = Σ Cmn
E
Σ (MMo))) 1=1
J m.
n
m=
H1D Dxm
Test loss. To evaluate the performance ICL tasks, measure the model's prediction error the last (n-th) output of the prompt. For continuous tasks, the test error is (Ma(Pa) fxx,)) For boolean tasks, the test accuracy is ((sign(M((P..)) = f(x,)., where 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., 2022; Chen et al., 2024; Li et al., 2024; Bhattamishra al., 2024; Park et al., 2024), image datascts (Fu et al., 2024; Kirsch et al., 2024; Singh et al., 2024; Reddy, 2024), and language datasets (Akyürek al., 2024; Olsson et al., 2022). In this section, present the following claim:
4

---

Published in Transactions on Machine Learning Research (08/2025)
Table 1: Task diversity shortens the ICL plateau. We train a transformer with various combinations of different tasks (d = 10). For each run, we report two metrics: the time to escape the plateau and the time 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 3, Sparse Parity(2)) shows the average of (6) results. We find that multi-task training shortens the ICL plateau. To further robustly support our claims, we repeated each experiment three times with varying random seeds and reported the mean values. Details are provided in Appendix C.
Boolean Ticks
Continuous Tasks
Number of tasks
Sparse (2)
Spazse Parity(3)
Linear Regression
Quadratic Regression
LeakyReLU Regression
Sparse Linear Regression
1000k
>> 1000k
35.6k (11.1)
17.1k (19.2k)
9.7k (10.7k)
33.8k (39.6k)
26.5k (43.9k)
60.33 (82.6k)
13.8k (15.8k)
9.9k (12.1k)
3.7k (5.0k)
17.5k (20.5k)
3
8.6k (7.1k)
9.6k (14.9k)
4.8kc (5.8k)
5.1k (7.2k)
2.3k (3.0k)
5.3k (6.4kk
4.7k (5.5k)
6.0k (10.5k)
3.5k (4.5k)
3.2k (5.8k)
2.6k (3.4k)
4.0k (5.0kk)
3.22 (3.9k)
4.8% (8.0k)
3.0k (3.9k)
3.5k (6.0k)
2.6k (3.5k)
3.0k (3.99k)
(3.0k)
(5.0k)
(2.4k)
2.4k (4.8k)
1.9k (2.4k)
1.9k (2.4kc)
Claim: Task diversity shortens the ICL plateau, making each task casier to learn.
Here, task diversity refers to learning multi-task TCL mixture of distinct function classes. The claim that multi-task ICL is casier to learn than single-task ICL is surprising it contradicts the natural intuition that the combined complexity of multiple ICL lasks would lengthen the learning process, not shorten it.
Quadratic Regression
Sparse Parity(2)
Sparse Parity(3)
1.4
1.4
100
1.4
100
1.0
1.0
75
Accuracy (%)
1.0
75
Accuracy (%6]
50
Loss
0.5
0.5
0.5
25
25
0.0
0.0
0
0.0
0
2k
20k Iterations
40k
5.8k
62.9k
100k
26k
242k
300k
Rerations
Iterations
Quadratic Regression
Sparse Parity(2)
Sparse Parity(3)
1.4
1.4
100
1.4
100
1.0
1.0
Accuracy (%6)
1.0
Accuracy (%))
Loss
50
50
0.5
0.5
25
25
0.0
0
7k
15k lberations
30k
21k
62k
80k
36k
60k
77k
Iterations
Iterations
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 vs. Quadratic Regression Lincar Regression. (Middle): Sparse Parity(2) vs. Sparse Parity(2)+ Linnar Regression. (Right): Sparse P3rity(3) vs. Sparse
Regression.
Gencrality of our claim. We provide comprchensive experimental evidence to support this claim. Table summarizes our experimental results on transformers. Figure 3 shows the results of state space models, specifically Mamba (Gu & Dao, 2024) and Hyena (Poli et al., 2023). Across the hundreds of task combinations
5

---

Published in Transactions on Machine Learning Research (08/2025)
Tasks
Complexity
Te
T1 + ........ + Tx
Reduced
Reduced
Reduced
Figure 4: Complexity model of multi-task training. The illustration demonstrates the theoretical framework and computational aspects of multi-task training.
in these setups, we consistently observed that task diversity allows training to escape platcaus more quickly. Table 2 presents results for multi-task learning with uneven probabilities of mixture. Figure 7 shows results in natural language TCT tasks supporting our claim. Refer to Appendix A for experimental details.
However, not all ICL tasks mutually reduce the duration of plateaus. For instance, we consider Regbench task (Akyürek al., 2024), which generated by a random automata. Our experiments in Appendix D.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 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.
Table 2: Muti-task ICL with uneven probabilities. We conducted experiments multi-task ICL with uneven mixture probabilities, where 5 ICL tasks are mixed with the ratio (2 8: 8 1). Under these uneven mixtures, we consistently observed a shortened plateau.
SP(2)
SP(3)
LR
QR
Sparse LR
OC
Do
2.8k (4.2k)
10.2k (11.5k)
2.3k (3.3k)
2.4k (2.9k)
3.0k (6.8k)
2.0k (2.5k)
2.4k (4.1k)
2.0k (2.5k)
Model of task complexity. We quickly describe our mental model of the aggregale 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 the complexity observed when training with just the single task. The complexity of a multi-task ICL is the sum the complexities of the constituent tasks, but the task diversity reduces it. This reduction makes the aggregate complexity of the multi-task ICT. less than the complexity of the individual single-tass ICLs. Section 4.2 discusses when and why task diversity could reduce individual complexities. Figure 4 illustrates this notion.
Escaping plateau & training completion. In the ICL sctups we consider, we observe that (i) there is only one plateau, and (ii) learning is very rapid once this plateau is escaped from, as illustrated in Appendix B. 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) platcau. The results of Table 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., 2022; Bictti al., 2022; Jin al., 2023; Wang & Ma, 2023: Berthier et al., 2024) and empirically (Nakkiran al., 2019: Refinetti al., 2023; Rubruck et al., 2024). Therefore, expect that ICL tasks involving complex hierarchical structures may exhibit multiple plateaus.

---

Published in Transactions on Machine Learning Research (08/2025)
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- no-context learning
LR + QR
LR SP(2)
1.4
1.4
1.0
1.0
1.0
1.0
Distance
Discance
Loss
Loss
0.5
0.5
0.5
0.0
0.0
0.0
15k
10k
lterations
lerations
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 = 00.5) +
Quadratic Regression(u = 0.5). (Right): Linear Regression(p = 1) + Sparse Parity(2).
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-contert function as
g == argmin, J~DD,eNDx [[(()) - f(x))]
i.e., 95 is the context-indepeenent function that minimizes the test error. For boolean ICL, tasks, g 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 F= (fl((()) wir} and Dj is given by w ~ NN,, Ia), then 9 (z) = nTx. As another example, if F = f=)
and Dy is given by W ~ NNU, laxa), then ()(r)= = vUr. The ICL plateau corresponds to task-wise no-contert learning, which we describe in the following. When we train a model Mg for ICL with function classes FI,...,Fe, the model's output during its plateau corresponds to
M&(P.) = (x,), Pa is sampled from Fm-
In other words, the model identifies the function class Fm E (Fi...., FR) and then applics the optimal no- context function corresponding to Fm. The in-context demonstrations (Iı, (nn),...,,,1111 ((nn are used to determine the function class Fm [FI...., } but not to determine the specific f f Fm. This claim can be verified by measuring the error between Mo(P.) and 8i, which we plot in Figure 5. Revisiting the Sparse P2rity(2) task, the accuracy during plateau is indeed attributed to 0.55 = EynDs (1()((x) = f(x))) Appendix G provides further details on no-context learning regime.
In Scction 2, we used the scaling factors c Cg to normalize cach task's empirical loss. Specifically, we sot Cm = 1/Ej~Dy..we D> [((g (x). fx))]. Our findings on no-context learning imply that the normalized loss will have a plateau of height
7

---

Published in Transactions on Machine Learning Research (08/2025)
Task
1.0
Loss
Checkpoint Task A
Task B
from scratch
Loss
0.2
LR
QR
(912)
S913)
from Task ckpl
Task
0.0
Iterations
Figure 6: (Left) Illustration. We pre-train on Task A and extract checkpoint + 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 (bluc color) indicating that the model pre-trained on Task A signilicantly aids the learning of Task B.
4.2 Common structure across ICL tasks
Consider a multi-task ICL setup, where model is trained set of tasks U Tm Suppose there exists "common structure" C shared across tasks. Denoting the remaining part of each task Ims we decompose each task as Tm = C + Im for m = 1,.... k. Thus, multi-task ICL training is decomposed into two sub-problems: [learning C] and [learning II....,ZAA 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 shared structure makes il 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 toy experiment on feature learning.
4.2.1 Checkpoint experiment
Consider the following chockpoint experiment. For each single-task training with task A, we save the model as it escapes the plateau as rllustrated 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 lask B with shortened plateau. This implies that (1) task A and task B share common structure and that (2) the plateau arises from the difficulty of learning the common structure. We conduct analogous experiment on natural language ICL tasks and observe similar results, shown in Figure 15 of the appendix. Further details are presented in Appendix E.
4.2.2 Generality of our hypothesis
Consider the following intuition. When U Tm are trained concurrently, the model receives multiple "views" of the common structure through the different compositions + Im for m = I....,k. We hypothesize (3): These multiple "views" of C make casier to learn in the sense of more favorable optimization landscape. This may be the key mechanism allowing task diversity shorten the ICL plateau, and this phenomenon may extend beyond the ICL setup.

---

Published in Transactions on Machine Learning Research (08/2025)
The following feature learning experiment makes this intuition more concrete. Consider the 2-layer feature learning setup, setup with large body of prior work (Damian et al., 2022; Ba et al., 2022; Dandi et al., 2023; Wang et al., 2023). For imput I E R2, the true function to learn is fa (z) = Uo(A,z), where U E g and a true feature matrix A. E RAxd, The goal is to learn f = f* with f(x) = Vo(wx), where h' > h, V R and W E RRx' The loss function is:
&
L(W = E
Σ
EAND
(o ( uTo(A.)))
m=1
U= [1
ux V = [ 2k]
The idea is that u1,--.., Ug represent & sub-tasks that share the common feature matrix A, If we sample "i...., Ug from & different distributions, this corresponds to multi-task training of k distinct tasks. Conversely, if we sample "I,..., Me 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 singlo-task training, as shown in Figure (right). Figure 19 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 platcau and makes our intuition more concrete through analogy. The results of this toy model, shown in Figure 7, lead to make the general hypothesis (3): Training multiple tasks jointly with a shared structure makes it casier to learn the common structure.
WS(4)
Single-task
10
WS(4) WS(2)
Multi-tasks
8-
WS(4) + WL
8
Loss
6-
Loss
4
4-
2-
2-
0-
30.0k
130.0k
180.0K
7.0k
20.0k
Iterations
Iterations
Figure 7: (Left) Language ICL task. Previous work Fu et al. (2024) 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 D.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 F for further details.
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 algorithu to identify the specific function being demonstrated by the in-context demonstrations.
A plausible candidate the induction head, a circuit that scarches over the sequence for previous instances of current loken and prediets the same completion again. Indeed, Olsson et al. (2022) 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. (2024); Singh et al. (2024); Reddy (2024).
Retrieval. Sample 1024 5-tuples of one key and four values: {(k, u/ of, vi. H 1034 The k,s and u/s are independently sampled from Dx and Dy, respectively. To generate prompts, we uniformly sample (n - 1) 5-tuples without replacement and uniformly choose one i per 5-tuple, resulting in [(k, v.) Next, we sample p ~ Unif ((I,..., n - -1}) and set q = kp- Finally, given P. = (k1, Viy.rr kn-1, Vn-1, q), the task is
9

---

Published in Transactions on Machine Learning Research (08/2025)
to predict Vp. We consider two Retricval tasks: Gaussian Retrieval with (Dc, Dy) = (VV(0, 1)) and Boolean Retrieval with (Dx, Dy) (Unil (+1]4, Unif (+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 induction head. (It is unclear whether an induction head is useful for our ICL tasks at all.) For further details, please refer to Appendix E.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, quickly list and define the ICL tasks that consider, which are primarily adapted from Garg al. (2022); Bhattamishra et al. (2024). Each ICL task is specified by F the function class, D) a probability distribution over the function class, and Dx a probability distribution over the inputs. For continuous ICL tasks, Dx = (0, Ia). For boolean ICL tasks, Dx = Unif (=1]d.
Continuous ICL lasks:
Linear Regression. F= (f(f(())) wtz}. Dy: Each element of w E I&d is independently sampled from N(u,1)
Quadratic Regression. F (flf(z) = xWW. Dz: Each element of W E gdxd is independently sampled from 1 N1, 1).
Sparse Linear Regression. F F= - *). Dy: Each element of w E Rd is independently sampled from 1), To sample Wsparse; we uniformly choose k coordinates and retain the corresponding coordinates of w.
LeakyReLU Regression. F= ( [f) LeakyReLU (wTx)} with negative slope a 0.5. Dz; Each elements of w e Rd is independently sampled from N(p,1) 1).
Boolean ICL tasks:
Sparse Parity(k)). ( = -Ilee x[i)). Dz: A A (I,...,, is uniformly sampled subset of size k.
Parity. F= (flf(x)= = IleA zli]}. Dy: AS (l....,. d} is uniformly sampled subset, regardless of the size.
6 Conclusion
In this work, we identify that training on diverse set of multiple TCT. tasks 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 tataallltysccaee" /scale" seen in the modern era of deep learning. Indeed, LLM training via highly diverse multi-tas tasti learning, requiring wide
next-token prediction can be thought of effectively range of reasoning skills for 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.
10

---

Published in Transactions on Machine Learning Research (08/2025)
7 Acknowledgements
This research was supported by a grant from KRAFTON AL. We thank Wonjong Rhee for initiating this project as part of a course at Seoul National University. We are grateful to Junsu Kim, Jaesung Park, and Hyunsik Chae for their valuable feedback, and to the participants of the KAIST-SN workshop for their insightful discussions at the early stages of this work. We also thank Chanwoo Park for sharing his empirical insights, and Minkyu Kim for providing feedback on the figures.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NewrIPS, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. ICLR, 2023.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architectures and algorithms. ICML, 2024.
Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. NewrIPS, 2022.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as slatisticians: Provable in-context learning with in-context algorithm selection. NewrIPS, 2023.
J. Baxter. model of inductive bias learning. Journal of Artificial Intelligence Research, 12:149 12:149-198, 198, 2000.
Raphaäl Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales two-layers neural networks. Foundations of Computational Mathematics, 2024.
Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and 1mms by learning to learn discrete functions. ICI.R, 2024.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. NewrIPS, 2022.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Lilwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Πya Sutskever, and Dario Amodei. Language models are few-shot learners. NewrIPS, 2020.
Rich Caruana. Multitask learning. Machine Learning, 28(1):41 75, 1997.
Stephanie C. Y. Chan, Adam Santoro, Andrew K. Lampinen, Jane X. Wang, Aaditya Singh, Pierre H. Richemond, Jay McClelland, and Felix Hill. Data distributional propertics drive emergent in-context learning in transformers. NewrIPS, 2022.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax atlention for in-context learning: Emergence, convergence, and optimality. COLT, 2024.
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. ICML, 2024.
Yuan Cheng, Songtao Feng, Jing Yang. Hong Zhang, and Yingbin Liang. Provable benefit of multitask representation learning in reinforcement learning, NeurIPS, 2022.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. TCML, 2021.
11

---

Published in Transactions on Machine Learning Research (08/2025)
Michael Crawshaw. Multi-task learning with deep neural networks: A survey. arXiv 2009.097900 2020.
Alex Damian, Jason D. Lee, and Mahdi Soltanolkotabi. Neural networks can learn reprosentations with gradient descent. COLT, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at time. Mathematics of Modern Machine Learning Workshop at NewrIPS 2023, 2023.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Baobao Chang, Xu Sun, Lei Li, and Zhifang Sui. A survey on in-context learning, aiv 2301. 2301.00234, 2024.
Simon Shaolei Du, Wei Hu, Sham M. Kakade, Jason D. Lee, and Qi Lei. Few-shot learning via learning the representation, provably. ICML, 2021.
Deqing Fu, Tian-Qi Chen, Robin Jía, and Vatsal Sharan. Transformers learn higher-order oplimization methods for in-context learning: A study with linear models. arXiv 2310.17086, 2023.
Jingwen Fu, Tao Yang, Yuwang Wang, Yan Lu, and Nanning Zheng. Breaking through the learning plateaus of in-context learning in transformer. ICML, 2024.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. NewrIPS, 2022.
Nikhil Ghosh, Song Mei, and Bin Yu. The three stages of learning dynamics in high-dimensional kernel methods. ICLR, 2022.
Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How well can transformers emulate in-context newton's method? arXiv 2403. 03183, 2024.
Riccardo Grazzi, Julien Niklas Siems, Simon Schrodi, Thomas Brax, and Frank Hutter. Is mamba capable of in-context learning? ICLR Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. ar 2312.00752, 2024.
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah, Harkirat Singh Behl, Xin Wang, Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need. arXiv 2306. 11644, 2023.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representations. ICLR, 2024.
Tianyu He, Darshil Doshi, Aritra Das, and Andrey Gromov. Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks. arXiv 2406.02550, 2024.
Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear rl JCML, 2021.
Junjie Hu, Sebastian Ruder, Aditya Siddhant, Graham Neubig, Orhan Firat, and Melvin Johnson. XTREME: A massively multilingual multi-task benchmark for evaluating cross-lingual generalisation. JCML, 2020.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I. Jordan. Provably efficient reinforcement learning with linear function approximation. COLT, 2020.
Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon S. Du, and Jason D. Lee. Understanding incremental learning of gradient descent: A fine-grained analysis of matrix sensing. ICML, 2023.
12

---

Published in Transactions on Machine Learning Research (08/2025)
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. arriv 2212.04458 2024.
Yingoong Li, Xupeng Wei, Haonan Zhao, and Taigao Ma. Can mamba in-context learn task mixtures? JCML 2024 Workshop on In- Contert Learning, 2024.
Licong Lin, Yu Bai, and Song Mei. Transformers decision makers: Provable in-context reinforcement learning via supervised pretraining. ICLR, 2024.
Ziqian Lin and Kangwook Lec. Dual operating modes of in-context learning. ICML, 2024.
Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? aix 2101.06804, 2021.
Rui Lu, Gao Huang, and Simon S. Du. On the power of multitask representation learning in linear mdp. NewrlPS, 2022.
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. ICLR, 2024.
Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes. The benefit of multitask representa- tion learning. JMLR, 17(81):1-32, 2016.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetaICT: Learning to learn in context. ACL, 2022a.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? EMNLP, 2022b.
Preetum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang, Benjamin L. Edelman, Fred Zhang, and Boaz Barak. Sgd on neural networks learns functions of increasing complexity. NeurIPS, 2019.
Kim Anh Nguyen, Sabine Schulte im Walde, and Ngoc Thang Vu. Distinguishing antonyms and synonyms a patter neural network. EACL, 2017.
Feng Nie, Meixi Chen, Zhirui Zhang, and Xu Cheng, Improving few-shot performance of language models via nearest neighbor calibration. ICML, 2022.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hoddield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndoussc, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Xr 2209. 11895, 2022.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to Icarn? a comparative study on in-context learning tasks. /CML, 2024.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. ICML, 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenA/ blog, 2019.
Allan Raventos, Manshooj Paul, Feng Chen, and Surya Ganguli. The effects of pretraining task diversity on in-context learning of ridge regression. ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
13

---

Published in Transactions on Machine Learning Research (08/2025)
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. NewrIPS, 2023.
Gautam Reddy, The mechanistic basis of data dependence and abrupt learning in an in-context classification task. RCLR, 2024.
Maria Refinetti, Alessandro Ingrosso, and Sebastian Goldt, Neural networks trained with sgd learn distributions of increasing complexity. ICML, 2023.
Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context learning. ACI, 2022.
Jirko Rubruck, Jan P. Bauer, Andrew Saxe, and Christopher Summerfield. Early learning of the optimal constant solution in neural networks and humans. arXiv 2406. 1.06.17460 17467, 2024.
Sebastian Ruder. An overview of multi-task learning in deep neural networks. orXiv 1 1706. 05098, 2017.
Brihat Sharma, Yanjun Gao, Timothy Miller, Matthew M. Churpek, Majid Afshar, and Dmitriy Dligach. Multi-task training with in-domain language models for diagnostic reasoning, arXiv 2306. 04551, 2023.
Aaditya K. Singh, Ted Moskovitz, Felix Hill, Stephanic C. Y. Chan, and Andrew M. Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. ICML, 2024.
Jiajun Song, Zhuoyan Xu, and Yigiao Zhong, Out-of-distribution generalization via composition: lens through induction heads in transformers. arXiv 2408.09503, 2024.
Linfeng Song, Kun Xu, Yue Zhang, Jianshu Chen, and Dong Yu. ZPR2: Joint zero pronoun recovery and resolution using multi-task learning and BERT. ACL, 2020.
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, and al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. TMLR, 2023.
Eric Todd, Millicent Li, Arnab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. ICLR, 2024.
Nilesh Tripuraneni, Michael I. Jordan, and Chi Jin. On the theory of transfer learning: The importance of task diversity. NewrIPS, 2020.
Nilesh Tripuraneni, Chi Jin, and Michael I. Jordan. Provable meta-learning of lincar representations. ICML, 2021.
Nilesh Tripuraneni, Lyric Doshi, and Steve Yadlowsky. Can transformers in-context learn task mixtures? NewrIPS Workshop on Distribution Shifts: New Prontiers with Foundalion Models, 2023.
Ashish Vaswani, Noam Shazcer, Niki Parmar, Jakob Uszkorcit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Atlention all you need. NewrIPS, 2017.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. TCML, 2023.
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zuechet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agũera y Arcas, Max Vladymyrov, Razvan Pascanu, and João Sacramento. Uncovering mosa-optimization algorithms in transformers. ICLR Workshop on Mathematioal and Empirical Underslanding of Foundation Models, 2024.
Mingze Wang and Chao Ma. Understanding multi-phase optimization dynamics and rich noulinear behaviors of relu networks. NeurIPS, 2023.
Yifei Wang, Yuyang Wu, Zeming Wei, Stefanie Jegelka, and Yisen Wang. A theoretical understanding of self-correction through in-context alignment. XrXiv 2405. 18634, 2024.
14

---

Published in Transactions on Machine Learning Research (08/2025)
Zhichao Wang, Andrew Engel, Anand Sarwate, loana Dumitriu, and Tony Chiang. Spectral evolution and invariance in linear-width neural networks. NewrIPS, 2023.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models. TMLR, 2022.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, and Tengyu Ma. Larger language models do in-context learning differently, arXiv 2303.03846, 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. ICLR, 2022.
Jiaqi Yang, Wei Hu, Jason D. Lee, and Simon S. Du. Impact of representation learning in linear bandits. ICLR, 2021.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. JMLR, 25(49)11-55, 2024.
Yi Zhang, Arturs Backurs, Sebastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with LEGO: syuthetic reasouing task. arXiv 2206. 040000 2023a.
Zhihan Zhang, Wenhao Yu, Mengxia Yu, Zhichun Guo, and Meng Jiang. A survey of multi-task learning in natural language processing: Regarding lask relatedness and training methods. ACL, 2023b.
Yulai Zhao, Jianshu Chen, and Simon S. Du. Blessing of class diversity in pre-training. AISTATS, 2023.
Wenjic Zhou, Minghua Zhang, and Yunfang Wu. Multi-task learning with language modeling for question generation. ACL, 2019.
15

---

Published in Transactions on Machine Learning Research (08/2025)
A Experimental details
Model Architectures. We focus on transformer (Vaswani et al., 2017), Mamba (Gu & Dao, 2024), and Hyena (Poli et al., 2023). To handle model architectures that process inputs and outputs as vectors in an embedding space, we first pad each f(x;)) with (d - 1) zeros to match the dimension of z4s. 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 a scalar. Each model's hyperparameters are dictated as follows. For transformer, we use GPT2 (Radford et. al., 2019), 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 the hidden state to output. To train a model from scratch, we randomly initialize the model parameters.
Configuration. We use the Adam optimizer (Kingma & Ba, 2015) with learning rate of 0.0001. We use B 64, making each task's batch size 64/k. Regarding the choice of ((),), we use mean-squared error loss for contimuou 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 pllateau is defined as fllateau min 100 Ei=t-g Training loss(t') < 0.8 This effectively captures the plateau es- cape as we have normalized each task's loss around 1. The training exit conditions are defined 100 Ei= Test error(t') 0.2 for continuous ICL tasks and texit
follows texit min42100
min 100 100 E 2 Test accuracy(1") > 0.95 for boolean TCL. tasks. (platean and texit measured for each individual task.
Batch generation process The following pseudo-algorithm summarizes the batch generation process of (multi)-task ICL training.
After constructing batch, minimize the following empirical loss. Denote flim) and a.fima) the corresponding in-context function and i-th input for each p(.mm)
n
1(0):
e (Ma (Pf.m()) (P(..))), platm)) (o(tmm)))
Cm
m=1
Algorithm Batch generation process
Require: Multi-tasks (Fi....,
Requíre: Corresponding domains {Dx, ........ Dx,
Require: Batch size for each task (BI..... Ba}
Ensure: Generated batch of prompts
1: Initialize Batch B <-6
2: for m = to do
3: for j= 1 to Bm do
Sample function f f~ Fm
Sample imputs 21,..., In 11D Dxi
Generate prompt: plim) = (I1, fer,((11))) 2n-11: In-1, f(zn-1), Ina f(xn))
7: Update batch B + BU {plim)
8: end for
end for
16

---

Published in Transactions on Machine Learning Research (08/2025)
B Training loss dynamics: One plateau, then rapid drop
In the ICL setups consider, we observe that (i) there is single platcau phase, and (ii) learning procoeds very rapidly once the model escapes this plateau. We provide several representative examples of training dynamics that support this observation below. For the cases Sparse Parity(2) and Sparse Parity(3), training was performed jointly on a combination of three tasks: Sparse Parity(2), Sparse Parity(3), and Sparse Linear Regression.
Sparse Parity(2)
Sparse Parity(3)
Training loss
Training loss
1
1
0.8
0.9
0.8
0.6
0.7
0.4
0.6
Iteratio
Iteration
0.2
0
8k
Linear Regression
Quadratic Regression
Training loss
Training loss
0.9
0.8
0.8
0.6
0.7
0.5
Iteration
Iteratio
10k
20k
30k
40k
0
5k
10k
15k
20k
Leaky ReLU Regression
Sparse Linear Regression
Training loss
Training loss
1
1.4
W/M /Nw
1.2
hy MAA MM Aowee
0.8
1
0.6
0.8
0.6
0.4
0.4
Iteration
Iteration
0.2
6k
10k
10k
20k
50k
Figure 8: Training loss curves showing a single plateau followed by rapid convergence.
17

---

Published in Transactions on Machine Learning Research (08/2025)
C Detailed analysis of table 1
In Table 1, we provided comprchensive experimental evidence demonstrating that task diversity shortens the in-context learning (ICL) plateau, facilitating faster learning. To further robustly support our claims, we repeated each experiment three times with varying random seeds. The figure below plots the mean values along with standard deviations. It clearly shows a strong tendency: as the number of tasks increases, both the time to escape the platcau (bluc) and the time to complete learning (orange) decrease significantly.
Sparse Parlty(2) - plateau
Sparse Parity(2) complete
escapee the plateeu
8
complete leatning
Reeation
80000
Iteration
100000
40000
75000
52000
20000
25000
Number of tasks
Numbar of tasks
Figure 9: Sparse Parity(2)
Sparse Parity(3) plateau
Sparse Parity(3) complete
time escape the plalesu
Urme tu complete the learning
=1 std
11 std
iteraton
150000
ittration
200000
100000
150000
100000
50000
50000
6
Number of tasks
Number of tasks
Figure 10: Sparse Parity(3)
Linear Regression plateau
Linear Regression complete
70300
60000
time escape the platesu
Uime to complete the learning
=1
60000
50000
50000
40000
Ieeatton
Iteration
40000
20000
30000
20000
20200
10000
10200
Koumbeer
Number
Figure 11: Linear Regression
18

---

Published in Transactions on Machine Learning Research (08/2025)
Quadretic Regression plateau
Quadratic Regression complete
20000
time the platesu
complete the learning
17500
20000
15000
Ieeratn
12300
Iteration
15000
10000
10000
7500
3000
5000
2500
Number of tasks
Number tasks
Figure 12: Quadratic Regression
Leaky Relu Regression plateau
Leaky Relu Regression complete
12000
time escape the platesu
12000
complete learning
=1 std
#1 std
10000
10000
Iterato
8000
6000
6200
4000
4000
2000
2000
Number of tasks
Number of tasks
Figure 13: Leaky ReLu Regression
Sparse Linear Regression platesu
Sparse Linear Regression - complete
50000
time escupp the plateeu
60000
Urn tu complete Uhe learning
+1 std
+1 std
50200
40000
40200
iteraton
30000
iteraton
X0000
20000
20200
10000
Number of tasks
Number of tasks
Figure 14: Sparse Linear Regression
19

---

Published in Transactions on Machine Learning Research (08/2025)
D Language ICL tasks
D.1 Natural Language ICL tasks
WS(4)
10
WS(4) from WS(2) ckpt
WS(4) from WL ckpt
Loss
25.0k
100.0k Iterations
180.0k
Figure 15: Checkpoint experiments language ICL tasks.
To further verify claim tasks with real-world language data, examine ICL in natural language processing (NLP) tasks.
Model. We use GPT2 (Radford al., 2019) 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. (2017). Following the approach in Todd et al. (2024), choose 5245 words that can be tokenized to 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. (2024), where words are separated by a space, and examples are separated by a line:
<bos>>. word wordh ... word
A. answer
Q. words word ... words
A. S
Q. word word word
A.
We devise WordSelection(d) by generalizing the WordSelection Task (Fu et al., 2024). The goal of this task to select a single word from given words. The modol's goal to learn in-context which word select.
WordSelection(d).
We also devise WordLength Task. Let len(w) be the length of word w and % be the remainder operator. Given four words Wi.... wa, the goal of this task is to learn in-context either Σi= len(w;) or Σi= (11 from given examples.
WordLength.
Result. Prior work Fu et al. (2024) reported the difficulty of learning WordSclection(4), by showing that it cannot learn within about 78k iterations. Indeed, as shown in (Figure 7, Right), single-task training
20

---

Published in Transactions on Machine Learning Research (08/2025)
of WordSelection(4) exhibits plateau for over 180k iterations. However, mixing WordSelection(4) and WordSelection(2) significantly shortens the plateau to 3.2k 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 15, using checkpoint models shortens the platean, demonstrating that NLP tasks also share certain common structure.
D.2 Formal Language ICL task: Regbench
We follow the configuration of Akyürek et al. (2024). To jointly train Quadratic Regression and DFA tasks, since their input formats differ, we use two different input embedding layers. Figure 16 illustrates the result.
DFA
Quadratic Regression
1.5
11
Cross Entrogy Loss
1.3
MSE Loss
1.2
1.1
1.0
0.8
20000
40000
60000
80000
20000
40000
60000
80000
Heerations
Iterations
Figure 16: 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.
E
Checkpoint experiment
E.1
Experimental details
In this section, provide the configurations for the checkpoint experiments in Section 4.2.1. The imput dimension is set to d = 10. For each single-task training, save the model once it escapes the plateau. The exact iteration numbers where the models are saved are provided in Table 3. Since Sparse Parity(2) and Sparse Parity(3) do not escape the plateau within observable training steps, no checkpoint models could be obtained.
Boolean Task
Continuous Task
Boolean
Lincar Regression
Quadratic Regression
Sparse Linear Regressiio
Gaussian
Retrieval
Retrieval
Iteration number
105k
19k
12k
7k
50k
Table 3: Checkpoint models.
After saving the bascline models, we train each task (Lincar Regression, Quadratic Regression, Sparse Lincar Regression, Sparse Parity(2), and Sparse Parity(3)) according to the configuration in Appendix A. To caleulate the ratio, divide the result by the plateau escape time summarized in the first row of Table 1.
21

---

Published in Transactions on Machine Learning Research (08/2025)
Learning time comparison of checkpoint models. In Figure 6, 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 17 confirms that this indeod the case, although the effect less robust compared to Figure
Checkpoint Task A
0.8
UR
0.6
9n
0.4
SpLR
0.2
Sp
5PE3
Task
0.0
Figure 17: Learning time ratios of checkpoint models.
22

---

Published in Transactions on Machine Learning Research (08/2025)
E.2
Experiment on Retrieval task
Recall the description of Retrieval lask.
Retrieval. Sample 1024 5-tuples of one key and four values: {(k,, 2/ vz uf, h The k,s and u/s are independently sampled from D and Dv, respectively. To generate prompts, we uniformly sample (n 1) 5-tuples without replacement. and uniformly choose one e per 5-luple, resulting {(k, v:)] Next, sample p Unif (fI...., 1}) and set q = kp. Finally, given Pa (k,, VI....,,, 1, 1n q), the task is to predict Vp- We consider two Retrieval tasks: Gaussian Retrieval and Boolean Retrieval, with (Dx, Dy) (0 (o,, 1)) and (Di, Dv) (Unif {=1]", Unif {=1}), respectively.
To train contimuous ICL tasks, we use chockpoint 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 Figure 18.
0.2
0.8
1.0
Retrieval
5
QR
902
Other Tasks
Figure 18: Checkpoint experiment Retrieval.
F
Feature learning experiment
The true function is (x) = Uo(A.z) where U E Rkxh and A, E Rhxd, Columns of U, denoted as um88 are independently sampled from N (pm Ia), with (H1...., being pre-determined vectors prior to training.
To design multi-task learning, we independently sample different vectors ;i...., from /h Sh-1, making N(mm la)s differ. In contrast, for single-task training, all ,1.-.., are set to the same vector μt, which is also sampled from vh x Sh-1.
The model that learns the true function is f(x) = Vo(wx) where V E ma and W E B xd. We set d = 150, h E {10, 20}, h E {100, 300), and & E {15, The Adam optimizer is used with a learning rate of Ir = 0.001 and we apply 0.1 for the model's heads, following the approach in Berthier et al. (2024). We use the sigmoid activation function (xx) T+c-x.
23

---

Published in Transactions on Machine Learning Research (08/2025)
Additional Figures. Figure 7 represents the result when (d, h1, k) (150, 10, 100, 15). To enhance the generality, 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).
20
22.06
borations
borazions
Iborotion
borctions
(a) (6', 1b) = (100, 15)
(b) (1', k) = (100, 30)
(c) (h',, .k) = (300, 15)
(d) (N', k) = (300, 30)
12.54
25
12.0
20/A
12.0
12.34
Berations
Berations
(e) kh', (100, 15)
(f) (h A)) (100, 30)
(g) (h', k) (300, 15)
(h) (N', R) (300, 30)
Figure 19: First row corresponds to the case when (d, h) (150, 10). Second row corresponds to the case when (d, h) = (150,20)
G No-context learning
G.1 Optimal no-context functions
In this section, we provide detailed explanation of the optimal no-context function. We have defined the optimal no-context function as follows. For continuous ICL. tasks, g argmin, Begniin,E-Dy,---20 ( f(x))), where the argmin is taken over {g (glg: Rd + R}. For boolean ICL tasks, o = argmax, DDDr,, [1(g(x) f(x)),, where the argmin is taken over {g [g: {+1) {=1}}.
For boolean ICL (asks, we can calculate 9 for each x {+1]d. The closed form of gx) can be expressed using G2: (f() 15 F} follows.
#(1 E Ga) > #(-1e (otherwise) gx(x) - {
-{1
=
For continuous ICL tasks, we cannot apply the same approach since the output space not discrete. However, can still derive the closed form for (x) for most tasks: (z): E,~Dy [f()]. This is because
(x): argmin [(((x) a) argmin Eソ-D>
For tasks other than LeakyReLU Regression, property, as listed follows.
obtain the closed forms of g (x) using the above
Linear Regression: (+) = (p, ( Tx)
Quadratic Regression: (x) T. Un, (U)ij = p.
Sparse Linear Regression: (x()) $/2 (H,...,, Tx.
Decision Tree: (x) μ.
For LeakyReLU Regression, we can estimate (x) with small error through sufficient number of samples from Dz.
24

---

Published in Transactions on Machine Learning Research (08/2025)
G.2 Experimental details and additional figures
Let tasks with in-conlext function classes Um= Fm be trained simultaneously. Note that we do not discern the single-task and multi-task training: & can be 1 as well. During training, we measure
E D~D
(P. 97,, (In)), P. = (21), ................................ fn-1:f(En-1) Fn)
UD Dam
for continuous ICL tasks and
E JnDs Bi Rn Dxm
H1D ((ssgg((Mo(P.)) + 93., (Ea))) P. - (())... (..,, In-1, (In- vn).
for boolcan ICL tasks. As shown by the following figures, these mcasurements are close to O, which indicates that the model's output corresponds to task wise optimal no-contezt function. We follow Appendix A for other configurations. Moreover, to prevent the optimal no-context be merely zero function, we μ to parameterize the in-context function classes.
1.4
1.5
1.0
1.0
Dratuncc
0.0
0.0
0.
100k
31.5k
36k
Berral AET
Banabonn
(a) Sparse (2)ity(2)
(b) Quadratic Regression (μ
(c) QR [-=-1) + SP
Figure 20: No-context learning regime. During the plateau, the model output matches the task-wisc optimal no-context function.
25


Reviews:
[{"review_id": "awyzS9tx0r", "reviewer": "Reviewer_UhrC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strenghs \n\nClear empirical result -  cutting optimisation time across six synthetic tasks and three very different architectures is very nice divercity in results.\n\nConvincing transfer test - a checkpoint that has escaped the plateau on Task A instantly accelerates Task B, backing the claim that multitask exposure uncovers reusable structure rather than just adding data\n\n\n# Weaknesses\n\nThe paper rules out induction heads via Retrieval checkpoints but offers no concrete mechanistic account.\n\nThe models are outdated, would the findings stay valid with more current LLMs, as well as reasoning models.\n\nICL papers turn to use real data for the minimum text classification (AGNews, sentiments and so on), the proxy experiment may not lead to real world findings and the authors should conduct this experiments.\n\nEscape time is defined by a fixed training-loss threshold (0.8). Different tasks have different loss scales even after normalisation—robustness to threshold can be shown by experiments? e.g. it's not task agnostic and need manual tune?\n\nStatistical confidence. Plots show single training runs; variance over random seeds / initialisations is not reported and known to have large impact in ICL due to the choise of examples. I would suggest to show mean+-std.\n\nBatch-size confound. Multi-task batches contain more distinct functions per update. Could the speed-up stem from effective larger gradient variety rather than diversity per se?\n\nLinks to broader optimisation literature (e.g. curriculum learning, data-ordering effects) are only briefly mentioned. This is not a question but a kindly note to make a greater link to literature.\n\nNormalising losses to plateau at 1 is convenient, but could mask dynamics; ablations without this heuristic are absent about the coefficients.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Training on a mixture of diverse in-context-learning tasks deal with the loss plateaus that plague single-task ICL.\n\nThe authors show this effect across six synthetic regression/parity tasks, two state-space models (Hyena, Mamba), a GPT-2-like transformer, and they demonstrate that a checkpoint that escapes the plateau on one task jump-starts learning on others.\nTogether these results position task diversity as a powerful implicit curriculum for faster ICL optimisation, while highlighting open questions about the latent structure being learned and whether the speed-up persists in modern, billion-parameter LLMs.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "please adress the weaknesses", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "X9rHlNd12A", "reviewer": "Reviewer_bBGm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This is an interesting (and counterintuitive) finding, and I think it is both explained well and studied carefully. To the best of my knowledge, the concept of the loss plateau as representing a period where the model provides the optimal no-context answer is also novel, and both helpful intuitively and substantiated empirically. I also think the paper is strengthened by the inclusion of non-transformer architectures and by the brief study of language tasks in the appendix. \n\nI think the claims about applicability to NLP in B.1 should be moderated slightly, as these tasks remain quite artificial and differ substantially from the complexities of tasks with real language data. Similarly, I think the claim in both the end of the abstract and end of the intro that this could be a factor in large-scale language model training should either be adjusted to indicate further uncertainty or substantiated with stronger evidence; I think this is an interesting phenomenon and it's certainly possible that it plays a role in language model pretraining, but there's really not a lot of evidence to draw this conclusion from studying much smaller models trained from scratch in limited data regimes.\n\nI also think it would be interesting to include more study of tasks that have non-reciprocal loss plateau shortening-- e.g. quadratic regression and DFA in B.2-- or task pairs that don't exhibit this behavior at all, like training Boolean Retrieval before Boolean ICL. Can you characterize any properties of such tasks? The titular claim is that task diversity is helpful, but it seems this is not always the case-- is it possible that tasks that are diverse _to a point_ are useful (e.g. as you hypothesize, to learn formatting), but tasks that are higher diversity in form do not show optimization benefits from multi-task training. More characterization of the situations where this phenomenon occurs would strengthen the work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies the optimization behavior of models trained from scratch on in-context learning tasks. Prior work documents a loss plateau followed by rapid learning on single-task ICL; this paper considers the relative optimization behavior of single-task and multi-task ICL. \n\nThe authors show that training with two or more ICL tasks at once results in shorter loss plateaus than training on any individual task. They demonstrate that this behavior is partially but not fully transferable to a sequential training setup by taking a checkpoint trained on one task to the end of the plateau and training it on a second task; a loss plateau for the second task exists in this setting but is very brief. They hypothesize that the loss plateau is a period of _no-context learning_, where the model outputs the optimal \"no-context\" answer -- e.g. the single value with minimal test error --  for all inputs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I'd most like to see more analysis of the situations where task diversity does not shorten the learning plateau for one or both of the tasks; however, even without this, I think the paper makes an interesting contribution.\n\nOther, more minor suggestions:\n- slightly more careful language around the applicability of the claims to NLP/LLMs\n- I think Appendix B could be a section of the main paper-- the results are quite interesting, and it seems you do have space to include it in the main body.\n- typo in Figure 5 caption: \"raining\" -> \"training\"", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "oq91USu9Bf", "reviewer": "Reviewer_wB6h", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper provides compelling evidence of and investigation into the benefits of task diversity for learning. \n\nWeaknesses:\n- Some of the terminology is rather vague, which makes the claims more difficult to understand. See Requested Changes for examples.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the observation that training a model on a diverse set of tasks/functions makes the model learn more quickly than training it on more data from a single task. The authors hypothesize that common structure shared among a set of functions/tasks is what allows it to learn this structure more easily, because that structure is present across multiple \"views\" from different tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please clarify the difference between \"statistical\" benefits of task diversity vs. \"optimization\" benefits.\n- I'm confused why the test error is lower than training error in Figure 1.\n- Please clarify how many runs are used for each training curve shown in the paper. \n- Please specify why the particular tasks were chosen (beyond that they were used in prior work). \n- In Table 2, please clarify which model was used in this experiment, and which task received 1/2 of the examples during training. Or, is each column a different training run? What is the difference between the number in and out of parentheses? \n- Typo: \"raining loss\" in Figure 5\n- Can you provide some analysis on the kinds of examples the model is learning to get correct in the Sparse Parity(2) task in Section 4.1 when its accuracy is roughly 55% for both train and test?\n- I don't quite understand the claim made in Section 4.1 -- is it that in single-task learning, the model has learned the function itself, rather than learned to use the context? \n- How is \"escaping\" the plateau identified?\n\nI realize some of the details above (e.g., how plateau escape is identified) may be experiment-specific and present in the supplementary -- in this case, I'd suggest including a general description of the method (e.g., identifying a point in time when the loss/accuracy has reached some threshold) in the main paper.", "broader_impact_concerns": "None, although I'd suggest elaborating at least a bit on \"potential societal consequences\" as mentioned in the current statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)