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

We thank the reviewer for acknowledging the novel perspective on understanding ICL of our paper. Considering the word limits of rebuttal, we notice that some issues might overlap, so we have organized our responses as follows.

Weakness 1 (The paper relies on several assumptions and simplifications...)
& Question 1 (How would the inclusion of layer normalization and residual connections affect the theoretical framework and findings?)

As for residual connection, intuitively, our conclusions can be naturally extended to versions with residual connections, which are analogous to introducing a residual connection structure in dual models, i.e., $\mathrm{Atten}(x) + x = \hat{W}\phi(W_Qx) + x$, where $\hat{W}$ is the weight in the dual model $f(x) = W\phi(x)$ trained under the original representation loss in Eq (9). This modification does not have a significant impact on the specific representation loss or generalization bound, but its effect on the GD process of the dual model still requires more detailed analysis.

As for layer normalization, from the representation learning perspective of the dual model, it normalizes $x$ to a scaled unit sphere to make the process more stable intuitively. However, we acknowledge the analytical challenges posed by the nonlinearity of layer normalization, and the more detailed effects still warrant further investigation. This may require more powerful theoretical tools, e.g. [1], which is also a direction for our future research.

[1] On the Nonlinearity of Layer Normalization.

Weakness 2 (Lack of discussions towards some related works that argue against the equivalence between ICL and gradient descent...)

Following the previous work [1], the first mentioned work [2] further explores potential shortcomings in the previously used evaluation metrics and further points out that the ICL output is more dependent on previous lower layers, which is different from the finetuning that relies on all model parameters. Inspired by this, it introduces a new finetuning method called LCGD and validates the claim with new metrics.

Unlike previous works including [3,4], which studies ICL and Gradient Descent (GD) under strong constraints (specific regression tasks and weight constructions, etc.), the second mentioned work [5] highlights significant gaps between these strong constraints and real-world LLMs. They observe that there are differences in how ICL and GD modify the output distribution and emphasize the distinction between naturally emergent ICL (pretrained on natural text data) and task-specific ICL (such as in [3,4]). It provides a clearer definition and experimental results in more realistic settings.

Although our work also relates ICL to GD, our focus is on how we can view the ICL inference process as a GD process on a dual model from a representation learning perspective:

• To simplify theoretical analysis, we concentrate on analyzing the core component of the Transformer architecture—the attention layer—while neglecting residual connections and layer normalization. This does introduce some divergence from real-world large language models;
• In terms of research approach, we focus on the GD process of dual model under representation learning loss, rather than analyzing GD directly on the original model as in [2,5]. This is also a significant departure from the mentioned works.
• Additionally, the representation learning process of the dual model resembles a self-supervised process, which we believe is reasonable: unlike fine-tuning, ICL inference is not explicitly provided specific metrics to adapt to the target task and this also aligns with certain aspects of previous work [5]. And we also compare this process with existing self-supervised learning methods in our work.
• Inspired by this representation learning perspective, we propose modifications to the attention mechanism, which were not sufficiently explored in previous work.

Thank you for your suggestion and we will add more supplementary information on the relevance between our work and the mentioned works in our future revision.

[1] Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers

[2] In-context Learning and Gradient Descent Revisited

[3] Transformers Learn In-context by Gradient Descent

[4] Trained Transformers Learn Linear Models In-context

[5] Do pretrained Transformers Learn In-Context by Gradient Descent?

Question 2: How do the proposed methods perform on more complex and diverse tasks beyond the linear regression tasks used in the experiments?

For more complex tasks, considering the limited computational resources and time constraints, we select the pre-trained BERT-base-uncased model to apply attention modifications, and validate the results on part of GLUE datasets. More experimental settings are detailed in the (global) author rebuttal and the results are presented in the Table 1 of PDF (To AC and All Reviewers).

For the regularized models, we select different values $\alpha$ and conclude that smaller absolute values of $\alpha$ are recommended. For the augmented models, we adopt a parallel MLP approach for data augmentation $g_1/g_2$, that is, $g_1(W_Vx) = W_Vx + c W_2\sigma(W_1x)$​. The results suggest that using $g_1$ alone may be a more effective choice under our settings. For the negative sample models, we continued to select tokens with lower attention scores as negative samples. The final results showed limited improvement in model performance. We think that this may be due to the ineffective selection of negative sample tokens, and better methods for selecting negative samples would be interesting to explore in the future. In conclusion, these results validate the potential of our proposed methods.

Final Note: We want to thank you again for all the questions you have provided. If there are any remaining questions, please do not hesitate to let us know.

We sincerely thank the reviewer for acknowledging the theoretical contribution of our paper. We have carefully addressed each of your concerns and considering the word limits of rebuttal, we notice that some issues might overlap, so we have organized our responses as follows.

Weakness 1: Slightly simplified transformer architecture is used (it would be interesting to see the skip-connection version) ...

The attention module, as one of the core components of the transformer, is crucial for extracting contextual information. Therefore, for the sake of analytical simplicity, we focus on comparing pure attention modules/mechanisms. Intuitively, our conclusions can be naturally extended to versions with residual connections, which are analogous to introducing a residual connection structure in dual models, i.e., $\mathrm{Atten}(x) + x = \hat{W}\phi(W_Qx) + x$, where $\hat{W}$ is the weight in the dual model $f(x) = W\phi(x)$​ trained under the original representation loss in Eq (9). Although it is relatively straightforward to connect residual connections with dual models in form, the role of such residual connections in the learning process of the dual model's weight from a representation learning perspective still requires further investigation. Therefore, a more reasonable explanation is currently lacking. We believe this is an interesting direction for future research.

Weakness 2 (The experimental setup and results are a bit limited ... as well as a somewhat "realistic" task would be interesting to see - and what the results mean qualitatively. )
& Question 1 （How could a "realistic" task look like for your evaluation?）

Due to constraints of page sizes, we have simplified the description of the experiments for simulation tasks in the main body. In fact, more detailed experimental setups are presented in the Appendix E. We will provide more detailed supplementary information on the experimental setups if space permits in future revisions.

For the evaluation on more realistic tasks, we have supplemented our work with additional experiments. More experimental settings are detailed in the (global) author rebuttal and the results is presented in the Table 1 of PDF (To AC and All Reviewers).

Regarding the qualitative results: For the regularized models, smaller absolute values of $\alpha$ are recommended. We interpret this as adding a small scaled identity matrix to the attention score matrix helps achieve full rank, thereby better preserving information. For the augmented models, to better retain the information from the pre-trained $W_V$ and $W_K$ weights, we adopted a parallel MLP approach for data augmentation $g_1$ and $g_2$. For example, $g_1(W_Vx) = W_Vx + c W_2\sigma(W_1x)$. The results suggest that enhancing the value mapping may be a more effective choice. For the negative sample models, we continued to select tokens with lower attention scores as negative samples while the final results showed limited improvement in model performance. We believe this may be due to the ineffective selection of negative sample tokens, and better methods for selecting negative samples would be interesting to explore in the future.

Weakness 3 (A small section on how these results can be applied for bigger models in practice.)
& Question 2(How could one practically extend your augementations to e.g. modern LLMs ...)

We believe that our supplementary experiments on pre-trained BERT provide more insights into this issue (see global rebuttal (To AC and All Reviewers)). Although the scale of the model and datasets selected is relatively small due to our limited computational resource and time constraints, we believe similar approaches would naturally extend to larger-scale experiments.

Take the augmented models as the example: we also considered using more complex data augmentations for the linear key/value mapping ($g_1(W_Vx)$ and $g_2(W_Kx)$). However, unlike previous methods used in simulation tasks, we do not choose $g_1/g_2$ as $g_1(W_Vx)=W_2σ(W_1W_Vx)$, because our experiments showed that this design struggles to fully leverage the pre-trained weights $W_V$ and $W_K$. Thus, we adopted a parallel approach, that is, $g_1(W_Vx) = W_Vx + c W_2\sigma(W_1x)$, where $c$ is a hyperparameter to control the influence of the new branch ($g_2$ is similar). This approach introduces nonlinear "augmentation" while preserving the knowledge from the original pre-trained weights, making it easier to train. The results in Table 1 also prove the effectiveness of this method especially using $g_1$ alone, which leads to the most significant improvement in model performance.

Additionally, we do not rule out the possibility of more powerful and efficient augmentation methods. The parallel MLP approach is chosen primarily to make better use of the pre-trained weights $W_V$/ $W_K$ and this design is quite general and does not consider specific characteristics of particular tasks. We still encourage the design of more task-specific augmentation format for different tasks. For example, in CV tasks, it might be natural to incorporate CNN-extracted features into the $g_1$/$g_2$ in ViTs.

Similarly, our discussion on the negative sample models further supports the need for more task-specific designs. As shown in Table 1, selecting tokens with low attention scores as negative samples is a rather crude approach, leading to relatively limited performance improvements. Exploring how to select or construct higher-quality negative samples is also an interesting direction. These detailed considerations go beyond the scope of this paper, and these interesting directions will be the focus of our future work.

Final Note: We are excited that you acknowledge our theoretical contribution. If there are any remaining questions, please do not hesitate to let us know. Thank you once again for your insightful comments and for your encouraging feedback!

We sincerely appreciate your recognition of the novelty of our paper. We have carefully addressed each of your concerns and considering the word limits of rebuttal, we notice that some issues might overlap, so we have organized our responses as follows.

Weakness 1 (The work may ignore the format of task input/output. And some of recent work...)
& Question 5 (The dynamics of token interactions may not be well supported in the experiments.)

The mentioned work [1] extends the original structured input format [2,3] to the unstructured setting and then explores the impact of various factors (one/two layers, look-ahead attention mask, PE) when learning linear regression tasks both theoretically and experimentally. The explanation for ICL derived from these interesting findings is more aligned with real-world scenarios (where the examples are usually stored in different tokens).

Although our work also relates ICL to GD, our focus is on how we can view the ICL inference process as a GD process on a dual model from a representation learning perspective. We do not make assumptions about the input format, such as $[x; y]$ or $[x; 0], [0; y]$, as we do not target specific tasks (e.g. linear regression tasks) but aim to provide more general explanations. To simplify the analysis, our input tokens can be viewed as embeddings of demonstration sentences (or intermediate states among multi layers).

Correspondingly, in the experimental section, our focus is not on the dynamics of token interactions as specific tasks are not being targeted. Instead, we are more concerned with whether the result of ICL inference $h'\_{T+1}$ is equivalent to the test prediction $\hat{y}_{test}$​ of the trained dual model, as illustrated in the left part of Figure 3. This result aligns with our analysis in Theorem 3.1.

[1] Theoretical Understanding of In-Context Learning in Shallow Transformers with Unstructured Data

[2] Transformers learn in-context by gradient descent

[3] Trained transformers learn linear models in-context

Weakness 2 (The experimental section is not well-organized and self-explained...)
& Question 2 (As mentioned in weeknesses part...)

Thank you for your careful suggestions! Due to the resrtiction of page sizes, we simplified the description of the experiments. In fact, more details of our experiments are presented in the Appendix E. We will reorganize our experimental part and provide more details in our future revision.

Question 1 (Will the ratio of the number of negative samples have an impact on generalization bound?)

The answer is yes. After introducing negative samples, we consider the following representation loss:

$$\mathcal{L}(f) = \mathbb{E}_{x\sim\mathcal{D}\_{\mathcal{T}}} \left[- \frac{1}{K} \sum\_{j = 1}^{K} \left( W\phi(W_Kx) \right)^T\left(W_Vx - W_Vx^{-}\_{j} \right)\right],$$

where $K$ is the number of negative samples for each $x_i$ and $x\_{j}^-$ denotes the $j$-th negative sample for token $x$. The empirical loss is modified correspondingly. Then, by the other definitions in Section 3.3, we can obtain the generalization bound as

$$\mathcal{L}(\hat{f})\le \mathcal{L}(f) + O\left(w\rho d\_o \sqrt{\mathrm{Tr}(K\_S)\left(\frac{5}{N^2} + \frac{1}{rN^3}\right)} + \sqrt{\frac{log\frac{1}{\delta}}{N}}\right),$$

where $r = \frac{K}{N}$ is the ratio of the number of negative samples. It can be observed that as the ratio increases, the generalization error will decrease. However, we also notice that $\frac{5}{N^2} > \frac{1}{rN^3}$​​ thus the reduction in generalization error due to an increased proportion of negative samples is limited. The proof is similar to that of Theorem 3.2 (see Appendix C). The main difference is that we need to define the new function class $G$ and $F$ according to above definition. Nevertheless, we do not rule out the possibility of a tighter generalization bound using better theoretical tools.

Question 3 (Since the authors ... which one could be the best ... need add these three together? why...?)

The answer is related to the specific tasks but generally speaking, using all three methods together is not necessary, especially when the design of augmentated or negative modification is not effective enough. We conduct experiments with their combinations on linear, trigonometric, and exponential tasks. The results are shown in Figure 1 of the PDF in To AC and All Reviewers. For the latter two tasks, the results are actually worse than using regularized or augmented modification individually (compared to Figures 11 and 12 in Appendix E). Therefore, when the design of augmented or negative mofification is not effective, we recommend using a single modification method.

Question 4 (Does the conclusion hold the same for both one-layer transformer and multi-layer transformer?)

For more scenarios, the conclusion still holds: using all three methods simultaneously is not necessary. We supplemented our experiments using BERT (multi-layer Transformer), see (To AC and All Reviewers). From Table 1 in the PDF, the combined models do not outperform the augmented version alone. We also observed that negative sample improvements remain quite limited, indicating that our method for selecting negative samples is not effective enough. In these cases, using augmented modification alone will be a better choice.

Question 6 (seems missing a conclusion section)

We will streamline our paper content and supplement the conclusion section based on the valuable questions raised by all you reviewers in the future revision. Thank you very much for your reminder!

Final Note: We want to thank you again for all the questions you have provided. If there are any remaining questions, please do not hesitate to let us know.

We thank the reviewer for the encouraging feedback, especially for recognizing the theoretical contribution and insights of our paper. Our response is detailed below.

Weakness 1: 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.

In fact, our findings can be naturally extended to more realistic tasks. We supplement our experiments on more general language tasks. More experimental settings are detailed in the (global) author rebuttal and the results is presented in the Table 1 of PDF (To AC and All Reviewers). Considering the limited computational resources and time constraints, we choose the BERT-base-uncased model to and select part of GLUE datasets to validate the effectiveness of modifications to the attention mechanism. More discussions on the results are provided in detail in our response to Weakness 2.

Weakness 2: 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.

To answer the questions, below we discuss the main experimental results presented in the Table 1 of PDF (To AC and All Reviewers).

For the regularized modification, we consider different values of $\alpha$, and as can be observed in Table 1, except for RTE, the best regularized models outperform the original model on the other three datasets. while we also note that when the absolute value of $\alpha$ is too large, the model's performance declines significantly, so smaller absolute values for $\alpha$ is recommended.

For the augmented modification, we also consider applying more complex ''augmentation'' functions to the linear key/value mappings. Specifically, we adopt a parallel approach, i.e., $g_1(W_Vx) = W_Vx + c W_2\sigma(W_1x)$, where $c$ is a hyperparameter to control the influence of the new branch, $\sigma$ is the GELU activation function and the hidden layer dimension is set to twice the original size of $W_Vx$. And $g_2(W_Kx) = W_Kx + c W_2\sigma(W_1x)$ follows the same format. Experimental results show that the best augmented models achieve better performance than the original model across all four datasets. Notably, using $g_1$ alone proves to be more effective than other methods and using both $g_1$ and $g_2$ introduces more parameters, which is particularly significant for larger models. Thus, under the augmentation methods and experimental settings we selected, using $g_1$​ alone is recommended.

For the negative modification, we continue to select tokens with lower attention scores as negative samples. The parameter $r$ represents the ratio of tokens used as negative samples, while $\beta$ indicates the overall reduction in attention scores. And the results show that the best negative models only outperform the original model on CoLA and STS-B, whereas their performance on MRPC and RTE is worse than the original one. This suggests that our simple approach of considering tokens with low attention scores as negative samples might be too coarse. A more effective method for constructing negative samples should be designed, which is a direction worth exploring in the future.

We also consider combining different modification methods. The results indicate that under our settings, the combination of augmented and negative modification achieves the best performance on CoLA, MRPC, and RTE, while the combination of regularized and augmented modification achieves the best performance on STS-B. However, their optimal performance is slightly inferior to the best performance achieved with augmented models alone. Therefore, we conclude that using all three modifications simultaneously is not necessary. With appropriate hyperparameter choices, using augmented modification alone or in combination with one other modification is sufficient.

In conclusion, the experimental results show the potential of our approach of improving the attention mechanism from a representation learning perspective. However, due to time and computational resource limitations, our experiments are conducted on a limited tasks and model size. More detailed parameter searches, validation across additional tasks and models, and the development of task-specific augmentation and negative sampling methods are all interesting directions worth exploring in the future.

Final Note: Thank you for your detailed review. We are excited that you found our novel insights into attention mechanism. If there are any remaining questions, please do not hesitate to let us know.

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!