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

Recently, large language models (LLMs) based on the Transformer architectures (Vaswani et al., 2017) has shown surprising in-context learning (ICL) capabilities (Brown et al., 2020; Wei et al., 2022; Dong et al., 2022; Liu et al., 2023). By prepending several training examples before query inputs without labels, the models can make predictions for the queries and achieve excellent performance without any parameter updates. This excellent capability enables pre-trained LLMs such as GPT models to be used in general downstream tasks conveniently. Despite the good performance of the ICL capabilities, the mechanism of ICL still remains an open question.

In order to better understand the ICL capabilities, many works began to give explanations from different aspects. Xie et al. (2021) propose a Bayesian inference framework to explain how ICL occurs between pretraining and test time, where the LLMs infers a shared latent concept among the demonstration examples. Garg et al. (2022) demonstrate through experiments that pre-trained Transformer-based models can learn new functions from in-context examples, including (sparse) linear functions, two-layer neural networks, and decision trees. Zhang et al. (2023b) adopt a Bayesian perspective and show that ICL implicitly performs the Bayesian model averaging algorithm, which is approximated by the attention mechanism. Li et al. (2023) define ICL as an algorithm learning problem where a transformer model implicitly builds a hypothesis function at inference-time and derive generalization bounds for ICL. Han et al. (2023) suggest that LLMs can emulate kernel regression algorithms and exhibit similar behaviors during ICL. These works have provided significant insights into the interpretation of ICL capabilities from various perspectives.

In addition to the above explorations, there are also some attempts to relate ICL capabilities to gradient descent. Inspired by the dual form of linear attention proposed in Aiserman et al. (1964) and Irie et al. (2022), the ICL process is interpreted as implicit fine-tuning in the setting of linear attention by Dai et al. (2022). However, there is still a certain noticeable gap between linear attention and the widely used softmax attention. Additionally, this comparison is more of a formal resemblance and the specific details of gradient descent, including the form of the loss function and training data, require a more fine-grained exploration. Akyürek et al. (2022) show that by constructing specific weights, Transformer layers can perform fundamental operations (mov, mul, div, aff), which can be combined to execute gradient descent. Von Oswald et al. (2023a) adopt another construction, such that the inference process on a single or multiple linear attention layers can be equivalently seen as taking one or multiple steps of gradient descent on linear regression tasks. Building upon this weight construction method, subsequent work has conducted a more in-depth exploration of the capabilities of ICL under a causal setting, noticing that the inference of such attention layers is akin to performing online gradient descent (Ding et al., 2023; Von Oswald et al., 2023b). However, these analyses are still conducted under the assumption of linear attention and primarily focus on linear regression tasks, adopting specific constructions for the input tokens (concatenated from features and labels) and model weights. This limits the explanation of the Transformer’s ICL capabilities in more general settings. Thus, the question arises: Can we relate ICL to gradient descent under the softmax attention setting, rather than the linear attention setting, without assuming specific constructions for model weights and input tokens?

Motivated by the aforementioned challenges and following these works that connect ICL with gradient descent, we explore the ICL inference process from a representation learning lens. First, by incorporating kernel methods, we establish a connection between the ICL inference process of one softmax attention layer and the gradient descent process of its dual model. The test prediction of the trained dual model will be equivalent to the ICL inference result. We analyze the training process of this dual model from the perspective of representation learning and compare it with existing representation learning methods. Then, we derive a generalization error bound of this process, which is related to the number of demonstration tokens. Our conclusions can be easily extended to more complex scenarios, including a single Transformer layer and multiple attention layers. Furthermore, inspired by existing representation learning methods especially contrastive learning, we propose potential modifications to the attention layer and experiments are designed to support our findings.

In this section, we will address two questions discussed above: (i) Without assuming specific constructions for model weights and input tokens, how to relate ICL to gradient descent in the setting of softmax attention instead of linear attention? (2) What are the specific forms of the training data and loss function in the gradient descent process corresponding to ICL? In addressing these two questions, we will explore the gradient descent process corresponding to ICL from the perspective of representation learning.

Analyzing the dual gradient descent process of ICL from the perspective of representation learning inspires us to consider that: Do existing representation learning methods, especially contrastive learning methods, also involve a dual attention inference process? Alternatively, can we modify the attention mechanism by drawing on existing methods? In fact, since there are lots of mature works in representation learning especially contrastive learning, it is possible for us to achieve this by drawing on these works (He et al., 2020; Chen et al., 2020c; Wu et al., 2018; Chen et al., 2020a; Chen and He, 2021). We will provide some simple perspectives from the loss function, data augmentations and negative samples to try to adjust attention mechanism. It is worth noting that these modifications are also applicable to the self-attention mechanism, and we will explore these variants in experiments. More details can be seen in Appendix D.

Attention Modification inspired by the Contrastive Loss: It can be observed that the unnormalized similarity in Eq (9) allows $\|{\bm{W}}\|$ to be optimized to infinity if we ignore the Layer Normalization (LN) layer to prevent this. As for one single attention layer without LN layer, to address this issue, we can introduce regularization term to constrain the norm of ${\bm{W}}$, specifically by

where $\alpha$ is a hyperparameter. Equivalently, the attention output Eq (1) will be modified as

This modification is equivalent to retaining less prompt information for query token during aggregation and relatively more demonstration information will be attended to.

Attention Modification inspired by the Data Augmentation: If we analogize the key and value mappings to data augmentations in contrastive learning, then for the representation learning process of ICL, these overly simple linear augmentations may limit the model’s ability to learn deeper representations. Thus, more complicated augmentations can be considered. Denoting these two augmentations as $g_{1}$ and $g_{2}$, the loss function will be modified as

Correspondingly, the attention layer can be adjusted as,

where $g_{1}(\cdot)$ and $g_{2}(\cdot)$ will be column-wise here. Here we add augmentations for all tokens instead of only demonstration ones to maintain uniformity in the semantic space. In experiments, we simply select MLP for $g_{1}$ and $g_{2}$. It’s worth noting that here we only propose the framework, and for different tasks, the augmentation approach should be specifically designed to adapt them.

Attention Modification inspired by the Negative Samples: Negative samples play a crucial role in preventing feature collapse in contrastive learning methods while the representation learning process of ICL only brings a single pair of features closer, lacking the modeling of what should be pushed apart, which could potentially limit the model’s ability to learn representations effectively. Therefore, we can introduce negative samples to address this:

where $\mathcal{N}(i)$ is the set of the negative samples for ${\bm{x}}_{i}$ and $\beta$ is a hyperparameter. Correspondingly, the attention layer is modified as

where $\tilde{{\bm{X}}}_{D}=[\tilde{{\bm{x}}}_{1},\tilde{{\bm{x}}}_{2},...,\tilde{{\bm{x}}}_{N}]$. Here we simply use other tokens as negative samples and we emphasize that for specific tasks, an appropriate design of negative samples will be more effective.

In this section, we design experiments on synthetic tasks to support our findings and more experiments including on more realistic tasks can be seen in Appendix E. The questions of interest are: (i) Is the result of ICL inference equivalent to the test prediction of the trained dual model? (ii) Is it potential to improve the attention mechanism from the perspective of representation learning?

Linear Task Setting: Inspired by Von Oswald et al. (2023a), to validate the equivalence and demonstrate the effectiveness of the modifications, we firstly train one softmax self-attention layer using linear regression tasks. We generate the task by ${\bm{s}}={\bm{W}}{\bm{t}}$ where every element of ${\bm{W}}\in\mathbb{R}^{d_{s}\times d_{t}}$ is sampled from a normal distribution ${\bm{W}}_{ij}\sim\mathcal{N}(0,1)$ and ${\bm{t}}$ from uniform distribution ${\bm{t}}\sim U(-1,1)^{d_{t}}$. We set $d_{t}=11$ and $d_{s}=1$. Then, at each step, we use generated $\{{\bm{x}}_{i}=[{\bm{t}}_{i};s_{i}]\}^{N+1}_{i=1}$ to form the input matrix ${\bm{X}}$ where the last token will be used as the query token and the label part will be masked, that is, ${\bm{x}}_{N+1}=[{\bm{t}}_{i};0]$. Here we consider only one query token ($T=0$) and we denote ${\bm{x}}^{\prime}_{T+1}={\bm{x}}_{N+1}$ to maintain consistency of notation in Section 2.1. Finally, the attention layer is trained to predict $\hat{s}_{N+1}$ to approximate the true label $s_{N+1}$ using mean square error (MSE) loss.

Model Setting: It is worth noting that to facilitate direct access to the dual model, we use positive random features as kernel mapping functions (Performer architecture (Choromanski et al., 2020)) to approximate the standard softmax attention, that is, $\phi({\bm{x}})=e^{{\bm{w}}^{T}{\bm{x}}-\|{\bm{x}}\|^{2}/2}$ where ${\bm{w}}\sim\mathcal{N}(0,I)$. We set the dimension of the random features as $d_{r}=100(d_{t}+d_{s})=1200$ to obtain relatively accurate estimation. After training, the weights of the attention layer have been determined. Thus, given specified input ${\bm{X}}$, we can construct the dual model $f({\bm{x}})={\bm{W}}\phi({\bm{x}})$ and its corresponding training data and test input according to Theorem 3.1.

We perform three experiments under different random seeds for linear regression tasks with the results of one presented in Figure 3. In addition, we also conduct more experiments including these on trigonometric, exponential synthetic regression tasks and more realistic tasks. More details of experiments setting and results can be found in Appendix E. We mainly discuss the results on the linear regression task as follows.

Equivalence Between ICL and Gradient Descent: To answer the first question, we generate the test input ${\bm{X}}_{test}$ using the same method as training and obtain the ICL result of the query token ${\bm{h}}^{\prime}_{T+1}$. On the other hand, we use ${\bm{X}}_{test}$ to train the dual model according to Theorem 3.1 and get the test prediction $\hat{{\bm{y}}}_{test}$. The result is shown in the left part part of Figure 3. It can be observed that after $N=15$ epochs training on the dual model, the test prediction $\hat{{\bm{y}}}_{test}$ is exactly equivalent to the ICL inference result ${\bm{h}}^{\prime}_{T+1}$ by one softmax attention layer, which aligns with our analysis in Theorem 3.1. More detailed experiments can be seen in Appendix E.1.

Analysis on the Modifications: In Section 4, we discussed different modifications to the attention mechanism from perspectives of contrastive loss, data augmentation and negative samples. Here we call these modifications regularized models, augmented models and negative models respectively. More details of modifications for self-attention mechanism can be seen in Appendix D.

For regularized models, we vary different $\alpha$ to investigate the impact on pretraining performance under the same setting, as shown in the center left part of Figure 3. It can be observed that when $\alpha>0$, the regularized models converges to a poorer result while when $\alpha<0$, the model converges faster and achieves final results comparable to the normal model without regularization ($\alpha=0$). At least for this setting, this is a little contrary to our initial intention of applying regularization to the contrastive loss where $\alpha$ should be positive. We explain it that the appropriate $\alpha$ contributes to achieving a full-rank attention matrix as stated in Appendix D, preserving information and accelerating convergence.

For augmented models, we simply choose a single-layer MLP for $g_{1}(\cdot)$ and $g_{2}(\cdot)$ as data augmentations to enhance the value and key embeddings respectively in Eq (14) and we choose GELU (Hendrycks and Gimpel, 2016) as the activation function. It can be observed in the center right part of Figure 3 that when we only use $g_{2}$, that is, only provide augmentation for keys, the model actually shows slightly faster convergence than other cases. Furthermore, when we use two-layer MLP as $g_{2}^{+}({\bm{x}})$ as a more complicated augmentation function, the result indicates that although the model initially converges slightly slower due to the increased number of parameters, it eventually accelerates convergence and achieves a better solution. This indicates that appropriate data augmentation indeed have the potential to enhance the capabilities of the attention layer.

For negative models, we select the $k$ tokens with the lowest attention scores as negative samples for each token. From Eq (15), we can see that it is equivalent to subtracting a certain value from the attention scores corresponding to those negative samples. We vary the number of negative samples $k$ and $\beta$ in Eq (15) and the results are shown in the right part of Figure 3. It can be found that the model has the potential to achieve slightly faster convergence with appropriate settings ($k=3$ and $\beta=0.1$). In fact, it can be noted that in the original attention mechanism, attention scores are always non-negative, indicating that some irrelevant information will always be preserved to some extent. However, in the modified structure, attention scores can potentially become negative, which makes the model more flexible to utilize information. Certainly, as we discussed in Section 4, for different tasks, more refined methods of selecting augmentations and constructing negative samples may be more effective and we also leave these aspects for future.

Since Transformers have shown remarkable ICL abilities (Brown et al., 2020), many works have aimed to analyze the underlying mechanisms (Garg et al., 2022; Wang et al., 2023). To explain how Transformers can learn new tasks without parameter updates given few demonstrations, an intuitive idea is to link ICL with (implicit) gradient updates. The most relevant work to ours is that of Dai et al. (2022), which utilizes the dual form to understand ICL as an implicit fine-tuning (gradient descent) of the original model under a linear attention setting (Aiserman et al., 1964; Irie et al., 2022). They design a specific fine-tuning setting where only the parameters for the key and value projection are updated and the causal language modeling objective is adopted. In this context, they find ICL will have common properties with fine-tuning. Based on this, Deutch et al. (2024) investigate potential shortcomings in the evaluation metrics used by Dai et al. (2022) in real model assessments and propose a layer-causal GD variant that performs better in simulating ICL. As a comparison, our research also uses the dual form to analyze the nonlinear attention layer and explores the specific form of the loss used in the training process. However, we link ICL to the gradient descent performed on the dual model rather than fine-tuning the original model. The former process utilizes a self-supervised representation learning loss formalized as Eq (9) determined by the attention structure itself while performing supervised fine-tuning on the original model is often determined by task-specific training objectives (or manually specified causal language modeling objective Dai et al. (2022)). A more formal and detailed comparison can be found in Appendix F.

Additionally, many other works also link ICL with gradient descent, aiming to explore the Transformer’s ability to perform gradient descent algorithms to achieve ICL (Bai et al., 2023; Schlag et al., 2021). Akyürek et al. (2022) reveal that under certain constructions, Transformer can implement simple basic operations (mov, mul, div and aff), which can be combined to further perform gradient descent. Von Oswald et al. (2023a) provide a simple and appealing construction for solving least squares solutions in the linear attention setting. Subsequently, Zhang et al. (2023a); Ahn et al. (2023); Mahankali et al. (2023) provide theoretical evidence showing that the local or global minima will have a form similar to this specific construction proposed by Von Oswald et al. (2023a) under certain assumptions. These works, both experimentally and theoretically, often focus on specific linear regression tasks ($y={\bm{w}}^{T}{\bm{x}}$) and specific structured input format where each token takes the form $[{\bm{x}},y]$ consisting of the input part ${\bm{x}}$ and the label part $y$. In addition, the label part of the final query to be predicted is masked, represented as $[{\bm{x}},0]$. Subsequent works have expanded this exploration under more complicated setups, including examining nonlinear attention instead of linear attention(Cheng et al., 2023; Collins et al., 2024), using unstructured inputs rather than structured ones(Xing et al., 2024), and considering casual or autoregressive setting(Ding et al., 2023; Von Oswald et al., 2023b). As a comparison to these works, 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 specific input forms. Instead, we aim to view the ICL inference process from the perspective of representation learning in the dual model. However, we would like to point out that under these specific weight and input settings, an intuitive explanation can also be provided from a representation learning perspective (see Appendix F). We also notice that Shen et al. (2023) experimentally show that there may exist differences between ICL inference in LLMs and the fine-tuned models in real-world scenarios from various perspectives and assumptions used in previous works may be strong. As mentioned earlier, our analysis primarily focus on linking ICL with gradient descent on the dual model of a simplified Transformer rather than fine-tuning the original model. Analyzing more realistic models will also be our future directions.

In this paper, we establish a connection between the ICL process of Transformers and gradient descent of the dual model, offering novel insights from a representation learning lens. Based on this, we propose modifications for the attention layer and experiments under our setup demonstrate their potential. Although we have made efforts in understanding ICL, there are still some limitations in our analysis: (1) our work primarily focuses on the simplified Transformer and the impact of structures like layer normalization, residual connections, and others requires more nuanced analysis; (2) for more tasks and settings, the proposed model modifications may require more nuanced design and validation. We leave these aspects for future exploration. And we believe that this work mainly studies the theory of in-context learning, which does not present any foreseeable societal consequence.

We sincerely appreciate the anonymous reviewers for their helpful suggestions and constructive comments. This research was supported by National Natural Science Foundation of China (No.62476277, No.6207623), Beijing Natural Science Foundation (No.4222029), CCF-ALIMAMA TECH Kangaroo Fund (No.CCF-ALIMAMA OF 2024008), and Huawei-Renmin University joint program on Information Retrieval. We also acknowledge the support provided by the fund for building worldclass universities (disciplines) of Renmin University of China and by the funds from Beijing Key Laboratory of Big Data Management and Analysis Methods, Gaoling School of Artificial Intelligence, Renmin University of China, from Engineering Research Center of Next-Generation Intelligent Search and Recommendation, Ministry of Education, from Intelligent Social Governance Interdisciplinary Platform, Major Innovation & Planning Interdisciplinary Platform for the “DoubleFirst Class” Initiative, Renmin University of China, from Public Policy and Decision-making Research Lab of Renmin University of China, and from Public Computing Cloud, Renmin University of China.