AISTATS’26 CauScale Workshop Best Paper reveals generalization mechanism of LLMs

July 1, 2026
New research led by Laura Balzano and Qing Qu establishes a mathematical framework explaining how LLMs generalize to novel, out-of-distribution tasks.

Recent studies have shown that large language models (LLMs) such as ChatGPT and Gemini exhibit remarkable zero-shot generalization, allowing them to perform novel tasks via prompting alone, without task-specific retraining. Specifically, they can recognize and apply a pattern in a series of consistently-formatted prompts, an ability called in-context learning (ICL).

A comic strip illustrates a robot learning examples of English to Spanish translation and using them to translate a new phrase.
After exposure to a few English–Spanish pairs as in-context examples, the LLM successfully generalizes the underlying mapping to translate a new query like “good morning.” Image courtesy of Qing Qu

Despite the remarkable empirical success of in-context learning, a fundamental question remains unanswered: When can large language models reliably generalize to out-of-distribution (OOD) tasks? Lead authors Soo Min Kwon and Alec S. Xu, both Ph.D. students in Electrical and Computer Engineering (ECE), addressed this challenge by developing a rigorous mathematical framework that characterizes the OOD generalization capabilities of in-context learning. Their work earned the Best Paper Award at the 2026 AISTATS Causality in the Age of AI Scaling Workshop.

Portrait of Soo Min Kwon
Soo Min Kwon
Portrait of Alec Xu
Alec Xu

The research team developed a mathematical framework based on low-dimensional subspaces whose simplified geometric structures capture intrinsic patterns in the data, thereby explaining when LLMs can and cannot generalize to OOD tasks. Their analysis revealed that when the training data are organized as a union of connected or overlapping subspaces, the model can exploit the latent relationships among these subspaces to generalize across their convex hull, extending well beyond the examples provided in the prompt. In contrast, when the training data are confined to a single isolated subspace, the trained model’s reasoning remains restricted to that subspace and fails to generalize to novel tasks. These findings provide one of the first rigorous theoretical explanations for the remarkable in-context learning capabilities of modern LLMs.

“A central finding of our work is that the seemingly OOD generalization of in-context learning is often not truly out of distribution: the evaluation distribution implicitly resides within the span of the pre-training data,” explained Kwon.

This work highlights the critical role of task diversity in pre-training LLMs, which is mathematically characterized by a union of low-dimensional subspaces. Building on this insight, the team’s subsequent work further demonstrated that task diversity in the pre-training data is also essential for accelerating LLM training, leading to faster convergence and more efficient learning. Together, these results highlighted the broad theoretical and practical significance of the proposed framework, showing that task diversity is a fundamental principle governing both the generalization and optimization of modern LLMs.

Three graphs illustrate two training phases for transformer models and their subsequent inference abilities.
When task vectors are drawn from a union of subspaces, a transformer via ICL can generalize to subspaces with zero density in the training data, whereas a model trained on a single subspace cannot. Image: Soo Min Kwon

“I have always found it fascinating that by studying simplified models, we can reveal the fundamental principles behind the behavior of today’s large-scale AI systems,” said ECE Professor Qing Qu, who co-advised the work. “These models give us a concrete, mathematically grounded explanation for phenomena that would otherwise be difficult to understand directly.”

The award-winning paper, “Towards Understanding Out-of-Distribution Generalization for In-Context Learning via Low-Dimensional Subspaces,” was also co-authored with ECE Ph.D. student Can Yaras and ECE Professor Laura Balzano. This work laid the theoretical foundation for the team’s subsequent paper, “The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces,” which further establishes task diversity as a fundamental principle governing both the generalization and training efficiency of large language models.