Technical Summary: When and How Unlabeled Data Provably Improve In-Context Learning

Problem Statement

The paper investigates the capabilities of Transformers to perform Semi-Supervised In-Context Learning (SS-ICL). Specifically, it addresses the scenario where a model is prompted with a dataset containing feature-label pairs $(x_i, y_i)$, but a fraction of the labels are missing (set to 0). The goal is to predict the label of a query feature $x$ using only the context provided, without updating model parameters.

The authors focus on a canonical setting where demonstrations are drawn from a binary Gaussian Mixture Model (GMM). In this setting, $x_i \sim \mathcal{N}(\mu y_i, \sigma^2 I)$, where the task mean $\mu$ is sampled from a prior distribution. The central question is: When and how can Transformers learn in context from unlabeled data?

Methodology

The study employs a rigorous theoretical framework analyzing the optimization landscape and expressivity of attention mechanisms under the SS-ICL setting.

1. Data Model: The input prompt consists of $n$ demonstrations where $m$ labels are observed ($y_i \in \{-1, 1\}$) and $n-m$ are missing ($y_i = 0$). The query $x$ and its true label $y$ are generated similarly.
2. Model Architecture: The authors analyze linear attention models, defined by the mechanism $att(Z; W) = (ZW_q W_k^\top Z^\top) M Z W_v$. They study both single-layer and multi-layer (or looped) variants.
3. Theoretical Approach:
   • Single-Layer Analysis: The authors characterize the global minimum of the loss landscape for a single-layer linear attention model trained with squared loss.
   • Multi-Layer Analysis: They investigate how deeper networks (or looped single networks) can construct complex estimators. They derive the form of the estimators as polynomials of the empirical covariance matrix $X^\top X$.
   • Connection to Classical Algorithms: The paper draws connections between the learned attention mechanisms and classical semi-supervised algorithms, specifically Expectation-Maximization (EM) and Belief Propagation.

Key Contributions

1. Limitations of Single-Layer Linear Attention

The paper proves that single-layer linear attention models fail to utilize unlabeled data, even when the data contains information about the feature distribution.

• Theorem 1: Under an isotropic task prior, the optimal single-layer linear attention model recovers the Supervised Plug-in (SPI) estimator: $\hat{\mu}_s = \frac{1}{|I|} \sum_{i \in I} y_i x_i$.
• Implication: The model effectively ignores the unlabeled samples (where $y_i=0$). Theoretical analysis shows that in the meta-learning setting (averaging over random tasks), the sample covariance of unlabeled data converges to a scalar multiple of the identity matrix, providing no task-specific signal to a single-layer model.

2. Depth Enables Semi-Supervised Learning

In contrast to single-layer models, multi-layer (or looped) linear attention can effectively leverage unlabeled data.

• Polynomial Expressivity: The authors show that an $L$-layer linear attention model can implement estimators of the form:
  $$\hat{\mu} = \sum_{i=0}^{K} a_i (X^\top X)^i X^\top y$$
  where $K = O(3^L)$. This allows the model to construct high-degree polynomials of the empirical covariance matrix.
• Logarithmic Depth Sufficiency: Since the degree of expressible polynomials grows exponentially with depth, a logarithmic depth ($O(\log K)$) is sufficient to represent high-degree monomials required for optimal semi-supervised estimation.
• Theorem 3 & 4: As the number of samples $n \to \infty$, multi-layer attention models can approximate the optimal semi-supervised estimator (specifically, a projection onto the top eigenvector of the covariance matrix), achieving classification errors close to the Bayes-optimal limit.

3. Mechanism: EM and Belief Propagation

The paper provides an interpretation of how these models utilize unlabeled data:

• Label Propagation: The iterative update of labels in a multi-layer network mimics the Expectation-Maximization (EM) algorithm. Each layer performs a "pseudo-labeling" step (Expectation) followed by a re-estimation of the task mean (Maximization).
• Softmax vs. Linear: While linear attention implements a linear EM-style update, softmax attention implements a non-linear, similarity-weighted update analogous to Belief Propagation.

4. Application to Tabular Foundation Models

Motivated by the theoretical finding that looping improves performance, the authors propose LoopTabFM, an algorithm that iteratively applies off-the-shelf tabular foundation models (e.g., TabPFN) to generate soft pseudo-labels for unlabeled data.

• Algorithm: The model is run in a loop where unlabeled data is assigned soft labels, added to the prompt, and the model is re-evaluated. A validation risk metric is used to select the best iteration.
• Empirical Validation: Experiments on real-world tabular datasets show that this looping strategy significantly improves test accuracy compared to single-pass inference, particularly in few-shot settings.

Results

Theoretical Results

• Single-Layer: Performance is identical to the supervised-only SPI estimator. Increasing the number of unlabeled samples does not improve accuracy (Figure 1b, 1c).
• Multi-Layer: Performance improves with depth. A 2-layer model significantly outperforms the 1-layer model and approaches the performance of the SSPI-1 estimator (a semi-supervised plug-in using the first power of the covariance). A 5-layer model approaches SSPI-$\infty$ (using the top eigenvector), closely matching the Bayes-optimal error in the asymptotic regime.
• Architecture Comparison: In experiments (Figure 2c), Transformers (MLP + Softmax) outperform linear attention, which in turn outperforms softmax-only attention. However, all architectures converge toward the Bayes-optimal curve as $n$ increases.

Empirical Results

• Synthetic Data: Experiments on GMMs confirm that 2-layer and 5-layer linear attention models achieve higher accuracy than 1-layer models as the number of unlabeled samples increases.
• Real-World Data: Using TabPFN-v2 on OpenML datasets, the LoopTabFM algorithm (iterative pseudo-labeling) achieved relative improvements of up to 19.38% over the baseline (single-pass) on specific datasets (e.g., OpenML ID 1464). On average, the method showed a 6.35% relative improvement. Most gains occurred within the first two iterations.

Significance and Claims

The paper claims to provide the first theoretical characterization of semi-supervised in-context learning in Transformers. Its primary significance lies in:

1. Bridging the Gap: It explains why recent empirical successes in "many-shot" ICL (which often include unlabeled or weakly labeled data) work: depth allows the model to implicitly perform semi-supervised learning algorithms like EM without parameter updates.
2. Depth as a Resource: It establishes that depth is not merely for representational capacity but is crucial for extracting signal from unlabeled data in the context window.
3. Practical Strategy: It proposes a simple, parameter-free strategy (looping) to enhance the semi-supervised capabilities of existing foundation models, demonstrating that "mild amount of depth/looping suffices" to achieve most of the theoretical benefits.

The authors note limitations, acknowledging that their core theoretical results are specific to linear attention and that the optimal semi-supervised estimator for general non-isotropic data remains an open problem. They also observe that while looping helps, performance can degrade on specific real-world datasets due to noise or class imbalance, suggesting a need for data-aware looping algorithms in future work.