Does learning the right latent variables necessarily improve in-context learning?

We thank the reviewer for their valuable comments and appreciate that they found that our work tackles an “interesting and important” question. We provide further clarifications and evidence below and hope that it alleviates the concerns raised by the reviewer.

…a theoretical framework or at least some mathematical justifications…

We refer the reviewer to works [3-6] which highlight that the use of an information bottleneck does lead to better generalization. The bottleneck which learns the true underlying latents can be seen as an information bottleneck since it is a low-dimensional embedding of the whole dataset. In particular, the explicit model could solve tasks by either distilling the information from context examples into the bottleneck (i.e. approximating the true latent) or by copying all the context examples in the bottleneck. The latter is disincentivized due to the lower dimensionality of the bottleneck; and thus we see that solely through maximum likelihood based training we are able to uncover the right underlying latents. Thus, based on the above reasoning, we stress that our arguments are not just heuristically motivated.

Further, one can see that the functions represented by the explicit model can actually be represented by the implicit model too; in particular the implicit model could first uncover the right latents in the query token’s representation and then use it for downstream prediction, e.g in [7]. However, with increasing context sizes it is less and less the case that an implicit model’s hypothesis space can be represented by the explicit model. Hence we can argue from the perspective that explicit models trim the potential hypothesis space and thus could enjoy better generalization, as is studied in multiple works on generalization bounds [8].

…careful with the general statement that using parametric assumptions inevitably leads to better generalization…

The reviewer has a valid point and we acknowledge that we weren’t completely clear about it in our draft. We will revise our writing to reflect that the kind of shortcuts we refer to are more kernel-regression / nearest neighbor style shortcuts, especially even when there is an underlying true parametric form for the solution.

It is true that for linear regression the ordinary least squares solution can be written as $w = (X^TX)^{-1}X^Ty$ which can equivalently be seen as both a parametric as well as a non-parametric model through the lens of prediction, i.e. $w^T x_* = y^TX(X^TX)^{-1} x_* = \sum_{i=1}^N k(x_i, x_*) y_i$ where $k(x_i, x_*) = x_i^T(X^TX)^{-1}x_*$. See section A.8 of [1] for a related discussion. However, while this equivalence exists for linear regression, it may not necessarily hold for other, more complex modeling approaches.

We refer the reviewer to the general response where we show that indeed by injecting shortcuts in an extreme fashion, the explicit model which relies less on kernel regression because of the bottleneck ends up generalizing better than the implicit one.

The paper would clearly benefit from providing more empirical evidence for the fact that the task vectors are actually always learned well by the explicit model.

We perform linear decoding and causal interventions to test for whether the explicit model does indeed infer the task vectors correctly. The only cases where we see the explicit model not being very successful at obtaining task latents is in nonlinear MLP and classification based experiments, where the latents are considerably high-dimensional. Here, the explicit models do “learn the correct latents” as best as possible, and both Figure 3 and 5 supports this result. The crucial component to understanding this is that the explicit model (even without the oracle prediction function; blue) achieves the same task performance as when the ground-truth latent variable is used to provide an auxiliary supervision signal (purple), which should be read as a performance ceiling. Note that perfect latent variable decoding accuracy for the classification tasks is impossible because with finite contexts there is insufficient evidence to infer the true underlying decision boundary, which accounts for the reviewer’s concerns about Figure 5(a) showing high latent variable decoding error in these plots. Indeed, the models trained to predict true latent variables (purple models) in Figure 5(a) show that the explicit model (blue) decodes the latents just as well across tasks.

[1/3]

Dear Authors,

some of my points have been clarified through your explanations and references to the literature as well as additional experiments. I therefore raised the score from 3 to 5. However, I would still appreciate a more thorough mathematical treatment going alongside your experiments.

Furthermore, releasing the code in the sublementary material is clearly better than not releasing it at all, but still not ideal for reproducibility.

We thank the reviewer for raising their score. However, we are a bit confused by the feedback provided and would like to request for additional information.

Mathematical Treatment: We thank the reviewer for their comment but are unclear as to what kind of additional mathematical formalism they are looking for. Could the reviewer clarify what exact formalism beyond the new Appendix E are they implying?

As an additional note, we are of the strong opinion that empirical evidence holds its own place in scientific contributions. In particular, a number of published works also formalize their theory in standard maximum likelihood estimation or empirical analysis of an emergent phenomena, just like us [1-5]. Could the reviewer provide further clarifications?

Reproducibility: We feel that we may have misunderstood the reviewer’s feedback. We plan to release our github repository publicly. Beyond experimental details in the Appendix and release of the code in supplementary, what else does the reviewer want us to provide?

We hope that our response has addressed the reviewer's concerns; and we would be happy to provide any further clarifications.

References

[1] Hendel, Roee, Mor Geva, and Amir Globerson. "In-context learning creates task vectors." arXiv preprint arXiv:2310.15916 (2023).

[2] Müller, Samuel, et al. "Transformers can do bayesian inference." arXiv preprint arXiv:2112.10510 (2021).

[3] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.

[4] Zhang, Chiyuan, et al. "Understanding deep learning (still) requires rethinking generalization." Communications of the ACM 64.3 (2021): 107-115.

[5] Mittal, Sarthak, Yoshua Bengio, and Guillaume Lajoie. "Is a modular architecture enough?." Advances in Neural Information Processing Systems 35 (2022): 28747-28760.

Dear authors,

Thank you for your reply.

However, I am not sure if I understand your first point

We first want to stress that learning the task vectors is paramount to generalization in in-context learning.

Given an optimal $g\_{z} : \\mathcal{X} \\rightarrow \\mathcal{Y}$ that is generated via some $z$, which is the case you consider. (i.e. $g\_{z}$ is your $g$ for fixed $z$). And given a model $f: \\mathcal{X} \\rightarrow \\mathcal{Y}$ that learns to approximate $g\_{z}$. I.e. the risk R = \\mathbb{E}\_{x \sim P^\\mathcal{D}}[l(g_{z}(x), f(x))] is small for some loss function $l$.

Where does the dependence of $f$ on $z$ come from??? Could you elaborate?

The model $f$ might in fact even learn an arbitrary task vector at some point. To see this, consider a bijection $h: \mathcal{X} \rightarrow \mathcal{Z}$ and define $g'(x) = g_{z}(h^{-1}(h(x)))$. Then $g'$ is still optimal, in fact even $g = g'$, but there is an arbitrary $h(x)$ to be found somewhere in the computation $g'$ does. I.e. what happens if $h$ is the encoder before the information bottleneck and $g_z \circ h^{-1}$ is the decoder after the information bottleneck?!

We believe that the reviewer answers this question themselves.

I do not agree with this statement. I would even like to point out that reviewer nwsg fully aggrees with my point that you do not show something very surprising in your paper:

However, it remains unclear whether the core question of the study is worth exploring. In my opinion, it is not surprising that implicit models perform as well as (or slightly better than) explicit models, given that much of modern AI progress has already been achieved implicitly. From this perspective, I think the paper(for the future work) could provide more valuable insights by either: (1) Identifying tasks where explicit models outperform implicit models, or (2) Exploring why implicit models perform better than explicit models.

I also do not agree with you saying

Then, wouldn’t the reviewer agree that it is an interesting and important scientific contribution to provide evidence as to whether the converse holds (does an optimal predictor need to learn explicitly?)?

As I said before, that the "converse" does not hold is not surprising.

We thank the reviewer for their valuable and detailed comments and are glad that they thought our work “brings a new perspective”, is “high quality” and brings a “novel perspective in a timely manner”. We now address their key questions and concerns.

I believe the paper would be very much improved if it verified empirically that the explicit model is itself resistant to shortcut learning in the presence of injected shortcuts.

The reviewer brings up an interesting and useful suggestion of studying how sensitive the explicit model is in the presence of injected shortcuts. To test this, we conduct an additional experiment based on sinusoidal regression where the queries were by design sampled close to the context during training. Our results highlight that the implicit model led to worse OOD generalization far from the context than the explicit model, highlighting that it was leveraging kernel-regression style shortcuts as opposed to learning the underlying sinusoid function. We refer the reviewer to the shared global response for details about the experiment and the results.

the use of "commonly" can be taken to imply that most people believe that shortcut learning is the reason for poor OOD generalization.

We refer the reviewer to works [1-5] which highlight how transformer models (and pre-trained large language models) fail to generalize because they rely on shortcuts as opposed to learning the underlying task structure. Since we study ICL in the context of transformer models, the kind of shortcuts preventing generalization studied in [1-5] transfer to our setting as well. More mechanistically, recent works [6-8] have identified pattern matching mechanisms like induction head to underlie ICL; something that further motivates our study of explicit models.

Based on the related work, we felt that this was a commonly held belief, but if the reviewer feels that it is not the case, we would be happy to omit it.

does your understanding of OOD generalization include robustness to variation of incidental features on ID data?

It is unclear to us what the reviewer means by “incidental features”. Could they please clarify?

We thank the reviewer for their time and constructive feedback and hope that our responses were sufficient in clarifying the excellent questions asked by the reviewer. We are more than happy to address any further questions that arise during this discussion phase.

References

[1] Anil, Cem, et al. "Exploring length generalization in large language models." Advances in Neural Information Processing Systems 35 (2022): 38546-38556.

[2] Zhang, Yi, et al. "Unveiling transformers with lego: a synthetic reasoning task." arXiv preprint arXiv:2206.04301 (2022).

[3] Liu, Bingbin, et al. "Transformers learn shortcuts to automata." arXiv preprint arXiv:2210.10749 (2022).

[4] Yang, Sohee, et al. "Do Large Language Models Latently Perform Multi-Hop Reasoning?." arXiv preprint arXiv:2402.16837 (2024).

[5] Bihani, Geetanjali, and Julia Taylor Rayz. "Learning shortcuts: On the misleading promise of nlu in language models." arXiv preprint arXiv:2401.09615 (2024).

[6] Crosbie, J., & Shutova, E. (2024). Induction heads as an essential mechanism for pattern matching in in-context learning. arXiv preprint arXiv:2407.07011.

[7] Tang, R., Kong, D., Huang, L., & Xue, H. (2023). Large language models can be lazy learners: Analyze shortcuts in in-context learning. arXiv preprint arXiv:2305.17256.

[8] Wang, L., Li, L., Dai, D., Chen, D., Zhou, H., Meng, F., ... & Sun, X. (2023). Label words are anchors: An information flow perspective for understanding in-context learning. arXiv preprint arXiv:2305.14160.

Based on your recent results, and your overall response, I have updated my score from 6 to 8. I would very much appreciate it if the authors abbreviated parts of the paper and included the results on injected shortcuts (Figure 8) and the update to Appendix D in response to "discussion on how observations differ across tasks" (response to reviewer nwsg).

Injected shortcuts

Thank you for acting on my suggestion. This result provides some empirical support for your presentation of the explicit model as parametric and the implicit model as non-parametric (p. on line 70, Section 3, and Figure 1). (To be fair to readers, please use the same range for the Y axis in both subfigures of Figure 8.) An additional improvement would be to demonstrate this on a few more tasks.

OOD generalization and "commonly".

Thank you for the additional references. Please cite these to support your framing.

Part of what I'm referring to wrt "commonly" is that generically, under empirical risk minimization, distributional assumptions hold, and there's no guarantee that a model will generalize beyond the distribution on which it was trained. I believe it's fair to argue that the brittleness of large transformers is, in part, due merely to their vast over-parameterization. They're high-variance models, although ones that happened to have been trained on a large amount of data. Please witness, e.g., how in Tang et al, 2023, the tendency of transformers to be susceptible to injected shortcuts increased with larger model. This suggests that the the particular shortcut behavior demonstrated in Tang et al, 2023, can be attributed in part to the sheer over-parameterization of the model(s).

That said, I believe it would help your work be properly understood in relation to other work if you were to state explicitly which kinds of shortcuts your work does not address. As you state in your response to reviewer q3Zg, you're addressing "kernel-regression / nearest neighbor style shortcuts". This is important, because it's possible that shortcuts demonstrated in some of the references you provided here in your response are of a different style altogether. You could mention in your framing that (1) you're addressing parametric vs. non-parametric styles of shortcuts and (2) not addressing shortcuts due to e.g. over-parameterization. That's just one example. One question is whether the [1-6, 8] you shared contain other styles of shortcuts. Even a rough and heuristic categorization of the styles of shortcuts could help readers to know in what context to understand your work.

Robustness to variation of incidental features on ID data

Apologies for my infelicitous phrasing. I asked this question to stimulate discussion to better understand your reasoning about OOD generalization (cf. my comment about ERM above). Robustness is the opposite of brittleness. In Tang et al, 2023, to continue with that reference, the injected shortcuts (such as the use of "movie" in their Figure 2) are variations that are incidental to the task. So I suspect that your answer to my question is affirmative, in a sense. Where to place the line between robustness to (noisy) ID data and OOD data can be argued in a different venue.

We thank the reviewer for their detailed comments and are glad that the reviewer found our work “addressing an interesting problem” and clearly written with a diversity of tasks. We take this opportunity to further clarify the concerns that were raised by the reviewer.

The finding that explicitly learning latent variables does not enhance performance may be limited to the specific tasks chosen for evaluation, especially those where implicit models already perform well

We thank the reviewer for bringing this point up and clarify that we actually tried to choose tasks for which we intuitively felt that the explicit model would do better, and not vice versa. For example, for all our tasks with low-dimensional task latents, eg. sinusoidal regression, linear regression, linear classification, raven’s progressive matrices, gene targeting and alchemy, the latents are actually low-dimensional and thus intuitively the explicit model should do better. However, counter-intuitively we see that the implicit model performs just as well, and often better; which really highlights the strength of our results supporting the hypothesis. However, to alleviate the reviewer’s concerns, we conduct an additional experiment based on a modular and compositional task, the results and details of which are provided in the shared global response.

discussion on how observations differ across tasks

We thank the reviewer for bringing this point forward and have now revised the draft and presented the requested analysis in Appendix D. Below we explain differences in results across tasks:

Classification tasks vs. regression tasks. OoD performance is generally strong (across all models) for classification because decision boundaries are within the training domain and do not change beyond it. In contrast, for regression tasks, the function continues to change beyond the observed training domain, making OoD prediction more difficult. This is also why known prediction functions give little benefit in classification tasks: they are already solved well OoD with ordinary implicit and explicit models.

The explicit model with known prediction function does not give benefits in nonlinear (MLP) regression. This is because the problem of inferring an MLP’s weights given some context examples is too difficult, so the explicit model potentially opts for a different, non-parametric solution. This is supported by the latent variable decoding results in Fig. 5 (previously Fig. 4), which show that even with a known prediction function the explicit model does not learn to infer the correct latent variable for the nonlinear (MLP) regression task.

Figure 4b, can you explain why MLP predictions are generally less effective than those from transformers?

We do not have a very clear answer as to why MLP predictions are generally less effective than those from transformers, and can only conjecture that it has to do with the inductive bias of the transformer decoder, in particular of the attention modules and dot-product similarity metric. These results are aligned with our point that the functional form of the decoder in the explicit model is crucial to reap the benefits of the learned latents.

We thank the reviewer for their time and constructive feedback and hope that our responses were sufficient in clarifying all the useful questions asked by the reviewer. We are more than happy to address any further questions that arise during this discussion phase.

Thank you for your response, which partially addressed my previous questions. I believe the additional tasks and discussions significantly enhance the paper. However, it remains unclear whether the core question of the study is worth exploring. In my opinion, it is not surprising that implicit models perform as well as (or slightly better than) explicit models, given that much of modern AI progress has already been achieved implicitly. From this perspective, I think the paper(for the future work) could provide more valuable insights by either: (1) Identifying tasks where explicit models outperform implicit models, or (2) Exploring why implicit models perform better than explicit models.

We thank the reviewer for their feedback, but it is not clear from their review as to what parts of the work were unclear. We provide a comprehensive assessment of the positioning, motivation and contribution of our work below, as well as an overview of our empirical analysis. We hope that it clarifies the reviewer’s concerns, and would be happy to address any further comments.

Motivation

Multiple works [6,7] claim that transformers perform ICL through pattern matching mechanisms such as induction heads, acting as shortcuts. Our motivation is to test if minimally modifying the transformer architecture to incentivize it to explicitly learn the task latents would prevent learning of shortcuts and lead to better OOD generalization. We test this theory by analyzing whether such a model actually learns the latents well, if it generalizes better than the implicit model and identify why it doesn’t.

Contributions and Novelty

We have now revised the draft to include a list of our contributions at the end of the introduction. The primary contribution of our work is that learning task latents correctly is not necessarily the only or biggest problem for generalizing in-context to out-of-distribution queries and tasks. This finding provides a different viewpoint to the importance and relevance of inferring the right latents (in contrast to [4,5]) and shows that even if this part is solved, it does not guarantee good OoD performance. Note that this negative finding is important to the field, as it significantly prunes the search space of solutions [3].

Impact

We argue that our findings also suggest promising directions. For example, it can lead to future work on improving the prediction module as opposed to context aggregation if OoD generalization is the goal (for instance, through task-specific architectures). Note that this approach is impossible using an “implicit” model, as context aggregation and prediction are entangled and we cannot use different inductive biases for each. Hence, we believe that even though we don’t develop these solutions in our work, our findings and analysis provide a thorough and comparative study to guide future research in OoD generalization for ICL.

Structure and Logic

While other reviewers found our work to be clearly written, we may have been unclear at some places. Could the reviewer clarify what in particular about the structure and logic was unclear, so that we can both rectify it in our draft and clarify it to the reviewer?

motivations seem like a course project rather than a publishable work

In this work, we study whether learning the right functional form of the solution through task latents leads to better OOD generalization in In-Context Learning, which as the reviewer themselves have pointed out is an “interesting” problem. We believe that our controlled experiments and systematic analysis provides answers and novel insights to an important question, which is valuable to science and beyond the scope of a course project.

how is the claim being rejected by the experiments

Our claims are being rejected by the experiments in the following manner

Learning shortcuts prevents OOD generalization: If learning shortcuts over the right parametric inference procedure was the main issue preventing OOD generalization in ICL, then the explicit transformer model should perform better than the implicit one. Our experiments show that this is not the case, even though explicit models do learn the true latent variables.

Failure of explicit models is due to suboptimal downstream prediction: Our experiments identify that the learned prediction function is unable to combine even the correct latents in a systematic way (Figure 3). In fact we know that the explicit model often does learn a linear representation of these right latents (Figure 5). This identifies the prediction problem as an important one for further research.

[1/2]

whether the experiments are scalable to LLMs

We also emphasize that the questions we investigated cannot be studied in complex modalities like language or vision because the true latents are generally unknown in these settings. Access to the true latents was necessary for the systematic analysis and results in Figure 3 and Figure 5, for instance, along with the insights they provided. It is exactly for this reason that we, as well as a number of related works that explore such questions [1,2], operate in this regime.

While understanding shortcut-learning and robust generalization would be interesting in pre-trained large language models, we consider it as relevant future work that is outside of our current scope. This is primarily because we learn less about the problem by directly evaluating on a modality where neither the nature of shortcuts nor the true latents are known. Naively testing in this domain can lead to conflicting results and evidence, as well as overconfidence in incorrect claims. Thus, our goal is to study these questions in a more controlled setting where we can make claims based on our results with more assurances.

We thank the reviewer for their time and feedback, and hope that our response clarifies their concerns. We would appreciate it if the reviewer would also let us know their specific concerns in more detail to give us a chance to improve our work and provide potential clarifications.

References

[1] Von Oswald, Johannes, et al. "Transformers learn in-context by gradient descent." International Conference on Machine Learning. PMLR, 2023.

[2] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.

[3] Karl, Florian, et al. "Position: Embracing Negative Results in Machine Learning." arXiv preprint arXiv:2406.03980 (2024).

[4] Xie, Sang Michael, et al. "An explanation of in-context learning as implicit bayesian inference." arXiv preprint arXiv:2111.02080 (2021).

[5] Todd, Eric, et al. "Function vectors in large language models." arXiv preprint arXiv:2310.15213 (2023).

[6] Crosbie, J., & Shutova, E. (2024). Induction heads as an essential mechanism for pattern matching in in-context learning. arXiv preprint arXiv:2407.07011.

[7] Tang, R., Kong, D., Huang, L., & Xue, H. (2023). Large language models can be lazy learners: Analyze shortcuts in in-context learning. arXiv preprint arXiv:2305.17256.

[2/2]

We thank the reviewer for their time and feedback, and are glad that they thought that the problem we study is “interesting” and “important for the community”. We address all of their questions below.

do you think that the controlled experiments are enough to draw such a strong conclusion?

We acknowledge the reviewer’s point and think that our controlled experiments are exactly what allows us to draw such a strong conclusion: that learning good latents only leads to better ICL performance when paired with an appropriate prediction function. Specifically, knowledge of the ground truth latents is crucial to most of our experiments; something we would not have in less controlled tasks. Nonetheless, we believe our claim is still important (if anything, more) in more natural forms of ICL, where the latents and prediction functions are simply much more complex.

We also conduct additional experiments to test for compositional generalization in a reusable mixture-of-experts task as well as analysis of the proposed models in the presence of shortcuts. We refer the reviewer to the shared global response for details about this setup and the corresponding results. Are there any specific additional experiments that would help solidify the validity of our conclusions for the reviewer?

will there be cases where learning good latent representation improves the prediction performance?

The reviewer brings up a very relevant question. We clarify that one of the conclusions of our empirical evidence is that learning good latent representations does in fact help in-context learning, but only when the architecture is able to use them optimally for the downstream prediction. Our experiments with explicit models indicate that the learned downstream prediction ended up being sub-optimal, and replacing it with the “known” prediction function instead (Figure 3) led to strong OOD performance.

We thank the reviewer for their time and constructive feedback. We hope that our responses were sufficient in clarifying all the great questions asked by the reviewer and are more than happy to address any further questions that arise during this discussion phase.