The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces

• URL: http://arxiv.org/abs/2606.06814v1
• Date: Fri, 05 Jun 2026 01:35:42 GMT
• Title: The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces
• Abstract summary: In-context learning (ICL) has sparked a wide range of studies designed to understand its underlying mechanisms.<n>Existing works often study how training task diversity shapes both the learning dynamics and generalization capabilities of ICL.<n>This paper presents a minimal analytical model under which these phenomena provably emerge from the properties of the training data.
• Score: 20.732738714750095
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: The transformer's emergent ability to perform in-context learning (ICL) has sparked a wide range of studies designed to understand its underlying mechanisms. Existing works often study how training task diversity, defined either as the number of ICL training task vectors or as the number of function classes from which the task vectors are drawn, shapes both the learning dynamics and generalization capabilities of ICL. While both definitions have uncovered many interesting phenomena, many observations under the latter definition remain theoretically unexplained. This paper presents a minimal analytical model under which these phenomena provably emerge from the properties of the training data. By modeling the training task vectors as a mixture of low-rank Gaussians, we show how training task diversity, defined by the number of non-overlapping columns between subspaces that parameterize the covariance matrices, improves both the generalization and optimization trajectory of ICL with linear attention. In particular, we show that our model can explain (i) why training with task diversity shortens the ICL plateau and (ii) why ICL appears to achieve out-of-distribution generalization. We conclude by empirically demonstrating how our results extend to nonlinear transformers and nonlinear function classes. Overall, our work presents a tractable framework to unify existing observations.

Related papers

• Provable In-Context Vector Arithmetic via Retrieving Task Concepts [53.685764040547625]
  We show how nonlinear residual transformers trained via gradient descent on cross-entropy loss perform factual-recall ICL tasks via vector arithmetic.<n>These results elucidate the advantages of transformers over static embedding predecessors.
  arXiv  Detail & Related papers  (2025-08-13T13:54:44Z)
• Provable In-Context Learning of Nonlinear Regression with Transformers [66.99048542127768]
  In-context learning (ICL) is the ability to perform unseen tasks using task specific prompts without updating parameters.<n>Recent research has actively explored the training dynamics behind ICL, with much of the focus on relatively simple tasks.<n>This paper investigates more complex nonlinear regression tasks, aiming to uncover how transformers acquire in-context learning capabilities.
  arXiv  Detail & Related papers  (2025-07-28T00:09:28Z)
• Understanding In-Context Learning on Structured Manifolds: Bridging Attention to Kernel Methods [45.94152084965753]
  We establish a novel connection between the attention mechanism and classical kernel methods.<n>We derive generalization error bounds in terms of the prompt length and the number of training tasks.<n>Our result characterizes how the generalization error scales with the number of training tasks.
  arXiv  Detail & Related papers  (2025-06-12T17:56:26Z)
• Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspective [9.249642973141107]
  We demystify the out-of-distribution capabilities of in-context learning (ICL) by studying linear regression tasks parameterized with low-rank covariance matrices.<n>We prove that a single-layer linear attention model incurs a test risk with a non-negligible dependence on the angle, illustrating that ICL is not robust to such distribution shifts.<n>This suggests that the OOD generalization ability of Transformers may actually stem from the new task lying within the span of those encountered during training.
  arXiv  Detail & Related papers  (2025-05-20T18:15:49Z)
• Illusion or Algorithm? Investigating Memorization, Emergence, and Symbolic Processing in In-Context Learning [50.53703102032562]
  Large-scale Transformer language models (LMs) trained solely on next-token prediction with web-scale data can solve a wide range of tasks.<n>The mechanism behind this capability, known as in-context learning (ICL), remains both controversial and poorly understood.
  arXiv  Detail & Related papers  (2025-05-16T08:50:42Z)
• Re-examining learning linear functions in context [4.126494564662494]
  In-context learning (ICL) has emerged as a powerful paradigm for easily adapting Large Language Models (LLMs) to various tasks.<n>We explore a simple model of ICL in a controlled setup with synthetic training data.<n>Our findings challenge the prevailing narrative that transformers adopt algorithmic approaches to learn a linear function in-context.
  arXiv  Detail & Related papers  (2024-11-18T10:58:46Z)
• Transformers are Minimax Optimal Nonparametric In-Context Learners [36.291980654891496]
  In-context learning of large language models has proven to be a surprisingly effective method of learning a new task from only a few demonstrative examples. We develop approximation and generalization error bounds for a transformer composed of a deep neural network and one linear attention layer. We show that sufficiently trained transformers can achieve -- and even improve upon -- the minimax optimal estimation risk in context.
  arXiv  Detail & Related papers  (2024-08-22T08:02:10Z)
• Asymptotic theory of in-context learning by linear attention [37.3817914656799]
  In-context learning is a cornerstone of Transformers' success.<n>Questions about the necessary sample complexity, pretraining task diversity, and context length for successful ICL remain unresolved.
  arXiv  Detail & Related papers  (2024-05-20T03:24:24Z)
• How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? [82.51626700527837]
  Transformer-based large language models displayed impressive in-context learning capabilities, where a pre-trained model can handle new tasks without fine-tuning. We analyze how the mechanics of how Transformer to achieve ICL contribute to the technical challenges of the training problems in Transformers.
  arXiv  Detail & Related papers  (2024-02-23T21:07:20Z)
• How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations [98.7450564309923]
  This paper takes initial steps on understanding in-context learning (ICL) in more complex scenarios, by studying learning with representations. We construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size.
  arXiv  Detail & Related papers  (2023-10-16T17:40:49Z)
• Low-Rank Multitask Learning based on Tensorized SVMs and LSSVMs [65.42104819071444]
  Multitask learning (MTL) leverages task-relatedness to enhance performance. We employ high-order tensors, with each mode corresponding to a task index, to naturally represent tasks referenced by multiple indices. We propose a general framework of low-rank MTL methods with tensorized support vector machines (SVMs) and least square support vector machines (LSSVMs)
  arXiv  Detail & Related papers  (2023-08-30T14:28:26Z)
• Adaptive Discrete Communication Bottlenecks with Dynamic Vector Quantization [76.68866368409216]
  We propose learning to dynamically select discretization tightness conditioned on inputs. We show that dynamically varying tightness in communication bottlenecks can improve model performance on visual reasoning and reinforcement learning tasks.
  arXiv  Detail & Related papers  (2022-02-02T23:54:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.