Toward Understanding In-context vs. In-weight Learning

We thank reviewer VrYi for their thoughtful reviews. We address your questions and comments here.

W1

We agree that our model does not exactly fit the training of large language models at the token level. However, if one interprets examples to represent documents, then assuming an i.i.d. distribution during training fits the setting of large language models quite well. Also note that few shot learning is another example of in-context learning where assuming that the examples presented are i.i.d. is again realistic. Another setting for ICL is having rare classes during training and testing. In this case, the training and testing sequences can be drawn i.i.d., but there is a large number of rare classes such that each class is seen only a couple of times. In this case, ICL may emerge due to insufficient in-weight data. Note that we do not assume that the data in the context and query would be i.i.d., and our lower bound for ICL (Proposition 2) and upper bound for Algorithm 1 (Proposition 3) hold under much more general settings, without making any assumption (e.g. independence) about the data distribution.

W2

We emphasize that the (input, label) setting still captures the sequential nature of ICL , since the ICL predictor is still a function of the entire sequence. However, we do not perform general sequence learning where the sequence length can vary. Our setting, also studied by Chan et al. (2022), is the simplest setting that allows us to readily investigate ICL vs. IWL in transformer networks. We can extend our study to problems with more complex sequence-space computational structure; indeed, this would correspond to having a more difficult concept class for IWL and ICL, as well as a different ICL predictor than the one instantiated in Section 2.2. We would also like to mention that Tong et al. (2024) is a concurrent work.

W3–W5

Thank you for the suggestion. We have reorganized the plots to address the raised concerns to improve clarity and reduce clutter—we have delegated the less important content to Appendix B and emphasized only the important messages. We have also included a schematic diagram to demonstrate data generation, training, and evaluation (Figure 6).

Regarding Figure 7 (now Figure 16 in the revised version), we can observe that when the input noise $\sigma \in \{0.0, 0.1\}$ the errors between the IW and IC predictors are similar on both the high-frequency ($C_H$) and low-frequency ($C_L$) classes, while when $\sigma = 1.0$, the IW error is clearly higher than the IC error on $C_L$. We note that for relevant context on $C_L$, the IW predictor can eventually achieve lower error than the IC predictor which is also when ICL begins to disappear.

W6 + Questions

Please take a look at the section “Significance of the LLM experiments” in the general response. In particular, the LLM experiments show that on inputs with in-weight data, the IWL predictor can potentially override the ICL predictor. We do not claim that after fine-tuning, the LLM will perform IWL on all inputs, in fact, ICL is very much present in the model for other inputs. This is consistent with our claim in the main text, where we show that ICL and IWL can emerge on different parts of the input space.

Thank you for the further discussions.

W1

Please note that the i.i.d. assumption is not applied to the elements of the context and the query but to the whole context and the query, that is $(\tilde{x}_i, y_i)$ are i.i.d. for all $i$, but there can be arbitrary correlation among the examples $(x^1_i,y^1_i), (x^2_i,y^2_i),\ldots,(x^L_i,y^L_i),(x,y)$ appearing in the context and the query (which allows the setting to model few-shot learning). The i.i.d. assumption here means that we randomly sample few-shot examples for various problems. We agree that it is possible that the training data may not follow this exactly, for example, it may follow a two-step procedure where first a problem is selected randomly and then several samples are selected from the same problem), but we believe that the i.i.d. assumption still gives a useful model.

Also note that our theory about ICL and IWL suggests (and also the experiments support this) that if some similar data is repeated too many times, the model will rather memorize the label than learn to predict based on the context, which suggests that ICL is more prominent the closer the data $(\tilde{x}_i, y_i)$ is to i.i.d.

To see this more formally, assume that the data distribution follows a two step generation procedure $P_{d|t} (data|task) P_{t} (task)$; then the options may be (i) the i.i.d. case where every sample is generated independently by first sampling a task $t \sim P_t$ and then a data $d=(\tilde{x},y)$ from $P_{d|t}(\cdot | t)$, in which case we obtain i.i.d. samples from the marginal for $d$; and (ii) the case where we take multiple samples from fewer tasks, that is, first we sample independently a few tasks $t$ from $P_t$, then for each task $t$ we sample several data points from $P_{d|t}(\cdot | t)$. In the latter case, if the points from a given task are more similar, then the data generated in this way is more likely to encourage IWL than in the i.i.d. case of (i).

To clarify, Proposition 3 in our initial response is now Proposition 2, where we have regret guarantees without distributional assumptions.

W2

One more comment about this question: We lightly discuss the effect of sequential training (of real LLMs) compared to our model, arguing that it actually may contribute towards the transience of ICL; see the discussion around line 1770 in our new Appendix D about the transience of ICL.

W3-W5

Should be Figure 17 now (apologies, we have uploaded another revision in the meantime).

W6

In our view, the LLM experiment is an efficient simulation for pretraining, as the fine-tuning procedure we are using is identical to what is used in pretraining, so we are not specifically optimizing for IWL vs. ICL. Of course we changed the data distribution, as instead of mixing our fine-tuning data into the usual (pretraining) data distribution, we only used the data we wanted the model to memorize, which is the most sample-efficient way to increase the number of times the fine-tuning data appears in the training, as required by our theory (although it violates the random order used in pretraining). This reduced the training time needed for the memorization effect to show up; mixing our finetuning data with other “usual” data would have slowed down the training. Also note that we do not consider the effect that the model is post-trained (using supervised finetuning and RLHF) between the normal pretraining and our “extension” of it. Nevertheless, based on the above considerations we think that our LLM experiment is sufficiently similar to pretraining with a changed data distribution (where the fine-tuning data appears more), and hence it is able to simulate how IWL and ICL changes during training (where the original and the fine-tuned models play the role of an earlier and a later checkpoint in the training, respectively).

Note that we also added a sentence to the beginning of Section 4 to clarify the role of the LLM experiment.

Clarity of the presentation

We are sorry to hear that we have not managed to address your concerns about the clarity of the presentation. If you have any suggestions that would help here, we would be grateful to hear about them.

We thank reviewer ueJX for their thoughtful reviews. We address your questions and comments here.

Descriptions of IC and IW predictors

To avoid any misunderstanding, we would like to point out that there is a typo in the reviewer’s interpretation of the experimental setting in the “Question”, where the descriptions of the IW and IC predictors should be swapped.

On Figure 2 and the reviewer’s criticisms that : (i) ICL emerges first on rare classes first; (ii) for large data size the transformer switches to IW learning.

Indeed, the OOBD plots (bottom row) in Figure 2 show the model’s ability to perform in-context learning: since the mapping from the inputs to the labels is changed (as described in lines 353-360), the only way to achieve good performance in this case is to use in-context learning where possible (i.e., when the context is relevant). As such, whenever we see high error rates for relevant contexts, it means that the model is not doing IC prediction. This supports our claim for (ii). In particular, the fact that the OOBD error of the transformer for relevant contexts (both for high- and low-frequency classes) becomes 1 for large training set sizes shows that the model is not able to do IC prediction, and the (close to) 0 IBD error (for all cases) shows that the model is able to do proper prediction for the base distribution, which – by exclusion – has to be IW prediction.

Regarding (i), we indeed did not make our claim precise in lines 208-210; and now have rephrased. We refer the reviewer to the general response for “Regarding the claim about ‘ICL emerging first in rare classes’”. Supporting our revised statements, on the two leftmost columns of the bottom row of Figure 2 (OOBD, Rel. Cont., $C_H$ and $C_L$), the transformer has a significantly lower error than random guessing at the beginning, showing the presence of IC prediction. As $N$ increases, we start to see that the transformer begins to degrade in OOBD performance for relevant contexts, eventually losing this IC-prediction ability for both high- and low-frequency classes, while achieving close to zero error for in-distribution data, showing the emergence of IW prediction.

Experiments with larger context lengths

We have conducted experiments with $L = 4$ while varying the number of relevant contexts (Figure 4 - please note that we have reorganized the figures and the references correspond to the revised version). Here increasing the number of relevant contexts, similar to burstiness in Chan et al. (2022) and Singh et al. (2023), can encourage the emergence of ICL as suggested by our theoretical findings as well. In the Omniglot experiments in Section 3.2 we have conducted experiments with $L = 2$ (Figure 5). Generally we see that the transformer exhibits some ICL capabilities on rare classes initially, which then disappears as $N$ increases. This corroborates the findings from Section 3.1 with the synthetic dataset using $L = 1$. Please also note that the original Figure 6 corresponded to cases where the relevant information appears only once in the context, which becomes very challenging for ICL even for relatively small context lengths (this has been observed in earlier work and is the reason why the relevant context is repeated several times in the experiments of, e.g., Chan et al., 2022a, and Singh et al., 2023).

Regarding using more transformer blocks and attention heads.

While we believe that using larger models may further demonstrate how the theoretical results align with scale, Singh et al. (2024) has demonstrated that induction heads can indeed be implemented with only two transformer blocks with a single head per layer. We further note that Singh et al. (2024) uses Omniglot to validate their results as well. Consequently, we do not believe that there is any confounder induced by the (potential) lack of model expressiveness.

We thank reviewer 8YMd for their detailed feedback, below we address their comments and questions. We have also fixed the typos, thank you for raising them.

W1

Please see the section “The theoretically optimal performance of IC and IW predictors, and their selection” in the general response.

W2

Please see sections “Claim on “Regarding the claim about ‘ICL emerging first in rare classes’” and “The theoretically optimal performance of IC and IW predictors, and their selection” in our general response.

Regarding Figure 1, the bounds properly describe the typical trends of the error of the predictors; the figure only illustrates the different regimes which show up when we compare IWL with induction-head-type ICL (that is why we neglected the constants when plotting the graphs). The figure also illustrates the typical behavior that ICL becomes harder when the context length and the amount of irrelevant information increase (and can occur anywhere in the context).

W3

Similar bounds can be obtained for the continuous input space using standard statistical learning theory results. Indeed, one can obtain generalization bounds for the IWL predictor that decay with $O(1/\sqrt{N_x})$ and depend on some complexity measures of the function class (e.g., Rademacher complexity) [1]. These bounds still show that as the sample size increases, the excess risk converges to 0 in the limit. Our regret bounds in Section 2.3 automatically hold on continuous domains. However, to ensure that the regret for learning $\alpha$ can be meaningfully optimized, we need some additional assumptions (in the tabular setting this is an online learning problem with linear loss functions). A sufficient assumption is that the function $\alpha$ can be parametrized with a finite number of parameters such that the loss function is convex in these parameters and the resulting function class is rich enough to partition the input space into two sets such that IC prediction is better on one set and IW is better on the other. We are happy to discuss these assumptions in the paper, but we thought that the technicalities needed to make these precise (e.g., defining Rademacher complexities) could perhaps draw the attention from the underlying simple phenomenon.

[1] Bousquet, Boucheron, Lugosi, Introduction to statistical learning theory. Advanced Lectures in Machine Learning, Springer, pp. 169–207, 2004

W4

Please see the section “Criticism that the stylized model/data distribution is too simplistic in the theory” in the general response.

W5

Please see the section “Regarding the point that the IC and IW predictors in the experiments might not exhibit only ICL and IWL respectively” in the general response

W6

Please see the section “Significance of the LLM experiments” in our general response

Comparison to Wei et al., 2023

Thanks for bringing up this paper. Indeed, it looks at how the semantics (most similar to in-weight knowledge in our terminology) effects in-context learning (in much larger language models), but they consider much higher-level features of the models (such as size), and not the actual “strength” of in-weight knowledge/generalization as we do. Most similar to our investigation is their experiment comparing ICL abilities of a language model and its instruction-tuned version, arguing that instruction-tuning helps ICL – in our experiments this would be similar to providing in-context examples to boost the ICL capabilities of the model. Our experiments are similar to their flipped-label scenario, where their in-context examples are in contradiction to the models’ semantic knowledge, but they examine dynamics between IWL and ICL as a function of simple model parameters, while here we are interested in much more refined connections, and our LLM experiments aims to control directly the strength of the model’s in-weight (semantic in their terminology) knowledge, as much as possible given the unknown training data and the biases already encoded in the trained model.

Thank you for your clarifications. You have answered most of my questions but I have some remaining follow-up questions. In the interest of prioritisation, my score is more based on the 'weaknesses' discussion above (to which I will reply shortly) than these questions, and I will understand if you might not get time to address all of these questions.

Q1. Tabular alpha: Thank you for clarifying that $\alpha$ is tabular. My follow-up question remains. Since $\alpha$ is tabular, the only way for it to learn to put mass on $h$ vs. $g$ given a context $\tilde x$ is for $\tilde x$ to come up during training. But if this context comes up during training, that also allows (tabular) sub-model $g$ to memorise it. As soon as $\alpha$ observes that $g$ achieves high loss on the context and starts putting mass on $h$, $g$'s loss will rapidly decrease, and it would actually be apropriate for $\alpha$ to put the mass back onto $g$. Is that right?

Q2. Sampling distribution for $y$: Thanks for clarifying. I see this is outlined in the preliminaries section.

Q3. Example risk vs. expected risk: I don't think you really answered my question. I'm concerned about the apparent comparison between:

1. the upper-bound on the expected risk of IWL, where the expectation is taken over all $y$ and holds for any $x$, and
2. upper- and lower-bounds on the cross entropy for a given $x$ and $y$ assuming that $x$ has a certain number of irrelevant labels (which depends on the concrete value of $y$.

So the first bound is an average and the second is a function of $y$ (through the variable $k$ which implicitly depends on $y$).

It seems that in order to perform a direct comparison with the IWL bound, you would want to compute the expected-risk version of the ICL bounds. Note that this would require comparing bounds for different values of $y$ and therefore different values of $k$. Does that make sense?

Q4 Assumption on $y$ in corollary 1: Sorry, I either misread or misquoted the statement of the corollary. My real question is, why is it necessary to explicitly state the assumption that $y$ is one-hot again in the statement of this corollary?

If it's a stylistic choice, fine---I'm just checking that I am not missing something. I thought this assumption is also needed and used for proposition 2 (it's cited in the proof) but it wasn't stated explicitly in the statement of proposition 2.

Q5. Upper and lower bounds crossing: Thanks for clarifying the intended range of $\epsilon$. It seems my original question was improperly motivated by looking at logarithm base 10, for which $-\log_{10}(\epsilon) < 1-\epsilon$ for some $\epsilon \in (0,1)$. My mistake.

Q6. The big-O constant in figure 1: Thanks. Let me check my understanding:

• The "second term" you refer to in your answer is $6 G_\infty \sqrt{\log(2|\mathcal{X}|C/\delta)/N_x}$ which can be rewritten as $6 G_\infty \sqrt{\log(2|\mathcal{X}|C/\delta)} \cdot 1/\sqrt{N_x}$.
• As you observe on line 211, this term is in $O(1/\sqrt{N_x})$. The corresponding "big-O constant" to which you refer in the figure is therefore $6 G_\infty \sqrt{\log(2|\mathcal{X}|C/\delta)}$ (this is the bit I missed motivating my original question).
• To avoid specifying each individual quantity determining this constant, you just replace it with 20 for the figure, in other words you just plot the IW risk bound as (first term of IW risk bound) $+ 20/\sqrt{N_x}$.
• It "does not substantially affect the graphs and the conclusions" in the sense that regardless of the constant, the IW upper bound still falls below the IC risk lower bound for large enough $N_x$ and $k$.

Is that right?

Another typo: On lines 346, 809, 1035, 1288, and 1295 (both lines near 1295) you have "transient of ICL", did you mean "transience of" or "transient nature of" in each case?

Notes:

• I did see appendix D. You discuss two experiments in Singh et al.'s setting where IWL and ICL have asymptotically equivalent risk. In the Omniglot experiment, you say that you replicate the finding of Singh et al.. In the synthetic experiment using your model, you say ICL does not turn out to be transient. To me, this appears to serve only to undermine the claim that your model has bearing on Singh et al.'s type of transience (because it doesn't arise in the closest variant of your model in an empirical experiment, though I am not holding this against you).
• As I said I still find your model informative. I think if you want to point to prior work on a phenomenon like transience that your model is consistent with, consider Panwar et al. (2023), who study under the name "forgetting" a phenomenon like "transience" driven instead by a difference in risk.
• A minor point, the Appendix D hyperlink does not take me all the way to appendix D. (Possibly this is due to the presence of figures, perhaps using \clearall before the start of appendix D would fix this, not sure...)

W2: Incomplete characterisation of emergence for rare classes: Now about the other interpretation of your analysis, in the rare classes regime. Here you say that in this setting, IWL performs worse than ICL due to insufficient in-weight data. As I noted, I understand this to be based on a comparison between the upper bound on the IWL expected risk and the upper bound on the ICL example risk (the IWL upper bound is larger for small $N_x$).

My concern is that, of course, a larger upper bounds does not, in general, imply a larger value. You need to see a lower bound higher than an upper bound to infer that the former's value is higher than the latter's. Or, it would be sufficient to assume that the upper bound is actually tight.

In your reply you mention that "the bounds properly describe the typical trends of the error of the predictors." What is this assertion based on?

W4. Specific model of in-context learning appears quite restrictive: Finally, I'm still concerned about the concrete softmax-similarity induction-head ICL predictor as being too restrictive.

You pointed out that your mixture model works for a black-box ICL predictor. I understand this but it does not address my concern, because your story for the ICL vs. IWL conditions of emergence depends on your risk bounds for the specific model, which are not shown to hold for general black box predictors.

I understand that the softmax-similarity induction-head ICL predictor model is a small change from induction heads to allow approximate matches with prior tokens. However, I am concerned about this change because this change is load-bearing for the risk analysis. It is "load-bearing" in the sense that the risk lower bound (which is crucial for the paper's claims because it's what you compare to the IWL risk bound) specifically comes from considering the contributions to the output from all irrelevant pairs in the context. Yet, this lower bound is sensitive to the specifics of the ICL model. An exact-match induction head that ignored prior labels unless the inputs were equal to the prompt would have a different lower bound, for example. More generally, any larger family of ICL models (containing this family as a subset) would have a smaller or equal risk lower bound.

I think it's fine to analyse some arbitrary model of ICL that permits analysis, and this model seems vaguely justified, but the choice is not that compelling to me and since the conclusions seem to depend on the choice, I would have liked to see the choice acknowledged as such so as to make it clear to readers that this is just 'one' model of ICL (and if your transformer performs ICL using a different mechanism, maybe this risk analysis won't apply).

Summary: I've tried to lay out my understanding of the paper's claims and contents, and how this has given rise to my concerns, in detail. While my understanding is based on my best effort given available time I acknowledge that I might have misunderstood something about the argument or the paper's claims. I would consider raising my score if you can show me this is the case. Alternatively, if my understanding of the analysis is true, I would consider raising my score if you can adjust the introduction to make the claims match the conclusions of the analysis. I welcome further discussion.

We greatly appreciate reviewer 8YMd for their thorough review and on-going discussions. We will first address the follow-up questions and later address the weaknesses tomorrow AoE.

Q1

You are right that this would happen in the tabular case with no noise. However, if noise is present, the number of times a given input/context is seen matters. On the other hand, as we mentioned before, we chose the tabular setting only for convenience and our results are applicable more generally – given your concerns it seems to make sense to write out this generalization more properly. We discussed these generalizations in our response to W3 above. In short, the results of Section 2.3 are applicable much more generally (without essentially any restrictions), but are only useful if the regret terms in the upper bound can be optimized meaningfully. A sufficient assumption is that $\alpha$ is parametrized with a finite number of parameters such that the loss function is convex in these parameters and the resulting function class is rich enough to partition the input space into two sets such that IC prediction is better on one set and IW is better on the other. Note, however, that although proper regret bounds are not possible, regret-minimization algorithms (i.e., usual gradient-based optimization algorithms) perform well in practice even in the non-convex setting, so the bound of Proposition 3 based on regrets is still meaningful in such cases.

Q3

You are correct that the risk of IWL is an expectation over $y$ given $x$, the query, since the output of the IWL predictor only depends on the query. However, the output of the ICL learner depends on the entire sequence $\tilde{x}$, and for some distributions of $\tilde{x}$ (depending on $k$), the ICL predictor has a higher expected risk.

More specifically, we have

$\mathbb{E}_{(\tilde{x}, y) \sim D} \left[ \ell(g(\tilde{x}), y) \right]$

$= \mathbb{E}_{(\tilde{x}, y) \sim D} \left[ \ell(g(x), y) \right]$

$= \mathbb{E}_{(x, y) \sim D’} \left[ \ell(g(x), y) \right]$,

where $D’$ is the marginal density of $(x, y)$. Our IWL risk upper bound holds pointwise for each $x$. For the ICL predictor, the risk is $\mathbb{E}_{(\tilde{x}, y) \sim D}\ell(h(\tilde{x}), y))$, and our bounds for ICL holds pointwise for each $\tilde{x}$ and $y$. Suppose $D$ is only supported on sequences with at least 1 irrelevant label, then the ICL predictor would have a minimum error that can potentially be higher than that of the IWL predictor. Note that under our data generating process, $y$ is sampled iid given the query $x$.

Q4

It is a stylistic choice to include this explicitly in the statement, and now we have also added it to Proposition 2; thanks for spotting this.

Q5

To avoid misunderstanding, we have added a footnote about the base of the logarithm on page 4.

Q6

Exactly. Of course, the final point where the IW upper bound falls below the IC lower bound depends on the actual value of the constant, but all we wanted to show in the figure was that this happens eventually. Please let us know if you think more explanation is needed about this in the text.

Typos

Thank you, hopefully we have corrected all occurrences now.

Thank you for your response and revisions. You have allayed some of my concerns (W3, W5, W6) but unfortunately I maintain my concerns regarding the other issues (W1, W2, W4). The revisions have come some way to indirectly addressing these concerns but ultimately I am still unsure that the framing of the paper in the introduction is accurate and I would like to discuss further to see if we can resolve this.

In your introduction, you mention that through your generalisation and regret bound analysis in your simple model:

• you "identify conditions where in-context and in-weight learning emerge,"
• these conditions are "consistent with" the empirical findings of Chan et al. (2022a), and
• the conditions "partially explain" transience (I note the revision from saying they "explain" transience).

In my understanding, this is the central case for the significance of the theoretical contribution of the paper (with the empirical contributions serving to support the theoretical contribution). Accordingly, the standard I think is appropriate for evaluating this paper is, to what extent are these claims substantiated soundly and explicitly in the main text?

My remaining concerns center around whether or not these claims are in fact substantiated by your theoretical analysis. I have looked at the answers in the top-level comment and the new sections in the paper, but my core technical concerns have not been directly addressed.

First, let me check my understanding: What are the "conditions" you identify, and how are they based on your model?

1. I think that when you refer to the conditions your analysis identifies, you refer to the passage you clarified on lines 228--232:

   Under this algorithm, it is possible that initially IWL performs worse than ICL due to insufficient in-weight data. For example, when the training data contains a large number of rare classes, and each class is seen only a few times. Eventually, with enough in-weight data, the model can memorize the solution to achieve as good or even better prediction accuracy by using IWL on the query, rather than using the context for prediction and exhibiting ICL.

2. By "this algorithm" you refer to the optimal IWL--ICL mixture model, whose predictions are based on ICL or IWL depending on which has lower expected risk.

   Therefore, if I understand correctly, this (the expected risk comparison) is the "condition" you have identified. Is that right?

3. Then you interpet this condition: "it is possible that initially IWL performs worse than ICL due to insufficient in-weight data."

   If I understand correctly, you consider this claim supported by your theoretical framework based on the fact that the IWL expected risk upper bound is initially higher than the ICL risk upper bound. Right?

4. Then you continue to interpret, "with enough in-weight data, the model can memorize the solution to achieve as good or even better prediction accuracy by using IWL ... rather than ... ICL."

   If I understand correctly, here you are referring to the fact that with enough examples, the IWL risk upper bound falls below (or as low as) the ICL risk lower bounds (for some $k$). Is that right?

Assuming this is right (please clarify if you would put it differently), let me restate my concerns.

W1. Mismatch with "transience" phenomenon from Singh et al.: You say this story "is consistent with the empirical observation that ICL ... can be transient" then cite Singh et al. (2023) and point to the new Appendix D.

I agree that your theoretical framework points to conditions under which IWL could arise during training (when the IWL lower bound goes under the ICL upper bound). One could call this "transience." However, I disagree that this is the "transience" studied by Singh et al. (2023). In my reading, the kind of "transience" studied in Singh et al. (2023) is specific to the setting where IWL has no risk advantage over IWL.

I still don't think your theory matches up with this setting. As far as I can tell, your model makes no prediction about what a transformer should prefer in the case where IWL expected risk = ICL expected risk.

Therefore, I am not sure in what sense your framework either offers a "partial explanation" or is even "consistent" with the findings of Singh et al. (2023), since it refers to a different setting.

We would really like to thank reviewer 8YMd again for the in-depth discussion and suggestions.

W1

Thanks for keeping this interesting question on the agenda. We think that we finally consolidated our findings properly, including a satisfactory explanation for the contradictory experiments: Our model predicts the transience of ICL whenever at some point IW prediction becomes better than IC prediction. This is easy to see to happen in cases where IWL is asymptotically better than ICL. However, when the two methods can both achieve the same asymptotic performance, it is hard to see which one can have any advantage in the large but finite sample regime. In fact, while Singh et al. (2023) shows an example where ICL is transient, in Figure 21 (left) we show an example where it is persistent and another where it is decreased over time but remains visible (this is similar to the experimental results of Singh et al.). Now we also provide some speculations for this case, which, together with our theory, correctly predicts the transience of ICL. We clarified these in the introduction (lines 64-68, and lines 71-74), at the end of Section 2.2, and in lines 1758-1792 of Appendix D.

W2

To be more concrete, we have replaced Proposition 1 with a special bound for the cross-entropy loss, for which one can obtain faster rates than for general convex functions. The new rate, which comes with a lower bound that even matches the leading constant, is $\Theta(\log(N_x)/N_x)$. We have also changed Figure 1 to reflect this convergence rate. Of course, these rates are not directly applicable to training a complex function class like transformers, but they are suggestive and align qualitatively with some of the experimental findings.

W4

While our ICL class $\mathcal{H}$ (used in the theory section) is indeed specific, it can be used to illustrate and analyze the potential difficulties with in-context learning: In general, the difficulty of learning $h$ and the quality of the solution depends on the hardness of the in-context learning problem. Our simple ICL function class makes the “prefix-matching” aspect (Olsson et al., 2022) of the induction head explicit, and hence allows to quantify the difficulty of ICL in terms of the number of irrelevant contexts (one of many possible measures to quantify hardness). While this is explicit in our analysis (lower bounds) for $\mathcal{H}$ due to our simple structure, the same difficulty is also reflected in our experiments, where we vary the number of relevant contexts and also the context length. We can see that irrelevant labels indeed present a challenge for ICL as we show in our transformer experiments: when we have only 1 relevant label out of 4 labels, the transformer does not learn to do ICL at all (Figure 4, first column). In general with fewer relevant contexts or more irrelevant contexts, it is more difficult for ICL to emerge (see also Figures 11 and 12). Chan et al. (2022) and Singh et al. (2023) also made similar observations, and had to repeat relevant data in the context several times (“bursty samples”) to achieve in-context learning. Thus, from this perspective, we think that the function class we consider to demonstrate limitations is both simple and meaningful. Note that we have also added in lines 190-192 that our ICL predictor is only an example of a mechanism under which ICL can occur.

Suggestions

W3

Thank you for your suggestion. We have included them in our appendices A.1.

W5

We added “qualitatively” to line 350 to strengthen this. We believe we have been clear about this---we don't claim that the theory would actually properly match practice.

W6

We added the first two points and, for the third one, we have emphasized that we use a few examples to demonstrate that memorization can affect IW vs IC learning. Regarding comparison to Wei et al. (2023), it seems that for similar-sized models (the smaller GPT models) their experiment is quite inconclusive (Figure 2 in their paper) as even for no label flipping the prediction accuracy is around 50%, so it is not clear what kind of comparison could be made. We added a paragraph explaining this to the end of Appendix C.1.

We thank reviewer gQxH for their insightful feedback. Here we address your comments and questions.

W1

Though real-world LLMs are more complex, we focus on the IWL vs. ICL aspect of in-context learning, as studied by Chan et al. (2022a), Singh et al. (2023), and the concurrent paper of Tong at el. (2024). We propose a stylized theoretical model that is amenable to analysis; given the theoretical model, we describe a possible mechanism for the emergence of ICL vs. IWL, and support our hypothesis with experiments.

Tong and Pehlevan. MLPs Learn In-Context. arXiv preprint arXiv:2405.15618 (2024).

W2

We refer the reviewer to the general response regarding “Criticism that the stylized model/data distribution is too simplistic in the theory”.

W3

Please see the section “Significance of the LLM experiments” in the general response.

Q1

Thank you for bringing this paper to our attention. It considers a related problem: while we consider the difference of ICL and IWL, this paper considers different variants of IWL predictors, and examines the problem that different IWL predictors (skills) can be selected by the model given the input. The difference in our approaches is that while we model the behavior of an LLM with general predictors and a combiner, Lin and Lee (2024) considers the model of optimal Bayesian prediction in linear regression with the squared loss over Gaussian mixtures. While they explain the potentially incorrect in-weight learning with task-similarity, we are mostly interested in understanding how the training process can affect where the model will apply in-context versus in-weight prediction, which is important in designing specific training methods to enable balancing factuality and grounding to the prompt.

Lin and Lee. Dual Operating Modes of In-Context Learning. ICML 2024. https://arxiv.org/abs/2402.18819

Criticism that the stylized model/data distribution is too simplistic in the theory (reviewers gQxH, 8YMd, and VrYi)

Our main message is that ICL can emerge when it outperforms IWL on an input sequence, for example when the model does not have enough samples to learn in-weight. We use a simple model that can learn to weigh ICL vs. IWL. The ICL predictor in Section 2.2 is instantiated with the kernel predictor just for concreteness. However, our learning results in Section 2.3 do not assume a particular form of the ICL predictor and can be applied more generally (we revised the text introducing the specific form of $\mathcal{H}$ in line 187 accordingly) . In particular, the proposed algorithm can treat the ICL and IWL predictors as black boxes and can learn to use ICL vs. IWL depending on the regret of the predictors.

Our particular ICL predictor is inspired by the induction-head mechanism where it chooses and copies the most similar token (Olsson et al., 2022), and this is one way to instantiate this. The ICL predictor is a natural generalization of the induction head: we do not assume exact match between the query and the examples, but instead weigh the labels by the similarity between the example and the query. Our combination of the ICL and IWL predictor closely follows the mechanistic working of a softmax attention model, and could be thought of as a more abstract model being applied at an arbitrary layer of a transformer, combining concepts rather than tokens.

Regarding the point that the IC and IW predictors in the experiments might not exhibit only ICL and IWL respectively (reviewers ueJX and 8YMd)

There is no easy way (and it might actually be impossible) to separate the IW and IC predictor circuits of a transformer model, since they share significant architectural components. To estimate the respective performances, we trained two transformers, ICL and IWL in the paper, with only in-context and, respectively, in-weight data. These transformers act as proxies for measuring whether it is possible to perform ICL and IWL in the idealized case.

To further bridge the gap between the theory and experiments, we have added new experiments in Appendix B.2 that follow the stylized model (Figures 13 and 14). To further generalize to other model architectures we have also trained RNN-based models (Figure 15) on the same synthetic datasets, which shows similar trends as the results with the transformer architecture.

Clarification of the experiment figures

In the revised version we improved the presentation of the experiments and their conclusions. Specifically we have updated the axes and labels to be consistent across all figures. To reduce clutter we have also reorganized some results and moved the less informative results to Appendix B.