Theoretical Insights into In-context Learning with Unlabeled Data

We thank the reviewer for the detailed feedback. Below, we provide our responses to the comments.

W1: Limited to Model structure: theoretical results focus only on linear attention models. While the paper does include empirical comparison to softmax-based Transformers, no theoretical treatment is provided for these more widely used architectures.

Response: We kindly refer the reviewer to our responses to W1 and W3 of Reviewer 9Hjq, and W1 of Reviewer GrWH for detailed discussions addressing this concern.

W2: Assumption of Isotropic Priors: The analysis assumes task means are uniformly sampled from the unit sphere, which may not reflect realistic downstream data distributions.

Response: We use the isotropic GMM to ensure theoretical tractability and to investigate the effects of model depth and unlabeled data in a controlled setting. While it may not capture the full complexity of real-world data, it serves as a principled and analytically tractable starting point for studying algorithms designed to improve semi-supervised performance.

Building on the core conclusion of our work that deeper or looped models enhance semi-supervised learning, we proposed the Looping TabPFN algorithm, which effectively leverages unlabeled data to improve performance in practical settings. Thus, while our theoretical framework assumes isotropic priors for analytical clarity, the key insights extend to broader scenarios. We agree that further exploration under more realistic or structured priors is valuable and constitutes a promising direction for future work.

W3: While effective, the looping approach introduces additional compute overhead, and the paper does not explore how to optimize this trade-off.

Response: We agree that the trade-off between accuracy and computational cost is an important consideration. While this work does not focus primarily on optimizing that trade-off, both our synthetic and TabPFN experiments show that only a small number of iterations (typically fewer than 5) are sufficient to achieve the best performance. In fact, the first few iterations contribute the most to performance gains, suggesting that the additional compute overhead is relatively modest given the improvement achieved.

We also refer the reviewer to our response to Question of Reviewer 9Hjq, where we further discuss the effectiveness and practicality of the looping strategy.

W4: Line 296-298 How to decide α is confused. does it need require domain knowledge or validation data to set optimally?

Response: In our SSPI algorithm (Lines 246–247), the parameter $\alpha$ is used under the assumption that its optimal value is known. In our experiments, $\alpha$ is selected based on domain knowledge of the data distribution. However, determining the optimal value of $\alpha$ in practice given only hyperparameters such as the distribution, $d$, $n$, and $p$ is a fundamental mathematical problem. We view this as an important direction for future theoretical investigation.

Q1: What is w.p in eq(1).

Response: "w.p." means "with probability".

Q2: Eq (3): is Z the demonstration example or instruction without query x or with query x but without x’s label?

Response: In Eq (3), $Z$ includes the query $x$ (as the last token in $Z$) but without $x$’s label. This setting is consistent with prior work in the linear attention literature (Ahn et al.(2023), Mahankali et al.(2024), Zhang et al.(2023), Li et al.(2024)).

Q3: Iine 138, why input of layer l only contains attention part without ffn part?

Q4: Line 142, I do not understand this question? The prediction of y is decided by res(attnZlast)+ffn(attnZlast+res(Zlast-1)) according to transformer block.

Response (Q3&Q4): In this work, our goal is to analyze the optimization behavior of the attention mechanism, and FFN will introduce additional nonlinearity and complicate theoretical analysis. To that end, we focus only on the attention block, excluding the FFN component and taking the output of the final attention layer as the model output for prediction. This setting is consistent with prior work on multi-layer linear attention (Ahn et al.(2023), Li et al.(2025)). We agree that studying the FFN is also important. As shown in our experiments (Fig. 2(c)), the full Transformer model (attention + MLP) achieves the best performance. However, we believe that a focused study on the attention mechanism, and in particular, the role of depth in attention, remains underexplored and is a necessary step toward a deeper understanding of the full Transformer architecture.

Notably, this restriction applies only to our theoretical study. In our Transformer (Fig 2c, cyan curve) and TabPFN experiments, we do not impose any architectural constraints. This demonstrates that the theory-inspired looping strategy remains effective in practice, even when applied to more complex models.

Q5:How sensitive is the performance of the looping-based TabPFN method to the accuracy of initial pseudo-labels? Would noisy pseudo-labels mislead the iteration? Does the portion between label, unlabel, psedu-label sensitive?

Response: Good point. To evaluate this, we examined the performance of TabPFN under varying levels of initial pseudo-label accuracy. The results, shown below, confirm that the quality of pseudo-labeled data plays an important role and higher-quality pseudo-labels lead to greater performance improvements. (Note: Due to time constraints, we conducted the experiment using a single data setting with a fixed labeled set. We randomly sample 100 groups of pseudo-labeled data and report the average relative improvement.)

Regarding the proportion of labeled, unlabeled, and pseudo-labeled data, our synthetic experiments (e.g., Fig. 1) show that increasing both labeled and unlabeled data consistently improves prediction performance. Motivated by the reviewer's question, we believe that more effective looping strategies could be developed. For example, by gradually assigning pseudo-labels to high-confidence unlabeled examples rather than labeling the entire unlabeled set. Such adaptive approaches may further enhance the performance of the looping method and represent a promising direction for future exploration.

Q6: How should one determine the optimal number of loops in practice? Does it correlate with data properties like noise level or number of unlabeled samples?

Response: In our TabPFN experiments, we implement a validation risk criterion to monitor the variance of assigned pseudo-labels (see Lines 710–715 in the Supplementary Material), computed as $Validation~Risk=\frac{1}{n}\sum_i \min(|y^{pseudo}_i-1|,|y^{pseudo}_i+1|).$ It is used to determine when to stop looping.

We also evaluated the optimal number of looping iterations with respect to different levels of initial accuracy, and observed an inverse correlation: lower initial accuracy tends to require more loops. However, due to the limited number of datasets and looping iterations, we are cautious about making a definitive conclusion. A more rigorous investigation is needed to establish this connection.

Additionally, we emphasize that based on both our synthetic results (Fig. 1) and TabPFN experiments, the number of loops (or layers) required for strong performance is relatively small (typically fewer than 5) with the initial iterations/layers contributing the most significant improvements.

After reviewing the comments, we suspect that the reviewer may have inadvertently submitted a review targeted at a different paper (also on the topic of ICL), as the content does not appear to relate to our submission. As a result, we were unable to prepare a rebuttal.

If possible, we kindly ask the reviewer to share any key concerns specific to our paper so that we can address them during the discussion period. We look forward to your feedback.

We thank the reviewer for the insightful comments and suggestions, which we found valuable and have incorporated to improve the clarity and quality of the manuscript.

W1: The core theoretical results are derived for linear attention models.

Response: We kindly refer the reviewer to our responses to W1 and W3 of Reviewer 9Hjq as well as our response to Q3 below (where we establish a connection between linear and softmax attention) for a detailed discussion addressing this concern.

W2: The analysis relies on an isotropic Gaussian Mixture Model where the task mean is sampled from a uniform sphere. This assumption is crucial for the argument that single-layer attention fails because the expected covariance $\mathbb{E}[X^\top X/n]$ becomes uninformative. It's unclear how the results would change under more structured or anisotropic data priors.

Response: We provide a detailed response and additional experimental results in our response to Q2 below.

Q1: Role of Softmax and MLPs: Your experiments show that a full transformer outperforms the linear attention models. What is your hypothesis on the specific roles that softmax attention and MLP layers play in improving semi-supervised performance? Do they allow the model to learn a better polynomial estimator of $X^\top X$ or do they enable a different, more powerful learning mechanism altogether?

Response: Great question. We believe the improved performance of the full Transformer model comes from mechanisms beyond simply learning higher-order polynomials of $X^\top X$.

1. As noted in our Lemma 1, a 5-layer linear attention model can already express polynomial functions of $X^\top X$ up to degree 120. Yet, full Transformers still outperform it (see Fig. 2(c)), suggesting that their advantage is not merely due to increased expressive power over $X^\top X$.
2. Attention mechanisms compute output based on pairwise similarities, and given input $X$, they operate as $att(X)=f(X^\top X)$. In contrast, MLPs act directly on the input $MLP(X)=f’(X)$, and therefore do not improve the polynomial estimator of $X^\top X$, but instead provide more flexible and nonlinear feature transformations.

Therefore, we hypothesize that the MLP layers contribute through a more complex, nonlinear data projection mechanism that complements the attention module by capturing feature-level representations beyond pairwise relationships. Understanding the exact theoretical role of MLPs in this context remains an open and promising direction for future work.

Q2: Beyond Isotropic Priors: Your proof for the failure of 1-layer attention hinges on the isotropic task prior, which makes $\mathbb{E}[X^\top X/n]$ uninformative. How would your results change for an anisotropic prior (e.g., $\mu\sim \mathcal{N}(0,\Sigma))$? Could a single-layer model then leverage unlabeled data, or would depth still be required?

Response: We thank the reviewer for this insightful question. First, when the task prior is changed to $\mu\sim\mathcal{N}(0,\Sigma)$, single-layer linear attention still fails to benefit from unlabeled data. This is because attention computes pairwise similarity as $x_i^\top W x_j$ where $W=W_kW_q^\top$. When $\mu$ has covariance $\Sigma$, the input data satisfies $\mathbb{E}[xx^\top]=\Sigma+\sigma^2 I:=\bar\Sigma$. This can be transformed into the equivalent isotropic case by: $x_i^\top W x_j=\bar x_i \bar\Sigma^{1/2}W\bar\Sigma^{1/2}\bar x_j^\top,$ where $\bar x$ has identity (up to a scalar) covariance. Thus, the anisotropic setting effectively reduces to the isotropic case under a linear transformation, and the conclusion that single-layer attention fails to leverage unlabeled data still holds. The first row of the table below supports this where we fix $d=m=10$ and generate non-identity covariance using $\Sigma=A^\top A/trace(A^\top A)$ with a random matrix $A$. Results show that performance remains unchanged with increasing numbers of unlabeled samples.

Motivated by the reviewer’s comment, we further explored when single-layer models can benefit from unlabeled data. From our analysis in the proof of Lemma 2 (Supplementary Material), the prediction from a single-layer model depends on $x^\top X^\top X$ (assuming $W=I$). If $\mathbb{E}[\mu]=0$, then $\mathbb{E}[x^\top X^\top X]=0$, and the unlabeled data does not help. However, if we use a prior with nonzero mean (e.g., $\mu\sim\mathcal{N}(\nu,I)$ where $\nu\neq 0$), the model can extract useful signal from $\mathbb{E}[x^\top X^\top X]\neq 0$. The second row in the table shows that under this setting, increasing unlabeled data improves single-layer model performance: (we set $d=m=10$)

This is a valuable discussion, and we will incorporate these findings into our revised manuscript.

Q3: Connection to Pseudo-Labeling: You compellingly argue that the label propagation process resembles iterative pseudo-labeling. Could you make this connection more formal?

Response: As discussed in Proposition 1 (Label Propagation) and the discussion that follows (Lines 216–221), multi-layer linear attention can be viewed as an EM-like algorithm, where each layer corresponds to one iteration of pseudo-label refinement. Here, we provide a formal discussion as suggested by the reviewer, and we also refer to our response to W3 of Reviewer 9Hjq for further details. We will incorporate it into our revised manuscript.

Specifically, for attention mechanism (subsuming both linear and softmax), at the $\ell$th layer, the pseudo-label assigned to the $i$th example is updated as follows $y_i^\ell\to y_i^\ell+c_\ell\sum_js_{ij}y_j^\ell$ where $s_{ij}$ is the weighting strategy that determines how much different tokens influence the token $i$. For linear attention, $s_{ij}=x_i^\top W_\ell x_j$; for softmax attention, $s_{ij}=e^{x_i^\top W_\ell x_j}/\sum_k e^{x_i^\top W_\ell x_k}$. This form resembles the update step in EM/label propagation algorithms (but with different update rules for different models), where the model iteratively updates predictions based on current pseudo-labels.

We thank the reviewer for recognizing our theoretical contributions. We address the comments below.

W1: The analyzed attention module is simplified as w/o softmax version (eq. (3)), differing from mainstream model settings. This may limit the application of the theory.

Response: In this work, our goal is to analyze the exact optimization behavior of attention mechanisms. Note that 1) attention captures token dependencies through pairwise similarities; and 2) even linear attention induces a non-convex landscape. Therefore, its theoretical analysis is already a challenging and meaningful problem. Prior works (Ahn et al.(2023), Mahankali et al.(2024), Li et al.(2024)) and ours begin with linear attention as a tractable starting point. In contrast, our work considers the semi-supervised setting and examines how model depth affects prediction performance. To the best of our knowledge, this perspective has not been theoretically explored.

While our analysis does not directly cover softmax attention, we empirically compare linear attention, softmax attention, and full Transformers (softmax attention + MLP) in Fig. 2(c). Under our GMM setting, softmax attention without the nonlinear MLP does not outperform linear attention. We provide further discussion and explanation in our response to W3, where we describe how attention mechanisms can generally be formalized as a pseudo-labeling mechanism, applicable to both softmax and linear. However emphasis is given to linear attention as we can establish concrete and nontrivial theoretical guarantees.

Lastly, although our theoretical results are based on linear attention, they motivate the idea of looping, which we successfully validate in realistic scenarios using real data and TabPFN. These results demonstrate that the theoretical insights carry over and provide practical benefits in more complex settings.

W2: The experimented dimension is only set to 10 (Line 285). Why not try a high dimensionality?

Response: First, since the GMM is a standard and generalizable data model, the dimensionality does not affect our theoretical results. To clearly validate our theoretical findings while keeping the computational cost low, we chose a representative dimension of $d = 10$. Notably, even with such a low dimension, it requires up to 10,000 in-context samples to empirically approach Theorem 2 (see Fig. 1(c), where the green curves approach the black dotted line).

To further address the reviewer's concern, we also conducted additional experiments at higher dimensions (presented in the table). The results are consistent with those in Fig. 1, where we observe that 1-layer=SPI=Theorem 1, and SSPI-1<2-layer<5-layer$\approx$SSPI-$\infty$. (Note: due to time constraints, each setting is evaluated based on a single run, which may introduce variance.)

In addition, for the Looping TabPFN experiments, the model employs an embedding dimension of 192.

W3: There is a gap between the discussed task setting and the real datasets. For example, it's been widely proven that the attention with softmax outperforms the attention without softmax. However, Figure 2(c) reaches an opposite conclusion. Therefore, the content of this paper may not directly apply to standard Transformer-structured models.

Response: We thank the reviewer for raising this important point. While linear and softmax attention differ in their operational functions, they share similarities from a label propagation perspective due to their attention-based structure, and both can be interpreted as a pseudo-labeling algorithm.

A generic attention mechanism with suitable value weights, regardless of linear or softmax, updates the labels via: $y_i^\ell\to y_i^\ell+c_\ell\sum_js_{ij}y_j^\ell.$ Here, $s$ is the attention map: $s_{ij}=x_i^\top W_\ell x_j$ for linear attention, and is the softmax similarity between $i$th and $j$th token for softmax attention, both determined by the attention weights at layer $\ell$, where $W_\ell=W_{k\ell}W_{q\ell}^\top$. From this, it can be seen that the first attention layer will impute the missing labels using labeled data based on attention similarities between labeled and unlabeled features. The future layers act as a label propagation mechanism that updates the labels by creating their data-dependent mixtures. We will revise the paper with this exposition to highlight how attention generically implements pseudo-labeling.

We agree that softmax attention is often better than linear attention for real datasets but it doesn’t have to be true for all datasets. Specifically, for the supervised GMM setting, a single linear attention emulates the Bayes optimal estimator as proven in the paper. This mostly explains why linear attention outperforms softmax attention in Fig 2(c). Although softmax attention underperforms linear attention in isolation, the full Transformer model (softmax attention + MLP) outperforms all other models. (See Fig. 2c).

We also emphasize the following

1. Linear attention is widely used in theoretical study of transformers (see Ahn et al.(2023), VonOswaldetal.(2023), Mahankali et al.(2024)) and also underlies state-of-the-art architectures like gated-linear attention and Mamba2 (Li et al.(2024), Li et al.(2025)).
2. GMM is a classical and established dataset model for the study of classification problems/algorithms.
3. The benefit of increased depth for both attention types is evident, supporting our main theoretical claim that depth enhances semi-supervised performance.
4. Our theory highlights the utility of looping, a mechanism we validate in real-world settings using TabPFN. These results demonstrate that the insights from our theory translate into practical benefits beyond the simplified setup.

We will include a subsection in our revised manuscript to incorporate this discussion.

Q: Does looping TabPFN-v2 mean that using the output of one model run as input to the next run? I have a concern about the reusage of parameters. I believe this will deteriorate the performance as compared to running the same number but with different layers.

Response: Model looping is not a new concept and it has been widely explored and applied in various contexts (see citations below). It offers several advantages: it enables smaller models to handle more complex tasks, allows models to self-correct and refine outputs over multiple passes, and supports adaptability to varying input complexities by adjusting the number of iterations.

The primary benefit of looping is its flexibility and efficiency: In our TabPFN experiments, we did not retrain or redesign a deeper model. Instead, we used the same pretrained model and applied it iteratively in-context. This looping strategy improved performance without requiring additional model capacity or retraining from scratch. At a high-level, an autoregressive LLM can also be viewed as text generation by looping the same parameters. That said, we agree with the reviewer that our results can be enhanced by either (1) training a deeper architecture or (2) optimizing looped TabPFN for semisupervised learning.

• Mohtashami, Amirkeivan, et al. "Cotformer: A chain-of-thought driven architecture with budget-adaptive computation cost at inference." arXiv preprint arXiv:2310.10845 (2023).

• Giannou, Angeliki, et al. "Looped transformers as programmable computers." International Conference on Machine Learning. PMLR, 2023.

• Yu, Qifan, et al. "Enhancing auto-regressive chain-of-thought through loop-aligned reasoning." arXiv preprint arXiv:2502.08482 (2025).