The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces
Provides a theoretical framework to unify existing empirical observations about the role of task diversity in ICL, relevant for researchers studying transformer mechanisms.
The paper develops a minimal analytical model to explain how training task diversity, defined by non-overlapping subspaces, improves in-context learning (ICL) in transformers. The model shows that diversity shortens the ICL plateau and enables out-of-distribution generalization, with empirical validation on nonlinear transformers.
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.