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

Authors: ruifeng  ren
yong  liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 09:07:50

Abstract: 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.

Main paper: 
 introduction: 
 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 inferencetime 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-context learning with transformers: 
 The model we consider is composed of many stacked Transformer decoder layers, each of which is composed of an attention layer and a FFN layer. For simplicity, we have omitted structures such as residual connections and layer normalization, retaining only the most essential parts. We consider the standard ICL scenario, where the model's input consists of demonstrations followed by query inputs, that is, the input can be represented as X = [X D , X T ] ∈ R di× (N +T ) , where X D = [x 1 , x 2 , ..., x N ] denotes N demonstration tokens, and X T = [x ′ 1 , x ′ 2 , ..., x ′ T ] denotes T query tokens. Here, we focus more on how tokens interact during model inference while ignoring the internal structure of demonstration tokens. For the query input at position T + 1, its output after one layer of Transformer can be represented asEQUATIONx ′ T +1 = W 2 ReLu(W 1 h ′ T +1 + b 1 ) + b 2 ,(2) where W K , W Q , W V ∈ R do×di are parameters for key, query, value projections andW 1 ∈ R d h ×do ,W 2 ∈ R do×d h ,b 1 ∈ R d h , b 2 × R doare FFN parameters. Our concern is how the query token x ′ T +1 learns in-context information from demonstrations. Unlike previous work [Von Oswald et al., 2023a , Zhang et al., 2023a , Bai et al., 2023 ], here we do not make additional assumptions about the structure of input matrix X and parameters to study the Transformer's ability to implement some specific algorithms. Instead, we adopt the same setting as [Dai et al., 2022] to study more general cases.Representation learning aims to learn embeddings of data to preserve useful information for downstream tasks. One class of methods most relevant to our work is probably contrastive learning methods without negative samples [Chen and He, 2021 , Grill et al., 2020 , Caron et al., 2020 , Tian et al., 2021] . Contrastive learning is a significant approach of self-supervised learning (SSL) which aims at learning representations by minimizing the distance between the augmentations of the same data point (positive samples) while maximizing the distance from different data points (negative samples) [He et al., 2020 , Chen et al., 2020b , Oord et al., 2018 , Oh Song et al., 2016] . To alleviate the burden of constructing a sufficient number of negative samples while avoiding representational collapse, some works propose architectures for contrastive learning without negative samples, which mainly use weight-sharing network known as Siamese networks [Chen and He, 2021 , Grill et al., 2020 , Caron et al., 2020 , Tian et al., 2021] . The architecture takes two augmentations x 1 , x 2 from the same data x as inputs, which will be processed by online network and target network respectively to obtain the corresponding representations, that is, x1 = f online (x 1 ), x2 = f target (x 2 ). The two encoder networks share weights directly or using Exponential Moving Average (EMA). Then, x1 will be input into a predictor head to obtain the predictive representation z 1 = g( x1 ). Finally, we minimize the distance between the predictive representation and target representation, that is, L (z 1 , StopGrad( x2 )) where StopGrad(•) means x2 is treated as a constant during backpropagation. For L(•), we often choose the cosine similarity or the l 2 -norm as a measure of distance, although they are equivalent when the vector is normalized. Another class similar to our work is kernel contrastive learning [Esser et al., 2024] . Given an anchor x and its positive and negative samples x + , x -, it aims to optimize the loss function L = f (x) T (f (x -) -f (x + )), where f (x) = W ϕ(x) and ϕ(x) is the feature mapping for some kernel. We will consider the gradient descent process corresponding to the inference process of ICL from the perspective of representation learning and compare it with the two aforementioned representation learning patterns.It has been found that the linear attention can be connected to the linear layer optimized by gradient descent [Aiserman et al., 1964 , Irie et al., 2022 , Dai et al., 2022] , that is, the gradient descent on linear layer can be seen as the dual form2 of linear attention. A simple linear layer can be defined as f L (x) = W x, where W ∈ R do×di is the projection matrix. Given training inputs [x i ] N i=1 ∈ R di with their labels [y i ] N i=1 ∈ R do , a linear layer can output the predictions [ ŷi ] N i=1 where ŷi = W x i and then compute certain loss L( ŷi , y i ) for training. Backpropagation signals [e i ] N i=1 ∈ R do will be produced to update W in gradient descent process where e i = -η (∇ ŷi L) if we set η as the learning rate. During test time, the trained weight matrix W can be represented by its initialization W 0 and the updated part ∆W , that is,EQUATIONwhere ⊗ denotes the outer product according to the chain rule of differentiation. On the other hand, this process can be viewed from the perspective of linear attention. Let[k i ] N i=1 , [v i ] N i=1 ∈ R didenote the N key and value vectors constituting matrices K, V ∈ R di×N respectively. For a given query input q ∈ R di , linear attention is typically defined as the weighted sum of these value vectorsLA(V , K, q) = V K T q = N i=1 v i k T i q = N i=1 v i ⊗ k i q.Then, we can rewrite the output of a linear layer during test time as i=1 can be considered as values while [x i ] N i=1 as keys and x test as the query. This interpretation uses gradient descent as a bridge to connect predictions of linear layers with linear attention, which can be seen as a simplified softmax attention used in Transformers.EQUATIONInspired by this relationship, Dai et al. [2022] understand ICL as implicit fine-tuning. However, this interpretation based on linear attention deviates from the softmax attention used in practical Transformers. Furthermore, this alignment is also ambiguous as the specific details of the gradient descent process, including the form of loss function and dataset, have not been explicitly addressed. In addition , Von Oswald et al. [2023a] , Ding et al. [2023] also connect ICL with gradient descent for linear regression tasks using weight construction methods, where parameters W K , W Q and W V of the self-attention layer need to roughly adhere to a specific constructed form. However, these analyses rely on the setting of linear regression tasks and assumptions about the form of input tokens (concatenated with features and labels), which limits the interpretability of ICL capabilities from the perspective of gradient descent. Thus, we attempt to address these issues in the following sections. 
connecting icl with gradient descent: 
 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.Before we begin establishing the connection between ICL and gradient descent, we need to firstly rethink softmax attention with kernel methods. Dai et al. [2022] connect ICL with gradient descent under the linear attention setting. In fact, it is completely feasible to interpret ICL under softmax attention with the help of kernel methods. We define the attention block asEQUATIONwhich can be viewed as the product of an unnormalized part A u and a normalizing multiplier D, that is,EQUATIONwhere exp(•) is element-wise. Similar in [Choromanski et al., 2020] , we define softmax kernel K sm :R do × R do → R + as K sm (x, y) = e x T y = e ∥x∥ 2 +∥y∥ 2 2K guass (x, y) where K guass = e -∥x-y∥ 2 /2 is the guassian kernel when the variance σ 2 = 1. According to Mercer's theorem [Mercer, 1909] , there exists some mapping function ϕ : R do → R dr satisfying that K sm (x, y) = ϕ(x) T ϕ(y). Thus, noting that when omitting the √ d o -renormalization and equivalently normalize key and value vectors in Eq (6), every entry in the unnormalized part A u can be seen as the output of softmax kernel K sm defined for the mapping ϕ, which can be formulated as:A u (i, j) = exp (W K x i ) T W Q x j = K sm (W K x i , W Q x j ) = ϕ(W K x i ) T ϕ(W Q x j ). (7)There have been many forms of mapping function ϕ(•) used in linear Transformers research to approximate this non-negative kernel [Choromanski et al., 2020 , Katharopoulos et al., 2020 , Peng et al., 2021 , Lu et al., 2021] . For example, we can choose ϕ(•) as positive random features which has the form ϕ(x) = e w T x-∥x∥ 2 /2 to achieve unbiased approximation [Choromanski et al., 2020] . Alternatively, we can also choose ϕ(x) = elu(x) + 1 proposed by Katharopoulos et al. [2020] .Now, we begin to establish the connection between the ICL inference process of a softmax attention layer and gradient descent. We focus on a softmax attention layer in a trained Transformer model, where the parameters {W Q , W K , W V } have been determined and the input X = [X D , X T ] has the form introduced in Section 2.1. Then, after the inference by one attention layer, the query token at position T + 1 will have the form h ′ T +1 formulated by Eq (1). On the other hand, given a specific softmax kernel mapping function ϕ(x) that satisfies Eq (7), we can define the dual model for the softmax attention layer asEQUATIONwhere W ∈ R do×dr is parameters. We assume that the dual model obtains its updated weights W after undergoing one step of gradient descent with some loss function L. Subsequently, when we takez test = W Q x ′ T +1as the test input, we can obtain its test prediction asŷtest = f (z test ) = f W Q x ′ T +1 = W ϕ W Q x ′ T +1 .We will show that h ′ T +1 in Eq (1), is strictly equivalent to the above test prediction ŷtest , which implies that the inference process of ICL involves a gradient descent step on the dual model. This can be illustrated by the following theorem: Theorem 3.1. The query token h ′ T +1 obtained through ICL inference process with one softmax attention layer, is equivalent to the test prediction ŷtest obtained by performing one step of gradient descent on the dual model f (x) = W ϕ(x). The form of the loss function L is:EQUATION)where η is the learning rate and D is a constant.Proof can be found in Appendix A. Theorem 3.1 demonstrates the equivalence between the ICL inference process and gradient descent. Below, we delve into more detailed discussions:Training Set and Test Input: In fact, once the attention layer has already been trained, that is, W K , W Q , W V has been determined, the demonstration tokens [x i ] N i=1 will be used to construct a training set for the dual model. Specifically, the training data has the form {z(i) std , y (i) std } N i=1 where z (i) std = W K x i as inputs and y (i) std = W V x i as their labels. During training stage, for each input z (i) std , the dual model outputs its prediction ŷ(i) = f z (i) std = W ϕ z (i) std = W ϕ (W K x i ).Then, the loss function Eq (9) can be rewritten as L = -1 ηD N i=1 (y (i) std ) T ŷ(i) , which can be regarded as the cosine similarity. Then, using this loss function and the training data, we can perform one step of Stochastic Gradient Descent (SGD) on the dual model and obtain the updated W . Finally, during the testing stage, we take z test = W Q x ′ T +1 as the test input to get its prediction which will be consistent with the ICL result h′ T +1 , that is, ŷtest = f (z test ) = W ϕ W Q x ′ T +1 = h ′ T +1. This process can be illustrated in Figure 1 . Demonstration tokens provide information about the training data points and the weight matrix W is optimized to learn sufficient knowledge about demonstrations. This gradient descent process using the loss function L applied to f (x) can be seen as the dual form of the ICL inference process of the attention layer.Representation Learning Lens: Even though we have now clarified the details of the gradient descent process of ICL, what does this process more profoundly reveal to us? In fact, for a encoded demonstration token x i , the key and value mapping will generate a pair of features W K x i and W V x i that exhibit a certain distance from each other, akin to positive samples in contrastive learning. And then, ϕ(x) projects W K x i into a higher-dimensional space to capture deeper features. Finally, the weight matrix W , which maps ϕ(W K x i ) back to the original space, is trained to make the mapped vector as close as possible to W V x i . This process is illustrated in Figure 2 . Below, we attempt to understand this process from the perspective of existing representation learning methods introduced in Section 2.2, although we emphasize that there are certain differences between them.Comparison with Contrastive Learning without Negative Samples: If we consider the key and value mapping as two types of data augmentation, then from the perspective of contrastive learning without negative samples, this process can be similarly formalized asmin W L ŷ(i) , y (i) std = L ŷ(i) , StopGrad(y (i) std ) ,where StopGrad(•) is naturally applicable because there are no learning parameters involved in the generation process of the representation y (i) std . However, it's important to note that the representation learning process of ICL is much simpler: Firstly, the online and target networks are absent while the augmentations W K x i , W V x i are directly used as online and target representations respectively. Secondly, the predictor head is useful and not discarded, which is then used during test stage.Comparison with Contrastive Kernel Learning: Given an anchor data x and its positive and negative samples x + , x -, contrastive kernel learning aims to optimize the loss functionL = f (x)(f (x -) -f (x + )) where f (x) = W ϕ(x).There are significant differences in the representation learning process of ICL: Firstly, it does not involve negative samples. Secondly, there is no corresponding processing for positive samples, leading to parameter updates being solely dependent on the processing of the anchor.Extension to More Complicated Scenarios: Theorem 3.1 can be naturally extended to one single Transformer layer and multiple attention layers. As for one Transformer layer formed in Section 2.1, its dual model f + (x) = W ϕ(x) + b introduces an additional bias b and only W is trained while b remains fixed. In addition, the labels of training set will be y(i) std = W F W K x iwhere W F has potential low-rankness property induced by Relu(•). As for multiple attention layers, the ICL inference process will be equivalent to sequentially performing gradient descent and making predictions on the dual model sequence. We provide more details in Appendix B.Compared to Dai et al. [2022] considering the connection under linear attention setting, Theorem 3.1 gives explanation for more generally used softmax attention and offers a more detailed exploration of the training process. Additionally, unlike Von Oswald et al. [2023a,b] , Ding et al. [2023] 's focus on particular linear regression task and specific configurations of token and parameters, we aim to explain the process of token interactions during ICL inference in a more general setting.In this part, we are interested in the generalization bound of the ICL gradient process. When ICL inference is performed for some task T , we cannot provide all demonstrations related to task T limited by the length of input tokens. We denote S T ⊆ R di as all possible tokens for the task T and assume that these tokens will be selected according to the distribution D T . During a particular instance of ICL inference, let S = {x i } N i=1 ⊆ S T represent the example tokens we selected. We define the function class as F := {f (x) = W ϕ(W K x) | ∥W ∥ ≤ w} where ∥ • ∥ denotes the Frobenius norm. Generally, ignoring constant term in Eq (9), we consider the representation learning loss asEQUATIONwhere f ∈ F and D T is the distribution for some ICL task T . Correspondingly, the empirical loss will be formulated asL(f ) = -1 N N i=1 (W V x i ) T f (x i ) and we have f = arg min f ∈F L(f ).In addition, we denote the kernel matrix of demonstration tokens S as K S ∈ R N ×N where (K S ) i,j = ⟨ϕ(W K x i ), ϕ(W K x j )⟩, that is, the inner product of the feature maps after W K projection between the i-th token and j-th token. We state our theorem as follows: Theorem 3.2. Define the function class as F := {f (x) = W ϕ(W K x) | ∥W ∥ ≤ w} and let the loss function defined as Eq (10). Consider the given demonstration set as S = {x i } N i=1 where S ⊆ S T and S T is all possible demonstration tokens for some task T . With the assumption that ∥W V x i ∥, ∥W ϕ(W K x i )∥ ≤ ρ, then for any δ > 0, the following statement holds with probability at least 1 -δ for anyEQUATIONProof of 3.2 can be found in Appendix C. Theorem 3.2 provides the generalization bound of the optimal dual model trained on a finite selected demonstration set under a mild assumption that ∥W ∥ is bounded. Intuitively, as the number of demonstration (and therefore the number of demonstration tokens) increases, the generalization error decreases, which is consistent with existing experimental observations [Xie et al., 2021 , Garg et al., 2022 , Wang et al., 2024] . 
attention modification inspired by the representation learning lens: 
 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 ∥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 W , specifically byEQUATION)where α is a hyperparameter. Equivalently, the attention output Eq (1) will be modified asEQUATION)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 asL = - 1 ηD N i=1 [g 1 (W V x i )] T W ϕ(g 2 (W K x i )).Correspondingly, the attention layer can be adjusted as,EQUATIONwhere g 1 (•) and g 2 (•) 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:L = - 1 ηD N i=1 (W V xi ) T W ϕ(W K x i ), xi = x i - β |N (i)| j∈N (i) x j ,where N (i) is the set of the negative samples for x i and β is a hyperparameter. Correspondingly, the attention layer is modified asEQUATIONwhereXD = [ x1 , x2 , ..., xN ].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. 
experiments: 
 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 s = W t where every element of W ∈ R ds×dt is sampled from a normal distribution W ij ∼ N (0, 1) and t from uniform distribution t ∼ U (-1, 1) dt . We set d t = 11 and d s = 1. Then, at each step, we use generated{x i = [t i ; s i ]} N +1 i=1to form the input matrix X where the last token will be used as the query token and the label part will be masked, that is,x N +1 = [t i ; 0].Here we consider only one query token (T = 0) and we denote x ′ T +1 = x N +1 to maintain consistency of notation in Section 2.1. Finally, the attention layer is trained to predict ŝ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, ϕ(x) = e w T x-∥x∥ 2 /2 where w ∼ 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 X, we can construct the dual model f (x) = W ϕ(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 X test using the same method as training and obtain the ICL result of the query token h ′ T +1 . On the other hand, we use X test to train the dual model according to Theorem 3.1 and get the test prediction ŷ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 ŷtest is exactly equivalent to the ICL inference result h ′ 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.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 α 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 α > 0, the regularized models converges to a poorer result while when α < 0, the model converges faster and achieves final results comparable to the normal model without regularization (α = 0). At least for this setting, this is a little contrary to our initial intention of applying regularization to the contrastive loss where α should be positive. We explain it that the appropriate α 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 (•) and g 2 (•) 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 (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 β 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 β = 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. 
related work: 
 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 = w T x) and specific structured input format where each token takes the form [x, y] consisting of the input part x and the label part y. In addition, the label part of the final query to be predicted is masked, represented as [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. 
conclusion and impact statements: 
 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. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We have provided sufficient information on the computer resources in Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We confirm that our research meets the requirements.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We have provided relevant explanations at the end of the main body.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We think that our paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: Our paper does not use existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: Our paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Our paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: Our paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We repeat Theorem 3.1 as follows and provide proof and more discussion for it. Theorem A.1. The query token h ′ T +1 obtained through ICL inference process with one softmax attention layer, is equivalent to the test prediction ŷtest obtained by performing one step of gradient descent on the dual model f (x) = W ϕ(x). The form of the loss function L is:EQUATION)where η is the learning rate and D is a constant.Proof. The derivative of L with respect to W should be:∂L ∂W = - N i=1 1 ηD W V x i ⊗ ϕ(W K x i ) .Thus, after one step of gradient descent , the learned W will beEQUATIONwhere W 0 is the initialization of the reference model and η is the learning rate. So the test prediction will beEQUATIONOn the other hand, from the perspective of ICL process with one attention layer, with Eq (7) in our mind, we can rewrite Eq (1) ash ′ T +1 = W V Xsoftmax (W K X) T W Q x ′ T +1 √ d o = 1 D ′ W V [X D , X T ] [ϕ(W K X D ), ϕ(W K X T )] T ϕ(W Q x ′ T +1 ) = 1 D ′ [V D , V T ] [ϕ(K D ), ϕ(K T )] T ϕ(q),where we use[V D , V T ] = W V [X D , X T ], [K D , K T ] = W K [X D , X T ], q = W Q x ′ T +1 for sim- plify and D ′ = 1 T N ϕ(K D ) T ϕ(q) + 1 T T ϕ(K T ) T ϕ(q)is a constant to normalize the equivalent attention block. Further, we expand the above equation to connect the inference process of ICL using softmax attention with the gradient descent as followsh ′ T +1 = 1 D ′ V T ϕ(K T ) T ϕ(q) + 1 D ′ V D ϕ(K D ) T ϕ(q) = W ′ 0 ϕ(q) + 1 D ′ N i=1 V (i) D ⊗ ϕ(K (i) D ) ϕ(q)whereEQUATIOND are the i-th column vetors respectively. Then, in Eq (18), when setting the initialization W 0 = W ′ 0 and the constant D = D ′ , we will find thatEQUATIONwhich means ŷtest is strictly equivalent to h ′ T +1 . Thus, we have completed our proof.Given a reference model f (x) = W ϕ(x), by comparing Eq (19) and Eq (4), we can easily observe that the gradient descent on the loss function L applied to f (x) is the dual form of the inference process of ICL, whereV (i) D , ϕ(K (i) D) and ϕ(q) play the roles of backpropagation signals, training inputs and test inputs respectively. Recalling the form of Eq (4), we can interpret the W 0 as the initialization of the weight matrix which provide the information under the zero-shot case while the second part in Eq (19) shows that the demonstration examples in ICL acts as the training samples in gradient descent. The reference model f (x) = W ϕ(x), initialized with W 0 , will have test prediction ŷtest = h ′ T +1 after training. This is also why we refer to it as the dual model of the softmax attention layer. We also note that for different demonstrations, even though the model has the same query input, the different given demonstrations will result in different output results. This is equivalent to the dual model performing gradient descent in different directions from the same initialization.In Theorem 3.1, we provided the dual form of gradient descent for the ICL of one softmax attention layer. Here, we extend the conclusion to more complex scenarios, including one Transformer layer (attention layer plus one FFN layer) and multiple attention layers.As for one Transformer layer introduced in Section 2.1, we define the new dual model asf + (x) = W ϕ(x) + b.(20) We will show that after performing gradient descent on W , the test output ŷtest = f + (W Q x ′ T +1 ) will be equivalent to x′ T +1 . Our theorem is given as follows. Theorem B.1. The output x′ N +1 of ICL inference process with one Transformer layer, is strictly equivalent to the test prediction of its dual model f + (x) = W ϕ(x) + b, where f (x) is trained under the loss function L formed asEQUATION)where η is the learning rate, D is a constant, and W F will be determined once the specified pre-trained model, demonstrations and query tokens are given.Proof. Recalling the proof of Theorem 3.1, we can rewrite Eq (1) asEQUATIONwhereD = 1 T N ϕ(W K X D ) T ϕ(W Q x ′ T +1 ) + 1 T T ϕ(W K X T ) T ϕ(W Q x ′ T +1) is a constant to normalize the attention scores andW 0 = 1 D (W V X T )ϕ(W K X T ) T . Furthermore, h ′ N +1will be taken as input for the FFN sublayer and the Eq (2) can be rewritten asx′ T +1 = W 2 I M (W 1 h ′ T +1 + b 1 ) + b 2 = W 2 I M W 1 h ′ T +1 + W 2 I M b 1 + b 2, where I M ∈ R d×d is a diagonal matrix whose i-th diagonal element will be one if (W 1 h ′ T +1 +b 1 ) i ≥ 0 otherwise be zero. We need to note that this process is reasonable: for given demonstration and query tokens, once the parameters {W Q , W K , W V , W 1 , b 1 } of the Transformer layer are fixed after training, I M will be determined implicitly (otherwise, I M would be a function that varies with these settings). For simplify, we rewrite x′T +1 as x′ T +1 = W F h T +1 + b F , where W F = W 2 I M W 1 and b F = W 2 I M b 1 + b 2 . Furthermore, expanding h ′ T +1in the above Equation, we get:EQUATIONOn the other hand, we define a reference model:f + (x) = W ϕ(x) + b,where ϕ(•) is exactly the mapping function satisfying Eq (7) to approximate the softmax kernel.Given the loss formed in Eq (21), we can note that the right part in L is exactly the output of this reference model when taking W K x i as input, that is,L = - 1 ηD N i=1 (W F W V x i ) T (W ϕ(W K x i ) + b) = - 1 ηD N i=1 (W F W V x i ) T f + (W K x i ) .We can calculate the derivative of L with respect to W as∂L ∂W = - 1 ηD N i=1 W F W V x i ⊗ ϕ(W K x i ) .Suppose that the weight matrix W in the reference model f (x) is initialized as W init , then using one step of stochastic gradient descent (SGD) [Amari, 1993] with learning rate η, the weight matrix W will be updated asW = W init -η ∂L ∂W = W init + N i=1 1 D W F W V x i ⊗ ϕ(W K x i ) .Compared to Eq (23), we can setW init = W F W 0 , b = b F and take W Q x ′ T +1 as test input.Then, after one step update to W , the output of the reference model will bef + (W Q x ′ T +1 ) = W ϕ(W Q x ′ T +1 ) + b = W init + N i=1 1 D W F W V x i ⊗ ϕ(W K x i ) ϕ(W Q x ′ T +1 ) + b = W F W 0 + N i=1 1 D W F W V x i ⊗ ϕ(W K x i ) ϕ W Q x ′ T +1 + b F = x′ T +1 ,which implies that if we initialize the reference modelf + (x) = W ϕ(x) + b with W init = W F W 0 , b = b F ,then after one step of gradient descent for W , the test output of f + (W Q x N +1 ) will be identical to the ICL result of one Transformer layer. Thus, we call the reference model with settingW init = W F W 0 , b = b Fas the dual model corresponding to the ICL inference process. Finally, we complete our proof. Now, we discuss Theorem B.1 from the following perspectives:• Training set and test input: In fact, we can observe that the loss function L can be seen as the sum of inner products of N vector-pairs. In Eq (21), the right vector happens to be the predicted outputŷ(i) std = f + (z (i) std ) = W ϕ(z (i) std ) + b of the dual model for training input z (i) std = W K x i .Correspondingly, the vector on the left can be regarded as the true label y(i) std = W F W V x i .In other words, it can be seen that the dual model performs one step SGD given training set {z(i) std , y (i) std } N i=1 on W using the loss L: L = 1 ηD N i=1 y (i) std T ŷ(i) std .And then taking z test = W Q x i as test input, it finally output the prediction y test , which achieves the ICL result xN+1 . Compared to Theorem 3.1, after introducing the FFN layer, the main difference is that the labels of the training data become y(i) std = W F W V x i instead of y (i) std = W V x i . Additionally, compared to f (x), an extra bias b is introduced in the new dual model f + (x), which also have a different initialization W init = W F W 0 rather than W 0 . We also need to note that in the dual model f + (x), only W is trained, while b remains unchanged after initialization. • Potential Low-rankness of W F :Noting that W F = W 2 I M W 1 where W 1 ∈ R d h ×d , W 2 ∈ R d×d h , I M ∈ R d h ×d h (here we assume that d i = d o = d for simplify), the rank of W F will satisfy Rank(W F ) ≤ min {Rank(W 1 ), Rank(W 2 ), Rank(I M )} .We observe that I M is a diagonal matrix with elements being zero or one, and its rank is determined by the number of non-zero elements. Here, we can make a mild assumption that we can set Rank(W 1 ) = Rank(W 2 ) = min{d, d h }. This assumption is quite mild as even for any random square matrix as it will be non-singular with probability 1. In addition, we also assume d h > d which is consistent with settings in practice. Therefore, we get Rank(W 1 ) = Rank(W 2 ) = d, and the upper bound of Rank(W F ) will beRank(W F ) ≤ min {d, Rank(I M )} .Thus, we can find that if we want to avoid losing information, W F should strive to maintain Rank(I M ) > d which will be more easily achieved as d h becomes larger than d. Otherwise, Rank(I M ) is likely to gradually decrease with an increase in the number of Transformer layers. This explains the necessity of setting d h > d in practice. In some cases where Rank(I M ) < d, meaning that the number of non-zero elements in I M or positive elements in W 1 h N +1 + b 1 is less than d, the upper bound of Rank(W F ) will be Rank(I M ) and the lower bound of Rank(W F ) will be given asRank(W F ) ≥ Rank(W 2 I M ) + Rank(I M W 1 ) -Rank(I M ) = Rank(I M ),which implies the rank of W F will exactly equal to Rank(I M ). We should note that this condition, i.e., Rank(I M ) < d, is easily satisfied when d h = d or when d h is slightly larger than d (for example, d < d h < 2d in an expected sense). Thus, we conclude that W F has the potential low-rank property.• Representation Learning Lens: For a encoded demonstration representation x i , the key and value projections will generate a pair of feature W K x i and W V x i to create a certain distance between data representations in space. And then, on the one hand, a potential lowrank transformation W F is applied to the W V x i , attempting to compress some information which increases the difficulty of contrastive learning and forces the model to learn better features; on the other hand, ϕ(•) projects W K x i into a higher-dimensional space to capture deeper-level features. Finally, we need to train the weight matrix W , which maps ϕ(W K x i ) back to the original space, aiming to make the mapped vector as close as possible to W F W V x i . This interpretation is illustrated in Figure 4 .In this part , we extend Theorem 3.1 to multiple attention layers. Here we adopt the attention layer based on PrefixLM [Roberts et al., 2019] , where the query tokens can compute attention with all preceding tokens (including itself), while for demonstration ones, attention can be computed between themselves, excluding the query tokens. Existing work [Ding et al., 2023] has theoretically and experimentally explained that PrefixLM achieves better results than CasualLM. In this paper, we assume we have only one query token, that is, there is no query input before the considered query token. With the assumption that T = 0, to maintain notational simplicity, we use x N +1 to represent the query token here instead of x ′ T +1 and the input will be X = [X D , x N +1 ]. We assume that there are L attention layers and the output of the l-th layer X (l) can be expressed as:EQUATIONN +1 = W (l) V H (l-1) Softmax (W (l) K H (l-1) ) T W (l) Q h (l-1) N +1 √ d .where we set H (0) = X = [X D , x N +1 ] as the initial input. And the final output of the query token is ĥ(L)N +1 = h (L) N +1 .Here, we assume that after training, the parametersEQUATIONV ∈ R do×di are fixed and we setd o = d i = d.Next, we extend Theorem 3.1 to the case of multiple softmax attention layers. Formally, we present our result in the following theorem.Theorem B.2. Given L softmax attention layers whose parameters {W(l) Q , W (l) K , W (l) V } Ll=1 are fixed after training, the ICL output of these layers is equivalent to sequentially performing one step gradient descent on a sequence of dual modelsf (l) (x) = W (l) ϕ(x) L l=1, where the loss function for the l-th dual model is:EQUATIONwhere h(l-1) iis the output of the (l -1)-th attention layer for the i-th token, η is the learning rate and D (l) is a constant. The input for the l-th dual model is generated by the trained (l -1)-th dual model.Proof. Given H (l-1) as the input for the l-th attention layer, the inference process of h(l) N +1 is h (l) N +1 = W (l) V H (l-1) Softmax (W (l) K H (l-1) ) T W (l) Q h (l-1) N +1 √ d = W (l) 0 ϕ W (l) Q h (l-1) N +1 + N i=1 1 D (l) W (l) V h (l-1) i ⊗ ϕ(W (l) K h (l-1) i ) ϕ W (l) Q h (l-1) N +1 ,whereD (l) = 1 T N +1 ϕ(W (l) K H (l-1) ) T ϕ(W (l) Q h (l-1) N +1) is a constant to normalize the attention scores and W(l) 0 = 1 D (l) (W (l) V h (l-1) N +1 )ϕ(W (l) K h (l-1)N +1 ) T . According to Theorem 3.1, we can easily get the dual model f(l) init (h) = W (l) init ϕ(h) where the initialization is W (l) init = W (l)0 . Given the loss function L (l) formed as Equation 24and training set z(i) std , y (i) std N i=1where z(i) std = W (l) K h (l-1) i and y (i) std = W (l) V h (l)i , we perform one step SGD with learning rate η on weight matrix W (l) and will get trained dual model:f (l) (x) = W (l) ϕ(x) = (W (l) init + ∆W (l) )ϕ(x) = W (l) 0 + N i=1 1 D (l) W (l) V h (l-1) i ⊗ ϕ(W (l) K h (l-1) i ) ϕ(x)Taking test input as z(l) test = W (l) Q h (l-1) N +1 , the prediction f (l) (z (l)test ) will exactly equal to h(l) N +1 .Next, we will show how to obtain HD through the trained dual model f (l) (x). And afterEQUATIOND will constitute the training set as well as the test input for the next dual model f(l+1) init (x).Keeping the initialized dual model f (l) init (x) and the trained one f (l) (x) in mind, we can compute the demonstration token output h (l) i (i = 1, 2, ..., N ) of l-th attention layer ash (l) i = W (l) V H (l-1) D Softmax (W (l) K H (l-1) D ) T W (l) Q h (l-1) i √ d = N i=1 1 D (l) i W (l) V h (l-1) i ⊗ ϕ(W (l) K h (l-1) i ) ϕ W (l) Q h (l-1) i = D (l) D (l) i N i=1 1 D (l) W (l) V h (l-1) i ⊗ ϕ(W (l) K h (l-1) i ) ϕ W (l) Q h (l-1) i = D (l) D (l) i W (l) init + N i=1 1 D (l) W (l) V h (l-1) i ⊗ ϕ(W (l) K h (l-1) i ) -W (l) init ϕ W (l) Q h (l-1) i = D (l) D (l) i W (l) -W (l) init ϕ W (l) Q h (l-1) i = D (l) D (l) i f (l) (W (l) Q h (l-1) i ) -f (l) init (W (l) Q h (l-1) i ) (25) where D (l) i = 1 T N ϕ(W (l) K H (l-1) D ) T ϕ(W (l) Q h (l-1) i) is a constant to normalize the attention scores for h(l)i . Therefore, once we obtain the trained dual model f (l) (h), we can use the Eq (25) to get the demonstration token outputEQUATIONN ]. These demonstration token outputs, along with the output h init (h). This process continues layer by layer until we obtain the ultimate ICL output h (L) N +1 . In summary, the ICL inference process across L attention layers is equivalent to performing gradient descent on L dual models sequentially. Thus, we complete our proof. This theorem is a natural extension of Theorem 3.1: when considering the stacking of multiple attention layers, a sequence of dual models is correspondingly generated. Although these dual models have the same form f (x) = W ϕ(x), they have different initializations and datasets. As the ICL inference process progresses layer by layer between attention layers, we equivalently perform gradient descent on the dual models one by one. The input H (l) for each attention layer, including demonstration tokens and query tokens, can be obtained from the test output of the dual models. This can be illustrated in Figure 5 .C.1 Proof of Theorem 3.2In this part, we provide the proof regarding the generalization boundary in Theorem 3.2. We restate our theorem as follows: Theorem C.1. Define the function class as F := {f (x) = W ϕ(W K x) | ∥W ∥ ≤ w} and let the loss function defined as Eq (10). Consider the given demonstration set as S = {x i } N i=1 where S ⊆ S T and S T is all possible demonstration tokens for some task T . With the assumption that ∥W V x i ∥, ∥W ϕ(W K x i )∥ ≤ ρ, then for any δ > 0, the following statement holds with probability at least 1 -δ for any f ∈ FEQUATION)Proof. Our proof is similar to the Lemma 4.2 in Saunshi et al. [2019] , but here we focus on a different function class. Firstly, we consider the classical generalization bound based on the Rademacher complexity of the function class which can refer to Theorem 3.1 in Mohri et al. [2018] . For a real function class G whose functions map from a set Z to [0, 1] and for any δ > 0, if S is a training set composed by N iid samples {x i } N i=1 , then with probability at leastEQUATIONwhere R S (G) is the traditional Rademacher complexity. By setting S exactly the demonstration set and G = g f (x) = -(W V x) T W ϕ(W K x) ∥W ∥ ≤ w , we can apply this bound to our case.Then, we construct a function classF = f (x) = [f (x); W V x] = [W ϕ(W K x); W V x] ∥W ∥ ≤ wwhose functions map from S to R 2do . Next, we will first prove R S (G) ≤ 2ρR S ( F) and to do this, we need to use the following Lemma:Lemma C.2 (Corollary 4 in Maurer [2016] ). Let Z be any set, andEQUATIONwhere f|S = ft (z j )t∈[n],j∈[M ].We apply Lemma C.2 to our case by setting Z = R di , S to be exactly the demonstration set, F to be the function class we constructed and n = 2d. We also use h : R 2do → R where h(x) = -⟨x 1:do , x do+1:2do ⟩ and thus we haveg f (x) = h( f (x)) = h([f (x); W V x]) = -(W V x) T W ϕ(W K x).We can find that g f (x) = g f (x) and the left side of inequality ( 28) is exactly R S (G).Then we can see that h is √ 2ρ-Lipschitz with the assumption that ∥W V x i ∥, ∥W ϕ(W K x i )∥ ≤ ρ and we have R S (G) ≤ 2ρR S ( F). Now using Lemma C.2 and the classical generalization bound (27), we have that with probability at least 1EQUATIONLet f * ∈ arg min f ∈F L(f ). According to Hoeffding's inequality, with probability at least 1 -δ 2 , we have that L(f * ) ≤ L(f * ) + 3 log 2 δ 2N . Combining this with (29), the fact that L( f ) ≤ L(f * ) and applying a union bound, we can get thatEQUATIONNext, we give the upper bound for R S ( F).R S ( F) = E σ∼{±1} 2N do sup ∥Wj ∥≤w 2N do t=1 σ t ( f|S ) t (Definition of Rademacher complexity) = E σ∼{±1} N do   sup ∥W ∥≤w do j=1 W j N i=1 σ i,j ϕ(W K x i )   (W V x i is independent of W j ) ≤ E σ∼{±1} N do   sup ∥W ∥≤w do j=1 ∥W j ∥ N i=1 σ i,j ϕ(W K x i )   (By Cauchy-Schwartz inequality) ≤ wd o E σ∼{±1} N N i=1 σ i ϕ(W K x i ) (Using the fact that ∥W j ∥ ≤ w) ≤ wd o E σ∼{±1} N   N i=1 σ i ϕ(W K x i ) 2   (By Jensen's inequality) = wd o Tr(K S )Substituting the upper bound of R S ( F) into (30), we will get thatEQUATIONThus we finish our proof.One may also wonder whether the ratio of negative samples mentioned in Section 4 will affect the generalization bounds. In fact, after introducing negative samples and ignoring constant term in Eq (9), we consider the following representation loss:L(f ) = E x∼D T   - 1 K K j=1 (W ϕ(W K x)) T W V x -W V x - j   ,where we consider sampling K negative samples for each x i and x - j denotes the j-th negative sample for token x . Correspondingly, the empirical loss will be considered asL(f ) = -1 N N i=1 1 K K j=1 (W ϕ(W K x i )) T W V x i -W V x - ijwhere x - ij is the j-th negative sample for x i . Then, by retaining the other definitions in Section 3.3, corresponding to Theorem 3.2, we can obtain the generalization bound asL( f ) ≤ L(f ) + O   wρd o Tr(K S ) 5 N 2 + 1 rN 3 + log 1 δ N   ,where r = K N is excatly the the ratio of the number of negative samples. It can be observed that as the ratio of negative samples increases, the generalization error decreases. However, we also notice that 5N 2 > 1 rN 3 thus the former term dominates, which means the reduction in generalization error due to an increased proportion of negative samples is limited. Nevertheless, we do not rule out the possibility of a tighter generalization bound, which is a promising direction for future research.Proof Sketch. The proof process is similar to that of Theorem 3.2.The main difference lies in the fact that we should firstly define the function classG = -1 K K j=1 (W ϕ(W K x i )) T W V x i -W V x - j∥W ∥ ≤ w to use the classical bound. In addition, we defineF = f (x) = [f (x); W V x; W V x - 1 ; ...; W V x - K ]∥W ∥ ≤ w whose functions map from S to R (K+2)do . Similarly, when using Lemma C.2, we set Z = R di , F be the above function class and n = (K + 2)d o . We also use h : R (K+2)do → R defined ash(x) = -1 K K j=1 x T 1:do (x do+1:2do -x (j+1)do+1:(j+2)do). Then we notice thatEQUATION)With the assumption that ∥W V x∥, ∥W ϕ(W K x)∥ ≤ ρ, we can get that the Frobenius norm of theJocabian J of h has ∥J∥ 2 F ≤ 4ρ 2 + ρ 2 + K K 2 ρ 2 = (5 + 1 K )ρ 2 . Thus we get that h is 5 + 1 K ρ- Lipschitz.The rest of the proof process is similar to that of Theorem 3.2. Ultimately, we will obtain the aforementioned generalization error.In this section, we provide a more detailed discussion on improving the model structure from the perspective of representation learning especially contrastive learning, which is presented in Section 4 of the main body. And we also point out the corresponding modifications in the self-attention mechanism, which are adopted in our experiments.Although we have figured out the representation learning loss of the implicit gradient updates, it can be observed that this loss function has a flaw: due to the lack of normalization for y (i) std and ŷ(i) when calculating the cosine distance, the loss can theoretically be optimized to negative infinity. To address this issue, we introduce regularization to constrain the norm of W , that is,L = - 1 ηD N i=1 (W V x i ) T W ϕ(W K x i ) + α 2η ∥W ∥ 2 F ,where α is a hyperparameter to balance the two parts. As a result, we can see that the gradient update for W will be in an exponentially smoothed manner meaning that a portion of the initial part will be discarded at every step, that is,W (t) = W (t-1) -η ∂L ∂W = (1 -α)W (t-1) + N i=1 D -1 W V h i ⊗ ϕ(W K h i ).Equivalently, the inference process of ICL can be seen as the first step of the aforementioned update, and the attention mechanism will be correspondingly adjusted as,h ′ T +1 = (1 -α)W 0 ϕ(q) + D -1 N i=1 V (i) D ⊗ ϕ(K (i) D ) ϕ(q),which means more demonstration information will be attended to. This will directly result in Eq 13.This result can be easily extended to self-attention mechanism. As for a self-attention layer, if all other tokens adopt the same modification, the self-attention layer will becomeH = W V Xsoftmax (W K X) T W Q X √ d o -αW V X = W V X softmax (W K X) T W Q X √ d o -αI ,which leads to the model structure incorporating an operation similar to skip connections. Furthermore, to ensure numerical stability, we normalize the attention scores yielding:H = W V X • Norm softmax (W K X) T W Q X √ d o -αI ,where Norm(•) is performed column-wise to ensure that the attention scores sum to 1. The above modification reduce the attention score of each token to its own information during aggregation.It is worth noting that, although our initial intention is to impose regularization on the contrastive loss where α > 0 to prevent it from diverging to negative infinity, we find in experiments that this modification remains effective even when α is less than 0. We interpret this as possibly stemming from the fact that an appropriate α helps the attention block become full-rank, thereby better preserving information, which can be illustrated by Lemma D.1: Lemma D.1. Let the attention block A ∈ R n×n . There exists some δ > 0 such that, for any 0 < |α| < δ, the attention block A + αI n will become full-rank.Proof. Define f (α) = det(αI n + A), which is a polynomial of degree n in α. Then, f (α) has only finitely roots. Let α 1 , α 2 , . . . , α r be the non-zero roots of f (t). Now, consider δ = min{|α 1 |, |α 2 |, . . . , |α r |}. For 0 < |α| < δ, we can claim that f (α) = det(αI n + A) ̸ = 0. Thus, A + αI n becomes non-singular (full-rank) and we complete the proof.Lemma D.1 provides one possible case for appropriate α. In fact, the selection of α can be quite flexible; for instance, similarly, when δ = max{|α 1 |, |α 2 |, . . . , |α r |} and |α| > δ holds, A + αI n also remains full-rank. Our experimental results related to regularized models will further illustrate the effectiveness of an appropriate α in enhancing model performance.We also acknowledge that our modification is relatively straightforward and may not be optimal. However, we believe that it may be a good choice to make structural improvements to the model from the perspective of the loss function, or more generally, from an optimization standpoint. For example, to address the issue of non-normalized y(i)std and ŷ(i) , we can also modify the loss function from the perspective of ridge regression as:L = 1 2ηD N i=1 ∥W V x i -W ϕ(W K x i )∥ 2 F + α 2η ∥W ∥ 2 F .And the optimal W * will beW * = ϕ(W K X)ϕ(W K X) T + αDI -1 W V Xϕ(W K X).Correspondingly, the attention mechanism will be modified toEQUATION)where we neglect the normalization operation. This result is very similar to the mesa-layer proposed by Von Oswald et al. [2023b] , which optimizes linear attention layers under the auto-regressive setting.Here, we presented its form on softmax self-attention setting using kernel methods and explained it from the perspectives of contrastive loss and ridge regression. Although the matrix inversion calculation in Eq (33) can be computationally expensive, effective methods for computing Eq (33), including both forward computation and backward propagation, have been thoroughly researched in Von Oswald et al. [2023b] , which contributes to making the above modification practically applicable.In addition to discussing the loss function, the contrastive learning paradigm also offers our some insights. In the corresponding representation learning process of ICL, we can easily notice that "data augmentation" is performed using a simple linear mapping, which may be not sufficient for learning deeper-level features. To address this, we can employ more complicated nonlinear functions for more complex augmentations. Denoting these two augmentations as g 1 and g 2 , consequently, the process of contrastive learning will be modified as followsL = - 1 ηD N i=1 [g 1 (W V x i )] T W ϕ(g 2 (W K x i )).Correspondingly, the gradient update for W will becomeW (t) = W (t-1) -η ∂L ∂W = W (t-1) + N i=1 D -1 g 1 (W V x i ) ⊗ ϕ(g 2 (W K x i )).And from the perspective of ICL, correspondingly, the last token will be updated ash ′ T +1 = W 0 ϕ(q) + D -1 N i=1 g 1 (V (i) D ) ⊗ ϕ(g 2 (K (i) D )) ϕ(q).And by reformulating the above equation we will get Eq (14) in the main body.Correspondingly, the modification for self-attention layer can be adjusted as,H = g 1 (W V X)softmax g 2 (W K X) T W Q X √ d o ,where g 1 (•) and g 2 (•) will be column-wise here. It is worth noting that here we have only presented the framework of using nonlinear functions as data augmentations to modify the self-attention layer and in the simplest case, we can set g 1 (x) and g 2 (x) as MLPs (Multi-Layer Perceptrons). However, in practice, it is encouraged to use data augmentation functions that are tailored to specific data structures. For example, in the case of CMT [Guo et al., 2022] , the used Convolutional Neural Networks (CNNs) can be considered as a form of "strong data augmentations" suitable for image datas within our framework. We consider the exploration of various augmentation methods tailored to different types of data as an open question for future research.Although the gradient descent process corresponding to ICL exhibits some similarities with traditional contrastive learning approaches without negative samples, there are also significant differences: In traditional Siamese networks, the augmented representations as positive pairs are further learned through target and online network that share weights (or at least influence each other using EMA).The output of the target network is then passed through a predictor to compute the contrastive loss.In contrast, the representation learning pattern corresponding to ICL indeed performs more simply, which may potentially limit the ability of the dual model to learn representations fully without negative samples. To address this, similar to most contrastive learning approaches, we can introduce negative samples forcing the model to separate the distances between positive and negative samples at the same time, that is,L = - 1 ηD N i=1 (W V x i ) T W ϕ(W K x i ) + β ηD N i=1 1 |N (i)| j∈N (i) (W V x j ) T W ϕ(W K x i ) = - 1 ηD N i=1   W V   x i - β |N (i)| j∈N (i) x j     T W ϕ(W K x i ) = - 1 ηD N i=1 (W V xi ) T W ϕ(W K x i ), where xi = x i -β |N (i)| j∈N (i) x j , N (i)is the set of the negative samples for x i and β is a hyperparameter. As a result, the gradient descent on W will be modified asW (t) = W (t-1) -η ∂L ∂W = W (t-1) + N i=1 D -1 W V xi ⊗ ϕ(W K x i ).Correspondingly, the ICL process for ĥN+1 will beh ′ T +1 = W 0 ϕ(W Q x ′ T +1 ) + D -1 N i=1 W V xi ⊗ ϕ(W K x i ) ϕ(W Q x ′ T +1 ).And this will directly result in Eq (15) in the main body.As for a self-attention layer, similarly, we can get the corresponding modification asEQUATION)where X(i) = xi . In corresponding experiments, for each token, we simply choose other the k least relevant tokens as its negative samples, i.e., the k tokens with the lowest attention scores. Noting that here we simply use other token representations as negative samples for x i . However, there are more ways to construct negative samples that are worth exploring (for instance, using noise vectors or tokens with low semantic similarity as negative samples). For specific data structures and application scenarios, customizing the selection or construction of negative samples may be more effective.In this part, we will discuss our experimental setup in more details and provide more results on linear regression task.Inspired by Garg et al. [2022] and Von Oswald et al. [2023a] , we choose to pretrain a softmax attention layer before exploring the equivalence proposed by Theorem 3.1. In fact, pretraining is not mandatory since our theoretical analysis does not depend on any specific weight construction. In other words, the inference results of ICL and the test prediction of the dual model will still remain consistent for an attention layer with arbitrary weights or even random initialization. However, for the convenience of further investigating the impact of subsequent modifications to the model structure and to better align with real-world scenarios, we still opted for pretraining to let the model acquire some task-specific knowledge. Additionally, our experiments are conducted in a self-attention setting.When we focus only on the last token, this is equivalent to considering the case with only one query token (T = 0) in Section 2.1. The experiments are completed on a single 24GB NVIDIA GeForce RTX 3090 and the experiments can be completed within one day.For the linear regression task, we generate the task by s = W t where every element of W ∈ R ds×dt is sampled from a normal distribution W ij ∼ N (0, 1) and t is sampled from a Gaussian distribution x ∼ U (-1, 1) dt . To facilitate more accurate estimation of attention matrices using random features and considering the limited learning capacity of a single attention layer, we only set a small value for d t = 11 and d s = 1. Then, at each step, we use generated {x i = [t i ; s i ]} N +1 i=1 to form the input matrix X while the label part of the query token is masked to be zero, that is, x N +1 = [t i ; 0] where we consider only one query token and we denote x ′ T +1 = x N +1 to maintain consistency of notation in Section 2.1. The softmax attention layer is expected to predict ŝN+1 to approximate the ground truth value s N +1 . We use mean square error (MSE) as the loss function, that is, for each epoch,L = 1 N step Nstep j=1 ∥ŝ (j) N +1 -s (j) N +1 ∥ 2 ,where ŝ(j) N +1 and s (j)N +1 are the prediction and ground truth value at j-th step and N step is the number of steps. We set N step = 1024 for N + 1 = 16 which means the total number of tokens remains 16384. We choose stochastic gradient descent (SGD) [Amari, 1993] as the optimizer and we set the learning rate to 0.003 for normal and regularized models, while the remaining experiments to 0.005. We also attempt the multi-task scenario, where the input token at each step is generated from a different task. However, we find it challenging for a single attention layer to effectively learn in this setting, resulting in disordered predictions. Therefore, our experiments are currently limited to single-task settings, and the multi-task scenario is worth further investigation in the future.It is worth noting that we approximate the attention matrix calculation using random features as kernel mapping function instead of using the traditional softmax function in the self-attention layer [Choromanski et al., 2020] . The mapping function ϕ : R do → R dr has the form of ϕ(x) = e w T x-∥x∥ 2 /2 where w ∼ N (0, I). Orthogonal random features [Yu et al., 2016 , Choromanski et al., 2020] or simplex random features [Reid et al., 2023] can be chosen to achieve better performance theoretically. We investigate the impact of changing the dimension of random features d r on the approximation of attention matrices and output, using Mean Squared Error (MSE) and Mean Absolute Error (MAE) as evaluation metrics, where we conduct 50 repeated experiments and calculated the average values for each value of d r , as shown in Figure 8 . It can be observed that as the dimension of random features increases, the approximation performance gradually improves, with both errors reaching a low level in the end. We visualize the exact attention matrix and compare it with the estimated attention matrices obtained using different values of d r , as shown in Figure 6 . Again, it can be seen that as d r increases, the approximation of the true attention matrix improves gradually and similar results can be observed for the analysis of output matrices in Figure 7 .To obtain a more accurate estimation of the attention matrix, we set the output dimension of the mapping function to be 100 times the input dimension, that is, d r = 100(d s + d t ) = 1200. Furthermore, we visualize the exact attention matrix and the output with the approximation results, which are shown in the Figure 9 . As we can see, although some larger values are not estimated accurately due to the limited dimension of the random features we select, the majority of the information is still estimated comprehensively well. These findings indicate that our choice of using positive random features as mapping functions to estimate the true softmax attention and conduct experiments is relatively feasible.After the weights W Q , W K , W V of the attention layer have been determined, we generate test N + 1 tokens in the same way where the s part of the (N + 1)-th token is also set to be zero and finally input the test tokens into the attention layer to obtain the corresponding predicted ĥN+1 = [ t(1)EQUATIONN +1 ]. Here, we also use h ′ T +1 = ĥN+1 to maintain the notation consistency in Section 2.1. ] of the dual model will strictly equal h ′ T +1 . We conduct experiments under the same setup using different random seeds to explore the effects of various model modifications. The data for all three experiments are generated under identical conditions. One set of experimental results is presented in the main text, while the results of the other two sets are shown in the Figure 10 . Similar to the discussion in the main text, we can achieve better performance than the normal model with appropriate parameter settings.Figure 10 : The performance for regularized models (Center Left), augmented models (Center Right) and negative models (Right) with different settings for different random seeds.In addition to conducting experiments on linear regression tasks, we also extended our experiments to involve trigonometric and exponential tasks. For trigonometric task, we generate the task by s = cos(W t) where cos(•) is element-wise, W ∈ R ds×dt is sampled from the normal distribution W ij ∼ N (0, 1) while t is sampled from the uniform distribution x ∼ U (0, π) dt . In experiments, we found that for one softmax attention layer, learning higher-dimensional tasks is challenging. Therefore, we only set d t = 7 and d s = 1. At each step, we useN + 1 = 128 tokens {x i = [t i ; s i ]} N +1i=1 and the total number of tokens remains unchanged at 16384. Compared to the setting N + 1 = 16 of linear tasks, we observed that for more complex tasks, the attention layer needs to use more tokens to provide information at each training step. Similarly, we mask the label part of the last token, that is, s N +1 = 0 and use mean square error (MSE) loss to train the attention layer. We choose SGD as the optimizer and the learning rate is set as 0.005. The rest of the settings remain consistent with those used in the linear task. The result for trigonometric regression task is shown in Figure 11 .Firstly, as shown in the left part of Figure 11 , the inference results is of ICL is strictly equivalent to the prediction of the dual model, that is, ĥN+1 = ŷtest as well as the label part ŝ(1) N +1 = ŝ(2) N +1 , aligning with our analysis in Theorem 3.1.The performance of modified model during training process can be seen in the remaining parts of Figure 11 . For regularized models, as seen in the center left part of figure 11 , the models when α < 0 converge slightly faster and reach better final results compared to the normal model (α = 0). For augmented models, we use as the same augmentation functions g 1 and g 2 as the ones in the linear regression task, that is, g 1 (x) = g 2 (x) = σ(W x) where σ(•) is GELU activation function. However, for g + 2 , we use ELU as the activation function. We can find from the center right part of Figure 11 that, compared to the normal model, using g 1 alone and using g 1 and g 2 simultaneously as data augmentations significantly degrade the model's performance, including convergence speed and final results. However, using g 2 alone yields comparable result with the normal model. Particularly, when using g + 2 , the model accelerates its convergence speed. However, for negative models, the performance with the selected number of negative samples k and the parameter β is worse than the normal model, which suggests that our simple approach of selecting those tokens low attention scores as negative samples is not a reasonable method. Just as we discussed in Section 4, for different tasks, a more refined strategy for selecting negative samples should be considered.For exponential task, we generate the task by s = exp(W t) where exp(•) is also element-wise, W ∈ R ds×dt is sampled from the normal distribution W ij ∼ N (0, 1) while t is sampled from the uniform distribution x ∼ U (-1, 1) dt . We only set d t = 6 and d s = 1 considering the limited learning capacity of one softmax attention layer. At each training step, we use N + 1 = 512 tokens{x i = [t i ; s i ]} N +1i=1 and the total number of tokens remains unchanged at 16384. Compared to the setting N + 1 = 16 of linear tasks and N + 1 = 128 of trigonometric tasks, we also find that for exponential tasks, the attention layer needs more tokens to provide in-context information at each training step. The rest of the settings remain consistent with those used in the trigonometric task. The result for exponential regression task is shown in Figure 12 .Similarly, as shown in the left part of Figure 12 , the result ĥN+1 of ICL inference is equivalent to the test prediction ŷtest of the dual model after training, just as stated in Theorem 3.1. For regularized models, it can be observed that when α = 16, the model converges faster and achieves better result. For augmented models, using g 1 or g 2 alone as data augmentations results in better performance. However, when both g 1 and g 2 are used simultaneously, the training process becomes unstable, so we did not show it in the center right part of Figure 12 . For negative model, similar to the case in the trigonometric task, the different combinations of negative samples' number k and parameter β do not show a significant improvement over the normal model, highlighting the importance of the strategy for selecting negative samples. We leave the exploration of a more refined negative sample selection strategy when facing various tasks for future consideration.In addition, we also conduct experiments with their combinations on linear tasks, trigonometric tasks , and exponential tasks. The results are shown in Figure 13 . For linear tasks, a combination of regularized and augmented modifications is sufficient. However, for the other two tasks, the results are actually worse than using regularized or augmented modification individually (compared to Figures 11 and 12 ). We think this may be due to the ineffective selection of negative samples, which is amplified when combined. Therefore, when the design of augmentation or negative sample improvement methods is not effective, we recommend using a single modification method. property of matrix W F , we explore its upper bound of rank. Noting thatW F = W 2 I M W 1 where W 1 ∈ R d h ×d , W 2 ∈ R d×d h , I M ∈ R d h ×d h , the upper bound of Rank(W F ) is Rank(W F ) ≤ min {d, d h , Rank(I M )} ,where Rank(I M ) is equivalent to the number of non-zero elements in I M . We fix d = 12 while varying the values of d h . We generate 1024 sets of X test for different tasks and repeat the experiments 5 times. Finally, we calculate the average upper bound of the rank of W F . The results are shown in the right part of Figure 14 , indicating that when d h ≥ 2.75d = 33, the upper bound remains stable and equals d = 12. Otherwise, when d h is set to a smaller value, W F exhibits clear low-rank property.We supplement our experiments on more more realistic NLP tasks. We choose the BERT-baseuncased model (can be downloaded from Huggingface library [Wolf, 2019] , hereafter referred to as BERT[Kenton and Toutanova, 2019] ) to validate the effectiveness of modifications to the attention mechanism and select four relatively smaller GLUE datasets (CoLA, MRPC, STS-B, RTE) [Wang, 2018] . We load the checkpoint of the pre-trained BERT model, where 'classifier.bias' and 'classifier.weight' are newly initialized, and then we fine-tune the model to explore the performance of three attention modifications as well as their combinations. In terms of more detailed experiment settings, we set the batch size to 32, the learning rate to 2e-5, and the number of epochs to 5 for all datasets.All experiments are conducted on a single 24GB NVIDIA GeForce RTX 3090. All experimental results are presented in Table 1 . Below, we discuss the various modifications and their performance.For the regularized modification, we consider different values of α, specifically selected from {-0.5, -0.1, 0.1, 0.5}. As can be observed in Table 1 , except for RTE, the best regularized models outperform the original model on the other three datasets. However, we also note that when the absolute value of α is too large, the model's performance declines significantly, so we recommend using smaller absolute values for α.For the augmented modification, we also consider applying more complex "augmentation" functions to the linear key/value mappings. However, unlike the previous methods used in simulation tasks, we do not simply select g 1 and g 2 as MLPs, i.e., g 1 (WV x) = W 2 σ(W 1 W V x).This design is avoided because it could undermine the effort made during pre-training to learn the weights W V and W K , leading to difficulties in training and challenges in comparison. Instead, we adopt a parallel approach, i.e., g 1 (WV x) = W V x + cW 2 σ(W 1 x),where c is a hyperparameter to control the influence of the new branch, σ is the GELU activation function and the hidden layer dimension is set to twice the original size of W V x. g 2 (W K x) = W K x + cW 2 σ(W 1 x) follows the same format.Experimental results show that the best augmented models achieve better performance than the original model across all four datasets. Notably, augmentation on the value mapping (i.e., using g 1 alone) proves to be more effective than other methods, both in terms of performance and the amount of additional parameters introduced. Using both g 1 and g 2 introduces more parameters, which is particularly undesirable for larger models. Thus, under the augmentation methods and experimental settings we selected, using g 1 alone is recommended.In addition, we do not rule out the possibility of more powerful and efficient augmentation methods. Our choice of g 1 and g 2 as parallel MLPs is primarily motivated by the desire to make better use of the pre-trained weights W K and W V . We have also noticed that this specific augmentation function design is structurally similar to the Parallel Adapter [He et al., 2021] . However, we would like to emphasize that our parallel design is just a specific case within this broader augmented modification framework and this is a new perspective for understanding the Parallel Adapter. As for practical implementation, the Parallel Adapter method focuses more on efficient training, so it uses fewer parameters, and the original W V and W K are freezed-only the newly introduced parameters are trained. In contrast, our approach aims to validate the benefits of introducing stronger nonlinear augmentation functions into the linear value/key mappings. Therefore, we set a higher hidden layer dimension (twice that of W V x or W K x) and also train W V and W K simultaneously. This design is relatively general and does not take into account the specific characteristics of individual tasks. We still encourage the development of more task-specific augmentation strategies tailored to different tasks.For the negative modification, we continue to select tokens with lower attention scores as negative samples. The parameter r represents the proportion of tokens used as negative samples, while β indicates the overall reduction in attention scores. We choose r from {0.1, 0.2, 0.3} and β from {0.1, 0.2}. Under these combinations, 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. Specifically, we choose α = 0.1, g 1 /c = 1 and r = 0.2/β = 0.1 respectively as the basis for combining the three types of modifications, considering their overall performance across all datasets. 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.Overall, the experimental results show that our modifications inspired by the representation learning process are helpful in enhancing performance. This further validates the potential of our approach of thinking about and improving the attention mechanism from a representation learning perspective.In addition, we would like to reiterate that more 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.In this section, we provide additional details about the related work in Section 6, especially those that involve formalization. Dai et al. [2022] interpret ICL as implicit fine-tuning: More specifically, let X = [X D , X T ] where X D = [x 1 , x 2 , . . . , x N ] denotes the demonstration tokens andX T = [x ′ 1 , x ′ 2 , . . . , x ′T ] be query tokens. On the one hand, for ICL, they consider the output of q = W Q x ′ T +1 under the linear attention setting asFICL (q) = W V [X D , X T ](W K [X D ; X T ]) T q = W V X T (W K X T ) T q + W V X D (W K X D ) T q = W ZSL q + LinearAtten(W V X D , W K X D , q) = W ZSL q + i ((W V x i ) ⊙ (W K x i )) T q = W ZSL q + ∆W ICL q,where W 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,FFT (q) = (W V + ∆W V )XX T (W K + ∆W K ) T q = (W ZSL + ∆W FT )qwhere ∆W K and ∆W V denote the parameter updates and they are acquired by back-propagation from task-specific training objectives [Dai et al., 2022] , which is a supervised learning process of the original model. Considering the similarity in form between FICL and FF T , 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):L = - 1 ηD N i=1 (W V x i ) T W ϕ(W K x 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 formed as Eq (9) 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, we can 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 Section 3.3), this transformation can provide reverse insights into analyzing the model mechanisms. • It can clearly observed that the dual model involves a self-supervised representation learning process from the dual perspective. 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 Section 4. • Intuitively, this explanation might be also 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.In addition, although we do not target specific tasks like linear regression as previous works mentioned in Section 6, we would like to point out that under those specific weight and input settings, an intuitive explanation can also be provided from a representation learning perspective. Here, we take the linear regression task as well as the weight constructions considered by Von Oswald et al. [2023a] as an example. Specifically, it assumes that the structured input is H= [h i ] N i=1 ∈ R (d+1)×(N )where h i = [x i , y i ] is sampled from some linear task y = w T x and the query token will be hN +1 = [x N +1 , -w T 0 x N +1 ].And the considered linear self-attention layer will take the constructed weights and query output as:EQUATIONwhere w 0 is the underlying initial matrix. Then the label part of hN+1 will has the form as ỹN+1 = -w T 0x N +1 + ∆w T x N +1 = -(w T 0 -η N N i=1 (w T 0 x i -y i )x T i )x N +1 = -ŷN+1, which is equivalent to the output -ŷ N +1 (multiplied by -1) of the linear layer y = w T x where w is initialized as w 0 after performing one step of gradient descent under mean squared loss L = 1 2N N i=1 ∥w T 0 x i -y i ∥ 2 . In practice, the underlying initial weight matrix w 0 is set to be approximately 0 thus the test input can be formed as h N +1 = [x i , 0] [Von Oswald et al., 2023a] . In addition, when reading out the label ŷN+1 , the test prediction ỹN+1 will be multiplied again by -1, which can be done by a final projection matrix (or equivalently, P = -η N I). In this case, we first note that the dual model of the linear attention layer can be written as f (z) = W z where W ∈ R (d+1)×(d+1) and similar to Eq (9), it will be trained under the loss below:EQUATIONBy substituting the corresponding weights in Eq (35) where we replace P = -η N I for the readout, the loss can be reformulated as:EQUATIONRecalling that h i = [x i , y i ] is sampled from some linear task y = w T x, we assume that ∥W ∥ F ≤ ∥w∥ 2 , it can then be easily seen that the optimal solution for Eq (37) will beEQUATIONFurthermore, similar to Section 3.2, we take W Q h N +1 as the input where W Q is constructed as Eq (35) and h N +1 = [x N +1 , 0], the optimal dual model will output the result f(W Q h N +1 ) = W * W Q h N +1 = [0, w T x N +1 ] = [0, y N +1 ]where the label part will be just the answer for the test query. Additionally, it would also be interesting to explore how these weights converge to the constructed form in Eq (35) or other forms under this special setting as previous works illustrated from the perspective of the dual model. Investigating this issue goes beyond the scope of this paper, and we will leave it for future exploration.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We have provided our code and instructions in the supplemental material. Our experiments are primarily simulation-based, and we have provided the code to generate the data in the code part. Guidelines:• The answer NA means that paper does not include experiments requiring code. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.It should be clarified that the term "dual" here is different from the one in mathematical optimization theory. Instead, it follows the terminology used in previous works[Irie et al., 2022, Dai et al., 2022], where the forward process of the attention layer and backward process on some model are referred to as a form of "dual". 


Decision: The paper seeks to understand in-context learning from a gradient descent perspective on a dual model. In addition to theoretical analysis, the authors also conduct empirical studies to verify the theoretical results. 

While some reviewers raised concerns regarding to the assumptions made for the analysis and the generalization ability, the authors' response  generally resolved most of the concerns.  Considering the contributions to understanding the mechanism of LLMs, I prefer an acceptance.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: * **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: questions: Please see the weakness section.
limitations: limitations: The paper has discussed the limitations sufficiently.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: (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: 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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.
limitations: limitations: The authors have clearly stated the limitations of this work, the selected transformer components of this work, and selected task settings.

Comments: title: Replying to Official Comment by Reviewer HnPE
Looking forward to your reply
Looking forward to your reply
Looking forward to your reply
Replying to Official Comment by Reviewer cRiK
Replying to New Official Comment by Reviewer cRiK
Replying to Official Comment by Reviewer cRiK
comment: 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!
Thank you for your detailed responses, my concerns have been addressed and I will maintain my score.
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!
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!
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!
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!
Thank you for all the clarifications and aiming for a clarified version.

The confidence score is now updated from 3 to 4.
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
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.
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
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).