Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens
Authors: Ruifeng Ren
Yong Liu
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformers
In-context Learning
Representation Learning
Primary area: interpretability_and_explainability
TL;DR: In this paper, we attempt to explore the ICL process in Transformers through a lens of representation learning.

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.

Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens
Ruifeng Ren
Yong Liu" Gaoling School of Artificial Intelligence Renmin University China
Gaoling School of Artificial Intelligence Remmin University of China Beijing, China
Beijing, China
renruifeng3200ruc.ede en
liuyonggsai&rus eragesa@@ruu..o edu. cn
Abstract
Pre-trained large language models based Transformers have demonstrated remarkable in-context learning (ICL) abilities. With just a few demonstration examples, the models However, it is still an open question to understand the mechanism of ICL. In this paper, we attempt 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 the dual model's test outputs. We delve into the training process of this dual model from representation learning standpoint and further derive a generalization error bound related the quantity 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.
implement
tasks without any parameter updates.
Introduction
Recently, large language models (LLMs) based on the Transformer architectures Vaswani et al. 2017] has shown surprising in- in-context learning (ICL) capabilities [Brown al. 2020 Wei 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 used in general downstream tasks conveniently. Despite the good performance of the ICI. capabilities, the mechanism of ICL still remains open question.
In order better understand the ICL capabilities, many works began give explanations from different aspects. Xie et al. [2021] propose Bayesian inference framework to explain how ICL occurs between pretraining and test time, where the LLMs infers shared latent concept among the demonstration examples. Garg et al. [2022] demonstrate through experiments that pre-trained Transformer-basee based models can learn new functions from in-context examples, including (sparse) linear functions, two-layer neural networks, and decision trees. Zhhang et al. [2023b) adopt 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 ICI as an algorithm
learning problem where transformer model implicitly builds a hypothesis function inference-
time and derive generalization bounds for ICL. Han et al. (003) suggest that LLMS can emulate
Corresponding Author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 [Aisemman 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 al. [202]]. However, there is still 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 more fine-gres grained exploration. Akyüüek 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 more in-depth exploration of the capabilities of ICL under 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 linear regression tasks, adopting specific constructions for the input tokens (concatenated from features and labels) and model weights. This limits the explanation of the Transforeer'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 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 generalization eπror bound of this process, which is related to the number of demonstration tokens. Our conclusions can be easily extended to more complex scenarios, including 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.
2 Preliminaries
2.1
In-context Learning with Transformers
The model we consider is composed of many stacked Transformer decoder layers, each of which composed of an attention layer and 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 [Xp, X,] E (N+T() where Xp [x1, ......, æN. denotes N demonstration tokens, and X = [x', Izzy .... x.,1 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 as
h
Wy X softmax (WWXX)* WD
(1)
1 = WyReLa Wihh + b1) + bb,
(2)
where WK, Wo, Wr E xdd, are parameters for key, query, value projections and W, E xd, bb E Rdh, b2 Rdd. are FFN parameters. Our concern is how the query token z learns in-context information from demonstrations. Unlike previous work (Von Oswald et al. 2023a Zhang et al. 2023a Bai et al. 2023), J, 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
dd xd. W₂ E Rd Rd,,
specific algorithms. Instead, we adopt the same setting as
(Dai et al. 2022 to study more general
cases.
2

---

2.2 Self-Supervised Representation Learning Using Contrastive Loss Functions
Representation learning aims learn embeddings of data to preserve useful information for down- stream tasks. One class of methods most relevant to our work is probably contrastive learning methods without negative samples [Chen and He [02T]] Grill al. 2020, Caron al. [020 (Tan et al. 202T]. Contrastive learning is significant approach of self-supervised learning (SSL) which aims 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 al., 2020b, (ord et al., 2018 Oh Song et al. 2016]. To alleviate the burden of constructing 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 21, 22 from the same data inputs, which will be processed by online network and target network respectively to obtain the corresponding representations, that is, Φ = fonline ( (21), 2e = ftarget (x2). The two encoder networks share weights directly using Exponential Moving Average (EMA). Then, &1 will be input into a predictor head to obtain the predictive representation 21 = g Finally, minimize the distance between the predictive representation and target representation, that is, C (z), StopGrad (2)) where StopGrad(() means े2 is treated as a constant during backpropagation. measure of distance, although they are equivalent when the vector is normalized. Another class similar to our work is kernel contrastive learning [Esser al. 2024). Given anchor æ and its positive and negative samples x+ ,,, aims to optimize the loss function L = f(x)"" (f(z the feature mapping for some kernel. We will consider the gradient descent process corresponding
For [(-), we often choose the cosine similarity or the (2-norm
f( )), where f(x) = W ((x) and 6(xx)
to the inference process of ICL from the perspective of representation learning and compare it with
the two aforementioned representation learning patterns.
2.3 Gradient Descent on Linear Layer is the Dual Form of Linear Attention
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. [022 Dai al. [022], that is, the gradient descent on linear layer can be seen as the dual form of linear attention. A simple linear layer can be defined as (x) Wx, where W e Rdoxd. is the projection matrix. Given training inputs [xi]]=] Rd with their labels [y
Rd,, a linear layer output the predictions [ where Ú: Wa a; and then compute certain loss L(U, y:) for training. Backpropagation signals [e Rd will be produced to update W in gradient descent process where e, = -7 (V C) if we set as the learning rate. During test time, the trained weight matrix W can be represented by its initialization Wo and the updated part AW, that is,
N W = W. AW = W. + Σ e, eeeaa Z
where @ denotes the outer product according the chain rule of differentiation. On the other hand, this process can be viewed from the perspective of linear attention. Let | E R denote the N key and value vectors constituting matrices K, V E Rdx.N respectively. For given query input q E Rd, linear attention is typically defined as the weighted sum of these value vectors
LA(V. K. = VK q
Σ v,k q
Σ
us ② k2 g.
Then, we can rewrite the output of a linear layer during test time
N
Sr. [r.eett)))= Wateet
WoWtess +
Σ
& 3;
Etest
WoWtest
LA(E, X, Btest),
Ft 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 some model are referred form of "dual".
3

---

fostunee
Dual Model /oo (in
Wx
Softmax attention layer with parameters (We, W., W.W
Thinned Dual baoke Heth Wote!
C(r sill)
W;
47
Test Input
Resmplo
antimal
Applee
fruit
Quaodiom
What
Figure 1:
The ICL output h 4+ of one softmax attention layer is equivalent to the test prediction
Dtest of its trained dual model f(x) W ((x). The training data and test input can be obtained by linear transformations of demonstration and query tokens, respectively.
where E R..xN. and X Rdxx are stacked by backpropagation signals [e,] and training inputs [x:]] respectively. We can find from Eq (4) that the trained weight W records all training datapoints and the test prediction of the linear layer indicates which training datapoints chosen to activate using LA(-) where [e,]/= can be considered as values while [a>]]- as keys and Deest as the query. This interpretation uses gradient descent bridge to connect predictions of linear layers with linear attention, which can be seen as simplified softmax attention used in Transformers.
Inspired by this relationship. Dai et al. [2022] understand ICI. 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 the specific details of the gradient descent process, including the form loss function and dataset, have not been explicitly addressed. In addition, Von Oswald al. 2023a Ding et al. [2023) also connect ICL with gradient descent for linear regression tasks using weight construction methods, where parameters WK. Wq and Wy of the self-atteni eeffaattentoon layer need roughly adhere to 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 ICI. capabilities from the perspective of gradient descent. Thus, we attempt to address these issues in the following sections.
3 Connecting ICL with Gradient Descent
In this section, we will address two questions discussed above: Without assuming specific constructions for model weights and input tokens, how relate ICL 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 ICL? In addressing these two questions, we will explore the gradient descent process corresponding to ICL from the perspective of representation learning.
3.1 Connecting Softmax Attention with Kernels
Before we begin establishing the connection between ICL and gradient descent, we need to firstly rethink softmax attention with kernel methods. Dai al. 2022 connect ICL with gradient descent under the linear attention setting. In fact, is completely feasible to interpret ICL under softmax attention with the help of kernel methods. We define the attention block as
A softmax (w X) WQ
which can be viewed as the product of an unnormalized part A,, and normalizing multiplier D, that is,
A = A. A, exp ( (W=x)"
D diag(1 Am),
(6)
where exp(-) is element-wise. Similar Choromanski et al. 2020), we define softmax kernel K km Rdd x Rdd + R4 as K (xm y) = "= the guassian kernel when the variance o² = According to Mercer's 's theorem Mercer 1909
|=]]" 1v]
K Kyuase( (x, y) where Kywass = e -|I= per-y/l 3/1
there exists some mapping function Id - d satisfying that K (x, y) d(x)² ø Thus.

---

noting that when omitting the vdg renormaaaliinon and equivalently normalize key and value vectors Eq (6), every entry in the unnormalized part A,, can be seen as the output of softmax kernel Km defined for the mapping &, which can be formulated as:
(i, exp (W x.) (VQ,) = Km W kE, Wo = o (Wx=,,) (W (7)
Woz;)
There have been many forms of mapping function o(.) used in linear Transformers research to approximate this non-negative kernel [Choromanski et al. 2020 Katharopoulos et al. 2020 Peng
al., 2021 Lu al. 2021). For example, we can choose ((-) as positive random features which
has the form (()) = eth x ||=||/ /2 to achieve unbiased approximation [Choromanski et al. 2020 Alternatively,
also choose ((x=) clu(xr) proposed by Katharopoulos et al. 2020
3.2 The Gradient Descent Process of ICL
Now, we begin to establish the connection between the ICL inference process of a softmax attention layer and gradient descent. We focus on softmax attention layer in a trained Transformer model, where the parameters @Wo WK. Wv} have been determined and the input X = [X D. Xrl has the form introduced in Section 2.I] Then, after the inference by one attention layer, the query token at position T + 1 will have the form h formulated by Eq I1)
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
f(x) Wo(#),
(8)
where W E [Rd, xd 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 take Zzest = W WQrT++1 as the test input, we can obtain its test prediction as
Jtest [(ztest) f (Woar++) We (WQar++1-
We will show that h441 in Eq ( is strictly equivalent 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' 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 the dual model f(x) W Wo(1)) The form of the loss function L is:
N
C= DD Σ (Wvx,)) Wo(W)o,),
(9)
where 1) 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, WK. Wo, Wv has been determined, the demonstration tokens [n:1M will be used to construct a
training set for the dual model. Specifically, the training data has the form (2(() y Maadi==1 where =td æ; inputs and aii = Wy.,, z2; as their labels. During training stage, for each input = the dual model outputs its prediction ( = std = We
(4)
(6)
std w (Wx*.;). Then, the
loss function Eq 9 can be rewritten as L = DD Σ Tg 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 2test Wo as the test input to get its prediction which will consistent
with the ICL result h/11 that is, ÿtest = f(tear)) We (Wo h/ This process can be illustrated in Figure Demonstration tokens provide information about the training data points and the weight matrix W is optimized to learn sufficient knowledge about demonstrations. This
5

---

Oigh MP
807
Af W.A.
Delty Tromets
Trap srveet
Figure 2: Left Part: The representation learning process for the ICL inference by one attention layer.
Remaining Part: Comparison of the ICL Representation Learning Process (Center Left), Contrastive Learning without Negative Samples (Center Right), and Contrastive Kernel Learning (Right).
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 æ; the key and value mapping will generate pair of features Wxæ, and Wyæ; that exhibit certain distance from each other, akin positive samples in contrastive learning. And then, ((x) projects W; into a higher-dimensional space to capture deeper features. Finally, the weight matrix W, which maps of (x) back the original space, is trained to make the mapped vector as close as possible to Wv This process is illustrated Figure 2 Below, we attempt to understand this process from the perspective of existing representation learming 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 two types of data augmentation, then from the perspective of contrastive learning without negative samples, this process can be similarly formalized as
(000) (1) ysid) ( (00")StopCrral(yba))
min W
where StopGrad(-) is naturally applicable because there are no learning parameters involved in the generation process of the representation you. However, it's important to note that the representation learning process of ICL is much simpler: Firstly, the online and target networks the augmentations W Wxæ;, Wvæ; are directly used as online and target representations respectively, Secondly, the predictor head is useful and not discarded, which then used during test stage.
absent while
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 function = ffx f(x*) where f(x) = W Wo(x). There are significant differences in the representation learning process of ICL: Firstly, 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 its dual model (x) = W <oxx) + b introduces an additional bias and only W is trained
(a) while b remains fixed. In addition. the labels of training set will be yotd WF Wxæ; where WF 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. [0022] considering the connection under linear attention setting, Theorem 3.]
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 more general setting.

---

3.3 Generalization Bound of the dual gradient descent process for ICL
In this part, we are interested in the generalization bound of the ICL gradient process. When ICL inference is performed for some task 7, we cannot provide all demonstrations related to task T limited by the length of imput tokens. We denote ST S R all possible tokens for the task 7 and assume that these tokens will be selected according to the distribution Dy. During particular instance of ICL inference, let S = {a: =1 S ST represent the example tokens we selected. We define the function class as F := (f()) Wø( w WKx) x [[[]| S w) where || denotes the Frobenius norm. Generally, ignoring constant term in Eq ( we consider the representation learning loss as
[()) = BDNDY [- (Wvz)* f(m) EREDD [- (Wvx)) W Wo((k=))) (
(10)
where E F and Dy is the distribution for some ICL. task T. Correspondingly, the empirical loss mig L(f). In
will be formulated as L(f) = N Σ (Wvæ,)* f(x.) and we have
addition, we denote the kernel matrix of demonstration tokens S as Ks E RRNNN xN where (Ks), = (o6 (o(Waa: ((W WRary)) )) that is, the inner product of the feature maps after Wx 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(z) W WKW ||| S w) and let the loss function defined as Eq (10 Consider the given demonstration set as S = { where S S ST and ST is all possible demonstration tokens for some task T. With the assumption that ax [| @W@lW (Wo)::) S p, for any 5 >> O, the following statement holds with probability at least 8 for any F
wpdo v Tr(Ks)
log 8 N
c(A S(f))
+
(11)
N
Proof of 3.2 can be found in Appendix C Theorem 3.2 provides the generalization bound of the
optimal dual model trained on finite selected demonstration set under mild assumption that [|]]] is bounded. Intuitively, as the number of demonstration (and therefore the number of demonstration tokens) increases, the generalization error decreases, which consistent with existing experimental observations [iie et al. 2021 Garg et al. 2022 Wang et al. 2024].
4 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 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 202T]. 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 || [ 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 by
nD (wvx,) W ( xx)) +
a 1 2m
(12)
where a is hyperparameter. Equivalently, the attention output Eq will be modified as
h Wv Xb. (1 ((11-)x,, x)x softmax (w (WxxX)" WoxT
(13)
This modification is equivalent to retaining less prompt information for query token during aggregation and relatively more demonstration information will be attended to.
7

---

Dix -i--lll each step
Regulariedd Models
Augmented Mode's
Negative Modells
35
2:00
05
1-44
owb
a-0.5
30
0,642
normal
ke1, Fw0.1
normal
5-0.5
1.00
1:00
kw3, Fw0.5
20005
eoochs
epochs
opochs
Figure 3: The equivalence between ICL of one softmax attention layer and gradient descent, along with analysis on different model modifications. Left Part: 1 h la as the gradient descent proceeds under setting N = 15; Remaining Part: the performance for regularized models (Center Left), augmented models (Center Right) and negative models (Right) with different settings.
Attention Modification inspired by the Data Augmentation: If 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 learn deeper representations. Thus, more complicated augmentations can be considered. Denoting these two augmentations 9I and G2. the loss function will be modified
C= DD Σ len (Wva:)) W (((g((Waa))))
Correspondingly, the attention layer can be adjusted as,
1/ 91 V X )softmax
(14)
where (-)-) and (-) 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 91 and 92- 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 crucial role in preventing feature collapse in contrastive learning methods while the representation learning process of ICL only brings single pair of features closer, lacking the modeling of what should be pushed apart, which could potentially limit the model's s ability learn representations effectively. Therefore, we can introduce negative samples to address this:
N Σ (Wv,,) W Wo.,K 3 = x2 L f D ==1
Σ Cj
W
JEN(O)
where (iii is the set of the negative samples for æ; and a hyperparameter. Correspondingly, the attention layer is modified
W WV {D. Xr) softmax (Wxx)* WWrTT///T
(15)
where XD 1 $2, n.d. 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.
5 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) it potential improve the attention mechanism from the perspective of representation learning?
Linear Task Setting: Inspired by Von Oswald et al. [023a]] 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 Wt where every element of W
dd, Xd, is sampled from normal distribution W N 1) and t from uniform distribution t ~
8

---

U1,, We set d, = 11 and d₂ = 1. Then, at each step, we use generated {z, = | s N+1 to form the imput matrix X where the last token will be used as the query token and the label part
N+
J=l i=1
will be masked, that is, NN++ = [4;; 0J. Here we consider only one query token (7 = 0) and we denote 4'711 = N1 to maintain consistency of notation in Section 2.1 Finally, the attention layer is trained to predict 1 to approximate the true label SN+1 using mean square error (MSH) loss.
Model Setting: It is worth noting that 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= x= ||= where ~ (0, I). We set the dimension of the random features as d, = 100 (d + dg) = 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 (Φ(z) and its corresponding training data and test input according Theorem 3.1
We perform three experiments under different random seeds for linear regression tasks with the results of one presented in Figure In addition, we also conduct more experiments including these on trigonometric, exponential synthetie 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.
To answer the first question, we generate the test input Xtstt using the same method as training and obtain the ICL result of the query token h On the other hand, we use X to train the dual model according to Theorem 3.1 get the test prediction ÿtest- The result is shown in the left part part of Figure B] It can be observed that after
Equivalence Between ICL and Gradient Descent:
N = 15 epochs training on the dual model, the test prediction ÿtest is exactly equivalent to the ICL inference result h1 by one softmax attention layer, which aligns with our analysis in Theorem 3.1 More detailed experiments can be seen in Appendix E.
Analysis on the Modifications: In Section 4 we discussed different modifications to the attention mechanism from perspectives of contrastive loss, data augmentation and negative samples. Here we call these modifications regularized models, augmented models and negative models respectively.
More details of modifications for self-attention mechanism can be seen in Appendix D
For regularized models, we vary different 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 a > 0, the regularized models converges to a poorer result while when a O, the model converges faster and achieves final results comparable to the normal model without regularization (a- 0). At least for this setting, this is little contrary to our initial intention of applying regularization to the contrastive loss where a should be positive. We explain it that the appropriate a 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 (-) and g2(-) as data augmen- tations enhance the value and key embeddings respectively in Eq (14) and choose GELU Hendrycks and Gimpel] [0166] as the activation function. It can be observed in the center right part of Figure 3 that when only use I2, 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 (1z) 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 (5)). 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 3 in Eq [5) 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 3 = 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 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 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. [022], which utilizes the dual form to understand ICL as an implicit fine- tuning (gradient descent) of the original model under linear attention setting (Aiserman et al.. 1964 Irie ctt al. [022]. 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 shorteomings in the evaluation metries 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- 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- lask-specffic training objectives (or manually specified causal language modeling objective Dai et al. 2022) A more formal and detailed comparison can be found in Appendix FF]
Additionally, many other works also link ICL with gradient descent, aiming to explore the Trans- former'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. [203aa] 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). Mahankkai 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. [203a] under certain assumptions. These works, both experimentally and theoretically, often focus on specific linear regression tasks (y- = wTa. x) and specific structured input format where each token takes the form [z. y) consisting of the input part and the label part y. In addition, the label part of the final query to be predicted is masked, represented as [ae, Oj. Subsequent works have expanded this exploration under more complicated setups, including examining nonlinear attention instead of linear OncentionChhen 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, 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.
7 Conclusion and Impact Statements
In this paper, we establish 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 in-context learning, which does not present any foresccable societal consequence.
10

---

8 Acknowledgements
We sincerely appreciate the anonymous reviewers for their helpful suggestions and constructive com- ments. This research was supported by National Natural Science Foundation of China (No.62476277, No.6207623), Beijing Natural Science Foundation (No.4222029), CCF-ALIMAMA TECH Kanga- roo Fund N.o.CCF-ALIMANA
OF 2024008), and Renmin University joint program on Information Retrieval. We also acknowledge the support provided by the fund for building worldelass universities (disciplines) of Renmin University of China and by the funds from Beijing Key Labo- ratory of Big Data Management and Analysis Methods, Gaoling School of Artificial Intelligence, Renmin University of China, from Engineering Research Center of Next-Generation Intelligent Search and Recommendation, Ministry of Education, from Intelligent Social Governance Interdisci- plinary Platform, Major Innovation & Planning Interdisciplinary Platform for the "DoubleFirst Class" Initiative, Renmin University of China, from Public Policy and Decision-making Research Lab of Renmin University of China, and from Public Computing Cloud, Renmin University of China.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, Advances in Neural Information Processing Systems,
2023.
MA Aiserman, Emmanuil M Braverman, and Lev I Rozonoer. Theoretical foundations of the potential function method in pattern recognition. Avtomat. i Telemeh, 25(6))917-936, 1964.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-context learnin?? investigations with linear models. arXiv preprint arXiv:2211.1566 2022.
Shun-ichi Amari. Backpropagation and stochastic gradient descent method. Neurocomputing, 5(4-5): 185-196, 1993.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306. 04667,, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33: 33:1877 117-17701,,200
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems, 33:9912 19-994,,9 2020.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1697-1607. PMLR, 2020a.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning. pages 1597-1607. PMLR, 2020b.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 15750-15758, 2021.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.04229 2020c.
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.065528 2023.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.1479. 2020.

---

Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In-context learning with transformers: Softmax attention adapts to function lipschitzness. arXiv preprint arXiv:2402.11639 2024.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in- context? language models secretly perform gradient descent as meta optimizers. arXiv preprini arXiv:2212. 10559, 2022.
Gilad Deutch, Nadav Magar, Tomer Natan, and Guy Dar. In-context learning and gradient descent revisited. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume Long Papers), pages 1017-1028, 2024.
Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. Causallm is not optimal for in-context learning. arXiv preprint arXiv:2308.06912, 2023.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey for in-context learning. arXiv preprini vrXiv:2301.00244 2022.
Pascal Esser, Maximilian Fleissner, and Debarghya Ghoshdastidar. Non-parametric representation learning with kernels. Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 11910-11918, 2024.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30555 2022.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self--uuprrised learning. Advances in neural information processing systems, 33:21271 13:21271-21224, 2020.
Jianyuan Guo, Kai Han, Han Wu, Yehui Tang, Xinghao Chen, Yunhe Wang, and Chang Xu. Cmt: Convolutional neural networks meet vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12175-12185, 2022.
Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. In-context learning of large language models explained as kernel regression. arXiv preprint arXiv:2305. 2023.
Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg- Kirkpatrick, and Graham Neubig. Towards a unified view of parameter-efficient transfer learning. arXiv preprint 1r.iv:2110.04366 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729-9738, 2020
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
Stti:1606.04415 2016.
Kazuki Iric, Róbert Csordás, and Jürgen Schmidhuber. The dual form of neural networks revisited: Connecting test time predictions to training patterns via spotlights of attention. In International Conference on Machine Learning, pages 9639-9659. PMLR, 2022.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. pages 5156-5165, 2020.
Jacob Devlin Ming-Wor Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of naacl--HI.T. volume 1, page 2. Minneapolis, Minnesota, 2019.
Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. arXiv preprint 1xXiv:2301.07065 2023.
12

---

Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 55(9):1-35, 2023.
Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. Soft: Softmax-fred transformer with linear complexity. Advances in Neural Information Processing Systems, 34:21297-21309, 2021.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576. 2023.
Andreas Maurer. A vector- pectr-contracttin inequality for rademacher complexities. In Algorithmic Learning Theory: 27th International Conference, ALT 2016, Bari, Italy, October 19-21, 2016, Proceedings 27, pages 3-17. Springer, 2016.
J. Mereer. Functions of positive and negative type. and their connection with the theory of integral equations. Philosophical Transactions of the Royal Society of London. Series A, Containing Papers of a Mathematical Physical Character, 209:415-446 1909. ISSN 02643952. URL http: // www. jstor..or... aarg/stabee/9100
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018.
Hyun Oh Song, Yu Xiang, Stefanie Jegelka, and Silvio Savarese. Deep metric learning via lifted structured feature embedding. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4004-4012, 2016.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1wXiv:187..0311 2018.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint 4rXiv:2103.02143, 2021.
Isaac Reid, Krzysztof Marcin Choromanski, Valerii Likhosherstov, and Adrian Weller. Simplex random features. In International Conference on Machine Learning, pages 28864-28888, PMLR, 2023.
Adam Roberts, Colin Raffel, Katherine Lee, Michael Matena, Noam Shazeer, Peter J. Liu, Sharan Narang, Wei Li, and Yanqi Zhou, Exploring the limits of transfer learning with a unified text-to-text transformer. Technical report, Google, 2019.
Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar. theoretical analysis of contrastive unsupervised representation learning. In International Confer- ence on Machine Learning. pages 5628- 5637. PMLR, 2019.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355-9366. PMLR, 2021.
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Do pretrained transformers really learn in-context by gradient descent? arXiv preprint arXiv:2310.08540, 2023.
Yuandong Tian, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning dynamics without contrastive pairs. In International Conference on Machine Learning, pages 10268-10278. PMLR, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Iliia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev. Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. pages 35151-35174, 2023a.
13

---

Johannes Von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint (rxii:2309905888,, 2023b.
Alex Wang. Glue: A multi-task benchmark and analysis platform for natural language understanding.
arXiv preprint 1rXiv:1804.000 000461, 2018.
Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. Label words are anchors: An information flow perspective for understanding in-context learning. arXiv preprint arXiv:2305. 1416600 2023.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are latent variable models: Explaining and finding good demonstrations for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sehastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.0704 20660782 2022.
T Wolf. Huggingface's 's transformers: State-of-the-ant natural language processing, arXiv preprint arXiv:1910.03771, 2019.
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non- parametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3733-3742, 2018
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.00 1.00000, 2021.
Yue Xing, Xiaofeng Lin, Namjoon Suh, Qifan Song, and Guang Cheng. Benefits of transformer: In- context learning in linear regression tasks with unstructured data. arXiv preprint arXiv:2402.00743, 2024.
Felix Xinnan X Yu, Ananda Theertha Suresh, Krzysztof M Choromanski, Daniel N Holtmann-Rice, and Sanjiv Kumar. Orthogonal random features. Advances in neural information processing systems, 29, 2016.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306. 300.:2927.00977, 2023a.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. / 19420, 2023b.
14

---

A Details of Theorem 3.1
We repeat Theorem 3.1 as follows and provide proof and more discussion for it.
Theorem A.I. The query token W1r41 obtained through ICL inference process with one softmax attention layer, is equivalent to the test prediction Gtest obtained by performing one step of gradient descent on the dual model ((x) = Wo(x). The form of the loss function C is:
1
C
(Wyx;) W WW(WK)))
(16)
nD
where 7i is the learning rate and D is constant.
Proof. The derivative of L with respect W should be:
ac
N =
Σ
1 DD
Wvz, @ &
ow
Thus, after one step of gradient descent, the learned W will be
ac W W. "ow Wo +
Σ WWrzo
(17)
where Wo the initialization of the reference model and 7) is the learning rate. So the test prediction will be
N Σ
Blest Wod (WQz'T+) +
Wvx, 3 (WWx))
(Waz'++1)
(18)
D
On the other hand, from the perspective of ICL process with one attention layer, with Eq 7 in our mind, can rewrite Eq
(WKX)² (VxX)"Wox'H
h/441 V ssofmma
Vdo
1
Wr W XT| di d( 6(WoaT41)
D
1
D' [V. Vzl (0( of Kr))² p(q).
X,l, K Ky Wx X Wq for sim- plify and D' o ) (())) + (Kr) ((q) 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 follows
where VVD Vzl
1
hT++1 D' Vro KT) ((g) + VDO(KD) KD) (o)(
(
D'
N
Wi)(g) + D'
Σ
v & ((K (8) ((9)
where W 1 (ro((r)²² and ((()) K/ D are the i-th column vetors respectively.
D7
Then, in Eq when setting the initialization Wo that
W/ and the constant D D', we will find
N
Viest = Woo(g) +
Σ
2 ok( ((q) = h/
(19)
D
which means Stest is strictly equivalent to h Thus, have completed our proof.
15

---

Given a reference model f(x) = Wo(z), by comparing Eq (9) and Eq we can easily observe that the gradient descent on the loss function L applied to f(x) the dual form of the inference process of ICL, where ( of D' and ((q) play the roles of backpropagation signals, training inputs and test inputs respectively. Recalling the form of Eq + we can interpret the Wo as the initialization of the weight matrix which provide the information under the zero-shot case while the second part in Eq (T9) shows that the demonstration examples in ICL acts as the training samples in gradient descent. The reference model f(x)) W (ax), initialized with Wo. will have test prediction Yteat = +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.
B Extensions to more complex scenarios
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.
In Theorem 3.1
B.1 Extension to one Transformer Layer
As for one Transformer layer introduced Section 2.1 we define the new dual model as
f+(x)= W (()) + b.
(20)
We will show that after performing gradient descent W, the test output ÿtest = I (W (WQrlr+1)
will be equivalent T Our theorem is given follows.
Theorem B.1. The output n1 of ICL inference process with one Transformer layer, is strictly equivalent to the test prediction of its dual model f+(x)= W ( + b, where f ((x) is trained under the loss function C formed as
N
C C=
(W=Wvz:)" (W (WW@WWaa kx:) b) Σ
(21)
DD
where 17 is the learning rate, is a constant, and WE will be determined once the specified pre-trained model, demonstrations and query tokens given.
Proof. Recalling the proof of Theorem
we can rewrite Eq I as
11411 Woo (W (QQ
+
Σ W
(WQRTT+)
(22)
where D 1 X (WQx
ize the attention scores and Wo = { Io WR k X T TW is a constant to normal- input for the FFN sublayer and the Eq ② can be rewritten
Furthermore, h'x41 will be taken as
+ b,) + b2 W2 Wi + W2 + b2.
where Las E Rdxd is diagonal matrix whose i-th diagonal element will be one (W (Wih'++++b) > O otherwise be zero. We need to note that this process is reasonable: for given demonstration and query tokens, once the parameters ( (We. WK. Wv, W1, b,} } of the Transformer layer are fixed after training, IMs will be determined implicitly (otherwise, IM would be function that varies with these settings). For simplify, we rewrite Φ'+1 as
1/111 = WF + bF,
where Wr WTT W1 and by = W.Inb + by. Furthermore, expanding h441 in the above Equation, we get:
#T WrWood (WQrT=1)
Σ
Wr Wva a (o o (Wou'T+1 + by
D
(23)
= - +
Σ
1
WvWWa @ (Wxx))
$ (Wowr4+) br.
D
1=1
16

---

On the other hand, we define reference model:
(*) W (x) + b,
where (b.) is exactly the mapping function satisfying Eq 7 to approximate the softmax kernel. Given the loss formed in Eq (2T). reference model when taking Wia x2; as input, that is,
note that the right part in & is exactly the output of this
C= -= TD Σ (W>Wva.) (Wo(Kx..) +b)
1 nD
Σ
(wrWva))
I (Wxv.).
We can calculate the derivative of L with respect to W as
OC
N
Σ
OW
DD
Suppose that the weight matrix W in the reference model (x) initialized as Winsts then using one step of stochastic gradient descent (SGD) will be updated as
[Amari 1993 with learning rate 7) the weight matrix
0C Winit init 7) oW
W
init +
Σ
Wr Wvwww @ & D
1
Compared to Eq 23 after one step update to W, the output of the reference model will be
we can set Winit WrWo, b- by and take Woa'44 test input. Then,
ft +W@zz++)) = W @ W WorT++1 +
N
W énit
1
WyWya, (WKa.)) @(WQ2'441) +
D
im1
W. Σ D 1=1
Wt Wva, 2 (WWx2;
& (WoaT++1)
by #
which implies that if we initialize the reference model f+(x)= Wo(x)++ + with Winit WF Wo. by. then after one step of gradient descent for W, the test output of f (W(WWNN+1) will be identical to the ICL result of one Transformer layer. Thus, call the reference model with setting init = WFWi Wo. b = by as the dual model corresponding to the ICL inference process. Finally, complete our proof.
Now, we discuss Theorem B.I from the following perspectives:
Training set and test input: In fact,
observe that the loss function L can be seen as the sum of inner products of N vector-pairs. In Eq (2T), the right vector happens to be the predicted output (6)
(2 (td)
W
std + b of the dual model for training input (( (6)
z std WKZ, Correspondingly, the vector on the left can be regarded as the true label
y Ustd- WP Wya, In other words, it be seen that the dual model performs one step LN SGD given training set fz std V i= on W using the loss L:
L
D
std
And then taking Stest Wor; as test input, it finally output the prediction ytests which achieves the ICL result N Compared to Theorem 3.1 after introducing the FFN layer, the main difference that the labels of the training data become y = Wr Wya;; instead (2) of Votd Wy z;,- Additionally, compared f(x). an extra bias b is introduced in the new dual model (æ=), which also have a different initialization W W, rather than Wo. We also need to note that in the dual model f (x), only W is trained, while b remains unchanged after initialization.
17

---

grod
M22
aW)
Higher Dimensionall Space
bise
Leoe-aak
@ W 2.)
W
Wi,
w,
preeiil fon-nak
W.dt
Figure 4: The representation learning process for the ICL inference by one Transformer layer.
Potential Low-rankness of Wp: Noting that WF W211s W1 where W1 E Rddx W2 E Rdxd, IM E Rd xdy (here we assume that d; = do = d for simplify), the rank of We will satisfy
Rank (WF) & min (Rank( W1), Rank W2), Rank(Iu))
We observe that IM is diagonal matrix with elements being zero one, and its rank is determined by the number of non-zero elements. Here, we can make a mild assumption that set Rank Wi) Rank even for any random square matrix as will be non-singular with probability 1. addition, also assume d, d which is consistent with settings in practice. Therefore, get Rank Wi) = Rank( 22) = d, and the upper bound of Rank (WF) will be
minfd, d. This assumption is quite mild as
Rank(Wr) S min {d, Rank([as))
Thus, we can find that if we want to avoid losing information, WP should strive to maintain Rank( JM) d which will be more easily achieved d2 becomes larger than d. Otherwise, Rank( (M) ) is likely to gradually decrease with an increase in the number of Transformer layers. This explains the necessity of setting d,, d in practice. In some cases where Rank Rank(IM) < d. meaning that the number of non zero elements in IM or positive elements in W, WihN+1+ + b less than the upper bound of Rank(Wz will be Rank (Ias) and the lower bound of Bank(Wa will be given as
Rank(Wr) 2 Rank W2 + W11 Rank (Ts) = Rank(Tar),
which implies the rank of Wp will exactly equal to Ra We should note that this condition, i.e., Rank (Ias) < d, is easily satisfied when d, or when d, is slightly larger than d (for example, d < d, < 2d in an expected sense). Thus, we conclude that WF has the potential low-rank property.
Representation Learning Lens: For encoded demonstration representation æ, the key and value projections will generate a pair of feature Wxæ; and Wvæ; to create a certain distance between data representations rank transformation WF is applied to the Wyz; attempting to compress some information which increases the difficulty of contrastive learning and forces the model learn better features; on the other hand, ((-) projects Wxæ, into a higher-dimensional space to capture deeper-level features. Finally, we need train the weight matrix W. which maps o( (WW x ax;) back to the original space, aiming to make the mapped vector as close as possible to Wr w, This interpretation is illustrated in Figure
space. And then, the one hand, potential low-
B.2 Extension to Multiple Attention Layers
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,
18

---

RP
R
Atterrio Laper
Arrec
Attinn (A8
11
11
Soqueaee of
Maai (
Go) 9 (
Plal whlp
e1 S (alliig yn1h
data pronse
training datases (Klet Mlle iw?
initialized (-th dual model for (as) (a))
b updale
12
test Inpat (#le ) ind
tamined
(x) ((a))
Figure 5: Illustrating the ICL inference process of multiple softmax attention layers from the perspective of dual models. The layer-wise process of ICL. can be viewed as a gradual gradient descent the dual model sequence. The datasets used for each gradient descent, including training data and test input, are obtained from the outputs of the previous dual model before and after training.
assume we have only query token, that is, there is query input before the considered query token. With the assumption that 7 = 0, to maintain notational simplicity, we use N++1 to represent the query token here instead of 4 1 and the input will be X = [X ØN+1]. We assume that there are attention layers and the output of the 1-th layer X can be expressed as:
[ 2 1 = (L-), w l w( (1)
(w!l) (1-1) w (1)
!1)
HD
((-1) (1-1)
H (
w H ll Softmax
Q
Vd
(w K
(1) (1-1)) w
H
w H (1-1) Softmax
N+1
N+1
dd
where we set (0) X Dx xN as the initial input. And the final output of the query token
(L) is h N+11 h N+1- Here, we assume that after training, the parameters (!) w w e Rdoxd, are fixed and we set d, = d, = d.
Next, extend Theorem B. the case of multiple softmax attention layers. Formally, present our result in the following theorem.
!) Theorem B.2. Given L softmax attention layers whose parameters (wil) w w W are fixed after training, the ICL output of these layers is equivalent to sequentially performing one step gradient descent on sequence of dual models (f)) = (() (())())) where the loss function for the 1-th dual model is:
N ( (( w ww
(24)
D
==1
(1-1), the output of the (1- 1) attention layer for the i-th token, rj the learning rate and D) constant. The inpus for the 1-th dual model is generated by the trained (1 - 1)-th dual model.
where by
19

---

Proof. Given H (1-1) the input for the 1-th attention layer, the inference process of h (1) N41 is
(W 11
11)= W
h
(1-1) N1
NN+1 W ( H Softmax
vd
Sownh--y h2 ((--)) W.
-W (W 1 (
o W
Φ (wonk=ll
D
where D (W (1)
H H (o 10 ( a constant to normalize the attention scores
and ( DUT ( (1-1) (W (Z According Theorem 3.1 we can easily get
(1)
the dual model fa hN+1
w (1) (h) where the initialization is w
(t)
W (l) Given the loss
iiii
function (()) formed as Equation 24 and training set {=
(6) where
(6) - h (1-1) and
stdd 3std i-1
std
(4) w h (1) we perform one step SGD with learning rate on weight matrix w (!) and will get std trained dual model:
f( W ((x) (W-l)
init + (o()))(x))
N
= W
W (-11) @ Φ @W (-1)) Φ(x)
D
(0) (1-1) Taking test input as ttsst w( h N-1 the prediction f()) (z will exactly equal to h N+1
H through the trained dual model (x)). And after
((+++) w(l-") w((t1) projections, H (1) will constitute the training well as the test in- Next, we will show how to obtain
put for the next dual model p((4+1) (x))
Keeping the initialized dual model find (w) and the trained one (x) in mind, we can compute the demonstration token output h
1,2, ... N) of (-th attention layer
(W (1) H D 10) (-
(-11
h
( D
Softmax
Vd
N
1
(1-1)
@ (0 (1111) 2 w
hi
D
N
1
w
a1 h
2 d mn h /1-1)] @ (w!
D
D
D
N
w
Σ
W 11 h2 (1-1)
@ Φ W 1 n ((1-))))
w!ll) snic
Φ (wi" (--11)
D
Exit
D
D
D
(
( h (1--))))
f(ll (w(( (1 Sinie (w(l)) 1-11)
h (1
D
init
(
(25)
where D 1 (W K f ((
h (11)) is a constant to normalize the attention scores for 1, Therefore, once we obtain the trained dual model (h). we can use the Eq 25 to get the demonstration token output ! [h h ...., h..,, These demonstration token outputs, along with the output h h for query tokens, will together constitute the training set and test input for the next dual model ((+1) (h).
(h). This process continues layer by layer until we obtain the ultimate ICL output h (L) - 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, sequence of dual models is correspondingly generated. Although these dual
20

---

models have the same form f()) = W Wo(æ), 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 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 Figure
C Proof of the Generalization bound
C.1 Proof of Theorem 3.2
In this part, we provide the proof regarding the generalization boundary in Theorem 3.2 our theorem as follows:
We restate
Theorem C.1. Define the function class as F -W((k)) [|
w) and let the loss function defined as Eq 110.. Consider the given demonstration set as S = { where S C ST and ST is all possible demonstration tokens for some lask T. With the assumption that [| (WWo(Waae least 8 for any EE
$ P, then for any 8 >> O, the following statement holds with probability
wpdo SypdoVTr(KKs Tr((Ks) N
log
< S +()++000
+
(26)
Proof. Our proof is similar to the Lemma 4.2 in Saunshi et al. 20 but here we focus on a different function class. Firstly, consider the classical generalization bound based the Rademacher complexity of the function class which can refer to Theorem 3. in Mohri et al. (2018). For a real function class G whose functions map from set Z to [0, 1] and for any 6 0, if S is training set composed by N iid samples {x; J.N then with probability at least 1 2. for all g G
N Σ g(z,)+ + (Rs(G)
log 2N
[g(())] S
+ +
(27)
N
where Rs(G) is the traditional Rademacher complexity. By setting S exactly the demonstration set and G = (x) (Wvx)" W (WWxx=))|ll W|| S w) we can apply this bound to our case.
Then, we construct function class f { [ = [f()): Wvai = (W W W S w
whose functions map from S to R2d., Next, we will first prove Rs(G) - 2p (F) and to do this, we need to use the following Lemma:
Lemma C.2 (Corollary 4 in Maurer (0166). Let Z be any set, and S = (2 1 E MM Let F be class of functions f Z + R and h 1 - R be L-Lipschitz. For all f E F, let 9j h Then
E a-(=1]M
sap (9F10))
S V2T
sup (a. (fis)) jef
(28)
feF
am(+1]mM
where j.s (F((z)) tE|m].jE[M]
We apply Lemma C.2 to our case by setting = Rd. S to be exactly the demonstration set, F to be the function class we constructed and n 2d. We also use : Rdd. R
where l((x) = -x Ld,+1:2a, and thus we have 9)(x²) h(f(x)) = h(f(x); Wvæ)
-(wxx) Wo(WRan We can find that (x) (x( and the left side of inequality 288 is exactly (G)
Then we can see that h is /2p-Lipschiie with the assumption that Wy ll W @W S and have Rs(G)= 2pRs Now using Lemma C.2 and the classical generalization bound (27
21

---

we have that with probability at least
(Rs(F) N
log
[()) 2
+
(29)
Let I E arg JeF [(). According to Hoeffding's inequality, with probability at least
2 we
have that E(f*) < C(F) + 3v
log R
Combining this with (2)), the fact that Ż(} - E(f*) and
2N
applying a union bound, we can get that
(F) +
log 8
c( S00++
(30)
Next, we give the upper bound for Rs(F).
2Nd.
Rs(F)=
E an ao(1l]2ad
sup
Σ oe
of(fis)e
(Definition of Rademacher complexity)
W Sw
1=1
do
sup
T f (Wxx,)
(Wyæ; is independent of W;)
a-((1) Ndo
|
==1
de IW
sup || 1=1
Σ A=1
J. (W wx)))
(By Cauchy-Schwarcs inequality)
wdo
E
0 WK,)
(Using the fact that [ S w)
N
S wdl.
BONSHIIN
Σ
a W k 2i
(By Jensen's inequality)
wd,,Tr(K
Substituting the upper bound of (F) into 30). we will get that
upd,y Tr Tr(Ks)
log
<()) s L()) +
+
(31)
N
Thus finish our proof.
C.2
Extension to negative models:
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 we consider the following representation loss:
1 K
C(F) = EZ~DA
Σ (Wo(k)))) (Wvz- Wv2,)
where we consider sampling negative samples for each 24 and x, denotes the j-th nega-
tive sample for token Correspondingly, the empirical loss will be considered L() (Wo(W<=))) (Wvz; Wyz,g) where Zij is the j-th negative sample for
- E
Jy. Then, by retaining the other definitions in Section 3.3 obtain the generalization bound as
corresponding Theorem 3.2 we can
log N
<() s C()) +0 wpdo Tr(Ks)
+
N2
rN3
22

---

where r K 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 N FN thus the former term dominates, which means the reduction in generalization error due to an increased proportion of negative samples limited. Nevertheless, we do not rule out the possibility of a tighter generalization bound, which is promising direction for future research.
Proof Sketch. The proof process is similar to that of Theorem 3.2 The main dif- ference lies in the fact that we should firstly define the function class G E (WW()))) (Wvx;- -Vvz) || S we
the classical bound. In
[ S w whose func- tions map from S to R(K+2)d,, Similarly, when using Lemma C.2 we set Z = Rd,, F be the above function class and n = (K + 2)d.- We also use h RKK+2)) + IR defined as h(x) = - k
addition, we define F = [f(x) = [f(x): Wyæ; WyaT:.... Wyak
E j=1 TT.d (Id,+1:22d - I(j+ +(4+))))) Then we notice that
Oh
1
K
(Eds 1:24. - 36+4))(++1(542)))
1z
K
jml
(32)
0h
0h
1
K 21:do
0zd.+1:20
With the assumption that ||wrll | (W S p. we can get that the Frobenius norm of the Jocabian J of has || S 4p² + p² R (5 + }pp Thus we get that is / + k0-
Lipschitz. The rest of the proof process is similar that of Theorem 3.2 Ultimately, we will obtain the aforementioned generalization enror.
D Details and More Discussions for Section 4
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 of the main body. And we also point out the corresponding modifications the self-attention mechanism, which are adopted in experiments.
D.1 More Discussion the Contrastive Loss
Although we have figured out the representation learning loss of the implicit gradient updates, it can (6) be observed that this loss function has flaw: due to the lack of normalization for you and ( when calculating the cosine distance, the loss can theoretically be optimized to negative infinity. To address this issue, we introduce regularization constrain the norm of W. that
D Σ (Wvz:)) Wo(Wkz;) (
a W W 2n
+
where a is 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,
ØL
w = ((t-1) -1 aw = (1 - w(-)) + Σ D D- Wyh, @ (W.h.).
i=1
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,
Wr = (1 ar) Wo)(()) + D-1
v 2 (K{()) olg).
23

---

which means more demonstration information will be attended to. This will directly result in Eq 13
This result be easily extended to self-attention mechanism. As for self-attention layer, if all other tokens adopt the same modification, the self-attention layer will become
(W aWr
H = Wr X Xsotmma
dd,
(WKX)* Vd
W
softmax
al
which leads to the model structure incorporating an operation similar skip connections. Further- more, to ensure numerical stability, we normalize the attention scores yielding:
(WKX)) ( al
H = Wv X Norm softmaa
Vd.
where Norm(-) is performed column-wiss 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 loss where a > 0 to prevent it from diverging to negative infinity, we find in experiments that this modification remains effective even when is less than O. We interpret this as possibly stemming from the fact that an appropriate a helps the attention block become full-rank, thereby better preserving information, which can be illustrated by Lemma D.I]
initial intention is impose regularization
the contrastive
Lemma D.1. Let the attention block A R=xx. There exists some such that, for any |a| < 6, the attention block A + aIn will become full-rank.
Proof. Define f(a) = I, i + +A), which is polynomial of degree n in a. Then, f(a) has only finitely roots. Let 01 02,..., Orp be the non-zero roots of ftt). Now, consider minf/as ||2l:|/2)......... jay)..... [ For < aal < 6, we can claim that f(a) = det(aI, + A) # 0. Thus, A aI, becomes non-singular (full-rank) and we complete the proof.
Lemma D.I provides possible for appropriate In fact, the selection of be quite flexible; for instance, similarly. when 8 = max{/a1) |, lazl..... |a,, } and |a] > holds, + aln also remains full-rank. Our experimental results related to regularized models will further illustrate the effectiveness of an appropriate a 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, more generally, from optimization standpoint. For example, (x) address the issue of non-normalized yord and (6), we can also modify the loss function from the perspective of ridge regression as:
L 2n
Σ
Wvx; W (
W W
24
And the optimal W will be
W = XoxXXoWxxXi) +DDI Wv (XxX).
Correspondingly, the attention mechanism will be modified to
== W= (Wox) [ x (W x)" + aDDI
x) (33)
where we neglect the normalization operation. This result is very similar to the mesa-layer proposed by Von Oswald al. (2023b). which optimizes linear attention layers under the auto-regressive setting. Here, we presented its form on softmax self aff-aattetion setting using kernel methods and explained from the perspectives
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.
24

---

D.2 More Discussion on the Data Augmentation
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 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 91 and 92, consequently, the process of contrastive learning will be modified as follows
N
C C
Σ ==1
19 (Wva:)) W (Wxe)))
DD
Correspondingly, the gradient update for W will become
( w(- - aw And from the perspective of ICI., correspondingly, the last token will be updated as
OC
w( + ED (Wya.) @ (l92(Wkx.))
h Woo(g) Dー1 g (()((()))) 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 = 91 ( X Jsoftmax 92 W X Wo )
Vd
where (1) and 92(-) will be column-wiss here. It is worth noting that here we have only presented the framework of using nonlinear functions data augmentations modify the self-attention layer and in the simplest case, we can set 91 (z) and g2(2)) as MLPs (Multi-Layer Perceptrons). However, in practice, it is encouraged to data augmentation functions that are tailored specific data structures. For example, in the case of CMT [Guo et al., 2022]. the used Convolutional Neural Networks (CNNs) can be considered as form of "strong data augmentations" suitable for image datas within our framework. We consider the exploration of various augmentation methods tailored different types of data as an open question for future research.
D.3 More discussion on the Negative Samples
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 most contrastive learning approaches, negative samples forcing the model to separate the distances between positive and negative samples the same time, that is,
introduce
B
C
(Wvx.) W Wo( + n
Σ (Wvx,) W (o(Wke))
D
[(1) (EN(()
1=1
==1
N
Wv
x; [i(1]
Σ æj JEN(
w Wo(WRaD
n
Σ (Wrā,)" W Wo(WK
25

---

(W()) EjEN() 2,, (i) is the set of the negative samples for æ and 3 is hyperparameter. As result, the gradient descent W will be modified
where 2;
Oc aw
W w(t-))
W((-1) D Wya; (WWx)))
Correspondingly, the ICL process for h2 will be
N
1 N Woo(Woar + D-1
Σ
@(Wox)--).
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 as
(( (WXX))
H Wy softmax
(34)
vdo
where (() .; In corresponding experiments, for each token, we simply choose other the &: least relevant tokens as its negative samples, i.e., the tokens with the lowest attention scores. Noting that here we simply use other token representations as negative samples for æ;, However, there are more ways to construct negative samples that are worth exploring (for instance, using noise vectors or tokens with low semantic similarity negative samples). For specific data structures and application scenarios, customizing the selection or construction of negative samples may be more effective.
E
More Experiments
E.1
More details of Experiments on Linear Task
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. [0023a], we choose to pretrain a softmax attention layer before exploring the equivalence proposed by Theorem 3. In fact, pretraining 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- specifcc knowledge. Additionally, our experiments are conducted in self-attention setting. When focus only on the last token, this is equivalent to considering the case with only one query token (I' = 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 Wt where every element of W E Rd.. is sampled from normal distribution Wi, ~ (N(0,)) 1) and is sampled from a Gaussian distribution ~ U (- 1)d: To facilitate more accurate estimation of attention matrices using random features and considering the limited learning capacity of single attention layer, we only a small value for d, = 11 and d, = 1. Then, at each step, use generated (æ; = [4,; si]) N+1 form the input matrix X while the label part of the query token is masked to be zero, that is, x [t;; 0] where we consider only one query token and denote 2/141 IN+1 1 to maintain consistency of notation in Section 2.1 The softmax attention layer is expected to predict & to approximate the ground truth value + We use mean square error (MSE) as the loss function, that is, for each epoch,
Natep
( 84+1
L
N Nstep
N+
where (0) and N+1 are the prediction and ground truth value at j-th step and Nstep the number of steps. We set N Natep 1024 for N + 16 which means the total number of tokens remains
26

---

(a) d 3
(b) d, = 12
(c) 120
(d) d, = 12000
(c) exact attention
Figure 6: The estimation the attention matrix by positive random features when varying d,
(a) d,
(b) d, 12
(c) d, 120
(d) d, 12000
(e) exact oulput
Figure 7: The estimation of the output matrix by positive random features when varying d,
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 different task. However, 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.
find it challenging for single attention layer effectively learn in
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 al. 2020]. The mapping function Φ RRd Rd. has the form of <(x))
2 /2 where ~ N(0, 7). 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, on the approximation of attention matrices and output, using Mean Squared Error (MSE) and Mean Absolute Error (MAB) as evaluation metrics, where we conduct 50 repeated experiments and calculated the average values for each value of d,, as shown Figure g can be observed that as the dimension of random features increases, the approximation performance gradually improves, with both errors reaching low level the end. We visualize the exact attention matrix and compare with the estimated attention matrices obtained using different values of d,, shown in Figure 6 Again, it can be seen that as d, increases, the approximation the true attention matrix improves gradually and similar results be observed for the analysis of output matrices in Figure 7
e
To obtain 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, 100(d, + de) 1200. Further- more, we visualize the exact attention matrix and the output with the approximation results, which are shown in the Figure - As can see, although some larger values are not estimated accurately due to the limited dimenston of the random features we select, the majority of the information still estimated comprehensively well. These findings indicate that choice of using positive random features as mapping functions to estimate the true softmax attention and conduct experiments is relatively feasible.
After the weights WQ. WK. Wy the attention layer have been determined, we generate test N tokens in the same way where the 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 h N+ N41
(1)
Here, also use 1 h to maintain the notation consistency in Section 2.1
27

---

Attention Error
Output Error d,
0.97
0.38
MSE MAE
0.30
MSE
MAE
0.30
0.05
0.20
0.04
0.03
0.15
0.01
0.30
0.00
400
200
1000
1200
400
500
1000
1200
(a) Estimation error of attention ma-
(b) Estimation error of attention ma-
trix vs. dr
trix vs. dr
Figure 8: The error of positive random features in estimating the attention and output matrices as d, varies.
attention
attention 1012
output 4 6 10
exect output
1012
e 10
10
10
14
14-
14
(a) exact attention matrix and its approxi- (b) the exact output and approximation mation
Figure 9: The comparison between the exact attention matrix, output and their estimated approxima- tions using random features under setting N 16 and d, 1200.
On the other hand, we construct a dual model f(æ) = W oox)) where o(-) is strictly equivalent to the kernel mapping function used in the attention layer. We transform the first N tokens the training set according to Theorem 3.1 and train the dual model using the loss formed by Eq ( In fact,
according to Theorem 3.1 after we perform one step of gradient descent on this training set, the test
prediction jtest = k
(2) N411 of the dual model will strictly equal h'y
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 the Figure 1O, Similar to the discussion in the main text, we can achieve better performance than the normal model with appropriate parameter settings.
Regularzedd Modals
Augmented Models
Negative Models
25-
emD
normel 9:5g:
normal
a-0.5
b=1,5=0:1
pw0.5
1055
$-0.1
5-0.
1.0
0.0
0:
0.5
spochs
opocis
epochs
Regularizdd Modals
Augmented
Nagative Models
2.5
2.0
0,50)
5-0.1
01
kel. 5-0.5
LON3
0.5
k-33, S-0.1
k=3, $-0.5
03
65
20
40
106
66
64
100
40 60
epochs
opochs
egochs
Figure 10: The performance for regularized models (Center Left), augmented models (Center Right) and negative models (Right) with different settings for different random seeds.
28

---

E.2
More details of Experiments on Different Tasks
In addition to conducting experiments on linear regression tasks, we also extended our experiments to involve trigonometric and exponential tasks.
E.2.1
More details of Experiments on Trigonometric Tasks
1iax - Ni- de each step
Regularized Models
Augmented Models
Negetive Mociels
0.273
sarmal
nermal
0350
0.23
d=2
9 9-6epp
2:
k=n B.o.1
kw6, f-0.5
1424, F-0.11
2
2w24, 3i00.
0.13
0.5
0.12
0.50
40 303
2.00
3.000
3.0
15C
-00
252
150
620
3:00
0140
246
260
22005
0pCCTS
opachs
opachs
Figure 11: The equivalence between ICL of one softmax attention layer and gradient descent, along with analysis different model modifications for trigonometric tasks. Left Part: |l - h l|2 as the gradient descent proceeds under setting N = 127; Remaining Part: the performance for regularized models (Center Left), augmented models (Center Right) and negative models (Right) with different settings.
For trigonometric task, we generate the task by 8 = cos Wt) where cos(.) is element-wise, W E Rd, is sampled from the normal distribution Wi ~ (00 1) while t sampled from the uniform distribution x ~ (0, π7 d In experiments, we found that for one softmax attention layer, learning nigher-dimensional tasks is challenging. Therefore,
only set d, 7 and d, = 1. At each step, use N +1 128 tokens {æ, [t;; { + and the total number of tokens remains unchanged at 16384. Compared to the setting + 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, $N+1 = O 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
= 16 of linear tasks, we observed that for more complex tasks,
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, hN+1 Vtest well as the label part (1) s(2) aligning with our analysis in Theorem [3.]]
N N
The performance of modified model during training process can be seen in the remaining parts of Figure 11 For regularized models, seen in the center left part of figure [1. the models when a < 0 converge slightly faster and reach better final results compared to the normal model (a = 0). For augmented models, we use as the same augmentation functions gi and g2 the ones in the linear regression task, that is, 91 (x)) 92 (x) = o( x) where (-) is GELU activation function. However, for 92 . we use ELU as the activation function. We can find from the center right part of Figure TT that, compared the normal model, using 91 alone and using 9i and 92 simultaneously as data augmentations significantly degrade the model's performance, including convergence speed and final results. However, using 92 alone yields comparable result with the normal model. Particularly, when using 93 the model accelerates its convergence speed. However, for negative models, the performance with the selected number of negative samples l and the parameter 3 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.
E.2.2 More details of Experiments on Exponential Tasks
For exponential task, we generate the task by & = exp( Wt) where exp(-) is also element-wise, W E [Rd, is sampled from the normal distribution Wy ~ _N(0, 1) while t is sampled from the uniform distribution æ ~ 1)4, We only set dç = and d2 = considering the limited learning capacity of one softmax attention layer. At each training step, use N + 512 tokens {æ, [t.;; s.]) and the total number of tokens remains unchanged 16384. Compared to the
29

---

1ix -/ each step
Regularied Models
Augmented Models
Negative Madels
2-16
0
mormal
3.
al
0.25
0.20
0.20
3.55
20005
e2ochs
apachs
opochs
Figure 12: The equivalence between ICL of one softmax attention layer and gradient descent, along with analysis on different model modifications for exponential tasks. Left Part: l|ütest - h ll2 as the gradient descent proceeds under setting N = 511; Remaining Part: the performance for regularized models (Center Left), augmented models (Center Right) and negative models (Right) with different settings.
setting N + 1 16 of linear tasks and N 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 1N+1 of ICL. inference is equivalent to the test prediction Ftest of the dual model after training, just as stated in Theorem 3.1 For regularized models, it can be observed that when a = 16, the model converges faster and achieves better result. For augmented models, using 91 or 92 alone data augmentations results in better performance. However, when both and 92 are used simultaneousll, the training process becomes unstable, so we did not show it in the center right part Figure [2] For negative model, similar the case in the trigonometric task, the different combinations of negative samples' number k and parameter 3 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.
E.3 More Experiments on Combinations
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, combination of regularized and augmented modifications sufficient. However, for the other two tasks, the results are actually worse than using regularized augmented modification individually (compared to Figures I] and I2). 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.
Exp Tasks
33
1.79
romes
23
38
ragu- rage
0.45
0.45
Hogi (AN)
normal
#ag tapo
0.2
0
1.3
2:25
0.25
0.3
0.00
20
382
150
333
3.55
30
130
epochs
epochs
Figure 13: Performance of different combinations on linear (Left), exponential (Center), and trigono-
metric (Right) tasks.
E,4 More Experiments One Transformer Layer
Similar to the experiments with one softmax attention layer, we also conduct experiments on Transformer layer (introducing one FFN layer after the attention layer) and trained its dual model based on Theorem B.1 As shown Figure 14 the inference result hx+1 ICL remains equivalent to the test prediction yrest of the trained dual model. Furthermore, validate the potential low-rank
30

---

Dost Dr+ step
Upper bound of Rank(W,)
(i, Rankllu))
Norm
liner
CO5
3.0
15.0
steps
The equivalence between ICL of one Transformer layer and gradient descent, along with analysis on upper bound of Rank(WF). Left: 1itest h/ lz as the gradient descent proceeds under setting N = 15; Right: the upper bound of (Rank(Wr p) when setting 12 and varying dh
Figure 14:
property of matrix WF, we explore its upper bound of rank. Noting that WF W2Ias W1 where W1 E Rd xd W2 E Rdxdi IM E Rd,, the upper bound of Rank( WF) is
Rank( (F) I min {d, dr, Bank(Ix))
where Rank( (M) ) is equivalent the number of non-zero elements in IM- We fix 12 while varying the values of dn. We generate 1024 sets of Xtest for different tasks and repeat the experiments times. Finally, we calculate the average upper bound of the rank of Wr. The results are shown in the right part of Figure 14] indicating that when d, 2 2.75d 33, the upper bound remains stable and equals d = 12. Otherwise, when d/, is set a smaller value, WF exhibits clear low-rank property.
E.5 More Experiments on More Realistic NLP Tasks
We supplement our experiments on more more realistic NLP tasks. We choose the BERT-base- uncased model (can be downloaded from Huggingface library [Wolf] 2011). hereafter referred to as BERT|Kenton and Toutanova 2019)) 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 "classi- fier.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 single 24GB NVIDIA GeForce RTX 3090. All experimental results are presented in Table Below, we discuss the various modifications and their performance.
For the regularized modification, (-0.5. 0011, 0, 0.5). As can be observed in Table except for RTE, the best regularized models outperform the original model on the other three datasets. However, 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 a.
consider different values of a, specifically selected from
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, do not simply select !1 and 92 as MLPs, i.e., I1 (Wyx) = Wzo W, Wyz). This design avoided because it could undermine the effort made during pre- pre-training to learn the weights Wy and Wrr leading to difficulties in training and challenges in comparison. Instead, we adopt a parallel approach, i.e., 91 (Wyx) new branch, a is the GELU activation function and the hidden layer dimension is set to twice the original size of Wvx. 92 (>WK))) WKZ cW (Wiz) follows the same format.
Wyx + cW (W,r). where is hyperparameter to control the inffuence of the
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 91 alone) proves to be more effective than other methods, both in terms of performance and the amount of additional parameters introduced. Using both 91 and 92 introduces more parameters, which is particularly undesirable for larger models. Thus, under the augmentation methods and experimental settings we selected, using 9i alone is recommended.
In addition, we do not rule out the possibility of more powerful and efficient augmentation methods. Our choice of 91 and 92 parallel MLPs is primarily motivated by the desire to make better use of
31

---

Model Types
Dataset
CoLA
MRPC
STS-B
RTE
Normal
Bert-base-uncased
56.82
90.24/86.27
88.29/87 .96
68.23
a = -1.0
0.0
79.01/68.87
57.23/60.16
52.71
a = 0.5
61.42
83.17/74.02
85.28/85.22
57.04
a -0.1
58.06
89.50/85.05
88.71/88.27
65.70
Regularized
= 0.1
58.34
90.59/86.76
88.12/87.88
64.98
Models
= 0.5
27.01
83.56/73.28
85.25/85.08
59.93
= 1.0
0.0
81.22/68.38 38
52.07/55.66 60
47.29
Local Best
61.42
90.59/86.76
88.71/88.27
65.70
g1 / 0.2
59.85
88. 11/33.33
88.56/88.22
68.59
g1 c
56.51
90.88/87.01
88.96/88.60
71.12
82
0.2
56.29
87.65/82.60
88.60/88.22 24
68.59
Augmented Models
g2
58.85
87.74/82.60
88.68/88.32
70.40
gi g2 / C = 0.2
57.32
89 89.62/85.29
88.48/88.11 19
71.12
gi & g2 c = 1
58.30
90.40/86.52
88.83/88.45
68.95
Local Best
59.85
90.88/87.01
88.96/88.60
71.12
0.1/ B 0.1 56.22
88.54/83.82
88.25/87.91
65.34
0.2/ B 0.1
57.92
90.00/85.78
88.22/87.8%
66.06
0.3/ [ 0.1
57.92
1/.3
88.26/87.90
67.15
Negative Models
0.1/ B 0.2
58.92
87.90/83.33
88.34/88.11
63.54
0.2/ 8 0.2
57.13
87.87/83.09
88.59/88.27
64.98
r 0.3/3 3 0.2
58.14
88.97/84.56
88.64/88.33
66.79
Local Best
58.92
90.00/85.78
88.64/88.33
67.15
Reg & Aug
56.56
88.54/83.82
88.86/88.66 60
68.59
Reg & Neg
58.11
88. 19/83.3
88.41/88.1 17
69.31
Combined Models
Aug & Neg
59.07
90,49/86.76
88.59/88
70.76
Reg & Aug & Neg
58.92
88.39/83.58
88.32/88.00
67.87
Local Best
59.07
90.49/86.76
88.86/88.60
70.76
Global Best
61.42
90.88/87.01
88.96/88.60
71.12
Table 1: Partial GLUE test results of different modifications. "Local Best" is used to display the best results for each modification type. where bolded results indicate the performance superior to the original model. "Global Best" is used to showcase the best results among all modifications, Matthews correlation, F1 scores/accuracy. Pearson/Spearman correlation, accuracy are reported for CoLA, MRPC, STS-B, RTE respectively.
the pre-trained weights W and Wv. 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 uses fewer parameters, and the original Wy and Wx are freezed- the newly introduced parameters are trained. In contrast, our approach aims to validate the benefits of introducing stronger nonlincar augmentation functions into the linear value/key mappings. Therefore, set higher hidden layer dimension (twice that of Wv or W k x) and also train Wy and WK simultaneously. This design 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 represents the proportion of tokens used negative samples, while B indicates the overall reduction in attention scores. We choose 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
32

---

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 a direction worth exploring in the future.
We also consider combining diffeerntt modification methods. Specifically, we choose a 0.1, g1/c and 0.2/3 0.1 respectively as the basis for combining the three types of modifica- tions, 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 the best performance achieved with augmented models alone. Therefore, we conclude that using all three modifications simultaneously 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 task-specific augmentation and negative sampling methods are all interesting directions worth exploring in the future.
F More Details about Related Work
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 implicit fine-tuning: More specifically,
let X = [X Xr) where X = [1 22.-.. .., xNn denotes the demonstration tokens and XT = [rr .z. -.... x be query tokens. On the hand, for ICL, they consider the output of W WQuTr+1 under the linear attention setting as
(q) Wv[XD. ( X
Wv XD WK X WZsL LinearAtten W X D, W Xx,))
WZSLq + Σ ((Wvz,) o (Wxz,) q
WZSL + AWicL4,
where Wzst.q is interpreted as the output in the zero-shot learning (ZSL) where no demonstrations given. On the other hand, they consider specific fine-tuning setting, which updates only the parameters for the key and value projection, that is,
(F) (Wv + AW, X + AW
( ZsL + AFFT)q
where AW k and AWy denote the parameter updates and they are acquired by back- pack-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 Frr, their focus is on establishing a connection between ICL and implicit fine-tuning on 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, consider the dual model f(x)= WØ(#) of the nonlinear attention layer, where the weight W are updated according to the following loss (presented Eq (9) in Section 3.2
N Σ D
C
where as 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
33

---

using this loss can be viewed from self-supervised learning lens: unlike in supervised fin-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 a; (so called self-suuer Therefore, modifications to this sule-supe learning loss will in turn cause modifications in the attention mechanism correspondingly, as we discussed in our work in Section offers several benefits:
We believe this perspective
By analyzing from the dual perspective, into an optimization process. Since optimization processes are well-known and have established thcoretical tools (for example, generalization error as mentioned in Section 3.3 this transformation can provide reverse insights into analyzing the model mechanisms.
transform the forward inference process
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 supervisdd) as mentioned above.
In addition, although we do not target specific tasks like linear regression as previous works mentioned in Section 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,] E R [[(++) where h, = [a,, y:] is sampled from some lincar task y = w x and the query token will be hx+1 [= mmm xN And the considered linear self-attention layer will take the constructed weights and query output as:
WK We [Tor
Taxd 0 W. =
Oaxd wo -1 P
1 N
(35)
hN+1 hN+1 + P(W H)WKH) WohN+1,
where wp is the underlying initial matrix. Then the label part of hy41 will has the form as #N+1 = ww #N++1 + Aw"xxx+ = - (wnj" N Σi= (wo x; y))=a JæN+1 = -NN+I, which is equivalent to the output -9N+1 (multiplied by 1) of the linear layer = w x where w initialized as wy after performing one step of gradient descent under mean squared loss Em
| xz - vil.a.
2N
In practice, the underlying initial weight matrix wo is set to be approximately 0 thus the test input can be formed as hN+1 [x, 0] (Von Oswald al. 2023a]. In addition, when reading out the label 1N- the test prediction nN will be multiplied again by -1, which can be done by a final N I). In this case, we first note that the dual model of the linear attention layer can be written as f(z) = Wz where W E (dd+1)x(d++) and similar to Eq
projection matrix (or equivalently, P
it will be trained under the loss below:
N
min W
Σ (PWyh)) WWKh,
(36)
By substituting the corresponding weights in Eq (35) where we replace P = N for the readout, the loss can be reformulated as:
N
x 0
min
L = W
Σ
10.
y W
(37)
N
Recalling that h; [x,, w/] is sampled from some linear task y TT.: x. we assume that [ || v ||w]] 1|l it can then be casily seen that the optimal solution for Eq (37) will be
0
W
(38)
w T
34

---

Furthermore, similar to Section 3.2 we take Wo as the input where Wo is constructed as Eq (35) and N 1 0,, the optimal dual model will output the result f (WohN+1) W Wo = [0, w" IN+1
00 yN+11 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 for future exploration.
35

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We confirm that the main claims made in the abstract and introduction accu- rately reflect our contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations at the end of the main body.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
36

---

We confirm that the paper provides the full set of assumptions and a complete (and correct) proof for each theoretical result.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: confirm that we have provided a detailed explanation our experimental setup.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
37

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
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 the code part.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
gwtp://wwiee/eeee Sawelee//oooSSbbbiicaa SonroddeummsssiooPPe.. www.issssoniiiiiiiooiiii for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have provided as detailed an explanation possible for all details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We conduct experiments with different random seeds for three times in the linear task setting, and the results indicate that the model modifications are effective to some extent.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
38

---

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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
We have provided sufficient information on the computer resources in Ap-
pendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate 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 (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We confirm that our research meets the requirements.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction..
10. 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.
39

---

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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include 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.
40

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
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 not available online, the authors are encouraged to reach out to the asset's s creators.
13. 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/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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.
14. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion 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.
15. 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 equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending 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.
41

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
42


Reviews:
[{"review_id": "mDOhr9r08K", "reviewer": "Reviewer_HnPE", "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.", "questions": "Please see the weakness section.", "limitations": "The paper has discussed the limitations sufficiently.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **Theoretical Depth**: It provides a rigorous theoretical analysis, including the derivation of a generalization error bound, which contributes to the theoretical foundation of Transformer models.\n* **Insights into Attention Mechanism**: The paper provides potential modifications to the attention layer, inspired by theory, which could potentially improve the learning capabilities of Transformers.\n* **Good Writing**: Clear Structure and Presentation: The paper is well-structured, with a clear abstract and a good introduction. It is easy for readers to understand.", "weaknesses": "* **Generalization to Other Tasks**: The paper's findings are based on specific tasks. It's unclear how well these insights would generalize to general language tasks.\n* **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the encouraging feedback, especially for recognizing the theoretical contribution and insights of our paper. Our response is detailed below.\n\n> **Weakness 1**: Generalization to Other Tasks: The paper's findings are based on specific tasks. It's unclear how well these insights would generalize to general language tasks.\n\nIn fact, our findings can be naturally extended to more realistic tasks. We supplement our experiments on more general language tasks. **More experimental settings are detailed in the (global) author rebuttal and the results is presented in the Table 1 of PDF** `(To AC and All Reviewers)`. Considering the limited computational resources and time constraints, we choose the BERT-base-uncased model to and select part of GLUE datasets to validate the effectiveness of modifications to the attention mechanism. More discussions on the results are provided in detail in our response to Weakness 2.\n\n> **Weakness 2**:  Empirical Validation: While the paper includes experiments, the extent of empirical validation is limited. We may need more experiments on real-world language tasks to verify the modification of the attention mechanism.\n\nTo answer the questions, below we discuss the main experimental results presented in the Table 1 of PDF `(To AC and All Reviewers)`.\n\n**For the regularized modification**, we consider different values of $\\alpha$, and as can be observed in Table 1, except for RTE, the best regularized models outperform the original model on the other three datasets. while we also note that when the absolute value of $\\alpha$  is too large, the model's performance declines significantly, so smaller absolute values for $\\alpha$ is recommended.\n\n**For the augmented modification**, we also consider applying more complex ''augmentation'' functions to the linear key/value mappings. Specifically, we adopt a parallel approach, i.e., $g_1(W_Vx) =  W_Vx  +  c W_2\\sigma(W_1x)$, where $c$ is a hyperparameter to control the influence of the new branch, $\\sigma$ is the GELU activation function and the hidden layer dimension is set to twice the original size of $W_Vx$. And $g_2(W_Kx) =  W_Kx  +  c W_2\\sigma(W_1x)$ follows the same format. Experimental results show that the best augmented models achieve better performance than the original model across all four datasets. Notably, using $g_1$ alone proves to be more effective than other methods and using both $g_1$ and $g_2$ introduces more parameters, which is particularly significant for larger models. Thus, under the augmentation methods and experimental settings we selected, using $g_1$​ alone is recommended. \n\n**For the negative modification**, we continue to select tokens with lower attention scores as negative samples. The parameter $r$ represents the ratio of tokens used as negative samples, while $\\beta$ indicates the overall reduction in attention scores. And the results show that the best negative models only outperform the original model on CoLA and STS-B, whereas their performance on MRPC and RTE is worse than the original one. This suggests that our simple approach of considering tokens with low attention scores as negative samples might be too coarse. A more effective method for constructing negative samples should be designed, which is a direction worth exploring in the future.\n\nWe also consider **combining different modification methods**. The results indicate that under our settings, the combination of augmented and negative modification achieves the best performance on CoLA, MRPC, and RTE, while the combination of regularized and augmented modification achieves the best performance on STS-B. However, their optimal performance is slightly inferior to the best performance achieved with augmented models alone. Therefore, we conclude that using all three modifications simultaneously is not necessary. With appropriate hyperparameter choices, using augmented modification alone or in combination with one other modification is sufficient.\n\nIn conclusion, the experimental results show the potential of our approach of improving the attention mechanism from a representation learning perspective. However, due to time and computational resource limitations, our experiments are conducted on a limited tasks and model size. More detailed parameter searches, validation across additional tasks and models, and the development of task-specific augmentation and negative sampling methods are all interesting directions worth exploring in the future.\n\n**Final Note**: Thank you for your detailed review. We are excited that you found our novel insights into attention mechanism. If there are any remaining questions, please do not hesitate to let us know.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZBNFAZHnR8", "reviewer": "Reviewer_cRiK", "summary": "The author's present a new way of linking in-context-learning (ICL) to gradient descent. \n\nThe author's are able to demonstrate that indeed a (simplified) transformer decoder layer ICL is equivalent to \"representation learning\".\n\nUsing the theoretical findings the authors are also able to propose extensions to the attention mechanism.\n\nThe paper shows theoretical proofs for these claims as well as additional experiments to demonstrate the claim.", "questions": "1. How could a \"realistic\" task look like for your evaluation?\n2. How could one practically extend your augementations to e.g. modern LLMs (even of `moderate' size, e.g. 7B)\n- I.e. what would these \"data augmentations\" or \"negative samples\" really look like?", "limitations": "Yes.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Strenghts:\n1. Excellent theoretical contribution to understanding in-context-learning\n2. Adding relevant proofs\n3. Using theoretical results to motivate better attention mechanisms", "weaknesses": "Weakness:\n1. 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.\n2. 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. \n3. A small section on how these results can be applied for bigger models in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for acknowledging the theoretical contribution of our paper. We have carefully addressed each of your concerns and considering the word limits of rebuttal, we notice that some issues might overlap, so we have organized our responses as follows.\n\n> **Weakness 1**: Slightly simplified transformer architecture is used (it would be interesting to see the skip-connection version) ...\n\nThe attention module, as one of the core components of the transformer, is crucial for extracting contextual information. Therefore, for the sake of analytical simplicity, we focus on comparing pure attention modules/mechanisms. Intuitively, our conclusions can be naturally extended to versions with residual connections, which are analogous to introducing a residual connection structure in dual models, i.e., $\\mathrm{Atten}(x) + x = \\hat{W}\\phi(W_Qx) + x$, where $\\hat{W}$ is the weight in the dual model $f(x) = W\\phi(x)$​ trained under the original representation loss in Eq (9). Although it is relatively straightforward to connect residual connections with dual models in form, the role of such residual connections in the learning process of the dual model's weight from a representation learning perspective still requires further investigation. Therefore, a more reasonable explanation is currently lacking. We believe this is an interesting direction for future research.\n\n> **Weakness 2** (The experimental setup and results are a bit limited ... as well as a somewhat \"realistic\" task would be interesting to see - and what the results mean qualitatively. )  \n> **& Question 1** （How could a \"realistic\" task look like for your evaluation?）\n\nDue to constraints of page sizes, we have simplified the description of the experiments for simulation tasks in the main body. In fact, more detailed experimental setups are presented in the Appendix E. We will provide more detailed supplementary information on the experimental setups if space permits in future revisions.\n\nFor the evaluation on more realistic tasks, we have supplemented our work with additional experiments. **More experimental settings are detailed in the (global) author rebuttal and the results is presented in the Table 1 of PDF** `(To AC and All Reviewers)`. \n\nRegarding the qualitative results: **For the regularized models**, smaller absolute values of $\\alpha$ are recommended. We interpret this as adding a small scaled identity matrix to the attention score matrix helps achieve full rank, thereby better preserving information. **For the augmented models**, to better retain the information from the pre-trained $W_V$ and $W_K$ weights, we adopted a parallel MLP approach for data augmentation $g_1$ and $g_2$. For example, $g_1(W_Vx) =  W_Vx  +  c W_2\\sigma(W_1x)$. The results suggest that enhancing the value mapping may be a more effective choice. **For the negative sample models**, we continued to select tokens with lower attention scores as negative samples while the final results showed limited improvement in model performance. We believe this may be due to the ineffective selection of negative sample tokens, and better methods for selecting negative samples would be interesting to explore in the future.\n\n> **Weakness 3** (A small section on how these results can be applied for bigger models in practice.)  \n> **& Question 2**(How could one practically extend your augementations to e.g. modern LLMs ...)\n\nWe believe that our supplementary experiments on pre-trained BERT provide more insights into this issue (see global rebuttal `(To AC and All Reviewers)`). Although the scale of the model and datasets selected is relatively small due to our limited computational resource and time constraints,  we believe similar approaches would naturally extend to larger-scale experiments.\n\nTake the augmented models as the example: we also considered using more complex data augmentations for the linear key/value mapping ($g_1(W_Vx)$ and $g_2(W_Kx)$). However, unlike previous methods used in simulation tasks, we do not choose $g_1/g_2$ as $g_1(W_Vx)=W_2σ(W_1W_Vx)$, because our experiments showed that this design struggles to fully leverage the pre-trained weights $W_V$ and $W_K$. Thus, we adopted a parallel approach, that is,  $g_1(W_Vx) =  W_Vx  +  c W_2\\sigma(W_1x)$, where $c$ is a hyperparameter to control the influence of the new branch ($g_2$ is similar). This approach introduces nonlinear \"augmentation\" while preserving the knowledge from the original pre-trained weights, making it easier to train. The results in Table 1 also prove the effectiveness of this method especially using $g_1$ alone, which leads to the most significant improvement in model performance.\n\nAdditionally, we do not rule out the possibility of more powerful and efficient augmentation methods. The parallel MLP approach is chosen primarily to make better use of the pre-trained weights $W_V$/ $W_K$ and this design is quite general and does not consider specific characteristics of particular tasks. We still encourage the design of more task-specific augmentation format for different tasks. For example, in CV tasks, it might be natural to incorporate CNN-extracted features into the $g_1$/$g_2$ in ViTs.\n\nSimilarly, our discussion on the negative sample models further supports the need for more task-specific designs. As shown in Table 1, selecting tokens with low attention scores as negative samples is a rather crude approach, leading to relatively limited performance improvements. Exploring how to select or construct higher-quality negative samples is also an interesting direction. These detailed considerations go beyond the scope of this paper, and these interesting directions will be the focus of our future work.\n\n**Final Note**: We are excited that you acknowledge our theoretical contribution. If there are any remaining questions, please do not hesitate to let us know. Thank you once again for your insightful comments and for your encouraging feedback!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9w8bOPa5pB", "reviewer": "Reviewer_k2vC", "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.", "questions": "- How would the inclusion of layer normalization and residual connections affect the theoretical framework and findings?\n- How do the proposed methods perform on more complex and diverse tasks beyond the linear regression tasks used in the experiments?", "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.", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(Reminder) I'm not a researcher working in the theory field, so I couldn't recognize this paper's theoretical contributions well. \n\n- The paper provides a fresh perspective on understanding ICL by connecting it to gradient descent and representation learning, which is a novel and insightful contribution to the field.\n- The authors extend their theoretical findings to various Transformer settings, including single and multiple attention layers, enhancing their conclusions' generalizability.", "weaknesses": "(Reminder) I'm not a researcher in the theory field, so I couldn't also recognize this paper's limitation from a perspective. \n\n- 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.\n- Lack of discussions towards some related works that argue against the equivalence between ICL and gradient descent [1,2].\n\n[1] In-context Learning and Gradient Descent Revisited, 2023\n[2] Do pretrained Transformers Learn In-Context by Gradient Descent, 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the novel perspective on understanding ICL of our paper. Considering the word limits of rebuttal, we notice that some issues might overlap, so we have organized our responses as follows.\n\n> **Weakness 1** (The paper relies on several assumptions and simplifications...)  \n> **& Question 1** (How would the inclusion of layer normalization and residual connections affect the theoretical framework and findings?)\n\nAs for residual connection, intuitively, our conclusions can be naturally extended to versions with residual connections, which are analogous to introducing a residual connection structure in dual models, i.e., $\\mathrm{Atten}(x) + x = \\hat{W}\\phi(W_Qx) + x$, where $\\hat{W}$ is the weight in the dual model $f(x) = W\\phi(x)$ trained under the original representation loss in Eq (9). \nThis modification does not have a significant impact on the specific representation loss or generalization bound, but its effect on the GD process of the dual model still requires more detailed analysis.\n\nAs for layer normalization, from the representation learning perspective of the dual model, it normalizes $x$ to a scaled unit sphere to make the process more stable intuitively. However, we acknowledge the analytical challenges posed by the nonlinearity of layer normalization, and the more detailed effects still warrant further investigation. This may require more powerful theoretical tools, e.g. [1],  which is also a direction for our future research.\n\n[1] On the Nonlinearity of Layer Normalization.\n\n> **Weakness 2** (Lack of discussions towards some related works that argue against the equivalence between ICL and gradient descent...)\n\nFollowing the previous work [1], the first mentioned work [2] further explores potential shortcomings in the previously used evaluation metrics and further points out that the ICL output is more dependent on previous lower layers, which is different from the finetuning that relies on all model parameters. Inspired by this, it introduces a new finetuning method called LCGD and validates the claim with new metrics.\n\nUnlike previous works including [3,4], which studies ICL and Gradient Descent (GD) under strong constraints (specific regression tasks and weight constructions, etc.), the second mentioned work [5] highlights significant gaps between these strong constraints and real-world LLMs. They observe that there are differences in how ICL and GD modify the output distribution and emphasize the distinction between naturally emergent ICL (pretrained on natural text data) and task-specific ICL (such as in [3,4]). It provides a clearer definition and experimental results in more realistic settings.\n\nAlthough our work also relates ICL to GD, our focus is on how we can view the ICL inference process as a GD process on a dual model from a representation learning perspective: \n- To simplify theoretical analysis, we concentrate on analyzing the core component of the Transformer architecture—the attention layer—while neglecting residual connections and layer normalization. This does introduce some divergence from real-world large language models; \n- In terms of research approach, we focus on the GD process of dual model under representation learning loss, rather than analyzing GD directly on the original model as in [2,5]. This is also a significant departure from the mentioned works.\n- Additionally, the representation learning process of the dual model resembles a self-supervised process, which we believe is reasonable: unlike fine-tuning, ICL inference is not explicitly provided specific metrics to adapt to the target task and this also aligns with certain aspects of previous work [5]. And we also compare this process with existing self-supervised learning methods in our work.\n- Inspired by this representation learning perspective, we propose modifications to the attention mechanism, which were not sufficiently explored in previous work. \n\nThank you for your suggestion and we will add more supplementary information on the relevance between our work and the mentioned works in our future revision.\n\n[1] Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers\n\n[2] In-context Learning and Gradient Descent Revisited\n\n[3] Transformers Learn In-context by Gradient Descent\n\n[4] Trained Transformers Learn Linear Models In-context\n\n[5] Do pretrained Transformers Learn In-Context by Gradient Descent?\n\n\n> **Question 2:** How do the proposed methods perform on more complex and diverse tasks beyond the linear regression tasks used in the experiments?\n\nFor more complex tasks, considering the limited computational resources and time constraints, we select the pre-trained BERT-base-uncased model to apply attention modifications, and validate the results on part of GLUE datasets.  **More experimental settings are detailed in the (global) author rebuttal and the results are presented in the Table 1 of PDF** `(To AC and All Reviewers)`. \n\n**For the regularized models**, we select different values $\\alpha$ and conclude that smaller absolute values of $\\alpha$ are recommended. **For the augmented models**, we adopt a parallel MLP approach for data augmentation $g_1/g_2$, that is, $g_1(W_Vx) =  W_Vx  +  c W_2\\sigma(W_1x)$​. The results suggest that using $g_1$ alone may be a more effective choice under our settings. **For the negative sample models**, we continued to select tokens with lower attention scores as negative samples. The final results showed limited improvement in model performance. We think that this may be due to the ineffective selection of negative sample tokens, and better methods for selecting negative samples would be interesting to explore in the future.\nIn conclusion, these results validate the potential of our proposed methods.\n\n**Final Note:** We want to thank you again for all the questions you have provided. If there are any remaining questions, please do not hesitate to let us know.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OpDyBn32Vo", "reviewer": "Reviewer_8J5b", "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.\n\n* Extend theoretical analysis to more complex scenarios with multiple attention layers\n* Propose potential modifications to the attention layer to enhance ICL capabilities\n\nExperiments 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.", "questions": "Theoretical Section\n* Will the ratio of the number of negative samples have an impact on generalization bound?\n\nExperimental Section\n\n* As mentioned in weeknesses part, the authors may need to add more experimental details to make it more self-explained.\n* 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? \n* Does the conclusion hold the same for both one-layer transformer and multi-layer transoformer?\n* The dynamics of token interactions may not be well supported in the experiments.\n\nConclusion Section?\n\n* seems missing a conclusion section due to the restriction of page sizes.", "limitations": "The authors have clearly stated the limitations of this work, the selected transformer components of this work, and selected task settings.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* While previous studies mainly on the linear attention, this work extends it to a more realistic case -> softmax attention, including its training dynamics.\n* Based on the dual gradient descent process, the authors also propose potential attention modifications, which could be benificial to real-world downstream application.", "weaknesses": "* The work may ignore the format of task input/output. And some of recent work such as (Theoretical Understanding of In-Context Learning in Shallow Transformers with Unstructured Data, Xing et al.) the success of ICL may be also because of the alignment of attention mechanism, especially in real-world nosiy cases.\n* The experimental section is not well-organized and self-explained, the readers may need to refer extra papers to get the protocols.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your recognition of the novelty of our paper. We have carefully addressed each of your concerns and considering the word limits of rebuttal, we notice that some issues might overlap, so we have organized our responses as follows.\n\n> **Weakness 1** (The work may ignore the format of task input/output. And some of recent work...)  \n> **& Question 5** (The dynamics of token interactions may not be well supported in the experiments.)\n\nThe mentioned work [1] extends the original structured input format [2,3] to the unstructured setting and then explores the impact of various factors (one/two layers, look-ahead attention mask, PE) when learning linear regression tasks both theoretically and experimentally. The explanation for ICL derived from these interesting findings is more aligned with real-world scenarios (where the examples are usually stored in different tokens).\n\nAlthough our work also relates ICL to GD, our focus is on how we can view the ICL inference process as a GD process on a dual model from a representation learning perspective. We do not make assumptions about the input format, such as $[x; y]$ or $[x; 0], [0; y]$, as we do not target specific tasks (e.g. linear regression tasks) but aim to provide more general explanations. To simplify the analysis, our input tokens can be viewed as embeddings of demonstration sentences (or intermediate states among multi layers). \n\nCorrespondingly, in the experimental section, our focus is not on the dynamics of token interactions as specific tasks are not being targeted. Instead, we are more concerned with whether the result of ICL inference $h'\\_{T+1}$ is equivalent to the test prediction $\\hat{y}_{test}$​ of the trained dual model, as illustrated in the left part of Figure 3. This result aligns with our analysis in Theorem 3.1.\n\n[1] Theoretical Understanding of In-Context Learning in Shallow Transformers with Unstructured Data\n\n[2] Transformers learn in-context by gradient descent\n\n[3] Trained transformers learn linear models in-context\n\n> **Weakness 2** (The experimental section is not well-organized and self-explained...)  \n> **& Question 2** (As mentioned in weeknesses part...)\n\nThank you for your careful suggestions! Due to the resrtiction of page sizes, we simplified the description of the experiments. In fact, more details of our experiments are presented in the Appendix E. We will reorganize our experimental part and provide more details in our future revision.\n\n> **Question 1** (Will the ratio of the number of negative samples have an impact on generalization bound?)\n\nThe answer is yes. After introducing negative samples, we consider the following representation loss:\n$$\n\\mathcal{L}(f) = \\mathbb{E}_{x\\sim\\mathcal{D}\\_{\\mathcal{T}}} \\left[- \\frac{1}{K} \\sum\\_{j = 1}^{K}  \\left( W\\phi(W_Kx) \\right)^T\\left(W_Vx - W_Vx^{-}\\_{j} \\right)\\right],\n$$\nwhere $K$ is the number of negative samples for each $x_i$ and  $x\\_{j}^-$ denotes the $j$-th negative sample for token $x$. The empirical loss is modified correspondingly. Then, by the other definitions in Section 3.3, we can obtain the generalization bound as\n$$\n\\mathcal{L}(\\hat{f})\\le \\mathcal{L}(f) + O\\left(w\\rho d\\_o \\sqrt{\\mathrm{Tr}(K\\_S)\\left(\\frac{5}{N^2} + \\frac{1}{rN^3}\\right)} + \\sqrt{\\frac{log\\frac{1}{\\delta}}{N}}\\right),\n$$\nwhere $r = \\frac{K}{N}$ is the ratio of the number of negative samples. It can be observed that as the ratio increases, the generalization error will decrease. However, we also notice that $\\frac{5}{N^2} > \\frac{1}{rN^3}$​​ thus the reduction in generalization error due to an increased proportion of negative samples is limited. The proof is similar to that of Theorem 3.2 (see Appendix C). The main difference is that we need to define the new function class $G$ and $F$ according to above definition. Nevertheless, we do not rule out the possibility of a tighter generalization bound using better theoretical tools.\n\n> **Question 3** (Since the authors ... which one could be the best ... need add these three together? why...?)\n\nThe answer is related to the specific tasks but generally speaking, using all three methods together is not necessary, especially when the design of  augmentated or negative modification is not effective enough. We conduct experiments with their combinations  on linear, trigonometric, and exponential tasks. **The results are shown in Figure 1 of the  PDF** in  `To AC and All Reviewers`. For the latter two tasks, the results are actually worse than using regularized or augmented modification individually (compared to Figures 11 and 12 in Appendix E). Therefore, when the design of augmented or negative mofification is not effective, we recommend using a single modification method.\n\n> **Question 4** (Does the conclusion hold the same for both one-layer transformer and multi-layer transformer?)\n\nFor more scenarios, the conclusion still holds: using all three methods simultaneously is not necessary. We supplemented our experiments using BERT (multi-layer Transformer), see ` (To AC and All Reviewers)`. From Table 1 in the PDF, the combined models do not outperform the augmented version alone. We also observed that negative sample improvements remain quite limited, indicating that our method for selecting negative samples is not effective enough. In these cases, using augmented modification alone will be a better choice.\n\n> **Question 6** (seems missing a conclusion section) \n\nWe will streamline our paper content and supplement the conclusion section based on the valuable questions raised by all you reviewers in the future revision. Thank you very much for your reminder!\n\n**Final Note:** We want to thank you again for all the questions you have provided. If there are any remaining questions, please do not hesitate to let us know.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
### **To AC and All Reviewers**

We thank the reviewers for providing valuable suggestions that help us improve our paper.

We are particularly encouraged that the reviewers have found that (i) the fresh perspective on understanding attention mechanisms `(HnPE, cRiK, k2vC,8J5b) ` ,  (ii) thorough theoretical analysis `(HnPE, cRiK, 8J5b)`,  (iii) the potential of proposed attention modifications `(8J5b, HnPE, cRiK, k2vC)`, and (iv) good writing `(HnPE)` of our work.

In response to the feedback, we've done our best to address each concern and have added new experiments and theoretical results. 

We notice that the reviewers share a common concern regarding the generalizability of our findings to broader NLP tasks. Below, we provide a detailed response to the additional experiments we've conducted:

- **Basic Experiment Setting**: We supplement our experiments on more realistic NLP tasks. Considering the limited computational resources and time constraints, we choose the BERT-base-uncased model (hereafter referred to as BERT) to validate the effectiveness of modifications to the attention mechanism. As for datasets, we select part of GLUE (CoLA, MRPC, STS-B, RTE). We load the checkpoint of the pre-trained BERT model (where the classifier are newly initialized) and then fine-tune the model to explore the performance of three modifications as well as their combinations. 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 of PDF**. Below, we discuss the setting of various modifications and their performance.

- **For the regularized modification**, we consider different values of $\alpha$. 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 $|\alpha|$  is too large, the model's performance declines significantly, so we recommend using smaller $|\alpha|$.

- **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(W_Vx)=W_2σ(W_1W_Vx)$ 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(W_Vx) =  W_Vx  +  c W_2\sigma(W_1x)$, where $c$ is a hyperparameter to control the influence of the new branch, $\sigma$ is GELU and the hidden dimension is set to twice the original size of $W_Vx$. And $g_2(W_Kx) =  W_Kx  +  c W_2\sigma(W_1x)$ follows the same format.

  Experimental results show that the best augmented models achieve better performance than the original model across all four datasets. Notably, 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 significant for larger models. Thus, under the augmentation methods and experimental settings we selected, using $u_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$ is primarily motivated by the desire to make better use of the pre-trained weights $W_K$ and $W_V$. 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 $\beta$ indicates the overall reduction in attention scores. The best negative models only outperform the original model on CoLA and STS-B, whereas their performance on MRPC and RTE is worse than the original one. This suggests that our simple approach of considering tokens with low attention scores as negative samples might be too coarse. A more effective method for constructing negative samples should be designed, which is a direction worth exploring in the future.

- We also consider **combining different modification methods**. The results indicate that under our settings, the combination of augmented and negative modification achieves the best performance on CoLA, MRPC, and RTE, while the combination of regularized and augmented modification achieves the best performance on STS-B. However, their optimal performance is slightly inferior to the best performance achieved with augmented models alone. Therefore, we conclude that using all three modifications simultaneously is not necessary. With appropriate settings, 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 even with rough parameter selections. This further validates the potential of our approach of thinking about and improving the attention mechanism from a representation learning perspective. However, due to time and computational resource limitations, our experiments are conducted on a limited set of datasets and model size. More detailed parameter searches, validation across additional tasks and models, and the development of task-specific augmentation and negative sampling methods are all interesting directions worth exploring in the future.

Meta-review: 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.
Decision: Accept (poster)