Title: Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformers

In-context Learning

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

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.

1 Introduction

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

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

kernel regression algorithms and exhibit similar behaviors during ICL. These works have provided significant insights into the interpretation of ICL capabilities from various perspectives.

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

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

2 Preliminaries

2.1 In-context Learning with Transformers

The model we consider is composed of many stacked Transformer decoder layers, each of which is composed of an attention layer and a FFN layer. For simplicity, we have omitted structures such as residual connections and layer normalization, retaining only the most essential parts. We consider the standard ICL scenario, where the model's input consists of demonstrations followed by query inputs, that is, the input can be represented as X = [ X D , X T ] ∈ R d i × ( N + T ) , where X D = [ x 1 , x 2 , ..., x N ] denotes N demonstration tokens, and X T = [ x ′ 1 , x ′ 2 , ..., x ′ T ] denotes T query tokens. Here, we focus more on how tokens interact during model inference while ignoring the internal structure of demonstration tokens. For the query input at position T +1 , its output after one layer of Transformer can be represented as

where W K , W Q , W V ∈ R d o × d i are parameters for key, query, value projections and W 1 ∈ R d h × d o , W 2 ∈ R d o × d h , b 1 ∈ R d h , b 2 × R d o are FFN parameters. Our concern is how the query token x ′ T +1 learns in-context information from demonstrations. Unlike previous work [Von Oswald et al., 2023a, Zhang et al., 2023a, Bai et al., 2023], here we do not make additional assumptions about the structure of input matrix X and parameters to study the Transformer's ability to implement some specific algorithms. Instead, we adopt the same setting as [Dai et al., 2022] to study more general cases.

2.2 Self-Supervised Representation Learning Using Contrastive Loss Functions

Representation learning aims to learn embeddings of data to preserve useful information for downstream tasks. One class of methods most relevant to our work is probably contrastive learning methods without negative samples [Chen and He, 2021, Grill et al., 2020, Caron et al., 2020, Tian et al., 2021]. Contrastive learning is a significant approach of self-supervised learning (SSL) which aims at learning representations by minimizing the distance between the augmentations of the same data point (positive samples) while maximizing the distance from different data points (negative samples) [He et al., 2020, Chen et al., 2020b, Oord et al., 2018, Oh Song et al., 2016]. To alleviate the burden of constructing a sufficient number of negative samples while avoiding representational collapse, some works propose architectures for contrastive learning without negative samples, which mainly use weight-sharing network known as Siamese networks [Chen and He, 2021, Grill et al., 2020, Caron et al., 2020, Tian et al., 2021]. The architecture takes two augmentations x 1 , x 2 from the same data x as inputs, which will be processed by online network and target network respectively to obtain the corresponding representations, that is, ˆ x 1 = f online ( x 1 ) , ˆ x 2 = f target ( x 2 ) . The two encoder networks share weights directly or using Exponential Moving Average (EMA). Then, ˆ x 1 will be input into a predictor head to obtain the predictive representation z 1 = g (ˆ x 1 ) . Finally, we minimize the distance between the predictive representation and target representation, that is, L ( z 1 , StopGrad(ˆ x 2 )) where StopGrad( · ) means ˆ x 2 is treated as a constant during backpropagation. For L ( · ) , we often choose the cosine similarity or the l 2 -norm as a measure of distance, although they are equivalent when the vector is normalized. Another class similar to our work is kernel contrastive learning [Esser et al., 2024]. Given an anchor x and its positive and negative samples x + , x -, it aims to optimize the loss function L = f ( x ) T ( f ( x -) -f ( x + )) , where f ( x ) = W ϕ ( x ) and ϕ ( x ) is the feature mapping for some kernel. We will consider the gradient descent process corresponding to the inference process of ICL from the perspective of representation learning and compare it with the two aforementioned representation learning patterns.

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., 2022, Dai et al., 2022], that is, the gradient descent on linear layer can be seen as the dual form 2 of linear attention. A simple linear layer can be defined as f L ( x ) = Wx , where W ∈ R d o × d i is the projection matrix. Given training inputs [ x i ] N i =1 ∈ R d i with their labels [ y i ] N i =1 ∈ R d o , a linear layer can output the predictions [ ˆ y i ] N i =1 where ˆ y i = Wx i and then compute certain loss L (ˆ y i , y i ) for training. Backpropagation signals [ e i ] N i =1 ∈ R d o will be produced to update W in gradient descent process where e i = -η ( ∇ ˆ y i L ) if we set η as the learning rate. During test time, the trained weight matrix ̂ W can be represented by its initialization W 0 and the updated part ∆ W , that is,

where ⊗ denotes the outer product according to the chain rule of differentiation. On the other hand, this process can be viewed from the perspective of linear attention. Let [ k i ] N i =1 , [ v i ] N i =1 ∈ R d i denote the N key and value vectors constituting matrices K , V ∈ R d i × N respectively. For a given query input q ∈ R d i , linear attention is typically defined as the weighted sum of these value vectors

Then, we can rewrite the output of a linear layer during test time as

features

h

h

hs

h.

hs

hv-3

hx-2

hy-1

hv

hi

h

hs

h

hr+1

ytes

equivalent

Dual Model

N)

WK

Softmax attention layer with parameters {Wo,Wk,W}

TrainedDualModel

C(g(i), ya)

f(a)=W（α)

grad

Z

Wv

Wa

23

2.

ZN-3

2N-2

2N-1

EN

T+

%tcst

Examples:Dogis

sanimal.".Appleis

fruit

.Question:Whatiscarrot

？

test input

labels

where E ∈ R d o × N and X ∈ R d i × N are stacked by backpropagation signals [ e i ] N i =1 and training inputs [ x i ] N i =1 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 are chosen to activate using LA( · ) where [ e i ] N i =1 can be considered as values while [ x i ] N i =1 as keys and x test as the query. This interpretation uses gradient descent as a bridge to connect predictions of linear layers with linear attention, which can be seen as a simplified softmax attention used in Transformers.

Inspired by this relationship, Dai et al. [2022] understand ICL as implicit fine-tuning. However, this interpretation based on linear attention deviates from the softmax attention used in practical Transformers. Furthermore, this alignment is also ambiguous as the specific details of the gradient descent process, including the form of loss function and dataset, have not been explicitly addressed. In addition, Von Oswald et al. [2023a], Ding et al. [2023] also connect ICL with gradient descent for linear regression tasks using weight construction methods, where parameters W K , W Q and W V of the self-attention layer need to roughly adhere to a specific constructed form. However, these analyses rely on the setting of linear regression tasks and assumptions about the form of input tokens (concatenated with features and labels), which limits the interpretability of ICL capabilities from the perspective of gradient descent. Thus, we attempt to address these issues in the following sections.

3 Connecting ICL with Gradient Descent

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

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 et al. [2022] connect ICL with gradient descent under the linear attention setting. In fact, it is completely feasible to interpret ICL under softmax attention with the help of kernel methods. We define the attention block as

which can be viewed as the product of an unnormalized part A u and a normalizing multiplier D , that is,

where exp( · ) is element-wise. Similar in [Choromanski et al., 2020], we define softmax kernel K sm : R d o × R d o → R + as K sm ( x , y ) = e x T y = e ∥ x ∥ 2 + ∥ y ∥ 2 2 K guass ( x , y ) where K guass = e -∥ x -y ∥ 2 / 2 is the guassian kernel when the variance σ 2 = 1 . According to Mercer's theorem [Mercer, 1909], there exists some mapping function ϕ : R d o → R d r satisfying that K sm ( x , y ) = ϕ ( x ) T ϕ ( y ) . Thus,

noting that when omitting the √ d o -renormalization and equivalently normalize key and value vectors in Eq (6), every entry in the unnormalized part A u can be seen as the output of softmax kernel K sm defined for the mapping ϕ , which can be formulated as:

There have been many forms of mapping function ϕ ( · ) used in linear Transformers research to approximate this non-negative kernel [Choromanski et al., 2020, Katharopoulos et al., 2020, Peng et al., 2021, Lu et al., 2021]. For example, we can choose ϕ ( · ) as positive random features which has the form ϕ ( x ) = e w T x -∥ x ∥ 2 / 2 to achieve unbiased approximation [Choromanski et al., 2020]. Alternatively, we can also choose ϕ ( x ) = elu( x ) + 1 proposed by Katharopoulos et al. [2020].

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 a softmax attention layer in a trained Transformer model, where the parameters { W Q , W K , W V } have been determined and the input X = [ X D , X T ] has the form introduced in Section 2.1. Then, after the inference by one attention layer, the query token at position T +1 will have the form h ′ T +1 formulated by Eq (1).

On the other hand, given a specific softmax kernel mapping function ϕ ( x ) that satisfies Eq (7), we can define the dual model for the softmax attention layer as

where W ∈ R d o × d r is parameters. We assume that the dual model obtains its updated weights ̂ W after undergoing one step of gradient descent with some loss function L . Subsequently, when we take z test = W Q x ′ T +1 as the test input, we can obtain its test prediction as

We will show that h ′ T +1 in Eq (1), is strictly equivalent to the above test prediction ˆ y test , which implies that the inference process of ICL involves a gradient descent step on the dual model. This can be illustrated by the following theorem:

Theorem 3.1. The query token h ′ T +1 obtained through ICL inference process with one softmax attention layer, is equivalent to the test prediction ˆ y test obtained by performing one step of gradient descent on the dual model f ( x ) = W ϕ ( x ) . The form of the loss function L is:

where η is the learning rate and D is a constant.

Proof can be found in Appendix A. Theorem 3.1 demonstrates the equivalence between the ICL inference process and gradient descent. Below, we delve into more detailed discussions:

Training Set and Test Input: In fact, once the attention layer has already been trained, that is, W K , W Q , W V has been determined, the demonstration tokens [ x i ] N i =1 will be used to construct a training set for the dual model. Specifically, the training data has the form { z ( i ) std , y ( i ) std } N i =1 where z ( i ) std = W K x i as inputs and y ( i ) std = W V x i as their labels. During training stage, for each input z ( i ) std , the dual model outputs its prediction ˆ y ( i ) = f ( z ( i ) std ) = W ϕ ( z ( i ) std ) = W ϕ ( W K x i ) . Then, the loss function Eq (9) can be rewritten as L = -1 ηD ∑ N i =1 ( y ( i ) std ) T ˆ y ( i ) , which can be regarded as the cosine similarity. Then, using this loss function and the training data, we can perform one step of Stochastic Gradient Descent (SGD) on the dual model and obtain the updated ̂ W . Finally, during the testing stage, we take z test = W Q x ′ T +1 as the test input to get its prediction which will be consistent with the ICL result h ′ T +1 , that is, ˆ y test = f ( z test ) = ̂ W ϕ ( W Q x ′ T +1 ) = h ′ T +1 . This process can be illustrated in Figure 1. Demonstration tokens provide information about the training data points and the weight matrix ̂ W is optimized to learn sufficient knowledge about demonstrations. This

Higher Dimensional Space

y

grad

C(y,g()-

C(p，y)-

igrad

C(p,p)

grad i

M

Predictor

M

W

M

Φ(W∞:)

(2()

(2)

Online Network

Target network

Data Augmentations

Data Augmentations

WKZ.

WZ

W2;

W2;

z()=Aug(z.)

=Aug(a)

z{) = Aug(z)

C;

gradient descent process using the loss function L applied to f ( x ) can be seen as the dual form of the ICL inference process of the attention layer.

Representation Learning Lens: Even though we have now clarified the details of the gradient descent process of ICL, what does this process more profoundly reveal to us? In fact, for a encoded demonstration token x i , the key and value mapping will generate a pair of features W K x i and W V x i that exhibit a certain distance from each other, akin to positive samples in contrastive learning. And then, ϕ ( x ) projects W K x i into a higher-dimensional space to capture deeper features. Finally, the weight matrix W , which maps ϕ ( W K x i ) back to the original space, is trained to make the mapped vector as close as possible to W V x i . This process is illustrated in Figure 2. Below, we attempt to understand this process from the perspective of existing representation learning methods introduced in Section 2.2, although we emphasize that there are certain differences between them.

Comparison with Contrastive Learning without Negative Samples: If we consider the key and value mapping as two types of data augmentation, then from the perspective of contrastive learning without negative samples, this process can be similarly formalized as

where StopGrad( · ) is naturally applicable because there are no learning parameters involved in the generation process of the representation y ( i ) std . However, it's important to note that the representation learning process of ICL is much simpler: Firstly, the online and target networks are absent while the augmentations W K x i , W V x i are directly used as online and target representations respectively. Secondly, the predictor head is useful and not discarded, which is then used during test stage.

Comparison with Contrastive Kernel Learning: Given an anchor data x and its positive and negative samples x + , x -, contrastive kernel learning aims to optimize the loss function L = f ( x )( f ( x -) -f ( x + )) where f ( x ) = W ϕ ( x ) . There are significant differences in the representation learning process of ICL: Firstly, it does not involve negative samples. Secondly, there is no corresponding processing for positive samples, leading to parameter updates being solely dependent on the processing of the anchor.

Extension to More Complicated Scenarios: Theorem 3.1 can be naturally extended to one single Transformer layer and multiple attention layers. As for one Transformer layer formed in Section 2.1, its dual model f + ( x ) = W ϕ ( x ) + b introduces an additional bias b and only W is trained while b remains fixed. In addition, the labels of training set will be y ( i ) std = W F W K x i where W F has potential low-rankness property induced by Relu( · ) . As for multiple attention layers, the ICL inference process will be equivalent to sequentially performing gradient descent and making predictions on the dual model sequence. We provide more details in Appendix B.

Compared to Dai et al. [2022] considering the connection under linear attention setting, Theorem 3.1 gives explanation for more generally used softmax attention and offers a more detailed exploration of the training process. Additionally, unlike Von Oswald et al. [2023a,b], Ding et al. [2023]'s focus on particular linear regression task and specific configurations of token and parameters, we aim to explain the process of token interactions during ICL inference in a more general setting.

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 T , we cannot provide all demonstrations related to task T limited by the length of input tokens. We denote S T ⊆ R d i as all possible tokens for the task T and assume that these tokens will be selected according to the distribution D T . During a particular instance of ICL inference, let S = { x i } N i =1 ⊆ S T represent the example tokens we selected. We define the function class as F := { f ( x ) = W ϕ ( W K x ) | ∥ W ∥ ≤ w } where ∥ · ∥ denotes the Frobenius norm. Generally, ignoring constant term in Eq (9), we consider the representation learning loss as

where f ∈ F and D T is the distribution for some ICL task T . Correspondingly, the empirical loss will be formulated as ˆ L ( f ) = -1 N ∑ N i =1 ( W V x i ) T f ( x i ) and we have ˆ f = arg min f ∈F ˆ L ( f ) . In addition, we denote the kernel matrix of demonstration tokens S as K S ∈ R N × N where ( K S ) i,j = ⟨ ϕ ( W K x i ) , ϕ ( W K x j ) ⟩ , that is, the inner product of the feature maps after W K projection between the i -th token and j -th token. We state our theorem as follows:

Theorem 3.2. Define the function class as F := { f ( x ) = W ϕ ( W K x ) | ∥ W ∥ ≤ w } and let the loss function defined as Eq (10). Consider the given demonstration set as S = { x i } N i =1 where S ⊆ S T and S T is all possible demonstration tokens for some task T . With the assumption that ∥ W V x i ∥ , ∥ W ϕ ( W K x i ) ∥ ≤ ρ , then for any δ > 0 , the following statement holds with probability at least 1 -δ for any f ∈ F

Proof of 3.2 can be found in Appendix C. Theorem 3.2 provides the generalization bound of the optimal dual model trained on a finite selected demonstration set under a mild assumption that ∥ W ∥ is bounded. Intuitively, as the number of demonstration (and therefore the number of demonstration tokens) increases, the generalization error decreases, which is consistent with existing experimental observations [Xie et al., 2021, Garg et al., 2022, Wang et al., 2024].

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 a dual attention inference process? Alternatively, can we modify the attention mechanism by drawing on existing methods? In fact, since there are lots of mature works in representation learning especially contrastive learning, it is possible for us to achieve this by drawing on these works [He et al., 2020, Chen et al., 2020c, Wu et al., 2018, Chen et al., 2020a, Chen and He, 2021]. We will provide some simple perspectives from the loss function, data augmentations and negative samples to try to adjust attention mechanism. It is worth noting that these modifications are also applicable to the self-attention mechanism, and we will explore these variants in experiments. More details can be seen in Appendix D.

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

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

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

test-h+ll2ateachstep

Regularized Models

Augmented Models

NegativeModels

0.8

2.00

3.5

2.00

Z

0.7

0.6

0.4

0.3

0.5

5

S

S

1.50

1.25

1.00

0.75

1.75

α=0

α=1

α=0.5

α=-0.5

α=-1

2.0

5

S

1.5

2.5

3.0

g

normal

g1&g2

g1

g2

S

5

S

1.75

1.50

1.25

1.00

0.75

normal

k=1,β=0.5

k=1,β=0.1

k=3,β=0.1

k=3,β=0.5

0.2

0.50

1.0

0.50

0.1

0.25

0.5

0.25

0.0

2.5

5.0

7.5

10.012.5

15.0

0.00

0

20

40

60

80

10(

0.0

0

20

40

60

80

100

0.00

0

20

40

60

80

100

steps

epochs

epochs

epochs

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

Correspondingly, the attention layer can be adjusted as,

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

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

where N ( i ) is the set of the negative samples for x i and β is a hyperparameter. Correspondingly, the attention layer is modified as

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

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) Is it potential to improve the attention mechanism from the perspective of representation learning?

Linear Task Setting: Inspired by Von Oswald et al. [2023a], to validate the equivalence and demonstrate the effectiveness of the modifications, we firstly train one softmax self-attention layer using linear regression tasks. We generate the task by s = Wt where every element of W ∈ R d s × d t is sampled from a normal distribution W ij ∼ N (0 , 1) and t from uniform distribution t ∼

U ( -1 , 1) d t . We set d t = 11 and d s = 1 . Then, at each step, we use generated { x i = [ t i ; s i ] } N +1 i =1 to form the input matrix X where the last token will be used as the query token and the label part will be masked, that is, x N +1 = [ t i ; 0] . Here we consider only one query token ( T = 0 ) and we denote x ′ T +1 = x N +1 to maintain consistency of notation in Section 2.1. Finally, the attention layer is trained to predict ˆ s N +1 to approximate the true label s N +1 using mean square error (MSE) loss.

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

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

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

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 α > 0 , the regularized models converges to a poorer result while when α < 0 , the model converges faster and achieves final results comparable to the normal model without regularization ( α = 0 ). At least for this setting, this is a little contrary to our initial intention of applying regularization to the contrastive loss where α should be positive. We explain it that the appropriate α contributes to achieving a full-rank attention matrix as stated in Appendix D, preserving information and accelerating convergence.

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

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

6 Related Work

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

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

7 Conclusion and Impact Statements

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

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 , 36:45614-45650, 2023.

MAAiserman, 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 algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 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.04637 , 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 are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

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-9924, 2020.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020a.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020b.

Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on 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.04297 , 2020c.

Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06528 , 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.14794 , 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 incontext? language models secretly perform gradient descent as meta optimizers. arXiv preprint 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 1: 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 preprint arXiv:2301.00234 , 2022.

Pascal Esser, Maximilian Fleissner, and Debarghya Ghoshdastidar. Non-parametric representation learning with kernels. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 38, pages 11910-11918, 2024.

Shivam Garg, Dimitris Tsipras, Percy S 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-30598, 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-supervised learning. Advances in neural information processing systems , 33:21271-21284, 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.12766 , 2023.

Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg-Kirkpatrick, and Graham Neubig. Towards a unified view of parameter-efficient transfer learning. arXiv preprint arXiv: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 arXiv:1606.08415 , 2016.

Kazuki Irie, 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 rnns: Fast autoregressive transformers with linear attention. pages 5156-5165, 2020.

Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of naacL-HLT , 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 arXiv:2301.07067 , 2023.

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-free 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-contraction 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.

Mercer. 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 or Physical Character , 209:415-446, 1909. ISSN 02643952. URL http://www.jstor.org/stable/91043 .

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 van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv: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. A theoretical analysis of contrastive unsupervised representation learning. In International Conference 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, Łukasz Kaiser, and Illia 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.

Johannes Von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05858 , 2023b.

Alex Wang. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 , 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.14160 , 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, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682 , 2022.

T Wolf. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 , 2019.

Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via nonparametric 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.02080 , 2021.

Yue Xing, Xiaofeng Lin, Namjoon Suh, Qifan Song, and Guang Cheng. Benefits of transformer: Incontext 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.09927 , 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.

A Details of Theorem 3.1

We repeat Theorem 3.1 as follows and provide proof and more discussion for it.

Theorem A.1. The query token h ′ T +1 obtained through ICL inference process with one softmax attention layer, is equivalent to the test prediction ˆ y test obtained by performing one step of gradient descent on the dual model f ( x ) = W ϕ ( x ) . The form of the loss function L is:

where η is the learning rate and D is a constant.

Proof. The derivative of L with respect to W should be:

Thus, after one step of gradient descent , the learned ̂ W will be

where W 0 is the initialization of the reference model and η is the learning rate. So the test prediction will be

On the other hand, from the perspective of ICL process with one attention layer, with Eq (7) in our mind, we can rewrite Eq (1) as

where we use [ V D , V T ] = W V [ X D , X T ] , [ K D , K T ] = W K [ X D , X T ] , q = W Q x ′ T +1 for simplify and D ′ = 1 T N ϕ ( K D ) T ϕ ( q ) + 1 T T ϕ ( K T ) T ϕ ( q ) is a constant to normalize the equivalent attention block. Further, we expand the above equation to connect the inference process of ICL using softmax attention with the gradient descent as follows

where W ′ 0 = 1 D ′ V T ϕ ( K T ) T and V ( i ) D , K ( i ) D are the i -th column vetors respectively.

Then, in Eq (18), when setting the initialization W 0 = W ′ 0 and the constant D = D ′ , we will find that

which means ˆ y test is strictly equivalent to h ′ T +1 . Thus, we have completed our proof.

Given a reference model f ( x ) = W ϕ ( x ) , by comparing Eq (19) and Eq (4), we can easily observe that the gradient descent on the loss function L applied to f ( x ) is the dual form of the inference process of ICL, where V ( i ) D , ϕ ( K ( i ) D ) and ϕ ( q ) play the roles of backpropagation signals, training inputs and test inputs respectively. Recalling the form of Eq (4), we can interpret the W 0 as the initialization of the weight matrix which provide the information under the zero-shot case while the second part in Eq (19) shows that the demonstration examples in ICL acts as the training samples in gradient descent. The reference model f ( x ) = W ϕ ( x ) , initialized with W 0 , will have test prediction ˆ y test = h ′ T +1 after training. This is also why we refer to it as the dual model of the softmax attention layer. We also note that for different demonstrations, even though the model has the same query input, the different given demonstrations will result in different output results. This is equivalent to the dual model performing gradient descent in different directions from the same initialization.

B Extensions to more complex scenarios

In Theorem 3.1, we provided the dual form of gradient descent for the ICL of one softmax attention layer. Here, we extend the conclusion to more complex scenarios, including one Transformer layer (attention layer plus one FFN layer) and multiple attention layers.

B.1 Extension to one Transformer Layer

As for one Transformer layer introduced in Section 2.1, we define the new dual model as

We will show that after performing gradient descent on W , the test output ˆ y test = f + ( W Q x ′ T +1 ) will be equivalent to ˆ x ′ T +1 . Our theorem is given as follows.

Theorem B.1. The output ˆ x ′ N +1 of ICL inference process with one Transformer layer, is strictly equivalent to the test prediction of its dual model f + ( x ) = W ϕ ( x ) + b , where f ( x ) is trained under the loss function L formed as

where η is the learning rate, D is a constant, and W F will be determined once the specified pre-trained model, demonstrations and query tokens are given.

Proof. Recalling the proof of Theorem 3.1, we can rewrite Eq (1) as

where D = 1 T N ϕ ( W K X D ) T ϕ ( W Q x ′ T +1 ) + 1 T T ϕ ( W K X T ) T ϕ ( W Q x ′ T +1 ) is a constant to normalize the attention scores and W 0 = 1 D ( W V X T ) ϕ ( W K X T ) T . Furthermore, h ′ N +1 will be taken as input for the FFN sublayer and the Eq (2) can be rewritten as

where I M ∈ R d × d is a diagonal matrix whose i -th diagonal element will be one if ( W 1 h ′ T +1 + b 1 ) i ≥ 0 otherwise be zero. We need to note that this process is reasonable: for given demonstration and query tokens, once the parameters { W Q , W K , W V , W 1 , b 1 } of the Transformer layer are fixed after training, I M will be determined implicitly (otherwise, I M would be a function that varies with these settings). For simplify, we rewrite ˆ x ′ T +1 as

where W F = W 2 I M W 1 and b F = W 2 I M b 1 + b 2 . Furthermore, expanding h ′ T +1 in the above Equation, we get:

On the other hand, we define a reference model:

where ϕ ( · ) is exactly the mapping function satisfying Eq (7) to approximate the softmax kernel. Given the loss formed in Eq (21), we can note that the right part in L is exactly the output of this reference model when taking W K x i as input, that is,

We can calculate the derivative of L with respect to W as

Suppose that the weight matrix W in the reference model f ( x ) is initialized as W init , then using one step of stochastic gradient descent (SGD) [Amari, 1993] with learning rate η , the weight matrix W will be updated as

Compared to Eq (23), we can set W init = W F W 0 , b = b F and take W Q x ′ T +1 as test input. Then, after one step update to W , the output of the reference model will be

which implies that if we initialize the reference model f + ( x ) = W ϕ ( x ) + b with W init = W F W 0 , b = b F , then after one step of gradient descent for W , the test output of f + ( W Q x N +1 ) will be identical to the ICL result of one Transformer layer. Thus, we call the reference model with setting W init = W F W 0 , b = b F as the dual model corresponding to the ICL inference process. Finally, we complete our proof.

Now, we discuss Theorem B.1 from the following perspectives:

Training set and test input: In fact, we can observe that the loss function L can be seen as the sum of inner products of N vector-pairs. In Eq (21), the right vector happens to be the predicted output ˆ y ( i ) std = f + ( z ( i ) std ) = W ϕ ( z ( i ) std ) + b of the dual model for training input z ( i ) std = W K x i . Correspondingly, the vector on the left can be regarded as the true label y ( i ) std = W F W V x i . In other words, it can be seen that the dual model performs one step SGD given training set { z ( i ) std , y ( i ) std } N i =1 on W using the loss L :

And then taking z test = W Q x i as test input, it finally output the prediction y test , which achieves the ICL result ˆ x N +1 . Compared to Theorem 3.1, after introducing the FFN layer, the main difference is that the labels of the training data become y ( i ) std = W F W V x i instead of y ( i ) std = W V x i . Additionally, compared to f ( x ) , an extra bias b is introduced in the new dual model f + ( x ) , which also have a different initialization W init = W F W 0 rather than W 0 . We also need to note that in the dual model f + ( x ) , only W is trained, while b remains unchanged after initialization.

y

>A

ysiad

b

W

Potential

bias

Low-rank

Φ(Wk∞i)

WF

WKCi

WvDi

Φ(Wkα:)

Higher Dimensional Space

1

1

1

M

WvDi

WKCi

potential low-rank

C

transformation

WeWvCi

Ci

Potential Low-rankness of W F : Noting that W F = W 2 I M W 1 where W 1 ∈ R d h × d , W 2 ∈ R d × d h , I M ∈ R d h × d h (here we assume that d i = d o = d for simplify), the rank of W F will satisfy

We observe that I M is a diagonal matrix with elements being zero or one, and its rank is determined by the number of non-zero elements. Here, we can make a mild assumption that we can set Rank( W 1 ) = Rank( W 2 ) = min { d, d h } . This assumption is quite mild as even for any random square matrix as it will be non-singular with probability 1. In addition, we also assume d h > d which is consistent with settings in practice. Therefore, we get Rank( W 1 ) = Rank( W 2 ) = d , and the upper bound of Rank( W F ) will be

Thus, we can find that if we want to avoid losing information, W F should strive to maintain Rank( I M ) > d which will be more easily achieved as d h becomes larger than d . Otherwise, Rank( I M ) is likely to gradually decrease with an increase in the number of Transformer layers. This explains the necessity of setting d h > d in practice. In some cases where Rank( I M ) < d , meaning that the number of non-zero elements in I M or positive elements in W 1 h N +1 + b 1 is less than d , the upper bound of Rank( W F ) will be Rank( I M ) and the lower bound of Rank( W F ) will be given as

which implies the rank of W F will exactly equal to Rank( I M ) . We should note that this condition, i.e., Rank( I M ) < d , is easily satisfied when d h = d or when d h is slightly larger than d (for example, d < d h < 2 d in an expected sense). Thus, we conclude that W F has the potential low-rank property.

Representation Learning Lens: For a encoded demonstration representation x i , the key and value projections will generate a pair of feature W K x i and W V x i to create a certain distance between data representations in space. And then, on the one hand, a potential lowrank transformation W F is applied to the W V x i , attempting to compress some information which increases the difficulty of contrastive learning and forces the model to learn better features; on the other hand, ϕ ( · ) projects W K x i into a higher-dimensional space to capture deeper-level features. Finally, we need to train the weight matrix W , which maps ϕ ( W K x i ) back to the original space, aiming to make the mapped vector as close as possible to W F W V x i . This interpretation is illustrated in Figure 4.

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, we

AttentionLayers

H(0)=X

Atten(H(o))

H(I)

Atten(H(l))

H(2)

H(-1)

Atten(H((-1))

H

H(-1)

Atten(H(l)

H=(）H

Sharing Weights

w9,w?,w!)

W2,w(2,w(2)

w9,w?,w!

W,w(,w()

Sequence of

H(0)=X

f(x)=W（c)

H

f(2）（c)=W(2（c)

H(2)

H(-1)

f(x)=W（c)

H(

H(L-)

f((a)=W(（a)

Dual Models

C=

1

（)g(

nD

data

training dataset

initialized l-thdualmodel

process

(2,ya}-1

finit(m)=Wnit(c)

W

update

W

H(-1)

test input

trained dual model

H(

W?

f0(m)=W(x)

assume we have only one query token, that is, there is no query input before the considered query token. With the assumption that T = 0 , to maintain notational simplicity, we use x N +1 to represent the query token here instead of x ′ T +1 and the input will be X = [ X D , x N +1 ] . We assume that there are L attention layers and the output of the l -th layer X ( l ) can be expressed as:

where we set H (0) = X = [ X D , x N +1 ] as the initial input. And the final output of the query token is ˆ h ( L ) N +1 = h ( L ) N +1 . Here, we assume that after training, the parameters W ( l ) Q , W ( l ) K , W ( l ) V ∈ R d o × d i are fixed and we set d o = d i = d .

Next, we extend Theorem 3.1 to the case of multiple softmax attention layers. Formally, we present our result in the following theorem.

Theorem B.2. Given L softmax attention layers whose parameters { W ( l ) Q , W ( l ) K , W ( l ) V } L l =1 are fixed after training, the ICL output of these layers is equivalent to sequentially performing one step gradient descent on a sequence of dual models { f ( l ) ( x ) = W ( l ) ϕ ( x ) } L l =1 , where the loss function for the l -th dual model is:

where h ( l -1) i is the output of the ( l -1) -th attention layer for the i -th token, η is the learning rate and D ( l ) is a constant. The input for the l -th dual model is generated by the trained ( l -1) -th dual model.

Proof. Given H ( l -1) as the input for the l -th attention layer, the inference process of h ( l ) N +1 is

where D ( l ) = 1 T N +1 ϕ ( W ( l ) K H ( l -1) ) T ϕ ( W ( l ) Q h ( l -1) N +1 ) is a constant to normalize the attention scores and W ( l ) 0 = 1 D ( l ) ( W ( l ) V h ( l -1) N +1 ) ϕ ( W ( l ) K h ( l -1) N +1 ) T . According to Theorem 3.1, we can easily get the dual model f ( l ) init ( h ) = W ( l ) init ϕ ( h ) where the initialization is W ( l ) init = W ( l ) 0 . Given the loss function L ( l ) formed as Equation 24 and training set { z ( i ) std , y ( i ) std } N i =1 where z ( i ) std = W ( l ) K h ( l -1) i and y ( i ) std = W ( l ) V h ( l ) i , we perform one step SGD with learning rate η on weight matrix W ( l ) and will get trained dual model:

Taking test input as z ( l ) test = W ( l ) Q h ( l -1) N +1 , the prediction ˆ f ( l ) ( z ( l ) test ) will exactly equal to h ( l ) N +1 .

Next, we will show how to obtain H ( l ) D through the trained dual model ˆ f ( l ) ( x ) . And after W ( l +1) K , W ( l +1) Q , W ( l +1) V projections, H ( l ) D will constitute the training set as well as the test input for the next dual model f ( l +1) init ( x ) .

Keeping the initialized dual model f ( l ) init ( x ) and the trained one ˆ f ( l ) ( x ) in mind, we can compute the demonstration token output h ( l ) i ( i = 1 , 2 , ..., N ) of l -th attention layer as

where D ( l ) i = 1 T N ϕ ( W ( l ) K H ( l -1) D ) T ϕ ( W ( l ) Q h ( l -1) i ) is a constant to normalize the attention scores for h ( l ) i . Therefore, once we obtain the trained dual model ˆ f ( l ) ( h ) , we can use the Eq (25) to get the demonstration token output H ( l ) D = [ h ( l ) 1 , h ( l ) 2 , ..., h ( l ) N ] . These demonstration token outputs, along with the output h ( l ) N +1 for query tokens, will together constitute the training set and test input for the next dual model f ( l +1) init ( h ) . This process continues layer by layer until we obtain the ultimate ICL output h ( L ) N +1 . In summary, the ICL inference process across L attention layers is equivalent to performing gradient descent on L dual models sequentially. Thus, we complete our proof.

This theorem is a natural extension of Theorem 3.1: when considering the stacking of multiple attention layers, a sequence of dual models is correspondingly generated. Although these dual

models have the same form f ( x ) = W ϕ ( x ) , they have different initializations and datasets. As the ICL inference process progresses layer by layer between attention layers, we equivalently perform gradient descent on the dual models one by one. The input H ( l ) for each attention layer, including demonstration tokens and query tokens, can be obtained from the test output of the dual models. This can be illustrated in Figure 5.

C 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. We restate our theorem as follows:

Theorem C.1. Define the function class as F := { f ( x ) = W ϕ ( W K x ) | ∥ W ∥ ≤ w } and let the loss function defined as Eq (10). Consider the given demonstration set as S = { x i } N i =1 where S ⊆ S T and S T is all possible demonstration tokens for some task T . With the assumption that ∥ W V x i ∥ , ∥ W ϕ ( W K x i ) ∥ ≤ ρ , then for any δ > 0 , the following statement holds with probability at least 1 -δ for any f ∈ F

Proof. Our proof is similar to the Lemma 4.2 in Saunshi et al. [2019], but here we focus on a different function class. Firstly, we consider the classical generalization bound based on the Rademacher complexity of the function class which can refer to Theorem 3.1 in Mohri et al. [2018]. For a real function class G whose functions map from a set Z to [0 , 1] and for any δ > 0 , if S is a training set composed by N iid samples { x i } N i =1 , then with probability at least 1 -δ 2 , for all g ∈ G

where R S ( G ) is the traditional Rademacher complexity. By setting S exactly the demonstration set and G = { g f ( x ) = -( W V x ) T W ϕ ( W K x ) ∣ ∣ ∥ W ∥ ≤ w } , we can apply this bound to our case.

Then, we construct a function class ˜ F = { ˜ f ( x ) = [ f ( x ); W V x ] = [ W ϕ ( W K x ); W V x ] ∣ ∣ ∥ W ∥ ≤ w } whose functions map from S to R 2 d o . Next, we will first prove R S ( G ) ≤ 2 ρ R S ( ˜ F ) and to do this, we need to use the following Lemma:

Lemma C.2 (Corollary 4 in Maurer [2016]) . Let Z be any set, and S = { z i } i M =1 ∈ Z M . Let ˜ F be a class of functions ˜ f : Z → R n and h : R n → R be L -Lipschitz. For all ˜ f ∈ ˜ F , let g ˜ f = h ◦ ˜ f . Then

where ˜ f |S = ( ˜ f t ( z j ) ) t ∈ [ n ] ,j ∈ [ M ]

We apply Lemma C.2 to our case by setting Z = R d i , S to be exactly the demonstration set, ˜ F to be the function class we constructed and n = 2 d . We also use h : R 2 d o → R where h ( x ) = -⟨ x 1: d o , x d o +1:2 d o ⟩ and thus we have g ˜ f ( x ) = h ( ˜ f ( x )) = h ([ f ( x ); W V x ]) = -( W V x ) T W ϕ ( W K x ) . We can find that g f ( x ) = g ˜ f ( x ) and the left side of inequality (28) is exactly R S ( G ) .

Then we can see that h is √ 2 ρ -Lipschitz with the assumption that ∥ W V x i ∥ , ∥ W ϕ ( W K x i ) ∥ ≤ ρ and we have R S ( G ) ≤ 2 ρ R S ( ˜ F ) . Now using Lemma C.2 and the classical generalization bound (27),

we have that with probability at least 1 -δ 2

Let f ∗ ∈ arg min f ∈F L ( f ) . According to Hoeffding's inequality, with probability at least 1 -δ 2 , we have that ˆ L ( f ∗ ) ≤ L ( f ∗ ) + 3 √ log 2 δ 2 N . Combining this with (29), the fact that ˆ L ( ˆ f ) ≤ ˆ L ( f ∗ ) and applying a union bound, we can get that

Next, we give the upper bound for R S ( ˜ F ) .

Substituting the upper bound of R S ( ˜ F ) into (30), we will get that

Thus we 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 (9), we consider the following representation loss:

where we consider sampling K negative samples for each x i and x -j denotes the j -th negative sample for token x . Correspondingly, the empirical loss will be considered as ˆ L ( f ) = -1 N ∑ N i =1 1 K ∑ K j =1 ( W ϕ ( W K x i )) T ( W V x i -W V x -ij ) where x -ij is the j -th negative sample for x i . Then, by retaining the other definitions in Section 3.3, corresponding to Theorem 3.2, we can obtain the generalization bound as

where r = K N is excatly the the ratio of the number of negative samples. It can be observed that as the ratio of negative samples increases, the generalization error decreases. However, we also notice that 5 N 2 > 1 rN 3 thus the former term dominates, which means the reduction in generalization error due to an increased proportion of negative samples is limited. Nevertheless, we do not rule out the possibility of a tighter generalization bound, which is a promising direction for future research.

Proof Sketch. The proof process is similar to that of Theorem 3.2. The main difference lies in the fact that we should firstly define the function class G = { -1 K ∑ K j =1 ( W ϕ ( W K x i )) T ( W V x i -W V x -j ) ∣ ∣ ∥ W ∥ ≤ w } to use the classical bound. In addition, we define ˜ F = { ˜ f ( x ) = [ f ( x ); W V x ; W V x -1 ; ... ; W V x -K ] ∣ ∣ ∥ W ∥ ≤ w } whose functions map from S to R ( K +2) d o . Similarly, when using Lemma C.2, we set Z = R d i , ˜ F be the above function class and n = ( K + 2) d o . We also use h : R ( K +2) d o → R defined as h ( x ) = -1 K ∑ K j =1 x T 1: d o ( x d o +1:2 d o -x ( j +1) d o +1:( j +2) d o ) . Then we notice that

With the assumption that ∥ W V x ∥ , ∥ Wϕ ( W K x ) ∥ ≤ ρ , we can get that the Frobenius norm of the Jocabian J of h has ∥ J ∥ 2 F ≤ 4 ρ 2 + ρ 2 + K K 2 ρ 2 = (5 + 1 K ) ρ 2 . Thus we get that h is √ 5 + 1 K ρ -Lipschitz. The rest of the proof process is similar to that of Theorem 3.2. Ultimately, we will obtain the aforementioned generalization error.

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 4 of the main body. And we also point out the corresponding modifications in the self-attention mechanism, which are adopted in our experiments.

D.1 More Discussion on the Contrastive Loss

Although we have figured out the representation learning loss of the implicit gradient updates, it can be observed that this loss function has a flaw: due to the lack of normalization for y ( i ) std and ˆ y ( i ) when calculating the cosine distance, the loss can theoretically be optimized to negative infinity. To address this issue, we introduce regularization to constrain the norm of W , that is,

where α is a hyperparameter to balance the two parts. As a result, we can see that the gradient update for W will be in an exponentially smoothed manner meaning that a portion of the initial part will be discarded at every step, that is,

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,

which means more demonstration information will be attended to. This will directly result in Eq 13.

This result can be easily extended to self-attention mechanism. As for a self-attention layer, if all other tokens adopt the same modification, the self-attention layer will become

which leads to the model structure incorporating an operation similar to skip connections. Furthermore, to ensure numerical stability, we normalize the attention scores yielding:

where Norm( · ) is performed column-wise to ensure that the attention scores sum to 1 . The above modification reduce the attention score of each token to its own information during aggregation. It is worth noting that, although our initial intention is to impose regularization on the contrastive loss where α > 0 to prevent it from diverging to negative infinity, we find in experiments that this modification remains effective even when α is less than 0 . Weinterpret this as possibly stemming from the fact that an appropriate α helps the attention block become full-rank, thereby better preserving information, which can be illustrated by Lemma D.1:

Lemma D.1. Let the attention block A ∈ R n × n . There exists some δ > 0 such that, for any 0 < | α | < δ , the attention block A + α I n will become full-rank.

̸

Proof. Define f ( α ) = det( α I n + A ) , which is a polynomial of degree n in α . Then, f ( α ) has only finitely roots. Let α 1 , α 2 , . . . , α r be the non-zero roots of f ( t ) . Now, consider δ = min {| α 1 | , | α 2 | , . . . , | α r |} . For 0 < | α | < δ , we can claim that f ( α ) = det( α I n + A ) = 0 . Thus, A + α I n becomes non-singular (full-rank) and we complete the proof.

Lemma D.1 provides one possible case for appropriate α . In fact, the selection of α can be quite flexible; for instance, similarly, when δ = max {| α 1 | , | α 2 | , . . . , | α r |} and | α | > δ holds, A + α I n also remains full-rank. Our experimental results related to regularized models will further illustrate the effectiveness of an appropriate α in enhancing model performance.

We also acknowledge that our modification is relatively straightforward and may not be optimal. However, we believe that it may be a good choice to make structural improvements to the model from the perspective of the loss function, or more generally, from an optimization standpoint. For example, to address the issue of non-normalized y ( i ) std and ˆ y ( i ) , we can also modify the loss function from the perspective of ridge regression as:

And the optimal W ∗ will be

Correspondingly, the attention mechanism will be modified to

where we neglect the normalization operation. This result is very similar to the mesa-layer proposed by Von Oswald et al. [2023b], which optimizes linear attention layers under the auto-regressive setting. Here, we presented its form on softmax self-attention setting using kernel methods and explained it from the perspectives of contrastive loss and ridge regression. Although the matrix inversion calculation in Eq (33) can be computationally expensive, effective methods for computing Eq (33), including both forward computation and backward propagation, have been thoroughly researched in Von Oswald et al. [2023b], which contributes to making the above modification practically applicable.

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 a simple linear mapping, which may be not sufficient for learning deeper-level features. To address this, we can employ more complicated nonlinear functions for more complex augmentations. Denoting these two augmentations as g 1 and g 2 , consequently, the process of contrastive learning will be modified as follows

Correspondingly, the gradient update for W will become

And from the perspective of ICL, correspondingly, the last token will be updated as

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,

where g 1 ( · ) and g 2 ( · ) will be column-wise here. It is worth noting that here we have only presented the framework of using nonlinear functions as data augmentations to modify the self-attention layer and in the simplest case, we can set g 1 ( x ) and g 2 ( x ) as MLPs (Multi-Layer Perceptrons). However, in practice, it is encouraged to use data augmentation functions that are tailored to specific data structures. For example, in the case of CMT [Guo et al., 2022], the used Convolutional Neural Networks (CNNs) can be considered as a form of "strong data augmentations" suitable for image datas within our framework. We consider the exploration of various augmentation methods tailored to different types of data as an open question for future research.

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 to most contrastive learning approaches, we can introduce negative samples forcing the model to separate the distances between positive and negative samples at the same time, that is,

where ˜ x i = x i -β |N ( i ) | ∑ j ∈N ( i ) x j , N ( i ) is the set of the negative samples for x i and β is a hyperparameter. As a result, the gradient descent on W will be modified as

Correspondingly, the ICL process for ˆ h N +1 will be

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

where ˜ X ( i ) = ˜ x i . In corresponding experiments, for each token, we simply choose other the k least relevant tokens as its negative samples, i.e., the k tokens with the lowest attention scores. Noting that here we simply use other token representations as negative samples for x i . However, there are more ways to construct negative samples that are worth exploring (for instance, using noise vectors or tokens with low semantic similarity as negative samples). For specific data structures and application scenarios, customizing the selection or construction of negative samples may be more effective.

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. [2023a], we choose to pretrain a softmax attention layer before exploring the equivalence proposed by Theorem 3.1. In fact, pretraining is not mandatory since our theoretical analysis does not depend on any specific weight construction. In other words, the inference results of ICL and the test prediction of the dual model will still remain consistent for an attention layer with arbitrary weights or even random initialization. However, for the convenience of further investigating the impact of subsequent modifications to the model structure and to better align with real-world scenarios, we still opted for pretraining to let the model acquire some task-specific knowledge. Additionally, our experiments are conducted in a self-attention setting. When we focus only on the last token, this is equivalent to considering the case with only one query token ( T = 0 ) in Section 2.1. The experiments are completed on a single 24GB NVIDIA GeForce RTX 3090 and the experiments can be completed within one day.

For the linear regression task, we generate the task by s = Wt where every element of W ∈ R d s × d t is sampled from a normal distribution W ij ∼ N (0 , 1) and t is sampled from a Gaussian distribution x ∼ U ( -1 , 1) d t . To facilitate more accurate estimation of attention matrices using random features and considering the limited learning capacity of a single attention layer, we only set a small value for d t = 11 and d s = 1 . Then, at each step, we use generated { x i = [ t i ; s i ] } N +1 i =1 to form the input matrix X while the label part of the query token is masked to be zero, that is, x N +1 = [ t i ; 0] where we consider only one query token and we denote x ′ T +1 = x N +1 to maintain consistency of notation in Section 2.1. The softmax attention layer is expected to predict ˆ s N +1 to approximate the ground truth value s N +1 . We use mean square error (MSE) as the loss function, that is, for each epoch,

where ˆ s ( j ) N +1 and s ( j ) N +1 are the prediction and ground truth value at j -th step and N step is the number of steps. We set N step = 1024 for N +1 = 16 which means the total number of tokens remains

(a)

d

r

= 3

(b)

d

r

= 12

(c)

d

r

= 120

(d)

d

r

= 12000

(e) exact attention

0

2

4

6

8

10

12

14

0

2

4

6

8

10

12

14

0

2

4

6

8

10

12

14

0

2

4

6

8

10

12

14

0

2

4

6

8

10

12

14

0

0

0

0

0

0.7

2

0.7

2

0.7

2

0.7

2

0.7

0.6

0.6

0.6

0.6

0.6

4

0.5

4

0.5

4

0.5

4

0.5

4

- 0.5

6

0.4

6

0.4

6 -

0.4

6

0.4

9

- 0.4

8

0.3

8

0.3

8-

0.3

8

0.3

8

0.3

10

0.2

10

- 0.2

10

0.2

10

0.2

10

0.2

12

0.1

12

- 0.1

12

0.1

12

0.1

12

0.1

14

14

14

14

14

(a)

d

r

= 3

(b)

d

r

= 12

(c)

d

r

= 120

(d)

d

r

= 12000

(e) exact output

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

0

0

0

0

0

2

2

2

2

2

4

1

4

1

4

1

4

1

4

1

6

6

6

6

6

8

0

8

0

0

8

0

8

0

10

-1

10

-1

10

-1

10

-1

10

12

12

12

12

12

14

14

2

14

Z

14

Z

14

16384 . We choose stochastic gradient descent (SGD) [Amari, 1993] as the optimizer and we set the learning rate to 0.003 for normal and regularized models, while the remaining experiments to 0 . 005 . We also attempt the multi-task scenario, where the input token at each step is generated from a different task. However, we find it challenging for a single attention layer to effectively learn in this setting, resulting in disordered predictions. Therefore, our experiments are currently limited to single-task settings, and the multi-task scenario is worth further investigation in the future.

It is worth noting that we approximate the attention matrix calculation using random features as kernel mapping function instead of using the traditional softmax function in the self-attention layer [Choromanski et al., 2020]. The mapping function ϕ : R d o → R d r has the form of ϕ ( x ) = e w T x -∥ x ∥ 2 / 2 where w ∼ N (0 , I ) . Orthogonal random features [Yu et al., 2016, Choromanski et al., 2020] or simplex random features [Reid et al., 2023] can be chosen to achieve better performance theoretically. We investigate the impact of changing the dimension of random features d r on the approximation of attention matrices and output, using Mean Squared Error (MSE) and Mean Absolute Error (MAE) as evaluation metrics, where we conduct 50 repeated experiments and calculated the average values for each value of d r , as shown in Figure 8. It can be observed that as the dimension of random features increases, the approximation performance gradually improves, with both errors reaching a low level in the end. We visualize the exact attention matrix and compare it with the estimated attention matrices obtained using different values of d r , as shown in Figure 6. Again, it can be seen that as d r increases, the approximation of the true attention matrix improves gradually and similar results can be observed for the analysis of output matrices in Figure 7.

To obtain a more accurate estimation of the attention matrix, we set the output dimension of the mapping function to be 100 times the input dimension, that is, d r = 100( d s + d t ) = 1200 . Furthermore, we visualize the exact attention matrix and the output with the approximation results, which are shown in the Figure 9. As we can see, although some larger values are not estimated accurately due to the limited dimension of the random features we select, the majority of the information is still estimated comprehensively well. These findings indicate that our choice of using positive random features as mapping functions to estimate the true softmax attention and conduct experiments is relatively feasible.

After the weights W Q , W K , W V of the attention layer have been determined, we generate test N +1 tokens in the same way where the s part of the ( N +1) -th token is also set to be zero and finally input the test tokens into the attention layer to obtain the corresponding predicted ˆ h N +1 = [ ˆ t (1) N +1 , ˆ s (1) N +1 ] . Here, we also use h ′ T +1 = ˆ h N +1 to maintain the notation consistency in Section 2.1.

(a) Estimation error of attention ma- trix vs. d r

(b) Estimation error of attention ma- trix vs. d r

AttentionError-dr

OutputError-dr

0.09

0.08

MSE

MSE

0.07

MAE

0.30

MAE

0.06

0.25

ror

0.05

Error

0.20

E

0.04

0.03

0.15

0.02

0.01

0.10

0.00

0

200

400

600

800

1000

1200

0

200

400

600

800

1000

1200

'p

'p

attention

exactattention

output

exactoutput

0

468101214

0

2468101214

0

2

4

6

810

0

2

4

6

810

0

0

0

0

2

0.6

2

2

2

4

4

4

1

4

6

0.4

6

6

6

8

8

8

0

8

0

10

0.2

10

10

1

10

-1

12

12

12

12

14

14

14

14

(a) the exact attention matrix and its approximation (b) the exact output and its approximation

On the other hand, we construct a dual model f ( x ) = W ϕ ( x ) where ϕ ( · ) is strictly equivalent to the kernel mapping function used in the attention layer. We transform the first N tokens as the training set according to Theorem 3.1 and train the dual model using the loss formed by Eq (9). In fact, according to Theorem 3.1, after we perform one step of gradient descent on this training set, the test prediction ˆ y test = [ ˆ t (2) N +1 , ˆ s (2) N +1 ] of the dual model will strictly equal h ′ T +1 .

We conduct experiments under the same setup using different random seeds to explore the effects of various model modifications. The data for all three experiments are generated under identical conditions. One set of experimental results is presented in the main text, while the results of the other two sets are shown in the Figure 10. Similar to the discussion in the main text, we can achieve better performance than the normal model with appropriate parameter settings.

Regularized Models

Augmented Models

Negative Models

2.5

α=0

2.5

normal

2.5

normal

α=0.5

g1&g2

k=1, β=0.1

2.0

α=1

2.0

91

2.0

k=1,β=0.5

α=-0.5

g2

k=3,β=0.1

α=-1

g

k=3,β=0.5

1.0

1.0

1.0

0.5

0.5

0.5

0.0

0.0

0.0

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

epochs

epochs

epochs

Regularized Models

Augmented Models

Negative Models

2.5

2.5

2.5

α=0

normal

normal

2.0

α=0.5

2.0

g1&g2

2.0

k=1, β=0.1

α=1

91

k=1,β=0.5

LoSS

1.5

α=-0.5

g2

Loss

1.5

k=3,β=0.1

α=-1

g

k=3,β=0.5

1.0

1.0

0.5

0.5

0.5

0.0

0.0

0.0

0

20

40

60

80

100

0

20

40

60

80

100

0

20

40

60

80

100

epochs

epochs

epochs

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

test-h+1ll2 at each step

RegularizedModels

Augmented Models

Negative Models

ION

2.0

1.5

1.0

1

1

5

S

0.200

0.300

0.275

0.250

0.225

0.175

α=0

α=1

α=-2

α=2

α=-1

5

S

S

0.30

0.25

0.45

0.40

0.35

g1&g2

normal

91

g2

g

S

S

0.50

0.45

0.40

0.35

0.30

0.25

normal

k=24,β=0.5

k=8,β=0.1

k=8,β=0.5

k=24,β=0.1

0.5

0.150

0.125

0.20

0.15

0.20

0.15

0.0

20

40

60

80

100

120

0.100

100

150

200

250

3

0.10

0

100

200

300

400

50C

0.10

0

50

100

150

200

250

300

steps

epochs

epochs

epochs

For trigonometric task, we generate the task by s = cos( Wt ) where cos( · ) is element-wise, W ∈ R d s × d t is sampled from the normal distribution W ij ∼ N (0 , 1) while t is sampled from the uniform distribution x ∼ U (0 , π ) d t . In experiments, we found that for one softmax attention layer, learning higher-dimensional tasks is challenging. Therefore, we only set d t = 7 and d s = 1 . At each step, we use N +1 = 128 tokens { x i = [ t i ; s i ] } N +1 i =1 and the total number of tokens remains unchanged at 16384 . Compared to the setting N +1 = 16 of linear tasks, we observed that for more complex tasks, the attention layer needs to use more tokens to provide information at each training step. Similarly, we mask the label part of the last token, that is, s N +1 = 0 and use mean square error (MSE) loss to train the attention layer. We choose SGD as the optimizer and the learning rate is set as 0.005. The rest of the settings remain consistent with those used in the linear task. The result for trigonometric regression task is shown in Figure 11.

Firstly, as shown in the left part of Figure 11, the inference results is of ICL is strictly equivalent to the prediction of the dual model, that is, ˆ h N +1 = ˆ y test as well as the label part ˆ s (1) N +1 = ˆ s (2) N +1 , aligning with our analysis in Theorem 3.1.

The performance of modified model during training process can be seen in the remaining parts of Figure 11. For regularized models, as seen in the center left part of figure 11, the models when α < 0 converge slightly faster and reach better final results compared to the normal model ( α = 0 ). For augmented models, we use as the same augmentation functions g 1 and g 2 as the ones in the linear regression task, that is, g 1 ( x ) = g 2 ( x ) = σ ( Wx ) where σ ( · ) is GELU activation function. However, for g + 2 , we use ELU as the activation function. We can find from the center right part of Figure 11 that, compared to the normal model, using g 1 alone and using g 1 and g 2 simultaneously as data augmentations significantly degrade the model's performance, including convergence speed and final results. However, using g 2 alone yields comparable result with the normal model. Particularly, when using g + 2 , the model accelerates its convergence speed. However, for negative models, the performance with the selected number of negative samples k and the parameter β is worse than the normal model, which suggests that our simple approach of selecting those tokens low attention scores as negative samples is not a reasonable method. Just as we discussed in Section 4, for different tasks, a more refined strategy for selecting negative samples should be considered.

E.2.2 More details of Experiments on Exponential Tasks

For exponential task, we generate the task by s = exp( Wt ) where exp( · ) is also element-wise, W ∈ R d s × d t is sampled from the normal distribution W ij ∼ N (0 , 1) while t is sampled from the uniform distribution x ∼ U ( -1 , 1) d t . We only set d t = 6 and d s = 1 considering the limited learning capacity of one softmax attention layer. At each training step, we use N +1 = 512 tokens { x i = [ t i ; s i ] } N +1 i =1 and the total number of tokens remains unchanged at 16384 . Compared to the

test-h+1ll2at each step

0.30

Regularized Models

0.351

Augmented Models

3.5

NegativeModels

ION

1.0

0.6

0.8

0.4

1

1

1

S

S

0.25

0.20

0.15

α=0

α=16

α=-16

α=-64

5

S

S

0.30

0.25

0.20

0.15

g2

g

normal

91

5

S

S

2.0

3.0

2.5

1.5

normal

k=8,β=0.1

k=8, β=0.5

k=32, β=0.1

k=32, β=0.5

1.0

0.2

0.10

0.10

0.5

0.0

100

200

300

400

500

0.05

50

100

150

200

250

30

0.05

50

100

150

200

250

30C

0.0

0

25

50

75

100

125

150

steps

epochs

epochs

epochs

setting N +1 = 16 of linear tasks and N +1 = 128 of trigonometric tasks, we also find that for exponential tasks, the attention layer needs more tokens to provide in-context information at each training step. The rest of the settings remain consistent with those used in the trigonometric task. The result for exponential regression task is shown in Figure 12.

Similarly, as shown in the left part of Figure 12, the result ˆ h N +1 of ICL inference is equivalent to the test prediction ˆ y test of the dual model after training, just as stated in Theorem 3.1. For regularized models, it can be observed that when α = 16 , the model converges faster and achieves better result. For augmented models, using g 1 or g 2 alone as data augmentations results in better performance. However, when both g 1 and g 2 are used simultaneously, the training process becomes unstable, so we did not show it in the center right part of Figure 12. For negative model, similar to the case in the trigonometric task, the different combinations of negative samples' number k and parameter β do not show a significant improvement over the normal model, highlighting the importance of the strategy for selecting negative samples. We leave the exploration of a more refined negative sample selection strategy when facing various tasks for future consideration.

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, a combination of regularized and augmented modifications is sufficient. However, for the other two tasks, the results are actually worse than using regularized or augmented modification individually (compared to Figures 11 and 12). We think this may be due to the ineffective selection of negative samples, which is amplified when combined. Therefore, when the design of augmentation or negative sample improvement methods is not effective, we recommend using a single modification method.

2.00

Linear Tasks

3.5

Exp Tasks

0.50

Tri Tasks

1.75

normal

3.0

normal

0.45

normal

5

S

S

1.00

1.50

1.25

regu + aug + nega

aug + nega

5

S

S

2.0

2.5

1.5

aug + nega

regu + aug + nega

5

0.30

S

0.40

0.35

0.25

regu + aug + nega

ebau+6ne

0.75

0.50

1.0

0.20

0.15

0.25

0.5

0.10

0.00

0

20

40

60

80

100

0.0

0

50

100

150

200

250

300

0.05

0

50

100

150

200

250

300

epochs

epochs

epochs

E.4 More Experiments on One Transformer Layer

Similar to the experiments with one softmax attention layer, we also conduct experiments on a Transformer layer (introducing one FFN layer after the attention layer) and trained its dual model based on Theorem B.1. As shown in Figure 14, the inference result ˆ h N +1 of ICL remains equivalent to the test prediction ˆ y test of the trained dual model. Furthermore, to validate the potential low-rank

IlItest -hT+1ll2 at each step

Upperbound of Rank(WF)

4

min{d,dn,Rank(m)}

12

10

Norm

m

1

8

2

6

4

linear

1

cOS

2

exp

2.5

5.0

7.5

steps

10.0

12.5

15.0

0

6

12

18

24

dn

30

36

42

48

property of matrix W F , we explore its upper bound of rank. Noting that W F = W 2 I M W 1 where W 1 ∈ R d h × d , W 2 ∈ R d × d h , I M ∈ R d h × d h , the upper bound of Rank( W F ) is

where Rank( I M ) is equivalent to the number of non-zero elements in I M . We fix d = 12 while varying the values of d h . Wegenerate 1024 sets of X test for different tasks and repeat the experiments 5 times. Finally, we calculate the average upper bound of the rank of W F . The results are shown in the right part of Figure 14, indicating that when d h ≥ 2 . 75 d = 33 , the upper bound remains stable and equals d = 12 . Otherwise, when d h is set to a smaller value, W F exhibits clear low-rank property.

E.5 More Experiments on More Realistic NLP Tasks

We supplement our experiments on more more realistic NLP tasks. We choose the BERT-baseuncased model (can be downloaded from Huggingface library[Wolf, 2019], hereafter referred to as BERT[Kenton and Toutanova, 2019]) to validate the effectiveness of modifications to the attention mechanism and select four relatively smaller GLUE datasets (CoLA, MRPC, STS-B, RTE) [Wang, 2018]. We load the checkpoint of the pre-trained BERT model, where 'classifier.bias' and 'classifier.weight' are newly initialized, and then we fine-tune the model to explore the performance of three attention modifications as well as their combinations. In terms of more detailed experiment settings, we set the batch size to 32, the learning rate to 2e-5, and the number of epochs to 5 for all datasets. All experiments are conducted on a single 24GB NVIDIA GeForce RTX 3090. All experimental results are presented in Table 1. Below, we discuss the various modifications and their performance.

For the regularized modification , we consider different values of α , specifically selected from {-0 . 5 , -0 . 1 , 0 . 1 , 0 . 5 } . As can be observed in Table 1, except for RTE, the best regularized models outperform the original model on the other three datasets. However, we also note that when the absolute value of α is too large, the model's performance declines significantly, so we recommend using smaller absolute values for α .

For the augmented modification , we also consider applying more complex 'augmentation' functions to the linear key/value mappings. However, unlike the previous methods used in simulation tasks, we do not simply select g 1 and g 2 as MLPs, i.e., g 1 ( W V x ) = W 2 σ ( W 1 W V x ) . This design is avoided because it could undermine the effort made during pre-training to learn the weights W V and W K , leading to difficulties in training and challenges in comparison. Instead, we adopt a parallel approach, i.e., g 1 ( W V x ) = W V x + c W 2 σ ( W 1 x ) , where c is a hyperparameter to control the influence of the new branch, σ is the GELU activation function and the hidden layer dimension is set to twice the original size of W V x . g 2 ( W K x ) = W K x + c W 2 σ ( W 1 x ) follows the same format.

Experimental results show that the best augmented models achieve better performance than the original model across all four datasets. Notably, augmentation on the value mapping (i.e., using g 1 alone) proves to be more effective than other methods, both in terms of performance and the amount of additional parameters introduced. Using both g 1 and g 2 introduces more parameters, which is particularly undesirable for larger models. Thus, under the augmentation methods and experimental settings we selected, using g 1 alone is recommended.

In addition, we do not rule out the possibility of more powerful and efficient augmentation methods. Our choice of g 1 and g 2 as parallel MLPs is primarily motivated by the desire to make better use of

the pre-trained weights W K and W V . We have also noticed that this specific augmentation function design is structurally similar to the Parallel Adapter [He et al., 2021]. However, we would like to emphasize that our parallel design is just a specific case within this broader augmented modification framework and this is a new perspective for understanding the Parallel Adapter. As for practical implementation, the Parallel Adapter method focuses more on efficient training, so it uses fewer parameters, and the original W V and W K are freezed-only the newly introduced parameters are trained. In contrast, our approach aims to validate the benefits of introducing stronger nonlinear augmentation functions into the linear value/key mappings. Therefore, we set a higher hidden layer dimension (twice that of W V x or W K x ) and also train W V and W K simultaneously. This design is relatively general and does not take into account the specific characteristics of individual tasks. We still encourage the development of more task-specific augmentation strategies tailored to different tasks.

For the negative modification , we continue to select tokens with lower attention scores as negative samples. The parameter r represents the proportion of tokens used as negative samples, while β indicates the overall reduction in attention scores. We choose r from { 0 . 1 , 0 . 2 , 0 . 3 } and β from { 0 . 1 , 0 . 2 } . Under these combinations, the best negative models only outperform the original model

on CoLA and STS-B, whereas their performance on MRPC and RTE is worse than the original one. This suggests that our simple approach of considering tokens with low attention scores as negative samples might be too coarse. A more effective method for constructing negative samples should be designed, which is a direction worth exploring in the future.

We also consider combining different modification methods . Specifically, we choose α = 0 . 1 , g 1 /c = 1 and r = 0 . 2 /β = 0 . 1 respectively as the basis for combining the three types of modifications, considering their overall performance across all datasets. The results indicate that under our settings, the combination of augmented and negative modification achieves the best performance on CoLA, MRPC, and RTE, while the combination of regularized and augmented modification achieves the best performance on STS-B. However, their optimal performance is slightly inferior to the best performance achieved with augmented models alone. Therefore, we conclude that using all three modifications simultaneously is not necessary. With appropriate hyperparameter choices, using augmented modification alone or in combination with one other modification is sufficient.

Overall, the experimental results show that our modifications inspired by the representation learning process are helpful in enhancing performance. This further validates the potential of our approach of thinking about and improving the attention mechanism from a representation learning perspective. In addition, we would like to reiterate that more validation across additional tasks and models, and the development of task-specific augmentation and negative sampling methods are all interesting directions worth exploring in the future.

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 as implicit fine-tuning: More specifically, let X = [ X D , X T ] where X D = [ x 1 , x 2 , . . . , x N ] denotes the demonstration tokens and X T = [ x ′ 1 , x ′ 2 , . . . , x ′ T ] be query tokens. On the one hand, for ICL, they consider the output of q = W Q x ′ T +1 under the linear attention setting as

where W ZSL q is interpreted as the output in the zero-shot learning (ZSL) where no demonstrations are given. On the other hand, they consider a specific fine-tuning setting, which updates only the parameters for the key and value projection, that is,

where ∆ W K and ∆ W V denote the parameter updates and they are acquired by back-propagation from task-specific training objectives [Dai et al., 2022], which is a supervised learning process of the original model. Considering the similarity in form between ˜ F ICL and ˜ F FT , their focus is on establishing a connection between ICL and implicit fine-tuning on the original model.

As a comparison, we turn our attention to establish a connection between ICL and the gradient descent process of the dual model, rather than the original model. More specifically, we consider the dual model f ( x ) = W ϕ ( x ) of the nonlinear attention layer, where the weight W are updated according to the following loss (presented as Eq (9) in Section 3.2):

where x i is the i -th demonstration token. The prediction output of the trained dual model will be consistent with the ICL output of the attention layer. The gradient descent process of the dual model

using this loss can be viewed from a self-supervised learning lens: unlike in supervised fine-tuning, where the original model is instructed to perform gradient descent using a given objective (loss), this loss formed as Eq (9) is determined (derived) by the attention mechanism itself and it also does not require additional "true label" to supervise each token x i (so called self-supervised). Therefore, modifications to this self-supervised learning loss will in turn cause modifications in the attention mechanism correspondingly, as we discussed in our work in Section 4. We believe this perspective offers several benefits:

By analyzing from the dual perspective, we can transform the forward inference process into an optimization process. Since optimization processes are well-known and have established theoretical tools (for example, generalization error as mentioned in Section 3.3), this transformation can provide reverse insights into analyzing the model mechanisms.

It can clearly observed that the dual model involves a self-supervised representation learning process from the dual perspective. Considering that there are lots of mature works in this area, we can draw on these works to reflect on the attention mechanism, which has also inspired attention modifications as illustrated in Section 4.

Intuitively, this explanation might be also reasonable as the original model is not explicitly instructed to provide the answer under some given objective (e.g., minimizing cross-entropy) during ICL inference process. Instead, the underlying criterion should be determined by the model's own structure (self-supervised) as we mentioned above.

In addition, although we do not target specific tasks like linear regression as previous works mentioned in Section 6, we would like to point out that under those specific weight and input settings, an intuitive explanation can also be provided from a representation learning perspective. Here, we take the linear regression task as well as the weight constructions considered by Von Oswald et al. [2023a] as an example. Specifically, it assumes that the structured input is H = [ h i ] N i =1 ∈ R ( d +1) × ( N ) where h i = [ x i , y i ] is sampled from some linear task y = w T x and the query token will be h N +1 = [ x N +1 , -w T 0 x N +1 ] . And the considered linear self-attention layer will take the constructed weights and query output as:

where w 0 is the underlying initial matrix. Then the label part of ˜ h N +1 will has the form as ˜ y N +1 = -w T 0 x N +1 +∆ w T x N +1 = -( w T 0 -η N ∑ N i =1 ( w T 0 x i -y i ) x T i ) x N +1 = -ˆ y N +1 , which is equivalent to the output -ˆ y N +1 (multiplied by -1 ) of the linear layer y = w T x where w is initialized as w 0 after performing one step of gradient descent under mean squared loss L = 1 2 N ∑ N i =1 ∥ w T 0 x i -y i ∥ 2 .

In practice, the underlying initial weight matrix w 0 is set to be approximately 0 thus the test input can be formed as h N +1 = [ x i , 0 ] [Von Oswald et al., 2023a]. In addition, when reading out the label ˆ y N +1 , the test prediction ˜ y N +1 will be multiplied again by -1 , which can be done by a final projection matrix (or equivalently, P = -η N I ). In this case, we first note that the dual model of the linear attention layer can be written as f ( z ) = Wz where W ∈ R ( d +1) × ( d +1) and similar to Eq (9), it will be trained under the loss below:

By substituting the corresponding weights in Eq (35) where we replace P = -η N I for the readout, the loss can be reformulated as:

Recalling that h i = [ x i , y i ] is sampled from some linear task y = w T x , we assume that ∥ W ∥ F ≤ ∥ w ∥ 2 , it can then be easily seen that the optimal solution for Eq (37) will be

Furthermore, similar to Section 3.2, we take W Q h N +1 as the input where W Q is constructed as Eq (35) and h N +1 = [ x N +1 , 0] , the optimal dual model will output the result f ( W Q h N +1 ) = W ∗ W Q h N +1 = [ 0 , w T x N +1 ] = [ 0 , y N +1 ] where the label part will be just the answer for the test query. Additionally, it would also be interesting to explore how these weights converge to the constructed form in Eq (35) or other forms under this special setting as previous works illustrated from the perspective of the dual model. Investigating this issue goes beyond the scope of this paper, and we will leave it for future exploration.

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 accurately reflect our contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We confirm that the paper provides the full set of assumptions and a complete (and correct) proof for each theoretical result.

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 crossreferenced.

All assumptions should be clearly stated or referenced in 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 a short proof sketch to 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: We confirm that we have provided a detailed explanation of 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have provided our code and instructions in the supplemental material. Our experiments are primarily simulation-based, and we have provided the code to generate the data in the code part.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a subset of experiments are 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 to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have provided as detailed an explanation as 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 of the paper to a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have provided sufficient information on the computer resources in Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We confirm that our research meets the requirements.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

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.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We think that our paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: Our paper does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\nTheoretical 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.\n\nLIMITATIONS:\nThe authors have clearly stated the limitations of this work, the selected transformer components of this work, and selected task settings.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n(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.\n\nWEAKNESSES:\n(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\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
The author's present a new way of linking in-context-learning (ICL) to gradient descent. 

The author's are able to demonstrate that indeed a (simplified) transformer decoder layer ICL is equivalent to "representation learning".

Using the theoretical findings the authors are also able to propose extensions to the attention mechanism.

The paper shows theoretical proofs for these claims as well as additional experiments to demonstrate the claim.

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

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

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

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

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

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

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

QUESTIONS:
Please see the weakness section.

LIMITATIONS:
The paper has discussed the limitations sufficiently.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-20