When and How Unlabeled Data Provably Improve In-Context Learning

TL;DR

This paper provides a theoretical analysis of how multilayer transformers can effectively leverage unlabeled data in in-context learning, especially with missing labels, and demonstrates practical improvements in semi-supervised tabular data tasks.

Contribution

It offers a theoretical framework explaining how deep transformers construct estimators from unlabeled data and proposes a method to enhance semi-supervised learning in real-world models.

Findings

Multilayer transformers can implicitly construct polynomial estimators from unlabeled data.

Deeper models exponentially increase the polynomial degree, improving semi-supervised learning.

Applying looping to foundation models enhances semi-supervised tabular data performance.

Abstract

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.