When and How Unlabeled Data Provably Improve In-Context Learning
This work addresses the challenge of semi-supervised learning in in-context learning for AI practitioners, offering a theoretical foundation and practical method to enhance performance with unlabeled data, though it is incremental as it builds on existing transformer models.
The paper tackles the problem of in-context learning with missing or incorrect labels in demonstrations, showing that multilayer or looped transformers can effectively leverage unlabeled data by constructing polynomial estimators, with mild depth sufficing for exponential gains, and evaluations on real-world datasets demonstrate significant improvements in semi-supervised tabular learning performance.
Recent research shows that in-context learning (ICL) can be effective even when demonstrations have missing or incorrect labels. To shed light on this capability, we examine a canonical setting where the demonstrations are drawn according to a binary Gaussian mixture model (GMM) and a certain fraction of the demonstrations have missing labels. We provide a comprehensive theoretical study to show that: (1) The loss landscape of one-layer linear attention models recover the optimal fully-supervised estimator but completely fail to exploit unlabeled data; (2) In contrast, multilayer or looped transformers can effectively leverage unlabeled data by implicitly constructing estimators of the form $\sum_{i\ge 0} a_i (X^\top X)^iX^\top y$ with $X$ and $y$ denoting features and partially-observed labels (with missing entries set to zero). We characterize the class of polynomials that can be expressed as a function of depth and draw connections to Expectation Maximization, an iterative pseudo-labeling algorithm commonly used in semi-supervised learning. Importantly, the leading polynomial power is exponential in depth, so mild amount of depth/looping suffices. As an application of theory, we propose looping off-the-shelf tabular foundation models to enhance their semi-supervision capabilities. Extensive evaluations on real-world datasets show that our method significantly improves the semisupervised tabular learning performance over the standard single pass inference.