Title: In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformers

self-attention

in-context learning
Abstract: A striking property of transformers is their ability to perform in-context learning (ICL), a machine learning framework in which the learner is presented with a novel context during inference implicitly through some data, and tasked with making a prediction in that context. As such, that learner must adapt to the context without additional training. We explore the role of *softmax* attention in an ICL setting where each context encodes a regression task. We show that an attention unit learns a window that it uses to implement a nearest-neighbors predictor adapted to the landscape of the pretraining tasks. Specifically, we show that this window widens with decreasing Lipschitzness and increasing label noise in the pretraining tasks. We also show that on low-rank, linear problems, the attention unit learns to project onto the appropriate subspace before inference. Further, we show that this adaptivity relies crucially on the softmax activation and thus cannot be replicated by the linear activation often studied in prior theoretical analyses.

In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness

Abstract

A striking property of transformers is their ability to perform in-context learning (ICL), a machine learning framework in which the learner is presented with a novel context during inference implicitly through some data, and tasked with making a prediction in that context. As such, that learner must adapt to the context without additional training. We explore the role of softmax attention in an ICL setting where each context encodes a regression task. We show that an attention unit learns a window that it uses to implement a nearest-neighbors predictor adapted to the landscape of the pretraining tasks. Specifically, we show that this window widens with decreasing Lipschitzness and increasing label noise in the pretraining tasks. We also show that on low-rank, linear problems, the attention unit learns to project onto the appropriate subspace before inference. Further, we show that this adaptivity relies crucially on the softmax activation and thus cannot be replicated by the linear activation often studied in prior theoretical analyses.

1 Introduction

One of the most compelling behaviors of pretrained transformers is their ability to perform in-context learning (ICL) [1]: determining how to solve an unseen task simply by making a forward pass on input context tokens. Arguably the most critical innovation enabling ICL is the self-attention mechanism [2], which maps each token in an input sequence to a new token using information from all other tokens. A key design choice in this self-attention architecture is of the activation function that controls how much 'attention" a token pays to other tokens. Softmax -activated self-attention (i.e. softmax attention) is most commonly, and successfully, used in practice [1, 3-6].

A natural approach to explain ICL adopted by the literature is to equate it with classical machine learning algorithms, primarily variants of gradient descent (GD). Several works have shown that

xquery

1.0

0.5

0.0

0.5

1.0

Position Relative to

xquery

0.000

0.005

0.010

0.015

0.020

Attention Weights

Softmax

Linear

1.0

0.5

0.0

0.5

1.0

1.0

0.5

0.0

0.5

1.0

100

200

300

400

Number of Context Samples

0.00

0.25

0.50

0.75

1.00

1.25

1.50

ICL Error

Softmax

Linear

100

200

300

400

100

200

300

400

when the ICL tasks are linear regressions and the activation in the attention unit is identity (referred to as linear attention), transformers that implement preconditioned GD during ICL are global optima of the pretraining loss, which is the population loss on ICL tasks [7-9]. In particular, the prediction of such transformers with l linear attention layers equals the prediction of a regressor trained with l preconditioned GD steps on the context examples. However, since these analyses are limited to linear attention and tasks, they do not explain the widespread success of softmax attention at ICL.

More recent work [10] extends these results by showing that for general regression tasks and any attention activation that is a kernel, ICL equates to training a kernel regressor via functional GD in the Reproducing Kernel Hilbert Space (RKHS) induced by the activation. However, this functional GD yields generalization guarantees only when the activation kernel is identical to a kernel that generates the labels, which does not apply to the softmax activation, as it is not a kernel. Further, like the aforementioned studies of the linear setting [7-9], this analysis only shows that pretraining leads to learning the covariate distribution, while the activation implicitly encodes the label distribution needed for accurate predictions. Thus, this line of work has not explained the very fundamental question of what softmax attention learns during pretraining that enables it to perform ICL on a wide variety of downstream tasks. Motivated by this gap in the literature, we ask the following question.

How does softmax attention learn to perform ICL?

To answer this question, we study general settings in which pretraining and evaluation ICL tasks are regressions that share only Lipschitzness and label noise variance . Specifically, the rate at which their ground-truth labels change along particular directions in the input space, and the variance in the label noise, is similar across tasks. In such settings, we observe that softmax attention acts as a nearest neighbors regressor with an attention window - i.e. neighborhood of points around the query that strongly influence, or 'attend to', the prediction - that adapts to the pretraining tasks. Specifically, our main result is as follows:

Main Claim: Softmax attention performs ICL by calibrating its attention window to the Lipschitzness and label noise variance of the pretraining tasks.

While this does not contradict the line of work showing that ICL manifests via a 'meta-learned" gradient-based algorithm, we show in a general setting that a simpler mechanism can explain the capabilities of a widely accepted model of ICL.

Outline. We substantiate the above claim via two streams of analysis. To our knowledge, these are the first results showing that softmax attention pretrained on ICL tasks recovers shared structure among the tasks that facilitates ICL on downstream tasks.

(1) Attention window scale adapts to Lipschitzness and noise variance - Section 3. We prove that the pretraining-optimal softmax attention estimator scales its attention window inversely with the task Lipschitzness and jointly with the noise level to optimally trade-off bias and variance in its prediction (Theorem 3.4). This requires tight upper and lower bounds on the pretraining ICL loss. While the upper bounds (Lemma C.8) hold for all L -Lipschitz tasks, the lower bounds (Lemma C.9) are more challenging and require considering specific classes of tasks. We consider two classes of generalized linear models (GLMs), and obtain lower bounds via novel concentrations for particular functionals on the distribution of the attention weights for tokens distributed on the hypersphere (Corollary G.5).

(2) Attention window directions adapt to direction-wise Lipschitzness - Section 4. We prove that when the target function class consists of linear functions that share a common low-dimensional structure, the optimal softmax attention weight matrix from pretraining projects the data onto this subspace (Theorem 4.4). In other words, softmax attention learns to zero-out the zero-Lipschitzness directions in the ambient data space, and thereby reduces the effective dimension of ICL. We prove this via a careful symmetry-based argument to characterize a particular gradient of the ICL loss as positive (Lemmas H.3 and H.4).

Tightness of results. Our results highlight the importance of shared Lipschitzness across training and test, as well as the critical role of the softmax activation, to ICL. We show that softmax attention pretrained on the setting from Section 3 in-context learns any downstream task with similar Lipschitzness to the pretraining tasks, while changing only the Lipschitzness of the evaluation tasks degrades performance (Theorem 3.5) - implying learning Lipschitzness is both sufficient and necessary for generalization . Further, to emphasize the necessity of the softmax , we show that the minimum ICL loss achievable by linear attention exceeds that achieved by pretrained softmax attention (Theorem 3.6). We verify all of these results with empirical simulations (Section 3.2 and Appendix J).

Notations. We use (upper-, lower-)case boldface for (matrices, vectors), respectively. We denote the (identity, zero) matrix in R d × d as ( I d , 0 d × d ), respectively, the set of column-orthonormal matrices in R d × k as O d × k , and the (column space, 2-norm) of a matrix B as (col ( B ) , ∥ B ∥ ), respectively. We indicate the unit hypersphere in R d by S d -1 and the uniform distribution over S d -1 as U d . We use asymptotic notation ( O , Ω ) to hide constants that depend only on the dimension d .

1.1 Additional Related Work

Numerous recent works have constructed transformers that can implement GD and other machine learning algorithms during ICL [11-15], but it is unclear whether pretraining leads to such transformers. [16] and [13] provide generalization bounds for ICL via tools from algorithmic stability and uniform concentration, respectively. [17] investigate the pretraining statistical complexity of learning a Bayes-optimal predictor for ICL on linear tasks with linear attention. [18-20] study the role of the pretraining data distribution, rather than the learning model, in facilitating ICL. [21] studies the dynamics of a softmax attention unit trained with GD on ICL tasks, but this analysis considers only linear tasks and orthogonal inputs. The connection between ICL with softmax attention and non-parametric regression has been noticed by other works that analyze the ICL performance of a softmax-like kernel regressor [22] and aim to improve upon softmax attention [23-27] rather than explain what it learns during pretraining. Please see Appendix A for further discussion of the large body of related works studying the theory of transformers, ICL and kernel regression.

2 Preliminaries

In-Context Learning (ICL) regression tasks. We study ICL in the regression setting popularized by [28], wherein each task is a regression problem in R d . The context for task t consists of a set of n feature vectors paired with noisy labels { x ( t ) i , f ( t ) ( x ( t ) i ) + ϵ ( t ) i } n i =1 , where f ( t ) : R d → R generates the ground-truth labels for task t and ϵ ( t ) i is label noise. Given this context, the model solves the task if it accurately predicts the label of a query x ( t ) n +1 . During pretraining, the model

observes many such tasks. Then, it is evaluated on a new task with context { x i , f ( ∗ ) ( x ( ∗ ) i ) + ϵ ( ∗ ) i } n i =1 and query x ( ∗ ) n +1 . We emphasize that the model is trained only on the pretraining tasks, not the evaluation context. Unlike traditional supervised learning, which would involve training on the context { x i , f ( ∗ ) ( x ( ∗ ) i ) + ϵ ( ∗ ) i } n i =1 in order to predict f ( ∗ ) ( x ( ∗ ) n +1 ) , ICL happens entirely in a forward pass , so there is no training using labels from f ( ∗ ) . Our inquiry focuses on how ICL is facilitated by the softmax activation in the self-attention unit, which we introduce next.

The Softmax Attention Unit. We consider a single softmax attention head H SA ( · ; θ ) : R ( d +1) × ( n +1) → R ( d +1) × ( n +1) parameterized by θ :=( W K , W Q , W V ) , where W K , W Q , W V ∈ R ( d +1) × ( d +1) are known as key, query, and value weight matrices, respectively. Intuitively, for a sequence of tokens Z = [ z 1 , . . . , z n +1 ] ∈ z ( d +1) × ( n +1) , the attention layer creates a 'hash map" where the key-value pairs come from key and value embeddings of the input tokens, { W K z i : W V z i } . Each token z i is interpreted as a query W Q z i , and during a pass through the attention layer, this query is matched with the keys { W K z j } j to return an average over the associated values { W V z j } j with a weight determined by the quality of the match (proportional to e ( W K z j ) ⊤ ( W Q z i ) ). Specifically, H SA ( Z ; θ ) = [ h SA ( z 1 , Z ; θ ) , · · · , h SA ( z n +1 , Z ; θ )] , where

With slight abuse of notation, we denote h SA ( z j ) = h SA ( z j , Z ; θ ) when it is not ambiguous. To study how this architecture enables ICL, we follow [28] to formalize ICL as a regression problem. Below we define the tokenization, pretraining objective and evaluation task.

Tokenization for regression. The learning model encounters token sequences of the form

where the ground-truth labelling function f maps from R d to R and belongs to some class F , each ϵ i is mean-zero noise, and the i -th input feature vector x i ∈ R d is jointly embedded in the same token with its noisy label f ( x i ) + ϵ i ∈ R . We denote this token z i . The ICL task is to accurately predict this label given the n context tokens { ( x i , f ( x i ) + ϵ i ) } n i =1 , where f may vary across sequences. The prediction for the label of the ( n +1) -th feature vector is the ( d +1) -th element of h SA ( z n +1 ) [10], denoted h SA ( z n +1 ) d +1 . Ultimately, the goal is to learn weight matrices such that h SA ( z n +1 ) d +1 is likely to approximate the ( n +1) -th label on a random sequence Z .

Pretraining protocol. We study what softmax attention learns when its weight matrices are pretrained using sequences of the form of (1). These sequences are randomly generated as follows:

where D ( F ) is a distribution over functions in F , D x is a distribution over R d , and D ϵ is a distribution over R with mean zero and variance σ 2 . The token embedding sequence Z is then constructed as in (1). Given this generative model, the pretraining loss of the parameters θ = ( W Q , W K , W V ) is the expected squared difference between the prediction of softmax attention and the ground-truth label of the ( n +1) -th input feature vector in each sequence, namely

We next reparameterize the attention weights to make (3) more interpretable. For the last column of W V , we show in Appendix B that any minimizer of (3) in the settings we consider must have the first d elements of this last column equal to zero. We follow [7, 9, 10] by setting the first n columns of W V to zero. As in [10], we fix the ( d +1 , d +1) -th element of W V , here as 1 for simplicity. In the same vein, we follow [7, 10] by setting the ( d +1) -th row and column of W K and W Q equal to zero. To summarize, the reparameterized weights are:

where M K , M Q ∈ R d × d . Now, since our goal is to reveal properties of minimizers of the pretraining loss, rather than study the dynamics of optimizing the loss, without loss of generality we can define M := M ⊤ K M Q and re-define the pretraining loss (3) as a function of M . Doing so yields:

Interpretation of the pretraining loss. The loss (ICL) clarifies how softmax attention can be inter- preted as a nearest neighbors regressor. When x ⊤ i M x n +1 is a proxy for the distance between x i and x n +1 (which we formally show in Section 3 as happening under reasonable assumptions), the softmax attention prediction is a convex combination of the noisy labels with weights determined by the closeness of x i to x n +1 , such that the labels of points closer to x n +1 have larger weight. Moreover, the decay in weights on points further from x n +1 is exponential and controlled by M , which effectively defines a neighborhood, or attention window, of points around x n +1 whose labels have non-trivial weight. More formally, we can think of the attention window defined for a query x n +1 as the set AttnWindow ( x n +1 ; M ) := { x : x ⊤ M x n +1 = Ω(1) } . As we have observed in Figure 1, our key insight is that pretrained M scales this attention window with the Lipschitzness of the function class. Generally speaking, larger M entails averaging over a smaller window and incurring less bias due to the function values of distant tokens in the estimate, while smaller M entails averaging over a larger window, resulting in larger bias due to distant token labels, but a smaller noise variance. Figure 2 further depicts this tradeoff.

Connection to non-parametric estimation and the Nadaraya-Watson estimator. A nonparamet-

ric estimation technique to interpolate between known values of a function is to use a kernel estimator. The Nadaraya-Watson (NW) estimator [29-31] is one such estimator, and interpolates the data as

where K ( r ) = e -r 2 /h for some bandwidth h . In Section B.1 we show that optimizing the pretraining loss (ICL) reduces to meta-learning the bandwidth of an NW estimator on a distribution of pretraining tasks. However, to our knowledge, the literature has not determined the optimal bandwidth for the kernel, as there has been no analysis of non-asymptotic lower bounds on the loss, which we need to characterize the optimal solution. A close work to ours is [32], which considers regression on general L -Lipschitz tasks, but this analysis provides only a tight upper bound on the loss.

Remark 2.1 (Extreme cases) . Consider the following two settings.

(i) Constant functions. If each of the functions the attention unit sees in pretraining is constant, as in the Left column of Figure 1, it is best to consider an infinite attention window, that is, take M = 0 d × d as this results in a uniform average over all the noisy token labels.

(ii) Rapidly changing functions. If the pretraining functions change rapidly, as in the Right column of Figure 1, attending to a distant token might only corrupt the estimate at the target. For example suppose the input tokens are used to construct Voronoi cells on the surface of the hypersphere, and the label for a new token in a cell is the label of the token used to construct that cell. The optimal estimator attends only to the single nearest token since this incurs error only from label noise.

Remark 2.2 (Softmax advantage) . To further highlight the utility of the softmax, we compare with linear attention [7, 9, 11], whose estimator can be written as h LA ( x ) = ∑ i ( f ( x i ) + ϵ i ) x ⊤ i M x , up to a universal scaling due to the value embedding. This is again a weighted combination of labels, but one that does not allow for adapting an attention window - any scaling of M does not change the relative weights placed on each label - unlike softmax attention. Please see Figure 1 (Middle Row) for a comparison of the weights used in the different estimators.

3 Pretraining Learns Scale of Attention Window

One of our observations of the attention estimator h SA is that it computes a nearest neighbours regression. We hypothesize that the role of pretraining is to select a neighbourhood within which to select tokens for use in the estimator. In this section we characterize the radius of this neighborhood.

Definition 3.1 (Lipschitzness) . A function f : X → R has Lipschitzness L if L is the smallest number satisfying f ( x ) -f ( x ′ ) ≤ L ∥ x -x ′ ∥ for all ( x , x ′ ) ∈ X 2 .

The general requirement for the function classes to which our results apply is that the class should be invariant to isometries, each function should be Lipschitz, and the function value at two points should be less correlated as those points get further. These are written formally in Assumption B.4. To be concrete, we work with the following two function classes that satisfy these assumptions (this is shown in Lemmas C.3 and C.7) to derive explicit bounds.

Definition 3.2 (Affine and ReLU Function Classes) . The function classes F aff L and F + L are respectively defined as:

D ( F aff L ) , D ( F + L ) are induced by drawing w ∼ Σ U d and b, l, l 1 , l 2 i.i.d. ∼ Unif ([ -L, L ]) for some Σ ≻ 0 d × d . Note that the max Lipschitzness of any function in these classes is L , and ( z ) + := max( z, 0) .

Next, we make the following assumption, similar to [7], on the covariate distribution.

Assumption 3.3 (Covariate Distribution) . The covariate distribution satisfies D x = Σ -1 U d .

Now we are ready to state our main theorem that characterizes minimizers of (ICL).

Theorem 3.4. Let Assumption 3.3 hold and tasks f be drawn from (Case 1) D ( F aff L ) or (Case 2) D ( F + L ) . For n = Ω(1) and Ω( n -d/ 2 ) ≤ σ 2 ≤ O ( nL 2 ) , any minimizer of the pretraining loss (ICL) satisfies 2 M ∗ = w KQ Σ , where for Λ := nL 2 σ 2 , α := 1 d +4 and β := 1 d +2 :

Theorem 3.4 shows that optimizing the pretraining population loss in Equation (ICL) leads to attention key-query parameters that scale with the Lipschitzness of the function class, as well as the noise level and number of in-context samples. These bounds align with our observations from Figures 1 and 2 that softmax attention selects an attention window that shrinks with the function class Lipschitzness, recalling that larger w KQ results in a smaller window. Further, the dependencies of the bounds on σ 2 and n are also intuitive, since larger noise should encourage wider averaging to average out the noise, and larger n should encourage a smaller window since more samples makes it more likely that there are samples very close to the query. To our knowledge, this is the first result showing that softmax attention learns properties of the task distribution during pretraining that facilitate ICL .

Learning Lipschitzness is critical to generalization. We next give the following generalization result for downstream tasks.

Theorem 3.5. Suppose softmax attention is first pretrained on tasks drawn from D ( F + L ) and then tested on an arbitrary L -Lipschitz task, then the loss on the new task is upper bounded as L ≤

0

500

1000

1500

2000

Pretraining Iterations

0

20

40

60

80

100

||M||

ReLU -- Isotropic

L

0

0.8

1.6

2.4

3.2

0

500

1000

1500

2000

Pretraining Iterations

0

20

40

60

80

100

ReLU -- Non-Isotropic

0

500

1000

1500

2000

Pretraining Iterations

0

20

40

60

80

100

Cos -- Isotropic

0

500

1000

1500

2000

Pretraining Iterations

0

25

50

75

100

125

Cos -- Non-Isotropic

O ( L 2 Λ β ) . Furthermore, if the new task is instead drawn from D ( F + L ′ ) , the loss is lower bounded as L ≥ Ω( L ′ 2 Λ 2 β ) for L ′ > L and L ≥ Ω( Λ βd/ 2 n ) for L ′ < L .

Theorem 3.5 shows that pretraining on D ( F + L ) yields a model that can in-context learn downstream tasks if and only if they have similar Lipschitzness as L . Thus, learning Lipschitzness is both sufficient and necessary for ICL. If the evaluation task Lipschitzness is much larger than that seen in pretraining, the pretrained model will give highly biased estimates. Conversely, if the evaluation Lipschitzness is much lower, the pretrained model will not optimally average the label noise.

Necessity of Softmax. To further emphasize the importance of the softmax in Theorem 3.4, we next study the performance of an analogous model with the softmax removed. We consider linear self-attention [7, 9, 11], which replaces the softmax activation with an identity operation. In particular, in the in-context regression setting we study, the prediction of f ( x n +1 ) by linear attention and the corresponding pretraining loss are given by:

As discussed in Remark 2.1, h LA ( x n +1 ) cannot adapt an attention window to the problem setting. We show below that this leads it to large ICL loss when tasks are drawn from D ( F + L ) .

Theorem 3.6 (Lower Bound for Linear Attention) . Consider pretraining on L LA with tasks f drawn from D ( F + L ) and covariates drawn from U d . Then for all M ∈ R d × d , L LA ( M ) = Ω( L 2 ) .

This lower bound on L LA is strictly larger than the upper bound on L from Theorem 3.5, up to factors in d , as long as σ 2 n ≤ 1 , which holds in all reasonable cases. Please see Appendix F for the proof.

3.1 Proof Sketch

To highlight the key insights of our analysis, in this section we consider a modification of the softmax attention that exhibits important properties of the original. Note that this approximation is for illustration only; the above results use the original softmax attention - see Appendices C, D, E. For now, consider a function class F L := { f : f ( x ) = L w ⊤ x , w ∈ S d -1 } of linear functions.

(Temporary) modification of the softmax attention. Rather than averaging over every token with a weight that decays exponentially with distance, we consider a modification which uniformly averages all tokens within a distance specified by w KQ = ∥ M ∥ . From Lemma B.5, without loss of generality (WLOG) we can consider M = w KQ I d . This means that, ignoring normalization, the weight assigned to f ( x i ) by the true soft-max attention is e -w KQ ∥ x -x i ∥ 2 . That is, for all x i satisfying ∥ x -x i ∥ < 1 / √ w KQ , the assigned weights within a constant factor of each other. Meanwhile, for x i satisfying ∥ x -x i ∥ = √ c / √ w KQ for c > 1 , the weights are e -c , decaying exponentially in c . This motivates us to consider a 'modified softmax attention" given by h MSA ( x ) := ∑ i f ( x i ) ✶ i ∑ j ✶ j , where

The In-context Loss. The pretraining loss in Equation ICL can be decomposed as:

We first upper and lower bound each of these terms separately, starting with L signal ( w KQ ) .

Noiseless Estimator Bias. (Please see Appendix C) This term is the squared difference between an unweighted average of the token labels within a radius of x , and the true label. Take w KQ = Ω(1) . Then for large d , most of the points x i satisfying ∥ x -x i ∥ ≤ 1 / √ w KQ lie on the boundary of the cap, that is, ∥ x -x i ∥ < 1 / √ w KQ = ⇒ ∥ x -x i ∥ ≈ 1 / √ w KQ . This motivates us to approximate the set of points x i satisfying the above as coming from a uniform distribution over just the boundary of the cap. The center of mass of a ring of radius 1 / √ w KQ embedded on the surface of a hyper-sphere, is O ( 1 / w KQ ) from the boundary of a sphere, so the squared bias is Θ( L 2 / w 2 KQ ) .

Noise. (Please see Appendix D for details) Since the noise is independent across tokens, we can write L noise ( w KQ ) = σ 2 ∑ j ✶ j , which is related to the number of tokens found within a 1 / √ w KQ radius of x . In Lemma G.1, we derive bounds for the measure of this region. For now, we replace the sum in the denominator with its expectation. We can bound 1 ∑ j j = Θ ( w d 2 KQ / n ) as long as w KQ ≲ n 2 /d .

✶ Combining the L signal and L noise terms. (Please see Appendix E for details) Overall, we have L = L signal + L noise with L signal = Θ ( L 2 / w KQ ) and L noise = Θ ( w d 2 KQ σ 2 / n ) . Minimizing this sum reveals that the optimal w KQ satisfies w KQ = Θ ( ( nL 2 / σ 2 ) 2 d +2 ) .

3.2 Experiments

We next empirically verify our intuitions and results regarding learning the scale of the attention window. In all cases we use the Adam optimizer with one task sampled per round, use the noise distribution D ϵ = N (0 , σ 2 ) , and run 10 trials and plot means and standard deviations over these 10 trials. Please see Appendix J for full details as well as additional results.

Ablations over L , σ and n . We verify whether the relationship between the attention window scale i.e. ∥ M ∥ -1 - and L , σ and n matches our bounds in Theorem 3.4 for the case when tasks are drawn from D ( F + L ) and the covariates are drawn from U d , as well as whether these relationships generalize to additional function classes and covariate distributions. We train on tasks drawn from D ( F + L ) and D ( F cos L ) , where F cos L := { f : f ( x ) = cos ( L w ⊤ x ) , w ∈ S d -1 } and D ( F cos L ) is induced by sampling w ∼ U d . In all cases we set d = 5 , and use ( L, σ, n ) = (1 , 0 . 01 , 20) if not ablating over these parameters, and vary only one of { L, σ, n } and no other hyperparameters within each plot.

Attention window scales inversely with L . Figure 3 shows that ∥ M ∥ increases with L in various settings. In Figure 3(Left, Middle-Left), tasks are drawn from D ( F + L ) , and in Figure 3(Middle-Right, Right), they are drawn D ( F cos L ) . In Figure 3(Left, Middle-Right), each x i is drawn from U d , whereas in Figure 3(Middle-Left, Right), each x i is drawn from a non-isotropic distribution ˜ U d on S d -1 defined as follows. First, let S d := diag ([1 , . . . , d ]) ∈ R d × d , then x ∼ ˜ U d is generated by sampling ˆ x ∼ N ( 0 d , I d ) , then computing x = S 1 / 2 d ˆ x ∥ S 1 / 2 d ˆ x ∥ . Although larger L implies larger ∥∇ x f ( x ) ∥ on average across f , it is not clear that it implies larger ∥∇ M K L ( W ⊤ K W Q ) ∥ nor ∥∇ M Q L ( W ⊤ K W Q ) ∥ , so it is surprising that larger L implies larger pretrained M (although it is consistent with our results).

Attention window scales with σ , inversely with n . Figure 4 shows that the dependence of ∥ M ∥ on σ and n also aligns with Theorem 3.4. As expected, ∥ M ∥ increases slower during pretraining for larger σ (shown in Figures 4(Left, Middle-Left)), since more noise encourages more averaging over a larger window to cancel out the noise. Likewise, ∥ M ∥ increases faster during pretraining for larger n (shown in Figures 4(Middle-Right, Right)), since larger n increases the likelihood that there is a highly informative sample in a small attention window. Here always the covariate distribution is U d .

0

500

1000

1500

2000

Pretraining Iterations

0

1

2

3

||M||

ReLU -- Varying

0.0

0.2

0.6

1

3

10

0

500

1000

1500

2000

Pretraining Iterations

0.0

0.5

1.0

1.5

2.0

Cos -- Varying

0

500

1000

1500

2000

Pretraining Iterations

0

1

2

3

4

ReLU -- Varying

n

n

5

10

50

100

200

500

0

500

1000

1500

2000

Pretraining Iterations

0

1

2

3

4

5

Cos -- Varying

n

Learning new tasks in-context. An implication of our work is that for the function classes we consider, the softmax attention estimator does not adapt to the function class beyond its Lipschitzness . We have already seen in Figures 3 and 4 that the growth of ∥ M ∥ during pretraining is similar across different function classes with the same Lipschitzness, as long as σ and n are fixed. Here we verify the conclusion from Theorem 3.5 that for fixed n and σ , the necessary and sufficient condition for downstream generalization, measured by small ICL error, is that the pretraining and downstream tasks have similar Lipschitzness. Figure 5 supports this conclusion. Here we set d = 5 , n = 200 , σ = 0 . 01 and draw each x i i.i.d. from U d . In Figure 5(Left, Middle-Left, MiddleRight), we train three attention units on tasks drawn from the 1-Lipschitz affine ( D ( F aff 1 ) ), ReLU ( D ( F + 1 ) ), and cosine ( D ( F cos 1 ) ) task distributions. Each plot shows the test ICL error on tasks drawn from a distribution in { D ( F aff 1 ) , D ( F + 1 ) , D ( F cos 1 ) }. Performance is similar regardless of the pairing of pretraining and test distributions, as the Lipschitzness is the same in all cases, demonstrating that pretraining on tasks with appropriate Lipschitzness is sufficient for generalization .

Moreover, Figure 5(Right) shows that when the Lipschitzness of the pretraining tasks does not match that of the test tasks, ICL performance degrades sharply, even when the tasks otherwise share similar structure. Here the test task distribution is D ( F cos 1 ) , and the pretraining task distributions are D ( F aff 1 ) , D ( F cos 0 . 1 ) , and D ( F cos 10 ) . The only pretraining distribution that leads to downstream generalization is D ( F aff 1 ) since its Lipschitzness matches that of the downstream tasks, despite the fact that it is not a distribution over cosine functions, unlike the other distributions. Thus, these results lend credence to the idea that in addition to being sufficient, pretraining on tasks with appropriate Lipschitzness is necessary for generalization .

4 Softmax Attention Learns Direction of Attention Window

Thus far, we have considered distributions over tasks that treat the value of the input data in all directions within the ambient space as equally relevant to its label. However, in practice the ambient dimension of the input data is often much larger than its information content - the labels may change very little with many features of the data, meaning that such features are spurious. This is generally true of embedded language tokens, whose embedding dimension is typically far larger than the minimum dimension required to store them (logarithmic in the vocabulary size) [1]. Motivated by this, we define a notion of 'direction-wise Lipschitzness' of a function class to allow for analyzing classes that may depend on some directions within the ambient input data space more than others.

Definition 4.1 (Direction-wise Lipschitzness of Function Class) . The Lipschitzness of a function class F with domain X ⊆ R d in the direction w ∈ S d -1 is defined as as the largest Lipschitz constant of all functions in F over the domain X projected onto w , that is:

Using this definition, we analyze function classes consisting of linear functions with parameters lying in a subspace of R d , as follows:

Definition 4.2 (Low-rank Linear Function Class) . The function class F lin B is defined as F lin B := { f : f ( x ) = a ⊤ B ⊤ x , a ∈ R k } , and D ( F lin B ) is induced by drawing a ∼ U k .

where B ∈ O d × k is a column-wise orthonormal matrix. Since our motivation is settings with lowdimensional structure, we can think of k ≪ d . Let B ⊥ ∈ O d × ( d -k ) denote a matrix whose columns

0

250

500

750

1000

Pretraining Iterations

10

2

10

1

10 0

Test ICL Error

Test on Affine,

L

=1

Trained on (all

L

=1)

ReLU

Cos

0

250

500

750

1000

Pretraining Iterations

10

2

10

1

10 0

Test on ReLU,

L

=1

Trained on (all

L

=1)

ReLU

Cos

0

250

500

750

1000

Pretraining Iterations

10

2

10

1

10 0

Test on Cos,

L

=1

Trained on (all

L

=1)

ReLU

Cos

0

250

500

750

1000

Pretraining Iterations

10

2

10

1

10 0

Test on Cos,

L

=1

Trained on

Affine,

L

=1

Cos,

L

=0.1

Cos,

L

=10

form an orthonormal basis for the subspace perpendicular to col ( B ) , and note that the Lipschitzness of F lin B in the direction w is L if w ∈ col ( B ) and 0 if w ∈ col ( B ⊥ ) . Observe that any function in F lin B can be learned by projecting the input onto the non-zero Lipschitzness directions, i.e. col ( B ) , then solving a k ≪ d -dimensional regression. To formally study whether softmax attention recovers col ( B ) , we assume the covariates are generated as follows.

̸

Assumption 4.3 (Covariate Distribution) . There are fixed constants c u = 0 and -∞ < c v < ∞ s.t. sampling x i ∼ D x is equivalent to x i = c u Bu i + c v B ⊥ v i where u i ∼ U k and v i ∼ U d -k .

Assumption 4.3 entails that the data is generated by latent variables u i and v i that determine labelrelevant and spurious features. This may be interpreted as a continuous analogue of dictionary learning models studied in feature learning works [33, 34]. We require no finite upper bound on | c v | nor 1 | c u | , so the data may be dominated by spurious features.

Theorem 4.4. Let B ∈ O d × k and consider the pretraining population loss (ICL) with f ∼ D ( F lin B ) . Suppose Assumption 4.3 holds, as well as at least one of two cases: (Case 1) σ = 0 , or (Case 2) n = 2 . Then among all M ∈ M := { M ∈ R d × d : M = M ⊤ , ∥ B ⊤ MB ∥ ≤ 1 c 2 u } , the minimizer of the pretraining population loss (ICL) is M ∗ = c BB ⊤ for some c ∈ (0 , 1 c 2 u ] .

Theorem 4.4 shows that softmax attention can achieve dimensionality reduction during ICL on any downstream task that has non-zero Lipschitzness only in col ( B ) by removing the zero-Lipschitzness features while pretraining on F lin B . Removing the zero-Lipschitzness features entails that the nearest neighbor prediction of pretrained softmax attention uses a neighborhood, i.e. attention window, defined strictly by projections of the input onto col ( B ) . To our knowledge, this is the first result showing that softmax attention pretrained on ICL tasks recovers a shared low-dimensional structure among the tasks . Please see Appendix J for empirical results verifying that softmax attention indeed recovers low-dimensional structure, even for tasks consisting of (nonlinear) generalized linear models.

5 Conclusion

We have presented, to our knowledge, the first results showing that softmax attention learns shared structure among pretraining tasks that facilitates downstream ICL. Moreover, we have provided empirical evidence suggesting that our conclusions about what softmax attention learns during pretraining generalize to function classes beyond those considered in our analysis.

Limitations and Future Work. 1. The model we use in this work is an attempt to understand a phenomenon that emerges in LLMs, which is that the output of the model can be 'primed' with some examples provided in the context that resembles few-shot learning, even though they are only trained on next token prediction. Establishing a mathematical framework for this remains an interesting question. 2. We consider the output of a single layer of attention. Studying the nature of the solution when this is iterated over multiple trained layers is an interesting future prospect.

6 Acknowledgments

This work was supported in part by NSF Grants 2127697, 2019844, 2107037, and 2112471, ARO Grant W911NF2110226, ONR Grant N00014-19-1-2566, the Machine Learning Lab (MLL) at UT Austin, and the Wireless Networking and Communications Group (WNCG) Industrial Affiliates Program.

References

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.

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.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research , 24(240):1113, 2023.

Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies . Association for Computational Linguistics, 2022.

Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446 , 2021.

Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, HengTze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239 , 2022.

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, 2023.

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.

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

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.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

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.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.

Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.17086 , 2023.

Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196 , 2023.

Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning , pages 19565-19594. PMLR, 2023.

Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.08391 , 2023.

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.

Xinyi Wang, Wanrong Zhu, and William Yang Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. arXiv preprint arXiv:2301.11916 , 2023.

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 , 2023.

Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249 , 2023.

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.

Yingyi Chen, Qinghua Tao, Francesco Tonin, and Johan AK Suykens. Primal-attention: Self-attention through asymmetric kernel svd in primal representation. arXiv preprint arXiv:2305.19798 , 2023.

Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: a unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv:1908.11775 , 2019.

Tan Nguyen, Minh Pham, Tam Nguyen, Khai Nguyen, Stanley Osher, and Nhat Ho. Fourierformer: Transformer meets generalized fourier integral theorem. Advances in Neural Information Processing Systems , 35:29319-29335, 2022.

Xing Han, Tongzheng Ren, Tan Minh Nguyen, Khai Nguyen, Joydeep Ghosh, and Nhat Ho. Designing robust transformers using robust kernel density estimation. arXiv preprint arXiv:2210.05794 , 2022.

Yichuan Deng, Zhihang Li, and Zhao Song. Attention scheme inspired softmax regression, 2023.

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.

Kathryn A Prewitt. A distribution-free theory of nonparametric regression. laszlo gyorfi, michael kohler, adam krzyzak, and harro walk. Journal of the American Statistical Association , 98(464):1084-1084, 2003.

Elizbar A Nadaraya. On estimating regression. Theory of Probability & Its Applications , 9(1):141-142, 1964.

Geoffrey S Watson. Smooth regression analysis. Sankhy¯ a: The Indian Journal of Statistics, Series A , pages 359-372, 1964.

Samuele Tosatto, Riad Akrour, and Jan Peters. An upper bound of the bias of nadaraya-watson kernel regression under lipschitz assumptions. Stats , 4(1):1-17, 2021.

Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning , pages 11112-11122. PMLR, 2021.

Zhenmei Shi, Junyi Wei, and Yingyu Liang. A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. arXiv preprint arXiv:2206.01717 , 2022.

Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. arXiv preprint arXiv:2306.15063 , 2023.

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 , 2023.

Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta optimizers. arXiv preprint arXiv:2212.10559 , 2022.

Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Do pretrained transformers really learn in-context by gradient descent? arXiv preprint arXiv:2310.08540 , 2023.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895 , 2022.

Asher Trockman and J. Zico Kolter. Mimetic initialization of self-attention layers, 2023.

Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:2310.00535 , 2023.

Enric Boix-Adsera, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind. Transformers learn through gradual rank increase, 2023.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. arXiv preprint arXiv:2303.04245 , 2023.

Samy Jelassi, Michael E. Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure, 2022.

Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015 , 2023.

Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308.16898 , 2023.

Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023.

Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. arXiv preprint arXiv:2310.08566 , 2023.

Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. arXiv preprint arXiv:2306.02896 , 2023.

Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representations. arXiv preprint arXiv:2310.10616 , 2023.

Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749 , 2022.

Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. arXiv preprint arXiv:2307.11353 , 2023.

Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research , 22(1):3463-3497, 2021.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 , 2019.

Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009.11264 , 2020.

Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint arXiv:2106.03764 , 2021.

Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems , 35:12071-12083, 2022.

Zhao Song, Guangyi Xu, and Junze Yin. The expressibility of polynomial based attention scheme, 2023.

Kevin Christian Wibisono and Yixin Wang. On the role of unstructured training data in transformers' in-context learning capabilities. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning , 2023.

Yichuan Deng, Zhao Song, and Tianyi Zhou. Superiority of softmax: Unveiling the performance edge over linear attention. arXiv preprint arXiv:2310.11685 , 2023.

James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. A mathematical theory of attention, 2020.

Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention, 2020.

Herbert Robbins. A remark on stirling's formula. The American Mathematical Monthly , 62(1):26-29, 1955.

Brian Knaeble. Variations on the projective central limit theorem. https://arxiv.org/pdf/0904.1048.pdf , 2015.

Elliott H Lieb and Michael Loss. Analysis , volume 14. American Mathematical Soc., 2001.

G.H. Hardy, J.E. Littlewood, and G. Pólya. Inequalities . Cambridge Mathematical Library. Cambridge University Press, 1952.

Contents

A Additional Related Work

Empirical study of ICL. Several works have studied ICL of linear tasks in the framework introduced by [28], and demonstrated that pretrained transformers can mimic the behavior of gradient descent [11-13, 28], Newton's method [14], and certain algorithm selection approaches [13, 16]. [35] studied the same linear setting with the goal of understanding the role of pretraining task diversity, while [36] argued via experiments on general auto-regressive tasks that ICL implicitly constructs a learning objective and optimizes it within one forward pass. Other empirical works have both directly supported [37] and contradicted [38] the hypothesis that ICL is a gradient-based optimization algorithm via experiments on real ICL tasks, while [39] empirically concluded that induction heads with softmax attention are the key mechanism that enables ICL in transformers. Lastly, outside of the context of ICL, [40] noticed that the attention parameter matrices of trained transformers are often close to scaled identities in practice, consistent with our findings on the importance of learning a scale to softmax attention training.

Transformer training dynamics. [21] and [41] studied the dynamics of softmax attention trained with gradient descent, but assumed orthonormal input features and either linear tasks [21] or that the softmax normalization is a fixed constant [41]. [42] proved that softmax attention with diagonal weight matrices incrementally learns features during gradient-based training. Other work has shown that trained transformers can learn topic structure [43], spatial structure [44], visual features [45] and support vectors [46, 47] in specific settings disjoint from ICL.

Expressivity of transformers. Multiple works have shown that transformers with linear [11, 36], ReLU [13, 14, 48], and softmax [12, 15] attention are expressive enough to implement generalpurpose machine learning algorithms during ICL, including gradient descent. A series of works have shown the existence of transformers that recover sparse functions of the input data [49-52]. [53] studied the statistical complexity the learning capabilities of attention with random weights. More broadly, [54-59] have analyzed various aspects of the expressivity of transformers.

Other studies of softmax attention. [60] hypothesized that the role of the softmax in attention is to facilitate a mixture-of-experts algorithm amenable to unstructured training data. [27] formulated a softmax regression problem and analyzed the convergence of a stylized algorithm to solve it. [22] showed that in a setting with ICL regression tasks a la [28], a kernel regressor akin to softmax attention with M equal to the inverse covariance of x converges to the Bayes posterior for a new ICL task - in this setting the conditional distribution of the label given the query and n labelled context samples - polynomially with the number of context samples, but did not study what softmax attention learns during pretraining. [61] also compared softmax and linear attention, but focused on softmax's greater capacity to separate data from two classes. [62] and [63] investigate the Lipschitz constant of attention rather than what attention learns.

Non-parametric regression. Our results imply that pretraining softmax attention reduces to the problem of meta-learning the bandwidth of a Nadaraya-Watson estimator with a Gaussian kernel. However, to our knowledge, the non-parametric regression literature has not addressed this problem. The closest work is [32], which only upper bounds the noiseless loss, and only in the limit n →∞ , whereas our result characterizes the optimal bandwidth, which requires upper and lower bounds on the noisy loss.

B Preliminaries

We first justify our claim that the first d rows of the last column of W V can be set to 0 d for any optimal choice of parameters.

Lemma B.1. If under the function distribution, a function f is equally likely as likely as -f , then any optimal solution to L ( W V , W K , W Q ) in 3 satisfies W V = ( 0 d × d 0 d × 1 0 1 × d c ) .

Proof. For readability we write β i = e -w KQ ∥ x i -x n +1 ∥ 2 ∑ j e -w KQ ∥ x j -x n +1 ∥ 2 Suppose W V = ( 0 d × d v 0 1 × d c ) was optimal, then the loss can be written

But because f and -f are equally likely, and because the noise is also symmetric about 0, we can write this as

We can couple the noise { ϵ i } and the data { x i } in the two summands above to write this as

where A = ∑ i cf ( x i ) β i -f ( x ) = -( ∑ i c ( -f )( x i ) β i ) , B = ∑ i v ⊤ x i β i , and C = ∑ i cϵ i β i . We can set B = 0 simply by setting v = 0 d × 1 , and this has loss

In all of the distributions over functions we consider for pretraining, f is equally likely as -f , so without loss of generality we set all elements of W V besides the ( d +1 , d +1) -th to 0. For simplicity, we set the ( d +1 , d +1) -th element to 1.

Assumption B.2 (Covariate Distribution) . For each token x , first we draw ˜ x as ˜ x ∼ U d . Then x is constructed as x = Σ 1 / 2 ˜ x .

Definition B.3 (Linear and 2-ReLU Function Classes) . The function classes F lin L and F + L are respectively defined as:

D ( F lin L ) , D ( F + L ) are induced by drawing w ∼ N ( 0 , Σ -1 ) and b, l, l 1 , l 2 ∼ Unif ([ -L, L ]) . We say that these classes are L -Lipschitz, because the maximum Lipschitz constant for any function in the class is L .

Note that because ∥ Σ -1 / 2 x i ∥ = 1 always, we have

2 x i M x n +1

Let M ′ = Σ 1 / 2 MΣ 1 / 2 . This means the attention estimator can be rewritten as

So the attention a token x n +1 places on another x i is related to the distance between

M ′ Σ -1 / 2 x n +1 and Σ -1 / 2 x i . It is natural to suppose under some symmetry conditions that M ′ is best chosen to be a scaled identity matrix so that the attention actually relates to a distance between tokens. Below we discus sufficient conditions for this.

!

"!

!′

Assumption B.4. The function class F and distribution D ( F ) satisfy

E f ∼ D ( F ) [ f ( x ) f ( y )] = ρ ( x ⊤ y ) ∀ x , y ∈ X 2 , for some monotonically increasing ρ .

For any isometry ϕ preserving the unit sphere, and f ∈ F , we have f ◦ ϕ ∈ F .

Lemma B.5. Under Assumption B.4, any minimizer of Equation ICL satisfies M ∗ = w KQ Σ -1 for some scalar w KQ ≥ 0 .

̸

Proof. Let { y i } = { Σ -1 / 2 x i } . Suppose M y n +1 = c y n +1 for any c > 0 for some y n +1 . Take c y n +1 = ∥ M y n +1 ∥ and y ′ n +1 = M y n +1 c y n +1 (the projection of y onto the sphere). Consider a function ω : R d → R d satisfying ω ( y n +1 ) = c y n +1 y n +1 . Note that this need not be linear. Let ϕ denote a rotation that sends y ′ n +1 to y n +1 .

We show that L ( M ) > L ( ω ) , that is, it is favorable to not rotate y n +1 . We have

Lets compare this with the loss of ω . For a depiction of this, please see Figure 6

There are three terms to compare. The first in each is identical. The second is also the same:

rotational symmetry of { y i } , y n +1

rotational symmetry of { y i }

The third takes some more work. For any choice of { y i } , let

We see that α y n +1 , { y i } ( y ∗ ) varies monotonically with y ⊤ n +1 y ∗ for all y n +1 , { y i } . That is, y ⊤ ∗ y n +1 > y ′⊤ ∗ y n +1 = ⇒ α y n +1 , { y i } ( y ∗ ) > α y n +1 , { y i } ( y ′ ∗ ) ,

Similarly, we have

Critically, for a given y n +1 , α y , { y i } ( y ∗ ) can be re-parameterized as

α y n +1 , { y i } ( y ∗ ) = α ′ { y i } ( y ∗ -y n +1 ) where α ′ { y i } is symmetric about 0 and decreasing. Similarly, ρ ( y ⊤ n +1 y ∗ ) can be re-parameterized as ρ ( y ⊤ n +1 y ∗ ) = ρ ′ ( y ∗ -y n +1 ) where α ′ , ρ ′ are symmetric decreasing rearrangement (that is, the set of points z such that ρ ( x ) > r is a ball about the origin). From Lemma I.2 we then have

So L ( ω ) < L ( M ) . Let

Observe that L ( ω ) = E y n +1 q ( c y n +1 ) . We might as well set ω to be such that c y n +1 is the same for all y n +1 and a minimizer of q , so we have ω ( y n +1 ) = c y n +1 for all y n +1 which implies ω = c I d for some c . Because the optimal M ′ is identity, the corresponding optimal M is Σ -1 .

B.1 Rewriting the Loss

As a result of this, we can take M = w KQ Σ -1 and write the attention estimator as

This allows us to make the transformation X → Σ -1 / 2 X . This has the effect of making both the data covariance and the induced function class covariance equal to the identity. Essentially, WLOG we will henceforth consider Σ = I d . Henceforth, the estimator will be taken to be

and the loss will be parameterized by w KQ as

Because the noise ϵ i is independent of everything else, we can decompose this into two terms, a signal term and a noise term as follows

We bound the first term in Appendix C and the second in Appendix D. A useful function that we bound in Lemma G.4 and Corrolary G.5 in Appendix G is

We will use this function, particularly for p = 0 and 1 .

C The Signal Term

The purpose of this section of the Appendix is to obtain upper and lower bounds on L signal ( w KQ ) . Because we work with two different distributions over functions, and because the bounds depend on the distributions, we will make the distribution explicit in the argument to the function

As a reminder, we consider the following two distributions over functions. Please see section B.1 to see why we have set the covariance of w to be identity.

Definition C.1 (Affine and 2-ReLU Function Classes) . The function classes F aff L and F + L are respectively defined as:

D ( F aff L ) , D ( F + L ) are induced by taking w ∼ U d , b, l, l 1 , l 2 ∼ Unif [ -L, L ] .

First we have the following trivial bound on L signal ( w KQ ) .

Lemma C.2. For all w KQ we have L signal ( w KQ ) ≤ 4 L 2 .

C.1 Affine functions

Here we consider the affine function class F aff L . First, we note that this class satisfies Assumption B.4.

Lemma C.3. The affine class F aff L in Definition 3.2 satisfies Assumption B.4.

Because b is independent of w , we have

w is isotropic, so ϕ ( w ) is also supported by the distribution on w .

Lemma C.4. For affine functions, the signal term is upper bounded as

Proof. In the interest of readability, we will denote x n +1 as x . Consider ˜ x such that ˜ x = ∑ i x i e -2 w KQ x ⊤ i x ∑ j e -2 w KQ x ⊤ i x . Then our loss is given by E [ l 2 w ⊤ ( x -˜ x ) ] 2 . First, since w is indepen- dent of x , { x i } , we have E l 2 ( w ⊤ ( x -˜ x ) ) 2 = E l 2 ww ⊤ ( x -˜ x )( x -˜ x ) ⊤ , Now w has a uniformly randomly chosen direction, so its covariance is a multiple of the identity. We have E Tr ( ww ⊤ ) = E ∥ w ∥ 2 = L 2 3 , so E l 2 ww ⊤ = L 2 3 d I d . Continuing, E ( w ⊤ ( x -˜ x ) ) 2 = L 2 3 d E ∥ x -˜ x ∥ 2 . Take any x ′ ⊥ x , we have

Decomposing ˜ x into an orthogonal and a parallel component, we have E ∥ x -˜ x ∥ 2 = E ∥ x -xx ⊤ ˜ x -x ′ x ′⊤ ˜ x ∥ 2 for some x ′ ⊥ x with ∥ x ′ ∥ = 1 . But

Case 1: w KQ ≥ d + √ d 2 .

Consider first the term E ∥ x (1 -x ⊤ ˜ x ) ∥ 2 = E (1 -x ⊤ ˜ x ) 2 . Here we have with probability 1 -1 n

The other term E ∥ x ′ x ′⊤ ˜ x ∥ 2 = E ( x ′⊤ ˜ x ) 2 is the component of the bias in the direction orthogonal to x .

Popoviciu's Variance inequality

With probability 1 -1 n , when w KQ ≥ d + √ d we have

Putting together Equations 11 and 12, we have with probability 1 -1 n ,

The signal bias is upper bounded by 4 L 2 always (Lemma C.2). The overall upper-bound on the expectation is

Case 2: w KQ < d + √ d 2 . We always have L ( w KQ ) ≤ 4 L 2

from Lemma C.2.

Lemma C.5. For affine functions, the signal term is lower bounded as

Proof. Similar to Equation (10), for ˜ x = ∑ i x i e -2 w KQ x ⊤ i x ∑ j e -2 w KQ x ⊤ i x , we have

Now consider the term 1 -x ⊤ ˜ x . We have

Case 1: w KQ ≥ d + √ d 2 . Here we have from Corollary G.5, with probability 1 -1 /n

With probability 1 /n ≤ 1 2 the lowest we can have is L signal ( w KQ ) = 0 , so overall we have

Case 2: d + √ d 4 ≤ w KQ ≤ d + √ d 2 . From Corollary G.5, with probability 1 -1 n

With probability 1 /n ≤ 1 2 the lowest we can have is L signal ( w KQ ; D ( F aff L )) = 0 , so overall we have

Case 3: d + √ d 4 > w KQ . From Corollary G.5, with probability 1 -1 n

With probability 1 /n ≤ 1 2 the lowest we can have is L signal ( w KQ ; D ( F aff L )) = 0 , so overall we have

Corollary C.6. Combining the above, we have

We can now perturb these bounds in the case ofthe ReLU-based function class F + L .

C.2 ReLU-based functions

Consider the function class

where ReLU ( z ) := ( z ) + := max( z, 0) . Consider a distributions on F + L , namely D ( F + L ) . Let D ( F + L ) be induced by w ∼ U d , b, l 1 , l 2 ∼ Unif [ -L, L ] . That is, a vector w is drawn uniformly on the unit hypersphere. Then two norms are selected, l 1 , l 2 , and the overall function is given by

so that it follows one affine rule in one halfspace, and another affine rule in the opposite halfspace. Please see section B.1 to see why we have set the covariance of w to be identity.

Lemma C.7. The class F + L and distribution D ( F + L ) defined above satisfy Assumption B.4.

Proof. 1. Each function is defined as being piece-wise L -Lipschitz, and it is continuous, so it is also L -Lipschitz overall.

With probability 1 -2 arccos ( x ⊤ y ) π the points x and y are such that ( w ⊤ x )( w ⊤ y ) < 0 (that is, they are on opposite sides of the hyperplane defining the two pieces of the ReLU). Because the bias b is independent of the other parameters, we have as in the proof of Lemma C.3

Let x = x ∥ x ∥ for any vector x . Consider a re-parameterization of the pair ( x , y ) as ξ θ ( x , y ) → ( x + y , x -y ) . Because x and y are on the unit sphere, this is a bijection as

That is, for any x , y , ξ -1 x ⊤ y ( ξ x ⊤ y ( x , y )) = ( x , y ) . The push-forward of ξ is also uniform, that is for x , y satisfying x ⊤ y = θ , ξ θ ( x , y ) is distributed as U d ×U d -1 . For any x , y , let ξ -1 θ ( x , y ) = ( x θ , y θ ) . Then we have E f [ f ( x θ ) f ( y θ )] is a decreasing function of θ . Finally, for θ ≤ θ ′ , L 2 x ⊤ θ ww ⊤ y θ > L 2 x ⊤ θ ′ ww ⊤ y θ ′ so x ⊤ θ ww ⊤ y θ < 0 = ⇒ x ⊤ θ ′ ww ⊤ y θ ′ < 0 . The product of two positive increasing functions is itself non-increasing. Since we have both E w [ L 2 x ⊤ ww ⊤ y ∣ ∣ x ⊤ ww ⊤ y > 0] and 2 arccos ( x ⊤ y ) π are increasing functions of x ⊤ y , we also have

is an increasing function of x ⊤ y since E w [ L 2 x ⊤ ww ⊤ y ∣ ∣ x ⊤ ww ⊤ y > 0] ≥ 0 and ( 2 arccos ( x ⊤ y ) π ) ≥ 0 .

w is distributed uniformly on the hypersphere, so ϕ ( w ) is also also distributed uniformly on the hypersphere for any isometry ϕ that preserves the origin.

Lemma C.8. The signal term is upper bounded as

Proof. We have

With probability 1 -1 n , when w KQ ≥ d + √ d 2 we have

We always have L signal ( w KQ ) ≤ 4 L 2 from Lemma C.2. So the overall upper bound is

For w KQ ≥ d + √ d 2 , as before, we always have L signal ( w KQ ; D ) ≤ 4 L 2 .

Lemma C.9. The signal term is lower bounded as

Proof. Again for readability we will write x n +1 as x . For any f ∈ F + L let f x , aff denote the corresponding affine function that is equal to f in the halfspace containing x , that is if f ( x ′ ) = l 1 ReLU ( w ⊤ x ′ ) + l 2 ReLU ( -w ⊤ x ′ ) + b , and WLOG w ⊤ x ′ > 0 , then f x , aff ( x ′ ) = l 1 w ⊤ x ′ + b . Note that f x , aff comes from a w selected from the unit sphere and b, l ∈ [ -L, L ] exactly as f ∼

D ( F L ) , so it is actually statistically indistinguishable from a sample from D ( F aff L ) , the distribution over affine functions in Definition 3.2 (and the object of Lemma C.5). The error of the nonlinear estimator can be written as

where γ i = e -w KQ ∥ x -x i ∥ 2 Σ -1 ∑ e -w KQ ∥ x -x j ∥ 2 Σ -1 Let us compare the two errors due to the two functions. Let

j

A = { i : ( x ⊤ i w )( x ⊤ w ) < 0 } denote the set of points on the opposite side to x of the hyperplane defining the function.

+

Here the third equality holds because f ( x i ) is independent of f x , aff ( x j ) if i ∈ A,j ̸∈ A .

Let q = E f, x , { x i } (∑ i ∈ A f ( x i ) γ i ) ) 2 = E f, x , { x i } (∑ i ∈ A f x , aff ( x i ) γ i ) ) 2 . Then from the above we have

which has minimum at q = L ( w ; D ( F )) / 2

signal KQ L , completing the proof.

D Bounds on Noise Variance

In this section we obtain upper and lower bounds on the variance of the estimator due to label noise. There are three relevant parameters: d , the ambient dimension of the data; w KQ , the scaling induced

by the attention layer; and n , the number of tokens. Recall that the noise term is

Because the ϵ i are independent, this can further be simplified as

Lemma D.1. The noise term is bounded for d + √ d ≤ w KQ ≤ ( n 45 √ d log n ) 2 d as

Proof. We have

Using Lemma G.5, we have with probability at least 1 -1 n

Finally, in the worst case, we have 0 ≤ L noise ( w KQ ) ≤ 1 .

Finally, we show that the noise term is monotonic in w KQ .

Lemma D.2. L noise ( w ) > L noise ( w ′ ) ⇐⇒ w > w ′

Proof. Let a i = e -w ′ ∥ x i -x n +1 ∥ 2 , b i = e -( w -w ′ ) ∥ x i -x n +1 ∥ 2 . The result follows from Lemma I.3 because { a i } and { b i } satisfy a i > a j ⇐⇒ b i > b j ⇐⇒ ∥ x i -x n +1 ∥ < ∥ x j -x n +1 ∥ .

E Optimizing the Loss

For the nonlinear function class F + L , we have the following.

Theorem E.1. Suppose the functions seen in pretraining are drawn from D ( F + L ) as in Definition 3.2, the covariates are drawn as Assumption 3.3, n = Ω ( L log n σ ) d and n 2 d +2 = Ω(1) , then the optimal M satisfies

and similarly where w KQ satisfies

bias upper bound

bias lower bound

noise

overall upper bound

overall lower bound

Proof. We consider three regions in which the optimal value could potentially lie and see that only the third region is viable.

√

Case 1. w KQ ≤ d + d : In this case, the signal term lower bounds the optimal loss by Lemma C.5 as Ω(1) .

Case 2. w KQ > Ω ( n log n ) 2 d . In this case, the noise term lower bounds the optimal loss. From Lemma

D.2 we know that the noise term is non-decreasing in w KQ so in the range w KQ > Ω( ( n log n ) 2 d is lower bounded by L noise ( w KQ ) at w KQ = Ω( ( n log n ) 2 d , which is Ω ( σ 2 log n ) .

Case 3. d + √ d ≤ w KQ ≤ Ω ( n log n ) 2 d By combining Lemmas C.8, C.9, and D.1 , we obtain the following overall bound on the loss:

for some constants c, c that only depend on d . In the range w KQ ≥ d + d , we have w KQ > 1 and w KQ ≤ n , so the upper bound can be relaxed as L noise ( w KQ ) ≤ 2 c ( L 2 n + σ 2 w d 2 KQ n ) , which is minimized at w KQ = ( nL 2 σ 2 d ) 2 d +2 . Here it is upper bounded by 4 c ( dL d σ 2 n ) 2 d +2 . We note first of all that for large enough n (as long as n = Ω ( σ log n L ) d and n 2 d +2 = Ω(1) ) this is lower than the lower bounds we got in Case 1 and Case 2 , so this is indeed the region of global optimal solution. From d

Lemma C.9 we have L noise ( w KQ ) ≥ L 2 w 2 KQ + σ 2 w 2 KQ n ≥ L 2 w 2 KQ which gives

for the upper bound, we similarly also have L noise ( w KQ ) ≥ L 2 w 2 KQ + σ 2 w d 2 KQ n ≥ σ 2 w d 2 KQ n which gives

where w KQ satisfies

Of course, for this to not be vacuous we need

We will again hide constants that depend only on d and write this as

which is true as long as n > ( L log n σ ) d .

For the affine function class F aff L , we have the following

Theorem E.2. If the functions seen in pretraining are drawn from D ( F aff L ) as in Definition 3.2, and the noise variance σ 2 and Liphscitz constant L satisfies n ≥ ( L log 2 n σ ) d +2 , and n 2 d ≥ Ω(1) , and the covariates are drawn as Assumption 3.3, the optimal M satisfies

Proof. Again we work with three cases.

Case 1. w KQ ≤ d + √ d . Again in this case we have a lower bound to the signal term of Ω(1) .

Case 2. w KQ ≥ Ω ( n log n ) 2 d . Again we have a lower bound of Ω ( σ 2 log n )

Case 3. d + √ d ≤ w KQ ≤ Ω ( n log n ) 2 d Combining Lemmas C.4, C.5, D.1 is

We will minimize the upper bound. First suppose L 2 σ 2 ≥ w KQ for the w KQ that minimizes the upper bound. Then we have

This upper bound is minimized at w KQ = n 2 d +2 . However, this contradicts the constraint that w KQ ≤ L 2 σ 2 , when n 2 d +2 ≥ L 2 σ 2 , as we assume. So we have w KQ ≥ L 2 σ 2 for the minimizer. This means the upper bound is no more than

This upper bound is minimized at w KQ = ( nL 2 σ 2 d ) 2 d +4 where it is upper bounded by

whenever n ≥ L 2 σ 2 . We see that

for the upper bound, we similarly also have

Of course, for this to not be vacuous we need

We will again hide constants that depend only on d and write this as

which again is true as long as n = Ω ( L log 2 n σ ) d +2

E.1 Generalization Bounds

We conclude this section with a proof of the generalization error on a new L -Lipschitz task.

Theorem E.3. Suppose our attention is first pretrained on tasks drawn from D ( F + L ) and then tested on an arbitrary L -Lipschitz task, then the loss on the new task is upper bounded as L ≤ O ( L 2 Λ β ) . Furthermore, if the new task is instead drawn from D ( F + L ′ ) , the loss is lower bounded as L ≥ min { Ω( L ′ 2 Λ 2 β ) , Ω( Λ βd/ 2 n ) }

Proof. We know from Theorem E.2 that Ω(Λ β ) ≤ w KQ ≤ O (Λ 2 β ) . The upper bound for L ( w KQ ) , which is O ( L 2 w KQ + w d 2 KQ n ) , is a convex function for d ≥ 2 , so in any range it attains its maximum value at the extreme points. We can check the cases to see that this is O (max { L 2 Λ β + Λ dβ/ 2 n , L 2 Λ 2 β + Λ dβ n } ) = O ( L 2 Λ β + Λ dβ/ 2 n + L 2 Λ 2 β + Λ dβ n ) = O ( L 2 Λ β ) for large enough n .

Now consider testing on a new task from D ( F + L ′ ) . The ICL loss for Ω ( Λ β ) ≤ w KQ ≤ O ( Λ 2 β ) is bounded below as Ω( L ′ 2 Λ 2 β ) and Ω( Λ βd/ 2 n ) .

The implication of this is that if L ′ ≫ L , the error scales as ( L ′ ) 2 rather than ( L ′ ) 2 d d +2 while for L ′ ≪ L , the error is lower bounded by a constant.

F Lower Bound for Linear Attention

In this section we prove Theorem 3.6.

Lemma F.1. Consider the function distributions D ( F L ) and D ( F + L ) described in Definition 3.2. We have L LA ≥ Ω( L 2 ) , that is, the ICL error is lower bounded as Ω( L 2 ) .

Proof. We start by decomposing the ICL loss into a bias dependent term and a cenetered term. For f ∈ F L ∈ {F aff L , F + L } , let f denote the centered function f -E x f . Let f ′ denote the flip of f about its expected value, so f ′ = E x f -f . We observe that f is independent of E x f . For linear attention, we have, for f ∼ D ( F L )

By symmetry, this is also equal to the same expression using f ′ instead of f , since f and f ′ are distributed identically. Besides, E x f = E x f ′ and ϵ is symmetric about the origin, so

Let A = ∑ n i =1 ( f ( x i ) x ⊤ i M x n +1 + ϵ i x ⊤ i M x n +1 ) -f ( x n +1 ) and B = E x f . Then we see that L LA ( M ) ≥ 1 2 E ( A + B ) 2 + 1 2 E ( -A + B ) 2 = E A 2 + E B 2 . Meanwhile, E ( E x f ) 2 is just the variance of the signal term in D ( F aff L ) or D ( F + L ) , which is L 2 3 . So L LA ( M ) ≥ L 2 3

G Bounds for g p ( r )

The purpose of this section is to obtain upper and lower bounds on

for p = 0 , 1 / 2 , 1 . For this, we will need high probability upper and lower bounds on the number of points in a spherical cap under a uniform distribution over the hypersphere. Consider n points { x i } drawn uniformly from σ d -1 , the uniform measure over S d -1 , the d -dimensional hypersphere. The measure of the ϵ -spherical cap around x ∈ S d -1 , C ( ϵ, x ) = { x ′ : x ′⊤ x > 1 -ϵ } is denoted by σ ϵ .

G.1 Bounds on Spherical Caps

1

-

i

r

1

-

i

+1

r

Lemma G.1. The area of the spherical cap C ( ϵ ) , σ ϵ is bounded as

Proof. We derive a lower bound as follows. We replace the surface area of a spherical cap in S d -1 with a d -1 dimensional ball of the same boundary. Let V d denote the volume of a d dimensional ball (that is, V 3 ( r ) = 4 3 πr 3 ), and let A d denote the surface area of a d dimensional sphere (so A 3 ( a ) = 4 πr 2 ). It is known that

Then we have

The upper bound is similar. This time we replace the cap with the surface of a hemisphere with the same boundary. We have

We will also need upper and lower bounds on a discretized version of the incomplete gamma function. Definition G.2. Denote by γ ( d, α, m ) the expression γ ( d, α, m ) = ∑ m i d e -αi .

We have the following

Lemma G.3. For d > 5 , 1 ≤ α ≤ 2 , the incomplete Gamma function is bounded as

Proof. We compare with the Gamma function

Note that ∫ ∞ 0 t d e -αt dt = 1 α d +1 ∫ ∞ 0 t d e -t dt = 1 α d +1 Γ( d + 1) . Because the function t d e -αt is uni-modal with maximum ( d αe ) d , we have from Lemma I.1

Now suppose m ≥ d + √ d α . Then we have

the first inequality follows because d + √ d α ≤ m , the second follows because 2 d +1 2 d ≥ 2 d + j +1 2 d + j , the last follows because ( 1 + √ d dα ) d ≤ e √ d α and 1 1 -e x -x ≤ 2 x for x ≤ 2 . Over all, we have

While for the upper bound we have

Finally, we use Lemma G.3, specifically that ( d αe ) d ≤ 1 α d +1 √ 2 πd ( d e ) d ≤ Γ( d +1) α d +1 to yield the desired result.

i =1

For m< d + √ d α , we have from Lemma G.7 that m d e -αm ≥ 1 √ e i d e -αi so

and

G.2 Bounds on g p ( r )

Lemma G.4. Suppose { x i } are drawn independently and uniformly from the unit hypersphere. For n log n ≥ 45 √ dr d 2 , n > 5 , d > 2 , p ≤ 2 , we have g p ( r ) = ∑ n i =1 ∥ x i -x ∥ p e -r ∥ x ⊤ i -x ∥ 2 satisfies

(1 -e p 2 -2 ) n 2 p 2 √ 8 e 4 πd ( 1 r ) d 2 + p 2 γ ( d 2 + p 2 , 2 , r ) ≤ g p ( r ) ≤ 3 n ( 2 r ) d 2 + p 2 γ ( d 2 + p 2 , 2 , r ) with probability at least 1 -1 2 n

Proof. For 0 ≤ i ≤ r let N i denote the number, and S i denote the set, of points satisfying 1 -i r ≤ x ⊤ i x ⇐⇒ ∥ x i -x ∥ ≤ ( 2 i r ) 1 2 . Also denote by N -1 the points satisfying x ⊤ i x < 0 , and let S -1 denote this set. Note that

Similarly,

Note that because N i > 0 ,

And similarly,

and so

By a Chernoff bound for Binomial random variables, we have with probability 1 -r n 2 :

N -1 ≤ n r and

Whenever and

Over all we have with probability 1 -n 2

We always have for p ≤ 2

So at last, we have

We obtain a lower bound in the same way.

n 2 n log n

It will be useful to simplify this bound in regimes that we are interested in

Corollary G.5. Suppose { x i } are drawn independently and uniformly from the unit hypersphere. For n log n ≥ 45 √ dr d 2 , n > 5 , p ≤ 2 ≤ d , we have g p ( r ) = ∑ n i =1 ∥ x i -x ∥ p e -r ∥ x ⊤ i -x ∥ 2 satisfies with probability 1 -1 2 n

The following bounds are known for the Gamma function.

Lemma G.6. The Gamma function satisfies

Proof. 1. Please see [64].

Please see [65].

Lemma G.7. The following inequality holds:

Proof. Take the logarithm of both sides, we have that this is equivalent to

A Taylor series expansion of log(1 + x ) demonstrates that log ( 1 + 1 √ d ) = ∑ i ( -1) i +1 1 i √ d i . For d > 1 , these terms are decreasing in absolute value beyond i = 2 , so we can upper bound the log with just the first two terms: log ( 1 + 1 √ d ) ≥ 1 √ d -1 2 d .

H Attention Window Captures Appropriate Directions

In this section we prove Theorem 4.4, which entails showing that if the Lipschitzness of the function class is zero in some directions, one-layer self-attention learns to ignore these directions when the function class consists of linear functions. First, we give a brief sketch of the proof.

H.1 Proof Sketch

We briefly sketch the proof of Theorem 4.4. WLOG we write M = BF + B ⊥ G where F := B ⊤ M and G := B ⊤ ⊥ M . Lemma H.2 leverages the rotational symmetry of F B in col ( B ) to show that the loss is minimized over ( F , G ) at ( F , G ) = ( c B ⊤ , c ′ B ⊥ ) for some constants c, c ′ . It remains to show that L ( c BB ⊤ + c ′ B ⊥ B ⊤ ⊥ ) > L ( c BB ⊤ ) whenever c ′ is nonzero. Intuitively, if the attention estimator incorporates the closeness of B ⊤ ⊥ x i and B ⊤ ⊥ x n +1 into its weighting scheme via nonzero Q , this may improperly up- or down-weight f ( x i ) , since projections of x i onto col ( B ⊥ ) do not carry any information about the closeness of f ( x i ) and f ( x n +1 ) .

̸

Using this intuition, we show that for any fixed c ′ and { v i } i such that c ′ v ⊤ i v n +1 = v ⊤ i ′ Qv n +1 for some i, i ′ , the attention estimator improperly up-weights f ( x 1 ) , where 1 ∈ arg max i c ′ v ⊤ i v n +1 WLOG. In particular, the version of the pretraining population loss (ICL) with expectation over a , { u i } i and { ϵ i } i is reduced by reducing c ′ v ⊤ 1 v n +1 . The only way to ensure all { c ′ v ⊤ i v n +1 } i are equal for all instances of { v i } i is to set c ′ = 0 , so this c ′ must be optimal.

To show that reducing c ′ v ⊤ 1 v n +1 reduces the loss with fixed { v i } i , we define α i := e c v c ′ v ⊤ i v n +1 for all i ∈ [ n ] and show the loss' partial derivative with respect to α 1 is positive, i.e.

̸

This requires a careful symmetry-based argument as the expectation over { u i } i cannot be evaluated in closed-form. To overcome this, we fix all u i but u 1 and one other u i ′ = u n +1 with α i ′ < α 1 . We show the expectation over ( u 1 , u i ′ ) can be written as an integral over ( y 1 , y 2 ) ∈ S k -1 × S k -1 of a sum of the derivatives at each of the four assignments of ( u 1 , u i ′ ) to ( y 1 , y 2 ) , and show that this sum is always positive. Intuitively, any 'bad' assignment for which increasing α 1 reduces the loss is outweighed by the other assignments, which favor smaller α 1 . For example, if y 1 = u n +1 = y 2 , and u 1 = y 1 and u i ′ = y 2 , we observe from (21) that increasing α 1 can reduce the loss. However, the cumulative increase in the loss on the other three assignments due to increasing α 1 is always greater.

H.2 Full Proof

We now prove Theorem 4.4 in full detail.

Lemma H.1. For any u ∈ S k -1 and α 1 , . . . , α n such that min i α i > 0 , and any c a , c u ∈ R \ { 0 } , define

Then for any δ > 0 , 0 / ∈ arg min 0 ≤ c ≤ δ J ( c ) .

Proof. We show that there exists some arbitrarily small ϵ > 0 such that J ( ϵ ) < J (0) by showing dJ ( c ) dc ∣ ∣ c =0 < 0 . We have dJ ( c ) dc

Setting c = 0 results in

̸

where (22) follows since E [ u i ] = 0 k , (23) similarly follows since odd moments of uniform random variables on the hypersphere are zero, (24) follows by the i.i.d.-ness of the u i 's, (25) follows since E [ u i u ⊤ i ] = 1 k I k and u ⊤ u = 1 , and (26) follows since min i α i > 0 . This completes the proof.

̸

Lemma H.2. Consider any B ∈ O d × k and resulting function class F lin B . Consider the training population loss L defined in (ICL) , and tasks drawn from D ( F lin B ) such that E a [ aa ⊤ ] = c 2 a I k for some c a = 0 and let M := M ⊤ K M Q be optimized over the domain M ˆ c := { M ∈ R d × d : M = M ⊤ , ∥ B ⊤ MB ∥ 2 ≤ ˆ c c 2 u } for any ˆ c > 0 . Then any

satisfies M ∗ = c ∗ 1 BB T + c ∗ 2 B ⊥ B ⊤ ⊥ for some c ∗ 1 : | c ∗ 1 | ∈ (0 , ˆ c c 2 u ] .

Proof. Without loss of generality (WLOG), we can decompose M = BF + B ⊥ G where F := B ⊤ M and G := B ⊤ ⊥ M . Recall that for each i ∈ [ n +1] , x i = c u Bu i + c v B ⊥ v i . Thus, for each i ∈ [ n ] ,

we have

where, for each i ∈ [ n ] , α i := e c v v ⊤ i Gx n +1 . For ease of notation, denote x = x n +1 .

We start by expanding the square and using the linearity of the expectation to re-write the population loss as:

WLOG we can write Fx = R ( Fx ) u ∥ Fx ∥ 2 for some rotation matrix R ( Fx ) ∈ O k × k . Denote C 1 ( Fx ) := ∥ Fx ∥ 2 . Then we have

where (30) follows since R ( Fx ) R ( Fx ) ⊤ = I k and (31) follows since the distribution of u i is the same as the distribution of R ( Fx ) ⊤ u i for any rotation R ( Fx ) ⊤ . Define

for any F ∈ R k × d . We have L signal ( M ) = c 2 a c 2 u E u , v [ g ( F , u , v )] , and Note that if F ′ = c B ⊤ , then R F ′ x = I k and C 1 ( F ′ x ) = c u c . Thus,

Define ˆ u := E { u i } , { v i } [ ∑ n i =1 u i e cuC 1 ( Fx ) u ⊤ i u α i ∑ n i =1 e cuC 1 ( Fx ) u ⊤ i u α i ] and WLOG write u i = p u i + q u i , where p u i := uu ⊤ u i and q u i := ( I k -uu ⊤ ) u i . Note that for any u i = p u i + q u i , u ′ i := p u i -q u i occurs with equal probability, and flipping q u i does not change any exponent or α i in (32). Thus

where ˜ c := E { u i } , { v i } [ ∑ n i =1 u ⊤ i u e cuC 1 ( Fx ) u ⊤ i u α i ∑ n i =1 e cuC 1 ( Fx ) u ⊤ i u α i ] . Note that for any u i , -u i occurs with equal probability, so

̸

̸

̸

Since α i > 0 and ¯ c u,v > 0 by definition, e c u C 1 ( Fx ) u ⊤ i u α i is monotonically increasing in u ⊤ i u . Also, f ( x ) := x x + c is monotonically increasing for x > 0 for all c > 0 . Thus we have that

̸

̸

̸

̸

and thereby ˜ c > 0 . Therefore, arg max u ′ ∈ S k -1 ( u ′ ) ⊤ ˆ u = u , in particular u ⊤ ˆ u > u ⊤ R ( Fx ) ⊤ ˆ u whenever R ( Fx ) u = u , so (32) is strictly positive if R ( Fx ) u = u . Thus, for any u , v such that R ( Fx ) u = u , g ( F , u , v ) > g ( C 1 ( Fx ) c u B ⊤ , u , v ) . Also, for any u , v such that R ( Fx ) u = u , g ( F , u , v ) = g ( C 1 ( Fx ) c u B ⊤ , u , v ) .

Next we need to account for ˜ L noise ( M , x ) . Again writing Fx = R ( Fx ) u ∥ Fx ∥ 2 and C 1 ( Fx ) = ∥ Fx ∥ 2 and using the rotational invariance of u i , we obtain

where (36) follows using the rotational invariance of u i . So, returning to (29), we have

̸

where (37) is strict if R ( Fx ) u = u for some u , v , which is equivalent to saying that F / ∈ { c ′ B ⊤ , c ′ > 0 } .

Next, recall that we have defined α i := e c v v ⊤ i Gx . Using a similar argument as earlier, by the rotational invariance of the v i 's, for any fixed x , we can write α i = e c v C 2 ( Gx ) v ⊤ i e 1 where C 2 ( Gx ) := ∥ Gx ∥ 2 and e 1 is the first standard basis vector.

Next, for c 1 , c 2 ≥ 0 and some fixed u , v , define

̸

̸

and let

Since H does not vary with v , we have ( C ∗ 1 ( u , v ) , C ∗ 2 ( u , v )) = ( C ∗ 1 ( u ) , C ∗ 2 ( u )) WLOG. In fact, H does not vary with u either, due to the rotational invariance of the u i 's. So, we have ( C ∗ 1 ( u , v ) , C ∗ 2 ( u , v )) = ( C ∗ 1 , C ∗ 2 ) WLOG, i.e. there is a single pair ( C ∗ 1 , C ∗ 2 ) that minimizes H ( u , v , c 1 , c 2 ) over c 1 , c 2 for all u ∈ S k -1 and v ∈ S d -k -1 .

Thus F ∗ = C ∗ 1 B ⊤ and G ∗ satisfies ∥ G ∗ x ∥ = C ∗ 2 for all x , which implies, using also that M is symmetric, that G ∗ = C ∗ 2 B ⊤ ⊥ .

Lemma H.3. Consider any α := [ α 1 , . . . , α n ] such that α 1 = max i α i and α 1 > min i α i > 0 . Further, let c ∈ (0 , 2] . Define

Then

Proof. We first compute ∂H signal ( u , α ) ∂α 1 . Using the linearity of the expectation and the quotient rule we obtain:

where

Note that terms with i = 1 do not appear in (41). We analyze S i, 1 + S i,i and each S i,j , j / ∈ { 1 , i } separately, and will ultimately show that each of these terms is positive. We start with the latter case as it is easier to handle. For j / ∈ { 1 , i } , we have

where the latter two terms in (42) are zero by the same argument as in (33): flipping the component of either u 1 or u i perpendicular to u does not change any of the values in any exponent, and each flip occurs with equal probability. Next, note that if α i = α 1 ,

thus S i,j = 0 . Otherwise, α i < α 1 by definition of α 1 , and there must be some such α i , since if not, there would be some c ′ ∈ R + such that α = c ′ α ∗ . For the case α i < α 1 , we use a symmetry argument to show that S i,j > 0 .

First we define additional notations. Let ¯ U 1 ,i := { u i ′ } i ′ ∈ [ n ] \{ 1 ,i } , and for any ( a, b ) ∈ [ -1 , 1] 2 , define

̸

In particular, for any a ∈ [ -1 , 1] , define p a := P u 1 [ u ⊤ 1 u = a ] . Since u 1 and u i are i.i.d., we have P u 1 , u i [ u ⊤ 1 u = a, u ⊤ i u = b ] = P u 1 , u i [ u ⊤ 1 u = b, u ⊤ i u = a ] = p a p b for any ( a, b ) ∈ [ -1 , 1] 2 Thus, by the law of total expectation we have

Next we show that for any instance of a, b and ¯ U 1 ,i , f a,b ( ¯ U 1 ,i ) + f b,a ( ¯ U 1 ,i ) is positive. We have:

̸

̸

with equality only if a = b or u j = u , since u ⊤ j u ≤ 1 with equality only if u j = u , and

̸

due to α 1 > α i and α i ′ > 0 for all i ′ . So we have S i,j > 0 .

Next we analyze S i, 1 + S i,i . In these cases we cannot immediately drop the components of u 1 and u i that are perpendicular to u . We have:

̸

Now we can split u ⊤ i u 1 into the product of the components of u i , u 1 in the direction u and the product of their components in the perpendicular subspace as before. Doing so yields

Next, define

We argue similarly as in the previous case, except that here we must include additional terms.

where

We show that for any ( a, b ) ∈ [ -1 , 1] 2 and any ¯ U 1 ,i , G a,b ( ¯ U 1 ,i ) is positive, which implies that S i, 1 + S i,i is positive by (46).

First, note that if b = a for any a ∈ [ -1 , 1] and ¯ U 1 ,i , we have

Using this notation, we can rewrite

Therefore,

Note that e c ( a + b ) d -3 a,b d -3 b,a > 0 , so it remains to show that the term inside the parentheses is positive. This term can be rearranged as:

̸

First we show that T 1 is positive by analyzing γ a,b and γ b,a . For any ( a, b ) ∈ [ -1 , 1] 2 such that a = b ,

̸

with equality holding if and only if a = -1 . If a = -1 , we have γ a,b = 1 + b -1 -b = 0 for all b ∈ [ -1 , 1] . Otherwise, (52) shows that γ a,b is strictly decreasing with b , so it is minimized over b ∈ [ -1 , 1] at b = 1 . When b = 1 , we have γ a,b = 1 -a + a -1 = 0 for all a . So, γ a,b ≥ 0 with equality holding if and only if a = -1 or b = 1 . Note that by symmetry, this implies γ b,a ≥ 0 with equality holding if and only if a = 1 or b = -1 . So, we can have both γ a,b = 0 and γ b,a = 0 if and only if a = b = -1 or a = b = 1 . However, we have a = b , so at least one of γ a,b and γ b,a are strictly positive, and T 1 is strictly positive (using also that α 1 > α i ).

We next show that T 2 is positive. Observe that

since α 1 > α i , so it remains to show

where

We first show the forward direction, namely b > a = ⇒ e ca γ b,a -e cb γ a,b > 0 .

̸

since each term inside the expectation is nonnegative, as a 2 ≤ 1 and α 1 > α i . Note that this implies G a,b ≥ 0 when a = b , so WLOG we consider b = a for the remainder of the proof. Now we focus on showing (61). Throughout, we will make use of the notation

which represents the cube root of the denominator in all terms when u ⊤ 1 u = a and u ⊤ i u = b , and

Note that if b = a , e ca γ b,a -e cb γ a,b = 0 . So, if we can show that for any fixed a , e ca γ b,a -e cb γ a,b is increasing with b as long as b ≥ a , then we will have e ca γ b,a -e cb γ a,b > 0 for b > a . To show e ca γ b,a -e cb γ a,b is increasing, we take its partial derivative with respect to b :

We would like to show that the RHS of (56) is nonnegative. To do so, we show that its partial derivative with respect to a is positive, so it achieves minimum value at a = -1 , at which point the value is positive. We have:

where q ( x ) := e cx (1 + cx -c ) . Note that q ( x ) is monotonically increasing in x ∈ [ -1 , 1] ; to see this, observe that

where the inequality follows since c ∈ (0 , 2] and x ∈ [ -1 , 1] . Therefore, since b > a , we have q ( b ) -q ( a ) ≥ 0 and ∂ ∂a ( ∂ ∂b ( e ca γ b,a -e cb γ a,b )) ≥ 0 from (57). As a result, ∂ ∂b ( e ca γ b,a -e cb γ a,b ) achieves minimum value at a = -1 . At this point, using (56) we have

This implies that the minimum value of e ca γ b,a -e cb γ a,b over b ∈ [ a, 1] is achieved at b = a , and we know this value is zero, so we have that e ca γ b,a -e cb γ a,b > 0 when b -a .

To show the backward direction of (54), namely e ca γ b,a -e cb γ a,b > 0 = ⇒ b > a , note that the converse, namely a > b = ⇒ e ca γ b,a -e cb γ a,b < 0 , follows by the same argument as above with a and b swapped. Therefore, we have T 2 > 0 as desired.

Lemma H.4. Consider any α := [ α 1 , α 2 ] such that α 1 > α 2 > 0 . Further, let c ∈ (0 , 1] . Define

Then

Proof. We have

Since n = 2 , we have

Now, we have

where

̸

̸

̸

̸

̸

We will show that for any ( a, b ) ∈ [ -1 , 1] 2 and ( p a , p b ) ∈ [0 , 1] 2 , H a,b is positive, which implies that N i is positive by (59). To do this, assuming h a,a is nonnegative for any a , it is sufficient to show

since this implies h a,b + h b,a > -√ h a,a h b,b and thus, from (60),

Before showing (61), we need to confirm that h a,a is not negative for all a ∈ [ -1 , 1] . We have

since each term inside the expectation is nonnegative, as α 1 > α 2 . Note that this implies H a,b ≥ 0 when a = b , so WLOG we consider a > b for the remainder of the proof. Note that

Using this, we have

To show that ˜ H a,b is positive, we need to show that P is positive. Without loss of generality we can consider α 1 = 1 and α 2 ∈ (0 , 1) by dividing the numerator and denominator of H noise by α 2 1 . Thus, for the remainder of the proof we treat α 1 as 1 and write α := α 2 for ease of notation. Using this notation we can expand P as follows:

Recall that 1 -α > 0 , so we need to show that the sum of the remaining terms is positive. These terms can be written as a polynomial in y := e c ( a -b ) as follows:

We know that y 6 > y 5 > · · · > 1 since a > b . We also have that α < 1 . Using these facts we next show that the sum of the third and smaller-order terms in the RHS of (66) is positive.

Next we show that the sum of the sixth-, fifth-, and fourth-order terms is positive. Let a 6 := α 3 , a 5 := α + α 2 +6 α 3 + α 4 + α 5 , and a 4 := 1 + 7 α +7 α 2 +17 α 3 +7 α 4 +7 α 5 + α 6 , so the sum of the sixth-, fifth-, and fourth-order terms is y 6 a 6 -y 5 a 5 + y 4 a 4 . Note that 32 a 6 < a 4 since α < 1 , and

where (68) follows by minimizing the terms inside the parentheses over y . Thus, we have ˜ H a,b > 0 , which completes the proof.

Now we can finally prove Theorem 4.4. We prove a slightly stronger result, formally stated as follows. Theorem H.5. Consider any B ∈ O d × k and the corresponding function class F lin B as defined in (4.2) . Suppose tasks are drawn from D ( F lin B ) and Assumption 4.3 holds. Recall the pretraining population loss:

Consider two cases:

Case 1: σ = 0 , n > 1 . Then define C p := 2 .

Case 2: σ > 0 , n = 2 . Then define C p := 1 .

Then in each case, among all M ∈ M := { M ∈ R d × d : M = M , ∥ B ⊤ MB ∥ 2 ≤ C p c 2 u } , any minimizer M ∗ of (70) satisfies M ∗ = c BB ⊤ for some c ∈ (0 , C p c 2 u ] .

̸

Proof. From Lemma H.2, we have M ∗ = c p BB ⊤ + ˜ cmathbfB ⊥ B ⊤ ⊥ for some ˜ c ∈ R and some c p ∈ (0 , C p c 2 u ] , where C p = 2 in Case 1 and C p = 1 in Case 2. Suppose that ˜ c = 0 . Then it remains to show that L ( c p BB ⊤ +˜ c B ⊥ B ⊤ ⊥ ) > L ( c p BB ⊤ ) .

We start by establishing the same notations as in the proof of Lemma H.2. For each i ∈ [ n +1] , x i = c u B u i + c v B ⊥ v i . Thus, for each i ∈ [ n ] , we have

where, for each i ∈ [ n ] , α i := e c 2 v ˜ c v ⊤ i v n +1 . For ease of notation, denote x = x n +1 , u := u n +1 and c = c p c 2 u . Also, note that for any x i , f ( x i ) = a ⊤ B ⊤ x i = c u a ⊤ u i , and that drawing f ∼ D ( F lin B ) is equivalent to drawing a ∼ D a for some distribution D a over R k such that E a ∼ D a [ aa T ] = c 2 a I k . Using this, we have:

where α := [ α 1 , . . . , α n ] and

H ( u , α )

Define α ∗ = [1 , . . . , 1] ∈ R n . We proceed by showing that for any u ∈ S d -1 , all α ∈ R n + satisfy

̸

(ii) if α = c ′ α ∗ for any c ′ ∈ R + , then H ( u , α ) > H ( u, α ∗ )

This implies L ( c p BB ⊤ +˜ c B ⊥ B ⊤ ⊥ ) > L ( c p BB ⊤ ) , since

which implies that ˜ c = 0 is the unique argument that achieves the minimal value of L ( c p BB ⊤ + ˜ c B ⊥ B ⊤ ⊥ ) over ˜ c ∈ R (and this value is E u [ H ( u , α ∗ )] ).

Proving ( i ) is trivial as it can be easily checked that H ( u , α ) = H ( u , c ′ α ) for all u ∈ S d -1 , α ∈ R n + , and c ′ ∈ R + .

̸

Proving ( ii ) is more involved. Consider any α = c ′ α ∗ for any c ′ ∈ R + . WLOGlet 1 ∈ arg max i α i . We show that the partial derivative of H ( u , α ) with respect to α 1 is strictly positive, which means that H ( u , α ) can be reduced by reducing α 1 by some ϵ > 0 . We can repeat this argument, repeatedly reducing max i α i at each step and thereby reducing the loss, until we reach an α ′ satisfying α ′ = c ′ α ∗ . Since the loss is reduced at each step, we have that H ( u , α ) > H ( u , α ∗ ) .

To show that the partial derivative of H ( u , α ) with respect to α 1 is strictly positive, we decompose ∂H ( u , α ) ∂α 1 = ∂H signal ( u , α ) ∂α 1 + ∂H noise ( u , α ) ∂α 1 , where

By Lemma H.3, we have ∂H signal ( u , α ) ∂α 1 > 0 . If σ = 0 we are done, otherwise we have n = 2 and ∂H noise ( u , α ) ∂α 1 > 0 by Lemma H.4. This completes the proof.

I Additional Lemmas

Lemma I.1. Consider a continuous unimodal function f . Then we have

Proof. Let T denote the point that achieves the maximum of f . Then we know that f ( t ) ≥ f ( ⌊ t ⌋ ) for t < T , while f ( t ) ≥ f ( ⌈ t ⌉ ) for t > T . This means ∫ i i -1 f ( t ) dt ≤ f ( i ) ≤ ∫ i +1 i f ( t ) dt for t ≤ ⌊ T ⌋ and ∫ i i -1 f ( t ) dt ≥ f ( i ) ≥ ∫ i +1 i f ( t ) dt for t ≥ ⌈ T ⌉ So

Similarly we have

Lemma I.2. If f and g are nonnegative measurable real functions, then

where f ∗ , g ∗ are the symmetric decreasing rearrangements of f and g .

Proof.

Lemma I.3. Suppose { a i } , { b i } are sorted the same way, a i > a j ⇐⇒ b i > b j . Then we have

Proof. Cross multiplying and expanding, we have

The last of which follows from the rearrangement inequality [67].

J Additional Experiments and Details

All experiments were run in Google Colab in a CPU runtime. We used a random seed of 0 in all cases. All training was executed in PyTorch with the Adam optimizer. We tuned learning rates in { 10 -3 , 10 -2 , 10 -1 } separately for linear and softmax attention, and we initialized M K and M Q by setting each to 0 . 001 I d , and tie the weights of M K and M Q to speed up training.

Figure 1. The upper row depicts our functions, which increase in Lipschitzness from left to right. The black curve depicts the ground truth, while the gray dots depict the noisy training samples. The shaded region represents the attention window. The middle row depicts the attention weights for softmax and linear attention. We remark that the softmax is able to adapt to the Lipschitzness while linear is not. The bottom row depicts the ICL error as a function of the context length n for Linear and ReLU pretraining using Linear and Softmax attention. That is, at each iteration, a context is drawn from a non-linear regression (defined below) consisting of a randomly phase shifted cosine function. The ICL task is to predict the function value at a randomly chosen query on the unit circle. Each point in the plot depicts the ICL error of a pretrained attention unit (using softmax (blue) or linear (orange) activation) at the end of 15000 iterations with learning rate 10 -3 . We use d = 2 and a distribution D ( F ν, hills ) . Here we define

and a distribution D ( F ν, hills ) is induced by drawing b uniformly from [ -π, π ] . We use ν = 0 , 1 . 5 , 6 for the left, middle and right plots in the bottom row, respectively.

Figures 3, 4, 5. In all cases, we use an exponentially decaying learning rate schedule with factor 0.999. In Figures 3 and 5 we use initial learning rate 0.1 and in Figure 4 we use an initial learning rate 0.01. Moreover, in all cases besides those with varying n in Figure 4, we compute gradients with respect to the ICL loss evaluated on N := ⌊ √ n ⌋ query samples per task (that is, each context input to the attention unit has n + N samples, of which n are labeled, and the other N labels are inferred). When n varies in Figure 4, we use N = 1 . In Figure 5 we show smoothed test ICL errors with smoothing rate 0.01.

J.1 Low-Rank Experiments

Due to our results in Section 3 showing that softmax attention can learn an appropriate attention window scale when pretrained on nonlinear tasks, we hypothesize that it can also learn the appropriate directions during pretraining on nonlinear tasks. To test this, we consider tasks drawn from lowrank versions of affine, quadratic and cosine function classes, in particular: F aff B := { f : f ( x ) = a ⊤ B ⊤ x +2 , a ∈ S k -1 } , F 2 B := { f : f ( x ) = ( a ⊤ B ⊤ x ) 2 , a ∈ S k -1 } and F cos B := { f : f ( x ) = cos ( 4 a ⊤ B ⊤ x ) , a ∈ S k -1 } . Each task distribution D ( F aff B ) , D ( F 2 B ) , D ( F cos B ) is induced by drawing a ∼ U k . We train M K and M Q with Adam with learning rate tuned separately for softmax and linear attention. We set d = 10 , k = 2 , n = 50 , and σ = 0 . 01 . We draw { x i } n +1 i =1 i.i.d. from a non-uniform distribution on S d -1 for each task, and draw one task per training iteration. We draw B randomly at the start of each trial, and repeat each trial 5 times and plots means and standard deviations over the 5 trials. We capture the extent to which the learned M = M ⊤ K M Q recovers col ( B ) via the metric ρ ( M , B ) := ∥ B ⊤ ⊥ MB ⊥ ∥ 2 σ min ( B ⊤ MB ) , where σ min ( A ) is the minimum singular value of A . For test error, we compute the average squared error on 500 random tasks drawn from the same distribution as the (pre)training tasks. Please see Appendix J for more details.

We randomly generate B on each trial by first sampling each element of ˆ B i.i.d. from the standard normal distribution, then take its QR decomposition to obtain B . To draw the covariates, we draw a random matrix ˜ J ∈ R d × d by sampling each element i.i.d. from the standard normal distribution. Then, we compute J = ( ˜ J ⊤ ˜ J ) 1 / 2 . Then we draw ˜ x i ∼ N ( 0 d , I d ) and set x i = J ˜ x i ∥ J ˜ x i ∥ .

Results. Figure 9 shows that softmax attention recovers the low-rank structure when tasks are drawn from each of the three function classes, which leads to test error improving with the quality of the learned subspace. In contrast, linear attention does not learn any meaningful structure in these cases.

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: Our abstract is consistent with our introduction. In the introduction, we point to the places in the paper in which we substantiate our claims.

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 provide a discussion in Section 5.

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: Assumptions are specified before all theorem statements.

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 provide details in Sections 3.2 and J.

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: Please see supplementary material.

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 provide details in Sections 3.2 and J.

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: Please see results in Sections 3.2 and J.

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: Please see Section J.

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 have read the Code of Ethics and ensured conformity.

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: [NA]

Justification: This is primarily an analysis of an already existing algorithm.

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: [NA]

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: [NA]

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: [NA]

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: [NA]

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: [NA]

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:\nThis study examines the impact of softmax attention on ICL regression, advancing beyond the typical linear treatment of the topic. In particular, the authors find that 1) softmax attention leads the model to adapt to the target function's Lipschitz constant and 2) enables the model to recover low dimensional structure in the target.\n\nSTRENGTHS:\nThe paper is very well-written and easy to read. It's great to see work that pushes beyond the usual linear treatment of attention, and explicitly considers the role of the softmax nonlinearity. The connection to kernel smoothing is very neat, and the resulting interpretation of softmax attention as an adaptive process sensitive to the target's Lipschitz constant is very insightful. I was particularly surprised that Lipschitz adaptation is both necessary and sufficient to illustrate the role of softmax attention in ICL regression, and that swapping function classes did not matter so long as the Lipschitz constant remained the same. Very cool work overall. Well done!\n\nWEAKNESSES:\nI found no substantive weaknesses in the analysis, but do have some follow-up question listed below. As with any theory paper, there's always more that can be done, but I think the level of work demonstrated in this paper is more than sufficient to merit a strong accept.\n\nQUESTIONS:\n- Could your results transfer in any way to an ICL classification setting, for instance like the one in [Reddy 2024](https://arxiv.org/abs/2312.03002) among others? Your Lipschitz adaptation argument, particularly as it controls a \"window size\" over inputs, sounds obliquely related to the way softmax attention is sometimes discussed as implementing \"selection\" or \"copy\" operations over context examples in ICL classification. I'm curious if there could be a deeper connection here?\n- When you train models in practice, do you find that the QKV matrices of the final model become like the M-parameterization of the matrices you highlight in eq. 4?\n\nLIMITATIONS:\nThe authors adequately address their limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper analyzes how softmax attention learns to perform in-context learning (ICL) through pretraining. The authors show that softmax attention adapts its \"attention window\" based on the Lipschitzness and noise characteristics of the pretraining tasks. They provide theoretical analysis for affine and ReLU-based function classes, demonstrating that softmax attention learns an optimal trade-off between bias and variance. The paper also explores how softmax attention can recover low-dimensional structure in the input space. Experiments are conducted to validate the theoretical findings.\n\nSTRENGTHS:\n- The paper addresses an important and timely topic in understanding the mechanisms behind in-context learning in transformer models. This is a crucial area of research given the widespread adoption of large language models. The paper makes a novel connection between softmax attention and nearest neighbor regression, providing an intuitive interpretation of how ICL works in this setting.\n- The analysis is clean and the main message is clear: softmax activation plays a critical role in enabling effective in-context learning. For example, Theorem 3.4 provides concrete bounds on how the attention window scales with task Lipschitzness, noise level, and context size. Also, I like Theorem 3.5, whose test data are any arbitrary L−Lipschitz task, which is pretty general ICL setting. \n- The experimental results generally support the theoretical analysis. Figure 3 nicely illustrates how the spectral norm of M (representing the inverse of the attention window size) increases with the Lipschitz constant L, aligning with the theory.\n\nWEAKNESSES:\n- There's a potential gap between the paper's main message and practical ICL scenarios. The analysis is based on a nearest-neighbor interpretation, but in real-world applications, LLMs often perform well even when there are no \"close\" examples in the context. \n- Some of the experimental writing and presentation could be improved. Figures 4 and 5 are particularly hard to follow without careful reading of the appendix. It would be helpful to have more detailed captions and clearer explanations in the main text.\n- The theoretical analysis is limited to relatively simple function classes (affine and ReLU-based). While this provides valuable insights, it's not clear how well these results generalize to the more complex functions learned by large language models in practice.\n\nQUESTIONS:\n- How do you think your analysis might extend to more complex, hierarchical function classes that might better represent the capabilities of large language models? How can it be generally applied when there is no neighbor data?\n- Have you considered how your findings might relate to or explain the emergent capabilities observed in large language models as they scale?\n\nLIMITATIONS:\nThe authors acknowledge some limitations in Section 5, including that their model only considers the output of a single layer of attention and that establishing a mathematical framework for priming effects in LLMs remains an open challenge. These are important limitations to note, as they highlight the gap between this analysis and the full complexity of modern language models.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores how softmax attention in transformer models enables in-context learning (ICL), where a model can adapt to solve new tasks using only a few input examples without additional training. The authors focus specifically on regression tasks, where the model must predict a continuous value given some input features. They show that during pretraining on a distribution of ICL regression tasks, softmax attention learns to implement a nearest neighbors predictor that is adapted to properties of the pretraining task distribution.\n\n The key insight is that softmax attention learns an \"attention window\" - a neighborhood around each input query point that determines which other input points influence the prediction. The size and shape of this attention window adapts based on properties of the pretraining tasks, specifically their Lipschitzness (how quickly function values can change) and the amount of label noise. Importantly, the authors demonstrate that learning this adapted attention window is crucial for generalization. The authors also prove that softmax attention can learn to project inputs onto a relevant low-dimensional subspace when the pretraining tasks depend only on projections of the inputs onto this subspace. \n\nTo validate their theoretical results, the authors conduct experiments on synthetic regression tasks with varying Lipschitzness, noise levels, and input dimensionality. Empirically, the authors demonstrate that softmax attention learns appropriate attention window scales across a range of nonlinear function classes, including ReLU networks and trigonometric functions.\n\nSTRENGTHS:\nA first strength is that the paper is not only well written but also excellently presented, which helps in understanding its mathematical content and putting it into a better light. The clarity of exposition is therefore a first great point. In terms of originality, it starts off with the fairly known/commonplace insight that there exists a connection between self-attention and Nadaraya-Watson kernel regression, thus establishing that learning the bandwidth of that estimator across multiple tasks is a necessity for ICL. In this sense and if it stopped there, the contribution wouldn't be particularly novel, as this is intuitive (if thinking of self-attention as learning a summary of the autocovariance function of data and therefore of its characteristic length) and fairly well understood in the literature already (see the cited works of Tosatto et al). One can argue - as in the paper - that Tosatto et al only give an upper bound on the bias, and this work provides a lower bound as well.\n\nHowever where the paper takes off in my opinion is when these arguments move away from purely the 'frequency cutoff / Lipschitzness' length-related realm, to then move into the *directional*, via using concentration arguments on the hypersphere. This is Theorem 4.4 which formalizes the intuition that ICL in transformers identifies low-dimensional subspaces shared by training tasks, an argument more typically found in the analysis of (non-contrastive) self supervised learning. The mathematical method of proof is elegant as well. The authors derive novel concentration inequalities for functionals of points uniformly distributed on high-dimensional spheres, and use a careful symmetry argument to show that any non-zero component in the orthogonal subspace increases the loss. Overall this represents a standout technical contribution well worthy of publication in my view.\n\nWEAKNESSES:\n- A small weakness in presentation is I believe that Theorem 4.4 should be emphasized, as to my knowledge this is the more novel part of the contribution. \n\n- Similarly, the theoretical guarantees are provided for relatively simple function classes (affine and ReLU-based), which may not represent the full range of tasks where ICL is effective.\n\n- The limited scope of experiments is understandable given theoretical assumptions (single-layer) but also important enough that it becomes a weakness, IMHO.\n\n- Finally, the paper could benefit from a more extensive comparison to other theoretical frameworks for understanding ICL.\n\nQUESTIONS:\nWould there be (not extremely involved) ways of moving towards more realistic settings ? i.e. moving to infinite width settings ? Would classification tasks with a cross-entropy loss be somewhat tractable if making the right Gaussian (concentration) assumptions ?\n\nLIMITATIONS:\nThere are some obvious, if hard to tackle, limitations in this work:\n\na. Simplicity of setting: The analysis focuses on single-layer models and synthetic tasks, which may limit its immediate applicability to more complex real-world scenarios.\n\nb. Gap to large language models: While insightful, the work doesn't fully explain emergent ICL in large language models trained on natural data.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies in-context learning (ICL) of one-layer attention-only transformers in a regression task. The paper argues that the product of query and key projection matrix is associated with the Liptchitzness of input data. The notion of attention windows is introduced based on this. The paper shows that attention windows will adapt to the Lipchitness and the noise level of data and perform dimensional reduction. Experiments are conducted to support theoretical claims.\n\nSTRENGTHS:\nThe paper provides a rigorous approach to the problem both with proof and empirical simulations.\n\nWEAKNESSES:\n- Although explaining ICL is important, the setting seems to be less realistic (see question 1).\n- The approach of understanding ICL through the lens of function Lipschitzness is somewhat limiting. For example, Figure 5 demonstrates a decline in generalization error, which is primarily attributed to the function's structure. While the Lipschitz property of a function is indeed a valuable aspect in studying ICL, it doesn't provide a comprehensive explanation or solution for ICL as a whole.\n\nQUESTIONS:\n- The assumption in Equation 4 can be understood as the optimal weights of values $W_V$ will ignore covariates $x$ while the attention will only consider covariates to decide where to attend in the sequence while disregarding predictor $f$. I wonder if this makes sense for hidden states of LLMs. In particular, there is no separation between covariates and predictors in LLMS.  Also, why is the optimal $W_V$ in the proof of Lemma B.1 assumed to have first columns to be zeros (in between Line 577-578)?\n\n\n- I agree that $x^\\top M y$ can be understood as the distance between $x$ and $y$. It might not straightforward to change $x^\\top M y$ into $||x - y||$ like used in Appendix. Can you clarify this more on how to obtain Equation (8) in Appendix from Equation (ICL) in the main text.\n\nLIMITATIONS:\nPlease see above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-17"}