Unlabeled Data Can Provably Enhance In-Context Learning of Transformers
Authors: Renpu Liu
Jing Yang
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Models
In‑Context Learning
Chain‑of‑Thought Reasoning
Primary area: theory

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.

Unlabeled Data Can Provably Enhance In-Context Learning of Transformers
Renpu Liu University of Virginia Charlottesville, VA 22903 renpu@virginia.c edu
Jing Yang University of Virginia Charlottesville, VA 22903 yangjingdvirginia. edu
Abstract
Large language models (LLMs) exhibit impressive in-context learning (ICL) ca- pabilities, yet the quality of their predictions is fundamentally limited by the few costly labeled demonstrations that can fit into a prompt. Meanwhile, there ex- 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 novel augmented ICL framework, in which the prompt in- cludes small set of labeled examples alongside block of unlabeled inputs. We focus on the multi-class linear classification setting and demonstrate that, with chain-of- of-thougttt (CoT) prompting, multi-layer transformer can effectively em- ulate expectation expectation-maxiiizatie (EM) algorithm. This enables the transformer to implicitly extract useful information from both labeled and unlabeled data, lead- ing 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.
Introduction
Since the introduction (Vaswani al.) 2017). transformers have become foundational models in diverse fields such as natural language processing (Radford 2018 Devlin et al.) 2019). computer
vision (Dosovitskiy) 2020), and reinforcement learning (Chen et al. 2021). A key driver of their impact is the remarkable capability for In-Context Learning (ICL) (Brown et al.) 2020). Without
requiring parameter updates, transformers performing ICL can adapt new tasks based solely on contextual examples provided within the prompt. This enables state-of-the-art few-shot performance multitude of applications, including reasoning and language understanding (Chowdhery et al. 2023), dialog generation (Thoppilan et al. 2022). and linear regression (Garg et al. 2022 Fu al. 2023). etc.
Despite the power of ICL, its reliance on labeled examples presents a significant bottleneck for large language models (LLMs).. Acquiring high-quality labeled data is in general expensive and time- consuming (Zhou et al., 2023 Chung et al.) [024] Sun et al. 2023 Wang et al.) (202).. For example, creating the instruction-tuning and RLHI datasets for models like GPT-3.5 and GPT-4 involved thousands of expert annotator hours, yet constituted less than 0.1% 1% of the tokens encountered during pre-training (Ouyang al. 2022 Achiam al. 2023).
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Some existing approaches attempt to mitigate labeled data scarcity in ICL. For instance, Wan et al. (2023); Chen et al. (2025a) use an LLM to automatically generate pseudo-demonstratioon at infer- ence time by pairing unlabeled queries with the model's own predictions as pseudo labels. However, model-generated sseudo-labels inevitably inherit the biases and error patterns of the teacher model, resulting in noisy demonstrations that may limit potential performance gains.
In this work, instead of synthesizing examples with pseudo-labels, we explore a different approach by directly utilizing abundant and continuously growing (Raffel et al. 22200 Touvron et al. 2023 unlabeled data during ICL. The fundamental question we aim to answer is:
Can we provably enhance the ICL performance of transformers by effectively leveraging plentiful unlabeled data alongside limited labeled examples?
We answer this question affirmatively from a new augmented in-context learning perspective. This paradigm involves prompting transformer with a mixture of few labeled examples and numerous unlabeled examples, aiming to infer the missing labels within a single forward pass. By reasoning over unlabeled examples directly in the prompt, it bypasses the need for potentially costly, time- consuming, and bias-introducing labeling or pscudo-label generation steps in conventional ICL. In this work, we focus on augmented ICL for multi-class linear classification. Our main contributions are as follows.
Expressiveness with CoT Prompting. First, we show that through Chain-of- ofoohhh (CoT) prompting, a multi-layer transformer can leverage both labeled and unlabeled data to effectively solve the multi-class linear classification problem during ICL, Essentially, the transformer is able to obtain an initial estimation of the mean vectors of classes using the labeled data, and then it- eratively refine the estimates by clustering the unlabeled data in an -epectatioonMaxiizzaatt (EM) fashion. We explicitly characterize the design of the transformer and theoretically prove that the class mean estimation will converge to the ground truth as the CoT steps increase. For a prompt consisting of N labeled and M unlabeled samples, the excess risk of our approach scales in O (1/NN- poly((M)) strictly improving the excess risk lower bound of O (I/N) for any classifier that utilizes N labeled samples only. Our results indicate that the augmented ICL can effectively utilize the information from the unlabeled data, enabling steady performance improve- ment as unlabeled data increases.
Training Convergence under Teacher Forcing, Second, we prove that, with proper initializa- tion, when applying gradient descent on the population loss defined through teacher forcing, the tunable parameters of the transformer converge to the desired solution linearly. Thus, the trained transformer can mimic the EM algorithm through CoT prompting during inference, theoretically demonstrating that the expressive solution for augmented ICL. is identifiable and learnable. Our proof involves a novel decomposition of the gradient of the CoT training loss into two analytically tractable terms. For each of them, we leverage the inherent isotropy of the involved quantities to simplify the analysis, which enables us to derive a tight upper bound on the critical inner-product term and obtain the linear convergence rate.
Empirical Results. Finally. we evaluate the performance of augmented ICL in transformers trained via teacher forcing. Our experimental results show that the augmented ICL approach sig- nificantly outperforms conventional ICI. in both class mean estimation and label prediction, with the advantage becoming more pronounced as the number of unlabeled data samples increases. Moreover, augmented ICL surpasses the Bayes-optimal classifier that relies solely on labeled data. These empirical observations are consistent with our theoretical findings.
2 Related Works
ICL with Transformers. Brown et al. (2020) first shows that GPT-3, a transformer-based LLM, can perform new tasks from onput-output pairs without parameter updates, suggesting its ICL ability. This intriguing phenomenon of transformers has attracted much attention, leading to various inter- pretations and hypotheses about its underlying mechanism. Research on ICL often demonstrates how transformers can emulate learning algorithms. For instance, several studies have designed transformers that execute gradient descent for linear and non-linear regression tasks (Akyürek et al.)
2023 Von Oswald et al. 2023a). Recent works demonstrate that transformers can implement more advanced optimization al algorithms other than vanilla gradient descent on various ICL. tasks (Bai et al.)

---

2024
2024a Ahn et al. 2024 Liu et al. (025). Another line of research adopts a statistical perspective: ICL can be viewed as an implicit form of Bayesian up- dating based on the examples provided in the prompt, with the diversity of pretraining data shaping the prior (Xie et al. 2022
Von Oswald et al. [203bb] Zhang et al.
Raventós et al.
2023. Garg et al. 2022).
Several studies (Gupta et al., 2024; Agarwal et al., 2024) investigate "unsupervised ICL", in which the prompt consists solely of unlabeled inputs. Another line of work leverages LLMs to generate pseudo-labels for unlabeled data, which are then used as demonstrations during ICL (Chen et al.
2023 Wan al. 2023; Yang et al. 2023] Chen al., 2025a). Our work leverages both labeled and unlabeled examples within the prompt to enhance ICL performance in a semi-superviied learning manner, which stands in sharp contrast the aforementioned studies.
Notably, a recent concurrent work (Li al. 2025) also investigates the impact of the semi-supervised data model the ICI. performance of transformers. Specifically, [i et al. (2025) focus a lin- ear transformer without nonlinear activations in a binary classification setting, and characterize the asymptotic ICL performance as the number of unlabeled samples approaches infinity. In contrast, we study a more realistic architecture that incorporates the softmax attention mechanism and establish a non-asymptotic convergence guarantee in the general multi-class setting.
Training Dynamics of Transformers. A number of recent works aim to provide theoretical char- acterizations of the training dynamics of transformers. [Ahn et al. (2024); Mahankali et al. (2023); Zhang et al. 4 Huang et al. (2023) investigate the training dynamics of transformers with single attention layer and a single head for in-context linear regression tasks. CCu et al. (2024) prove that transformers with multi-hear attention layers outperform those with single- head atten- tion. Cheng et al. (2024) show that local optimal solutions in transformers can perform gradient descent in-context for non-linear functions. Kim and Suzuki (2024) study the non-convex mean- field dynamics of transformers, and Nichani et al. (2024) characterize the convergence rate for the training loss in learning causal graph. Additionally, (Chen et al. (0244) investigate the gradient flow in training multi-head single-layer transformers for multi-task linear regression. [Chen and Li (2025) propose supervised training algorithm for multi-head transformers. The training dynamics of transformers for binary classification (Tarzanagh et al. 2023b a Vasudeva et al. 2024 Li et al.)
2023 Deora et al.] 2023 Li et al. 2024a), multi-class classification (Shen et al.) 2025) and next-
token prediction (Tian et al. 2023 2024 Li et al. 2024b Huang et al., 2024) have also been studied
recently.
Transformers with CoT. In language modeling tasks, transformers have been proven to be pow- erful across various downstream tasks. However, transformers struggle to solve mathematical or sci- entific problems with a single generation, particularly when multiple reasoning steps are required. generate intermediate results autoregres-
CoT prompting is introduced to enable transformers sively before reaching the final answer, and has been shown to boost performance on arithmetic, commonsense, and scientific tasks (Wei et al. 2022 Kojima et al.. 2022).
Recently, the training dynamics of transformers with CoT have been studied in Huang et al. (2025a) for weight prediction in linear regression, in Li et al.] (2024a) for in-context supervised learning, in Kim and Suzuki (2055); Wen et al. (2025) for the parity problems, and in Huang et al. (2025) for the even pairs problem. None of these studies, however, address whether the multi-step reasoning capacity through CoT can be utilized to extract information from unlabeled inputs.
Preliminaries
Notations. For matrix X, we use [X] (x)p.gr to denote the submatrix that contains rows p to q and columns r to and we use [X]., and [ j.: to denote the i-th column and j-th row of X, respectively. For convenience, we occasionally denote the i-th column X by [X]; when no ambiguity arises. [X).-c:-1 means the last C' columns of matrix X. We use || | to denote its Frobenius norm. For vector x, we use ||x||], ||||| and ||x |x to denote its l,, l2 and f.x. norms, respectively. We denote by 1 and O2 the d-dimensional all-1 and all-O column vectors, respectively. Laxb and Oaxb denote the all-1 and all-O matrices of size a x respectively. We denote the indicator function as 1 [(A),
which equals 1 if event A is true.
3

---

3.1 Transformer Architecture
In this work, we consider the encoder-based transformer architecture /Vaswani al. 2017). where each transformer layer consists of an attention layer followed by multi-layer perception (MLP) layer.
Definition 3.1 (Attention layer). Denote an M-head attention layer parameterized by
{(V m, Qm: km) mmmmmna as alin m. Q..ii......K..... k.))()),, where Vm, nn, Km E RDXD, Vm E [M]. Then, given an input sequence H E IDI(N+1), the output sequence of the attention layer is
(Vt m. Qm. Km)) (H) - H Σ
H) x ((KmH) ((mmH))(Qmm
Wim
where a is non-linear activation function.
Definition 3.2 (MLP layer). Given W D'D, W2 DxD and bias vector b E R MLP layer following the attention layer, denoted as MLP Wabb) maps each token in the input sequence (i.e, each column h; in H E DxN to another token as
MLP (W,. W2.b) (h.) = h, + W20 W +b),
where a is non-linear activation function.
3.2 Augmented In-context Learning
Conventional In-Context Learning (ICL). For an ICL task, a trained transformer given an ICL. instance I = (D+x,+1), where DD = {(xj, M5)])E[[)) and Xw+1 is a query. Here, Xj E IR is an in-context example, and yj the corresponding label for x,. For each instance, {(x). Mj) / N+1 generated independently accordingly predict UN+1 without any parameter updating of the transformer.
j=1
underlying distribution. The objective of ICL is to
Augmented ICL. In this work, we consider new unlabeled data augmented ICL framework.
y==1" and a set of unlabeled examples, Dunlabel (x, J=N++1 i.c., I Diabel U Dunlabel- Similar to conventional ICL, all (x,, Uj) pairs follow the same distribution. The objective of augmented ICL is then to predict labels for all the M unlabeled samples in Dunlabel.
Specifically, each ICL instance now comprises a set of labeled examples, Diabel := {(xx, (3)) 3y))
N-+M
We note that the augmented ICL generalizes the conventional ICL framework, and reduces to it when M = 1. While the conventional ICL can utilized to solve the prediction for those M unlabeled samples individually in parallel. by augmenting them in the same ICL instance, it provides an opportumity for the transformer to extract common statistical information in those unlabeled data, which can be utilized to improve the joint prediction accuracy.
Augmented ICL for Multi-class Linear Classification. We consider augmented ICL for a multi- class linear classification problem. We assume there exist C' classes, and the label space ) consists of one-hot vectors
where each e; E RCC is the i-th unit vector. For each ICL instance IM. the samples are randomly generated according
M PM: Yj Uniform())), Eg N0 E), Xj My,+ j J (3.1)
where M R and PM is prior distribution over dxC Denote the columns of M as {ui)e Then, each Xy essentially follows a C-component mixture Gaussian distribution parametrized by mean vectors [uIC and shared covariance matrix Σ. In this work, we assume Σ is isotropic. We adopt this assumption for theoretical tractability, it is crucial for deriving the dlseed-frmm update rules for the transformer. This approach is a standard and widely adopted practice related literature to facilitate theoretical analysis (He al. 2025 Zhang al. 2024bb Chen et al. 2025b).
3.3
Chain-of-Though Prompting for Augmented ICL.
The core challenge in augmented ICL is leveraging both unlabeled data and labeled examples to infer task structure from a single instance. Unlike standard few-shot ICL, which often uses direct pattern matching, the augmented ICL. requires more complex inference to effectively utilize the larger unlabeled set, making simple one-step prediction insufficient.

---

Chain-of- Thought (CoT) reasoning offers promising way to enhance transformer's ICL capabil- ities. This is crucial for augmented ICL, as it enables the transformer to effectively utilize unlabeled data through iterative latent parameter estimation and refinement.
To implement augmented in-context learning via CoT prompting, we first encode task instance I into an embedding matrix H by concatenating three column blocks: the labeled example block, the unlabeled example block, and the reasoning block as follows:
X1
XN XN+1 YN 0
XN+M 0
X
X
a
0
H
0
...
A
Y
0
(3.2)
P
PN
PN+1
PN+M
q1
9c
P
P
Q (0)
where Py IR is auxiliary embedding that stores the (predicted) classification probability vector for the j-th sample, as well as binary indicator to distinguish the labeled and unlabeled data.
qi Rd serves the initial CoT token for class which contains the one-hot vector e; indicate the corresponding class, and an all-zero vector representing the transformer's initial estimate for the mean vector p4z.
Denote a trained transformer with parameter as TF With CoT, we will use the transformer to generate intermediate steps before it outputs the prediction. Specifically, let H be the input sequence at the t-th step of CoT, where ()) H, and TF (H) as the corresponding output of the transformer. Then, will take out the last C columns of TF (H(+)))), and append them to the end of (t-1) form the imput for the next CoT step. Specifically, we have
X Ye
X
H 1), (F (H (1- -:-1
0
(3.3)
P
P
Q (0)
Q
(1)
(+)
where
e
ec
Q
(t)
Pe
(1(
(3.4)
Here is placeholder for dummy tokens, is the i-th unit vector, and (() is the estimated mean vector for class at the 1-th CoT step.
After iterations, we read out p (T) P from Q(T) as the final estimation of the class mean vectors. Then, the label of each unlabeled data be estimated through maximum likelihood estimation, i.e.,
8, ee:i arg min eE(C] |x,- - (x--m(11) 11 je [N
(3.5)
Expressiveness with CoT Prompting for Augmented ICL
In this section, show that multi-layer transformer implement Expectation-Maxiiiizee (EM)-style algorithm to extract useful statistical information from the unlabeled data, which will be combined with information extracted from the labeled data to jointly estimate the class means and improve the augmented ICL performance. Specifically, we have the following result.
Theorem 4.1. There exists 4-layer transformer, such that its output sequence the (t+ 1)-th CoT step satisfies
(1+1) = m P
17
N+M M
Σ
(0) Pij
x1) + 1 (M
C N
Σ
(eTyy))
(4.1)
M
j=N+1
for any E [C], where (1) a (T' +t) some positive constants a and T", p.j is the normalized weight
E exp
- -x,11 - BT
(4.2)
Σt E exp - x, - BT
and B is positive constant.
5

---

We outline the construction of each layer of the transformer below, and defer the detailed derivation and specific parameter implementation Appendix
The four-layer architecture is designed to mirror an EM iteration for Gaussian mixture model clus- tering (Zhao et al., 2020; Sula and Zheng 2022) within the transformer's forward pass. The EM algorithm operates iteratively. First, in the E-step, it utilizes the current class mean estimates em- bedded in the input sequence to compute the estimated class membership for each unlabeled data point. Subsequently, the M-step updates the class mean estimates by performing a maximum like- lihood estimation of the unlabeled data, and then combining them with the estimates obtained from the labeled data. Through this iterative process, the algorithm converges to accurate estimates of the underlying class means, enabling reliable classification.
The first layer. The first transformer layer includes a softmax- attention layer followed by MLP layer. We construct its parameters that it outputs the class membership esti- mate for the each unlabeled sample as in the form of Equation (4.2) where the mean estimates (p(")... .... p the parameter 3 embedded the first layer as well. This probability represents how likely sample j is estimated be in class i. Since the temperature parameter /rr is proportional the step index r, estimates from earlier CoT steps carry less importance. In the limit of B - oc, the weight vector depends only the latest CoT step, i.e.,
7-=1 embedded in the reasoning blocks Q()) ((t) in the input sequence, and
exp - (1))(()
2 p
-xyl
P
(4.3)
EC exp
- x5
The second and third layers. The second and third transformer layers consist of a linear attention layer followed by an MLP layer. These layers are designated for the M-step of the EM algorithm. In this step, the class mean estimates { are updated by maximizing the overall likeliklddooo of the unlabeled data with the estimated class membership probabilities Pij It aims to solve
N+M Σ
P1 {
(1(++)))
arg max (Me)e
Σ
(0 Pig
log N(2;; p. E).
The implementation for these two layers is equivalent to tasking one step gradient descent over i...
N
M
(1+1)
Po ( -x²)
(4.4)
M
J=
The fourth layer. Finally, the last transformer layer includes a ReLU-activated attention layer followed by an MLP layer. This layer calculates the initial class mean estimates for the labeled dataset and is only activated at the first CoT step. It implements the following updating rule:
C
N
p (++1) = 1 (1=0)
Σ y,)x, N
(4.5)
which initialize ji! to be the average of xys for the labeled data samples in class i. This initialization will be refined iteratively through the CoT steps by leveraging the information from the unlabeled data.
We note that the parameters of the last three layers are data-independent and can be explicitly con- structed beforehand, and only the parameters of the first layer depend the distribution of the data, which can be obtained through CoT training, as elaborated in Section 5
Next, we will show that the transformer specified in Theorem 4.1 will recover {u)C accurately with high probability, and explicitly characterize the benefit of unlabeled data in this augmented ICL.
Theorem 4.2 (Class Mean Estimation Error). Given the transformer described in Theorem 4.1 when N
36a L? log 1/e and M 2 max{36a² 2 (1/e() and 2 max VM, MMT' T'}, with

---

probability at least - e, the output of the transformer after CoT steps satisfies
log(1/c)
M -MI
N+ N+
where 01 C, a, L, K are positive constants.
Proof sketch. The proof of Theorem 4.2 contains three major steps. In Step 1, we utilize the Ho- effding's inequality ensure that with a sufficient number of labeled data N, the initial class mean (1) estimates (( ,... Mc are in a small neighborhood of the ground truth class means pa; ... M4C-
In Step 2, need to bound the gap between the gradient descent updating step for > 1 in
Equation (4.4). and one gradient descent step for the expected log-likelihood loss C( ((()()) - -3
}) =
Ex log Σ exp
To ensure that the gap is sufficiently small, we need
design the temperature parameter [3r that the normalized weight is biased heavily toward the class mean estimation obtained from the current CoT step, and the influence of previous CoT steps minimized. Then, utilizing Bernstein's 's inequality, this gap is bounded. In Step 3, we utilize Lipschitz continuity of L( L((u))) combing the bound on the gradient gap in Step 2, to show that [M M| < O(1/ /N + po)y(M)) for t large enough if M is in small neighborhood of M, which is guaranteed in Step 1. The complete proof can be found in Appendix
Based the smoothness of the Bayes risk, have the following corollary direct consequence of Theorem 4.2
Corollary 4.1 (Label Prediction Error Bound). Let V; be the predicted label for x; according to
Equation 3.5). .. Let R be the prediction error under the Bayes-optimal classifier with known class
mean vectors p4,"" PC.. Then, under the same conditions as described in Theorem 4.2
we have
(< (poloyyMT)
Ply, + "yl.....,,a -cc
Remark 1. The advantage of utilizing unlabeled data in the augmented ICL becomes evident when comparing Corollary 4.1 with the existing lower bound on the excess risk for classical binary clas- sification. It has been shown that the excess risk for any classifier trained on N labeled data scales in (()//) in the worst case of M (Li et al. 2017), which is in stark contrast to the upper bound (W/NN O poly(M)) in Corollary 4.1 This result indicates that the designed transformer can effectively utilize the unlabeled data through CoT prompting, and strictly improves the prediction accuracy of any classifier that utilizes the labeled data only.
Training Dynamics with Teacher Forcing
While Section 4 indicates that there exists transformer that is able to implement EM-type algorithm to utilize unlabeled data and improve the ICL performance through CoT prompting. in this section, show that such transformer can be obtained through teacher forcing training (Kim and Suzuki 2025 Huang et al. 2025b).
The training objective of teacher forcing is to ensure that the transformer can mimic the trajec- tory of iterative updating under an EM algorithm during the CoT inference. Formally, we re- quire that, the unlabeled set, the cross-entropy between the class distributions induced by the CoT estimates (p(()
) j and those induced by the reference method fref remains small for all I. t= 1....,, T. Specifically, given X,, Y, and X,,, we denote the generated reference trajectory as [X,, Y,, Xx) (1 1 Then, we construct the reference embedding sequence the 1-th CoT step
X X,
eC Vr E [t].
Y:
0
0
Q
ref
P
P
(t)
ref,11
Pef,C C
ref
We note that the reference embedding shares the same structure as the embedding defined in Equa- except that now the mean estimates are generated by the reference algorithm instead of
tion 3.3
7

---

the transformer itself. We then feed H (9) from its output TF (H))) rot
to the transformer, and extract the updated mean estimates
ref
(H
The corresponding CoT training loss can be defined as:
N- M
ÉCoT-train (0IIM)
1
ICE
( [TFF (H (t-1))
(5.1)
T
ref
+=
where lce is the cross-entropy loss and
qj
(t)
[pi]
(t)
PPC)) with
(1)
exp
1 2
Illl M ref,i
(t)
X, -
(t)
Pij
exp
2 |ef,e X3 xy112-1)
(1)
Similar to Ahn et al. (2024); Huang et al. (2025a), in this work, we analyze the training convergence of the population loss defined as:
CCoT-mmin + BTsa
CCoT (θ; (m;)
(5.2)
where the expectation taken over the randommess in the generation process of IM-
Directly analyzing the training dynamics of all layers of the transformer is intractable. On the other hand, as we mentioned in Section 4 the last three layers of the transformer can be constructed explicitly beforehand, as their parameters are data-independent. As a result, in the following, we will freeze these three layers and train the first layer only.
Assumption (Initializattonn. We initialize the first layer of the three-layer transformer described in Theorem 4.7
as follows:
0dx(d+20)
W(0)
(0) K
0 0(4C+4+2)x2C
1
O 01x2
g(0)
( diag (iiaa(((a+))) (00022))x(()C))) Ics 0(44C44)((422444)
where ( is a d x matrix whose entries are randomly sampled from standard Gaussian distribution, g(0) is a constant, and all the unspecified entries are equal to zero.
Theorem 5.1 (Training Convergence). Let (1(1), (k(), V be the parameters of the first attention layer of the transformer after applying k iterations of gradient descent on the population loss defined in Equation (5.2) Then, with the initialization specified in Assumption 7 we huve
[w(k) - 1ー S ( -0-0))--11 -1
for some positive constant while the other parameters in ((), (0) and ( remain unchanged.
Theorem 5.1 indicates that under teacher forcing training, the parameter matrix (k) of the first layer converges to Σ-1 the inverse of the noise covariance matrix, linearly. Combining with other parameters in (()), K( and (0), we observe that the teacher forcing training recovers the trans- former described in Thcorem 4.1 theoretically demonstrating that the expressive solution for aug- mented ICL. is identifiable and learnable.
Proof sketch. We use the superscript (ks, t) to denote the t-th CoT step in the k-th gradient descent (k.t) iteration. First, we drop the temperature term 3r in the definition of Piş given in Equation (.2
and approximate it as Equation (4.3), noting the approximation error can be made arbitrarily small by taking 3 sufficiently large. Next, define
exp 11 (R.1)
x, (
(k.t) 4i
Di=1 exp( (k.A)
x (1
which corresponds to replacing Σ by ((k) in the approximation of Pij
(k.0)

---

Class Mean Estimation Error (c=0.7)
Prediction Accuracy 0.7)
Prediction Accuracy (e
1-1
28
20
3.60
24
M =1
23
10
02-
20
M 20
4206
5030
6600
20036
14000
34600
0:02
2000
5000
0320
10000
10003
34600
4036
6002
6066
13000
10000 14030
Iteration
Iteration
on
Figure 1: Inference performance of the transformer trained via teacher forcing versus number of gradient descent iterations during training. Number classes 3, number of labeled examples N CoT steps = 5. The solid line shows the average results across 5 runs, and the shaded region represents +2 standard deviations.
To prove one-step improvement of gradient descent on the population loss under teacher forc- ing, we must exhibit constant > 0 such that -(w(k) Σ-1 (cor-teeaa)) S (l() - Σ 11 Our proof proceeds in three major steps. Step Since direct analysis of W is intractable, we propose a novel decomposition by applying Stein's lemma to break the gradient into two analytically tractable terms: one is the posterior-difference term involving Bip, (k.)) q, [ and the other is the -iffbieendiffeeencce term involving E[Vp, (k.4) - V9,h.01. Step 2. We show that an isotropic initialization of W(k) remains isotropic under gradient de- scent. The preservation of isotropy enforces alignment between Pj and qj in expectation, i.e., Elp (1 Elq (. Therefore, the posterior-differeen term vanishes. Step 3. We ama- lyze Vp, (k.1) and Vq, Shat based on the their inherent symmetric structure. This analysis shows the Jacobian difference term degenerates the following symmetric matrix under expectation: (diag(1/d) 111 ( Σ- which enables us to avoid complicated analysis di- rectly on the Jacobian difference term. Combining Steps and 3, we obtain the following upper bound for the inner product term (( Σ-1, V < -a'l [W()) E- which provides the desired result. The detailed proof can be found in Appendix D
(W(k)
6 Experimental Results
Compute resources. All experiments are conducted NVIDIA H100 GPU with 80 GB of memory. The experiments require roughly five hours to complete.
Problem setup. In the following experiments, the augmented ICL instances generated fol- lows. We set the number of classes = 3 and the data dimension d' = 3. The class mean vectors {pi)e= are randomly sampled from d-dimensional standard normal distribution. The covariance matrix Σ = eI is shared across classes, where I is the d-dimensional identity matrix. We set {0.7, 1.5}. Each instance contains N 5 labeled data points and M unlabeled data points, where M E f1, 10, 20). The M = 1 case recovers the conventional ICL setting.
Transformer structure. We construct transformer with the architecture specified in Theorem 4. This model features layers, with each layer composed of an attention module followed by an MLP module. Activation functions for the attention layers are configured as follows: softmax for the first layer, linear for the second and third layers, and ReLU for the fourth layer. We set dp = 16, and the number of CoT steps 7 5. During training. in each iteration, we randomly generate 64 augmented ICL instances, and perform one gradient descent (GD) on the average empirical CoT training loss defined in Equation (5.I) over the batch. In total, we perform 15, 000 GD iterations during training.
Results. We evaluate the performance of the trained transformer after every 100 GD iterations. For evaluation, we randomly generated 100 augmented ICL instances, and obtained the corresponding class mean estimates from the trained transformer through CoT prompting. We then utilize these estimated class means to obtain the label prediction results according to Equation (3.5) For each M E (1, 10, 20), we conduct 5 runs. We track the the class mean estimation error and prediction accuracy, and plot the average performance and standard deviation across these 5 runs in Figure] [

---

From Figure I we observe that augmented ICL outperforms conventional ICL significantly after sufficient number of training iterations. As M increases, the advantage becomes more prominent: the transformer's class mean estimation emor decrease and the classification accuracy increase, as predicted by our theoretical results Theorem and Corollary 4.1
We notice that the advantage of augmented ICL is more significant when is relatively small. This because when is small, the data distribution is less noisy, meaning that the features carry more information relevant to the labels. Therefore, the unlabeled data provides clearer structure that the transformer can leverage through augmented ICL to estimate class means more accurately.
Conclusion
In this work, introduced augmented ICL, framework in which models process mixture of labeled and unlabeled examples within the prompt. We provided theoretical insights showing that transformers equipped with CoT reasoning implement an EM-style algorithm for augmented ICI. in a multi-class linear classification task, with provably decreasing prediction error as the amount of unlabeled data increases. Moreover, showed that such transformer hehavior can emerge through standard teacher forcing training. Our empirical results support the theory.
Acknowledgments
The authors thank Li Fan, Wei Shen, Hadi Daneshmand and Cong Shen for their helpful discussions during the preparation and finalization of this work. RL and JY were partially supported by the U.S. NSF under grants 2318759, 2531023 and 2531789.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, L,, Aleman, L., Almeida, D., Al- tenschmidt, J., Altman, S., Anadkat, S., et al. (2023). Gpt-4 technical report. arXiv preprint arXiv:2303.08774.
Agarwal, R., Singh, A., Zhang, L., Bohnet, B., Rosias, L., Chan, S., Zhang, B., Anand, A., Abbas, Z., Nova, A., et al. (2024). Many-shot in-context learning. Proceedings of the 38th Interna- tional Conference on Neural Information Processing Systems, pages 76930-76966
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. (2024). Transformers learn to implement pre- conditioned gradient descent for in-context learning. Advances Neural Information Processing Systems, 36.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2023). What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations.
Bai, Y., Chen, F,, Wang, H., Xiong, C., and Mei, S. (2024). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information process- ing systems, 36.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mor- datch, I. (2021). Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 10011-55055999
Chen, S. and Li, Y. (2025). Provably learning a multi-head attention layer. In Proceedings of the 57th Annual ACM Symposium Theory Computing, pages 1754-1754.
Chen, S., Sheen, H., Wang, T., and Yang, 7. (2024). Training dynamies of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. CoRR.
10

---

Chen, W.-L.,, Wu, C.-K., Chen, Y.-N., and Chen, H.-H. (2023). Self-icl: Zero-shot in in-context learn- ing with self- geff-reeerted demonstrations. In The 2023 Conference on Empirical Methods Natural Language Processing.
Chen, Z., Wang, S., Tan, Z. Li, J., and Shen, C. (2025a). Maple: Many-shot adaptive pseudo- labeling for in-context learning. In Forty-second International Conference on Machine Learning-
Chen, Z., Wu, R., and Fang, G. (2025b). Transformers as unsupervised learning algorithms: study on gaussian mixtures. arXiv preprint arXiv:2505.11918.
Cheng, X., Chen, Y., and Sra, (2024). Transformers implement functional gradient descent to learn non-linear functions in context. In Proceedings of the 41st International Conference on Machine Learning, pages 8002-8037 8037.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. (2023). Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240)1-113.
Chung, H. W., Hou, L., Longpre, S., Zoph, B., Tay, Y., Fedus, W., Li, Y., Wang. X., Dehghani, M., Brahma, S., al. (2024). Scaling esstuuctionfiietuned language models. Journal of Machine Learning Research, 2577
Cui, Y., Ren, J., He, P., Tang, J., and Xing, Y. (2024). Superiority of multi-head attention in in- context linear regression. CoRR.
Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. (2023). On the optimization and general- ization of multi-head attention. Transactions on Machine Learning Research.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2019). Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguisties: human language technologies, volume (long and short papers), pages 4171-4186.
Diakonikolas, I., Kane, D. M., and Stewart, A. (2018), List-decodable robust mean estimation and learning mixtures of spherical gaussians. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing, pages 1047-1060.
Dosovitskiy, A. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 1999..
Fu, D., Chen, T., Jia, R., and Sharan, V. (2023). Transformers learn higher-order optimization methods for in-context learning: A study with linear models. In NeurIPS 2023 Workshop Mathematics of Modern Machine Learning.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. (2022). What can transformers learn in-context? case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598.
Gupta, S., Jegelka, S., Lopez-Paz, D., and Ahuja, K. (2024). Context is environment. In The Twelfth International Conference on Learning Representations.
He, Y., Chen, H.-Y., Cao, Y., Fan, J., and Liu, H. (2025), Transformers versus the em algorithm in multi-class clustering. arXiv preprint arXiv:2502.066
Huang, J., Wang, Z.,, and Lee, J. D. (2025a). Transformers learn to implement multi-step gradient descent with chain of thought. In The Thirteenth International Conference on Learning Repre- sentations.
Huang, R., Liang, Y., and Yang. J. (2024). Non-asymptotic convergence of training transformers for next-token prediction. Advances in Neural Information Processing Systems, 37:80634-80673.
Huang, R., Liang, Y., and Yang, J. (2025b). How transformers learn regular language recogni- tion: A theoretical study on training dynamics and implicit bias. In Forty-second International Conference Machine Learning.

---

Huang, Y., Cheng, Y., and Liang, Y. (2023). In-context convergence of transformers. arXiv preprint arXiv:2310.05249.
Kim, J. and Suzuki, T. (2024). Transformers learn nonlinear features in context: nonconvex mean- field dynamics the attention landscape. In Proceedings of the 4 International Conference on Machine Learning, pages 24527-24561.
Kim, J. and Suzuki, T. (2025). Transformers provably solve parity efficiently with chain of thought. In The Thirteenth International Conference on Learning Representations.
Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. (2022). Large language models are zero-shot reasoners. Advances in neural information processing systems, 35:22199-22215.
Li, H., Wang, M., Lu, S., Cui, X., and Chen, P.-.. (2024a). Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607
Li, H., Weng, M., Liu, S., and Chen, P.-Y. (2023). A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In International Conference on Learning Representations.
Li, T., Yi, X., Carmanis, C., and Ravikumar, P. (2017). Minimax gaussian classification & clustering. In Artificial Intelligence and Statistics, pages 1-9, PMLR.
Li, Y., Chang, X., Kara, M., Liu, X., Noy-Chowdhury, A., and Oymak, (2025). When and how unlabeled data provably improve in-context learning. arXiv preprint S5KV::2500.153299
Li, Y., Huang, Y., Ildiz, M. E., Rawat, A. S., and Oymak, S. (2024b). Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pages 685-693. PMLR.
Liu, R., Zhou, R., Shen, C., and Yang. J. (2025). On the learn-to-optimize capabilities of trans- formers in in-context sparse recovery. In The Thirteenth International Conference on Learning Representations.
Mahankali, A. V., Hashimoto, T., and Ma, T. (2023). One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. The Twelfth International Conference Learning Representations.
Nichani, E., Damian, A., and Lee, D. (2024). How transformers learn causal structure with gradient descent. In Proceedings of the 41st International Conference on Machine Learning, pages 38018-38070.
Ouyang, L.,, Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., al. (2022). Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744.
Radford, A. (2018). Improving language understanding by generative pre-training.
Raffel, C., Shazeer, N., Roberts, A.,
Lee, K., Narang, S.. Matena, M., Zhou, Y., Li, W., and Liu, J. (2020). Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-67.
Raventós, A., Paul, M., Chen, F,, and Ganguli, S. (2023). Pretraining task diversity and the emer- gence of non-bayesian in-context learning for regression. Advances in neural information pro- cessing systems, 36:14228-14246
Shen, W., Zhou, R., Yang. J., and Shen, C. (2025). On the training convergence of transformers for in-context classification of gaussian mixtures. In Forty-second International Conference on Machine Learning.
Sula, F. and Zheng, (2022). On the semi-super expectation maximization. arXiv preprint arXiv:2211.0053
12

---

Sun. Z., Shen, Y. Zhou, Q. Zhang, H., Chen, Z., Cox, D., Yang, Y. and Gan, C. (2023). Principle- driven aelf-alignment of language models from scratch with minimal human supervision. Ad- vances in Neural Information Processing Systems,
Tarzanagh, A., Li, Y., Thrampoulidis, C., and Oymak, S. (2023a). Transformers as support vector machines. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning.
Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. (2023b). Max-margin token selection in attention mechanism. In Thirty-seventh Conference on Neural Information Processing Systems.
Thoppilan, R., De Freitas, D., Hall, J., Shazeer, N., Kulshreshtha, A., Cheng, H.-T., Jin, A., Bos, T., Baker, L., Du, Y., et al. (2022). Lamda: Language models for dialog applications. CoRR.
Tian, Y., Wang, Y., Chen, B., and Du, S. S. (2023). Scan and snap: Understanding training dynam- ics and token composition I-layer transformer. Advances Neural Information Processing Systems, 16:71911-7199
Tian, Y. Wang. Y., Zhang. Z., Chen, B., and Du, S. S. (2024). Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. In The Twelfth International Conference on Learning Representations.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N.. Hambro, E., Azhar, F., et al. (2023). Llama: Open and efficient foundation language models. arXiv preprint 2rX1v:2302.13971
Vasudeva, B., Deora, P., and Thrampoulidis, C. (2024). Implicit bias and fast convergence rates for self-attention. Transactions Machine Learning Research.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polo- sukhin, I. (2017). Attention all you need. Advances in neural information processing systems, 30.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. (2023a). Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174, PMLR.
Von Oswald, J., Niklasson, E., Schlegel, M., Kobayashi, S., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., Vladymyrov, M., Pascanu, R., et al. (2023b). Uncovering mesa-optimization algo- rithms in transformers. arXiv preprint arXiv:2300.00 05858.
Wan, X., Sun, R., Nakhost, H., Dai, H., Eisenschlos, J. M., Arik, S. O., and Pfister, (2023). Universal self-adaptive prompting. arXiv preprint 4rX1v:2305.14926.
Wang, Y., Kordi, Y., Mishra, S., Liu, A., Smith, N. A., Khashabi, D., and Hajishirzi, H. (2023). Self-instruct: Aligning language models with self-generated instructions. In Proceedings of the s/ Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 13484-13508.
Wei, J., Wang. X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. (2022). srain-of-thoure prompting elicits reasoning in large language models. Advances in neural formation processing systems, 35:24824-24833
Wen, K., Zhang, H., Lin, H., and Zhang, J. (2025). From sparse dependence to sparse attention: Unveiling how chain-of -fthough enhances transformer sample efficiency. In The Thirteenth Inter- national Conference Learning Representations.
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. (2022). An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations.
Yang, J., Ma, S., and Wei, F. (2023). Auto-icl: In-context learning without human supervision. arXiv preprint XrXiv:2311.009 1.02633.
Zhang, R., Frei, S., and Bartlett, L.. (2024a). Trained transformers learn linear models in-context. Journal of machine learning research, 25(49).
13

---

Zhang, R., Wu, J., and Bartlett, P. (2024b). In-context learning of a linear transformer block: Benefits of the mlp component and one-step gd initialization. Advances in Neural Information Processing Systems, 8310-1111...1.
Zhao, Li, Y., and Sun, Y. (2020). Statistical convergence of the em algorithm gaussian mixture models. Electronic Journal of Statistics, 14:632-666 660.
Zhou, C., Liu, P. Xu, P., Iyer, S., Sun, J., Mao, Y., Ma, X., Efrat, A., Yu, P., Yu, L,, et al. (2023). Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36:55006-55011.
14

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We mentioned our main contribution and the scope in the abstract and intro-
duction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitation in the Appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specificatton, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't i acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [es]]
15

---

Justification:
We give the assumption and proof sketch for our theorems in the main paper, and give complete proof in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide all experiment details in the Experimental Results section.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data also be provided via
often one good way accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
16

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide the code in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines
cc/
mublic//gudes//CddSuubiiiss
for more details.
While we encourage the release code and data, we understand that this might not be possible, so "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
"niipe.cc/pulic/gu...oo Sume:// ississsion onPolicy) for more details.
Cubm
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only a subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We specify all the training and test details.
Guidelines:
The answer NA means that the paper does not include experiments. The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
We provide 2-sigma confidence interval in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
17

---

The assumptions made should be given (e.g,, Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of eπrоrs is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Provided in the Experimental Results section.
Guidelines:
The answer means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:://///e.repscccccccccc///ea
Answer: [Yes]
Justification:
The authors have read and fully understood the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The broader impacts of this work are discussed Appendix Due to the theoretical nature of this work, we do not foresee major negative impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
18

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All sources are cited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, wwwppppwpiiiiaeee thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
19

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: No asset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: Does not involve any human subject.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Does not involve any human subject.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
20

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Core method does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https://neutl httpps///eeerrpppeec///ooeeeerrr / for what should or should not be described.
21

---

Supplementary Materials
Contents
Introduction
2
Related Works
3
Preliminaries
Transformer Architecture
3.2
Augmented In-context Learning
3.3
Chain-of-Thought Prompting for Augmented ICL
Expressiveness with CoT Prompting for Augmented ICL
Training Dynamics with Teacher Forcing
Experimental Results
Conclusion
10
Broader Impacts
23
Limitations and Future Directions
23
Proof of Expressiveness
23
Proof of Theorem
23
Proof of Theorem
26
Proof of Corollary|4.
34
Proof of Training Dynamics
35
Auxiliary Lemmas
38
22

---

A Broader Impacts
This work provides theoretical insights into how transformers can leverage unlabeled data to improve in-context learning, a core capability underlying many recent advances in language models. By improving data efficiency and adaptability, findings could enable Al systems, partieularly in low-resource settings where labeled data is limited. These advances may benefit range of applications, including next-generation wireless communications and networking, healthcare, and financial services. Given the theoretical nature of this work, we anticipate minimal direct negative societal impact. Nonetheless, recognize that future practical implementations inspired by this research should adhere to responsible Al principles.
accessible and capable
B Limitations and Future Directions
Our analysis and experiments possess certain limitations. Below, we outline these limitations and propose directions for future research.
First, our analysis tracks parameter updates only in the first transformer layer, leaving all other layers frozen. As result, it remains unclear how weights in non-linear hidden layers evolve under teacher forcing. To the best of our knowledge, the training dynamics of multi-layer transformer with non- linear activation is still lacking investigation. A full, multi-layer treatment for end-to-end training remains an open problem.
Second, this paper is the first theoretical investigation of the influence of unlabeled data in in-context learning, therefore, we restricted the experiments to a synthetic data set. However, whether the same behavior emerges in real-world tasks, and how unlabeled examples influence in-context learning for large, fully-trained transformers, still unknown. Empirically understanding such impact is a promising future direction.
C Proof of Expressiveness
C.1 Proof of Theorem 4.1
We start from the proof of Theorem 4.1 which shows the transformer's capability of implementing an EM-style algorithm.
First, we restate the theorem below. Theorem C.1. There exists 4-layer transformer, such that its output sequence at the (t + 1)-th CoT step satisfies
P m
(1)
(0
N+M
Σ
Po
(1) (T - xx) + ((1=0)
C
Σ
(C.1)
M
N
N=
for any i [C], where (1) a (T' for some positive constants a and weight
(2) Psg is the normalized
Σ
exp
-1111
x,113_1 + BT
P3
(C.2)
ΣI= E exp
X, x1113-1 BT)
and B is a positive constant.
Recall that the input sequence at the t-th CoT step is formulated as
X X
H (t-))
Y:
0
P
P
Q
(0)
(1)
(1-1)
where
P = Pl,
P2:
pNl.
(C.3)
P, PN+1,
PN+2:
PN+M
(C.4)
Q(7) = 1q
4
(
d 00 --00 1).
(r)
(C.5)
23

---

We specify Pj and - as follows. For each data sample j E [N + MJ. we denote
0c
0a
f ((()
Oc
0c
Oc
(T)
0g
Pj
0
qp
JE[N+IN4MM
where ( stores the estimate of the mean vector of class from the 7-th CoT step, and uz stores
rescaled norm of (7), i.e., u. m/(T - |
Next, we specify the parameters of each layer of the transformer follows.
Layer 1:
The first layer of the transformer consists of an attention layer with a softmax activation function, and an MLP layer. Let the parameters of the attention layer satisfy
0 0x(d+2C)
Σ
0 04C4442)x22
1 O1x2
0 ((4420xx(++C)
V1
Ic
0 (d+C+))x(a-2--D)
Denote attm (p,) as the output token after passing Py through the first attention layer, and let %/ == attni [d) d + C 1: d 2C). Then, we have
Y; Σre(U:0-1 Eiercy exp 2 12
Xj + Br
Erejo:-1j exp
B + (1)())T TT),Tx,,))) x, + (T
Other entries in (t remain unchanged after this attention layer.
Subsequent to the first attention layer, token-wise MLP is applied. Similar to Kim and Suzuki (2025), in this work, we assume the MLP layer realize any deterministic token-wise link function with negligible error. The first MLP layer transforms input representations p such that
mlp (attnr (p;)) = - attn (P,) + +d+3
(1p ( = |13
Since 13/3C' + d d = 0 for j E [N] and p, |3C 3 for [N + +1: N + M). and the corresponding entries remain unchanged after passing through the first attention layer, this MLP layer only keeps "/ for tokens corresponding to the unlabeled dataset (i.e., j E [N]. and sets 7 to zero for all other tokens (i.e., j E [N + M))
Layer 2: The second layer of the transformer consists of an attention layer with a linear activation function, and an MLP layer. The parameters of the attention layer are set to satisfy
0 (224+40+21x22444c0
Q₂Kº
an
0
(24+30))x(4200)
V,
Ic
(4x(4+C+1)
We denote ((²) attna (7) (d) 20 1: 3C) as the vector extracted from the output token
after passing 9 through the second attention layer. Then, s,²) aa Σ N fixed scalar embedded in K₂
N+N++1 Yg. where aj is

---

We let the subsequent MLP layer realize the following token- wise Lipschitz function:
N+
(lpy(ū?))) = p -(T + 02)
p e s a
01
p
e e
Σ
Yi
+ 02
JEIN+1]
Q1
(l:y(e:)
T + 02 e;
Layer 3: Similar to the second transformer layer, the third layer also consists of a linear attention layer and an MLP layer. Consider the following parameterization for the attention layer:
((4+C)((++20)
33
Ic
0(420+0x((H++1)
10 (d+
V3
La
0c+4:+4c+al
Therefore, this attention layer realizes the following updating process:
attns (p(") mlp T T 02
Σ
x,c 7 (7)
+
JEIN
+I,+
After this linear attention layer, let the MLP layer realize the following function
+ 01
mlpy(e:)
01
Layer 4: For the last layer, introduce a transformer layer with a ReLU-activated attention layer followed by an MLP layer. We parameterize the attention layer as:
(d+C))
Ic
Q4K4
0(d-2++1x((+30+3))
0g
(d--2C)xd
V4-
La
(20+))x(1
The corresponding updating rule of this layer gives
C
alln)(m!)) atina (1,
Σ -
N
JE(N]
Therefore, we can further reformulate it as
C
Σ
Xy (eTy,), i
attny(u!))
N
JEIM
aling(u( (T)
if
Given the above 4-yayer transformer structure, by setting a1 = a/M and a2 T' for fixed > 0, I" > O, the output sequence corresponding to the ((t-1) block in the input sequence that satisfies:
(1+1)
p
(1)
N+M M
Σ
(m x;) ((==0)
C
(C.6)
M
N
j=N+1
Ø=1
for any i E [C], where (1()) = a (T' +t) for some positive constants a and T", Pig the normalized weight
Σt
- -x,113 (T
(0)
E.=0 E
- - -x, : BTT
exp
The proof is thus complete.
25

---

C.2 Proof of Theorem 4.2
In this section, we show the detailed proof of Theorem We start by restating the theorem. Theorem C.2 (Class Mean Estimation Error). Given the transformer described in Theorem 4.1 (3a>> L2K,,, K, (1/e)) and z max (M,, MM,TM T with probability at least e, the output of the transformer after CoT steps satisfies
when 2 36a log I//e and 2 max{36a2
log(L/cc
M M
N + VM
where 01, c, a. L,T", K are positive constants.
Step 1: First, we ensure that the initial estimation of the class mean vectors obtained from the labeled data gives a small estimation error.
Lemma (Initial estimation error from labeled data). Consider the initial class mean estimates
(1) m
C Σ x, . ye) e [C].
JEIN]
Then, for fixed K 2 1 and any positive constant
T 4K, we have
P
1 (1)
p.
K
S <xpp-cNNK/T)
T"
where e a positive constant.
Proof. We denote n, the number samples drawn from class in the N labeled data. Under the assumption that y; Uniform()), Vj E [N), we have n, ~ Binomial 1/C). Then, according the Chernoff's inequality, for any c E (0, 1), we have
PP
((
S 2exp -
30
For any u 2 0, let = u. K/T', we obtain
P P(mm m - E > u
K To
s 2
x2NKK
3CT
Therefore,
P [[)m: 1/]n-11>
u T"
cxp
N 3CT"
(C.7)
Conditional on n,, we have 1 Ejiy,=e, xj μr N(0, D/n.). We assume Σ is an isotropic matrix in the form of 21. Then, |Σ| 1 0², and we obtain the following inequality based on the Hoeffding's 's inequality.
ni
P(
2f Σ x, Mil > oo
n4
S 20
ni
j.y, me
For any v 2 0, by setting
= v2n, K/ (202TT),
we have
1
Σ
Xj
ui
u
K
<
2eepp(oonn,K//ee
Thi
=
K
NK
2exp
21
T
Co2T
NK
2exp
(C.8)
2Co2T
26

---

Then,
[1
K
P
C N
Σ
Xj
Cna p4i
(1
Cz. N
Mi
>
V 7
K
y.y,e.
SP
Cn:
X, μi
+
N
|| ス V
#
j,y,
Ca.
Σ ,
x
pi
T
or
Ca.
|| 2
/
(a)
exp
NK (-
for positive constant c. The inequality (a) holds by setting u = p/ |I N/Cn, in Equation C.22 The proof is thus complete.
in Equation (C.8) and setting
Step 2: Next, we bound the discrepancy between the gradient obtained from each CoT step for given input sequence, and the gradient of the population loss.
We define the population loss for any given set of class mean vectors (pilieto (i.e., any given M) as:
c(fMit) Ex log C Σ
exp(-221 M m:llB-
(C.9)
where the expectation is taken over the randomly generated data for given M, specified in Equation (3.1.
We first characterize an important property of [((p))) follows.
Lemma 2. The Jacobian of Vu,Ch f at μeg for all i E (C) is negative definite, i.e., n. L < 0.
Proof. Define
exp ( - x -1il1
Px(U:)
E exp - 2111x Me 12-
so that (p.) is softmax weight depending x and the centers (1e) Note that V2 L is the Hessian of VL at pt, given by
M.. L Ex (px(m))) px
x) - I -
where and the expectation is taken with respect to the distribution of Therefore, there exists constant 0 & a < 1 such that
vh c
(m -x) Px(())E
Now, for any nonzero vector Rdd, consider the quadratic form 42 Lu, using the above matrix inequality, we have
2 Lu Su TE (apx(())) - -)) T Px px(1))) -p
Therefore, rewriting the expectation as integral yields
1
/
pss I) To - x) x) - 1 - aN(x
u V22 Lu < C Rd
udx j"
C u T
T (m: x)(1: - xx x)>| ? m -
<. u00
Thus, the quadratic form negative for every nonzero u, and the matrix V2 negative definite. This completes the proof.
27

---

We note that for each CoT step t >> O, the updating induced by the constructed transformer is
(/(()
N
(++1)
(1)
Σ
Pij
(t)
(i,
x,),
(C.10)
M
N+1+1
10 where Pij is defined in Equation (4.2)
To simplify notation, denote
N+M Σ J=N+1
Puj ( xx)== V 2
(C.11)
We note that C itself is not explicit loss function. We use the notation V ( C to represent the equivalent gradient for the updating determined by the t-th CoT step. Lemma states that Vj, C is negative definite for each μıg. In the following lemma, we show that V L is negative definite for the concatenate vector μ when {
c are well seperated.
Lemma 3. The Jacobian of μL at ş negative definite, i.e., ViLC
Proof. Recall that E E Rdxd is symmetric positive definite. For E Fd and μi (HI,...., Hc) E (R define
((x; p1) log à Σ exp( |||x
1=112-1)-
Therefore, we have L(u) Ex [e(x; m The quadratic form of the Hessian of in direction
A A A....., Ac) can be written as
A (x)
Σ (p((m))lA + Var,, Tarp p4), A)).)))
(C.12)
Define
Pgg = [44 - 1
P. = min Pays Aij(x) = =||== -|xx- 11511 - ||x -17113-1 151
Then, we obtain that
A1j(x) = pi; + 22zj.
where Zij ~ (O,PP)))
Define the event
::= min Ax(()) 1
(C.13)
J#
and its complement E. Then, based the Gaussian tail bound and taking union bound over j f i, we have
Pr(E) 2 1 - (C 1)e 02/8 -08/8
n
(C.14)
Then, under event E, px (11) is lower bounded by
(x(4)) 2
1 ++ E AAu(X) 2 1+ (C (C-
1)e-12/1 = =:
(C.15)
Using C.144 C.15). we obtain
Ex Px (M::) Ex (p: Pr(E) + [x(()))) 2 n.) 0.
(C.16)
28

---

Therefore, based on Equation C. the expectation of the first term in Equation C can be upper bounded as
E px(())) -- E-EO 7.)8..1 A A-----
(C.17)
Define
U1 (E- (x pe:). A
(C.18)
Then, the second term in Equation
(C.12 can be expressed as
Exelmndu 2 0.
Varp (Hu:)i=)) = Px(li)uf -
Note that
Varp. (fu:
Σ
(x
Σ Px(()) )²
i=1
Px (n:)uz
( Po(()))
Σ peluu) (x + +ut))
( pe((ii)us)
2
Pa(m))) =2²²
(H, uz)
Σ (1 - (x()))))
(C.19)
Next, we condition Varp, [ 1C 1) on event & and its complement E, respectively, When event holds, combining Equation (C.15) and Equation (C.19) gives
Var (f4,) S Σ Px(mx)(1 Px (x,))))
2 Σ (1 1B.)u²
Under event E. have
Var (Lu) CC S -2 ut.
Σ
Define
k. max Ex | μ
Isico
Then, based on the definition of in Equation
(C.8). by using spectral norm inequality, we have
Ex [u]] = A E (x)(( (4) Σ A
S k.
Taking expectation over obtain
Ex[Varp (fust. s (2(1- 3.) + 1/2) k. Σ | 1
(C.20)
29

---

Plugging Equation (C.17) and Equation (C.20) into the expectation ofEquation C.122 we have
Ex
A
v
(()) A] S - -((1 -111-1 (1-m))8. (2(1 +.)
k. 2))
Σ
| E-1-
Recall that
77. (C - 1)e- B.= 1+ (C 1)e-12/14 12./4
and P+ is defined as
Piy = [c: 1
P. = min Pij. it
min
Then, for sufficiently large P. such that
(1 - n,) - (2(1 - 3.) + 3)) K. > 0,
4 ((p) negative definite.
In the following, characterize V p formed on the population loss. We have the following lemma. Lemma (Properties of the CoT gradient descent). Fix an epoch and a component index i [C].
E and compare it with VL, i.e., the gradient if GD is per-
there exist constants C1, C2 >> O such that, for every M 2 1,
Pr V p V 2 Sa M M-1 2 -1--e ex-exp(-VM)
and
Pr(1IV 210 Ell S 02 c2 cgM-
21 (1-eep(-VM)
Proof. Recall that
N+M
V n
E.
1
Σ
Pij
(
(p! -x,);
M
f=N+1
where p is given by
Σ 7=1 exp
-b x,,13 11/25 + +BT
(0) Pūj
E1=0 Ee exp
- x 1 + BT)
By choosing 3 + 0o, we further have
1111 x1,13
(0 exp
P5
E exp
-11 x, x1
where the samples {x,}s2N+1 are drawn from Gaussian mixture distribution.
Therefore, given p! the random variable (() (p x,) admits a uub-Gaussian tail bound since
Xj are Guassian random vectors and pi (m! x,) is Lipschitz continuous over Xj.
Then, by the Bernstein's 's inequality, for any fixed 8 (0, 1), with probability at least 5, we have
N+M
V F C C =
Pij (ft xj) Ex,
(0) Pej (p! x,)
M
J=++1
C1
log
30

---

where > 0 is some absolute constant.
By choosing 5 = exp(- M), we obtain that with probability least 1 exp(-evppp))epe M),
2 V N0 S S
(C.21)
for another constant d' > 0. This completes the proof of the first inequality.
Next, we show that || x itself is bounded with high probability. Consequently,
N+M
V =
Σ
Pij
10
(m
-x,)
M
j=N+1
W+
Σ
(p
x5)
Pij
M
j=N+1
Σ
X,
M
12N+1
(1 || + ||2} Il)
S M
j2N+1
|| M
|
(C.22)
J2N+1
where inequality (a) holds since Pig
(1)
< 1. Note that
(/(+- N+M
p!!)
(1-1)
Σ
(t-1)
(p((-1) x,)
pu;
Pij (4:
j=N+1
(2
(t 1}
N+M
p
(l 1e-1)
+
Σ
Pi①
(1-1)
M
M
N+1
Therefore, we have
|
S |1(!(-11)
11
+
|I,|||
M
J2
Σ
||y|||
M
J2N+1
Combining with Equation C.22 have
Σ
||x;||-
M
J2N+1
Applying the Bernstein's 's inequality, with probability at least 1 exp(e (VM,, we have
1
ps + -
1x
S
M
+2
1=1
where C; is positive constant.
Therefore, for any S T where T is total number of CoT steps, we have
C
|/ p
S
+
T
pi +
c
31

---

which implies
IV (0) 21 s c2 + c
where C2 and 03 are positive constants depends on T, |1 and T E Mi: The proof is thus complete.
Step 3: Finally, we show the convergence of the class mean estimation error.
Expanding the squared error || [(a](tt) 1iill gives
||
1ill² = [ 14l12 + (1)
( p Mi, V 22 2) //// / 21
< <| (t) 14:1² + v1) (1))) Mi, V BooL) + 2 (1) VC v2
+ ) V 2 21
(C.23)
Denote ( and fe the vectors obtained by stacking (7: fore, we have
CC and { respectively. There-
(
u1l S
22 (f((0 Me, V nio C) +2 "0 VLC voI + (1)")) pmo 2
To control the inner product term (īi(() ,,, V ( C), we perform first-order Taylor expansion of V Vite C around μ as
aan C V VuC+ (V,L) ( -w) + RR((") 1)) (a) (V (7 11) + ++R(( ,1),
where equality (a) holds since is the global minimizer of L and L is differentiable on Rd, thus V = 0, and R 14) is the remainder term
For the remainder term, have
(R( 11), 1), 14)
1. (a - H1T (V (Vuteraco-pp) -
-V2C (a =.)
1/¹ | V -2-c |1/(00 -m1/2
LEI1 cll = L L1 41ll
where Inequality (b) follows from the fact that v is twice continuously differentiable, its Jacobian is Lipchitz continuous in a neighborhood of μa, and L is the Lipchitz constant.
Therefore, there exists a constant À 0 such that
(1)(() - 1al12 + (2(0)) ,4 21) (2) V v)C())
< ||((0 1112 2 + (p (pt 11) T V (((!)) -p) +
())) L1 ull
(n (1 (1 >) (1
1ll1 + 2n [1 1alls
(C.24)
where Inequality (c) follows from Lemma which proves V L is negative definite.
Meanwhile, Lemma ensures with probability at least 1 - exp( - MM),
(1) VVC vद Snnn (0) M 1
(C.25)
(0 | n. 211² < c2 () M 2 + c3 (0)
1
(C.26)
32

---

Substituting CC.44)
C:25
and (C.26) into (C.23 then yields the one-step error recursion
(1(2))) pal2 = (1 2n (1) >) p 1ell + (1) L u1l
+ ( M 3 + (1) " ())
(C.27)
Next, we aim prove (1 - 1e|13 S K/(1+ T*) for a positive constant K by induction.
For some p 2 4, let
1/0) a t + 7
0 = (t + TT).
(C.28)
First, assume ( ul12 < K/(t + T for fixed t 2 1. From Equation (C.27 we note that
there exists a constant C₁ > such that
(1(1t) - -11l S
1 1-2 at
11 - -112 + c3 +t+ Th2
a²
+
aL
11 -u1l + C4
a
1
t ++
t
+ T
1 2 al
K
aL
K
$
a²
T+ +2
+ + (1) (1+ + C3
t + T"
(t + T')²
Therefore, we have
K
a
(++1)
111
++
+1
ai
K
aL
K (t+ T'
a²
K
K
< 1 2
+2
+ Caa(t++T)--1++
+ c3 (1+
t T
T
++T T"
++
(t (t+
K
aL
K
a²
= (-2ax 1)
+ +ca(l+ -(++) + C3
(C.29)
(1 Tiy2
1+
| T
(1 + Trp2
By choosing a 2 1/1, K 2 max 13cza 2 3c4aa and T 2 36a L,2,, we have
K
K
(-2ax + 1)
5
(1+
(1 + Tr)
2 aL K
38
K
t T' ++ T'
3(t + +')22
(C.30)
qja(t T'> (1+ (1++)
K
3(t + T')2'
a²
K
C3 (t + T' 3(t + T'y2"
Therefore, by substituting Equation
C330 into Equation C.29
we have
K
|| (t+1)
11l2
< 0,
1+ T +1
which implies
(++1)
-1el2-
K 1. T' +1 K +1
K
t + +1 K+1
K+1 ++1*
14 +1 1+ + T
K
K+1
t+ + T 41 K
33

---

Recall Lemma indicates that, with probability at least 1 - exp(-cNK/TT) for some constant it holds that [| Mlll - K/T". Therefore, for any fixed E [0, 1), let
N T (og((1/6))
36a
ck = M (t + T*)*
log(1/c);
(C.31)
max L²K, log (1/e)).
Then, with probability least c e- where the term arises from the condition required for Equation
CC22 to hold, the estimation error is upper bounded by
K +1
[a(t(() 11
+ T + K
K+1
(C.32)
M +
No. log(1/e)
For sufficiently small such that log(1/c) c, denoting d' K+1 we obtain
K
+
1
(K+ log(1/e)
M
No
min((1/ c) ( +
log(1/c)
log(1/e)
VM +
Since we let /M- +t² > log(1/e), we conclude that there exist constant c" such that with probability 1 - e, the estimation error is upper bounded by
ull d" log(1/e)
N + VM
This completes the proof of Theorem 4.2
C.3 Proof of Corollary 4.1
First, we restate the corollary below.
Corollary C.1 (Restatement of Corollary 4.1 Let V; be the predicted label for xj according to
Equation (3.5.. Let R be the prediction error under the Bayes-optimal classifier with known class
mean vectors M1...... M.C. Then, under the same conditions as described in Theorem 4.2 have
Ply, # y/p1, ,,c....,,c - R. $ 001 2011/// poly(M)).
Proof. First, we define
M M||F, define ÿ as the Bayes-optimal classifier given estimated
class means M and define g the Bayes-optimal classifier given ground truth class means M
Suppose ğ(x) # gxx). Then, there exist indices * such that g(x) and g(x) k. Because g(x)= is Bayes-optimal, we have
x mill S |x uall and ||x - --xmlll S ||x p.1l-
Denote < ||4 mal-- Therefore, from the geometric observation, the misclassification only happens when x is in the dihedral with angle where tan(0) A/C (Diakonikolas al. 2018).Tuss Thus, the probability for misclassification is upper bounded
( * 9(x)] < c'O.
a positive constant d. Since P + pI.
mc] R = [x]((x) + g(x)) and from Theo-
rem 4.2 we have A s d' d//1/(N VM) for positive constant c', the proof is thus complete.
34

---

D Proof of Training Dynamics
First, we restate Theorem 5.1 below.
120 be the parameters of the first attention layer of the transformer after applying k iterations of gradient descent on the pop- ulation loss defined in Equation (5.2). Then, with the initialization specified in Assumption | have
Theorem D.1 (Restatement of Theorem
5.1 Let (Q(h), K(k), V(
[W - S 1 IE
for some positive constant 6, while the other parumeters in Q(),, K(O) and V remain unchanged.
We assume ground truth means are IID sampled from standard Gaussian distribution: μ4, (O, for all i. Then, we introduce the following quantities: 1) the formulation of class mean estimations given by the transformer during teacher forcing training; 2) the reference class mean estimations given by the reference policy; and 3) the formulation of the gradient of the teacher forcing training loss.
At the k-th GD iteration during training, we denote the set of reference class mean estimations (R,t) (R.t) as 1 CrefCC for the CoT steps [T]. Given the reference class mean estimations, the estimation given by the transformer throughout teacher forcing satisfies
17(0)
N+M
(k(+1))
(k.e)
(k.)) (k.t) - -x)).
p
Mof
Pij
M
N+1
(k.11)
where Pij is given by
Er exp( -
112 x w p( (r) + (3)
(k.t)
Pij
Et Ee exp(
+ x w pt) + (T)
By choosing 8 + co, we further have
exp 10
1
(T)
+ x ( pp
Pif
EE exp
(
12+ x W W pz
We choose the reference policy under which
(7())
Peef Tof
(k.++))
(8.1)
Σ
Pij
(k,1) (k.l)
(1 -x,);
M
j=N+1
with
exp
- 1 Hef
(k.t) - 11112:-11
(1.0)
Pij
EC exp - ref.c x, 112.-1
To simplify the notation, when there no ambiguity, we drop the superscript (k) for the training iteration. Denote q
(t) At the k-th training iteration, the (t) p Peyl and q." [p]) PPell
(tt)
CoT training loss with teacher forcing is
N+M
ECoT (0:IM)
1
Σ
Σ
CE
[ (H (t-))))
ref 1
T
t=1 N=
1 T N+M
Σ
Σ
CE
T N
where CE is the cross entropy loss function.
35

---

Define (1) % written as
| 1 1|2 + x p and 5 = -11/. cyl. Note that the derivative can be
=
0 expla!!))
0CE
{ 10 ap.a)")
2 exple
(() (e) Pii Pij
0s (1)
Furthermore, since /s! /0 = Ma, where a, b E [d]. by the chain rule, we have
acE 0
dcE
Σ
(2) Pig Ma,
(D.1)
OW ab
0
OWas ab
Based on the notations, we will prove Theorem 5.1 as follows.
Step 1: Given the gradient of the eross entropy loss with respect the learnable parameter matrix W, our first step is to provide a decomposition of the gradient so that it becomes analytically tractable.
In the matrix form, Equation D.1 can be written as
Vw CE qj
(2)
.qa')
M(q))) -4, )x
By the Stein's 's lemma, we have
BMM
qj
(1)
}x
EM
M
Ex, 1
-q [Ex]
Ex, Vq. |
EM
ME
(8)
9, [[
EEM
MEx,
Vqg"" (1)
Vq
A
Step 2: Based on the decomposition, we aim to show that A1 = 0.
We
note
that when taking the expectation μ4. Therefore, pler, p4. When the reference class mean Σ JJENN Xy (e (eTv;)]
the labeled dataset, have
E
C
estimations are generated by gradient descent over the population loss, we have p = μi for any i E [C] and t e [Z'] the gradient over the population loss is zero:
(t)
exp
(t)
x1,13
Fx
(u (or, x;))
c=l
exp
11 -x,l1z-1
exp(
- 1111400, - -x1
Σ
(xx(())
(u
x)dx,
R2
>c exp( - - /1ll x|| )-1)
C
k=1
(0)
(P(( - -x)x((x=xx===---
where sp (x) is the pdf of Gaussian distribution with mean μ,, and covariance matrix E, and equal- (k.1)
ity (a) holds since Pref, p4. Given the above-diseussed property of the reference class mean estimations, for Ex, 9 q in A1, its is obvious that Ex, [ 1/C. For Ex, (4 let ( initialize form a isotropic matrix wI, and we assume at training iteration step kt preserve the isorooicc as (Ak), Therefore, since the ground truth Σ an isotropic matrix, the temperature acts identically on all classes:
exp
-x,1(2
exp -
x3 x1113-11
= E
Do=1 exp
- - x,112)
- exp - ||:
x
36

---

Therefore, we have xx, (0) p O, which gives A1 = 0.
Step 3: Finally, we analyze the properties of A2, and obtain the final results afterwards. We will prove that W preserves isotropic by induction. Note that we assume training iteration step k w(k) is isotropic. Besides, we initialize W(0) an isotropic matrix.
Ag can be rewritten as
A = EM MEx, (Vg. vq [
= EM [ ((diag(E(q))) - Ex (9 (T) M T w - -(ding(Blaf))) () - Ex (q (a (af)) M M.
Because the class prior is uniform and the isotropic initialisation, we have Ex, [a]]]] Ex, [q 1/C. Since each coordinate of q; coordinates have the same joint distribution, we have
(() (or a has the same marginal distribution and any two distinct
(1)) (i(g(B()))) = diag( ( = 1 gz,
Exqq [q = Ex "a (q c 11
Therefore, we have
A = EM M diag diag(1/O) ca 117) MJ 10110=11]
Note that V Lcor Loor (W())) A2, therefore, we obtain
1 MT (w/+=>=-)] -1)]||
Woor (W
EM
M
(diag(1/C)-
Che 11 12
1
-MI (ovl---))1
1
EM
[ MM
=0² (1-C) ()) (
Since all columns in M are sampled from (0. I) and (()) is assumed to be an isotropic matrix, it's obvious that A2 is also an isotropic matrix. It follows that
(w()) E LcT)
= =BM trace (M
(M(
C2 11) MJ (WO---)) E (w(k) (0-))---1']
1
(0) rrace
[d
(Wr(t) - (wll))
(w -2-1)))-11)) T
EM
c MIIJMT] (W())) -2-1)----- (wix)) (w ( - )
EM
-oaceace ((w(() (w xー - -/4 - (wrt)) (1(())-11)
=o 1 (1-2)
1 2 (w(( -
where equation (a) follows from the assumption that Σ = I'I.
Let = (1 - 1/C).. Then,
[ --
E < (()) E -1-
--mm ((k) - Σ 11 + v277' (()) -
<(1 - <1) (()) 1
Select step size such that (1 21 and let (1 - 2, we obtain
[W Σ
1| < * | Σ 11.
37

---

E
Auxiliary Lemmas
Lemma (Stein's Lemma). Let X E Rdd be a random vector with
X X~N(u, Σ),
where μ Fd and Σ Rdxd is a positive definite matrix. Let f Rd R& be a continuously differentiable function such that
E [llf(())II
< 0o and E [[]Vf(X)]l=] < <0,
where || denotes the Euclidean norm in R& and || || is the Frobenius norm. Then, the following identity holds:
E[(x -e--No p) -E
where V f() is the d Jacobian matrix of f evaluated at X.
38


Reviews:
[{"review_id": "iJHEjnJEHS", "reviewer": "Reviewer_RyNA", "summary": "The paper introduces a new framework called augmented in-context learning (ICL) that leverages unlabeled data to improve the performance of transformers in classification task.\nAugmented ICL Framework: The research proposes an augmented ICL framework where a transformer's prompt includes both a small set of labeled examples and a larger block of unlabeled inputs. This approach allows the transformer to infer missing labels within a single forward pass, bypassing the need for costly and time-consuming data labeling or pseudo-label generation.\nExpressiveness with Chain-of-Thought (CoT) Prompting: The paper demonstrates that a multi-layer transformer, when prompted with Chain-of-Thought (CoT) reasoning, can effectively emulate an Expectation-Maximization (EM) algorithm for multi-class linear classification. This enables the transformer to extract useful information from both labeled and unlabeled data, leading to provable improvements in ICL accuracy. The study theoretically shows that the class mean estimation converges to the ground truth as CoT steps increase, and the approach's excess risk scales more favorably compared to methods using only labeled data.\nEmpirical Validation: Experiments show that the augmented ICL framework consistently outperforms conventional few-shot ICL in class mean estimation and label prediction, with the benefits becoming more significant as the amount of unlabeled data increases. Augmented ICL even surpasses the Bayes-optimal classifier that relies solely on labeled data.", "questions": "see", "limitations": "Does language models trained on real-world data, where patterns are not as clear as in the synthetic data, align with the proposed theoretical explanation of unlabled ICL?\n\nDoes the proposed theory generalize to other architectures than transformer?", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "I am not familiar with this field. So my review will provide limited value.\n\n## Strengths\n\nThis paper shows that transformer language model can learn from unlabeled ICL data from a theoretical perspective.\n\nThe paper is generally well-written and clear. The appendix provides comprehensive details about the proposed theory.\n\nThe paper's originality is high. The concept of integrating unlabeled data into ICL in a theoretically sound manner, specifically by showing how a transformer can emulate an EM algorithm through CoT prompting, is novel. While ICL and EM are not new, their combination in this specific, provable framework for transformers is a unique contribution that pushes the boundaries of our understanding of in-context learning.\n\n## Weaknesses\n\nLacking experiments on real-world pre-training data.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful reading and thoughtful comments. Please see our responses below.\n\n***Weakness 1:*** Lacking experiments on real-world pre-training data.\n\n***Response to Weakness 1:***\nWe thank the reviewer’s helpful comment. As a theoretical work, in our submission, we rely on a synthetic experiment to validate our theoretical claim, i.e., a transformer trained with standard teacher‑forcing can learn to exploit unlabeled data in addition to labeled examples to enhance its in‑context learning (ICL) accuracy.\n\n\nTo validate the theoretical claim in more practical settings with a real-world dataset, as a first step, we conduct experiments on the AG News classification dataset with the Qwen3-8B model. The token embedding dimension is 4096, which is much higher than the dimension we consider in the synthetic experiment. We consider 4 candidate classes, and for each prompt, it consists of 4 labeled examples along with $1, 4$ or $8$ unlabeled samples. To encourage CoT reasoning, we use the following prompt:\n\n\n*First, think step-by-step, compare them with the patterns you saw \n in both the labeled and unlabeled blocks, and decide the best category for each one. \n You can compare with other unlabeled samples before making a final prediction.*\n \nFor each prompt with $M$ unlabeled samples, we evaluate the model's performance after it completes its CoT reasoning process. We extract the final prediction for each of the $M$ samples from the model's generated output. These $M$ predictions are then compared against their ground-truth labels to calculate the overall prediction accuracy. The report results in the following table.\n\n| M | Test Accuracy |\n|-----|---------------|\n| 1   | 0.880 |\n| 4   | 0.896 |\n| 8   | 0.926 |\n\nWe note that the average test accuracy, calculated over 100 prompts for each M, increases monotonically as M increases, which is consistent with our theoretical results. \n\nWe will conduct more comprehensive experiments on a range of ICL tasks using real-world datasets and will report the results in the revised version.\n\n\n***Limitation 1:*** Does language models trained on real-world data, where patterns are not as clear as in the synthetic data, align with the proposed theoretical explanation of unlabeled ICL?\n\n***Response to Limitation 1:*** We thank the reviewer for the insightful question. As elaborated in our response to Weakness 1, we have performed an additional experiment on a real-world dataset to verify the impact of unlabeled samples on the ICL performance of a pre-trained LLM. We observe that although the assumptions we made for the theoretical analysis may not hold in this setting, the ICL performance monotonically improves as the number of unlabeled samples increases, which is aligned with our theoretical explanation of the augmented ICL.\nWe will conduct more comprehensive experiments on a range of ICL tasks using real-world datasets to verify this and will report the results in the revised version.\n\n***Limitation 2:*** Does the proposed theory generalize to other architectures than the transformer?\n\n***Response to Limitation 2:*** We thank the reviewer for this question. Our work is intentionally focused on the transformer architecture, as ICL is a unique capability that is closely associated with and demonstrated by this class of models. Therefore, it is not straightforward to extend our theory to other deep learning architectures.\n\nHowever, we emphasize that our theory is not limited to the specific multi-layer transformer with CoT reasoning analyzed in the paper. It is generalizable across the transformer family. For example, our framework can be extended to a recursive transformer (or looped transformer). This would involve reformulating the explicit reasoning steps of CoT into implicit reasoning steps handled by the model's recurrent structure, which can be achieved by designing a suitable embedding scheme. \n\n\n----------\n\nWe hope our responses have addressed your concerns satisfactorily. We sincerely appreciate your constructive feedback and would be grateful if you would consider raising your evaluation based on our responses. We would also be happy to address any further questions you may have."}, {"review_id": "500pvHR80F", "reviewer": "Reviewer_J6se", "summary": "The paper presents some theoretical insights leveraging unlabeled data for ICL, in the multi-class linear classification setting. Key findings include that transformers can emulate EM-style algorithm with CoT prompting and such models can be trained using teacher forcing, resulting in provable ICL performance improvement. Empirical experiments are also conducted to validate the theoretical insights.", "questions": "1. \\[Line 130\\] It seems the normalization layers, which are commonly used, are missing here. Is there any justification?   \n2. \\[Line 160\\] The assumption that covariance is isotropic and shared across classes seems a strong one. How will the insights change if we relax these two assumptions?   \n3. Or if these two assumptions are essential for derivation, please elaborate in Section 3.2  \n4. \\[Line 323\\] The data dimension of 3 seems very small. Is there any justification for this? In practical settings, having tens or even hundreds of features seems typical.", "limitations": "na", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "1. Unlike traditional semi-supervised settings, this paper explores leveraging unlabeled data directly, which is a new and interesting setting.   \n2. There has been previous work showing empirically that unlabeled data can boost ICL performance, even simply batching some queries (arxiv.org/abs/2405.09798). Providing theoretical insights is significant.   \n3. The implementation addressed in Section 5 is a good addition to the theoretical existence proof.\n\nWeaknesses\n\n1. The work is limited to the linear classification task, and there’s little discussion on the applicability of the insights beyond this setting.  \n2. The setup for empirical experiments (Section 6\\) is not sufficiently justified. Many decisions seem important but not explained.", "was_revised": true, "final_justification": "The rebuttal has addressed my concerns, so I've raised my rating. However, my confidence remains low as I don't have much expertise in theoretical research", "rebuttal": "We thank the reviewer for the careful reading and thoughtful comments. Please see our responses below.\n\n***Weakness 1:*** The work is limited to the linear classification task, and there’s little discussion on the applicability of the insights beyond this setting.\n\n***Response to Weakness 1:*** We appreciate the reviewer's valuable feedback. While our theoretical analysis indeed focuses on the multi-class linear classification task, this choice allows for rigorous analytical insights into the transformer's capability to utilize unlabeled data effectively. Importantly, our theoretical framework sheds light on the fundamental mechanism by which the transformer is able to leverage unlabeled data during ICL through CoT prompting. We believe our analysis can be extended beyond the linear classification task considered in this work.\n\nFor example, a natural extension is to consider *generalized linear classification* tasks. In this framework, the label $ y \\in \\\\{-1, 1\\\\}$ is modeled through a nonlinear decision boundary using a link function $\\sigma: \\mathbb{R} \\to [0,1]$, such that the conditional probability of the label is given by\n\n$$\n\\mathbb{P}(y = 1 \\mid \\mathbf{x}) = \\sigma(\\mathbf{w}^{*\\top} \\mathbf{x}),\n$$\n\nwhere  $\\sigma$ is commonly chosen as the sigmoid function $ \\sigma(z) = 1 / (1 + e^{-z})$, and $\\mathbf{w}^* \\in \\mathbb{R}^d $ is an unknown parameter vector. This generalized linear model (GLM) setting introduces a richer hypothesis space while retaining mathematical tractability. In the context of augmented ICL, a key question is whether a transformer can utilize both the labeled and unlabeled examples to infer the latent structure governed by $\\mathbf{w}^* $ and improve the ICL performance. \n\nWe can construct a transformer and show that it can mimic a gradient descent step on the cross-entropy loss defined on the labeled examples plus an entropy regularization term defined on the unlabeled data. Such a regularization term can be used to improve the estimate of the optimal weights  $\\mathbf{w}^*$ and prevent overfitting, especially when labeled data is scarce.\nThen, we can show that the transformer can be obtained through a CoT-based teacher-forcing training.  The training loss is the $\\ell_2$ distance between the transformer's output at the $t$-th CoT step and the $t$-th gradient-descent step.\nWe will explore such extensions as the next step of our research.\n\n***Weakness 2:*** The setup for empirical experiments is not sufficiently justified.\n\n***Weakness 2:***  We thank the reviewer for this helpful suggestion. The key parameter choices of the experiments are motivated as follows: The number of classes ($C=3$) and data dimension ($d=3$) are chosen to make the task more complicated than the binary case ($C=2$) while maintaining a low computational cost.\nNoise levels ($\\epsilon \\in \\{0.7, 1.5\\}$) are selected to assess our method's robustness across varying signal-to-noise ratios (SNRs). The choice of the two distinct levels allows us to test the model in both low-noise (high SNR) and high-noise (low SNR) regimes. Transformer input dimension ($d_p=16$) is a direct consequence of our proposed embedding architecture, as specified in Equation (3.2). This choice ensures consistency between our model's theoretical formulation and its practical implementation.\n\nWe will integrate these justifications into the revised manuscript to improve clarity and reproducibility.\n\n\n\n\n***Question 1:*** The commonly used normalization layers are not considered.\n\n***Response to Question 1:***\nWe thank the reviewer for highlighting this important point. Indeed, normalization layers are commonly utilized in transformer architectures. However, in our theoretical framework, we omit the normalization layers to simplify the analysis and provide clearer insights into the essential mechanics of the transformer's learning and inference processes. This choice follows recent theoretical ICL studies [1–4]. We will add this clarification about the normalization layers in our revised manuscript.\n\n\n\n***Question 2:*** The assumption that covariance is isotropic and shared across classes seems a strong one. How will the insights change if we relax these two assumptions?\n\n\n***Response to Question 2:*** We thank the reviewer for this insightful question. The assumption of a shared, isotropic covariance matrix is a strategic choice made for theoretical tractability. Specifically, this simplification ensures that the update rules within the attention block maintain an isotropic structure, which is crucial for deriving a tractable, closed-form solution.\n\nWe note that this assumption is common and widely adopted in related literature to facilitate theoretical analysis [5, 6, 7].\n\nIf we were to relax this condition, the model would need to estimate class-specific and anisotropic covariance structures. This alternative presents significant theoretical challenges, making the update rules difficult to analyze. \n\n\n***Question 3:*** Need to elaborate if these two assumptions are essential.\n\n***Response to Question 3:*** We thank the reviewer for the suggestion. We will revise Section 3.2 as follows (from line 176):\n\n*In this work, we assume $\\mathbf{\\Sigma}$ is isotropic. We adopt this assumption for theoretical tractability, as it is crucial for deriving the closed-form update rules for the transformer. This approach is a standard and widely adopted practice in related literature to facilitate theoretical analysis (He et al., 2025; Zhang et al., 2024; Chen et al., 2025).*\n\n\n***Question 4:*** The data dimension is small.\n\n***Response to Question 4:*** We thank the reviewer for this question. We choose $d=3$ to validate our theoretical results efficiently and with low computational cost. Importantly, our theoretical framework and conclusions are not limited to this specific dimensionality.\n\nWe consider a more practical setting with much higher data dimension and conduct an additional experiment on a real-world dataset with a pre-trained LLM.  Please refer to our responses to Reviewer RyNA (Response to Weakness 1) for the experimental setup and results.\n\n\n\nReferences:\n\n[1] Uncovering mesa-optimization algorithms in transformers.\n\n[2] Transformers learn to implement pre-conditioned gradient descent for in-context learning.\n\n[3] Transformers as statisticians: Provable in-context learning with in-context algorithm selection.\n\n[4] When and How Unlabeled Data Provably Improves In-Context Learning\n\n[5] He et al., \"Transformers versus the em algorithm in multi-class clustering.\" arXiv preprint arXiv:2502.06007 (2025).\n\n[6] Zhang et al., \"In-context learning of a linear transformer block: Benefits of the MLP component and one-step DG initialization.\" Advances in Neural Information Processing Systems 37 (2024): 18310-18361.\n\n[7] Chen et al., \"Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures.\" arXiv preprint arXiv:2505.11918 (2025).\n\n\n----------\n\nWe hope our responses have addressed your concerns satisfactorily. We sincerely appreciate your constructive feedback and would be grateful if you would consider raising your evaluation based on our responses. We would also be happy to address any further questions you may have."}, {"review_id": "m5A6Zs3RYe", "reviewer": "Reviewer_ZDrZ", "summary": "The paper introduces augmented in-context learning (ICL): a prompt that mixes a few labeled examples with many unlabeled inputs. Focusing on multi-class linear classification under an isotropic Gaussian-mixture model, the authors prove that a 4-layer transformer, when driven by chain-of-thought (CoT) prompting, can emulate an expectation–maximization (EM) algorithm. Experiments on synthetic 3-class, 3-dim Gaussian data confirm that accuracy improves consistently with more unlabeled points and surpasses the Bayes classifier trained on the labeled set alone.", "questions": "Please refer to the weaknesses section above.", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths** \n- The core idea of augmenting a prompt with unlabeled data to improve In-context learning is interesting.\n- The paper provides rigorous theory under multi-class linear classification setting for a multi-layer transformer \n- A major strength of the work is the explicit construction of a 4-layer transformer that can emulate an Expectation-Maximization (EM) algorithm using chain-of-thought.\n- The writing is very clear and easy to follow. Notation is kept minimal, derivations are step-by-step, and appendices contain complete proofs, making the work easy to follow and verify.\n\n**Weaknesses**\n\n- The 4-layer network is rigidly engineered: later layers are hard-wired to act only at specific CoT steps (e.g., Layer 4 fires solely on step 1 to initialize class means). The analysis confirms that if such a network is trained under teacher-forcing it will converge, but doesn't show that a standard autoregressive transformer would naturally converge to this this circuit.\n\n- Theorem 4.2 requires the number of CoT steps (t) to grow with the number of unlabeled samples ($t \\geq M^{1/4}$) to achieve the improved accuracy.  Since each CoT step is a forward pass over a sequence whose length also grows with $M$ and $t$, the total inference cost grows substantially with the amount of unlabeled data. The paper also doesn't provide any analysis of memory or latency, which is critical for practical tasks. \n\n- All experiments use synthetic 3-dimensional, 3-class mixtures with just five labelled points. There are no language, vision, or large-scale tests, and no comparison against standard semi-supervised baselines (e.g., pseudo-labelling). Its unclear how well the theory translates to and how well this approach is suited for real world tasks.  \n\n- Figure 1 reports that the transformer exceeds the Bayes oracle trained on the labeled set for certain $\\epsilon$ values, but that oracle ignores unlabeled data by construction. A fairer comparator would be an explicit EM or self-training classifier that does exploit the same unlabeled dataset.\n\n- This work misses some key related works. For instance, the paper \"Context is Environment (Gupta et al., ICLR 2024)\" already provides theory and experiments showing that unlabeled context tokens improve ICL under distribution shift. The present manuscript neither cites nor contrasts this closely related work. Similarly as mentioned above, 'Many-Shot In-Context Learning (Agarwal et al.) also uses unlabelled queries to increase the context length and subsequently performance.", "was_revised": true, "final_justification": "The authors addressed most of my concerns. However, despite the additional experiments on AG News dataset, there is no comparison with existing baselines such as standard semi-supervised approached (e.g., pseudo-labelling). Its unclear how well their approach is suited for real world tasks.  Thus, I'll maintain my score at 4.", "rebuttal": "We thank the reviewer for the careful reading and thoughtful comments. Please see our responses below.\n\n***Weakness 1:*** The analysis doesn't show the convergence of the standard autoregressive transformer.\n\n***Response to Weakness 1:*** We thank the reviewer's insightful comment. Indeed, our current theoretical construction employs a designed 4-layer transformer architecture, with specific layers activated only at certain CoT steps, to facilitate precise theoretical analysis and ensure provable convergence under teacher-forcing training. \n\nWe note that explicitly constructing a transformer architecture to demonstrate specific capabilities during ICL is a common approach in the theoretical study of transformers [1–4], and our methodology follows this line of work. Building on this, we further show that such a transformer can be explicitly obtained through CoT training, with a rigorous analytical characterization of the non-asymptotic training dynamics.\n\nAnalyzing the training dynamics of a general autoregressive transformer remains analytically intractable due to the inherent complexity and non-convexity of the activation functions. Even for single-layer transformers, tractable analysis often relies on simplified or specialized architectures. For instance, [5] investigates a single-layer linear transformer in the context of linear regression. Other works have similarly focused on linear transformers or restricted settings such as single-layer, single-head architectures with ReLU-activated attention for classification tasks [1,6], diverging from the standard autoregressive architecture.\n\nMoreover, we believe that starting with a general autoregressive transformer is not necessary to demonstrate our primary claim, i.e., unlabeled data can provably enhance in-context learning.\n\nFor these reasons, in this work, we focus on a specifically constructed 4-layer transformer architecture. Extending our theoretical analysis to the standard autoregressive transformer is an important direction we leave for future work.\n \n***Weakness 2:*** Need analysis of memory and latency.\n\n***Response to Weakness 2:*** We thank the reviewer for this insightful comment. During inference with $M$ unlabeled samples, our method requires approximately $\\mathcal{O}(M^{1/4})$ CoT steps.\nAt any given step $t$, the input sequence length is on the order of $\\mathcal{O}(N+M+M^{1/4})=\\mathcal{O}(N+M)$. Since previous steps do not need to be stored, the peak memory requirement remains $\\mathcal{O}(N+M)$. \n\nIn terms of latency, the time complexity for a single forward pass under a standard transformer scales in $\\mathcal{O}((N+M)^2)$ due to the self-attention mechanism. As established in Theorem 4.2, the time complexity under our method is $\\mathcal{O}(N^2M^{1/4}+NM^{5/4}+M^{9/4})$. This is because under our method, the inference time for each forward pass in $\\mathcal{O}((N+M+M^{1/4})^2)$, and we need in total $\\mathcal{O}(M^{1/4})$ passes. When $M\\gg N$, our time complexity becomes $\\mathcal{O}(M^{9/4})$ and that of a single forward pass for a standard transformer is $\\mathcal{O}(M^{2})$. We view this moderate polynomial overhead as a reasonable trade-off: Theorem 4.2 shows a non-asymptotic guarantee that the test error converges without bias. In contrast, running a standard transformer for one (or any fixed number of) passes leaves a persistent bias term in the accuracy, as shown in the concurrent work [7, Thm. 4]. \n\nWe thank the reviewer again for this valuable comment. We will incorporate this detailed complexity analysis into the revised manuscript.\n\n***Weakness 3:*** It's unclear how well the theory translates to and how well this approach is suited for real-world tasks.\n\n***Response to Weakness 3:*** We thank the reviewer for this important comment. We have performed an additional experiment on a real-world dataset to verify the impact of unlabeled samples on the ICL performance of a pre-trained LLM. Please refer to our responses to Reviewer RyNA (Response to Weakness 1) for the experimental setup and results.\n\n***Weakness 4:*** Figure 1 missed the unlabeled classifier oracle.\n\n***Response to Weakness 4:*** We thank the reviewer for this constructive feedback. \n\nOur primary theoretical contribution is to demonstrate that a transformer is capable of using unlabeled data to improve in-context learning. Our comparison to the standard Bayes oracle, which by definition cannot access unlabeled data, was intended to isolate and highlight this specific capability. \n\n\nWe agree that including a semi-supervised baseline (such as an EM-based oracle) is a valuable addition to help us understand how well a transformer can utilize unlabeled data. We will add such baselines to our experiments in the revised manuscript to provide a more comprehensive evaluation.\n\n\n\n***Weakness 5:*** Missing some related works. \n\n***Response to Weakness 5:*** We thank the reviewer for pointing out these important related works. \n\nWe note that Gupta et al. (2023) use unlabeled data in an **unsupervised** ICL setting, where the prompt only contains unlabeled data. In contrast, we utilize unlabeled data in addition to the labeled examples to augment ICL in a **semi-supervised learning** fashion. Our method also differs from Agarwal et al. (2024), which it proposes two methods for many-shot ICL, namely, 1) Reinforced ICL and 2) unsupervised ICL. For 1), it synthesizes new, fully-labeled examples by generating and selecting its own rationales. For 2), it simply inputs unlabeled demonstrations in the prompt. Both methods differ from our approach, where we mix both labeled and unlabeled examples in the prompt and theoretically analyze its performance.\n\nWe will include and compare with those references and other relevant papers in the revision. \n\nReferences:\n\n[1] Ahn, Kwangjun, et al. \"Transformers learn to implement preconditioned gradient descent for in-context learning.\" Advances in Neural Information Processing Systems 36 (2023)\n\n[2] Bai, Yu, et al. \"Transformers as statisticians: Provable in-context learning with in-context algorithm selection.\" Advances in neural information processing systems 36 (2023)\n\n[3] Von Oswald, Johannes, et al. \"Transformers learn in-context by gradient descent.\" International Conference on Machine Learning. PMLR, 2023.\n\n[4] Akyürek, Ekin, et al. \"What learning algorithm is in-context learning? investigations with linear models.\" arXiv preprint arXiv:2211.15661 (2022).\n\n[5] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. \"One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention.\" arXiv preprint arXiv:2307.03576 (2023).\n\n[6] Li, Hongkang, et al. \"How do nonlinear transformers learn and generalize in in-context learning?.\" arXiv preprint arXiv:2402.15607 (2024).\n\n[7] Li, Yingcong, et al. \"When and How Unlabeled Data Provably Improve In-Context Learning.\" arXiv preprint arXiv:2506.15329 (2025).\n\n----------\n\nWe hope our responses have addressed your concerns satisfactorily. We sincerely appreciate your constructive feedback and would be grateful if you would consider raising your evaluation based on our responses. We would also be happy to address any further questions you may have."}, {"review_id": "FhrtSCpSKP", "reviewer": "Reviewer_uxcu", "summary": "This paper proposes a novel augmented ICL framework, where the prompt contains some labeled examples as well as unlabeled inputs. This framework is motivated by scarcity of labeled data and high cost to get labels for the readily available unlabeled data. Considering the setting of multi-class linear classification, the paper first presents a construction showing that the transformer output (class mean estimates) converges to the ground truth as the number of chain-of-thought (CoT) steps is increased. This is accomplished by obtaining initial estimates from the labeled examples and refining them using the unlabeled inputs emulating the expectation-maximization (EM) algorithm. This leads to an excess risk bound of $O(1/\\sqrt{N \\text{poly}(M)})$. Next, it shows that such a transformer can be trained by teacher forcing using a specific initialization. The theoretical result is supported by numerical experiments.", "questions": "Other suggestions about the paper writing:\n\n- There are several typos which should be corrected, e.g., it should be ‘simplify’ in line 73, ‘transformers’ in line 93, ‘exists’ in line 193, ‘taking’ in line 221, ‘\\mathbf{\\mu}’ in line 232, ‘combining’ in line 248, to list some of them.\n- The proof sketches are helpful, but the initial discussion in lines 72-75 in the Introduction is not very clear and can be refined.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "The main strengths of the paper are:\n\n- It proposes a novel augmented ICL framework (apart from the concurrent work [1]).\n- It is the first theoretical study showing the benefit of unlabeled data for ICL in transformers.\n\nHowever, the paper has some important weaknesses that should be addressed. \n\n- As the paper mainly focuses on theory, it should clearly discuss the novelty or relation with prior work for the proof techniques for the main results. Specifically, this prior work [2] comparing transformers with EM for multi-class clustering seems quite relevant and should be discussed.\n- It is unclear how effective the proposed augmented ICL framework can be in practice. I strongly suggest that the paper should include experimental results for at least one setting when prompting a pretrained LLM with this framework.\n- The paper is missing references to several relevant related works on understanding ICL in transformers in Section 2, e.g., [3-6] to list a few. Only four papers are discussed, which does not include the seminal work Garg et al., 2022, which is strange. More relevant papers should be mentioned/discussed. Also, a minor comment about line 29: Fu et al., 2023 is not a seminal work on ICL in transformers and can be directly cited in Section 2 rather than the introduction.\n\nReferences:\n\n[1] Li et al., “When and How Unlabeled Data Provably Improve In-Context Learning”, 2025.\n\n[2] He et al., “Transformers versus the EM Algorithm in Multi-class Clustering”, 2025.\n\n[3] von Oswald et al., “Transformers learn in-context by gradient descent”, ICML 2023.\n\n[4] Akyurek et al., “What learning algorithm is in-context learning? investigations with linear models”, ICLR 2023.\n\n[5] Xie et al., “An explanation of in-context learning as implicit bayesian inference”, ICLR 2022.\n\n[6] Raventos et al., “Pretraining task diversity and the emergence of non-bayesian in-context learning for regression”, NeurIPS 2023.", "was_revised": true, "final_justification": "The rebuttal addressed my concerns, so I increased my score from 3 to 4.", "rebuttal": "We thank the reviewer for the careful reading and thoughtful comments. Please see our responses below.\n\n\n***Weakness 1:*** Lack of discussion of the novelty or relation with prior work for the proof techniques for the main results.\n\n***Response to Weakness 1:*** We appreciate the reviewer's insightful comment and for highlighting these highly relevant works. We clarify our theoretical contributions compared with [1] and [2] as follows.\n\nOur work significantly differs from [1] in both problem setting and theoretical guarantees/proof novelty. [1] primarily focuses on *unsupervised learning* of Gaussian mixture models, aiming to mimic a clustering algorithm (Lloyd's algorithm). Its analysis centers on showing that a transformer's output can approximate an iterative algorithm, which is susceptible to local optima, since for such a problem, the EM algorithm is prone to convergence at local optima and is highly sensitive to initialization. In contrast, our work addresses **semi-supervised learning** during ICL, where a small amount of labeled data is available. Our goal is to demonstrate how unlabeled data can effectively augment traditional ICL that typically relies solely on labeled data. This fundamental difference in setting leads to **distinct theoretical contributions**. Our key technical contribution is proving a non-asymptotic convergence to the *ground-truth means*, not just an algorithm's output. Our novel two-stage proof technique first shows how the transformer can leverage labeled data to learn a high-quality initialization for the means. Next, we prove that with this strong initialization, the transformer's EM-style updates on the unlabeled data lead to monotonous improvement and convergence to the true parameters. This method of using labeled data to bypass the local optima problem, which is common in unsupervised EM, is central to our theoretical framework.\n\nCompared with the concurrent work [2], which also aims to show that transformers can leverage unlabeled data in ICL, our theoretical contributions and technical analysis differ in several aspects:\n\n**Theoretical contributions:** **First**, [2] studies a linear transformer without non‑linear activations, whereas we analyze a more realistic architecture that includes the soft‑max attention mechanism. **Second**, their theory is restricted to a Gaussian‑mixture model with exactly two classes. Our framework allows an arbitrary (finite) number of classes, covering multi‑class scenarios. **Third**, Theorem 3 in [2] provides only an asymptotic guarantee as the number of unlabeled samples tends to infinity, and their finite‑sample bound includes a non‑vanishing bias term. By contrast, in our work, **we establish a rigorous, non-asymptotic convergence guarantee for a transformer with softmax non-linearity in a general multi-class setting**. \n\n**Technical novelty:** Our technical proofs rely on several novel ingredients. First, we formally characterize the initial condition required to show the non-asymptotic convergence. Then, and most critically, we tackle the challenge of analyzing the gradient dynamics over the unlabeled data. **A key innovation in our analysis is the strategic design of the temperature parameter $\\beta\\tau$.** We show how it can be set to ensure the model's attention weights are sharply focused on the most recent class mean estimate, effectively isolating the gradient updates at each step. By combining this with Bernstein's inequality, we derive a tight bound on the deviation between the true stochastic gradient and the gradient of the ideal expected log-likelihood. Finally, by leveraging the Lipschitz continuity of the loss landscape and integrating our novel gradient bound, we prove the finite-sample convergence of our model. \n\nOur analysis differs fundamentally from that of [2] in several ways. First, we analyze a multi-layer transformer with standard softmax attention, a more complex and general setup. In contrast, the theory in [2] is restricted to transformers with the simpler linear attention mechanism. Second, our construction enables the transformer to perform EM-style updates on the unlabeled data. The model in [2], however, is shown to implement a different learning rule: a semi-supervised plug-in estimator. As a result of these differences, our analysis establishes a more promising convergence property for the model compared to the results shown in [2].\n\nWe'll incorporate this detailed comparison into the revised manuscript to clarify our theoretical contributions.\n\n\n***Weakness 2:*** It is unclear how effective the proposed augmented ICL framework can be in practice.\n\n***Response to Weakness 2:*** We have performed an additional experiment on a real-world dataset to verify the impact of unlabeled samples on the ICL performance of a pre-trained LLM. Please refer to our responses to Reviewer RyNA (Response to Weakness 1) for the experimental setup and results.\n\n\n***Weakness 3:*** Missing references to several relevant related works in Section 2.\n\n***Response to Weakness 3:*** We thank the reviewer for pointing out these important related works. \n\nWe will thoroughly revise Section 2 based on your suggestion. We will add and discuss the seminal work by Garg et al. (2022), the additional references [1-6], and other relevant papers. Furthermore, we will move the citation for Fu et al. (2023) from the introduction to the related works section as advised. \n\n\n\n***Question 1:*** Typos need to be fixed.\n\n***Response to Question 1:***  We thank the reviewer for their careful reading and for pointing out these typos. We will correct all the typos mentioned and perform a thorough proofreading to eliminate any remaining typos.\n\n***Question 2:*** Initial discussion in lines 72-75 is unclear.\n\n\n***Response to Question 2:*** We appreciate the reviewer’s helpful comment. We will revise the parts in the introduction you mentioned as follows:\n\n*Our proof introduces a novel decomposition of the gradient of the CoT training loss by applying Stein's lemma, separating it into two analytically tractable components: a posterior-difference term and a Jacobian-difference term. Leveraging the inherent isotropy preserved during gradient descent, we show that the posterior-difference term vanishes, while symmetry properties simplify the Jacobian-difference term, enabling us to derive a tight upper bound on the key inner-product expression and establish the desired sublinear convergence rate.*\n\nReferences:\n\n[1] He et al., \"Transformers versus the EM Algorithm in Multi-class Clustering\", 2025.\n\n[2] Li et al., \"When and How Unlabeled Data Provably Improve In-Context Learning\", 2025.\n\n[3] von Oswald et al., “Transformers learn in-context by gradient descent”, ICML 2023.\n\n[4] Akyurek et al., “What learning algorithm is in-context learning? investigations with linear models”, ICLR 2023.\n\n[5] Xie et al., “An explanation of in-context learning as implicit Bayesian inference”, ICLR 2022.\n\n[6] Raventos et al., “Pretraining task diversity and the emergence of non-bayesian in-context learning for regression”, NeurIPS 2023.\n\n----------\n\nWe hope our responses have addressed your concerns satisfactorily. We sincerely appreciate your constructive feedback and would be grateful if you would consider raising your evaluation based on our responses. We would also be happy to address any further questions you may have."}]

Author rebuttal:

Meta-review: This paper proposes a novel augmented ICL framework that integrates labeled examples and unlabeled data into prompts, demonstrating theoretically and empirically that transformers can emulate the EM algorithm via CoT prompting to enhance ICL performance in multi-class linear classification. This is the first theoretical study on unlabeled data’s impact on transformer-based ICL.

Although the reviewers with highest ratings have the lowest confidence scores, all reviewers give positive final ratings. After reading the paper, reviews and discussions, I recognize the contribution and technical soundness of this paper. During rebuttal, the authors added experiments on real-world dataset with Qwen3-8B, clarified extensions to generalized linear models and experimental parameter choices, added complexity analysis to justify overhead. Given the theoretical novelty, rigorous analysis, and empirical validation (including real-world data), along with the authors’ efforts to address the reviewers' concerns, the paper is recommended for acceptance.
Decision: Accept (poster)