Paper: Explaining Emergent In-Context Learning as Kernel Regression (ICLR-2024)

Reviewer comment (Theory):
Another reason why ICL does not implement KR directly are the label flipping experiments of Wei et al. (2023) and Kossen et al. (2023). They show that LLM performance generally suffers when 'flipping' the label relationship, e.g. changing labels from positive/negative to negative/positive for a sentiment analysis task. However, the kernel regression equation would predict that the performance of ICL does not depend on what tokens are used to encode the labels. I would appreciate it if the authors would add these observations to their discussion section.

Author response:
Thanks for pointing to this interesting phenomenon under label flipping. We see it as a great discussion to improve clarity. In the revised Section 4.2, we explain that Theorem 1 implies that the final prediction depends both on (1) in-context examples and (2) prior knowledge from pre-training. This can be seen from the $\eta^2 \epsilon_\theta$ term in Equation (6), which comes from the pre-training information in Equation (16-18) in Appendix A. Specifically, this term describes that the bias given by the pre-training distribution affects the ICL prediction, even if labels are flipped. This gives a partial elucidation of why ICL still depends on what tokens are because they might conflict with the bias in pre-training corpora. This will also align with Kossen et al.’s finding that “ICL cannot overcome prediction preferences from pre-training.”

Reviewer comment (Writing):
Relatedly, I cannot follow the authors when they say that kernel regression predicts that 'ICL performance relies on the label format'. I would appreciate if you could specify which of the experiments of Min et al. (2022b) you are referring to here. If it is their experiments in S. 5.2 which they call 'removing the format' (Impact of Input-Label Pairing) I would suggest you rephrase the paragraph here. In their experiments, Min et al. completely remove either the labels or demonstrations from the input. Such input is entirely incompatible with KR.

Author response:
The referred-to part is Section 5.2, “Impact of the label space” in Min et al., where they substitute labels with a random set of English words. As responded in (B1), using a label space different from the pre-training distribution will make the KR term diverge from or even conflict with the pre-training bias, thus influencing the final performance.

Reviewer comment (Theory):
I agree that KR provides an explanation for why similar samples help and why OOD samples (with high distance in input space to the test query) may be harmful. However, I think it is highly misleading to claim that Min et al. (2022b) show that 'out-of-distribution (OOD) demonstrative samples will degrade ICL accuracy'. It is important to note that what they call 'OOD' does not conform to standard expectations. In their OOD experiments they construct a set of input features that are 'randomly sampled from an external corpus', i.e. they completely destroy the input label relationship. This is different from classical OOD, e.g. samples from the same p(X, Y) but for rare or unusual X.

Author response:
Thanks for suggesting a correction, and we agree that a better experiment on evaluating the effect of ODD examples is beneficial. We conduct an experiment by converting test inputs to semantically similar but rarer sentences. Specifically, on the SST2 dataset, we prompt GPT-3.5-turbo to generate semantically similar while rarer (i.e. OOD) expressions of inputs and use them to substitute the original dataset inputs. We consider 3 types of OOD types. “Rare word” is where words are substituted with rare synonyms. “Complex” is to express the original sentence in a more complex structure. “Typo” is where we require adding typos to the original input. The results are listed below. We see that compared with original inputs (“None” category), these OOD types more or less decrease the ICL accuracy, with “typo” having the largest effect, dropping the accuracy to near random. This accords with our insight in Section 4.2 that OOD samples can negatively affect ICL performance. We add these results in the revision.

| OOD-type | None | rare word | complex | typo |
| --- | --- | --- | --- | --- |
| accuracy | 0.805 | 0.677 | 0.788 | 0.534 |

Reviewer comment (Theory):
In their 'Remaining Challenges' section, the authors refer to Min et al. (2022b)'s claim of 'LLM’s robustness to perturbed or random labels'. Luckily for the authors, it seems these claims do not always hold up to scrutiny (which is good for analogies between ICL and KR), see Kossen et al. (2023).

Author response:
We thank the reviewer for referring to this valuable evidence and will add this to the discussion.

Reviewer comment (Experiments):
The MNLI ICL performance numbers you report in Table 2 are not better than random guessing performance (~30%). Yet one of the attention head KR variants obtains performance much better than random guessing (~97%). If ICL performance is not good overall, I would not expect it to be outperformed by the 'head KR' variant. This makes me suspicious that there is some sort of overfitting going on here. I.e. the space of attention values is large (num_layers x num_heads) and one of them just happens to predict the true labels well. Do you have train/test splits for selecting and evaluating the KR head variant? Do you evaluate this over the full datasets or just a subset?

Author response:
The reviewer raises an interesting concern about the performance of the MNLI dataset. We follow the reviewer’s advice to eliminate the overfitting problem. As some MNLI does not provide test labels, and some tasks only have ~1000 pieces of data, we uniformly use 700 data in each validation set to represent tasks in a balanced way. We select the attention head that has the highest correlation with model predictions. It is then evaluated on model prediction reconstruction and task performance on a held-out set of 300 data per task. We see that the single-head reconstruction has an accuracy from 68\% to 80\% except on MNLI which is harder. The task performance of reconstructed outputs matches the model’s performance level. We revised the paper and merged Tables 1 and 2 accordingly.

| metric | sst2 | mnli | rotten-tomatoes | tweet_eval (hate) | tweet_eval (irony) | tweet_eval (offensive) |
| ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- |
| prediction reconstruction | 0.797 | 0.597 | 0.683 | 0.713 | 0.722 | 0.753 |
| task performance | 0.750 | 0.360 | 0.527 | 0.587 | 0.522 | 0.513 |

Reviewer comment (Experiments):
It would be great to see experiments with additional models, e.g. LLaMA/Falcon/Mistral models, to show your experimental observations do not depend on the GPT-j architecture.

Author response:
We went on to experiment on the LLaMa model. Results generally show similar trends and are added to Appendix D.

Reviewer comment (Theory):
I am not too familiar with proofs around KR and Bayesian inference. However, could you elaborate where Theorem 1 depends on examples being presented 'in-context'? It seems to me that, perhaps, Theorem 1 could hold true for all inference problems?

Author response:
The connection between Theorem 1 and “in-context” examples is Equation (19-24), which describes the effect of “in-context” examples $[\mathbf{x}_i, y_i]$ on the conditional probability of the test input. This is a unique challenge in LLM ICL, where demonstrative examples are combined in the same sequence as the testing input. In other words, LLMs are fed with the whole sequence in a uniform manner, but we provide little information on how the ICL sequence is constructed.

Reviewer comment (Theory):
The connection between kernel regression and Transformer attention is currently quite hand-wavy, with one paragraph pointing at similarities in the equation. Perhaps the biggest difference is the change from a kernel to the $\exp(q, k)$ from the attention. The authors say this "can be regarded as a kernel trick". I think that more discussion here would strengthen the submission. How do the differences here change prediction behavior? Does softmax attention imply a particular form for the kernel?

Author response:
The reviewer points to a valuable place for clarification. As $e^{q^\top k}$ (i.e., the Gaussian kernel or exponential kernel) is positive definite, it satisfies Mercer’s condition (https://en.wikipedia.org/wiki/Mercer's_theorem#Mercer's_condition), and Mercer’s theorem guarantees the existence of a kernel function. In case the kernel might be of infinite dimension (as is also the case for many other fields), a feature function that approximates the real kernel reasonably well is usually regarded as a “kernel trick”. We revised and added this part to the discussion.

Reviewer comment (Writing):
What is $T_x$ in Equation 4? I don't see a definition for it in the main text. Later you write 'semantic vectors of sample inputs', but that does not help. Is it related to the transition matrix?

Author response:
$T_{\mathbf{x}} = \prod_{i=0}^{l-1} \text{diag}(\mathbf{p}_{x_i})T $ is defined in the paragraph under Equation (3). It is indeed related to the transition matrix and can be used to compute the conditional probability after a sequence of observations $\mathbf{x}$.

Reviewer comment (Theory):
Figure 2: We observe that attention is high around the labels of the examples, as well as the input features of the test query. You write that 'This conforms to the intuition in kernel regression explanation in Theorem 1 that the inference on in-context learning prompts is a weighted average over sample labels.' However, looking at kernel regression, I would have expected the attention to compare _different input features_, i.e. the kernel is $K(x, x_i)$ and not $K(x, y)$. I would appreciate if you could comment on this.

Author response:
This is another interesting question. First, your intuition is right in that the attention “compares different input features”. This does not conflict with the current attention pattern because the label position contains information from the whole example $[\mathbf{x}_i, y_i]$ instead of only information $y_i$, as investigated in Section 5.4. Second, the attention weights need to be centered at label positions because it is the only place that $v$ access the label information $y_i$ because of the causal mask.

About the second question on how representations of tokens $\mathbf{x}_i$ can be computed in initial layers, this question requires describing the mathematical structure of linguistic data, which is an open challenge due to the lack of a suitable math framework. This is also the reason why many previous papers, such as Akyürek et al. (https://arxiv.org/abs/2211.15661), only study length-1 inputs.

Reviewer comment (Presentation):
It would be interesting to see Figure 2 for individual examples.

Author response:
We plot the first example in the dataset for Figure 2 in Appendix E. It shows a similar but slightly sparser pattern to Section 5.1.

Reviewer comment (Experiments):
"This is similar to our prediction in Section 4.1: not many attention layers are needed for kernel regression, as long as the required features have been computed by preceding layers" --> That makes sense. However, how many layers you need to compute adequate features could vary a lot between tasks. Therefore, the KR hypothesis does not allow you to make any hypothesis about the number of layers needed for ICL. Maybe if you had a toy experiment where the embeddings were perfect already?

Author response:
We agree with the reviewer in this regard, in the sense that the theory is only an existence statement but not constructive, so in practice, how many layers of features are needed should depend on the specific data/task in hand.

Reviewer comment (Presentation):
Please position the captions for tables above the table as per style guide.

Author response:
Thank you for pointing out the typos, and we revise the submission accordingly.

Reviewer comment (Writing):
Please cite GPT-j instead of just providing a link. The authors prefer the citation given at https://github.com/kingoflolz/mesh-transformer-jax.

Author response:
Thank you for pointing out the typos, and we revise the submission accordingly.

Reviewer comment (Experiments):
"Specifically, for each attention head, we use the maximal attention value ai within the range of [xi, yi] as the kernel weight" --> Interesting, based on your previous motivation I would have thought you would just take the attention value at y_i? Why maximal value? If this is an ad-hoc choice, make that clear. How well does average work?

Author response:
The reviewer raises another interesting question. This technical choice of using max attention was from the intuitive observation in Figure 2, that only a few high-value attention values exist in the attention map. We also tried using mean or sum in place of the maximum attention and found it to have little effect on Figures 3 and 4. This is probably due to the fact that the majority of attention weights are concentrated in the few high-value positions in each sample.

Reviewer comment (Writing):
The penultimate sentence in the discussion on Sensitivity to Label Format may contain grammatical errors.

Author response:
Thanks for pointing out the typos. We corrected it in the submission.

Reviewer comment (Theory):
The paper does not address how the Kernel Regression hypothesis explains the influence of pre-trained models on the similarity between $x_i$ and $x_{test}$, as well as the label space of $y_i$; specifically, in Section 4.2, this point is not discussed, and it is unclear whether the hypothesis can provide insights into this aspect.

Author response:
We agree that this is an important point where our explanation can provide insights. Intuitively, Theorem 1 implies a prediction output on $\mathbf{x}_{test}$ that is closer to labels $y_i$ of similar $\mathbf{x}_i$, and such similarity is related to how similar their conditional probabilities $P(\cdot | \mathbf{x})$are in the pre-trained LLM.

To be specific, in Equation (4) the similarity kernel measures similarity between $\mathbf{x}$ and $\mathbf{x}’$ on the space of $\text{vec}(T_{\mathbf{x}})$ and $\text{vec}(T_{\mathbf{x}’})$. This is a flattened vector of the matrix $T_{\mathbf{x}}$, which defines the “belief state” in HMMs. The belief state determines the conditional probability $P(\cdot | \mathbf{x})$ after prefix $\mathbf{x}$. This is exactly the pre-training objective and functionality of LLMs. Therefore, if $x_i$ and $x_{test}$ have similar conditional probabilities (i.e., similar follow-up generations), their similarity kernel will have a larger value and the predicted label will be more similar to $y_i$, and vice versa. We add this part to Section 4.1.

Reviewer comment (Theory):
The authors' primary assumption in Section 2.1 states that “as the number of samples n increase ... converges to a kernel-regression for”. The reviewer thinks this is a very strong assumption since this begs the question: How does ICL work so well for only 2-3 examples, especially for large models?

Author response:
We thank the reviewer for the time and review and are pleased to see that they like the insights in Section 4.2. Regarding the comments and questions:

**W1 & Q1: Understanding the main theorem**

The reviewer raises an interesting concern. If we understand correctly, the statement appears in Section 4.1 instead of Section 2.1. Besides, it is an explanation of the proven Theorem 1 results instead of an assumption. As for the fact that LLMs work well for only 2\~3 examples and saturate after 10\~15 examples, we would like to point out that this phenomenon is a different problem from Theorem 1. The mentioned phenomenon is about “how ICL outputs converge to ground truth,” while Theorem 1 studies “how the kernel-regression interpretation converges to ICL outputs”. Therefore, it does not conflict with Theorem 1.

Reviewer comment (Theory):
The authors' primary assumption in Section 2.1 states that “as the number of samples n increase ... converges to a kernel-regression for”. The reviewer thinks this is a very strong assumption since this begs the question: Why does ICL performance saturate usually after 10-15 examples?

Author response:
We thank the reviewer for the time and review and are pleased to see that they like the insights in Section 4.2. Regarding the comments and questions:

**W1 & Q1: Understanding the main theorem**

The reviewer raises an interesting concern. If we understand correctly, the statement appears in Section 4.1 instead of Section 2.1. Besides, it is an explanation of the proven Theorem 1 results instead of an assumption. As for the fact that LLMs work well for only 2\~3 examples and saturate after 10\~15 examples, we would like to point out that this phenomenon is a different problem from Theorem 1. The mentioned phenomenon is about “how ICL outputs converge to ground truth,” while Theorem 1 studies “how the kernel-regression interpretation converges to ICL outputs”. Therefore, it does not conflict with Theorem 1.

Reviewer comment (Experiments):
Connection with OOM formalism is intriguing but not explored in what seems like it's full quantitative potential to provide practical guidance for LLMs pretraining or prompt engineering.

Author response:
We agree with the reviewer other algorithms in OOMs could potentially be useful in partially explaining LLMs. Still, in this paper, we only need the forward prediction algorithm during the inference stage (because the setting of ICL is inference only, no learning involved), so many techniques in OOMs regarding the learning stage will not be used. Furthermore, we kindly note that OOMs share many similarities with HMMs, so the analysis results by using OOMs should be similar to the current status.

Reviewer comment (Theory):
As mentioned also in the paper, the kernel regression formalism does not capture some crucial phenomenology observed in in-context learning, like the sensitivity to in-context samples order, and there are no indications provided in the paper of how this shortcoming in the formalism could be addressed moving forward.

Author response:
We agree that these challenges are important to deal with in future work. Answers to these questions require a deeper analysis of the theoretical structure of linguistic data and their mathematical characterization. Specifically, future work should explore a theoretical framework more suitable for studying natural language, which is beyond the capability of conventional tools like HMMs and mixtures of distributions.

Reviewer comment (Theory):
The notation developed in the paper mostly results in an analysis (based on the heuristic similarity between QKV-attention and kernel regression) of the last tranformer layer.

Author response:
We thank the reviewer for this interesting question. Our explanation does not specify or limit the model from implementing kernel regression in a certain layer. In fact, this kernel regression can be implemented in one or more layers as long as these (possibly redundant) results can be passed to the final layer and aggregated. This “passing” can be easily implemented by the skip connections. We add this explanation to Section 5.1

Reviewer comment (Writing):
Formatting and notation of the article is somehow problematic with some symbols not being defined by the time they are used in the text. This makes understanding the details of the paper potentially tough, diminishing its impact. For instance, $\Sigma_{p_{pre-train}}$ appearing below eq (3) in the definition of $\epsilon_\theta$ is not defined.

Author response:
Thank you for suggesting improvements on notations. The covariance matrix $\Sigma_{p_\text{pre-train}}$ was a typo and supposed to be $\Sigma_{p_\text{pre-train}, l}$, which is defined in the last paragraph of Section 3.2 to be the covariance of elements in $T_{\mathbf{o}}$, when the sequence $\mathbf{o}$ is sampled from the initial distribution $p_\text{pre-train}$. In the same paragraph, $\epsilon_\theta$ is defined as $\inf_l \rho(\Sigma_{p_\text{pre-train}, l}^{-1} - \Sigma_{s_\theta, l}^{-1})$, where $s_\theta$ is the task-specific initial state, and $\rho$ is the spectral radius of a matrix. This quantifies the difference between pre-training and task-specific distributions. We also updated the revision accordingly.

Reviewer comment (Experiments):
Is this specific to transformers? How about LMs based on RNNs or attention-free architectures?

Author response:
Our explanation provides a mechanism of how attention-based models, such as Transformers and attention-RNNs, might be tackling the ICL problem. Other architecture might process the ICL inputs differently.

Reviewer comment (Experiments):
Did the authors test a toy benchmark where the pretraining sequences are generated from a known HHM where one can control size of the pretraining datasets and complexity of the HHM to empirically verify the convergence in Theorem 1? This seems like it would be a great way to verify the main theory result.

Author response:
Thanks for the constructive suggestion, and we agree that this helps the clarity and convincingness of the analysis. We experimented on randomly parameterized synthetic HMMs with 8 tasks, 80 states, and 100 observations. We vary the number of samples and evaluate the proposed kernel regression on fitting the Bayesian posterior, which is the intuition in Theorem 1. Results are listed in the table below. We indeed see a decreasing loss and increasing accuracy with more demonstrative examples. The loss also converges to a non-zero value according to Equation 6. These results are added to Appendix C.

| #samples | 1 | 2 | 4 | 8 | 16 | 32 | 64 | 128 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| distance | 1.322 | 0.942 | 0.519 | 0.218 | 0.163 | 0.085 | 0.093 | 0.083 |
| accuracy | 0.337 | 0.532 | 0.745 | 0.901 | 0.926 | 0.968 | 0.961 | 0.964 |

Reviewer comment (Writing):
Theorem 1 mentions that triangular brackets indicate the inner product, but this notation is not used in the theorem.

Author response:
Thanks for pointing out the typos. The inner product in Eq (4) is an unused notation. The summation notation in Eq (7) also contains a typo. We corrected them in the revision.

Reviewer comment (Writing):
Eq (7): sum over j but the index used is i.

Author response:
Thanks for pointing out the typos. The inner product in Eq (4) is an unused notation. The summation notation in Eq (7) also contains a typo. We corrected them in the revision.