Unlabeled Data Can Provably Enhance In-Context Learning of Transformers
Renpu Liu, Jing Yang
TL;DR
This work tackles the data scarcity challenge in in-context learning by proposing augmented ICL, which prompts transformers with a mix of a small labeled set and a large unlabeled set to infer missing labels without parameter updates. It shows that Chain-of-Thought prompting enables a multi-layer transformer to implement an EM-like algorithm for multi-class linear classification, yielding provable improvements in excess risk and linear convergence under teacher forcing. The paper provides a rigorous analysis of expressiveness and training dynamics, including a gradient-decomposition approach via Stein's lemma and a linear convergence proof for the first-layer weights, and confirms the theory with synthetic experiments where unlabeled data improves both class-mean estimation and label prediction. Overall, the results establish a theoretical and empirical foundation for leveraging unlabeled data to enhance ICL, with potential for broader semi-supervised reasoning in transformers.
Abstract
Large language models (LLMs) exhibit impressive in-context learning (ICL) capabilities, yet the quality of their predictions is fundamentally limited by the few costly labeled demonstrations that can fit into a prompt. Meanwhile, there exist vast and continuously growing amounts of unlabeled data that may be closely related to the ICL task. How to utilize such unlabeled data to provably enhance the performance of ICL thus becomes an emerging fundamental question. In this work, we propose a novel augmented ICL framework, in which the prompt includes a small set of labeled examples alongside a block of unlabeled inputs. We focus on the multi-class linear classification setting and demonstrate that, with chain-of-thought (CoT) prompting, a multi-layer transformer can effectively emulate an expectation-maximization (EM) algorithm. This enables the transformer to implicitly extract useful information from both labeled and unlabeled data, leading to provable improvements in ICL accuracy. Moreover, we show that such a transformer can be trained via teacher forcing, with its parameters converging to the desired solution at a linear rate. Experiments demonstrate that the augmented ICL framework consistently outperforms conventional few-shot ICL, providing empirical support for our theoretical findings. To the best of our knowledge, this is the first theoretical study on the impact of unlabeled data on the ICL performance of transformers.