Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens

Pre-trained large language models based on Transformers have demonstrated remarkable in-context learning (ICL) abilities. With just a few demonstration examples, the models can implement new tasks without any parameter updates. However, it is still an open question to understand the mechanism of ICL. In this paper, we attempt to explore the ICL process in Transformers through a lens of representation learning. Initially, leveraging kernel methods, we figure out a dual model for one softmax attention layer. The ICL inference process of the attention layer aligns with the training procedure of its dual model, generating token representation predictions that are equivalent to the dual model's test outputs. We delve into the training process of this dual model from a representation learning standpoint and further derive a generalization error bound related to the quantity of demonstration tokens. Subsequently, we extend our theoretical conclusions to more complicated scenarios, including one Transformer layer and multiple attention layers. Furthermore, drawing inspiration from existing representation learning methods especially contrastive learning, we propose potential modifications for the attention layer. Finally, experiments are designed to support our findings.

Paper

Similar papers

Peer review

Reviewer 8J5b5/10 · confidence 4/52024-07-12

Summary

The paper explores the ICL capabilities of LLMs, focusing on understanding the mechanisms underlying ICL. The authors aim to investigate the ICL process using representation learning principles. They mainly use kernel methods to develop a dual model for one softmax attention layer, demonstrating that the ICL inference process of this layer aligns with the training procedure of its dual model. * Extend theoretical analysis to more complex scenarios with multiple attention layers * Propose potential modifications to the attention layer to enhance ICL capabilities Experiments also support the theoretical findings and improvements on ICL. This paper may give a deeper understanding of ICL and suggests practical approaches for improving the ICL capabilities of LLMs in real-world applications.

Strengths

* While previous studies mainly on the linear attention, this work extends it to a more realistic case -> softmax attention, including its training dynamics. * Based on the dual gradient descent process, the authors also propose potential attention modifications, which could be benificial to real-world downstream application.

Weaknesses

* The work may ignore the format of task input/output. And some of recent work such as (Theoretical Understanding of In-Context Learning in Shallow Transformers with Unstructured Data, Xing et al.) the success of ICL may be also because of the alignment of attention mechanism, especially in real-world nosiy cases. * The experimental section is not well-organized and self-explained, the readers may need to refer extra papers to get the protocols.

Questions

Theoretical Section * Will the ratio of the number of negative samples have an impact on generalization bound? Experimental Section * As mentioned in weeknesses part, the authors may need to add more experimental details to make it more self-explained. * Since the authors propose three attention modifications, which one could be the best, and do we need add these three together? why or why don't need? * Does the conclusion hold the same for both one-layer transformer and multi-layer transoformer? * The dynamics of token interactions may not be well supported in the experiments. Conclusion Section? * seems missing a conclusion section due to the restriction of page sizes.

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have clearly stated the limitations of this work, the selected transformer components of this work, and selected task settings.

Reviewer k2vC5/10 · confidence 1/52024-07-12

Summary

The paper explores the in-context learning (ICL) abilities of Transformer-based models. The authors propose an interpretation of ICL through the lens of representation learning. They establish a connection between the inference process of softmax attention layers in Transformers and the gradient descent process of a dual model, providing theoretical insights and generalization bounds. The paper also suggests potential modifications to the attention mechanism inspired by contrastive learning techniques and supports its findings with experiments.

Strengths

(Reminder) I'm not a researcher working in the theory field, so I couldn't recognize this paper's theoretical contributions well. - The paper provides a fresh perspective on understanding ICL by connecting it to gradient descent and representation learning, which is a novel and insightful contribution to the field. - The authors extend their theoretical findings to various Transformer settings, including single and multiple attention layers, enhancing their conclusions' generalizability.

Weaknesses

(Reminder) I'm not a researcher in the theory field, so I couldn't also recognize this paper's limitation from a perspective. - The paper relies on several assumptions and simplifications, such as ignoring the impact of layer normalization and residual connections. These might limit the applicability of the findings in more complex Transformer architectures. - Lack of discussions towards some related works that argue against the equivalence between ICL and gradient descent [1,2]. [1] In-context Learning and Gradient Descent Revisited, 2023 [2] Do pretrained Transformers Learn In-Context by Gradient Descent, 2023

Questions

- How would the inclusion of layer normalization and residual connections affect the theoretical framework and findings? - How do the proposed methods perform on more complex and diverse tasks beyond the linear regression tasks used in the experiments?

Rating

5

Confidence

1

Soundness

3

Presentation

3

Contribution

3

Limitations

The experiments are primarily focused on linear regression tasks, and it remains to be seen how well the proposed methods generalize to a wider range of tasks and datasets.

Reviewer cRiK8/10 · confidence 4/52024-07-13

Summary

The author's present a new way of linking in-context-learning (ICL) to gradient descent. The author's are able to demonstrate that indeed a (simplified) transformer decoder layer ICL is equivalent to "representation learning". Using the theoretical findings the authors are also able to propose extensions to the attention mechanism. The paper shows theoretical proofs for these claims as well as additional experiments to demonstrate the claim.

Strengths

Strenghts: 1. Excellent theoretical contribution to understanding in-context-learning 2. Adding relevant proofs 3. Using theoretical results to motivate better attention mechanisms

Weaknesses

Weakness: 1. Slightly simplified transformer architecture is used (it would be interesting to see the skip-connection version). As recent work, such as "mechanistic interpretability" often relies on these skip connections. 2. The experimental setup and results are a bit limited. Slightly more description of the task would be good, as well as a somewhat "realistic" task would be interesting to see - and what the results mean qualitatively. 3. A small section on how these results can be applied for bigger models in practice.

Questions

1. How could a "realistic" task look like for your evaluation? 2. How could one practically extend your augementations to e.g. modern LLMs (even of `moderate' size, e.g. 7B) - I.e. what would these "data augmentations" or "negative samples" really look like?

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

Yes.

Reviewer HnPE6/10 · confidence 4/52024-07-20

Summary

The paper investigates the in-context learning (ICL) capabilities of Transformers, explaining it through a representation learning lens. It establishes a theoretical connection between ICL and gradient descent, deriving a generalization error bound tied to demonstration tokens. The authors also suggest modifications to the attention mechanism inspired by theory and support their findings with experiments.

Strengths

* **Theoretical Depth**: It provides a rigorous theoretical analysis, including the derivation of a generalization error bound, which contributes to the theoretical foundation of Transformer models. * **Insights into Attention Mechanism**: The paper provides potential modifications to the attention layer, inspired by theory, which could potentially improve the learning capabilities of Transformers. * **Good Writing**: Clear Structure and Presentation: The paper is well-structured, with a clear abstract and a good introduction. It is easy for readers to understand.

Weaknesses

* **Generalization to Other Tasks**: The paper's findings are based on specific tasks. It's unclear how well these insights would generalize to general language tasks. * **Empirical Validation**: While the paper includes experiments, the extent of empirical validation is limited. We may need more experiments on real-world language tasks to verify the modification of the attention mechanism.

Questions

Please see the weakness section.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper has discussed the limitations sufficiently.

Authorsrebuttal2024-08-13

Looking forward to your reply

We sincerely appreciate your time and effort in reviewing our manuscript and providing valuable feedback! As the author-reviewer discussion phase is nearing its end, we would like to confirm whether our response has effectively addressed your concerns. We provided a detailed response to your concerns a few days ago and hope that they have adequately addressed your concerns. If there are any remaining questions, please do not hesitate to let us know. We would greatly appreciate any additional feedback you may have!

Reviewer cRiK2024-08-08

Thank you for your rebuttal and additional comments. At this stage there are no questions on your rebuttal to my questions. With regards to a response to another reviewer (k2cV). > Weakness 2 (Lack of discussions towards some related works that argue against the equivalence between ICL and gradient descent...) The answer is not very clear. Could you elaborate on: 1. A 1-1 (quick) comparison against each of the five papers. [1-5]. 2. A more detailed explanation why using the dual is important. (As the dual is simply a mathematical tool).

Authorsrebuttal2024-08-09

Replying to Official Comment by Reviewer cRiK

Thank you sincerely for your feedback. Here is our response: > **Question 1: ** A 1-1 (quick) comparison against each of the five papers. [1-5]. - Oswald et al. [1] focus on the ability of the linear attention layer to perform gradient descent during ICL when faced with linear regression tasks. Their observations are based on specific assumptions, including the constructed forms for $W_Q, W_K, W_V$ and the input tokens (concatenated $[x,y]$). Our work does not target specific tasks like linear regression; therefore, we do not make detailed assumptions about the model weights (simply treated as weights after pre-training) or construct specific input forms. Instead, we aim to view the ICL inference process from the perspective of representation learning in the dual model. - Zhang et al. [2] also theoretically analyze the gradient flow of linear self-attention modules under certain initialization conditions when dealing with linear regression tasks and provide the forms of weights at the global minimum as well as the prediction error. Similar to the comparison above, the aim of our work is not to investigate the expressive capabilities of the model's structure for these specific tasks but to interpret the ICL inference process from a representation learning lens in a more general setting, thus we do not use these assumptions. - Shen et al. [3] experimentally illustrates that the assumptions used in previous works including Oswald et al. [1] and Zhang et al. [2] may be too strong in real LLMs. They analyze the differences between ICL inference in LLMs and the fine-tuned models in real-world scenarios from various perspectives, including weight sparsity, order stability, and output distribution. In comparison, our work: (1) for the sake of theoretical analysis, considers not a "complete" real model but a simplified one that omits structures like residual connections; and (2) our interpretation of the ICL inference is not linked to fine-tuning on the original model but rather to training on the dual model. - Dai et al. [4] use the dual form to interpret ICL as an implicit fine-tuning (gradient descent) of the original model under the linear attention setting and this alignment is ambiguous as the specific details of the gradient descent process is not clear. Thus, our work extends this analysis to the nonlinear attention setting and delve deeper into more details of this process (exploring the specific form of the loss function). The main difference is that we consider this as training the dual model under a self-supervised representation learning loss, rather than performing supervised fine-tuning on the original model. - Natan et al. [5] investigate potential shortcomings in the evaluation metrics used by Dai et al.[4] in real model assessments and propose a layer-causal GD variant that performs better in simulating ICL. However, similar to the comparison above [3], their study discusses complex gradients in the original model in real scenarios, whereas we turn our attention to the gradient descent of the dual model of the attention layer. > **Question 2: **A more detailed explanation why using the dual is important. (As the dual is simply a mathematical tool). By using the dual, or rather, by analyzing the potential gradient descent process of the dual model, we can gain new insights into the model mechanisms in reverse. Specifically, through the dual perspective, - we transform the forward inference process into an optimization process. Since optimization processes are well-known and have established theoretical tools (for example, generalization error as mentioned in our work), this transformation can provide reverse insights into analyzing the model mechanisms. - we can clearly observe that the dual model involves a self-supervised representation learning process. Considering that there are lots of mature works in this area, we can draw on these works to reflect on the attention mechanism, which has also inspired attention modifications as illustrated in our work. Finally, as for "the dual is simply a mathematical tool", we want to clarify that the term "dual" we use is different from the one in optimization within the mathematical field. Instead, it follows the terminology used in previous work [4], where the forward process of the attention layer and backward process on some model are referred to as a form of "dual". We are not sure if this has also led to some misunderstandings. If there are any remaining questions, please do not hesitate to let us know! [1] Von Oswald J, et al. Transformers learn in-context by gradient descent. [2] Zhang R, et al. Trained transformers learn linear models in-context. [3] Shen L, et al. Do pretrained Transformers Really Learn In-context by Gradient Descent? [4] Dai D, et al. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers [5] Nathan T B, et al. In-context learning and gradient descent revisited

Reviewer cRiK2024-08-11

Thank you for clarifying the points and taking the time to do a comparison against each paper. The “dual” has indeed led to slight confusion and perhaps the link to previous should be made clearer. Also, generally, having this concise summary of comparisions would be handy in the paper as well. (As it helps understand the contribution very precisely.) —- Also, as a follow-up to the comparison to [4] 1. Could you elaborate on the exact difference between “supervised learning of the original model” and “self-representation learning of the ‘dual’ model”? Thank you.

Authorsrebuttal2024-08-12

Replying to New Official Comment by Reviewer cRiK

We sincerely appreciate your suggestions, which are greatly beneficial in improving our work! - We will further clarify in Section 2.3 that the term "dual" we use is linked to prior work [1], but differs from the "dual" concept in mathematical optimization, to avoid causing any potential misunderstanding for readers. - Additionally, we will also include comparisons with the aforementioned related works in the main body of our future revision to make our contributions more precisely. > **Question:** Could you elaborate on the exact difference between “supervised learning of the original model” and “self-representation learning of the ‘dual’ model”? Dai et al. [1] interpret ICL as implicit fine-tuning: More specifically, let $X = [X_D;X_T]$ where $X_D = [x_1;...;x_N]$ denotes the demonstration tokens and $X_T= [x'\_1;...;x'\_{T}]$ be query tokens. On the one hand, for ICL, they consider the output of $q = W_Qx’\_{T+1}$ under the linear attention setting as $$ \begin{aligned} \tilde{F}\_{\mathrm{ICL}} (q) &= W\_{V}[X_D; X_T](W_K[X_D; X_T])^T q \\\\ &= W_{V}X_{T}(W_KX_T)^Tq + W_{V}X_{D}(W_KX_D)^Tq \\\\ &= W_{\mathrm{ZSL}}q + \mathrm{LinearAtten}(W_VX_{D}, W_{K}X_{D}, q) \\\\ &= W_{\mathrm{ZSL}}q + \sum_{i}\left( (W_{V}x_{i}) \odot (W_{K}x_{i}) \right)q \\\\ &= W_{\mathrm{ZSL}}q + \Delta W_{\mathrm{ICL}}q , \end{aligned} \\ $$ where $W\_{\mathrm{ZSL}}q$ is interpreted as the output in the zero-shot learning (ZSL) where no demonstrations are given. On the other hand, they consider a specific fine-tuning setting, which updates only the parameters for the key and value projection, that is, $$ \begin{align} \tilde{F}\_{\mathrm{FT}}(q) &= (W_{V} + \Delta W_{V})XX^T(W_K + \Delta W_K)^Tq \\\\ &= (W_{\mathrm{ZSL}} + \Delta W_{\mathrm{FT}})q, \end{align} $$ where $\Delta W_K$ and $\Delta W_V$ denote the parameter updates and **they are "acquired by back-propagation from task-specific training objectives" [1], which is a supervised learning process of the original model**. Considering the similarity in form between $\tilde{F}\_{\mathrm{ICL}}$ and $\tilde{F}\_{\mathrm{FT}}$​, their focus is on establishing a connection between ICL and implicit fine-tuning on the original model. ------ As a comparison, we turn our attention to establish a connection between ICL and the gradient descent process of the dual model, rather than the original model. More specifically, we consider the dual model $f(x)=Wϕ(x)$ of the nonlinear attention layer, where the weight $W$ are updated according to the following loss (presented as Eq (9) in Section 3.2): $$ \mathcal{L} = -\frac{1}{\eta D}\sum_{i=1}^{N}(W_Vx_i)^TW\phi(W_Kx_i), $$ where $x_i$ is the i-th demonstration token. The prediction output of the trained dual model will be consistent with the ICL output of the attention layer. The gradient descent process of the dual model using this loss can be viewed from a self-supervised learning lens: unlike in supervised fine-tuning where the original model is instructed to perform gradient descent using a given objective (loss), **this loss is determined (derived) by the attention mechanism itself and it also does not require additional "true label" to supervise each token $x_i$​ (so called self-supervised)**. Therefore, modifications to this self-supervised learning loss will in turn cause modifications in the attention mechanism correspondingly, as we discussed in our work in Section 4. We believe this perspective offers several benefits: - By analyzing from the dual perspective, our analysis shifts to an optimization process, which makes it easier for us to use existing theoretical tools, as mentioned in our response to previous question ("A more detailed explanation why using the dual is important"). - This self-supervised learning process of the dual model can also inspire us to draw on existing works to modify the attention mechanism, as also mentioned in our previous response. - Intuitively, this explanation might be more reasonable as the original model is not explicitly instructed to provide the answer under some given objective (e.g., minimizing cross-entropy) during ICL inference process. Instead, the underlying criterion should be determined by the model's own structure (self-supervised) as we mentioned above. Thank you once again for your valuable suggestions! [1] Dai D, et al. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers

Reviewer cRiK2024-08-12

Thank you for all the clarifications and aiming for a clarified version. The confidence score is now updated from 3 to 4.

Authorsrebuttal2024-08-12

Replying to Official Comment by Reviewer cRiK

We are deeply encouraged by your acknowledgement, and we promise to improve our work in the future version as discussed previously. Once again, we sincerely thank you for your suggestions and contributions to the discussion process!

Authorsrebuttal2024-08-13

Looking forward to your reply

We are truly grateful for the time and effort you have invested in reviewing our manuscript and offering your valuable feedback. As the author-reviewer discussion phase is drawing to a close, we kindly wish to ensure that our recent detailed response, submitted a few days ago, has effectively addressed all of your concerns. If any issues remain or if further clarification is needed, please do not hesitate to reach out. We would greatly appreciate any additional feedback you may have!

Authorsrebuttal2024-08-13

Looking forward to your reply

We sincerely appreciate the time and effort you have dedicated to reviewing our manuscript and providing us with your valuable feedback! As the author-reviewer discussion phase is drawing to a close, we kindly wish to confirm whether our response has adequately addressed your concerns. A few days ago, we submitted a detailed response addressing your concerns and hope that they have adequately resolved any issues. If there are any remaining questions, please do not hesitate to let us know. We would greatly appreciate any additional feedback you may have!

Authorsrebuttal2024-08-14

Replying to Official Comment by Reviewer HnPE

We have received your feedback and would like to express our sincere gratitude for the time you dedicated to the review and for the valuable suggestions you shared with us!

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC