How does Multi-Task Training Affect Transformer In-Context Capabilities? Investigations with Function Classes (2024)

Abstract: Large language models (LLM) have recently shown the extraordinary ability to perform unseen tasks based on few-shot examples provided as text, also known as in-context learning (ICL). While recent works have attempted to understand the mechanisms driving ICL, few have explored training strategies that incentivize these models to generalize to multiple tasks. Multi-task learning (MTL) for generalist models is a promising direction that offers transfer learning potential, enabling large parameterized models to be trained from simpler, related tasks. In this work, we investigate the combination of MTL with ICL to build models that efficiently learn tasks while being robust to out-of-distribution examples. We propose several effective curriculum learning strategies that allow ICL models to achieve higher data efficiency and more stable convergence. Our experiments reveal that ICL models can effectively learn difficult tasks by training on progressively harder tasks while mixing in prior tasks, denoted as mixed curriculum in this work.

## introduction recently, the emergence of in-context-learning capabilities in llms has revolutionized the field of nlp @xcite . by pre-training with nextword predictions, these models can be prompted with few-shot examples and make accurate incontext predictions during inference @xcite . the icl capability demonstrated even by smaller transformer models presents an alternative way to understand llms @xcite @xcite . to empirically understand this phenomenon, @xcite focus on learning a single function class in-context by a transformer model. their model achieves competitive normalized mean-squared error (mse) compared to the optimal ordinary least squares estimator when performing in-context linear regression. nevertheless, these models sometimes fail to converge and often struggle to generalize to more challenging function classes. while the follow-up studies @xcite @xcite have extensively analyzed how these models conduct icl, little work exists exploring how training on multiple function classes can enable transformer models to generalize and perform icl more efficiently. as we believe that these generalist models are designed to perform multiple tasks, there is a need to study the multitask icl capability of these models, which is missing in the literature. from prior multi-task learning studies @xcite @xcite , models can be trained on multiple related tasks to improve their performance on individual tasks. despite its popularity, mtl has been difficult to understand in transformer models when trained on natural language, most likely due to the difficulty of ranking and scheduling language tasks @xcite . however, the newly introduced framework of learning function classes in context @xcite provides an easier way to study this mtl paradigm. for example, the difficulty of a polynomial function class can be scaled by changing its degree (e.g., linear to quadratic), or changing the input distribution (e.g., gaussian to gaussian with decaying eigenvalues). motivated by this new paradigm, we conduct a systematic analysis by training a transformer on varying function class families and input distributions in a multi-task manner to examine if the same principles from mtl carry over into icl. during training, we explore different curriculum learning strategies to schedule the icl tasks of multiple function classes: sequential, mixed, and random ( §3.3). for benchmarking, we train another set of models only on a single function class family following @xcite . we quantitatively and qualitatively compare our models trained with and without curriculum across all tasks and analyze the normalized mse and attention matrices ( §4). our experiments show that curriculum learning is more data-efficient, achieving comparable performance to single-task models using only 1/9 of the training data. these curriculum models can also obtain an optimal mse in function classes where none of the single-task models converge. ## related work in-context learning in-context learning has garnered increasing attention in the past few years @xcite , and many papers have analyzed icl with natural language @xcite @xcite . it was not until @xcite that the analysis of icl through the paradigm of function class learning emerged. @xcite showed that transformers can learn linear regression close to the optimal ordinary least squares estimator, and other more complex function classes with respectable accuracy. however, they found that some function classes (e.g. gaussian with decaying eigenvalues) were hard to learn by transformers, as the training loss failed to converge. @xcite investigated a framework similar to @xcite , where they explored training models on a mixture of function classes; however, they did not delve into curriculum learning strategies. many papers also explored how icl works, with current literature pointing to it being a fuzzy gradient descent @xcite @xcite . additional theoretical work has examined how transformers can implement near-optimal regression algorithms and has analyzed stability conditions for icl @xcite . @xcite first introduced curriculum learning as a way to train models similar to the way that humans learned, by learning tasks in order from easy to hard. this work inspired a new area of research focused on utilizing curriculum learning in different contexts @xcite @xcite . among this exploration has been more complex curriculum learning strategies in well studied contexts @xcite . the novel function learning problem formulation in @xcite has encouraged us to focus on simple curriculum learning strategies that have been well-studied. our sequential curriculum aligns with the definition provided by @xcite . we have adapted our mixed curriculum from the standard curriculum, which is referred to as a "balanced curriculum" in @xcite . finally, our random curriculum serves as a baseline approach, as described in @xcite . by conducting the first exploratory study on these simple, widely-used curriculum learning strategies, we pave the way for more sophisticated strategies. ## results figure 1 : comparison of the moving average of all three curriculum learning strategies when evaluated on a quadratic function class dataset during test time. the mixed curriculum is the only model that is able to achieve an accurate normalized mse. the random curriculum performs comparatively worse, whereas the sequential curriculum performs substantially worse (yaxis is limited in order for mixed and random curricula to be differentiated). ## discussion in this paper, we examine how different curriculum learning strategies affect a transformer's icl capability. we compare these curriculum models against their respective single-task models and evaluate them across related tasks. this reveals that the mixed curriculum provides the best results, with increased data efficiency and model convergence.