Research map / Large language models / Language, safety and interpretability
Unlabeled Data Can Provably Enhance In-Context Learning of Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that a 4-layer softmax transformer with CoT can emulate EM on labeled plus unlabeled prompt data for multi-class linear classification, and that teacher-forcing training converges linearly to the EM-implementing solution.
Problem
Whether unlabeled data in the prompt can provably improve in-context learning of transformers when only a few labeled demonstrations are available, given that labeled demonstrations are costly.
Approach
Augmented ICL places labeled and unlabeled examples in one prompt; CoT steps let a transformer run EM-style mean estimation (softmax E-step, linear-attention M-step, ReLU layer for labeled initialization), so the mean-estimation error shrinks with the number of unlabeled samples M.
Key results
- Theory: class mean estimation error bounded by c log(1/eps)/(N + M^{1/4}), so excess risk scales as O(1/sqrt(N+poly(M))), beating the Omega(1/sqrt(N)) lower bound for labeled-only classifiers.
- Teacher-forced training converges linearly: ||W^(k) - Sigma^{-1}||_F^2 <= c^k ||W^(0) - Sigma^{-1}||_F^2.
- Experiments (C=3, d=3, N=5) show augmented ICL with larger M outperforms conventional ICL (M=1) in mean estimation error and accuracy, and surpasses the labeled-only Bayes-optimal classifier.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add unlabeled examples to the prompt alongside labeled examples | Augmented ICL transformer · Multi-class linear classification | class mean estimation error | Helped |
| add unlabeled examples to the prompt alongside labeled examples | Augmented ICL transformer · Multi-class linear classification | prediction accuracy | Helped |
| increase the number of unlabeled examples | Augmented ICL transformer · Multi-class linear classification | class mean estimation error | Helped |
| increase the number of unlabeled examples | Augmented ICL transformer · Multi-class linear classification | classification accuracy | Helped |
| use lower noise ($\epsilon=0.7$ rather than $\epsilon=1.5$) | Augmented ICL transformer · Multi-class linear classification | advantage over conventional ICL | Helped |
| use augmented ICL with unlabeled examples | Transformer trained via teacher forcing · Multi-class linear classification | prediction error | Helped |
Surprising findings
- Augmented ICL advantage is more pronounced when the noise level epsilon is smaller.
Limitations
- Restricted to multi-class linear classification with isotropic Gaussian noise and a fixed known covariance structure
- Theory relies on a 4-layer hand-constructed transformer and approximations (beta to infinity, arbitrary MLP link functions)
- Training analysis freezes the last three layers and trains only the first attention layer
Experimental setup
Datasets: synthetic Gaussian mixture ICL instances
Benchmarks: synthetic Gaussian mixture ICL instances (C=3, d=3)
Metrics: class mean estimation error label prediction accuracy
Baselines: conventional few-shot ICL (M=1) Bayes-optimal classifier using labeled data only
Builds on: chain-of-thought prompting teacher forcing EM algorithm for Gaussian mixtures
Closest papers by meaning
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.92)
- How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR 2024 · similarity 0.91)
- Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens (NeurIPS 2024 · similarity 0.90)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.90)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.90)
- Exact Conversion of In-Context Learning to Model Weights in Linearized-Attention Transformers (ICML 2024 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.