How training task diversity boosts in-context learning in transformers

A new paper by Soo Min Kwon, Alec S. Xu, Can Yaras, and colleagues from the University of Michigan and University of Illinois Urbana-Champaign presents a minimal analytical model to explain how training task diversity shapes in-context learning (ICL) in transformers. The authors model training task vectors as a mixture of low-rank Gaussians and define diversity by the number of non-overlapping columns between the subspaces that parameterize the covariance matrices. Using linear attention as a tractable proxy, they prove that higher diversity (more distinct subspaces) provably improves both generalization and optimization dynamics.

The model uncovers two key phenomena. First, diverse training tasks shorten the ICL plateau—the period where the model shows minimal improvement during training—by enabling faster alignment of attention weights. Second, it explains why ICL achieves out-of-distribution generalization: the subspace structure forces the model to learn compositional features rather than memorizing task-specific patterns. The team then empirically validates these results on nonlinear transformers and nonlinear function classes, showing the theory generalizes beyond linear attention. This work provides a rigorous, unified framework for many previously unexplained observations about ICL and offers practical guidance for designing training data to build more capable few-shot learners.